"""
GaugeGroup and derived objects, used primarily in gauge optimization
"""
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************

import numpy as _np

from pygsti.modelmembers import operations as _op

class GaugeGroup(object):
    """
    A parameterized set (ideally a group) of gauge transformations.

    Specifies the "optimization space" explored by gauge optimization
    algorithms.  This base class is used to define the common interface of all
    types of gauge "groups" (even though they need not be groups in the
    mathematical sense).

    Parameters
    ----------
    name : str
        A name for this group - used for reporting what type of
        gauge optimization was performed.
    """

    def __init__(self, name):
        """
        Creates a new gauge group object

        Parameters
        ----------
        name : str
            A name for this group - used for reporting what type of
            gauge optimization was performed.
        """
        self.name = name

    @property
    def num_params(self):
        """
        Return the number of parameters (degrees of freedom) of this gauge group..

        Returns
        -------
        int
        """
        return 0

    def compute_element(self, param_vec):
        """
        Retrieve the element of this group corresponding to `param_vec`

        Parameters
        ----------
        param_vec : numpy.ndarray
            A 1D array of length :method:`num_params`.

        Returns
        -------
        GaugeGroupElement
        """
        return GaugeGroupElement()

    @property
    def initial_params(self):
        """
        Return a good (or standard) starting parameter vector, used to initialize a gauge optimization.

        Returns
        -------
        numpy.ndarray
            A 1D array of length :method:`num_params`.
        """
        return _np.array([], 'd')

class GaugeGroupElement(object):
    """
    The element of a :class:`GaugeGroup`, which represents a single gauge transformation.
    """

    def __init__(self):
        """Creates a new GaugeGroupElement"""
        pass

    @property
    def transform_matrix(self):
        """
        The gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        return None

    @property
    def transform_matrix_inverse(self):
        """
        The inverse of the gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        return None

    def deriv_wrt_params(self, wrt_filter=None):
        """
        Computes the derivative of the gauge group at this element.

        That is, the derivative of a general element with respect to the gauge
        group's parameters, evaluated at this element.

        Parameters
        ----------
        wrt_filter : list or numpy.ndarray, optional
            Indices of the gauge group parameters to differentiate with respect to.
            If None, differentiation is performed with respect to all the group's parameters.

        Returns
        -------
        numpy.ndarray
        """
        return None

    def to_vector(self):
        """
        Get the parameter vector corresponding to this transform.

        Returns
        -------
        numpy.ndarray
        """
        return _np.array([], 'd')

    def from_vector(self, v):
        """
        Reinitialize this `GaugeGroupElement` using the the parameter vector `v`.

        Parameters
        ----------
        v : numpy.ndarray
            A 1D array of length :method:`num_params`

        Returns
        -------
        None
        """
        pass

    @property
    def num_params(self):
        """
        Return the number of parameters of this gauge group element.

        (This is equivalent to the number of parameters of the parent gauge group.)

        Returns
        -------
        int
        """
        return 0

    def inverse(self):
        """
        Creates a gauge group element that performs the inverse of this element.

        Returns
        -------
        InverseGaugeGroupElement
        """
        return InverseGaugeGroupElement(self)

class InverseGaugeGroupElement(GaugeGroupElement):
    """
    A gauge group element that represents the inverse action of another element.

    Parameters
    ----------
    gauge_group_el : GaugeGroupElement
        The element to invert.
    """

    def __init__(self, gauge_group_el):
        self.inverse_element = gauge_group_el

    @property
    def transform_matrix(self):
        """
        The gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        return self.inverse_element.transform_matrix_inverse

    @property
    def transform_matrix_inverse(self):
        """
        The inverse of the gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        return self.inverse_element.transform_matrix

    def deriv_wrt_params(self, wrt_filter=None):
        """
        Computes the derivative of the gauge group at this element.

        That is, the derivative of a general element with respect to the gauge
        group's parameters, evaluated at this element.

        Parameters
        ----------
        wrt_filter : list or numpy.ndarray, optional
            Indices of the gauge group parameters to differentiate with respect to.
            If None, differentiation is performed with respect to all the group's parameters.

        Returns
        -------
        numpy.ndarray
        """
        #Derivative of inv(M): d(inv_M) = inv_M * dM * inv_M
        Tinv = self.transform_matrix  # inverse of *original* transform
        dT = self.inverse_element.deriv_wrt_params(wrt_filter)  # shape (d*d, n)
        d, n = int(round(_np.sqrt(dT.shape[0]))), dT.shape[1]

        dT.shape = (d, d, n)  # call it (d1,d2,n)
        dT = _np.rollaxis(dT, 2)  # shape (n, d1, d2)
        deriv = -_np.dot(Tinv, _np.dot(dT, Tinv))  # d,d * (n,d,d * d,d) => d,d * n,d,d => d,n,d
        return _np.swapaxes(deriv, 1, 2).reshape(d * d, n)  # d,n,d => d,d,n => (d*d, n)

    def to_vector(self):
        """
        Get the parameter vector corresponding to this transform.

        Returns
        -------
        numpy.ndarray
        """
        return self.inverse_element.to_vector()

    def from_vector(self, v):
        """
        Reinitialize this `GaugeGroupElement` using the the parameter vector `v`.

        Parameters
        ----------
        v : numpy.ndarray
            A 1D array of length :method:`num_params`

        Returns
        -------
        None
        """
        return self.inverse_element.from_vector()

    @property
    def num_params(self):
        """
        Return the number of parameters of this gauge group element.

        (This is equivalent to the number of parameters of the parent gauge group.)

        Returns
        -------
        int
        """
        return self.inverse_element.num_params

    def inverse(self):
        """
        Creates a gauge group element that performs the inverse of this element.

        Returns
        -------
        GaugeGroupElement
        """
        return self.inverse_element  # inverting an inverse => back to original

class OpGaugeGroup(GaugeGroup):
    """
    A gauge group based on the parameterization of a single `LinearOperator`.

    The parameterization of this linear operator is used to parameterize the
    gauge-transform matrix.  This class is used as the base class for sevearl
    other of gauge group classes.

    Parameters
    ----------
    operation : LinearOperator
        The LinearOperator to base this Gauge group on.

    elementcls : class
        The element class to use when implementing the `element` method.

    name : str
        A name for this group - used for reporting what type of
        gauge optimization was performed.
    """

    def __init__(self, operation, elementcls, name):
        """
        Create a new `OpGaugeGroup`.

        Parameters
        ----------
        operation : LinearOperator
            The LinearOperator to base this Gauge group on.

        elementcls : class
            The element class to use when implementing the `compute_element` method.

        name : str
            A name for this group - used for reporting what type of
            gauge optimization was performed.
        """
        if not isinstance(operation, _op.LinearOperator):
            operation = _op.StaticArbitraryOp(operation, evotype='default', state_space=None)
        self._operation = operation
        self.element = elementcls
        GaugeGroup.__init__(self, name)

    @property
    def num_params(self):
        """
        Return the number of parameters (degrees of freedom) of this gauge group.

        Returns
        -------
        int
        """
        return self._operation.num_params

    def compute_element(self, param_vec):
        """
        Retrieve the element of this group corresponding to `param_vec`

        Parameters
        ----------
        param_vec : numpy.ndarray
            A 1D array of length :method:`num_params`.

        Returns
        -------
        GaugeGroupElement
        """
        elgate = self._operation.copy()
        elgate.from_vector(param_vec)
        return self.element(elgate)

    @property
    def initial_params(self):
        """
        Return a good (or standard) starting parameter vector, used to initialize a gauge optimization.

        Returns
        -------
        numpy.ndarray
            A 1D array of length :method:`num_params`.
        """
        return self._operation.to_vector()

class OpGaugeGroupElement(GaugeGroupElement):
    """
    The element type for `OpGaugeGroup`-derived gauge groups

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Create a new element based on `operation`

        Parameters
        ----------
        operation : LinearOperator
            The operation to base this element on. It provides both parameterization
            information and the gauge transformation matrix itself.
        """
        if not isinstance(operation, _op.LinearOperator):
            operation = _op.StaticArbitraryOp(operation, evotype='default', state_space=None)
        self._operation = operation
        self._inv_matrix = None
        GaugeGroupElement.__init__(self)

    @property
    def transform_matrix(self):
        """
        The gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        return self._operation.to_dense(on_space='minimal')

    @property
    def transform_matrix_inverse(self):
        """
        The inverse of the gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        if self._inv_matrix is None:
            self._inv_matrix = _np.linalg.inv(self._operation.to_dense(on_space='minimal'))
        return self._inv_matrix

    def deriv_wrt_params(self, wrt_filter=None):
        """
        Computes the derivative of the gauge group at this element.

        That is, the derivative of a general element with respect to the gauge
        group's parameters, evaluated at this element.

        Parameters
        ----------
        wrt_filter : list or numpy.ndarray, optional
            Indices of the gauge group parameters to differentiate with respect to.
            If None, differentiation is performed with respect to all the group's parameters.

        Returns
        -------
        numpy.ndarray
        """
        return self._operation.deriv_wrt_params(wrt_filter)

    def to_vector(self):
        """
        Get the parameter vector corresponding to this transform.

        Returns
        -------
        numpy.ndarray
        """
        return self._operation.to_vector()

    def from_vector(self, v):
        """
        Reinitialize this `GaugeGroupElement` using the the parameter vector `v`.

        Parameters
        ----------
        v : numpy.ndarray
            A 1D array of length :method:`num_params`

        Returns
        -------
        None
        """
        self._operation.from_vector(v)
        self._inv_matrix = None

    @property
    def num_params(self):
        """
        Return the number of parameters (degrees of freedom) of this element.

        Returns
        -------
        int
        """
        return self._operation.num_params

