Repository: Pyomo/pyomo

Issue:
DAE Transformations for Block-Derived components
The DAE discretization transformations currently work for ContinuousSet indexed Blocks but not components derived from Blocks. This was intentional because the user could overwrite the Block construct method and then the _update_block function in the discretization would no longer be valid. We should add a flag to Block-derived components that signals to the discretization transformation if it is safe to update the component like a normal Block. This could be as simple as setting/checking the ctype but there might be edge cases where a discretization-specific flag would be better. 

In the case where the Block construct method is overwritten, it might be useful to give the user a way to pass their own _update_me function to the discretizer.


Discussion:

Resolving patch:
diff --git a/pyomo/dae/misc.py b/pyomo/dae/misc.py
--- a/pyomo/dae/misc.py
+++ b/pyomo/dae/misc.py
@@ -127,10 +127,10 @@ def update_contset_indexed_component(comp):
                     _update_constraint(comp)
                 elif comp.type() == Expression:
                     _update_expression(comp)
-                elif type(comp) is IndexedBlock: # Only catch Blocks and not components derived
-                    _update_block(comp)          # from Blocks
-                elif isinstance(comp, Piecewise):
+                elif isinstance(comp, Piecewise): 
                     _update_piecewise(comp)
+                elif comp.type() == Block:
+                    _update_block(comp)   
                 else:
                     raise TypeError("Found component %s of type %s indexed "\
                         "by a ContinuousSet. Components of this type are "\
@@ -150,10 +150,10 @@ def update_contset_indexed_component(comp):
                     _update_constraint(comp)
                 elif comp.type() == Expression:
                     _update_expression(comp)
-                elif type(comp) is IndexedBlock: # Only catch Blocks and not components derived
-                    _update_block(comp)          # from Blocks
                 elif isinstance(comp, Piecewise):
                     _update_piecewise(comp)
+                elif comp.type() == Block: 
+                    _update_block(comp)    
                 else:
                     raise TypeError("Found component %s of type %s indexed "\
                         "by a ContinuousSet. Components of this type are "\
@@ -205,9 +205,37 @@ def _update_expression(expre):
 def _update_block(blk):
     """
     This method will construct any additional indices in a block
-    resulting from the discretization of a ContinuousSet.
+    resulting from the discretization of a ContinuousSet. For
+    Block-derived components we check if the Block construct method has
+    been overridden. If not then we update it like a regular block. If
+    construct has been overridden then we try to call the component's
+    update_after_discretization method. If the component hasn't
+    implemented this method then we throw a warning and try to update it
+    like a normal block. The issue, when construct is overridden, is that
+    anything could be happening and we can't automatically assume that
+    treating the block-derived component like a normal block will be
+    sufficient to update it correctly.
+
     """
     
+    # Check if Block construct method is overridden
+    # getattr needed below for Python 2, 3 compatibility
+    if blk.construct.__func__ is not getattr(IndexedBlock.construct,
+                                             '__func__', IndexedBlock.construct):
+        # check for custom update function
+        try :
+            blk.update_after_discretization()
+            return
+        except AttributeError:
+            logger.warning('DAE(misc): Attempting to apply a discretization transformation to '
+                  'the Block-derived component "%s". The component overrides the '
+                  'Block construct method but no update_after_discretization() '
+                  'function was found. Will attempt to update as a standard Block '
+                  'but user should verify that the component was expanded correctly. '
+                  'To suppress this warning, please provide an '
+                  'update_after_discretization() function on Block-derived '
+                  'components that override construct()' %(blk.name))            
+
     # Code taken from the construct() method of Block
     missing_idx = set(blk._index)-set(iterkeys(blk._data))
     for idx in list(missing_idx):


Test patch:
diff --git a/pyomo/dae/tests/test_misc.py b/pyomo/dae/tests/test_misc.py
--- a/pyomo/dae/tests/test_misc.py
+++ b/pyomo/dae/tests/test_misc.py
@@ -11,16 +11,40 @@
 # Unit Tests for pyomo.dae.misc
 #
 
+import logging
 import os
 from os.path import abspath, dirname
 currdir = dirname(abspath(__file__))+os.sep
 
+from six import StringIO
+
 import pyutilib.th as unittest
 
 from pyomo.environ import *
 from pyomo.dae import *
 from pyomo.dae.misc import *
 
+# FIXME: This class was copied from test_block.py. It should be moved
+# to a more general location
+class LoggingIntercept(object):
+    def __init__(self, output, module=None, level=logging.WARNING):
+        self.handler = logging.StreamHandler(output)
+        self.handler.setFormatter(logging.Formatter('%(message)s'))
+        self.handler.setLevel(level)
+        self.level = level
+        self.module = module
+
+    def __enter__(self):
+        logger = logging.getLogger(self.module)
+        self.level = logger.level
+        logger.setLevel(self.handler.level)
+        logger.addHandler(self.handler)
+
+    def __exit__(self, et, ev, tb):
+        logger = logging.getLogger(self.module)
+        logger.removeHandler(self.handler)
+        logger.setLevel(self.level)
+
 class TestDaeMisc(unittest.TestCase):
     
     # test generate_finite_elements method
@@ -708,6 +732,154 @@ def test_update_contset_indexed_component_unsupported_multiple(self):
 
         except TypeError:
             pass  
+
+    def test_update_block_derived(self):
+        class Foo(Block):
+            pass
+
+        m = ConcreteModel()
+        m.t = ContinuousSet(bounds=(0,10))
+        def _block_rule(b, t):
+            m = b.model()
+            def _init(m,j):
+                return j*2
+            b.p1 = Param(m.t, default=_init)
+            b.v1 = Var(m.t, initialize=5)
+        m.foo = Foo(m.t, rule=_block_rule)
+        generate_finite_elements(m.t, 5)
+        update_contset_indexed_component(m.foo)
+
+        self.assertEqual(len(m.foo),6)
+        self.assertEqual(len(m.foo[0].p1),6)
+        self.assertEqual(len(m.foo[2].v1),6)
+        self.assertEqual(m.foo[0].p1[6],12)
+
+    def test_update_block_derived_override_construct_nofcn(self):
+        class Foo(Block):
+
+            def construct(self, data=None):
+                Block.construct(self, data)
+
+        m = ConcreteModel()
+        m.t = ContinuousSet(bounds=(0,10))
+        def _block_rule(b, t):
+            m = b.model()
+            def _init(m,j):
+                return j*2
+            b.p1 = Param(m.t, default=_init)
+            b.v1 = Var(m.t, initialize=5)
+        m.foo = Foo(m.t, rule=_block_rule)
+        generate_finite_elements(m.t, 5)
+
+        OUTPUT = StringIO()
+        with LoggingIntercept(OUTPUT, 'pyomo.core'):
+            update_contset_indexed_component(m.foo)
+        self.assertIn('transformation to the Block-derived component',
+                      OUTPUT.getvalue())
+        self.assertEqual(len(m.foo),6)
+        self.assertEqual(len(m.foo[0].p1),6)
+        self.assertEqual(len(m.foo[2].v1),6)
+        self.assertEqual(m.foo[0].p1[6],12)
+
+    def test_update_block_derived_override_construct_withfcn(self):
+        class Foo(Block):
+            updated = False
+            def construct(self, data=None):
+                Block.construct(self, data)
+            
+            def update_after_discretization(self):
+                self.updated = True
+        m = ConcreteModel()
+        m.t = ContinuousSet(bounds=(0,10))
+        def _block_rule(b, t):
+            m = b.model()
+            def _init(m,j):
+                return j*2
+            b.p1 = Param(m.t, default=_init)
+            b.v1 = Var(m.t, initialize=5)
+        m.foo = Foo(m.t, rule=_block_rule)
+        generate_finite_elements(m.t, 5)
+        update_contset_indexed_component(m.foo)
+
+        self.assertTrue(m.foo.updated)
+        self.assertEqual(len(m.foo),2)
+        self.assertEqual(len(m.foo[0].v1),2)
+
+
+    def test_update_block_derived2(self):
+        class Foo(Block):
+            pass
+
+        m = ConcreteModel()
+        m.t = ContinuousSet(bounds=(0,10))
+        m.s = Set(initialize=[1,2,3])
+        def _block_rule(b, t, s):
+            m = b.model()
+            def _init(m,j):
+                return j*2
+            b.p1 = Param(m.t, default=_init)
+            b.v1 = Var(m.t, initialize=5)
+        m.foo = Foo(m.t, m.s, rule=_block_rule)
+        generate_finite_elements(m.t, 5)
+        update_contset_indexed_component(m.foo)
+
+        self.assertEqual(len(m.foo),18)
+        self.assertEqual(len(m.foo[0,1].p1),6)
+        self.assertEqual(len(m.foo[2,2].v1),6)
+        self.assertEqual(m.foo[0,3].p1[6],12)
+
+    def test_update_block_derived_override_construct_nofcn2(self):
+        class Foo(Block):
+
+            def construct(self, data=None):
+                Block.construct(self, data)
+
+        m = ConcreteModel()
+        m.t = ContinuousSet(bounds=(0,10))
+        m.s = Set(initialize=[1,2,3])
+        def _block_rule(b, t, s):
+            m = b.model()
+            def _init(m,j):
+                return j*2
+            b.p1 = Param(m.t, default=_init)
+            b.v1 = Var(m.t, initialize=5)
+        m.foo = Foo(m.t, m.s, rule=_block_rule)
+        generate_finite_elements(m.t, 5)
+
+        OUTPUT = StringIO()
+        with LoggingIntercept(OUTPUT, 'pyomo.core'):
+            update_contset_indexed_component(m.foo)
+        self.assertIn('transformation to the Block-derived component',
+                      OUTPUT.getvalue())
+        self.assertEqual(len(m.foo),18)
+        self.assertEqual(len(m.foo[0,1].p1),6)
+        self.assertEqual(len(m.foo[2,2].v1),6)
+        self.assertEqual(m.foo[0,3].p1[6],12)
+
+    def test_update_block_derived_override_construct_withfcn2(self):
+        class Foo(Block):
+            updated = False
+            def construct(self, data=None):
+                Block.construct(self, data)
+            
+            def update_after_discretization(self):
+                self.updated = True
+        m = ConcreteModel()
+        m.t = ContinuousSet(bounds=(0,10))
+        m.s = Set(initialize=[1,2,3])
+        def _block_rule(b, t, s):
+            m = b.model()
+            def _init(m,j):
+                return j*2
+            b.p1 = Param(m.t, default=_init)
+            b.v1 = Var(m.t, initialize=5)
+        m.foo = Foo(m.t, m.s, rule=_block_rule)
+        generate_finite_elements(m.t, 5)
+        update_contset_indexed_component(m.foo)
+
+        self.assertTrue(m.foo.updated)
+        self.assertEqual(len(m.foo),6)
+        self.assertEqual(len(m.foo[0,1].v1),2)
         
     # test add_equality_constraints method
     # def test_add_equality_constraints(self):