Merge 9bb3990 into ec2871b

docs/source/examples/evaluated_fixed_variables.py

@@ -1,10 +1,11 @@
 "Example pre-solve evaluated fixed variable"
-from gpkit import Variable, Model
+from gpkit import Variable, Model, units
 
 # code from t_GPSubs.test_calcconst in tests/t_sub.py
 x = Variable("x", "hours")
 t_day = Variable("t_{day}", 12, "hours")
-t_night = Variable("t_{night}", lambda c: 24 - c[t_day], "hours")
+t_night = Variable("t_{night}",
+                   lambda c: 1*units.day - c(t_day), "hours")
 
 # note that t_night has a function as its value
 m = Model(x, [x >= t_day, x >= t_night])

docs/source/examples/issue_1513.py

@@ -0,0 +1,85 @@
+"Tests non-array linked functions & subs in a vectorization environment"
+import numpy as np
+from gpkit import Variable, Model, ConstraintSet, Vectorize
+
+class Vehicle(Model):
+    "Vehicle model"
+    def setup(self):
+        self.a = a = Variable("a")
+        constraints = [a >= 1]
+        return constraints
+
+class System(Model):
+    "System model"
+    def setup(self):
+        with Vectorize(1):
+            self.Fleet2 = Fleet2()
+        constraints = [self.Fleet2]
+        self.cost = sum(self.Fleet2.z)
+        return constraints
+
+class Fleet2(Model):
+    "Fleet model (composed of multiple Vehicles)"
+    def setup(self):
+        x = Variable("x")
+        lambdafun = lambda c: [c[x]-1, np.ones(x.shape)]
+        with Vectorize(2):
+            y = Variable("y", lambdafun)
+            self.Vehicle = Vehicle()
+
+        self.z = z = Variable("z")
+        substitutions = {"x": 4}
+        constraints = [
+            z >= sum(y/x*self.Vehicle.a),
+            self.Vehicle,
+        ]
+        return constraints, substitutions
+
+m = System()
+sol = m.solve()
+print(sol.table())
+
+# now with more fleets per system
+class System2(Model):
+    "System model"
+    def setup(self):
+        with Vectorize(3):
+            self.Fleet2 = Fleet2()
+        constraints = [self.Fleet2]
+        self.cost = sum(self.Fleet2.z)
+        return constraints
+
+m = System2()
+sol = m.solve()
+print(sol.table())
+
+
+# now testing substitutions
+
+class Simple(Model):
+    "Simple model"
+    def setup(self):
+        self.x = x = Variable("x")
+        y = Variable("y", 1)
+        z = Variable("z", 2)
+        constraints = [
+            x >= y + z,
+        ]
+        return constraints
+
+class Cake(Model):
+    "Cake model"
+    def setup(self):
+        with Vectorize(3):
+            s = Simple()
+        c = ConstraintSet([s])
+        self.cost = sum(s.x)
+        return c
+
+m = Cake()
+m.substitutions.update({
+    "y": ("sweep", [1, 2, 3]),
+    "z": lambda v: v("y")**2,
+})
+sol = m.solve()
+print(sol.table())

gpkit/constraints/prog_factories.py

@@ -6,7 +6,9 @@
 from ..solution_array import SolutionArray
 from ..keydict import KeyDict
 from ..small_scripts import maybe_flatten
+from ..small_classes import FixedScalar
 from ..exceptions import Infeasible
+from ..globals import SignomialsEnabled
 
 
 def evaluate_linked(constants, linked):
@@ -20,20 +22,38 @@ def evaluate_linked(constants, linked):
             del key.descr["gradients"]
     for v, f in linked.items():
         try:
-            if v.veckey and v.veckey.original_fn:
+            if v.veckey and v.veckey.vecfn:
                 if v.veckey not in array_calulated:
-                    ofn = v.veckey.original_fn
-                    array_calulated[v.veckey] = np.array(ofn(kdc))
+                    with SignomialsEnabled():  # to allow use of gpkit.units
+                        vecout = v.veckey.vecfn(kdc)
+                    if not hasattr(vecout, "shape"):
+                        vecout = np.array(vecout)
+                    array_calulated[v.veckey] = vecout
+                    if (any(vecout != 0) and v.veckey.units
+                            and not hasattr(vecout, "units")):
+                        print("Warning: linked function for %s did not return"
+                              " a united value. Modifying it to do so (e.g. by"
+                              " using `()` instead of `[]` to access variables)"
+                              " would reduce the risk of errors." % v.veckey)
                 out = array_calulated[v.veckey][v.idx]
             else:
