1#  ___________________________________________________________________________
2#
3#  Pyomo: Python Optimization Modeling Objects
4#  Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
5#  Under the terms of Contract DE-NA0003525 with National Technology and
6#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
7#  rights in this software.
8#  This software is distributed under the 3-clause BSD License.
9#  ___________________________________________________________________________
10
11import weakref
12from pyomo.common.collections import ComponentMap
13from pyomo.core.base.component import ModelComponentFactory
14from pyomo.core.base.set import UnknownSetDimen
15from pyomo.core.base.var import Var
16from pyomo.dae.contset import ContinuousSet
17
18__all__ = ('DerivativeVar', 'DAE_Error',)
19
20
21def create_access_function(var):
22    """
23    This method returns a function that returns a component by calling
24    it rather than indexing it
25    """
26    def _fun(*args):
27        return var[args]
28    return _fun
29
30
31class DAE_Error(Exception):
32    """Exception raised while processing DAE Models"""
33
34
35@ModelComponentFactory.register("Derivative of a Var in a DAE model.")
36class DerivativeVar(Var):
37    """
38    Represents derivatives in a model and defines how a
39    :py:class:`Var<pyomo.environ.Var>` is differentiated
40
41    The :py:class:`DerivativeVar <pyomo.dae.DerivativeVar>` component is
42    used to declare a derivative of a :py:class:`Var <pyomo.environ.Var>`.
43    The constructor accepts a single positional argument which is the
44    :py:class:`Var<pyomo.environ.Var>` that's being differentiated. A
45    :py:class:`Var <pyomo.environ.Var>` may only be differentiated with
46    respect to a :py:class:`ContinuousSet<pyomo.dae.ContinuousSet>` that it
47    is indexed by. The indexing sets of a :py:class:`DerivativeVar
48    <pyomo.dae.DerivativeVar>` are identical to those of the :py:class:`Var
49    <pyomo.environ.Var>` it is differentiating.
50
51    Parameters
52    ----------
53    sVar : ``pyomo.environ.Var``
54        The variable being differentiated
55
56    wrt : ``pyomo.dae.ContinuousSet`` or tuple
57        Equivalent to `withrespectto` keyword argument. The
58        :py:class:`ContinuousSet<pyomo.dae.ContinuousSet>` that the
59        derivative is being taken with respect to. Higher order derivatives
60        are represented by including the
61        :py:class:`ContinuousSet<pyomo.dae.ContinuousSet>` multiple times in
62        the tuple sent to this keyword. i.e. ``wrt=(m.t, m.t)`` would be the
63        second order derivative with respect to ``m.t``
64    """
65
66    # Private Attributes:
67    # _stateVar   The :class:`Var` being differentiated
68    # _wrt        A list of the :class:`ContinuousSet` components the
69    #             derivative is being taken with respect to
70    # _expr       An expression representing the discretization equations
71    #             linking the :class:`DerivativeVar` to its state :class:`Var`.
72
73    def __init__(self, sVar, **kwds):
74
75        if not isinstance(sVar, Var):
76            raise DAE_Error(
77                "%s is not a variable. Can only take the derivative of a Var"
78                "component." % sVar)
79
80        if "wrt" in kwds and "withrespectto" in kwds:
81            raise TypeError(
82                "Cannot specify both 'wrt' and 'withrespectto keywords "
83                "in a DerivativeVar")
84
85        wrt = kwds.pop('wrt', None)
86        wrt = kwds.pop('withrespectto', wrt)
87
88        try:
89            num_contset = len(sVar._contset)
90        except AttributeError:
91            # This dictionary keeps track of where the ContinuousSet appears
92            # in the index. This implementation assumes that every element
93            # in an indexing set has the same dimension.
94            sVar._contset = ComponentMap()
95            sVar._derivative = {}
96            if sVar.dim() == 0:
97                num_contset = 0
98            else:
99                sidx_sets = list(sVar.index_set().subsets())
100                loc = 0
101                for i, s in enumerate(sidx_sets):
102                    if s.ctype is ContinuousSet:
103                        sVar._contset[s] = loc
104                    _dim = s.dimen
105                    if _dim is None:
106                        raise DAE_Error(
107                            "The variable %s is indexed by a Set (%s) with a "
108                            "non-fixed dimension.  A DerivativeVar may only be "
109                            "indexed by Sets with constant dimension"
110                            % (sVar, s.name))
111                    elif _dim is UnknownSetDimen:
112                        raise DAE_Error(
113                            "The variable %s is indexed by a Set (%s) with an "
114                            "unknown dimension.  A DerivativeVar may only be "
115                            "indexed by Sets with known constant dimension"
116                            % (sVar, s.name))
117                    loc += s.dimen
118            num_contset = len(sVar._contset)
119
120        if num_contset == 0:
121            raise DAE_Error(
122                "The variable %s is not indexed by any ContinuousSets. A "
123                "derivative may only be taken with respect to a continuous "
124                "domain" % sVar)
125
126        if wrt is None:
127            # Check to be sure Var is indexed by single ContinuousSet and take
128            # first deriv wrt that set
129            if num_contset != 1:
130                raise DAE_Error(
131                    "The variable %s is indexed by multiple ContinuousSets. "
132                    "The desired ContinuousSet must be specified using the "
133                    "keyword argument 'wrt'" % sVar)
134            wrt = [next(iter(sVar._contset.keys())), ]
135        elif type(wrt) is ContinuousSet:
136            if wrt not in sVar._contset:
137                raise DAE_Error(
138                    "Invalid derivative: The variable %s is not indexed by "
139                    "the ContinuousSet %s" % (sVar, wrt))
140            wrt = [wrt, ]
141        elif type(wrt) is tuple or type(wrt) is list:
142            for i in wrt:
143                if type(i) is not ContinuousSet:
144                    raise DAE_Error(
145                        "Cannot take the derivative with respect to %s. "
146                        "Expected a ContinuousSet or a tuple of "
147                        "ContinuousSets" % i)
148                if i not in sVar._contset:
149                    raise DAE_Error(
150                        "Invalid derivative: The variable %s is not indexed "
151                        "by the ContinuousSet %s" % (sVar, i))
152            wrt = list(wrt)
153        else:
154            raise DAE_Error(
155                "Cannot take the derivative with respect to %s. "
156                "Expected a ContinuousSet or a tuple of ContinuousSets" % i)
157
158        wrtkey = [str(i) for i in wrt]
159        wrtkey.sort()
160        wrtkey = tuple(wrtkey)
161
162        if wrtkey in sVar._derivative:
163            raise DAE_Error(
164                "Cannot create a new derivative variable for variable "
165                "%s: derivative already defined as %s"
166                % (sVar.name, sVar._derivative[wrtkey]().name))
167
168        sVar._derivative[wrtkey] = weakref.ref(self)
169        self._sVar = sVar
170        self._wrt = wrt
171
172        kwds.setdefault('ctype', DerivativeVar)
173
174        Var.__init__(self,sVar.index_set(),**kwds)
175
176
177    def get_continuousset_list(self):
178        """ Return the a list of :py:class:`ContinuousSet` components the
179        derivative is being taken with respect to.
180
181        Returns
182        -------
183        `list`
184        """
185        return self._wrt
186
187    def is_fully_discretized(self):
188        """
189        Check to see if all the
190        :py:class:`ContinuousSets<pyomo.dae.ContinuousSet>` this derivative
191        is taken with respect to have been discretized.
192
193        Returns
194        -------
195        `boolean`
196        """
197        for i in self._wrt:
198            if 'scheme' not in i.get_discretization_info():
199                return False
200        return True
201
202    def get_state_var(self):
203        """ Return the :py:class:`Var` that is being differentiated.
204
205        Returns
206        -------
207        :py:class:`Var<pyomo.environ.Var>`
208        """
209        return self._sVar
210
211    def get_derivative_expression(self):
212        """
213        Returns the current discretization expression for this derivative or
214        creates an access function to its :py:class:`Var` the first time
215        this method is called. The expression gets built up as the
216        discretization transformations are sequentially applied to each
217        :py:class:`ContinuousSet` in the model.
218        """
219        try:
220            return self._expr
221        except:
222            self._expr = create_access_function(self._sVar)
223            return self._expr
224
225    def set_derivative_expression(self, expr):
226        """ Sets``_expr``, an expression representing the discretization
227        equations linking the :class:`DerivativeVar` to its state
228        :class:`Var`
229        """
230        self._expr = expr
231
232