Merge pull request #331 from Kenneth-T-Moore/ken3

Setup refactor parts 2 and 3

openmdao/components/tests/test_exec_comp.py

@@ -63,6 +63,9 @@ def test_mixed_type(self):
                                                      x=np.arange(10,dtype=float)))
         prob.setup(check=False)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         self.assertTrue('x' in C1._inputs)
         self.assertTrue('y' in C1._outputs)
 
@@ -77,6 +80,9 @@ def test_simple(self):
 
         prob.setup(check=False)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         self.assertTrue('x' in C1._inputs)
         self.assertTrue('y' in C1._outputs)
 
@@ -91,6 +97,9 @@ def test_for_spaces(self):
 
         prob.setup(check=False)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         self.assertTrue('x' in C1._inputs)
         self.assertTrue('y' in C1._outputs)
         self.assertTrue('pi' not in C1._inputs)
@@ -122,6 +131,9 @@ def test_math(self):
 
         prob.setup(check=False)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         self.assertTrue('x' in C1._inputs)
         self.assertTrue('y' in C1._outputs)
 
@@ -138,6 +150,9 @@ def test_array(self):
 
         prob.setup(check=False)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         self.assertTrue('x' in C1._inputs)
         self.assertTrue('y' in C1._outputs)
 
@@ -154,6 +169,9 @@ def test_array_lhs(self):
 
         prob.setup(check=False)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         self.assertTrue('x' in C1._inputs)
         self.assertTrue('y' in C1._outputs)
 
@@ -214,6 +232,9 @@ def test_complex_step(self):
 
         prob.setup(check=False)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         self.assertTrue('x' in C1._inputs)
         self.assertTrue('y' in C1._outputs)
 

openmdao/components/tests/test_meta_model.py

