Allow initialization of indexed sets via a function returning a dict

doc/OnlineDocs/pyomo_modeling_components/Sets.rst

@@ -451,8 +451,37 @@ for this model, a toy data file (in AMPL "``.dat``" format) would be:
 
    >>> inst = model.create_instance('src/scripting/Isinglecomm.dat')
 
-This can also be done somewhat more efficiently, and perhaps more clearly,
-using a :class:`BuildAction` (for more information, see :ref:`BuildAction`):
+This can also be done much more efficiently using initialization functions
+that accept only a model block and return a ``dict`` with all the information
+needed for the indexed set:
+
+.. doctest::
+   :hide:
+
+   >>> model = inst
+   >>> del model.NodesIn
+   >>> del model.NodesOut
+
+.. testcode::
+
+    def NodesIn_init(m):
+        # Create a dict to show NodesIn list for every node
+        d = {i: [] for i in m.Nodes}
+        # loop over the arcs and record the end points
+        for i, j in model.Arcs:
+            d[j].append(i)
+        return d
+    model.NodesIn = pyo.Set(model.Nodes, initialize=NodesIn_init)
+
+    def NodesOut_init(m):
+        d = {i: [] for i in m.Nodes}
+        for i, j in model.Arcs:
+            d[i].append(j)
+        return d
+    model.NodesOut = pyo.Set(model.Nodes, initialize=NodesOut_init)
+
+This can also be done efficiently, and perhaps more clearly, using a
+:class:`BuildAction` (for more information, see :ref:`BuildAction`):
 
 .. doctest::
    :hide:
@@ -463,8 +492,8 @@ using a :class:`BuildAction` (for more information, see :ref:`BuildAction`):
 
 .. testcode::
 
-   model.NodesOut = pyo.Set(model.Nodes, within=model.Nodes)
    model.NodesIn = pyo.Set(model.Nodes, within=model.Nodes)
+   model.NodesOut = pyo.Set(model.Nodes, within=model.Nodes)
 
    def Populate_In_and_Out(model):
        # loop over the arcs and record the end points

examples/doc/samples/case_studies/diet/DietProblem.tex

@@ -62,7 +62,7 @@ \subsection*{Build the model}
 
 At this point we must start defining the rules associated with our parameters and variables.  We begin with the most important rule, the cost rule, which will tell the model to try and minimize the overall cost.  Logically, the total cost is going to be the sum of how much is spent on each food, and that value in turn is going to be determined by the cost of the food and how much of it is purchased.  For example, if three \$5 hamburgers and two \$1 apples are purchased, than the total cost would be $3 \cdot 5 + 2 \cdot 1 = 17$.  Note that this process is the same as taking the dot product of the amounts vector and the costs vector.
 
-To input this, we must define the cost rule, which we creatively call costRule as 
+To input this, we must define the cost rule, which we creatively call costRule as
 
 \begin{verbatim}def costRule(model):
     return sum(model.costs[n]*model.amount[n] for n in model.foods)
@@ -75,10 +75,10 @@ \subsection*{Build the model}
 
 This line defines the objective of the model as the costRule,  which Pyomo interprets as the value it needs to minimize; in this case it will minimize our costs.  Also, as a note, we defined the objective as ``model.cost'' which is not to be confused with the parameter we defined earlier as ``model.costs,'' despite their similar names.  These are two different values and accidentally giving them the same name will cause problems when trying to solve the problem.
 
-We must also create a rule for the volume consumed.  The construction of this rule is similar to the cost rule as once again we take the dot product, this time between the volume and amount vectors. 
+We must also create a rule for the volume consumed.  The construction of this rule is similar to the cost rule as once again we take the dot product, this time between the volume and amount vectors.
 
 \begin{verbatim}def volumeRule(model):
