Merge pull request #450 from roryyorke/rory/ss-repr-latex

murrayrm · web-flow · commit 6ed3f74de8ba · 2020-12-30T21:39:50.000-08:00
Added IPython LaTeX representation method for StateSpace objects
diff --git a/control/statesp.py b/control/statesp.py
@@ -74,6 +74,8 @@
     'statesp.use_numpy_matrix': False,  # False is default in 0.9.0 and above
     'statesp.default_dt': None,
     'statesp.remove_useless_states': True,
+    'statesp.latex_num_format': '.3g',
+    'statesp.latex_repr_type': 'partitioned',
     }
 
 
@@ -128,6 +130,33 @@ def _ssmatrix(data, axis=1):
     return arr.reshape(shape)
 
 
+def _f2s(f):
+    """Format floating point number f for StateSpace._repr_latex_.
+
+    Numbers are converted to strings with statesp.latex_num_format.
+
+    Inserts column separators, etc., as needed.
+    """
+    fmt = "{:" + config.defaults['statesp.latex_num_format'] + "}"
+    sraw = fmt.format(f)
+    # significand-exponent
+    se = sraw.lower().split('e')
+    # whole-fraction
+    wf = se[0].split('.')
+    s = wf[0]
+    if wf[1:]:
+        s += r'.&\hspace{{-1em}}{frac}'.format(frac=wf[1])
+    else:
+        s += r'\phantom{.}&\hspace{-1em}'
+
+    if se[1:]:
+        s += r'&\hspace{{-1em}}\cdot10^{{{:d}}}'.format(int(se[1]))
+    else:
+        s += r'&\hspace{-1em}\phantom{\cdot}'
+
+    return s
+
+
 class StateSpace(LTI):
     """StateSpace(A, B, C, D[, dt])
 
@@ -158,6 +187,24 @@ class StateSpace(LTI):
     time. The default value of 'dt' is None and can be changed by changing the
     value of ``control.config.defaults['statesp.default_dt']``.
 
+    StateSpace instances have support for IPython LaTeX output,
+    intended for pretty-printing in Jupyter notebooks.  The LaTeX
+    output can be configured using
+    `control.config.defaults['statesp.latex_num_format']` and
+    `control.config.defaults['statesp.latex_repr_type']`.  The LaTeX output is
+    tailored for MathJax, as used in Jupyter, and may look odd when
+    typeset by non-MathJax LaTeX systems.
+
+    `control.config.defaults['statesp.latex_num_format']` is a format string
+    fragment, specifically the part of the format string after `'{:'`
+    used to convert floating-point numbers to strings.  By default it
+    is `'.3g'`.
+
+    `control.config.defaults['statesp.latex_repr_type']` must either be
+    `'partitioned'` or `'separate'`.  If `'partitioned'`, the A, B, C, D
+    matrices are shown as a single, partitioned matrix; if
+    `'separate'`, the matrices are shown separately.
+
     """
 
     # Allow ndarray * StateSpace to give StateSpace._rmul_() priority
@@ -306,6 +353,136 @@ def __repr__(self):
             C=asarray(self.C).__repr__(), D=asarray(self.D).__repr__(),
             dt=(isdtime(self, strict=True) and ", {}".format(self.dt)) or '')
 
