Code Quality

- switched to google style docstrings
- removed dead code
- started to annotate types

__init__.py

@@ -2,4 +2,15 @@ from pystencils.field import Field, FieldType, extractCommonSubexpressions
 from pystencils.data_types import TypedSymbol
 from pystencils.slicing import makeSlice
 from pystencils.kernelcreation import createKernel, createIndexedKernel
-from pystencils.display_utils import showCode
+from pystencils.display_utils import showCode, toDot
+from pystencils.equationcollection import EquationCollection
+from sympy.codegen.ast import Assignment as Assign
+
+
+__all__ = ['Field', 'FieldType', 'extractCommonSubexpressions',
+           'TypedSymbol',
+           'makeSlice',
+           'createKernel', 'createIndexedKernel',
+           'showCode', 'toDot',
+           'EquationCollection',
+           'Assign']

astnodes.py

@@ -473,6 +473,11 @@ class SympyAssignment(Node):
     def __repr__(self):
         return repr(self.lhs) + " = " + repr(self.rhs)
 
+    def _repr_html_(self):
+        printed_lhs = sp.latex(self.lhs)
+        printed_rhs = sp.latex(self.rhs)
+        return f"${printed_lhs} = {printed_rhs}$"
+
 
 class ResolvedFieldAccess(sp.Indexed):
     def __new__(cls, base, linearizedIndex, field, offsets, idxCoordinateValues):

display_utils.py

@@ -42,141 +42,3 @@ def showCode(ast):
         def __repr__(self):
             return generateC(self.ast)
     return CodeDisplay(ast)
-
-
-
-# ----------------- Embedding of animations as videos in IPython notebooks ---------------------------------------------
-
-
-# -------   Version 1: Animation is embedded as an HTML5 Video tag  ---------------------------------------
-
-VIDEO_TAG = """<video controls width="100%">
- <source src="data:video/x-m4v;base64,{0}" type="video/mp4">
- Your browser does not support the video tag.
-</video>"""
-
-
-def __anim_to_html(anim, fps):
-    from tempfile import NamedTemporaryFile
-    import base64
-
-    if not hasattr(anim, '_encoded_video'):
-        with NamedTemporaryFile(suffix='.mp4') as f:
-            anim.save(f.name, fps=fps, extra_args=['-vcodec', 'libx264', '-pix_fmt',
-                                                   'yuv420p', '-profile:v', 'baseline', '-level', '3.0'])
-            video = open(f.name, "rb").read()
-        anim._encoded_video = base64.b64encode(video).decode('ascii')
-
-    return VIDEO_TAG.format(anim._encoded_video)
-
-
-def disp_as_video(anim, fps=30, show=True, **kwargs):
-    import matplotlib.pyplot as plt
-    from IPython.display import HTML
-
-    try:
-        plt.close(anim._fig)
-        res = __anim_to_html(anim, fps)
-        if show:
-            return HTML(res)
-        else:
-            return HTML("")
-    except KeyboardInterrupt:
-        pass
-# -------   Version 2: Animation is shown in extra matplotlib window ----------------------------------
-
-
-def disp_extra_window(animation, *args,**kwargs):
-    import matplotlib.pyplot as plt
-
-    fig = plt.gcf()
-    try:
-      fig.canvas.manager.window.raise_()
-    except Exception:
-      pass
-    plt.show()
-
-
-# -------   Version 3: Animation is shown in images that are updated directly in website --------------
-
-def disp_image_update(animation, iterations=10000, *args, **kwargs):
-    from IPython import display
-    import matplotlib.pyplot as plt
-
-    try:
-        fig = plt.gcf()
-        animation._init_draw()
-        for i in range(iterations):
-            display.display(fig)
-            animation._step()
-            display.clear_output(wait=True)
-    except KeyboardInterrupt:
-        pass
-
-
-# Dispatcher
-
-animation_display_mode = 'imageupdate'
-display_animation_func = None
-
-
-def disp(*args, **kwargs):
-    from IPython import get_ipython
-    ipython = get_ipython()
-    if not ipython:
-        return
-
-    if not display_animation_func:
-        raise Exception("Call set_display_mode first")
-    return display_animation_func(*args, **kwargs)
-
-
-def set_display_mode(mode):
-    from IPython import get_ipython
-    ipython = get_ipython()
-    if ipython is None:
-        return
-    global animation_display_mode
-    global display_animation_func
-    animation_display_mode = mode
-    if animation_display_mode == 'video':
-        ipython.magic("matplotlib inline")
-        display_animation_func = disp_as_video
-    elif animation_display_mode == 'window':
-        ipython.magic("matplotlib qt")
-        display_animation_func = disp_extra_window
-    elif animation_display_mode == 'imageupdate':
-        ipython.magic("matplotlib inline")
-        display_animation_func = disp_image_update
-    else:
-        raise Exception("Unknown mode. Available modes 'imageupdate', 'video' and 'window' ")
-
-
-set_display_mode('video')
-
-
-# --------------------- Convenience functions --------------------------------------------------------------------------
-
-
-def makeSurfacePlotAnimation(runFunction, frames=90, interval=30):
-    from mpl_toolkits.mplot3d import Axes3D
-    import matplotlib.animation as animation
-    import matplotlib.pyplot as plt
-    from matplotlib import cm
-
-    fig = plt.figure()
-    ax = fig.add_subplot(111, projection='3d')
-    X, Y, data = runFunction(1)
-    ax.plot_surface(X, Y, data, rstride=2, cstride=2, color='b', cmap=cm.coolwarm,)
-    ax.set_zlim(-1.0, 1.0)
-
-    def updatefig(*args):
-        X, Y, data = runFunction(1)
-        ax.clear()
-        plot = ax.plot_surface(X, Y, data, rstride=2, cstride=2, color='b', cmap=cm.coolwarm,)
-        ax.set_zlim(-1.0, 1.0)
-        return plot,
-
-    return animation.FuncAnimation(fig, updatefig, interval=interval, frames=frames, blit=False)
-
-

