create external api solver with support for matrix IO (and possibly higher order)

src/condor/backends/casadi/__init__.py

@@ -162,9 +162,9 @@ def __post_init__(self, *args, **kwargs):
         # might potentially want to overwrite for casadi-specific validation, etc.
         super().__post_init__(self, *args, **kwargs)
 
-    def flatten_value(self, value):
+    def flatten_value(self, value, force_asymetric=False):
         """flatten a value to the appropriate representation for the backend"""
-        if self.symmetric:
+        if self.symmetric and not force_asymetric:
             unique_values = symmetric_to_unique(
                 value, symbolic=isinstance(value, symbol_class)
             )
@@ -326,9 +326,9 @@ def get_symbol_data(symbol, symmetric=None):
             symmetric = symbol_is(symbol, symbol.T) and size > 1
         else:
             symmetric = (
-                np.isclose(symbol, symbol.T).all()
-                and len(shape) == 2
+                len(shape) == 2
                 and shape[0] == shape[1]
+                and np.isclose(symbol, symbol.T).all()
             )
 
     return BackendSymbolData(

src/condor/implementations/simple.py

@@ -101,7 +101,10 @@ def flatten(self):
                 )
                 for match_elem in elem.match
             ]
-            to_mat[tuple(mat_select)] = elem.backend_repr
+            to_shape = to_mat[tuple(mat_select)].shape
+            to_mat[tuple(mat_select)] = np.atleast_1d(
+                elem.flatten_value(elem.backend_repr, force_asymetric=True)
+            ).reshape(to_shape)
         return to_mat
 
     def key_to_matched_elements(self, k):
@@ -141,26 +144,21 @@ def key_to_matched_elements(self, k):
             )
             raise ValueError(msg)
 
-        # TODO document order/shape?
-        for match_i in match:
-            for dim_len in match_i.shape:
-                if dim_len not in (match_i.size, 1):
-                    msg = "each component of match must be a vector"
-                    raise ValueError(msg)
-
         return match, swap_axes
 
     def __setitem__(self, keys, value):
         match, swap_axes = self.key_to_matched_elements(keys)
         name = self.get_name_for_match(match)
 
-        expected_shape = tuple([int(match_i.size) for match_i in match])
+        expected_shape = tuple()
+        for match_i in match:
+            expected_shape += tuple(match_i.shape)
         if len(expected_shape) == 1:
             value = np.atleast_1d(value).squeeze()
         elif len(expected_shape) == 2:
             value = np.atleast_2d(value)
         symbol_data = get_symbol_data(value)
-        if expected_shape != symbol_data.shape:
+        if np.prod(expected_shape) != np.prod(symbol_data.shape):
             breakpoint()
             msg = "mismatching shape"
             raise ValueError(msg)
@@ -255,6 +253,7 @@ def func(inputs):
                 for elem, v in zip(self.model.output, user_out):
                     return_structure[elem.name] = v
                 out = return_structure.flatten()
+
             else:
                 out = user_out
 

tests/test_backends.py

@@ -6,6 +6,9 @@
 ops = backend.operators
 
 
