.. _Internals:

Internals
=========

In this section we provide some notes concerning the code intended to help
developers or maintainers get started.

``nonant_vardata_list``
-----------------------

An important internal structure is ``nonant_vardata_list``, which contains references to the Pyomo vardata objects
for the nonanticipative variables. For indexed Vars, there is one entry for each index. This list enables querying or
setting the values of the nonanticipative variables in an abstract way.

One way or the other, the scenario creator function
needs to attached to each scenario model a list of `objects for non-leaf scenario tree nodes (the list will have only one member for two-stage problems)
called ``_mpisppy_node_list``.
Each node object contains the ``nonant_vardata_list`` for the scenario at that scenario tree node. The list indexes are not meaningful, but the
lists are assumed to be in the same order for every scenario at the tree node.

Some uses
^^^^^^^^^

A simple use is found in ``sputils.py``

.. code-block:: python

   def first_stage_nonant_npy_serializer(file_name, scenario, bundling):
        # write just the nonants for ROOT in an npy file (e.g. for Conf Int)
       root = scenario._mpisppy_node_list[0]
       assert root.name == "ROOT"
       root_nonants = np.fromiter((pyo.value(var) for var in root.nonant_vardata_list), float)
       np.save(file_name, root_nonants)

Another simple use can be found function ``fix_nonants_upto_stage`` function in `xhat_eval.py`:


Standard Construction
^^^^^^^^^^^^^^^^^^^^^

If you are going to use the ``vardatalist`` you don't really need to know how it is constructed, but
we will illustrate one common way here.
Many modelers/users have a call to ``attach_root_node`` in their scenario creator function (e.g., see the scenario
creator in the farmer example), which is a utility
for two-stage problems only. Here is the function signature:

.. code-block:: python

    def attach_root_node(model, firstobj, varlist, nonant_ef_suppl_list=None):

inside the function, there is this call

.. code-block:: python

    model._mpisppy_node_list = [
        scenario_tree.ScenarioNode("ROOT", 1.0, 1, firstobj, varlist, model,
                                   nonant_ef_suppl_list = nonant_ef_suppl_list)
    ]

that attaches a node list with one ROOT node to the model object of the scenario. The ``ScenarioNode`` constructor
creates the ``nonant_vardata_list`` with this call

.. code-block:: python

    import mpisppy.utils.sputils as sputils

    self.nonant_vardata_list = sputils.build_vardatalist(scen_model, self.nonant_list)

where ``build_vardatalist`` function converts strings to vardata objects and expands indexed Vars.

nonant_ef_suppl_list
^^^^^^^^^^^^^^^^^^^^

The ``nonant_ef_suppl_list`` supplied (optionally) to the ``ScenarioNode`` constructor
sets up a nonant var data list for Vars whose nonanticipativity is enforced only in extensive
forms (which includes bundles). This can be useful for supplemental variables whose values
are implied by other nonanticipative variables (e.g. indicator variables).


``nonant_indices``
------------------

The ``nonant_indices`` dictionary has the same information as the ``nonant_vardata_list`` but in a slightly more convenient format, so it is used
more often in classes derived from ``SPBase``.  It is attached by a function in ``spbase.py`` (so `self` refers to ``SPBase``)

.. code-block:: python

    def _attach_nonant_indices(self):
        for (sname, scenario) in self.local_scenarios.items():
            _nonant_indices = dict()
            nlens = scenario._mpisppy_data.nlens        
            for node in scenario._mpisppy_node_list:
                ndn = node.name
                for i in range(nlens[ndn]):
                    _nonant_indices[ndn,i] = node.nonant_vardata_list[i]
            scenario._mpisppy_data.nonant_indices = _nonant_indices
        self.nonant_length = len(_nonant_indices)


Note that the dictionary is indexed by a pair that is node name and the index into ``vardata_list`` and these
indexes are used in various places, such as xbar.

applications examples
^^^^^^^^^^^^^^^^^^^^^

A direct example is in ``_fix_nonants_at_value`` in ``xhat_eval.py``. 

Here is a more subtle snippet from ``phbase.py`` that takes advantage of the fact that many other structures use the same indexes. The
only direct use of ``nonant_indices`` in this snippet is the reference to `nonant._value` to get the variable's current value. As an aside, we note that
the use of direct reference to the "protected" `_value` element in a Pyomo var data object is common in ``mpi-sppy``.

.. code-block:: python

        for k,s in self.local_scenarios.items():
            for ndn_i, nonant in s._mpisppy_data.nonant_indices.items():

                xdiff = nonant._value \
                        - s._mpisppy_model.xbars[ndn_i]._value
                s._mpisppy_model.W[ndn_i]._value += pyo.value(s._mpisppy_model.rho[ndn_i]) * xdiff
                if verbose and self.cylinder_rank == 0:



varid mapping
^^^^^^^^^^^^^

There is a mapping from the vardata object's varid back to the (node name, i) pair that is the key
in the ``nonant_indidices`` dictionary. 
When used carefully, this map allows other programs to quickly communicate about nonanticipative Vars.
The mapping is created by this funcion in ``spbase.py``:

.. code-block:: python

    def _attach_varid_to_nonant_index(self):
        """ Create a map from the id of nonant variables to their Pyomo index.
        """
        for (sname, scenario) in self.local_scenarios.items():
            # In order to support rho setting, create a map
            # from the id of vardata object back its _nonant_index.
            scenario._mpisppy_data.varid_to_nonant_index =\
                {id(var): ndn_i for ndn_i, var in scenario._mpisppy_data.nonant_indices.items()}