-                out = f(kdc)
+                with SignomialsEnabled():  # to allow use of gpkit.units
+                    out = f(kdc)
+            if isinstance(out, FixedScalar):  # to allow use of gpkit.units
+                out = out.value
+            if hasattr(out, "units"):
+                out = out.to(v.units or "dimensionless").magnitude
+            elif out != 0 and v.units and not v.veckey:
+                print("Warning: linked function for %s did not return"
+                      " a united value. Modifying it to do so (e.g. by"
+                      " using `()` instead of `[]` to access variables)"
+                      " would reduce the risk of errors." % v)
             if not hasattr(out, "x"):
                 constants[v] = out
                 continue  # a new fixed variable, not a calculated one
             constants[v] = out.x
             v.descr["gradients"] = {adn.tag: grad
-                                    for adn, grad in out.d().items()
-                                    if adn.tag}  # else it's user-created
+                                    for adn, grad in out.d().items()}
         except Exception as exception:  # pylint: disable=broad-except
             from .. import settings
             if settings.get("ad_errors_raise", None):

gpkit/keydict.py

@@ -3,7 +3,7 @@
 from collections.abc import Hashable
 import numpy as np
 from .small_classes import Numbers, Quantity, FixedScalar
-from .small_scripts import is_sweepvar, isnan
+from .small_scripts import is_sweepvar, isnan, veclinkedfn
 
 DIMLESS_QUANTITY = Quantity(1, "dimensionless")
 INT_DTYPE = np.dtype(int)
@@ -226,6 +226,14 @@ def __setitem__(self, key, value):
             super().__getitem__(key)[idx] = value
             return  # successfully set a single index!
         if key.shape: # now if we're setting an array...
+            if hasattr(value, "__call__"):  # a linked vector-function
+                key.vecfn = value
+                value = np.empty(key.shape, dtype="object")
+                it = np.nditer(value, flags=['multi_index', 'refs_ok'])
+                while not it.finished:
+                    i = it.multi_index
+                    it.iternext()
+                    value[i] = veclinkedfn(key.vecfn, i)
             if getattr(value, "shape", None):   # is the value an array?
                 if value.dtype == INT_DTYPE:
                     value = np.array(value, "f")  # convert to float

gpkit/nomials/math.py

@@ -179,7 +179,7 @@ def __ge__(self, other):
     def __add__(self, other, rev=False):
         other_hmap = getattr(other, "hmap", None)
         if isinstance(other, Numbers):
-            if not other:  # other is zero
+            if other == 0:
                 return Signomial(self.hmap)
             other_hmap = NomialMap({EMPTY_HV: mag(other)})
             other_hmap.units_of_product(other)

gpkit/nomials/variables.py

@@ -8,7 +8,7 @@
 from ..globals import NamedVariables, Vectorize
 from ..varkey import VarKey
 from ..small_classes import Strings, Numbers, HashVector
-from ..small_scripts import is_sweepvar
+from ..small_scripts import is_sweepvar, veclinkedfn
 
 
 def addmodelstodescr(descr, addtonamedvars=None):
@@ -133,29 +133,26 @@ def __new__(cls, shape, *args, **descr):  # pylint: disable=too-many-branches, t
         if "name" not in descr:
             descr["name"] = "\\fbox{%s}" % VarKey.unique_id()
 