+    def _latex_partitioned_stateless(self):
+        """`Partitioned` matrix LaTeX representation for stateless systems
+
+        Model is presented as a matrix, D.  No partition lines are shown.
+
+        Returns
+        -------
+        s : string with LaTeX representation of model
+        """
+        lines = [
+            r'\[',
+            r'\left(',
+            (r'\begin{array}'
+             + r'{' + 'rll' * self.inputs + '}')
+            ]
+
+        for Di in asarray(self.D):
+            lines.append('&'.join(_f2s(Dij) for Dij in Di)
+                         + '\\\\')
+
+        lines.extend([
+            r'\end{array}'
+            r'\right)',
+            r'\]'])
+
+        return '\n'.join(lines)
+
+    def _latex_partitioned(self):
+        """Partitioned matrix LaTeX representation of state-space model
+
+        Model is presented as a matrix partitioned into A, B, C, and D
+        parts.
+
+        Returns
+        -------
+        s : string with LaTeX representation of model
+        """
+        if self.states == 0:
+            return self._latex_partitioned_stateless()
+
+        lines = [
+            r'\[',
+            r'\left(',
+            (r'\begin{array}'
+             + r'{' + 'rll' * self.states + '|' + 'rll' * self.inputs + '}')
+            ]
+
+        for Ai, Bi in zip(asarray(self.A), asarray(self.B)):
+            lines.append('&'.join([_f2s(Aij) for Aij in Ai]
+                                  + [_f2s(Bij) for Bij in Bi])
+                         + '\\\\')
+        lines.append(r'\hline')
+        for Ci, Di in zip(asarray(self.C), asarray(self.D)):
+            lines.append('&'.join([_f2s(Cij) for Cij in Ci]
+                                  + [_f2s(Dij) for Dij in Di])
+                         + '\\\\')
+
+        lines.extend([
+            r'\end{array}'
+            r'\right)',
+            r'\]'])
+
+        return '\n'.join(lines)
+
+    def _latex_separate(self):
+        """Separate matrices LaTeX representation of state-space model
+
+        Model is presented as separate, named, A, B, C, and D matrices.
+
+        Returns
+        -------
+        s : string with LaTeX representation of model
+        """
+        lines = [
+            r'\[',
+            r'\begin{array}{ll}',
+            ]
+
+        def fmt_matrix(matrix, name):
+            matlines = [name
+                        + r' = \left(\begin{array}{'
+                        + 'rll' * matrix.shape[1]
+                        + '}']
+            for row in asarray(matrix):
+                matlines.append('&'.join(_f2s(entry) for entry in row)
+                                + '\\\\')
+            matlines.extend([
+                r'\end{array}'
+                r'\right)'])
+            return matlines
+
+        if self.states > 0:
+            lines.extend(fmt_matrix(self.A, 'A'))
+            lines.append('&')
+            lines.extend(fmt_matrix(self.B, 'B'))
+            lines.append('\\\\')
+
+            lines.extend(fmt_matrix(self.C, 'C'))
+            lines.append('&')
+        lines.extend(fmt_matrix(self.D, 'D'))
+
+        lines.extend([
+            r'\end{array}',
+            r'\]'])
+
+        return '\n'.join(lines)
+
+    def _repr_latex_(self):
+        """LaTeX representation of state-space model
+
+        Output is controlled by config options statesp.latex_repr_type
+        and statesp.latex_num_format.
+
+        The output is primarily intended for Jupyter notebooks, which
+        use MathJax to render the LaTeX, and the results may look odd
+        when processed by a 'conventional' LaTeX system.
+
+        Returns
+        -------
+        s : string with LaTeX representation of model
+
+        """
+        if config.defaults['statesp.latex_repr_type'] == 'partitioned':
+            return self._latex_partitioned()
+        elif config.defaults['statesp.latex_repr_type'] == 'separate':
+            return self._latex_separate()
+        else:
+            cfg = config.defaults['statesp.latex_repr_type']
+            raise ValueError("Unknown statesp.latex_repr_type '{cfg}'".format(cfg=cfg))
+
     # Negation of a system
     def __neg__(self):
         """Negate a state space system."""
diff --git a/control/tests/statesp_test.py b/control/tests/statesp_test.py
@@ -20,6 +20,7 @@
 from control.tests.conftest import ismatarrayout, slycotonly
 from control.xferfcn import TransferFunction, ss2tf
 
+from .conftest import editsdefaults
 
 class TestStateSpace:
     """Tests for the StateSpace class."""
@@ -840,3 +841,64 @@ def test_statespace_defaults(self, matarrayout):
         for k, v in _statesp_defaults.items():
             assert defaults[k] == v, \
                 "{} is {} but expected {}".format(k, defaults[k], v)