-    return sum(model.volumes[n]*model.amount[n] for n in 
+    return sum(model.volumes[n]*model.amount[n] for n in
         model.foods) <= model.max_volume
 
 model.volume = Constraint(rule=volumeRule)
@@ -90,7 +90,7 @@ \subsection*{Build the model}
 
 \begin{verbatim}
 def nutrientRule(n, model):
-    value = sum(model.nutrient_value[n,f]*model.amount[f] 
+    value = sum(model.nutrient_value[n,f]*model.amount[f]
         for f in model.foods)
     return (model.min_nutrient[n], value, model.max_nutrient[n])
 
@@ -160,7 +160,7 @@ \subsection*{Data entry}
 
 The amount of spaces between each element is irrelevant (as long as there is at least one) so the matrix should be formatted for ease of reading.
 
-Now that we have finished both the model and the data file save them both. It's convention to give the model file a .py extension and the data file a .dat extension.  
+Now that we have finished both the model and the data file save them both. It's convention to give the model file a .py extension and the data file a .dat extension.
 
 \subsection*{Solution}
 

examples/doc/samples/case_studies/diet/README.txt

@@ -14,7 +14,7 @@ to import the Pyomo package for use in the code.  The next step is to create an
 
 {{{
 #!python
-model = AbstractModel() 
+model = AbstractModel()
 }}}
 The rest of our work will be contained within this object.
 
@@ -79,7 +79,7 @@ We restrict our domain to the non-negative reals.  If we accepted negative numbe
 
 At this point we must start defining the rules associated with our parameters and variables.  We begin with the most important rule, the cost rule, which will tell the model to try and minimize the overall cost.  Logically, the total cost is going to be the sum of how much is spent on each food, and that value in turn is going to be determined by the cost of the food and how much of it is purchased.  For example, if three !$5 hamburgers and two !$1 apples are purchased, than the total cost would be 3*5 + 2*1 = 17.  Note that this process is the same as taking the dot product of the amounts vector and the costs vector.
 
-To input this, we must define the cost rule, which we creatively call costRule as 
+To input this, we must define the cost rule, which we creatively call costRule as
 
 {{{
 #!python
@@ -95,7 +95,7 @@ model.cost=Objective(rule=costRule
 
 This line defines the objective of the model as the costRule,  which Pyomo interprets as the value it needs to minimize; in this case it will minimize our costs.  Also, as a note, we defined the objective as "model.cost" which is not to be confused with the parameter we defined earlier as `"model.costs" despite their similar names.  These are two different values and accidentally giving them the same name will cause problems when trying to solve the problem.
 
-We must also create a rule for the volume consumed.  The construction of this rule is similar to the cost rule as once again we take the dot product, this time between the volume and amount vectors. 
+We must also create a rule for the volume consumed.  The construction of this rule is similar to the cost rule as once again we take the dot product, this time between the volume and amount vectors.
 
 {{{
 #!python
@@ -112,7 +112,7 @@ Finally, we need to add the constraint that ensures we obtain proper amounts of
 {{{
 #!python
 def nutrientRule(n, model):
-    value = sum(model.nutrient_value[n,f]*model.amount[f] 
+    value = sum(model.nutrient_value[n,f]*model.amount[f]
         for f in model.foods)
     return (model.min_nutrient[n], value, model.max_nutrient[n])
 
@@ -179,7 +179,7 @@ vc          0     30    0;
 
 The amount of spaces between each element is irrelevant (as long as there is at least one) so the matrix should be formatted for ease of reading.
 
-Now that we have finished both the model and the data file save them both. It's convention to give the model file a .py extension and the data file a .dat extension.  
+Now that we have finished both the model and the data file save them both. It's convention to give the model file a .py extension and the data file a .dat extension.
 
 == Solution ==
 
@@ -193,7 +193,7 @@ Using Pyomo we quickly find the solution to our diet problem.  Simply run Pyomo
 # ----------------------------------------------------------
 #   Problem Information
 # ----------------------------------------------------------
-Problem: 
+Problem:
 - Lower bound: 29.44055944
   Upper bound: inf
   Number of objectives: 1
@@ -205,28 +205,28 @@ Problem:
 # ----------------------------------------------------------
 #   Solver Information
 # ----------------------------------------------------------
-Solver: 
+Solver:
 - Status: ok
   Termination condition: unknown
   Error rc: 0
 
 # ----------------------------------------------------------
 #   Solution Information
 # ----------------------------------------------------------
-Solution: 
+Solution:
 - number of solutions: 1
   number of solutions displayed: 1
 - Gap: 0.0
   Status: optimal
-  Objective: 
-    f: 
+  Objective:
+    f:
       Id: 0
       Value: 29.44055944
-  Variable: 
-    amount[rice]: 
+  Variable:
+    amount[rice]:
       Id: 0
       Value: 9.44056
-    amount[apple]: 
+    amount[apple]:
       Id: 2
       Value: 10
 

pyomo/core/base/set.py

@@ -47,6 +47,7 @@
     ParameterizedIndexedCallInitializer,
     ParameterizedInitializer,
     ParameterizedScalarCallInitializer,
+    ScalarCallInitializer,
 )
 from pyomo.core.base.range import (
     NumericRange,
@@ -2299,14 +2300,31 @@ def construct(self, data=None):
                     # scalar sets (including set operators) to be
                     # initialized (and potentially empty) after construct().
                     self._getitem_when_not_present(None)
-            elif self._init_values.contains_indices():
-                # The index is coming in externally; we need to validate it
-                for index in self._init_values.indices():
-                    IndexedComponent.__getitem__(self, index)
             else:
-                # Bypass the index validation and create the member directly
-                for index in self.index_set():
-                    self._getitem_when_not_present(index)
+                # If this is an IndexedSet but the initializer is a function
+                # that does not accept indices, call the function and initialize
+                # from the object it provides (e.g., a dict with all members of
+                # the indexed set). This is similar to
+                # IndexedComponent._construct_from_rule_using_setitem.
+                if (
+                    self.is_indexed()
+                    and type(self._init_values._init) is ScalarCallInitializer
+                ):
+                    self._init_values = TuplizeValuesInitializer(
+                        Initializer(
+                            self._init_values._init(self.parent_block(), None),
+                            treat_sequences_as_mappings=False,
+                        )
+                    )
+
+                if self._init_values.contains_indices():
+                    # The index is coming in externally; we need to validate it
+                    for index in self._init_values.indices():
+                        IndexedComponent.__getitem__(self, index)
+                else:
+                    # Bypass the index validation and create the member directly
+                    for index in self.index_set():
+                        self._getitem_when_not_present(index)
         finally:
             # Restore the original initializer (if overridden by data argument)
             if data is not None:

pyomo/core/tests/unit/test_sets.py

@@ -3339,6 +3339,18 @@ def test_setargs5(self):
         model.Y = RangeSet(model.C)
         model.X = Param(model.C, default=0.0)
 
+    def test_setargs6(self):
+        # Test that we can create an indexed set from a function that returns
+        # a dict to define the set
+        model = ConcreteModel()
+        model.A = Set(initialize=[1, 2])
+        model.B = Set(model.A, initialize={1: [2, 3], 2: [3, 4]})
+        model.C = Set(model.A, initialize=lambda m: {x: [x + 1, x + 2] for x in m.A})
+        # convert to native data types for easier comparison
+        B = {k: v.ordered_data() for (k, v) in model.B.items()}
+        C = {k: v.ordered_data() for (k, v) in model.C.items()}
+        self.assertEqual(B, C)
+
     @unittest.skip("_verify was removed during the set rewrite")
     def test_verify(self):
         a = Set(initialize=[1, 2, 3])