class FullGaugeGroup(OpGaugeGroup):
    """
    A fully-parameterized gauge group.

    Every element of the gauge transformation matrix is an independent parameter.

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.

    evotype : Evotype or str, optional
        The evolution type.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.
    """

    def __init__(self, state_space, evotype='default'):
        operation = _op.FullArbitraryOp(_np.identity(state_space.dim, 'd'), evotype, state_space)
        OpGaugeGroup.__init__(self, operation, FullGaugeGroupElement, "Full")

class FullGaugeGroupElement(OpGaugeGroupElement):
    """
    Element of a :class:`FullGaugeGroup`

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Creates a new gauge group element based on `operation`, which
        is assumed to have the correct parameterization.
        """
        OpGaugeGroupElement.__init__(self, operation)

class TPGaugeGroup(OpGaugeGroup):
    """
    A gauge group spanning all trace-preserving (TP) gauge transformations.

    Implemented as a gauge transformation matrix whose first row is locked
    as `[1,0,0...0]` and where every other element is an independent parameter.

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.

    evotype : Evotype or str, optional
        The evolution type.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.
    """

    def __init__(self, state_space, evotype='default'):
        operation = _op.FullTPOp(_np.identity(state_space.dim, 'd'), evotype, state_space)
        OpGaugeGroup.__init__(self, operation, TPGaugeGroupElement, "TP")

class TPGaugeGroupElement(OpGaugeGroupElement):
    """
    Element of a :class:`TPGaugeGroup`

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Creates a new gauge group element based on `operation`, which
        is assumed to have the correct parameterization.
        """
        OpGaugeGroupElement.__init__(self, operation)

    @property
    def transform_matrix_inverse(self):
        """
        The inverse of the gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        if self._inv_matrix is None:
            self._inv_matrix = _np.linalg.inv(self._operation.to_dense())
            self._inv_matrix[0, :] = 0.0  # ensure invers is *exactly* TP
            self._inv_matrix[0, 0] = 1.0  # as otherwise small variations can get amplified
        return self._inv_matrix

class DiagGaugeGroup(OpGaugeGroup):
    """
    A gauge group consisting of just diagonal gauge-transform matrices.

    (Each diagonal element is a separate parameter.)

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.

    evotype : Evotype or str, optional
        The evolution type.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.
    """

    def __init__(self, state_space, evotype='default'):
        dim = state_space.dim
        ltrans = _np.identity(dim, 'd')
        rtrans = _np.identity(dim, 'd')
        baseMx = _np.identity(dim, 'd')
        parameterArray = _np.zeros(dim, 'd')
        parameterToBaseIndicesMap = {i: [(i, i)] for i in range(dim)}
        operation = _op.LinearlyParamArbitraryOp(baseMx, parameterArray,
                                                 parameterToBaseIndicesMap,
                                                 ltrans, rtrans, real=True,
                                                 evotype=evotype, state_space=state_space)
        OpGaugeGroup.__init__(self, operation, DiagGaugeGroupElement, "Diagonal")

class DiagGaugeGroupElement(OpGaugeGroupElement):
    """
    Element of a :class:`DiagGaugeGroup`

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Creates a new gauge group element based on `operation`, which
        is assumed to have the correct parameterization.
        """
        OpGaugeGroupElement.__init__(self, operation)

class TPDiagGaugeGroup(TPGaugeGroup):
    """
    A gauge group consisting of just trace-preserving (TP) diagonal gauge-transform matrices.

    That is, where the first (`[0,0]`) element is fixed at 1.0,
    and each subsequent diagonal element is a separate parameter.

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.

    evotype : Evotype or str, optional
        The evolution type.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.
    """

    def __init__(self, state_space, evotype='default'):
        """
        Create a new gauge group with gauge-transform dimension `dim`, which
        should be the same as `mdl.dim` where `mdl` is a :class:`Model` you
        might gauge-transform.
        """
        dim = state_space.dim
        ltrans = _np.identity(dim, 'd')
        rtrans = _np.identity(dim, 'd')
        baseMx = _np.identity(dim, 'd')
        parameterArray = _np.zeros(dim - 1, 'd')
        parameterToBaseIndicesMap = {i: [(i + 1, i + 1)] for i in range(dim - 1)}
        operation = _op.LinearlyParamArbitraryOp(baseMx, parameterArray,
                                                 parameterToBaseIndicesMap,
                                                 ltrans, rtrans, real=True,
                                                 evotype=evotype, state_space=state_space)
        OpGaugeGroup.__init__(self, operation, TPDiagGaugeGroupElement, "TP Diagonal")

class TPDiagGaugeGroupElement(TPGaugeGroupElement):
    """
    Element of a :class:`TPDiagGaugeGroup`

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Creates a new gauge group element based on `operation`, which
        is assumed to have the correct parameterization.
        """
        TPGaugeGroupElement.__init__(self, operation)

class UnitaryGaugeGroup(OpGaugeGroup):
    """
    A gauge group consisting of unitary gauge-transform matrices.

    This group includes those (superoperator) transformation matrices that
    correspond to unitary evolution.  Parameterization is performed via a
    Lindblad parametrizaton with only Hamiltonian terms.

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.

    basis : Basis or {"pp", "gm", "std"}
        The basis to use when parameterizing the Hamiltonian Lindblad terms.

    evotype : Evotype or str, optional
        The evolution type.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.
    """

    def __init__(self, state_space, basis, evotype='default'):
        errgen = _op.LindbladErrorgen.from_operation_matrix(
            _np.identity(state_space.dim, 'd'), "H", basis, mx_basis=basis, evotype=evotype)
        operation = _op.ExpErrorgenOp(errgen)
        OpGaugeGroup.__init__(self, operation, UnitaryGaugeGroupElement, "Unitary")

class UnitaryGaugeGroupElement(OpGaugeGroupElement):
    """
    Element of a :class:`UnitaryGaugeGroup`

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Creates a new gauge group element based on `operation`, which
        is assumed to have the correct parameterization.
        """
        OpGaugeGroupElement.__init__(self, operation)

class SpamGaugeGroup(OpGaugeGroup):
    """
    Gauge transformations which scale the SPAM and non-unital portions of the gates in a gate set.

    A 2-dimensional gauge group spanning transform matrices of the form:
    [ [ a 0 ... 0]   where a and b are the 2 parameters.  These diagonal
      [ 0 b ... 0]   transform matrices do not affect the SPAM operations
      [ . . ... .]   much more than typical near-unital and TP operations, and
      [ 0 0 ... b] ] so we call this group of transformations the "SPAM gauge".

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.

    evotype : Evotype or str, optional
        The evolution type.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.
    """

    def __init__(self, state_space, evotype='default'):
        """
        Create a new gauge group with gauge-transform dimension `dim`, which
        should be the same as `mdl.dim` where `mdl` is a :class:`Model` you
        might gauge-transform.
        """
        dim = state_space.dim
        ltrans = _np.identity(dim, 'd')
        rtrans = _np.identity(dim, 'd')
        baseMx = _np.identity(dim, 'd')
        parameterArray = _np.zeros(2, 'd')
        parameterToBaseIndicesMap = {0: [(0, 0)],
                                     1: [(i, i) for i in range(1, dim)]}
        operation = _op.LinearlyParamArbitraryOp(baseMx, parameterArray,
                                                 parameterToBaseIndicesMap,
                                                 ltrans, rtrans, real=True,
                                                 evotype=evotype, state_space=state_space)
        OpGaugeGroup.__init__(self, operation, SpamGaugeGroupElement, "Spam")

class SpamGaugeGroupElement(OpGaugeGroupElement):
    """
    Element of a :class:`SpamGaugeGroup`

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Creates a new gauge group element based on `operation`, which
        is assumed to have the correct parameterization.
        """
        OpGaugeGroupElement.__init__(self, operation)

class TPSpamGaugeGroup(OpGaugeGroup):
    """
    Similar to :class:`SpamGaugeGroup` except with TP constrains.

    This means the `[0,0]` element of each transform matrix is fixed at 1.0
    (so all gauge transforms are trace preserving), leaving just a single degree
    of freedom.

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.

    evotype : Evotype or str, optional
        The evolution type.  The special value `"default"` is equivalent
        to specifying the value of `pygsti.evotypes.Evotype.default_evotype`.
    """

    def __init__(self, state_space, evotype='default'):
        """
        Create a new gauge group with gauge-transform dimension `dim`, which
        should be the same as `mdl.dim` where `mdl` is a :class:`Model` you
        might gauge-transform.
        """
        dim = state_space.dim
        ltrans = _np.identity(dim, 'd')
        rtrans = _np.identity(dim, 'd')
        baseMx = _np.identity(dim, 'd')
        parameterArray = _np.zeros(1, 'd')
        parameterToBaseIndicesMap = {0: [(i, i) for i in range(1, dim)]}
        operation = _op.LinearlyParamArbitraryOp(baseMx, parameterArray,
                                                 parameterToBaseIndicesMap,
                                                 ltrans, rtrans, real=True,
                                                 evotype=evotype, state_space=state_space)
        OpGaugeGroup.__init__(self, operation, TPSpamGaugeGroupElement, "TP Spam")

class TPSpamGaugeGroupElement(OpGaugeGroupElement):
    """
    Element of :class:`TPSpamGaugeGroup`

    Parameters
    ----------
    operation : LinearOperator
        The operation to base this element on. It provides both parameterization
        information and the gauge transformation matrix itself.
    """

    def __init__(self, operation):
        """
        Creates a new gauge group element based on `operation`, which
        is assumed to have the correct parameterization.
        """
        OpGaugeGroupElement.__init__(self, operation)

class TrivialGaugeGroup(GaugeGroup):
    """
    A trivial gauge group with no degrees of freedom.

    Useful for telling pyGSTi that you don't want to do any gauge optimization
    within the framework common to the other gauge groups. Using a
    `TrivialGaugeGroup` instead of `None` in gauge optimization will prevent
    pyGSTi from wondering if you meant to not-gauge-optimize and displaying
    warning messages.

    Parameters
    ----------
    state_space : StateSpace
        The state space for this gauge group.  This is the state space that
        elements of the gauge group act on.  This should be the same as `mdl.state_space`
        where `mdl` is a :class:`Model` you want to gauge-transform.
    """

    def __init__(self, state_space):
        self.state_space = state_space
        GaugeGroup.__init__(self, "Trivial")

    @property
    def num_params(self):
        """
        Return the number of parameters (degrees of freedom) of this gauge group.

        Returns
        -------
        int
        """
        return 0

    def compute_element(self, param_vec):
        """
        Retrieve the element of this group corresponding to `param_vec`

        Parameters
        ----------
        param_vec : numpy.ndarray
            A 1D array of length :method:`num_params`.

        Returns
        -------
        TrivialGaugeGroupElement
        """
        assert(len(param_vec) == 0)
        return TrivialGaugeGroupElement(self.state_space.dim)

    @property
    def initial_params(self):
        """
        Return a good (or standard) starting parameter vector, used to initialize a gauge optimization.

        Returns
        -------
        numpy.ndarray
            A 1D array of length :method:`num_params`.
        """
        return _np.empty(0, 'd')

class TrivialGaugeGroupElement(GaugeGroupElement):
    """
    Element of :class:`TrivialGaugeGroup`

    Parameters
    ----------
    dim : int
        The Hilbert-Schmidt space dimension of the gauge group.
    """

    def __init__(self, dim):
        """
        Creates a new trivial gauge group element of dimension `dim`.
        (so transform matirx is a `dim` by `dim` identity matrix).
        """
        self._matrix = _np.identity(dim, 'd')
        GaugeGroupElement.__init__(self)

    @property
    def transform_matrix(self):
        """
        The gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        return self._matrix

    @property
    def transform_matrix_inverse(self):
        """
        The inverse of the gauge-transform matrix.

        Returns
        -------
        numpy.ndarray
        """
        return self._matrix  # inverse of identity is itself!

    def deriv_wrt_params(self, wrt_filter=None):
        """
        Computes the derivative of the gauge group at this element.

        That is, the derivative of a general element with respect to the gauge
        group's parameters, evaluated at this element.

        Parameters
        ----------
        wrt_filter : list or numpy.ndarray, optional
            Indices of the gauge group parameters to differentiate with respect to.
            If None, differentiation is performed with respect to all the group's parameters.

        Returns
        -------
        numpy.ndarray
        """
        return _np.empty(0, 'd')

    def to_vector(self):
        """
        Get the parameter vector corresponding to this transform.

        Returns
        -------
        numpy.ndarray
        """
        return _np.empty(0, 'd')

    def from_vector(self, v):
        """
        Reinitialize this `GaugeGroupElement` using the the parameter vector `v`.

        Parameters
        ----------
        v : numpy.ndarray
            A 1D array of length :method:`num_params`

        Returns
        -------
        None
        """
        assert(len(v) == 0)

    @property
    def num_params(self):
        """
        Return the number of parameters (degrees of freedom) of this element.

        Returns
        -------
        int
        """
        return 0


