add shape check

gpkit/keydict.py

@@ -4,6 +4,7 @@
 import numpy as np
 from .small_classes import Numbers, Quantity, FixedScalar
 from .small_scripts import is_sweepvar, isnan
+from .varkey import GLOBAL_KEYMAP
 
 DIMLESS_QUANTITY = Quantity(1, "dimensionless")
 INT_DTYPE = np.dtype(int)
@@ -46,6 +47,7 @@ class KeyMap:
     keymap = []
     log_gets = False
     varkeys = None
+    vks = None
 
     def __init__(self, *args, **kwargs):
         "Passes through to super().__init__ via the `update()` method"
@@ -62,9 +64,11 @@ def parse_and_index(self, key):
                 return key.veckey, key.idx
             return key, None
         except AttributeError:
-            if not self.varkeys:
+            if not self.vks:
                 return key, self.update_keymap()
         # looks like we're in a substitutions dictionary
+        if not self.varkeys:
+            self.varkeys = KeySet(self.vks)
         if key not in self.varkeys:  # pylint:disable=unsupported-membership-test
             raise KeyError(key)
         newkey, *otherkeys = self.varkeys[key]  # pylint:disable=unsubscriptable-object
@@ -195,6 +199,8 @@ def __setitem__(self, key, value):
                 raise e
             if not hasattr(value, "shape"):
                 value = np.full(key.shape, value)
+            elif key.shape != value.shape:
+                raise KeyError("Key and value have different shapes") from e
             for subkey, subval in zip(key.flat, value.flat):
                 self[subkey] = subval
             return