jupytersetup.py

@@ -5,6 +5,9 @@ from tempfile import NamedTemporaryFile
 import base64
 import sympy as sp
 
+__all__ = ['log_progress', 'makeImshowAnimation', 'makeSurfacePlotAnimation',
+           'disp', 'setDisplayMode']
+
 
 def log_progress(sequence, every=None, size=None, name='Items'):
     """Copied from https://github.com/alexanderkuk/log-progress"""
@@ -64,7 +67,6 @@ def log_progress(sequence, every=None, size=None, name='Items'):
         )
 
 
-
 VIDEO_TAG = """<video controls width="80%">
  <source src="data:video/x-m4v;base64,{0}" type="video/mp4">
  Your browser does not support the video tag.
@@ -96,6 +98,28 @@ def makeImshowAnimation(grid, gridUpdateFunction, frames=90, **kwargs):
     return animation.FuncAnimation(fig, partial(updatefig, image=grid), frames=frames)
 
 
+def makeSurfacePlotAnimation(runFunction, frames=90, interval=30):
+    from mpl_toolkits.mplot3d import Axes3D
+    import matplotlib.animation as animation
+    import matplotlib.pyplot as plt
+    from matplotlib import cm
+
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection='3d')
+    X, Y, data = runFunction(1)
+    ax.plot_surface(X, Y, data, rstride=2, cstride=2, color='b', cmap=cm.coolwarm,)
+    ax.set_zlim(-1.0, 1.0)
+
+    def updatefig(*args):
+        X, Y, data = runFunction(1)
+        ax.clear()
+        plot = ax.plot_surface(X, Y, data, rstride=2, cstride=2, color='b', cmap=cm.coolwarm,)
+        ax.set_zlim(-1.0, 1.0)
+        return plot,
+
+    return animation.FuncAnimation(fig, updatefig, interval=interval, frames=frames, blit=False)
+
+
 # -------   Version 1: Embed the animation as HTML5 video --------- ----------------------------------
 
 def displayAsHtmlVideo(anim, fps=30, show=True, **kwargs):
@@ -162,6 +186,8 @@ def disp(*args, **kwargs):
 def setDisplayMode(mode):
     from IPython import get_ipython
     ipython = get_ipython()
+    if not ipython:
+        return
     global animation_display_mode
     global display_animation_func
     animation_display_mode = mode