from __future__ import annotations

import json
import warnings
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import emcee
import h5py
import numpy as np
import pandas as pd
from scipy.special import factorial as fact

import lymph

def lyprox_to_lymph(
    data: pd.DataFrame,
    method: str = "unilateral",
    modalities: List[str] = ["MRI", "PET"],
    convert_t_stage: Optional[Dict[int, Any]] = None
) -> pd.DataFrame:
    """Convert LyProX output into pandas :class:`DataFrame` that the lymph
    package can use for sampling.

    `LyProX <https://lyprox.org>`_ is our online interface where we make
    detailed patterns of involvement on a per-patient basis available and
    visualize it in useful ways.

    Args:
        data: Patient data exported from the LyProX interface.
        method: Can be ``"unilateral"``, ``"bilateral"`` or ``"midline"``. It
            corresponds to the three lymphatic network classes that are
            implemented in the lymph package.
        modalities: List of diagnostic modalities that should be extracted from
            the exported data.
        convert_t_stage: For each of the possible T-categories (0, 1, 2, 3, 4)
            this dictionary holds a key where the corresponding value is the
            'converted' T-category. For example, if one only wants to
            differentiate between 'early' and 'late', then that dictionary
            would look like this:

            .. code-block:: python

                convert_t_stage = {
                    0: 'early',
                    1: 'early',
                    2: 'early',
                    3: 'late',
                    4: 'late'
                }
    Returns:
        A converted pandas :class:`DataFrame` that can then be used with the
        lymph package.
    """
    t_stage_data = data[("tumor", "1", "t_stage")]
    midline_extension_data = data[("tumor", "1", "extension")]
    diagnostic_data = data[modalities].drop(columns=["date"], level=2)

    if convert_t_stage is not None:
        diagnostic_data[("info", "tumor", "t_stage")] = [
            convert_t_stage[t] for t in t_stage_data.values
        ]
    else:
        diagnostic_data[("info", "tumor", "t_stage")] = t_stage_data

    if method == "midline":
        diagnostic_data[("info", "tumor", "midline_extension")] = midline_extension_data
    elif method == "unilateral":
        diagnostic_data = diagnostic_data.drop(columns=["contra"], level=1)
        diagnostic_data.columns = diagnostic_data.columns.droplevel(1)

    return diagnostic_data

class EnsembleSampler(emcee.EnsembleSampler):
    """A custom wrapper of emcee's ``EnsembleSampler`` that adds a sampling
    method that automatically tracks convergence.
    """
    def __init__(
        self,
        nwalkers,
        ndim,
        log_prob_fn,
        pool=None,
        moves=None,
        args=None,
        kwargs=None,
        backend=None,
        vectorize=False,
        blobs_dtype=None,
        parameter_names: Optional[Union[Dict[str, int], List[str]]] = None
    ):
        """Just define a default mixture of moves.
        """
        if moves is None:
            moves = [
                (emcee.moves.DEMove(),        0.8),
                (emcee.moves.DESnookerMove(), 0.2)
            ]

        super().__init__(
            nwalkers,
            ndim,
            log_prob_fn,
            pool,
            moves,
            args,
            kwargs,
            backend,
            vectorize,
            blobs_dtype,
            parameter_names
        )

    def run_sampling(
        self,
        max_steps: int = 10000,
        check_interval: int = 100,
        trust_threshold: float = 50.,
        rel_acor_threshold: float = 0.05,
        verbose: bool = True,
        **kwargs
    ) -> np.ndarray:
        """Extract ``start`` from settings of the sampler and perform sampling
        while monitoring the convergence.

        Args:
            max_steps: Maximum number of sampling steps to perform.
            check_interval: Number of sampling steps after which to check for
                convergence.
            trust_threshold: The autocorrelation estimate is only trusted when
                it is smaller than the number of samples drawn divided by this
                parameter.
            rel_acor_threshold: The relative change of two consequtive trusted
                autocorrelation estimates must fall below.
            verbose: Show progress during sampling and success at the end.
            **kwargs: Any other ``kwargs`` are directly passed to the ``sample``
                method.

        Returns:
            A list of mean autocorrelation times, computed every
            ``check_interval`` samples.
        """
        if verbose:
            print("Starting sampling")

        start = np.random.uniform(
            low=0., high=1.,
            size=(self.nwalkers, self.ndim)
        )

        acor_list = []
        old_acor = np.inf
        idx = 0
        is_converged = False

        for sample in self.sample(start, iterations=max_steps, progress=verbose, **kwargs):
            # after `check_interval` number of samples...
            if self.iteration % check_interval:
                continue

            # ...compute the autocorrelation time and store it in an array.
            new_acor = self.get_autocorr_time(tol=0)
            acor_list.append(np.mean(new_acor))
            idx += 1

            # check convergence based on two criterions:
            # - has the acor time crossed the N / `trust_theshold` line?
            # - did the acor time stay stable?
            is_converged = np.all(new_acor * trust_threshold < self.iteration)
            rel_acor_diff = np.abs(old_acor - new_acor) / new_acor
            is_converged &= np.all(rel_acor_diff < rel_acor_threshold)

            # if it has converged, stop
            if is_converged:
                break

            old_acor = new_acor

        if verbose:
            if is_converged:
                print(f"Sampler converged after {self.iteration} steps")
            else:
                print("Max. number of steps reached")

            acc_frac = 100 * np.mean(self.acceptance_fraction)
            print(f"Acceptance fraction = {acc_frac:.2f}%")
            print(f"Mean autocorrelation time = {np.mean(old_acor):.2f}")

        return acor_list

