Make ModelicaSystem.linearize() return all values

OMPython/ModelicaSystem.py

@@ -43,6 +43,8 @@
 import numpy as np
 import importlib
 import pathlib
+from dataclasses import dataclass
+from typing import Optional
 
 from OMPython.OMCSession import OMCSessionBase, OMCSessionZMQ
 
@@ -54,6 +56,57 @@ class ModelicaSystemError(Exception):
     pass
 
 
+@dataclass
+class LinearizationResult:
+    """Modelica model linearization results.
+
+    Attributes:
+        n: number of states
+        m: number of inputs
+        p: number of outputs
+        A: state matrix (n x n)
+        B: input matrix (n x m)
+        C: output matrix (p x n)
+        D: feedthrough matrix (p x m)
+        x0: fixed point
+        u0: input corresponding to the fixed point
+        stateVars: names of state variables
+        inputVars: names of inputs
+        outputVars: names of outputs
+    """
+
+    n: int
+    m: int
+    p: int
+
+    A: list
+    B: list
+    C: list
+    D: list
+
+    x0: list[float]
+    u0: list[float]
+
+    stateVars: list[str]
+    inputVars: list[str]
+    outputVars: list[str]
+
+    def __iter__(self):
+        """Allow unpacking A, B, C, D = result."""
+        yield self.A
+        yield self.B
+        yield self.C
+        yield self.D
+
+    def __getitem__(self, index: int):
+        """Allow accessing A, B, C, D via result[0] through result[3].
+
+        This is needed for backwards compatibility, because
+        ModelicaSystem.linearize() used to return [A, B, C, D].
+        """
+        return {0: self.A, 1: self.B, 2: self.C, 3: self.D}[index]
+
+
 class ModelicaSystem:
     def __init__(self, fileName=None, modelName=None, lmodel=None, commandLineOptions=None,
                  variableFilter=None, customBuildDirectory=None, verbose=True, raiseerrors=False,
@@ -967,13 +1020,22 @@ def optimize(self):  # 21
 
         return optimizeResult
 
-    # to linearize model
-    def linearize(self, lintime=None, simflags=None):  # 22
-        """
-        This method linearizes model according to the linearized options. This will generate a linear model that consists of matrices A, B, C and D.  It can be called:
-        only without any arguments
-        usage
-        >>> linearize()
+    def linearize(self, lintime: Optional[float] = None, simflags: Optional[str] = None) -> LinearizationResult:
+        """Linearize the model according to linearOptions.
+
+        Args:
+            lintime: Override linearOptions["stopTime"] value.
+            simflags: A string of extra command line flags for the model
+              binary.
+
+        Returns:
+            A LinearizationResult object is returned. This allows several
+            uses:
+            * `(A, B, C, D) = linearize()` to get just the matrices,
+            * `result = linearize(); result.A` to get everything and access the
+              attributes one by one,
+            * `result = linearize(); A = result[0]` mostly just for backwards
+              compatibility, because linearize() used to return `[A, B, C, D]`.
         """
 
         if self.xmlFile is None:
@@ -1043,7 +1105,8 @@ def linearize(self, lintime=None, simflags=None):  # 22
             self.linearinputs = inputVars
             self.linearoutputs = outputVars
             self.linearstates = stateVars
-            return [A, B, C, D]
+            return LinearizationResult(n, m, p, A, B, C, D, x0, u0, stateVars,
+                                       inputVars, outputVars)
         except ModuleNotFoundError:
             raise Exception("ModuleNotFoundError: No module named 'linearized_model'")
 

OMPython/__init__.py

@@ -39,7 +39,7 @@
 import logging
 
 from OMPython.OMCSession import OMCSessionBase, OMCSessionZMQ
-from OMPython.ModelicaSystem import ModelicaSystem, ModelicaSystemError
+from OMPython.ModelicaSystem import ModelicaSystem, ModelicaSystemError, LinearizationResult
 
 # Logger Defined
 logger = logging.getLogger('OMPython')
@@ -61,6 +61,7 @@
 __all__ = [
     'ModelicaSystem',
     'ModelicaSystemError',
+    'LinearizationResult',
 
     'OMCSessionZMQ',
     'OMCSessionBase',

tests/test_linearization.py

@@ -71,7 +71,7 @@ def test_getters(self):
         mod.setLinearizationOptions("stopTime=0.02")
         assert mod.getLinearizationOptions("stopTime") == ["0.02"]
 
-        mod.setInputs(["u1=0", "u2=0"])
+        mod.setInputs(["u1=10", "u2=0"])
         [A, B, C, D] = mod.linearize()
         g = float(mod.getParameters("g")[0])
         l = float(mod.getParameters("l")[0])
@@ -82,3 +82,27 @@ def test_getters(self):
         assert np.isclose(B, [[0, 0], [0, 1]]).all()
         assert np.isclose(C, [[0.5, 1], [0, 1]]).all()
         assert np.isclose(D, [[1, 0], [1, 0]]).all()
+
+        # test LinearizationResult
+        result = mod.linearize()
+        assert result[0] == A
+        assert result[1] == B
+        assert result[2] == C
+        assert result[3] == D
+        with self.assertRaises(KeyError):
+            result[4]
+
+        A2, B2, C2, D2 = result
+        assert A2 == A
+        assert B2 == B
+        assert C2 == C
+        assert D2 == D
+
+        assert result.n == 2
+        assert result.m == 2
+        assert result.p == 2
+        assert np.isclose(result.x0, [0, np.pi]).all()
+        assert np.isclose(result.u0, [10, 0]).all()
+        assert result.stateVars == ["omega", "phi"]
+        assert result.inputVars == ["u1", "u2"]
+        assert result.outputVars == ["y1", "y2"]