@@ -23,6 +23,9 @@ def test_sin_metamodel(self):
         testlogger = TestLogger()
         prob.setup(logger=testlogger)
 
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         msg = ("No default surrogate model is defined and the "
                "following outputs do not have a surrogate model:\n"
                "['f_x']\n"
@@ -72,6 +75,10 @@ def test_sin_metamodel_preset_data(self):
         # check that missing surrogate is detected in check_setup
         testlogger = TestLogger()
         prob.setup(logger=testlogger)
+
+        # Conclude setup but don't run model.
+        prob.final_setup()
+
         msg = ("No default surrogate model is defined and the "
                "following outputs do not have a surrogate model:\n"
                "['f_x']\n"

openmdao/core/problem.py

@@ -24,6 +24,7 @@
 from openmdao.utils.general_utils import warn_deprecation
 from openmdao.utils.logger_utils import get_logger
 from openmdao.utils.mpi import MPI, FakeComm
+from openmdao.utils.name_maps import prom_name2abs_name
 from openmdao.utils.graph_utils import all_connected_edges
 from openmdao.vectors.default_vector import DefaultVector
 from openmdao.vectors.default_multi_vector import DefaultMultiVector
@@ -66,6 +67,14 @@ class Problem(object):
         If True, allocate vectors to store ref. values.
     _solver_print_cache : list
         Allows solver iprints to be set to requested values after setup calls.
+    _initial_condition_cache : dict
+        Any initial conditions that are set at the problem level via setitem are cached here
+        until they can be processed.
+    _setup_status : int
+        Current status of the setup in _model.
+        0 -- Newly initialized problem or newly added model.
+        1 -- The `setup` method has been called, but vectors not initialized.
+        2 -- The `final_setup` has been run, everything ready to run.
     """
 
     def __init__(self, model=None, comm=None, use_ref_vector=True, root=None):
@@ -114,6 +123,14 @@ def __init__(self, model=None, comm=None, use_ref_vector=True, root=None):
 
         recording_iteration_stack = []
 
+        self._initial_condition_cache = {}
+
+        # Status of the setup of _model.
+        # 0 -- Newly initialized problem or newly added model.
+        # 1 -- The `setup` method has been called, but vectors not initialized.
+        # 2 -- The `final_setup` has been run, everything ready to run.
+        self._setup_status = 0
+
     def __getitem__(self, name):
         """
         Get an output/input variable.
@@ -128,14 +145,70 @@ def __getitem__(self, name):
         float or ndarray
             the requested output/input variable.
         """
-        if name in self.model._outputs:
-            return self.model._outputs[name]
+        # Caching only needed if vectors aren't allocated yet.
+        if self._setup_status == 1:
+
+            # We have set and cached already
+            if name in self._initial_condition_cache:
+                val = self._initial_condition_cache[name]
+
+            # Vector not setup, so we need to pull values from saved metadata request.
+            else:
+                proms = self.model._var_allprocs_prom2abs_list
+                meta = self.model._var_abs2meta
+                if name in meta['input']:
+                    if name in self.model._conn_abs_in2out:
+                        src_name = self.model._conn_abs_in2out[name]
+                        val = meta['input'][src_name]['value']
+                    else:
+                        val = meta['input'][name]['value']
+
+                elif name in meta['output']:
+                    val = meta['output'][name]['value']
+
+                elif name in proms['input']:
+                    abs_name = proms['input'][name][0]
+                    if abs_name in self.model._conn_abs_in2out:
+                        src_name = self.model._conn_abs_in2out[abs_name]
+                        # So, if the inputs and outputs are promoted to the same name, then we
+                        # allow getitem, but if they aren't, then we raise an error due to non
+                        # uniqueness.
+                        if name not in proms['output']:
+                            # This triggers a check for unconnected non-unique inputs, and
+                            # raises the same error as vector access.
+                            abs_name = prom_name2abs_name(self.model, name, 'input')
+                        val = meta['output'][src_name]['value']
+                    else:
+                        # This triggers a check for unconnected non-unique inputs, and
+                        # raises the same error as vector access.
+                        abs_name = prom_name2abs_name(self.model, name, 'input')
+                        val = meta['input'][abs_name]['value']
+
+                elif name in proms['output']:
+                    abs_name = prom_name2abs_name(self.model, name, 'output')
+                    val = meta['output'][abs_name]['value']
+
+                else:
+                    msg = 'Variable name "{}" not found.'
+                    raise KeyError(msg.format(name))
+
+                self._initial_condition_cache[name] = val
+
+        elif name in self.model._outputs:
+            val = self.model._outputs[name]
+
         elif name in self.model._inputs:
-            return self.model._inputs[name]
+            val = self.model._inputs[name]
+
         else:
             msg = 'Variable name "{}" not found.'
             raise KeyError(msg.format(name))
 
+        # Need to cache the "get" in case the user calls in-place numpy operations.
+        self._initial_condition_cache[name] = val
+
+        return val
+
     def __setitem__(self, name, value):
         """
         Set an output/input variable.
@@ -147,13 +220,27 @@ def __setitem__(self, name, value):
         value : float or ndarray or list
             value to set this variable to.
         """
-        if name in self.model._outputs:
-            self.model._outputs[name] = value
-        elif name in self.model._inputs:
-            self.model._inputs[name] = value
+        # Caching only needed if vectors aren't allocated yet.
+        if self._setup_status == 1:
+            self._initial_condition_cache[name] = value
         else:
-            msg = 'Variable name "{}" not found.'
-            raise KeyError(msg.format(name))
+            if name in self.model._outputs:
+                self.model._outputs[name] = value
+            elif name in self.model._inputs:
+                self.model._inputs[name] = value
+            else:
+                msg = 'Variable name "{}" not found.'
+                raise KeyError(msg.format(name))
+
+    def _set_initial_conditions(self):
+        """
+        Set all initial conditions that have been saved in cache after setup.
+        """
+        for name, value in iteritems(self._initial_condition_cache):
+            self[name] = value
+
+        # Clean up cache
+        self._initial_condition_cache = {}
 
     @property
     def root(self):
@@ -199,6 +286,7 @@ def run_model(self):
         if self._mode is None:
             raise RuntimeError("The `setup` method must be called before `run_model`.")
 
+        self.final_setup()
         return self.model.run_solve_nonlinear()
 
     def run_driver(self):
@@ -213,6 +301,7 @@ def run_driver(self):
         if self._mode is None:
             raise RuntimeError("The `setup` method must be called before `run_driver`.")
 
+        self.final_setup()
         with self.model._scaled_context_all():
             return self.driver.run()
 
@@ -257,7 +346,12 @@ def cleanup(self):
     def setup(self, vector_class=DefaultVector, check=True, logger=None, mode='auto',
               force_alloc_complex=False, multi_vector_class=None):
         """
-        Set up everything.
+        Set up the model hierarchy.
+
+        When `setup` is called, the model hierarchy is assembled, the processors are allocated
+        (for MPI), and variables and connections are all assigned. This method traverses down
+        the model hierarchy to call `setup` on each subsystem, and then traverses up te model
+        hierarchy to call `configure` on each subsystem.
 
         Parameters
         ----------
@@ -304,8 +398,42 @@ def setup(self, vector_class=DefaultVector, check=True, logger=None, mode='auto'
             mode = 'rev'
         self._mode = mode
 
-        model._setup(comm, vector_class, 'full', force_alloc_complex=force_alloc_complex,
-                     mode=mode, multi_vector_class=multi_vector_class)
+        model._setup(comm, 'full')
+
+        # Cache all args for final setup.
+        self._vector_class = vector_class
+        self._check = check
+        self._logger = logger
+        self._force_alloc_complex = force_alloc_complex
+        self._multi_vector_class = multi_vector_class
+
+        self._setup_status = 1
+        return self
+
+    def final_setup(self):
+        """
+        Perform final setup phase on problem in preparation for run.
+
+        This is the second phase of setup, and is done automatically at the start of `run_driver`
+        and `run_model`. At the beginning of final_setup, we have a model hierarchy with defined
+        variables, solvers, case_recorders, and derivative settings. During this phase, the vectors
+        are created and populated, the drivers and solvers are initialized, and the recorders are
+        started, and the rest of the framework is prepared for execution.
+        """
+        vector_class = self._vector_class
+        check = self._check
+        logger = self._logger
+        force_alloc_complex = self._force_alloc_complex
+        multi_vector_class = self._multi_vector_class
+
+        model = self.model
+        comm = self.comm
+        mode = self._mode
+
+        if self._setup_status < 2:
+            model._final_setup(comm, vector_class, 'full', force_alloc_complex=force_alloc_complex,
+                               mode=mode, multi_vector_class=multi_vector_class)
+
         self.driver._setup_driver(self)
 
         if isinstance(model, Group):
@@ -321,7 +449,9 @@ def setup(self, vector_class=DefaultVector, check=True, logger=None, mode='auto'
         if check and comm.rank == 0:
             check_config(self, logger)
 
-        return self
+        if self._setup_status < 2:
+            self._setup_status = 2
+            self._set_initial_conditions()
 
     def check_partials(self, logger=None, comps=None, compact_print=False,
                        abs_err_tol=1e-6, rel_err_tol=1e-6, global_options=None,
@@ -368,6 +498,9 @@ def check_partials(self, logger=None, comps=None, compact_print=False,
             For 'J_fd', 'J_fwd', 'J_rev' the value is: A numpy array representing the computed
                 Jacobian for the three different methods of computation.
         """
+        if self._setup_status < 2:
+            self.final_setup()
+
         if not global_options:
             global_options = DEFAULT_FD_OPTIONS.copy()
             global_options['method'] = 'fd'
@@ -726,6 +859,9 @@ def compute_total_derivs(self, of=None, wrt=None, return_format='flat_dict'):
         derivs : object
             Derivatives in form requested by 'return_format'.
         """
+        if self._setup_status < 2:
+            self.final_setup()
+
         with self.model._scaled_context_all():
             if self.model._owns_approx_jac:
                 totals = self._compute_total_derivs_approx(of=of, wrt=wrt,