def tupledict_to_jsondict(dict: Dict[Tuple[str], List[str]]) -> Dict[str, List[str]]:
    """Take a dictionary that has tuples as keys and stringify those keys so
    that it can be serialized to JSON.
    """
    jsondict = {}
    for k, v in dict.items():
        if np.any([',' in s for s in k]):
            raise ValueError("Strings in in key tuple must not contain commas")

        jsondict[",".join(k)] = v
    return jsondict

def jsondict_to_tupledict(dict: Dict[str, List[str]]) -> Dict[Tuple[str], List[str]]:
    """Take a serialized JSON dictionary where the keys are strings of
    comma-separated names and convert them into keys of tuples.
    """
    tupledict = {}
    for k, v in dict.items():
        tupledict[tuple(n for n in k.split(","))] = v
    return tupledict

class HDFMixin(object):
    """Mixin for the :class:`Unilateral`, :class:`Bilateral` and
    :class:`MidlineBilateral` classes to provide the ability to store and load
    settings to and from an HDF5 file.
    """
    graph: Dict[Tuple[str], List[str]]
    patient_data: pd.DataFrame
    modalities: Dict[str, List[float]]

    def to_hdf(
        self,
        filename: str,
        name: str = "",
    ):
        """Store some important settings as well as the loaded data in the
        specified HDF5 file.

        Args:
            filename: Name of or path to HDF5 file.
            name: Name of the group where the info is supposed to be
                stored.
        """
        filename = Path(filename).resolve()

        with h5py.File(filename, 'a') as file:
            group = file.require_group(f"{name}")
            group.attrs["class"] = self.__class__.__name__
            group.attrs["graph"] = json.dumps(tupledict_to_jsondict(self.graph))
            group.attrs["modalities"] = json.dumps(self.modalities)
            group.attrs["base_symmetric"] = getattr(
                self, "base_symmetric", "None"
            )
            group.attrs["trans_symmetric"] = getattr(
                self, "trans_symmetric", "None"
            )

        with pd.HDFStore(filename, 'a') as store:
            store.put(
                key=f"{name}/patient_data",
                value=self.patient_data,
                format="fixed",     # due to MultiIndex this needs to be fixed
                data_columns=None
            )

def system_from_hdf(
    filename: str,
    name: str = "",
    **kwargs
):
    """Create a lymph system instance from the information saved in an HDF5
    file.

    Args:
        filename: Name of the HDF5 file where the info is stored.
        name: Subgroup where to look for the stored settings and data.

    Any other keyword arguments are passed directly to the constructor of the
    respective class.

    Returns:
        An instance of :class:`lymph.Unilateral`, :class:`lymph.Bilateral` or
        :class:`lymph.MidlineBilateral`.
    """
    filename = Path(filename).resolve()
    recover_None = lambda val: val if val != "None" else None

    with h5py.File(filename, 'a') as file:
        group = file.require_group(f"{name}")
        classname = group.attrs["class"]
        graph = jsondict_to_tupledict(json.loads(group.attrs["graph"]))
        modalities = json.loads(group.attrs["modalities"])
        base_symmetric = recover_None(group.attrs["base_symmetric"])
        trans_symmetric = recover_None(group.attrs["trans_symmetric"])

    with pd.HDFStore(filename, 'a') as store:
        patient_data = store.get(f"{name}/patient_data")

    if classname == "Unilateral":
        new_cls = lymph.Unilateral
    elif classname == "Bilateral":
        new_cls = lymph.Bilateral
    elif classname == "MidlineBilateral":
        new_cls = lymph.MidlineBilateral
    else:
        raise RuntimeError(
            "The classname loaded from the file does not correspond to an "
            "implemented class in the `lymph` package."
        )

    new_sys = new_cls(
        graph=graph,
        base_symmetric=base_symmetric,
        trans_symmetric=trans_symmetric,
        **kwargs
    )
    new_sys.modalities = modalities
    new_sys.patient_data = patient_data
    return new_sys

def fast_binomial_pmf(k: int, n: int, p: float):
    """Compute the probability mass function of the binomial distribution.
    """
    q = (1. - p)
    binom_coeff = fact(n) / (fact(k) * fact(n - k))
    return binom_coeff * p**k * q**(n - k)

def change_base(
    number: int,
    base: int,
    reverse: bool = False,
    length: Optional[int] = None
) -> str:
    """Convert an integer into another base.

    Args:
        number: Number to convert
        base: Base of the resulting converted number
        reverse: If true, the converted number will be printed in reverse order.
        length: Length of the returned string. If longer than would be
            necessary, the output will be padded.

    Returns:
        The (padded) string of the converted number.
    """
    if number < 0:
        raise ValueError("Cannot convert negative numbers")
    if base > 16:
        raise ValueError("Base must be 16 or smaller!")
    elif base < 2:
        raise ValueError("There is no unary number system, base must be > 2")

    convertString = "0123456789ABCDEF"
    result = ''

    if number == 0:
        result += '0'
    else:
        while number >= base:
            result += convertString[number % base]
            number = number//base
        if number > 0:
            result += convertString[number]

    if length is None:
        length = len(result)
    elif length < len(result):
        length = len(result)
        warnings.warn("Length cannot be shorter than converted number.")

    pad = '0' * (length - len(result))

    if reverse:
        return result + pad
    else:
        return pad + result[::-1]

def comp_state_dist(table: np.ndarray) -> Tuple[np.ndarray, List[str]]:
    """Compute the distribution of distinct states/diagnoses from a table of
    individual diagnoses detailing the patterns of lymphatic progression per
    patient.

    Args:
        table: Rows of patients and columns of LNLs, reporting which LNL was
            involved for which patient.

    Returns:
        A histogram of unique states and a list of the corresponding state
        labels.

    Note:
        This function cannot deal with parts of the diagnose being unknown. So
        if, e.g., one level isn't reported for a patient, that row will just be
        ignored.
    """
    _, num_cols = table.shape
    table = table.astype(float)
    state_dist = np.zeros(shape=2**num_cols, dtype=int)
    for row in table:
        if not np.any(np.isnan(row)):
            idx = int(np.sum([n * 2**i for i,n in enumerate(row[::-1])]))
            state_dist[idx] += 1

    state_labels = []
    for i in range(2**num_cols):
        state_labels.append(change_base(i, 2, length=num_cols))

    return state_dist, state_labels

def draw_diagnose_times(
    num_patients: int,
    stage_dist: Dict[Any, float],
    diag_times: Optional[Dict[Any, int]] = None,
    time_dists: Optional[Dict[Any, List[float]]] = None,
) -> Tuple[List[int], List[Any]]:
    """Draw T-stages from a distribution over them and determine the
    corresponding diagnose time or draw a one from a distribution over diagnose
    times defined for the respective T-stage.

    Args:
        num_patients: Number of patients to draw diagnose times for.
        stage_dist: Distribution over T-stages.
        diag_times: Fixed diagnose time for a given T-stage.
        time_dists: Holds a distribution over diagnose times for each T-stage
            from which the diagnose times will be drawn if it is given. If this
            is ``None``, ``diag_times`` must be provided.

    Returns:
        The drawn T-stages as well as the drawn diagnose times.
    """
    if num_patients < 1:
        raise ValueError("Number of patients to draw must be 1 or larger")
    if not np.isclose(np.sum(stage_dist), 1.):
        raise ValueError("Distribution over T-stages must sum to 1.")

    # draw the diagnose times for each patient
    if diag_times is not None:
        t_stages = list(diag_times.keys())
        drawn_t_stages = np.random.choice(
            t_stages,
            p=stage_dist,
            size=num_patients
        )
        drawn_diag_times = [diag_times[t] for t in drawn_t_stages]

    elif time_dists is not None:
        t_stages = list(time_dists.keys())
        max_t = len(time_dists[t_stages[0]]) - 1
        time_steps = np.arange(max_t + 1)

        drawn_t_stages = np.random.choice(
            t_stages,
            p=stage_dist,
            size=num_patients
        )
        drawn_diag_times = [
            np.random.choice(time_steps, p=time_dists[t])
            for t in drawn_t_stages
        ]

    else:
        raise ValueError(
            "Either `diag_times`or `time_dists` must be provided"
        )

    return drawn_t_stages, drawn_diag_times

def draw_from_simplex(ndim: int, nsample: int = 1) -> np.ndarray:
    """Draw uniformly from an n-dimensional simplex.

    Args:
        ndim: Dimensionality of simplex to draw from.
        nsample: Number of samples to draw from the simplex.

    Returns:
        A matrix of shape (nsample, ndim) that sums to one along axis 1.
    """
    if ndim < 1:
        raise ValueError("Cannot generate less than 1D samples")
    if nsample < 1:
        raise ValueError("Generating less than one sample doesn't make sense")

    rand = np.random.uniform(size=(nsample, ndim-1))
    unsorted = np.concatenate(
        [np.zeros(shape=(nsample,1)), rand, np.ones(shape=(nsample,1))],
        axis=1
    )
    sorted = np.sort(unsorted, axis=1)

    diff_arr = np.concatenate([[-1., 1.], np.zeros(ndim-1)])
    diff_mat = np.array([np.roll(diff_arr, i) for i in range(ndim)]).T
    res = sorted @ diff_mat

    return res


