Merge pull request #305 from eric-wieser/transformation-pretty

Add _repr_pretty_ to transformations

clifford/test/test_transformations.py

@@ -1,5 +1,8 @@
+import textwrap
+
 import numpy as np
 import pytest
+from IPython.lib import pretty
 
 from clifford import Layout, ConformalLayout, BasisVectorIds
 from clifford import transformations
@@ -78,6 +81,14 @@ def test_invariants(self, g2):
         # test that distributivity is respected
         assert f(g2.scalar + 2*e1 + 3*(e1^e2)) == f(g2.scalar) + 2*f(e1) + 3*f(e1^e2)
 
+        assert pretty.pretty(f) == textwrap.dedent("""\
+        OutermorphismMatrix(array([[ 0,  1],
+                                   [-1,  0]]),
+                            Layout([1, 1],
+                                   ids=BasisVectorIds(['u', 'v']),
+                                   order=BasisBladeOrder.shortlex(2),
+                                   names=['', 'eu', 'ev', 'euv']))""")
+
     def test_between_layouts(self, g2, g3):
         matrix = np.array([[1, 0],
                            [0, 1],
@@ -93,6 +104,19 @@ def test_between_layouts(self, g2, g3):
         assert f(e1) == ex
         assert f(e2) == ey
 
+        assert pretty.pretty(f) == textwrap.dedent("""\
+        OutermorphismMatrix(array([[1, 0],
+                                   [0, 1],
+                                   [0, 0]]),
+                            layout_src=Layout([1, 1],
+                                              ids=BasisVectorIds(['u', 'v']),
+                                              order=BasisBladeOrder.shortlex(2),
+                                              names=['', 'eu', 'ev', 'euv']),
+                            layout_dst=Layout([1, 1, 1],
+                                              ids=BasisVectorIds(['x', 'y', 'z']),
+                                              order=BasisBladeOrder.shortlex(3),
+                                              names=['', 'ex', 'ey', 'ez', 'exy', 'exz', 'eyz', 'exyz']))""")
+
 
 class TestLinearMatrix:
     def test_same_layout(self, g2):
@@ -108,6 +132,16 @@ def test_same_layout(self, g2):
         x = 1 + 2*e2 + 3*e1*e2
         assert f_dual(x) == x.dual()
 
+        assert pretty.pretty(f_dual) == textwrap.dedent("""\
+        LinearMatrix(array([[ 0,  0,  0,  1],
+                            [ 0,  0,  1,  0],
+                            [ 0, -1,  0,  0],
+                            [-1,  0,  0,  0]]),
+                     Layout([1, 1],
+                            ids=BasisVectorIds(['u', 'v']),
+                            order=BasisBladeOrder.shortlex(2),
+                            names=['', 'eu', 'ev', 'euv']))""")
+
     def test_between_layouts(self, g2, g3):
         matrix = np.array([[1, 0, 0, 0],
                            [0, 1, 0, 0],

clifford/transformations.py

@@ -147,6 +147,21 @@ def __call__(self, mv: MultiVector) -> MultiVector:
             raise ValueError("Object is from the wrong layout")
         # superclass should produce the result
 
+    def _repr_pretty_layout_args_helper(self, p):
+        """ helper function to be called by _repr_pretty_ in subclasses to print the layout arguments """
+        if self.layout_src is not self.layout_dst:
+            # only show both if they're different
+            prefix = 'layout_src='
+            with p.group(len(prefix), prefix, ''):
+                p.pretty(self.layout_src)
+                p.text(',')
+            p.breakable()
+            prefix = 'layout_dst='
+            with p.group(len(prefix), prefix, ''):
+                p.pretty(self.layout_dst)
+        else:
+            p.pretty(self.layout_src)
+
 
 class LinearMatrix(FixedLayout, Linear):
     """
@@ -274,6 +289,18 @@ def from_rotor(cls, rotor: MultiVector) -> 'LinearMatrix':
             rotor.layout, rotor.layout
         )
 
+    def _repr_pretty_(self, p, cycle):
+        if cycle:
+            raise RuntimeError("Should not be cyclic")
+
+        prefix = '{}('.format(type(self).__name__)
+
+        with p.group(len(prefix), prefix, ')'):
+            p.pretty(self._matrix)
+            p.text(',')
+            p.breakable()
+            self._repr_pretty_layout_args_helper(p)
+
 
 class OutermorphismMatrix(LinearMatrix):
     r"""
@@ -345,6 +372,19 @@ def __init__(self, matrix, layout_src: Layout, layout_dst: Layout = None):
         )
 
         super().__init__(full_matrix, layout_src, layout_dst)
+        self._vector_matrix = matrix
+
+    def _repr_pretty_(self, p, cycle):
+        if cycle:
+            raise RuntimeError("Should not be cyclic")
+
+        prefix = '{}('.format(type(self).__name__)
+
+        with p.group(len(prefix), prefix, ')'):
+            p.pretty(self._vector_matrix)
+            p.text(',')
+            p.breakable()
+            self._repr_pretty_layout_args_helper(p)
 
 
 def between_basis_vectors(layout_src: Layout, layout_dst: Layout, mapping: Dict[Any, Any]=None) -> OutermorphismMatrix: