Repository: hoburg/gpkit
Pull request #830 (2016-09-02 19:19:03)

get vector variables out of a model properly
related to #822


Patch:
diff --git a/gpkit/constraints/set.py b/gpkit/constraints/set.py
index 07ed370..9c3a819 100644
--- a/gpkit/constraints/set.py
+++ b/gpkit/constraints/set.py
@@ -1,16 +1,12 @@
 "Implements ConstraintSet"
+import numpy as np
+
 from ..small_classes import HashVector
 from ..keydict import KeySet, KeyDict
 from ..small_scripts import try_str_without
 from ..repr_conventions import _str, _repr, _repr_latex_
 
 
-def _sort_by_num_models(var):
-    "return integer for Variable sorting"
-    mods = var.key.models
-    return len(mods) if mods else 0
-
-
 def _sort_by_name_and_idx(var):
     "return tuplef for Variable sorting"
     return (var.key.str_without(["units", "idx"]), var.key.idx)
@@ -60,14 +56,27 @@ class ConstraintSet(list):
             return list.__getitem__(self, key)
         else:
             from ..nomials import Variable
-            variables = [Variable(**key.descr) for key in self.varkeys[key]]
+            variables = [Variable(**k.descr) for k in self.varkeys[key]]
             if len(variables) == 1:
                 return variables[0]
-            else:
-                variables.sort(key=_sort_by_num_models)
-                variable = variables[0]
-                # note: doesn't work for vector variables
-                return variable
+            elif variables[0].key.idx and variables[0].key.shape:
+                # maybe it's all one vector variable!
+                vector = np.full(variables[0].key.shape, np.nan, dtype="object")
+                goaldict = dict(variables[0].key.descr)
+                del goaldict["idx"]
+                for variable in variables:
+                    testdict = dict(variable.key.descr)
+                    del testdict["idx"]
+                    if testdict == goaldict:
+                        vector[variable.key.idx] = variable
+                    else:
+                        vector = None
+                        break
+                if vector is not None:
+                    return vector
+            raise ValueError("multiple variables are called '%s'; use"
+                             " variable_byname('%s') to see all of them"
+                             % (key, key))
 
     def variables_byname(self, key):
         "Get all variables with a given name"
diff --git a/gpkit/interactive/widgets.py b/gpkit/interactive/widgets.py
index fbec2d3..7afcec7 100644
--- a/gpkit/interactive/widgets.py
+++ b/gpkit/interactive/widgets.py
@@ -46,8 +46,7 @@ def modelinteract(model, ranges=None, fn_of_sol=None, **solvekwargs):
                 floatslider = widgets.FloatSlider(min=vmin, max=vmax,
                                                   step=vstep, value=v,
                                                   description=varkey_latex)
-                # TODO: find way to set width across ipython versions
-                # floatslider.width = "20ex"
+                floatslider.width = "20ex"
                 floatslider.varkey = k
                 ranges[str(k)] = floatslider
 
diff --git a/gpkit/nomials/nomial_math.py b/gpkit/nomials/nomial_math.py
index 31cd5cc..f6df7ad 100644
--- a/gpkit/nomials/nomial_math.py
+++ b/gpkit/nomials/nomial_math.py
@@ -691,29 +691,15 @@ class SignomialInequality(ScalarSingleEquationConstraint):
             raise TypeError("SignomialInequality could not simplify to"
                             " a PosynomialInequality")
 
-    def _fill_default_x0(self, x0, varkeys):
-        """For all keys in varkeys, updates x0 with default values and
-        substitutions. Returns x0.
-
-        Order of precedence for x0 default:
-            - substitution value
-            - x0 value
-            - sp_init value
-            - 1.0
-        """
-        if x0 is None:
-            x0 = {}
-        x0.update({
-            vk: self.substitutions.get(vk, x0.get(vk, vk.descr.get("sp_init",
-                                                                   1.0)))
-            for vk in varkeys})
-        return x0
-
     def as_gpconstr(self, x0):
         "Returns GP approximation of an SP constraint at x0"
         siglt0, = self.unsubbed
         posy, negy = siglt0.posy_negy()
-        x0 = self._fill_default_x0(x0, negy.varlocs)
+        if x0 is None:
+            x0 = {vk: vk.descr["sp_init"] for vk in negy.varlocs
+                  if "sp_init" in vk.descr}
+        x0.update({var: 1 for var in negy.varlocs if var not in x0})
+        x0.update(self.substitutions)
         pc = PosynomialInequality(posy, "<=", negy.mono_lower_bound(x0))
         pc.substitutions = self.substitutions
         return pc
@@ -735,7 +721,11 @@ class SignomialEquality(SignomialInequality):
     def as_gpconstr(self, x0):
         "Returns GP approximation of an SP constraint at x0"
         siglt0, = self.unsubbed
-        x0 = self._fill_default_x0(x0, siglt0.varlocs)
+        if x0 is None:
+            x0 = {vk: vk.descr["sp_init"] for vk in siglt0.varlocs
+                  if "sp_init" in vk.descr}
+        x0.update({var: 1 for var in siglt0.varlocs if var not in x0})
+        x0.update(self.substitutions)
         posy, negy = siglt0.posy_negy()
         mec = (posy.mono_lower_bound(x0) == negy.mono_lower_bound(x0))
         mec.substitutions = self.substitutions