"""多维多时间序列神经网络滚动训练的数据工具.

version: 0.0.19

author: Congyu Wang

date: 2021-08-26
"""
import numpy as np
import tensorflow as _tf
import numpy as _np
from numba import njit
from typing import List as _List

__all__ = ["TimeSeriesData", "TrainValData"]

class TimeSeriesData:
    """单个时间序列信息.

    Notes:
        用于储存个股的数据信息及预测label，全部使用numpy，日期格式为整数: ``YYYYMMDD``。
        数据分三个部分：时间，数据，标签，第一个维度都是时间，数据的第二个维度为特征。

    """

    def __init__(self,
                 dates: _np.ndarray,
                 data: _np.ndarray,
                 labels: _np.ndarray):
        """储存个股的数据信息及预测目标，全部使用numpy，日期格式为整数: ``YYYYMMDD``.

        Args:
            dates: 日期列, 1D ``numpy.ndarray``, 整数
            data: 训练输入的X，2D ``numpy.ndarray``, (日期长度 x 特征数量)
            labels: 训练标签Y, 1D ``numpy.ndarray``, 长度与dates相同。
                如果为分类问题则是2D, (日期长度 x 类别数量)

        """
        # 检查参数类型
        if (not type(dates) is _np.ndarray or
                not type(data) is _np.ndarray or
                not type(labels) is _np.ndarray):
            raise ValueError("Data should be numpy arrays")
        # 检查日期、数据、标签长度是否一致
        if len(dates) != len(data) or len(dates) != len(labels):
            raise ValueError("Bad data shape")
        # 检查维度是否正确
        if dates.ndim != 1 or data.ndim != 2 or not 1 <= labels.ndim <= 2:
            raise ValueError("Wrong dimensions")
        self.dates = dates.astype(_np.int32)
        self.data = data
        self.labels = labels

class TrainValData:
    """根据训练天数、验证天数、样本历史长度、训练起点生成不同训练阶段的数据."""

    def __init__(self,
                 time_series_list: _List[TimeSeriesData],
                 train_length: int = 1200,
                 validate_length: int = 300,
                 history_length: int = 30,
                 train_val_gap: int = 10,
                 sample_step: int = 2,
                 fill_na: _np.float = _np.NaN,
                 normalize: bool = False):
        """用于获取不同阶段的训练集和验证集.

        Notes:
            ``time_series_list``储存全部的时间序列信息，
            其中每支股票序列为一个单独``TimeSeriesData``，
            完整数据为``List[TimeSeriesData]``类型。

            此外需要提供训练集总交易天数(``train_length``)、
            验证集总交易天数(``validate_length``)、
            单个样本用到的历史长度(``history_length``)、
            采样步进大小(``sample_step``)。

            使用方法为：通过get(start_date)方法获取从start_date
            开始的训练机和验证集。通过逐渐增大start_date训练多个模型回测。

            ``train_val_gap``参数为验证集第一天与训练集最后一天中间间隔的天数，
            如果是相临，则train_val_gap = 0。设置该参数的目的如下：

            如果希望预测未来十天的累计收益，则预测时用到的输入数据为最近的历史数据来预测
            未来十天的累计收益，即用t(-history)到t(0)的数据来预测t(1)到t(11)的累计收益
            而训练时因为要用到十天累计收益做标签，最近的一个十天累计收益是从t(-10)到t(0)，
            用到的历史数据则必须是t(-history-11)到t(-11)的数据。
            而validation时，如果第一个预测点是t(1)(明天收盘价)至t(11)的累计收益，
            则与最后一个训练的数据即：t(-10)至t(0)之间间隔了10天，
            使用``train_val_gap=10``。

            可选项为fill_na，缺失数据填充值，默认为np.Na
            训练时跳过所有有缺失数据的样本。

        Args:
            time_series_list: TimeSeriesData 列表
            train_length: 训练集天数
            validate_length: 验证集天数
            history_length: 每个样本的历史天数
            train_val_gap: 训练集与验证集的间隔
            sample_step: 采样sample时步进的天数
            fill_na: 默认填充为np.NaN，训练时会跳过有确实数据的样本
            normalize: 是否对非率值做每个历史片段的max/min标准化

        """
        # 检查参数类型
        if type(time_series_list) is not list:
            raise ValueError("time_series_list should be a list")
        # 不允许空列表
        if len(time_series_list) == 0:
            raise ValueError("Empty time_series_list")
        # 检查列表元素类型
        for t in time_series_list:
            if type(t) is not TimeSeriesData:
                raise ValueError("time_series_data should be a list "
                                 "of TimeSeriesData objects")
        # 检查参数数值
        if (type(history_length) is not int or
                type(validate_length) is not int or
                type(sample_step) is not int or
                type(train_length) is not int or
                type(train_val_gap) is not int or
                history_length < 1 or
                validate_length < 1 or
                sample_step < 1 or
                train_val_gap < 0 or
                train_length < history_length):
            raise ValueError("bad arguments")

        if type(fill_na) is not _np.float:
            raise ValueError("fill_na should be numpy float")

        # 确保数据特征数量一致
        self.__feature_counts = time_series_list[0].data.shape[1]
        for series in time_series_list:
            if series.data.shape[1] != self.__feature_counts:
                raise ValueError("time series do not have "
                                 "the same number of features")

        # 确保标签维度一致
        label_dims = time_series_list[0].labels.ndim
        for series in time_series_list:
            if series.labels.ndim != label_dims:
                raise ValueError("time labels do not have "
                                 "the same number of dimensions")

        # 标签类别数量
        class_num = 0
        if label_dims == 2:
            class_num = time_series_list[0].labels.shape[1]
            for series in time_series_list:
                if series.labels.shape[1] != class_num:
                    raise ValueError("time series labels do not have "
                                     "the same number of classes")

        self.__class_num = class_num

        # 获取日期列表（所有时间序列日期的并集）
        self.__distinct_dates = _np.unique([date for stock in time_series_list
                                            for date in stock.dates])
        self.__distinct_dates.sort()

        # 聚合数据为(序列(股票), 时间, 特征数量)的张量，缺失数据为np.NaN
        # 标签维度为(序列(股票), 时间)
        self.__data = _np.empty((len(time_series_list),
                                 len(self.__distinct_dates),
                                 self.__feature_counts),
                                dtype=_np.float32)
        if self.__class_num == 0:
            self.__labels = _np.empty((len(time_series_list),
                                       len(self.__distinct_dates)),
                                      dtype=_np.float32)
        else:
            self.__labels = _np.empty((len(time_series_list),
                                       len(self.__distinct_dates),
                                       self.__class_num),
                                      dtype=_np.float32)

        self.__series_date_matrix = _np.empty((len(time_series_list),
                                               len(self.__distinct_dates), 2),
                                              dtype=_np.int32)
        self.__data[:] = fill_na
        self.__labels[:] = fill_na

        # 根据日期序列的位置向张量填充数据
        dates_positions = {date: index
                           for index, date in enumerate(self.__distinct_dates)}
        dates_position_mapper = _np.vectorize(lambda d: dates_positions[d])
        for i, series in enumerate(time_series_list):
            # 找到该序列series.dates日期在日期列表中的位置
            # 将第i个序列填充至tensor的第i行
            position_index = dates_position_mapper(series.dates)
            self.__data[i, position_index, :] = series.data
            if self.__class_num == 0:
                self.__labels[i, position_index] = series.labels
            else:
                self.__labels[i, position_index, :] = series.labels
            self.__series_date_matrix[i, position_index, 0] = series.dates
            self.__series_date_matrix[i, position_index, 1] = i

        self.__train_length = train_length
        self.__validate_length = validate_length
        self.__history_length = history_length
        self.__sample_step = sample_step
        self.__train_val_gap = train_val_gap
        self.__normalize = normalize

    def get(self,
            start_date: int,
            order="by_date",
            validate_only=False,
            validate_length=None,
            normalize=False):
        """获取从某天开始的训练集和验证集.

        Notes:
            根据设定的训练集天数以及验证集天数，从start_date开始获取正确的
            训练集以及验证集，以及他们各自的日期范围信息(该信息以字典形式返回)。

            需要注意:
            训练集的的开始和结束是指其data, label时间范围并集，而
            验证集的开始和结束则只是指其label的时间范围。
            验证集的输入数据可以与训练集重叠，只要其标签数据的包含的时间范围
            与训练集数据包含的时间范围没有交集即可。

            具体时间信息参考函数返回的最后一个元素，是一个包含时间信息的``dict``。

        Args:
            start_date: 该轮训练开始日期，整数``YYYYMMDD``
            order: 有三种顺序 ``shuffle``, ``by_date``, ``by_series``。
                分别为随机打乱股票和时间，按时间顺序优先，按股票顺序优先，默认by_date。
            validate_only: 如果设置为True，则只返回validate set
                和训练集、验证集时间信息。可以用于训练后的分析。
            validate_length (int): override class validate_length
            normalize (bool): override class normalize

        Returns:
            如果``validate_only=False``，返回训练集、验证集、日期信息：
            (train, val, dates_info(dict))。
            如果为``True``，则返回验证集、日期信息。

        Raises:
            ValueError: 日期范围超出最大日期会报ValueError。

        """
        if validate_length is not None:
            if type(validate_length) is not int:
                raise ValueError("`validate_length` must be an integer")
            if validate_length < 1:
                raise ValueError("`validate_length` should be at least 1")
        else:
            validate_length = self.__validate_length
        if not normalize:
            normalize = self.__normalize
        return self.__get_in_memory__(start_date,
                                      order,
                                      validate_only,
                                      validate_length,
                                      normalize)

    def __get_in_memory__(self,
                          start_date,
                          order="by_date",
                          validate_only=False,
                          validate_length=None,
                          normalize=False):
        """使用显存生成历史数据.

        使用tensorflow from_tensor_slices，通过传递完整的tensor进行训练，
        股票数量大时，需要较大内存

        """
        # 获取用于构建训练集、验证集的相关信息
        kwargs = {"start_date": start_date, "order": order}
        if validate_length is not None:
            kwargs.update({"validate_length": validate_length})
        train_args, val_args, dates_info = self.__get_period_info__(**kwargs)
        train_args = (*train_args, normalize)
        val_args = (*val_args, normalize)
        # 将输入的数据、标签片段转化为单个sample包含history日期长度的历史信息
        (val_x,
         val_y,
         val_dates_series) = __full_tensor_generation__(*val_args)
        # 转化为tensorflow DataSet
        val = _tf.data.Dataset.from_tensor_slices((val_x, val_y))
        val_dates_list = val_dates_series[:, 0].tolist()
        val_series_list = val_dates_series[:, 1].tolist()
        dates_info["validation"]["dates_list"] = val_dates_list
        dates_info["validation"]["series_list"] = val_series_list

        if validate_only:
            return val, dates_info

        (train_x,
         train_y,
         train_dates_series) = __full_tensor_generation__(*train_args)
        train = _tf.data.Dataset.from_tensor_slices((train_x, train_y))
        train_dates_list = train_dates_series[:, 0].tolist()
        train_series_list = train_dates_series[:, 1].tolist()
        dates_info["training"]["dates_list"] = train_dates_list
        dates_info["training"]["series_list"] = train_series_list
        return train, val, dates_info

    def __get_period_info__(self, start_date, order="by_date",
                            validate_length=None):
        """根据开始时间计算用于构建训练集、验证集的相关信息."""
        if type(start_date) is not int:
            raise ValueError("start date should be an integer YYYYMMDD")

        # 找到大于等于start_date的最小日期
        after_start_date = self.__distinct_dates >= start_date
        first_date = _np.min(self.__distinct_dates[after_start_date])

        # 查看剩余日期数量是否大于等于训练集验证集总长度
        if _np.sum(after_start_date) < (self.__train_length +
                                        validate_length +
                                        self.__train_val_gap):
            raise ValueError("date range exceeded end of dates")

        # 计算各个时间节点在时间列表中的位置
        # 训练集开始位置(data的开始位置)
        train_start_index = __first_index__(self.__distinct_dates, first_date)
        # 训练集结束位置(不包含)
        train_end_index = train_start_index + self.__train_length
        # 验证集开始位置(data的开始位置)
        val_start_index = (train_end_index -
                           self.__history_length +
                           self.__train_val_gap + 1)
        # 验证集结束位置(不包含)
        val_end_index = (train_end_index +
                         validate_length +
                         self.__train_val_gap)

        # 根据各个数据集的开始结束位置以及训练数据的顺序选项，获取构建数据的参数
        train_args = self.__get_generator_args__(train_start_index,
                                                 train_end_index,
                                                 order=order)
        val_args = self.__get_generator_args__(val_start_index,
                                               val_end_index,
                                               order=order)
        dates_info = self.__dates_info__(train_start_index,
                                         val_start_index,
                                         train_args,
                                         val_args)

        return train_args, val_args, dates_info

    def __get_generator_args__(self, start_index, end_index, order="by_date"):
        """获取单个generator需要的数据片段.

        Notes:
            根据数据集的开始、结束位置以及的顺序选项，获取该训练集的数据、标签片段
            以及用于生成训练数据的(序列, 日期)pair列表的顺序信息(generation_list)。

            generation_list第一列为序列编号，第二列为日期。

            generation_list中的日期数字代表*每个历史数据片段*的第一个日期相对
            该数据集片段（start_index:end_index）的位置。

            注意：

                - 该处的日期列表不代表每个历史片段的结束位置

                - 也不是相对TrainValData类日期列表的位置
        """
        length = end_index - start_index
        data = self.__data[:, start_index:end_index, :]
        label = self.__labels[:, start_index:end_index]
        dates_series = self.__series_date_matrix[:, start_index:end_index, :]
        generation_list = [[series_i, t]
                           for t in range(0,
                                          length - self.__history_length + 1,
                                          self.__sample_step)
                           for series_i in range(len(data))]

        if order == "shuffle":
            generation_list = _np.array(generation_list, dtype=_np.int32)
            _np.random.shuffle(generation_list)
        elif order == "by_date":
            generation_list = _np.array(generation_list, dtype=_np.int32)
        elif order == "by_series":
            generation_list = sorted(generation_list, key=lambda k: k[0])
            generation_list = _np.array(generation_list, dtype=_np.int32)
        else:
            raise ValueError("wrong order argument, choose from `shuffle`, "
                             "`by_date`, and `by_series`")

        history_length = self.__history_length

        return data, label, generation_list, history_length, dates_series

    def __dates_info__(self,
                       train_start_index,
                       val_start_index,
                       train_args,
                       val_args):
        """根据生成数据的列表，计算用于显示的日期信息."""
        # 获取generation_list(生成数据的顺序信息)的时间列
        # 该时间列为相对数据片段开头的时间位置
        train_generation_list = train_args[2]
        val_generation_list = val_args[2]
        train_time_index = train_generation_list[:, 1]
        val_time_index = val_generation_list[:, 1]

        # 加上片段的开始位置，得到相对TrainValData类日期列表的位置
        train_time_index = train_time_index + train_start_index
        val_time_index = val_time_index + val_start_index

        # 训练集在日期列表中的开始位置
        training_beginning = _np.min(train_time_index)

        # 结束位置：加上历史长度减去一，获取最大日期位置(inclusive)
        training_ending = _np.max(train_time_index)
        training_ending += self.__history_length - 1

        # validation集每次取的都是某历史片段末尾的数据
        # 所以加上历史减去一
        validation_index = val_time_index + self.__history_length - 1
        validation_beginning = _np.min(validation_index)
        validation_ending = _np.max(validation_index)

        dates_info = {
            "training": {
                "start_date": int(self.__distinct_dates[training_beginning]),
                "end_date": int(self.__distinct_dates[training_ending]),
            },
            "validation": {
                "start_date": int(self.__distinct_dates[validation_beginning]),
                "end_date": int(self.__distinct_dates[validation_ending]),
            }
        }
        return dates_info

def __full_tensor_generation__(data,
                               label,
                               generation_list,
                               history,
                               dates_series,
                               normalize):
    """将输入的数据、标签片段转化为单个sample包含history日期长度的历史信息."""
    # 先将该数据片段的历史维度展开

    # 根据generation_list指定的series，日期，获取标签及数据片段
    total_time_length = data.shape[1]
    expanded = [__gather_2d__(data[:, i: (total_time_length + i
                                          - history + 1), :], generation_list)
                for i in range(history)]

    # date_all dimensions: (series * dates, features, history)
    data_all = _np.stack(expanded, axis=-1)
    label_all = __gather_2d__(label[:, history - 1:], generation_list)
    dates_series_all = __gather_2d__(dates_series[:, history - 1:],
                                     generation_list)

    # 去掉所有包含缺失数据的某股票某时间历史片段
    label_nan = _np.isnan(label_all)
    if _np.ndim(label_all) == 2:
        # 如果是分类问题, 需要特殊处理
        label_nan = _np.any(label_nan, axis=1)
    data_nan = _np.isnan(data_all)
    series_time_nan = _np.any(
        _np.any(data_nan, axis=2),
        axis=1
    )
    not_nan = _np.logical_not(_np.logical_or(series_time_nan, label_nan))

    data_all = data_all[not_nan]
    label_all = label_all[not_nan]
    dates_series_all = dates_series_all[not_nan]

    # max/min standardization for series greater than 1
    if normalize:
        max_val = _np.max(data_all, axis=-1, keepdims=True)
        non_rate_mask = _np.less(1.0, _np.squeeze(max_val))
        min_val = _np.min(data_all, axis=-1, keepdims=True)
        with _np.errstate(divide='ignore', invalid='ignore'):
            normalized = _np.nan_to_num((data_all - min_val)
                                        / (max_val - min_val))
        data_all[non_rate_mask] = normalized[non_rate_mask]

    data_all = _np.transpose(data_all, axes=(0, 2, 1))
    data_all = _tf.constant(data_all, dtype=_tf.float32)
    label_all = _tf.constant(label_all, dtype=_tf.float32)

    return data_all, label_all, dates_series_all

def __first_index__(array, element):
    """计算第一个出现的元素的位置."""
    return _np.min(_np.where(array == element))

@njit
def __gather_2d__(array: _np.ndarray, generation_list: _np.ndarray):
    out_shape = (len(generation_list), *array.shape[2:])
    out_array = np.empty(shape=out_shape,
                         dtype=array.dtype)
    pointer = 0
    for i, j in generation_list:
        out_array[pointer] = array[i, j]
        pointer += 1
    return out_array


"""
run MD, replica exchange (parallel tempering) and calcaulte energy using OpenMM
"""

from sys import stdout

import copy

import numpy as np
import netCDF4

import simtk.openmm
import simtk.openmm.app
import simtk.unit

from rotation import random_rotation

openmm_solvent_models = {  "OpenMM_Gas":None,
                            "OpenMM_GBn":simtk.openmm.app.GBn,
                            "OpenMM_GBn2":simtk.openmm.app.GBn2,
                            "OpenMM_HCT":simtk.openmm.app.HCT,
                            "OpenMM_OBC1":simtk.openmm.app.OBC1,
                            "OpenMM_OBC2":simtk.openmm.app.OBC2 }

KB = 0.001987204134799235  # kcal/mol/K

class OpenMM_MD(object):
    def __init__(self, prmtop, inpcrd, phase="OpenMM_Gas", temperature=300.):
        """
        """
        self._prmtop = simtk.openmm.app.AmberPrmtopFile(prmtop)
        inpcrd = simtk.openmm.app.AmberInpcrdFile(inpcrd)

        selected_solvent = openmm_solvent_models[phase]
        system = self._prmtop.createSystem(nonbondedMethod = simtk.openmm.app.NoCutoff, \
                constraints = simtk.openmm.app.HBonds, implicitSolvent = selected_solvent)
        integrator = simtk.openmm.LangevinIntegrator(temperature*simtk.unit.kelvin, 1/simtk.unit.picosecond, 0.002*simtk.unit.picoseconds)

        self._simulation = simtk.openmm.app.Simulation(self._prmtop.topology, system, integrator)
        self._simulation.context.setPositions(inpcrd.positions)
        print("Energy minimizing")
        self._simulation.minimizeEnergy()

    def _initialize_nc(self, nc_file_name, niterations):
        nc_handle = netCDF4.Dataset(nc_file_name, mode="w", format="NETCDF4")
        natoms = len( list( self._prmtop.topology.atoms() ) )

        nc_handle.createDimension("three", 3)
        nc_handle.createDimension("natoms", natoms)
        nc_handle.createDimension("nconfs", niterations)

        nc_handle.createVariable("positions", "f8", tuple(["nconfs", "natoms", "three"]))
        return nc_handle

    def run(self, nc_file_name, steps_per_iteration=500, niterations=1000):
        """
        """
        self._simulation.reporters.append(simtk.openmm.app.StateDataReporter(stdout, steps_per_iteration,
                                            step=True, potentialEnergy=True))

        nc_handle = self._initialize_nc(nc_file_name, niterations)

        for iteration in range(niterations):
            self._simulation.step(steps_per_iteration)
            state = self._simulation.context.getState(getPositions=True, getEnergy=True)

            positions = copy.deepcopy(state.getPositions())
            conf = np.array(positions.value_in_unit(simtk.unit.angstrom), dtype=float)
            conf = random_rotation(conf)

            nc_handle.variables["positions"][iteration,:,:] = conf
        nc_handle.close()

        return None

class OpenMM_TREMD(object):
    def __init__(self, prmtop, inpcrd, phase, temperatures):
        """
        :param prmtop: str, name of AMBER prmtop file
        :param inpcrd: str, name of AMBER coordinate file
        :param phase: str
        :param temperatures: list or ndarray of float
        """
        self._temperatures = list(temperatures)
        self._simulations  = self._create_simulations(prmtop, inpcrd, phase, temperatures)
        self._acepted_exchange = np.zeros([len(self._temperatures)-1], dtype=float)
        self._nr_attempts = np.zeros([len(self._temperatures)-1], dtype=float)
        self._start = 0

    def run(self, nc_file_name, steps_per_iteration, niterations, rotations_per_iteration):
        """
        """
        nc_handle = self._initialize_nc(nc_file_name, niterations, rotations_per_iteration)

        nrotations = 0
        for iteration in range(niterations):

            self._md_evolve(steps_per_iteration)

            energies = self._get_potential_energies()
            nc_handle.variables["energies"][iteration, :] = np.array(energies, dtype=float)

            positions = self._get_positions()
            for state, p in enumerate(positions):
                crd = np.array(p.value_in_unit(simtk.unit.angstrom), dtype=float)
                nc_handle.variables["positions"][iteration, state, :, :] = crd

            crd = np.array(positions[0].value_in_unit(simtk.unit.angstrom), dtype=float)
            for rotation in range(rotations_per_iteration):
                rotated_crd = random_rotation(crd)
                nc_handle.variables["rotated_positions"][nrotations, :, :] = rotated_crd
                nrotations += 1

            self._exchange(self._start)
            self._switch_start()

            print("iteration ", iteration)
            print("energies ", energies)

        acceptance_rate = self._acepted_exchange / self._nr_attempts
        nc_handle.variables["acceptance_rate"][:] = acceptance_rate
        nc_handle.close()
        return None

    def _switch_start(self):
        if self._start == 0:
            self._start = 1
        elif self._start == 1:
            self._start = 0
        else:
            raise RuntimeError("self._start is %d"%self._start)
        return

    def _create_simulation(self, prmtop, inpcrd, phase, temperature):
        prmtop = simtk.openmm.app.AmberPrmtopFile(prmtop)
        inpcrd = simtk.openmm.app.AmberInpcrdFile(inpcrd)
        selected_solvent = openmm_solvent_models[phase]

        system = prmtop.createSystem(nonbondedMethod=simtk.openmm.app.NoCutoff, constraints=simtk.openmm.app.HBonds, implicitSolvent=selected_solvent)
        integrator = simtk.openmm.LangevinIntegrator(temperature*simtk.unit.kelvin, 1/simtk.unit.picosecond, 0.002*simtk.unit.picoseconds)

        simulation = simtk.openmm.app.Simulation(prmtop.topology, system, integrator)
        simulation.context.setPositions(inpcrd.positions)
        simulation.minimizeEnergy()
        return simulation

    def _create_simulations(self, prmtop, inpcrd, phase, temperatures):
        simulations = []
        for temperature in temperatures:
            sim = self._create_simulation(prmtop, inpcrd, phase, temperature)
            simulations.append(sim)
        return simulations

    def _md_evolve(self, steps):
        for sim in self._simulations:
            sim.step(steps)
        return None

    def _get_potential_energies(self):
        energies = []
        for sim in self._simulations:
            state = sim.context.getState(getEnergy=True)
            energy = copy.deepcopy(state.getPotentialEnergy())
            e = energy.value_in_unit(simtk.unit.kilocalorie_per_mole)
            energies.append(e)
        return energies

    def _get_positions(self):
        positions = []
        for sim in self._simulations:
            state = sim.context.getState(getPositions=True)
            pos = copy.deepcopy(state.getPositions())
            positions.append(pos)
        return positions

    def _get_velocities(self):
        velocities = []
        for sim in self._simulations:
            state = sim.context.getState(getVelocities=True)
            vel = copy.deepcopy(state.getVelocities())
            velocities.append(vel)
        return velocities

    def _metropolis_prob(self, T1, E1, T2, E2):
        assert T1 > 0 and T2 > 0, "T1 and T2 must be possitive"
        deltaE = ((1./T1 - 1./T2) / KB) * (E2 - E1)
        return np.exp(-deltaE)

    def _exchange(self, start):
        assert start in [0, 1], "start must be either 0 or 1"
        simulation_pairs = zip( self._simulations[start : -1 : 2], self._simulations[start+1 : : 2] )

        energies = self._get_potential_energies()
        energy_pairs = zip( energies[start : -1 : 2], energies[start+1 : : 2] )

        temperature_pairs = zip( self._temperatures[start : -1 : 2], self._temperatures[start+1 : : 2] )

        positions = self._get_positions()
        position_pairs = zip( positions[start : -1 : 2], positions[start+1 : : 2] )

        velocities = self._get_velocities()
        velocity_pairs = zip( velocities[start : -1 : 2], velocities[start+1 : : 2] )

        pair_indices = range(start, len(self._temperatures)-1, 2)

        for i in range(len(simulation_pairs)):
            E1 = energy_pairs[i][0]
            E2 = energy_pairs[i][1]
            T1 = temperature_pairs[i][0]
            T2 = temperature_pairs[i][1]

            x_prob = self._metropolis_prob(T1, E1, T2, E2)
            if x_prob >= 1. or x_prob > np.random.random():
                simulation_pairs[i][0].context.setPositions( position_pairs[i][1] )
                simulation_pairs[i][1].context.setPositions( position_pairs[i][0] )

                vel_unit = simtk.unit.nanometer/simtk.unit.picosecond

                v1 = np.sqrt(T1 / T2 ) * np.array( velocity_pairs[i][1].value_in_unit(vel_unit) )
                v2 = np.sqrt(T2 / T1 ) * np.array( velocity_pairs[i][0].value_in_unit(vel_unit) )

                simulation_pairs[i][0].context.setVelocities( simtk.unit.quantity.Quantity(v1, unit=vel_unit) )
                simulation_pairs[i][1].context.setVelocities( simtk.unit.quantity.Quantity(v2, unit=vel_unit) )

                self._acepted_exchange[pair_indices[i]] += 1.

            self._nr_attempts[pair_indices[i]] += 1.
        return

    def _initialize_nc(self, nc_file_name, niterations, rotations_per_iteration):
        """
        """
        nc_handle = netCDF4.Dataset(nc_file_name, mode="w", format="NETCDF4")
        natoms = len( list( self._simulations[0].topology.atoms() ) )

        nc_handle.createDimension("three", 3)
        nc_handle.createDimension("natoms", natoms)
        nc_handle.createDimension("nstates", len(self._simulations) )
        nc_handle.createDimension("niterations", niterations)
        nc_handle.createDimension("nrotations", niterations * rotations_per_iteration)
        nc_handle.createDimension("npairs", len(self._simulations) - 1)

        nc_handle.createVariable("energies", "f8", tuple(["niterations", "nstates"]) )
        nc_handle.createVariable("positions", "f4", tuple(["niterations", "nstates", "natoms", "three"]))
        nc_handle.createVariable("rotated_positions", "f4", tuple(["nrotations", "natoms", "three"]))
        nc_handle.createVariable("acceptance_rate", "f8", tuple(["npairs"]))
        nc_handle.createVariable("temperatures", "f8", tuple(["nstates"]))

        nc_handle.variables["temperatures"][:] = np.array(self._temperatures)
        return nc_handle

def openmm_energy(prmtop_file, crd, phase):
    """
    crd in angstroms
    crd can be an inpcrd file name, a array with shape (natom, 3) or shape (nframe, natom, 3)
    """
    implicit_solvents = openmm_solvent_models
    selected_solvent = implicit_solvents[phase]

    if type(crd) == str:
        inpcrd = simtk.openmm.app.AmberInpcrdFile(crd)
        position = inpcrd.positions
        crd = position.value_in_unit(simtk.unit.angstrom)

    crd_ensemble = np.array(crd, dtype=float)
    if len(crd_ensemble.shape) == 2:
        crd_ensemble = np.array([crd_ensemble], dtype=float)

    if len(crd_ensemble.shape) != 3:
        raise RuntimeError("crd_ensemble has wrong shape")

    prmtop = simtk.openmm.app.AmberPrmtopFile(prmtop_file)
    system = prmtop.createSystem(nonbondedMethod=simtk.openmm.app.NoCutoff,
                                 constraints=None, implicitSolvent=selected_solvent)
    dummy_integrator = simtk.openmm.VerletIntegrator(0.002*simtk.unit.picoseconds)
    simulation = simtk.openmm.app.Simulation(prmtop.topology, system, dummy_integrator)

    pot_energies = []
    for conf in crd_ensemble:
        simulation.context.setPositions(conf / 10.)  # convert to nano-meter which is internal unit of OpenMM
        state = simulation.context.getState(getEnergy=True)
        energy = state.getPotentialEnergy()
        pot_energies.append(energy.value_in_unit(simtk.unit.kilocalorie_per_mole))

    return np.array(pot_energies, dtype=float)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

###############################################################################
#                                                                             #
# RMG - Reaction Mechanism Generator                                          #
#                                                                             #
# Copyright (c) 2002-2019 Prof. William H. Green (whgreen@mit.edu),           #
# Prof. Richard H. West (r.west@neu.edu) and the RMG Team (rmg_dev@mit.edu)   #
#                                                                             #
# Permission is hereby granted, free of charge, to any person obtaining a     #
# copy of this software and associated documentation files (the 'Software'),  #
# to deal in the Software without restriction, including without limitation   #
# the rights to use, copy, modify, merge, publish, distribute, sublicense,    #
# and/or sell copies of the Software, and to permit persons to whom the       #
# Software is furnished to do so, subject to the following conditions:        #
#                                                                             #
# The above copyright notice and this permission notice shall be included in  #
# all copies or substantial portions of the Software.                         #
#                                                                             #
# THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  #
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,    #
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE #
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER      #
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING     #
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER         #
# DEALINGS IN THE SOFTWARE.                                                   #
#                                                                             #
###############################################################################

"""
This module contains the TransportData class for storing transport properties.
"""

import numpy as np

import rmgpy.constants as constants
from rmgpy import quantity
from rmgpy.quantity import DipoleMoment, Energy, Length, Volume
from rmgpy.rmgobject import RMGObject

class TransportData(RMGObject):
    """
    A set of transport properties.

    The attributes are:
    =================  ============================================================
    Attribute          Description
    =================  ============================================================
    `shapeIndex`        0 for monoatomic, 1 for linear, 2 for nonlinear
    `epsilon`           Lennard-Jones well depth
    `sigma`             Lennard-Jones collision diameter
    `dipoleMoment`      Dipole Moment
    `polarizability`    Polarizability Volume
    `rotrelaxcollnum`   Rotational relaxation number at 298 K, saved as a double.
    =================  ============================================================

    """

    def __init__(self, shapeIndex=None, epsilon=None, sigma=None, dipoleMoment=None, polarizability=None,
                 rotrelaxcollnum=None, comment=''):
        self.shapeIndex = shapeIndex
        try:
            self.epsilon = Energy(epsilon)
        except quantity.QuantityError:
            self.epsilon = quantity.Temperature(epsilon)
            self.epsilon.value_si *= constants.R
            self.epsilon.units = 'kJ/mol'
        self.sigma = Length(sigma)
        self.dipoleMoment = DipoleMoment(dipoleMoment)
        self.polarizability = Volume(polarizability)
        self.rotrelaxcollnum = rotrelaxcollnum
        self.comment = comment

    def __repr__(self):
        """
        Return a string representation that can be used to reconstruct the
        TransportData object.
        """
        attributes = []
        if self.shapeIndex is not None:
            attributes.append('shapeIndex={0!r}'.format(self.shapeIndex))
        if self.epsilon is not None:
            attributes.append('epsilon={0!r}'.format(self.epsilon))
        if self.sigma is not None:
            attributes.append('sigma={0!r}'.format(self.sigma))
        if self.dipoleMoment is not None:
            attributes.append('dipoleMoment={0!r}'.format(self.dipoleMoment))
        if self.polarizability is not None:
            attributes.append('polarizability={0!r}'.format(self.polarizability))
        if self.rotrelaxcollnum is not None:
            attributes.append('rotrelaxcollnum={0!r}'.format(self.rotrelaxcollnum))
        if self.comment:
            attributes.append('comment="""{0!s}"""'.format(self.comment))
        string = 'TransportData({0!s})'.format(', '.join(attributes))
        return string

    def __reduce__(self):
        """
        A helper function used when picking a TransportData object.
        """
        return (TransportData, (self.shapeIndex, self.epsilon, self.sigma, self.dipoleMoment,
                                self.polarizability, self.rotrelaxcollnum, self.comment))

    def get_collision_frequency(self, T, M, mu):
        """
        Return the value of the Lennard-Jones collision frequency in Hz at the
        given temperature `T` in K for colliders with the given concentration
        `M` in mol/m^3 and reduced mass `mu` in amu.

        This seems to also exist in rmgpy.pdep.configuration.calculate_collision_frequency
        Why the redundancy?
        """
        sigma = self.sigma.value_si
        epsilon = self.epsilon.value_si
        M *= constants.Na  # mol/m^3 -> molecules/m^3
        Tred = constants.R * T / epsilon
        omega22 = 1.16145 * Tred ** (-0.14874) + 0.52487 * np.exp(-0.77320 * Tred) + 2.16178 * np.exp(-2.43787 * Tred)
        mu *= constants.amu
        return omega22 * np.sqrt(8 * constants.kB * T / constants.pi / mu) * constants.pi * sigma * sigma * M

    def to_cantera(self):
        """
        Returns a Cantera GasTransportData object.

        The Cantera usage is as follows:

        GasTransportData().set_customary_units(self, geometry, diameter, well_depth, dipole=0.0, polarizability=0.0, rotational_relaxation=0.0, acentric_factor=0.0)
        Set the parameters using customary units: diameter in Angstroms, well depth in Kelvin, dipole in Debye, rotational relaxiation at 298 K, and polarizability in Angstroms^3.
        These are the units used in in CK-style input files.
        """
        import cantera as ct

        ct_transport = ct.GasTransportData()

        if self.shapeIndex == 0:
            geometry = 'atom'
        elif self.shapeIndex == 1:
            geometry = 'linear'
        elif self.shapeIndex == 2:
            geometry = 'nonlinear'

        # collision diameter in angstroms
        diameter = self.sigma.value_si * 1e10
        # Well depth in Kelvins
        well_depth = self.epsilon.value_si / constants.R
        # Dipole in debye
        dipole = self.dipoleMoment.value_si * constants.c * 1e21 if self.dipoleMoment else 0.0
        # polarizability in cubic angstroms
        polarizability = self.polarizability.value_si * 1e30 if self.polarizability else 0.0
        rotational_relaxation = self.rotrelaxcollnum if self.rotrelaxcollnum else 0.0
        acentric_factor = 0.0

        ct_transport.set_customary_units(geometry, diameter, well_depth, dipole, polarizability,
                                         rotational_relaxation, acentric_factor)

        return ct_transport