Added IPython LaTeX representation method for StateSpace objects

roryyorke · roryyorke · commit 96cd45a09783 · 2020-08-15T09:55:54.000+02:00
Added StateSpace method `_repr_latex_`, which returns a
MathJax-centric LaTeX representation of the instance.

Added two `statesp` configuration options for this representation.
One affects number formatting, and the other chooses the output type.
diff --git a/control/statesp.py b/control/statesp.py
@@ -73,6 +73,8 @@
     'statesp.use_numpy_matrix': True,
     'statesp.default_dt': None,
     'statesp.remove_useless_states': True, 
+    'statesp.latex_num_format': '.3g',
+    'statesp.latex_repr_type': 'partitioned',
     }
 
 
@@ -122,6 +124,34 @@ 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])
 
@@ -152,6 +182,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['latex_num_format']` and
+    `control.config.defaults['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['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['latex_repr_type']` must either be
+    `'partioned'` 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
@@ -292,6 +340,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
@@ -561,6 +561,69 @@ def test_str(self):
         assert str(sysdt1) == tref + "\ndt = 1.0\n"
 
 
+    def test_latex_repr(self):
+        # 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'g1_p3_p = {g1._repr_latex_()!r}')
+        #   control.set_defaults('statesp', latex_num_format='.5g')
+        #   print(f'g1_p5_p = {g1._repr_latex_()!r}')
+        #   control.reset_defaults()
+        #   control.set_defaults('statesp', latex_repr_type='separate')
+        #   print(f'g1_p3_s = {g1._repr_latex_()!r}')
+        #   control.set_defaults('statesp', latex_num_format='.5g')
+        #   print(f'g1_p5_s = {g1._repr_latex_()!r}')
+        from control import set_defaults, reset_defaults
+        g1 = StateSpace([[np.pi, 1e100],[-1.23456789, 5e-23]],
+                        [[0], [1]],
+                        [[987654321, 0.001234]],
+                        [[5]])
+        g2 = StateSpace([],
+                        [],
+                        [],
+                        [[1.2345,-2e-200],[-1,0]])
+
+        g1_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\\]'
+        g1_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\\]'
+        g1_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\\]'
+        g1_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\\]'
+
+        g2_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\\]'
+        g2_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\\]'
+        g2_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\\]'
+        g2_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\\]'
+
+
+        try:
+            reset_defaults()
+            self.assertEqual(g1_p3_p, g1._repr_latex_())
+            set_defaults('statesp', latex_num_format='.5g')
+            self.assertEqual(g1_p5_p, g1._repr_latex_())
+
+            reset_defaults()
+            set_defaults('statesp', latex_repr_type='separate')
+            self.assertEqual(g1_p3_s, g1._repr_latex_())
+            set_defaults('statesp', latex_num_format='.5g')
+            self.assertEqual(g1_p5_s, g1._repr_latex_())
+
+            reset_defaults()
+            self.assertEqual(g2_p3_p, g2._repr_latex_())
+            set_defaults('statesp', latex_num_format='.5g')
+            self.assertEqual(g2_p5_p, g2._repr_latex_())
+
+            reset_defaults()
+            set_defaults('statesp', latex_repr_type='separate')
+            self.assertEqual(g2_p3_s, g2._repr_latex_())
+            set_defaults('statesp', latex_num_format='.5g')
+            self.assertEqual(g2_p5_s, g2._repr_latex_())
+
+        finally:
+            reset_defaults()
+
+
+
 class TestRss(unittest.TestCase):
     """These are tests for the proper functionality of statesp.rss."""
 