diff --git a/gpkit/nomials/variables.py b/gpkit/nomials/variables.py
index 46c443e..bbb992e 100644
--- a/gpkit/nomials/variables.py
+++ b/gpkit/nomials/variables.py
@@ -102,10 +102,7 @@ class ArrayVariable(NomialArray):
         if "idx" in descr:
             raise KeyError("the description field 'idx' is reserved")
 
-        if isinstance(shape, Numbers):
-            shape = (shape,)
-        else:
-            shape = tuple(shape)
+        shape = (shape,) if isinstance(shape, Numbers) else tuple(shape)
 
         descr["shape"] = shape
 
@@ -120,7 +117,7 @@ class ArrayVariable(NomialArray):
             elif isinstance(arg, Strings) and "label" not in descr:
                 descr["label"] = arg
 
-        values = []
+        values, value_option = [], None
         for value_option in ["value", "sp_init"]:
             if value_option in descr:
                 values = descr.pop(value_option)
@@ -142,6 +139,7 @@ class ArrayVariable(NomialArray):
         if "name" not in descr:
             descr["name"] = "\\fbox{%s}" % VarKey.new_unnamed_id()
 
+        arraykey = VarKey(**descr)
         vl = np.empty(shape, dtype="object")
         it = np.nditer(vl, flags=['multi_index', 'refs_ok'])
         while not it.finished:
@@ -153,6 +151,7 @@ class ArrayVariable(NomialArray):
             elif valuetype == "list":
                 descr.update({value_option: values[i[0]]})
             vl[i] = Variable(**descr)
+            vl[i].key.arraykey = arraykey
 
         obj = np.asarray(vl).view(cls)
         obj.descr = descr
diff --git a/gpkit/small_classes.py b/gpkit/small_classes.py
index d853d81..7b5e7cc 100644
--- a/gpkit/small_classes.py
+++ b/gpkit/small_classes.py
@@ -140,9 +140,10 @@ def _enray_and_unit_dict(d_in, d_out, unitless_keys=(), united=False):
                                             unitless_keys, united)
         else:
             if hasattr(v[0], "units"):
-                # if the first element of the list already has units,
+                # if the first element of the list alrady has units,
                 # ensure unit consistency across the entire list
-                v = [e.to(k.units).magnitude for e in v]*k.units
+                qty = Quantity(1, v[0].units)
+                v = [e.to(v[0].units).magnitude for e in v]*qty
             else:
                 v = np.array(v)
                 if (united and hasattr(k, "units")


Test patch:
diff --git a/gpkit/tests/t_constraints.py b/gpkit/tests/t_constraints.py
index 67e523e..051942a 100644
--- a/gpkit/tests/t_constraints.py
+++ b/gpkit/tests/t_constraints.py
@@ -12,6 +12,22 @@ import gpkit
 class TestConstraint(unittest.TestCase):
     """Tests for Constraint class"""
 
+    def test_constraintget(self):
+        x = Variable("x")
+        x_ = Variable("x", model="_")
+        xv = VectorVariable(2, "x")
+        xv_ = VectorVariable(2, "x", model="_")
+        self.assertEqual(Model(x, [x >= 1])["x"], x)
+        with self.assertRaises(ValueError):
+            _ = Model(x, [x >= 1, x_ >= 1])["x"]
+        with self.assertRaises(ValueError):
+            _ = Model(x, [x >= 1, xv >= 1])["x"]
+        self.assertTrue(all(Model(xv.prod(), [xv >= 1])["x"] == xv))
+        with self.assertRaises(ValueError):
+            _ = Model(xv.prod(), [xv >= 1, xv_ >= 1])["x"]
+        with self.assertRaises(ValueError):
+            _ = Model(xv.prod(), [xv >= 1, x_ >= 1])["x"]
+
     def test_link_conflict(self):
         "Check that substitution conflicts are flagged during linking."
         x_fx1 = Variable("x", 1, models=["fixed1"])
diff --git a/gpkit/tests/t_solution_array.py b/gpkit/tests/t_solution_array.py
index 635a809..7994110 100644
--- a/gpkit/tests/t_solution_array.py
+++ b/gpkit/tests/t_solution_array.py
@@ -79,21 +79,6 @@ class TestSolutionArray(unittest.TestCase):
         tab = sol.table()
         self.assertTrue(isinstance(tab, str))
 
-    def test_units_sub(self):
-        # issue 809
-        if not gpkit.units:
-            return
-        T = Variable("T", "N", "thrust")
-        Tmin = Variable("T_{min}", "N", "minimum thrust")
-        m = Model(T, [T >= Tmin])
-        tminsub = 1000 * gpkit.units.lbf
-        m.substitutions.update({Tmin: tminsub})
-        sol = m.solve(verbosity=0)
-        self.assertEqual(sol(Tmin), tminsub)
-        self.assertFalse(
-            "1000N" in
-            sol.table().replace(" ", "").replace("[", "").replace("]", ""))
-
 
 class TestResultsTable(unittest.TestCase):
     """TestCase for results_table()"""