+rng = np.random.default_rng(12345)
+
+
 def test_reshaping():
     class Reshape(co.ExplicitSystem):
         """
@@ -26,8 +29,8 @@ class ReshapeEmbed(co.ExplicitSystem):
         output.reshape_mat = reshape.mat
         output.reshape_mat_y = reshape.mat @ y
 
-    x = np.arange(3) * 1.0
-    y = np.arange(4) * 1.0
+    x = rng.random(3)
+    y = rng.random(4)
 
     reshape_mat = Reshape(x, y).mat
     assert np.all(reshape_mat == (x[:, None] @ y[None, :]))

tests/test_external.py

@@ -2,10 +2,11 @@
 import pytest
 
 import condor
+from condor import backend
 from condor.backend import operators as ops
 
 
-class Numeric(condor.ExternalSolverWrapper):
+class NumericMisc(condor.ExternalSolverWrapper):
     def __init__(self, output_mode):
         self.output_mode = output_mode
 
@@ -57,19 +58,122 @@ def jacobian(self, input):
             )
 
 
-class Condoric(condor.ExplicitSystem):
+class CondoricMisc(condor.ExplicitSystem):
     a = input()
     b = input(shape=3)
     output.x = a**2 + 2 * b**2
     output.y = ops.sin(a)
 
 
+def simple_rot(th, axis):
+    non_axis = [i for i in range(3) if i != axis]
+    if isinstance(th, backend.symbol_class):
+        dcm = ops.zeros((3, 3))
+    else:
+        dcm = np.zeros((3, 3))
+    dcm[axis, axis] = 1
+    dcm[non_axis[0], non_axis[0]] = np.cos(th)
+    dcm[non_axis[0], non_axis[1]] = -np.sin(th)
+    dcm[non_axis[1], non_axis[1]] = np.cos(th)
+    dcm[non_axis[1], non_axis[0]] = np.sin(th)
+    return dcm
+
+
+def rot_der(th, axis):
+    non_axis = [i for i in range(3) if i != axis]
+    if isinstance(th, backend.symbol_class):
+        dcm = ops.zeros((3, 3))
+    else:
+        dcm = np.zeros((3, 3))
+    dcm[non_axis[0], non_axis[0]] = -np.sin(th)
+    dcm[non_axis[0], non_axis[1]] = -np.cos(th)
+    dcm[non_axis[1], non_axis[1]] = -np.sin(th)
+    dcm[non_axis[1], non_axis[0]] = np.cos(th)
+    return dcm
+
+
+class NumericRotation(condor.ExternalSolverWrapper):
+    def __init__(self, output_mode):
+        self.output_mode = output_mode
+
+        self.input(name="x")
+        self.input(name="y")
+        self.output(name="DCM", shape=(3, 3))
+
+    def function(self, inputs):
+        dcm = simple_rot(inputs.y, 1) @ simple_rot(inputs.x, 0)
+
+        if self.output_mode == 0:
+            return dcm.flatten()
+        if self.output_mode == 1:
+            return (dcm,)
+        if self.output_mode == 2:
+            return dict(DCM=dcm)
+
+    def jacobian(self, inputs):
+        dcm__y = rot_der(inputs.y, 1) @ simple_rot(inputs.x, 0)
+        dcm__x = simple_rot(inputs.y, 1) @ rot_der(inputs.x, 0)
+        if self.output_mode == 0:
+            # must be flattened output in first dimension so multiple outptus can be
+            # supported.
+            jac = np.zeros((9, 2))
+            jac[..., 0] = dcm__x.flatten()
+            jac[..., 1] = dcm__y.flatten()
+            return jac
+        elif self.output_mode == 1:
+            return dict(DCM__y=dcm__y, DCM__x=dcm__x)
+        else:
+            return {
+                ("DCM", "y"): dcm__y,
+                ("DCM", "x"): dcm__x,
+            }
+
+
+class CondoricRotation(condor.ExplicitSystem):
+    x = input()
+    y = input()
+    output.DCM = simple_rot(y, 1) @ simple_rot(x, 0)
+
+
+class NumericProd(condor.ExternalSolverWrapper):
+    def __init__(self, output_mode):
+        self.output_mode = output_mode
+        self.input(name="x", shape=3)
+        self.input(name="y", shape=4)
+        self.output(name="prod", shape=(3, 4))
+
+    def function(self, inputs):
+        prod = inputs.x @ inputs.y.T
+        if self.output_mode == 0:
+            return prod.flatten()
+        if self.output_mode == 1:
+            return (prod,)
+        if self.output_mode == 2:
+            return dict(prod=prod)
+
+
+class CondoricProd(condor.ExplicitSystem):
+    x = input(shape=3)
+    y = input(shape=4)
+
+    output.prod = x @ y.T
+
+
 rng = np.random.default_rng(12345)
 
 
 @pytest.mark.parametrize("output_mode", range(3))
-def test_external_output(output_mode):
-    kwargs = dict(a=rng.random(1), b=rng.random(3))
+@pytest.mark.parametrize(
+    "models",
+    [
+        (NumericMisc, CondoricMisc),
+        (NumericRotation, CondoricRotation),
+        (NumericProd, CondoricProd),
+    ],
+)
+def test_external_output(output_mode, models):
+    Numeric, Condoric = models  # noqa: N806
+    kwargs = {input_.name: rng.random(input_.shape) for input_ in Condoric.input}
     nsys = Numeric(output_mode)
     nout = nsys(**kwargs)
     cout = Condoric(**kwargs)
@@ -79,8 +183,16 @@ def test_external_output(output_mode):
 
 
 @pytest.mark.parametrize("output_mode", range(3))
-def test_external_jacobian(output_mode):
-    kwargs = dict(a=rng.random(1), b=rng.random(3))
+@pytest.mark.parametrize(
+    "models",
+    [
+        (NumericMisc, CondoricMisc),
+        (NumericRotation, CondoricRotation),
+    ],
+)
+def test_external_jacobian(output_mode, models):
+    Numeric, Condoric = models  # noqa: N806
+    kwargs = {input_.name: rng.random(input_.shape) for input_ in Condoric.input}
     nsys = Numeric(output_mode)
 
     class Jac(condor.ExplicitSystem):
@@ -111,9 +223,12 @@ class Jac(condor.ExplicitSystem):
     out_jac = Jac(**kwargs)
     for output_ in Condoric.output:
         for input_ in Jac.input:
-            assert np.all(
-                getattr(out_jac, f"nsys_d{output_.name}_d{input_.name}")
-                == getattr(out_jac, f"csys_d{output_.name}_d{input_.name}")
+            assert getattr(
+                out_jac, f"nsys_d{output_.name}_d{input_.name}"
+            ) == pytest.approx(
+                getattr(out_jac, f"csys_d{output_.name}_d{input_.name}"),
+                abs=1e-18,
+                rel=1e-15,
             )