Added validate_input_sizes function to Model class

src/condor/models.py

@@ -1,6 +1,8 @@
 import dataclasses as dc
 import logging
 
+import numpy as np
+
 from condor import backend, implementations
 
 # TODO: figure out how to make this an option/setting like django?
@@ -1382,6 +1384,8 @@ def __init__(self, *args, name="", **kwargs):
 
         self.bind_input_fields(*args, **kwargs)
 
+        self.validate_input_sizes()
+
         implementation_class = cls.get_implementation_class(cls)
         self.implementation = implementation_class(self)
         # self, list(self.input_kwargs.values())
@@ -1398,6 +1402,49 @@ def __init__(self, *args, name="", **kwargs):
 
         self.bind_embedded_models()
 
+    def validate_input_sizes(self):
+        cls = self.__class__
+        err_flag = 0
+        err_string = ""
+
+        for input_field in cls._meta.input_fields:
+            for name in list(input_field.keys()):
+                desired_shape = getattr(cls, name).shape
+                input_value = getattr(self, name)
+
+                if not isinstance(input_value, backend.symbol_class):
+                    input_value = np.atleast_2d(input_value)
+
+                input_shape = input_value.shape
+
+                # if input is not a vector, then it has to
+                # match desired_shape exactly
+                # if desired_input is a vector, then
+                # input can be row or column vector
+
+                base_shape_correct = input_shape == desired_shape
+                input_is_vector = len(input_shape) == 2 and (
+                    input_shape[0] == 1 or input_shape[1] == 1
+                )
+                vector_input_incorrect_size = (
+                    input_is_vector
+                    and not base_shape_correct
+                    and (input_value.T.shape != desired_shape)
+                )
+                non_vector_input_incorrect_size = (
+                    not input_is_vector and not base_shape_correct
+                )
+
+                if non_vector_input_incorrect_size or vector_input_incorrect_size:
+                    err_flag = 1
+                    err_string += (
+                        f"{input_field._name} {name} was assigned with shape "
+                        f"{input_shape}, but should be shape {desired_shape}\n"
+                    )
+
+        if err_flag:
+            raise ValueError(err_string)
+
     def bind_input_fields(self, *args, **kwargs):
         cls = self.__class__
         self.input_kwargs = {}

tests/test_model_api.py

@@ -146,3 +146,46 @@ class Sys(co.ExplicitSystem):
     out = Sys(x=x)
 
     assert out.y == 2 * x[1, 0]
+
+
+@pytest.mark.parametrize(
+    ("x_input", "y_input"),
+    [
+        (((0, 1), [[1, 2], [3, 4]])),
+        (np.array([0, 1]), np.zeros((2, 2))),
+        (np.ones((1, 2)), co.backend.operators.zeros((2, 2))),
+    ],
+)
+def test_valid_size_check(x_input, y_input):
+    class MySys(co.ExplicitSystem):
+        x = input(shape=(2, 1))
+        y = input(shape=(2, 2))
+        output.z = y @ x
+
+    sim = MySys(
+        x=x_input,
+        y=y_input,
+    )
+
+    assert sim.z.shape == (2, 1)
+
+
+@pytest.mark.parametrize(
+    ("x_input", "y_input"),
+    [
+        ((0.1, 0.1, 0.1), 0),
+        (np.array([0]), (1, 2, 3)),
+        (co.backend.operators.zeros((3, 3)), co.backend.operators.zeros((2, 1))),
+    ],
+)
+def test_invalid_size_check(x_input, y_input):
+    class MySys(co.ExplicitSystem):
+        x = input(shape=(2, 1))
+        y = input(shape=(2, 2))
+        output.z = y @ x
+
+    with pytest.raises(ValueError, match="was assigned with shape"):
+        sim = MySys(
+            x=x_input,
+            y=y_input,
+        )