added a call to an initialize function to generic_cylinders per a sug…

doc/src/generic_cylinders.rst

@@ -42,3 +42,46 @@ ignored.
    probably needed on the command line. Consistency with the files in the
    pickle directory might not be checked by the program.
 
+A Class in the module
+---------------------
+
+If you want to have a class in the module to provide helper functions,
+your module still needs to have an ``inparse_adder`` function and the module will need
+to have a function called ``get_mpisppy_helper_object(cfg)`` that returns
+the object.  It is called by ``generic_cylinders.py`` after cfg is
+populated and can be used to create a class. Note that the function
+``inparser_adder`` cannot make use of the class because that function
+is called before ``get_mpisppy_helper_object``.
+
+The class definition needs to include all helper functions other than
+``inparser_adder``.  The example ``examples.netdes.netdes_with_class.py``
+demonstrates how to implement a class in the module (although in this
+particular example, there is no advantage to doing that).
+
+
+custom_writer
+-------------
+
+This is an advanced topic. 
+Advanced users might want to write their own solution output function. If the
+module contains a function called ``custom_writer()``, it will be passed
+to the solution writer. Up to four functions can be specified in the module (or the
+class if you are using a class):
+
+   - ef_root_nonants_solution_writer(file_name, representative_scenario, bundling_indicator)
+   - ef_tree_solution_writer(directory_name, scenario_name, scenario, bundling_indicator)
+   - first_stage_solution_writer(file_name, scenario,bundling_indicator)
+   - tree_solution_writer(directory_name, scenario_name, scenario, bundling_indicator)
+
+The first two, if present, will be used for the EF if that is select
+and the second two for hub and spoke solutions.  For further
+information, look at the code in ``mpisppy.generic_cylinders.py`` to
+see how these are used and in ``mpisppy.utils.sputils`` for example functions
+such as ``first_stage_nonant_npy_serializer``.  There is a very simple
+example function in ``examples.netdes.netdes_with_class.py''.
+
+.. Warning::
+   These functions will only be used if cfg.solution_base_name has been given a value by the user.
+
+.. Warning::
+   Misspelled function names will not result in an error message, nor will they be called, of course.

examples/generic_tester.py

@@ -36,7 +36,7 @@
     solver_name = sys.argv[1]
 
 # Use oversubscribe if your computer does not have enough cores.
-# Don't use this unless you have to.
+# Don't use oversubscribe unless you have to.
 # (This may not be allowed on some versions of mpiexec)
 mpiexec_arg = ""  # "--oversubscribe" or "-envall"
 if len(sys.argv) > 2:
@@ -151,6 +151,13 @@ def do_one(dirname, modname, np, argstring, xhat_baseline_dir=None, tol=1e-6):
 #rebaseline_xhat("farmer", "farmer", 1, farmeref, "test_data/farmeref_baseline")
 do_one("farmer", "farmer", 1, farmeref, xhat_baseline_dir = "test_data/farmeref_baseline")
 
+# we need slammax and cross-scenario to make this work well
+netdesC = (f"--max-iterations=60 --instance-name=network-10-20-L-01 --netdes-data-path ./data "
+           f"--solver-name={solver_name} --rel-gap=0.0 --default-rho=10000 --presolve "
+           f"--subgradient-hub --xhatshuffle --max-solver-threads=2 "
+           f"--solution-base-name delete_me")
+do_one("netdes", "netdes_with_class", 2, netdesC, xhat_baseline_dir=None)
+# TBD: put in a baseline
 
 hydroef = (f"--EF --branching-factors '3 3' --EF-solver-name={solver_name}")
 #rebaseline_xhat("hydro", "hydro", 1, hydroef, "test_data/hydroef_baseline")

examples/netdes/netdes_with_class.py

