Commit eadeaadf authored by Martin Bauer's avatar Martin Bauer
Browse files

Code Quality

- switched to google style docstrings
- removed dead code
- started to annotate types
parent 170490ab
......@@ -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']
......@@ -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):
......
......@@ -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)
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment