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
11from pyomo.common import DeveloperError
12from pyomo.common.collections import ComponentMap
13from pyomo.common.dependencies import attempt_import
14from pyomo.common.errors import NondifferentiableError
15from pyomo.core.expr import current as EXPR, native_types
16from pyomo.core.expr.numvalue import value
17
18#
19# Sympy takes a significant time to load; defer importing it unless
20# someone actually needs the interface.
21#
22
23_operatorMap = {}
24_pyomo_operator_map = {}
25_functionMap = {}
26
27def _configure_sympy(sympy, available):
28    if not available:
29        return
30
31    _operatorMap.update({
32        sympy.Add: _sum,
33        sympy.Mul: _prod,
34        sympy.Pow: lambda x, y: x**y,
35        sympy.exp: lambda x: EXPR.exp(x),
36        sympy.log: lambda x: EXPR.log(x),
37        sympy.sin: lambda x: EXPR.sin(x),
38        sympy.asin: lambda x: EXPR.asin(x),
39        sympy.sinh: lambda x: EXPR.sinh(x),
40        sympy.asinh: lambda x: EXPR.asinh(x),
41        sympy.cos: lambda x: EXPR.cos(x),
42        sympy.acos: lambda x: EXPR.acos(x),
43        sympy.cosh: lambda x: EXPR.cosh(x),
44        sympy.acosh: lambda x: EXPR.acosh(x),
45        sympy.tan: lambda x: EXPR.tan(x),
46        sympy.atan: lambda x: EXPR.atan(x),
47        sympy.tanh: lambda x: EXPR.tanh(x),
48        sympy.atanh: lambda x: EXPR.atanh(x),
49        sympy.ceiling: lambda x: EXPR.ceil(x),
50        sympy.floor: lambda x: EXPR.floor(x),
51        sympy.sqrt: lambda x: EXPR.sqrt(x),
52        sympy.Abs: lambda x: abs(x),
53        sympy.Derivative: _nondifferentiable,
54        sympy.Tuple: lambda *x: x,
55    })
56
57    _pyomo_operator_map.update({
58        EXPR.SumExpression: sympy.Add,
59        EXPR.ProductExpression: sympy.Mul,
60        EXPR.NPV_ProductExpression: sympy.Mul,
61        EXPR.MonomialTermExpression: sympy.Mul,
62    })
63
64    _functionMap.update({
65        'exp': sympy.exp,
66        'log': sympy.log,
67        'log10': lambda x: sympy.log(x)/sympy.log(10),
68        'sin': sympy.sin,
69        'asin': sympy.asin,
70        'sinh': sympy.sinh,
71        'asinh': sympy.asinh,
72        'cos': sympy.cos,
73        'acos': sympy.acos,
74        'cosh': sympy.cosh,
75        'acosh': sympy.acosh,
76        'tan': sympy.tan,
77        'atan': sympy.atan,
78        'tanh': sympy.tanh,
79        'atanh': sympy.atanh,
80        'ceil': sympy.ceiling,
81        'floor': sympy.floor,
82        'sqrt': sympy.sqrt,
83    })
84
85sympy, sympy_available = attempt_import('sympy', callback=_configure_sympy)
86
87
88def _prod(*x):
89    ans = x[0]
90    for i in x[1:]:
91        ans *= i
92    return ans
93
94def _sum(*x):
95    return sum(x_ for x_ in x)
96
97def _nondifferentiable(*x):
98    if type(x[1]) is tuple:
99        # sympy >= 1.3 returns tuples (var, order)
100        wrt = x[1][0]
101    else:
102        # early versions of sympy returned the bare var
103        wrt = x[1]
104    raise NondifferentiableError(
105        "The sub-expression '%s' is not differentiable with respect to %s"
106        % (x[0], wrt) )
107
108class PyomoSympyBimap(object):
109    def __init__(self):
110        self.pyomo2sympy = ComponentMap()
111        self.sympy2pyomo = {}
112        self.i = 0
113
114    def getPyomoSymbol(self, sympy_object, default=None):
115        return self.sympy2pyomo.get(sympy_object, default)
116
117    def getSympySymbol(self, pyomo_object):
118        if pyomo_object in self.pyomo2sympy:
119            return self.pyomo2sympy[pyomo_object]
120        # Pyomo currently ONLY supports Real variables (not complex
121        # variables).  If that ever changes, then we will need to
122        # revisit hard-coding the symbol type here
123        sympy_obj = sympy.Symbol("x%d" % self.i, real=True)
124        self.i += 1
125        self.pyomo2sympy[pyomo_object] = sympy_obj
126        self.sympy2pyomo[sympy_obj] = pyomo_object
127        return sympy_obj
128
129    def sympyVars(self):
130        return self.sympy2pyomo.keys()
131
132# =====================================================
133# sympyify_expression
134# =====================================================
135
136class Pyomo2SympyVisitor(EXPR.StreamBasedExpressionVisitor):
137
138    def __init__(self, object_map):
139        sympy.Add  # this ensures _configure_sympy gets run
140        super(Pyomo2SympyVisitor, self).__init__()
141        self.object_map = object_map
142
143    def initializeWalker(self, expr):
144        return self.beforeChild(None, expr, None)
145
146    def exitNode(self, node, values):
147        if node.__class__ is EXPR.UnaryFunctionExpression:
148            return _functionMap[node._name](values[0])
149        _op = _pyomo_operator_map.get(node.__class__, None)
150        if _op is None:
151            return node._apply_operation(values)
152        else:
153            return _op(*tuple(values))
154
155    def beforeChild(self, node, child, child_idx):
156        #
157        # Don't replace native or sympy types
158        #
159        if type(child) in native_types:
160            return False, child
161        #
162        # We will descend into all expressions...
163        #
164        if child.is_expression_type():
165            return True, None
166        #
167        # Replace pyomo variables with sympy variables
168        #
169        if child.is_potentially_variable():
170            return False, self.object_map.getSympySymbol(child)
171        #
172        # Everything else is a constant...
173        #
174        return False, value(child)
175
176class Sympy2PyomoVisitor(EXPR.StreamBasedExpressionVisitor):
177
178    def __init__(self, object_map):
179        sympy.Add  # this ensures _configure_sympy gets run
180        super(Sympy2PyomoVisitor, self).__init__()
181        self.object_map = object_map
182
183    def initializeWalker(self, expr):
184        return self.beforeChild(None, expr, None)
185
186    def enterNode(self, node):
187        return (node._args, [])
188
189    def exitNode(self, node, values):
190        """ Visit nodes that have been expanded """
191        _sympyOp = node
192        _op = _operatorMap.get( type(_sympyOp), None )
193        if _op is None:
194            raise DeveloperError(
195                "sympy expression type '%s' not found in the operator "
196                "map" % type(_sympyOp) )
197        return _op(*tuple(values))
198
199    def beforeChild(self, node, child, child_idx):
200        if not child._args:
201            item = self.object_map.getPyomoSymbol(child, None)
202            if item is None:
203                item = float(child.evalf())
204            return False, item
205        return True, None
206
207def sympyify_expression(expr):
208    """Convert a Pyomo expression to a Sympy expression"""
209    #
210    # Create the visitor and call it.
211    #
212    object_map = PyomoSympyBimap()
213    visitor = Pyomo2SympyVisitor(object_map)
214    return object_map, visitor.walk_expression(expr)
215
216
217def sympy2pyomo_expression(expr, object_map):
218    visitor = Sympy2PyomoVisitor(object_map)
219    return visitor.walk_expression(expr)
220