@@ -0,0 +1,180 @@
+###############################################################################
+# mpi-sppy: MPI-based Stochastic Programming in PYthon
+#
+# Copyright (c) 2024, Lawrence Livermore National Security, LLC, Alliance for
+# Sustainable Energy, LLC, The Regents of the University of California, et al.
+# All rights reserved. Please see the files COPYRIGHT.md and LICENSE.md for
+# full copyright and license information.
+###############################################################################
+# modified April 2025 by DLW to illustrate the use of a class for helper functions
+''' Implementation of a simple network design problem. See README for more info
+
+    Created: 9 November 2019 by DTM
+
+    Scenario indices are ZERO based
+'''
+
+import os
+import mpisppy.utils.sputils as sputils
+import pyomo.environ as pyo
+import numpy as np
+from parse import parse
+
+#=========
+def inparser_adder(cfg):
+    # add options unique to sizes
+    # we don't want num_scens from the command line
+    cfg.mip_options()
+    cfg.add_to_config("instance_name",
+                        description="netdes instance name (e.g., network-10-20-L-01)",
+                        domain=str,
+                        default=None)                
+    cfg.add_to_config("netdes_data_path",
+                        description="path to detdes data (e.g., ./data)",
+                        domain=str,
+                        default=None)                
+
+
+def get_mpisppy_helper_object(cfg):
+    return NetDes(cfg)
+
+
+class NetDes:
+    def __init__(self, cfg):
+        self.cfg = cfg
+
+
+    def scenario_creator(self, scenario_name, path=None):
+        if path is None:
+            raise RuntimeError('Must provide the name of the .dat file '
+                               'containing the instance data via the '
+                               'path argument to scenario_creator')
+
+        scenario_ix = self._get_scenario_ix(scenario_name)
+        model = self.build_scenario_model(path, scenario_ix)
+
+        # now attach the one and only scenario tree node
+        sputils.attach_root_node(model, model.FirstStageCost, [model.x[:,:], ])
+
+        return model
+
+
+    def build_scenario_model(self,fname, scenario_ix):
+        data = parse(fname, scenario_ix=scenario_ix)
+        num_nodes = data['N']
+        adj = data['A']    # Adjacency matrix
+        edges = data['el'] # Edge list
+        c = data['c']      # First-stage  cost matrix (per edge)
+        d = data['d']      # Second-stage cost matrix (per edge)
+        u = data['u']      # Capacity of each arc
+        b = data['b']      # Demand of each node
+        p = data['p']      # Probability of scenario
+
+        model = pyo.ConcreteModel()
+        model.x = pyo.Var(edges, domain=pyo.Binary)           # First stage vars
+        model.y = pyo.Var(edges, domain=pyo.NonNegativeReals) # Second stage vars
+
+        model.edges = edges
+        model._mpisppy_probability = p
+
+        ''' Objective '''
+        model.FirstStageCost  = pyo.quicksum(c[e] * model.x[e] for e in edges)
+        model.SecondStageCost = pyo.quicksum(d[e] * model.y[e] for e in edges)
+        obj_expr = model.FirstStageCost + model.SecondStageCost
+        model.MinCost = pyo.Objective(expr=obj_expr, sense=pyo.minimize)
+
+        ''' Variable upper bound constraints on each edge '''
+        model.vubs = pyo.ConstraintList()
+        for e in edges:
+            expr = model.y[e] - u[e] * model.x[e] 
+            model.vubs.add(expr <= 0)
+
+        ''' Flow balance constraints for each node '''
+        model.bals = pyo.ConstraintList()
+        for i in range(num_nodes):
+            in_nbs  = np.where(adj[:,i] > 0)[0]
+            out_nbs = np.where(adj[i,:] > 0)[0]
+            lhs = pyo.quicksum(model.y[i,j] for j in out_nbs) - \
+                  pyo.quicksum(model.y[j,i] for j in in_nbs)
+            model.bals.add(lhs == b[i])
+
+        return model
+
+
+    def scenario_denouement(self, rank, scenario_name, scenario):
+        pass
+
+    def _get_scenario_ix(self, sname):
+        ''' Get the scenario index from the given scenario name by strpiping all
+            digits off of the right of the scenario name, until a non-digit is
+            encountered.
+        '''
+        i = len(sname) - 1
+        while (i > 0 and sname[i-1].isdigit()):
+            i -= 1
+        return int(sname[i:])
+
+
+    def first_stage_solution_writer(self, file_name, scenario, bundling):
+        # A custom first stage writer
+        # adapted from sputils.py; look there for sample tree printer code
+        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)
+        print(f"TEST: first stage writer: {file_name=}, {root_nonants[0]=}")
+        ## np.save(file_name, root_nonants)
+
+
+    ########## helper functions ########
+
+    #=========
+    def scenario_names_creator(self, num_scens,start=None):
+        # if start!=None, the list starts with the 'start' labeled scenario
+        if (start is None) :
+            start=0
+        return [f"Scenario{i}" for i in range(start,start+num_scens)]
+
+
+    #=========
+    def kw_creator(self, cfg):
+        # linked to the scenario_creator and inparser_adder
+        # side-effect is dealing with num_scens
+        inst = cfg.instance_name
+        ns = int(inst.split("-")[-3])
+        if hasattr(cfg, "num_scens"):
+            if cfg.num_scens != ns:
+                raise RuntimeError(f"Argument num-scens={cfg.num_scens} does not match the number "
+                                   "implied by instance name={ns} "
+                                   "\n(--num-scens is not needed for netdes)")
+        else:
+            cfg.add_and_assign("num_scens","number of scenarios", int, None, ns)
+        path = os.path.join(cfg.netdes_data_path, f"{inst}.dat")
+        kwargs = {"path": path}
+        return kwargs
+
+    def sample_tree_scen_creator(self, sname, stage, sample_branching_factors, seed,
+                                 given_scenario=None, **scenario_creator_kwargs):
+        """ Create a scenario within a sample tree. Mainly for multi-stage and simple for two-stage.
+            (this function supports zhat and confidence interval code)
+        Args:
+            sname (string): scenario name to be created
+            stage (int >=1 ): for stages > 1, fix data based on sname in earlier stages
+            sample_branching_factors (list of ints): branching factors for the sample tree
+            seed (int): To allow random sampling (for some problems, it might be scenario offset)
+            given_scenario (Pyomo concrete model): if not None, use this to get data for ealier stages
+            scenario_creator_kwargs (dict): keyword args for the standard scenario creator funcion
+        Returns:
+            scenario (Pyomo concrete model): A scenario for sname with data in stages < stage determined
+                                             by the arguments
+        """
+        # Since this is a two-stage problem, we don't have to do much.
+        sca = scenario_creator_kwargs.copy()
+        sca["seedoffset"] = seed
+        sca["num_scens"] = sample_branching_factors[0]  # two-stage problem
+        return self.scenario_creator(sname, **sca)
+
+    ######## end helper functions #########
+
+
+if __name__=='__main__':
+    print('netdes.py has no main')

mpisppy/generic_cylinders.py

@@ -387,10 +387,18 @@ def _do_decomp(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_
     wheel.spin()
 
     if cfg.solution_base_name is not None:
+        root_writer = getattr(module, "first_stage_solution_writer",
+                                     sputils.first_stage_nonant_npy_serializer)
+        tree_writer = getattr(module, "tree_solution_writer", None)
+
         wheel.write_first_stage_solution(f'{cfg.solution_base_name}.csv')
         wheel.write_first_stage_solution(f'{cfg.solution_base_name}.npy',
-                first_stage_solution_writer=sputils.first_stage_nonant_npy_serializer)
-        wheel.write_tree_solution(f'{cfg.solution_base_name}_soldir')    
+                first_stage_solution_writer=root_writer)
+        if tree_writer is not None:
+            wheel.write_tree_solution(f'{cfg.solution_base_name}_soldir',
+                                      scenario_tree_solution_writer=tree_writer)
+        else:
+            wheel.write_tree_solution(f'{cfg.solution_base_name}_soldir')
         global_toc("Wrote solution data.")
 
 
@@ -500,11 +508,25 @@ def _do_EF(module, cfg, scenario_creator, scenario_creator_kwargs, scenario_deno
         print("Warning: non-optimal solver termination")
 
     global_toc(f"EF objective: {pyo.value(ef.EF_Obj)}")
+
     if cfg.solution_base_name is not None:
+        root_writer = getattr(module, "ef_root_nonants_solution_writer", None)
+        tree_writer = getattr(module, "ef_tree_solution_writer", None)
+
         sputils.ef_nonants_csv(ef, f'{cfg.solution_base_name}.csv')
         sputils.ef_ROOT_nonants_npy_serializer(ef, f'{cfg.solution_base_name}.npy')
-        sputils.write_ef_tree_solution(ef,f'{cfg.solution_base_name}_soldir')
+        if root_writer is not None:
+            sputils.write_ef_first_stage_solution(ef, f'{cfg.solution_base_name}.csv',   # might overwite
+                                                  first_stage_solution_writer=root_writer)
+        else:
+            sputils.write_ef_first_stage_solution(ef, f'{cfg.solution_base_name}.csv')            
+        if tree_writer is not None:
+            sputils.write_ef_tree_solution(ef,f'{cfg.solution_base_name}_soldir',
+                                          scenario_tree_solution_writer=tree_writer)
+        else:
+            sputils.write_ef_tree_solution(ef,f'{cfg.solution_base_name}_soldir')
         global_toc("Wrote EF solution data.")
+
 
 def _model_fname():
     def _bad_news():
@@ -556,9 +578,13 @@ def _proper_bundles(cfg):
         fname = os.path.basename(model_fname)
         sys.path.append(dpath)
         module = importlib.import_module(fname)
-    
+
     cfg = _parse_args(module)
 
+    # Perhaps use an object as the so-called module.
+    if hasattr(module, "get_mpisppy_helper_object"):
+        module = module.get_mpisppy_helper_object(cfg)
+
     bundle_wrapper = None  # the default
     if _proper_bundles(cfg):
         # TBD: remove the need for dill if you are not reading or writing