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 __future__ import print_function
12import pyomo.common.unittest as unittest
13
14from pyomo.environ import (Var, Set, ConcreteModel,
15                           TransformationFactory)
16from pyomo.dae import ContinuousSet, DerivativeVar
17from pyomo.dae.diffvar import DAE_Error
18
19from io import StringIO
20
21from pyomo.common.log import LoggingIntercept
22
23from os.path import abspath, dirname, normpath, join
24currdir = dirname(abspath(__file__))
25exdir = normpath(join(currdir, '..', '..', '..', 'examples', 'dae'))
26
27
28class TestFiniteDiff(unittest.TestCase):
29    """
30    Class for testing the pyomo.DAE finite difference discretization
31    """
32
33    def setUp(self):
34        """
35        Setting up testing model
36        """
37        self.m = m = ConcreteModel()
38        m.t = ContinuousSet(bounds=(0, 10))
39        m.v1 = Var(m.t)
40        m.dv1 = DerivativeVar(m.v1)
41        m.s = Set(initialize=[1, 2, 3], ordered=True)
42
43    # test backward finite difference discretization
44    # on var indexed by single ContinuousSet
45    def test_disc_single_index_backward(self):
46        m = self.m.clone()
47        disc = TransformationFactory('dae.finite_difference')
48        disc.apply_to(m, nfe=5)
49
50        self.assertTrue(hasattr(m, 'dv1_disc_eq'))
51        self.assertEqual(len(m.dv1_disc_eq), 5)
52        self.assertEqual(len(m.v1), 6)
53
54        expected_disc_points = [0, 2.0, 4.0, 6.0, 8.0, 10]
55        disc_info = m.t.get_discretization_info()
56
57        self.assertEqual(disc_info['scheme'], 'BACKWARD Difference')
58
59        for idx, val in enumerate(list(m.t)):
60            self.assertAlmostEqual(val, expected_disc_points[idx])
61
62        self.assertTrue(hasattr(m, '_pyomo_dae_reclassified_derivativevars'))
63        self.assertIn(m.dv1, m._pyomo_dae_reclassified_derivativevars)
64
65        output = \
66"""\
67dv1_disc_eq : Size=5, Index=t, Active=True
68    Key : Lower : Body                               : Upper : Active
69    2.0 :   0.0 :   dv1[2.0] - 0.5*(v1[2.0] - v1[0]) :   0.0 :   True
70    4.0 :   0.0 : dv1[4.0] - 0.5*(v1[4.0] - v1[2.0]) :   0.0 :   True
71    6.0 :   0.0 : dv1[6.0] - 0.5*(v1[6.0] - v1[4.0]) :   0.0 :   True
72    8.0 :   0.0 : dv1[8.0] - 0.5*(v1[8.0] - v1[6.0]) :   0.0 :   True
73     10 :   0.0 :   dv1[10] - 0.5*(v1[10] - v1[8.0]) :   0.0 :   True
74"""
75        out = StringIO()
76        m.dv1_disc_eq.pprint(ostream=out)
77        self.assertEqual(output, out.getvalue())
78
79    # test backward finite difference discretization on
80    # second order derivative indexed by single ContinuousSet
81    def test_disc_second_order_backward(self):
82        m = self.m.clone()
83        m.dv1dt2 = DerivativeVar(m.v1, wrt=(m.t, m.t))
84        disc = TransformationFactory('dae.finite_difference')
85        disc.apply_to(m, nfe=2)
86
87        self.assertTrue(hasattr(m, 'dv1dt2_disc_eq'))
88        self.assertEqual(len(m.dv1dt2_disc_eq), 1)
89        self.assertEqual(len(m.v1), 3)
90
91        self.assertTrue(hasattr(m, '_pyomo_dae_reclassified_derivativevars'))
92        self.assertIn(m.dv1, m._pyomo_dae_reclassified_derivativevars)
93        self.assertIn(m.dv1dt2, m._pyomo_dae_reclassified_derivativevars)
94
95        output = \
96"""\
97dv1dt2_disc_eq : Size=1, Index=t, Active=True
98    Key : Lower : Body                                           : Upper : Active
99     10 :   0.0 : dv1dt2[10] - 0.04*(v1[10] - 2*v1[5.0] + v1[0]) :   0.0 :   True
100"""
101        out = StringIO()
102        m.dv1dt2_disc_eq.pprint(ostream=out)
103        self.assertEqual(output, out.getvalue())
104
105    # test forward finite difference discretization
106    # on var indexed by single ContinuousSet
107    def test_disc_single_index_forward(self):
108        m = self.m.clone()
109        disc = TransformationFactory('dae.finite_difference')
110        disc.apply_to(m, nfe=5, scheme='FORWARD')
111
112        self.assertTrue(hasattr(m, 'dv1_disc_eq'))
113        self.assertEqual(len(m.dv1_disc_eq), 5)
114        self.assertEqual(len(m.v1), 6)
115
116        expected_disc_points = [0, 2.0, 4.0, 6.0, 8.0, 10]
117        disc_info = m.t.get_discretization_info()
118
119        self.assertEqual(disc_info['scheme'], 'FORWARD Difference')
120
121        for idx, val in enumerate(list(m.t)):
122            self.assertAlmostEqual(val, expected_disc_points[idx])
123
124        self.assertTrue(hasattr(m, '_pyomo_dae_reclassified_derivativevars'))
125        self.assertIn(m.dv1, m._pyomo_dae_reclassified_derivativevars)
126
127        output = \
128"""\
129dv1_disc_eq : Size=5, Index=t, Active=True
130    Key : Lower : Body                               : Upper : Active
131      0 :   0.0 :     dv1[0] - 0.5*(v1[2.0] - v1[0]) :   0.0 :   True
132    2.0 :   0.0 : dv1[2.0] - 0.5*(v1[4.0] - v1[2.0]) :   0.0 :   True
133    4.0 :   0.0 : dv1[4.0] - 0.5*(v1[6.0] - v1[4.0]) :   0.0 :   True
134    6.0 :   0.0 : dv1[6.0] - 0.5*(v1[8.0] - v1[6.0]) :   0.0 :   True
135    8.0 :   0.0 :  dv1[8.0] - 0.5*(v1[10] - v1[8.0]) :   0.0 :   True
136"""
137        out = StringIO()
138        m.dv1_disc_eq.pprint(ostream=out)
139        self.assertEqual(output, out.getvalue())
140
141    # test forward finite difference discretization
142    # second order derivative indexed by single ContinuousSet
143    def test_disc_second_order_forward(self):
144        m = self.m.clone()
145        m.dv1dt2 = DerivativeVar(m.v1, wrt=(m.t, m.t))
146        disc = TransformationFactory('dae.finite_difference')
147        disc.apply_to(m, nfe=2, scheme='FORWARD')
148
149        self.assertTrue(hasattr(m, 'dv1dt2_disc_eq'))
150        self.assertEqual(len(m.dv1dt2_disc_eq), 1)
151        self.assertEqual(len(m.v1), 3)
152
153        self.assertTrue(hasattr(m, '_pyomo_dae_reclassified_derivativevars'))
154        self.assertIn(m.dv1, m._pyomo_dae_reclassified_derivativevars)
155        self.assertIn(m.dv1dt2, m._pyomo_dae_reclassified_derivativevars)
156
157        output = \
158"""\
159dv1dt2_disc_eq : Size=1, Index=t, Active=True
160    Key : Lower : Body                                          : Upper : Active
161      0 :   0.0 : dv1dt2[0] - 0.04*(v1[10] - 2*v1[5.0] + v1[0]) :   0.0 :   True
162"""
163        out = StringIO()
164        m.dv1dt2_disc_eq.pprint(ostream=out)
165        self.assertEqual(output, out.getvalue())
166
167    # test central finite difference discretization
168    # on var indexed by single ContinuousSet
169    def test_disc_single_index_central(self):
170        m = self.m.clone()
171        disc = TransformationFactory('dae.finite_difference')
172        disc.apply_to(m, nfe=5, scheme='CENTRAL')
173
174        self.assertTrue(hasattr(m, 'dv1_disc_eq'))
175        self.assertEqual(len(m.dv1_disc_eq), 4)
176        self.assertEqual(len(m.v1), 6)
177
178        expected_disc_points = [0, 2.0, 4.0, 6.0, 8.0, 10]
179        disc_info = m.t.get_discretization_info()
180
181        self.assertEqual(disc_info['scheme'], 'CENTRAL Difference')
182
183        for idx, val in enumerate(list(m.t)):
184            self.assertAlmostEqual(val, expected_disc_points[idx])
185
186        output = \
187"""\
188dv1_disc_eq : Size=4, Index=t, Active=True
189    Key : Lower : Body                                : Upper : Active
190    2.0 :   0.0 :   dv1[2.0] - 0.25*(v1[4.0] - v1[0]) :   0.0 :   True
191    4.0 :   0.0 : dv1[4.0] - 0.25*(v1[6.0] - v1[2.0]) :   0.0 :   True
192    6.0 :   0.0 : dv1[6.0] - 0.25*(v1[8.0] - v1[4.0]) :   0.0 :   True
193    8.0 :   0.0 :  dv1[8.0] - 0.25*(v1[10] - v1[6.0]) :   0.0 :   True
194"""
195        out = StringIO()
196        m.dv1_disc_eq.pprint(ostream=out)
197        self.assertEqual(output, out.getvalue())
198
199    # test central finite difference discretization
200    # second order derivative indexed by single ContinuousSet
201    def test_disc_second_order_central(self):
202        m = self.m.clone()
203        m.dv1dt2 = DerivativeVar(m.v1, wrt=(m.t, m.t))
204        disc = TransformationFactory('dae.finite_difference')
205        disc.apply_to(m, nfe=2, scheme='CENTRAL')
206
207        self.assertTrue(hasattr(m, 'dv1dt2_disc_eq'))
208        self.assertEqual(len(m.dv1dt2_disc_eq), 1)
209        self.assertEqual(len(m.v1), 3)
210
211        output = \
212"""\
213dv1dt2_disc_eq : Size=1, Index=t, Active=True
214    Key : Lower : Body                                            : Upper : Active
215    5.0 :   0.0 : dv1dt2[5.0] - 0.04*(v1[10] - 2*v1[5.0] + v1[0]) :   0.0 :   True
216"""
217        out = StringIO()
218        m.dv1dt2_disc_eq.pprint(ostream=out)
219        self.assertEqual(output, out.getvalue())
220
221    # test collocation discretization on var indexed by ContinuousSet and Set
222    def test_disc_multi_index(self):
223        m = self.m.clone()
224        m.v2 = Var(m.t, m.s)
225        m.dv2 = DerivativeVar(m.v2)
226
227        disc = TransformationFactory('dae.finite_difference')
228        disc.apply_to(m, nfe=5)
229
230        self.assertTrue(hasattr(m, 'dv1_disc_eq'))
231        self.assertTrue(hasattr(m, 'dv2_disc_eq'))
232        self.assertEqual(len(m.dv2_disc_eq), 15)
233        self.assertEqual(len(m.v2), 18)
234
235        expected_disc_points = [0, 2.0, 4.0, 6.0, 8.0, 10]
236        disc_info = m.t.get_discretization_info()
237
238        self.assertEqual(disc_info['scheme'], 'BACKWARD Difference')
239
240        for idx, val in enumerate(list(m.t)):
241            self.assertAlmostEqual(val, expected_disc_points[idx])
242
243    # test collocation discretization on var indexed by multiple ContinuousSets
244    def test_disc_multi_index2(self):
245        m = self.m.clone()
246        m.t2 = ContinuousSet(bounds=(0, 5))
247        m.v2 = Var(m.t, m.t2)
248        m.dv2dt = DerivativeVar(m.v2, wrt=m.t)
249        m.dv2dt2 = DerivativeVar(m.v2, wrt=m.t2)
250
251        disc = TransformationFactory('dae.finite_difference')
252        disc.apply_to(m, nfe=2)
253
254        self.assertTrue(hasattr(m, 'dv2dt_disc_eq'))
255        self.assertTrue(hasattr(m, 'dv2dt2_disc_eq'))
256        self.assertEqual(len(m.dv2dt_disc_eq), 6)
257        self.assertEqual(len(m.dv2dt2_disc_eq), 6)
258        self.assertEqual(len(m.v2), 9)
259
260        expected_t_disc_points = [0, 5.0, 10]
261        expected_t2_disc_points = [0, 2.5, 5]
262
263        for idx, val in enumerate(list(m.t)):
264            self.assertAlmostEqual(val, expected_t_disc_points[idx])
265
266        for idx, val in enumerate(list(m.t2)):
267            self.assertAlmostEqual(val, expected_t2_disc_points[idx])
268
269    # test collocation discretization on var indexed by ContinuousSet and
270    # multi-dimensional Set
271    def test_disc_multidimen_index(self):
272        m = self.m.clone()
273        m.s2 = Set(initialize=[('A', 'B'), ('C', 'D'), ('E', 'F')])
274        m.v2 = Var(m.t, m.s2)
275        m.dv2 = DerivativeVar(m.v2)
276        m.v3 = Var(m.s2, m.t)
277        m.dv3 = DerivativeVar(m.v3)
278
279        disc = TransformationFactory('dae.finite_difference')
280        disc.apply_to(m, nfe=5)
281
282        self.assertTrue(hasattr(m, 'dv1_disc_eq'))
283        self.assertTrue(hasattr(m, 'dv2_disc_eq'))
284        self.assertTrue(hasattr(m, 'dv3_disc_eq'))
285        self.assertEqual(len(m.dv2_disc_eq), 15)
286        self.assertEqual(len(m.v2), 18)
287        self.assertEqual(len(m.dv3_disc_eq), 15)
288        self.assertEqual(len(m.v3), 18)
289
290        expected_disc_points = [0, 2.0, 4.0, 6.0, 8.0, 10]
291        disc_info = m.t.get_discretization_info()
292
293        self.assertEqual(disc_info['scheme'], 'BACKWARD Difference')
294
295        for idx, val in enumerate(list(m.t)):
296            self.assertAlmostEqual(val, expected_disc_points[idx])
297
298    # test passing the discretization invalid options
299    def test_disc_invalid_options(self):
300        m = self.m.clone()
301
302        with self.assertRaises(TypeError):
303            TransformationFactory('dae.finite_difference').apply_to(m, wrt=m.s)
304
305        with self.assertRaises(ValueError):
306            TransformationFactory('dae.finite_difference').apply_to(m, nfe=-1)
307
308        with self.assertRaises(ValueError):
309            TransformationFactory('dae.finite_difference').apply_to(m,
310                                                                    scheme='foo')
311
312        with self.assertRaises(ValueError):
313            TransformationFactory('dae.finite_difference').apply_to(m,
314                                                                    foo=True)
315
316        TransformationFactory('dae.finite_difference').apply_to(m, wrt=m.t)
317        with self.assertRaises(ValueError):
318            TransformationFactory('dae.finite_difference').apply_to(m, wrt=m.t)
319
320        m = self.m.clone()
321        disc = TransformationFactory('dae.finite_difference')
322        disc.apply_to(m)
323        with self.assertRaises(ValueError):
324            disc.apply_to(m)
325
326    # test discretization using fewer points than ContinuousSet initialized
327    # with
328    def test_initialized_continuous_set(self):
329        m = ConcreteModel()
330        m.t = ContinuousSet(initialize=[0, 1, 2, 3, 4])
331        m.v = Var(m.t)
332        m.dv = DerivativeVar(m.v)
333
334        log_out = StringIO()
335        with LoggingIntercept(log_out, 'pyomo.dae'):
336            TransformationFactory('dae.finite_difference').apply_to(m, nfe=2)
337        self.assertIn('More finite elements', log_out.getvalue())
338
339    # Test discretizing an invalid derivative
340    def test_invalid_derivative(self):
341        m = ConcreteModel()
342        m.t = ContinuousSet(bounds=(0, 10))
343        m.v = Var(m.t)
344        m.dv = DerivativeVar(m.v, wrt=(m.t, m.t, m.t))
345
346        with self.assertRaises(DAE_Error):
347            TransformationFactory('dae.finite_difference').apply_to(m)
348
349    # test trying to discretize a ContinuousSet twice
350    def test_discretize_twice(self):
351        m = self.m.clone()
352
353        disc1 = TransformationFactory('dae.finite_difference')
354        disc1.apply_to(m, nfe=5)
355
356        disc2 = TransformationFactory('dae.finite_difference')
357
358        with self.assertRaises(DAE_Error):
359            disc2.apply_to(m, nfe=5)
360
361
362if __name__ == "__main__":
363    unittest.main()
364