+
+
+# test data for test_latex_repr below
+LTX_G1 = StateSpace([[np.pi, 1e100], [-1.23456789, 5e-23]],
+                    [[0], [1]],
+                    [[987654321, 0.001234]],
+                    [[5]])
+
+LTX_G2 = StateSpace([],
+                    [],
+                    [],
+                    [[1.2345, -2e-200], [-1, 0]])
+
+LTX_G1_REF = {
+    'p3_p' : '\\[\n\\left(\n\\begin{array}{rllrll|rll}\n3.&\\hspace{-1em}14&\\hspace{-1em}\\phantom{\\cdot}&1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{100}&0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n-1.&\\hspace{-1em}23&\\hspace{-1em}\\phantom{\\cdot}&5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-23}&1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\hline\n9.&\\hspace{-1em}88&\\hspace{-1em}\\cdot10^{8}&0.&\\hspace{-1em}00123&\\hspace{-1em}\\phantom{\\cdot}&5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\]',
+
+    'p5_p' : '\\[\n\\left(\n\\begin{array}{rllrll|rll}\n3.&\\hspace{-1em}1416&\\hspace{-1em}\\phantom{\\cdot}&1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{100}&0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n-1.&\\hspace{-1em}2346&\\hspace{-1em}\\phantom{\\cdot}&5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-23}&1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\hline\n9.&\\hspace{-1em}8765&\\hspace{-1em}\\cdot10^{8}&0.&\\hspace{-1em}001234&\\hspace{-1em}\\phantom{\\cdot}&5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\]',
+
+    'p3_s' : '\\[\n\\begin{array}{ll}\nA = \\left(\\begin{array}{rllrll}\n3.&\\hspace{-1em}14&\\hspace{-1em}\\phantom{\\cdot}&1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{100}\\\\\n-1.&\\hspace{-1em}23&\\hspace{-1em}\\phantom{\\cdot}&5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-23}\\\\\n\\end{array}\\right)\n&\nB = \\left(\\begin{array}{rll}\n0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\\\\nC = \\left(\\begin{array}{rllrll}\n9.&\\hspace{-1em}88&\\hspace{-1em}\\cdot10^{8}&0.&\\hspace{-1em}00123&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n&\nD = \\left(\\begin{array}{rll}\n5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\end{array}\n\\]',
+
+    'p5_s' : '\\[\n\\begin{array}{ll}\nA = \\left(\\begin{array}{rllrll}\n3.&\\hspace{-1em}1416&\\hspace{-1em}\\phantom{\\cdot}&1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{100}\\\\\n-1.&\\hspace{-1em}2346&\\hspace{-1em}\\phantom{\\cdot}&5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-23}\\\\\n\\end{array}\\right)\n&\nB = \\left(\\begin{array}{rll}\n0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\\\\nC = \\left(\\begin{array}{rllrll}\n9.&\\hspace{-1em}8765&\\hspace{-1em}\\cdot10^{8}&0.&\\hspace{-1em}001234&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n&\nD = \\left(\\begin{array}{rll}\n5\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\end{array}\n\\]',
+}
+
+LTX_G2_REF = {
+    'p3_p' : '\\[\n\\left(\n\\begin{array}{rllrll}\n1.&\\hspace{-1em}23&\\hspace{-1em}\\phantom{\\cdot}&-2\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-200}\\\\\n-1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}&0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\]',
+
+    'p5_p' : '\\[\n\\left(\n\\begin{array}{rllrll}\n1.&\\hspace{-1em}2345&\\hspace{-1em}\\phantom{\\cdot}&-2\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-200}\\\\\n-1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}&0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\]',
+
+    'p3_s' : '\\[\n\\begin{array}{ll}\nD = \\left(\\begin{array}{rllrll}\n1.&\\hspace{-1em}23&\\hspace{-1em}\\phantom{\\cdot}&-2\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-200}\\\\\n-1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}&0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\end{array}\n\\]',
+
+    'p5_s' : '\\[\n\\begin{array}{ll}\nD = \\left(\\begin{array}{rllrll}\n1.&\\hspace{-1em}2345&\\hspace{-1em}\\phantom{\\cdot}&-2\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\cdot10^{-200}\\\\\n-1\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}&0\\phantom{.}&\\hspace{-1em}&\\hspace{-1em}\\phantom{\\cdot}\\\\\n\\end{array}\\right)\n\\end{array}\n\\]',
+}
+
+refkey_n = {None: 'p3', '.3g': 'p3', '.5g': 'p5'}
+refkey_r = {None: 'p', 'partitioned': 'p', 'separate': 's'}
+
+
+@pytest.mark.parametrize(" g,      ref",
+                         [(LTX_G1, LTX_G1_REF),
+                          (LTX_G2, LTX_G2_REF)])
+@pytest.mark.parametrize("repr_type", [None, "partitioned", "separate"])
+@pytest.mark.parametrize("num_format", [None, ".3g", ".5g"])
+def test_latex_repr(g, ref, repr_type, num_format, editsdefaults):
+    """Test `._latex_repr_` with different config values
+
+    This is a 'gold image' test, so if you change behaviour,
+    you'll need to regenerate the reference results.
+    Try something like:
+        control.reset_defaults()
+        print(f'p3_p : {g1._repr_latex_()!r}')
+    """
+    from control import set_defaults
+    if num_format is not None:
+        set_defaults('statesp', latex_num_format=num_format)
+
+    if repr_type is not None:
+        set_defaults('statesp', latex_repr_type=repr_type)
+
+    refkey = "{}_{}".format(refkey_n[num_format], refkey_r[repr_type])
+    assert g._repr_latex_() == ref[refkey]
+
