diff --git a/src/condor/models.py b/src/condor/models.py index eff72506..806669cd 100644 --- a/src/condor/models.py +++ b/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 = {} diff --git a/tests/test_model_api.py b/tests/test_model_api.py index 1c7d3262..53909ecc 100644 --- a/tests/test_model_api.py +++ b/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, + )