-        value_option = None
-        if "value" in descr:
-            value_option = "value"
-        if value_option:
-            values = descr.pop(value_option)
-        if value_option and not hasattr(values, "__call__"):
-            if Vectorize.vectorization:
-                if not hasattr(values, "shape"):
-                    values = np.full(shape, values, "f")
-                else:
-                    values = np.broadcast_to(values, reversed(shape)).T
-            elif not hasattr(values, "shape"):
-                values = np.array(values)
-            if values.shape != shape:
-                raise ValueError("the value's shape %s is different than"
-                                 " the vector's %s." % (values.shape, shape))
+        values = descr.pop("value", None)
+        if values is not None:
+            if not hasattr(values, "__call__"):
+                if Vectorize.vectorization:
+                    if not hasattr(values, "shape"):
+                        values = np.full(shape, values, "f")
+                    else:
+                        values = np.broadcast_to(values, reversed(shape)).T
+                elif not hasattr(values, "shape"):
+                    values = np.array(values)
+                if values.shape != shape:
+                    raise ValueError("value's shape %s is different from the"
+                                     " vector's %s." % (values.shape, shape))
 
         veckeydescr = descr.copy()
         addmodelstodescr(veckeydescr)
-        if value_option:
+        if values is not None:
             if hasattr(values, "__call__"):
-                veckeydescr["original_fn"] = values
-            veckeydescr[value_option] = values
+                veckeydescr["vecfn"] = values
+            veckeydescr["value"] = values
         veckey = VarKey(**veckeydescr)
 
         descr["veckey"] = veckey
@@ -165,11 +162,11 @@ def __new__(cls, shape, *args, **descr):  # pylint: disable=too-many-branches, t
             i = it.multi_index
             it.iternext()
             descr["idx"] = i
-            if value_option:
-                if hasattr(values, "__call__"):
-                    descr[value_option] = veclinkedfn(values, i)
+            if values is not None:
+                if hasattr(values, "__call__"):  # a vector function
+                    descr["value"] = veclinkedfn(values, i)
                 else:
-                    descr[value_option] = values[i]
+                    descr["value"] = values[i]
             vl[i] = Variable(**descr)
 
         obj = np.asarray(vl).view(NomialArray)
@@ -178,14 +175,6 @@ def __new__(cls, shape, *args, **descr):  # pylint: disable=too-many-branches, t
         return obj
 
 
-def veclinkedfn(linkedfn, i):
-    "Generate an indexed linking function."
-    def newlinkedfn(c):
-        "Linked function that pulls out a particular index"
-        return np.array(linkedfn(c))[i]
-    return newlinkedfn
-
-
 class VectorizableVariable(Variable, ArrayVariable):  # pylint: disable=too-many-ancestors
     "A Variable outside a vectorized environment, an ArrayVariable within."
     def __new__(cls, *args, **descr):

gpkit/small_scripts.py

@@ -3,6 +3,14 @@
 import numpy as np
 
 
+def veclinkedfn(linkedfn, i):
+    "Generate an indexed linking function."
+    def newlinkedfn(c):
+        "Linked function that pulls out a particular index"
+        return np.array(linkedfn(c))[i]
+    return newlinkedfn
+
+
 def initsolwarning(result, category="uncategorized"):
     "Creates a results dictionary for a particular category of warning."
     if "warnings" not in result:

gpkit/solution_array.py

@@ -570,7 +570,7 @@ def todataframe(self, showvars=None,
                     ", ".join("%s=%s" % (k, v) for (k, v) in key.descr.items()
                               if k not in ["name", "units", "unitrepr",
                                            "idx", "shape", "veckey",
-                                           "value", "original_fn",
+                                           "value", "vecfn",
                                            "lineage", "label"])])
         return pd.DataFrame(rows, columns=cols)
 

gpkit/tests/t_examples.py

@@ -48,6 +48,9 @@ def test_dummy_example(self, example):
     #     import matplotlib.pyplot as plt
     #     plt.close("all")
 
+    def test_issue_1513(self, example):
+        pass
+
     def test_autosweep(self, example):
         from gpkit import ureg
         bst1, tol1 = example.bst1, example.tol1

gpkit/tests/t_sub.py

@@ -227,7 +227,8 @@ def test_vector_sweep(self):
     def test_calcconst(self):
         x = Variable("x", "hours")
         t_day = Variable("t_{day}", 12, "hours")
-        t_night = Variable("t_{night}", lambda c: 24 - c[t_day], "hours")
+        t_night = Variable("t_{night}",
+                           lambda c: 1*gpkit.ureg.day - c(t_day), "hours")
         _ = pickle.dumps(t_night)
         m = Model(x, [x >= t_day, x >= t_night])
         sol = m.solve(verbosity=0)