diff --git a/.travis.yml b/.travis.yml
index 7c996575c..8324c838f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -17,6 +17,7 @@ addons:
- libopenmpi-dev
- mpi-default-bin
- openmpi-bin
+ - ffmpeg
##################################################
# environment variables applied to all build cases
@@ -64,7 +65,7 @@ matrix:
- *common_deps
- libhdf5-serial-dev
install:
- - pip install numpy mpi4py scipy h5py
+ - pip install numpy mpi4py scipy h5py matplotlib
- python: "2.7"
env:
- MPICONF="--with-mpi"
@@ -79,7 +80,7 @@ matrix:
- *common_deps
- libhdf5-serial-dev
install:
- - pip install numpy mpi4py scipy h5py
+ - pip install numpy mpi4py scipy h5py matplotlib
- python: "2.7"
env:
- MPICONF="--with-mpi"
@@ -95,7 +96,7 @@ matrix:
- *common_deps
- libhdf5-openmpi-dev
install:
- - pip install numpy mpi4py scipy
+ - pip install numpy mpi4py scipy matplotlib
- pip install --no-binary=h5py h5py
- python: "3.7"
env:
@@ -110,7 +111,7 @@ matrix:
- *common_deps
- libhdf5-serial-dev
install:
- - pip install numpy mpi4py scipy h5py coverage coveralls
+ - pip install numpy mpi4py scipy h5py coverage coveralls matplotlib
- python: "3.7"
env:
- MPICONF="--with-mpi"
@@ -125,7 +126,7 @@ matrix:
- *common_deps
- libhdf5-serial-dev
install:
- - pip install numpy mpi4py scipy h5py
+ - pip install numpy mpi4py scipy h5py matplotlib
- python: "3.7"
env:
- MPICONF="--with-mpi"
@@ -141,7 +142,7 @@ matrix:
- *common_deps
- libhdf5-openmpi-dev
install:
- - pip install numpy mpi4py scipy
+ - pip install numpy mpi4py scipy matplotlib
- pip install --no-binary=h5py h5py
diff --git a/doc/docs/Python_User_Interface.md b/doc/docs/Python_User_Interface.md
index 996ed88e0..cf05391bf 100644
--- a/doc/docs/Python_User_Interface.md
+++ b/doc/docs/Python_User_Interface.md
@@ -1525,10 +1525,112 @@ fr = mp.FluxRegion(volume=mp.GDSII_vol(fname, layer, zmin, zmax))
### Data Visualization
+This module attempts to provide some simple visualization functions that should just work in most cases to provide some graphical intuition for what you're working with. The spirit of the module is to provide functions that can be called with _no customization options whatsoever_ and will do useful relevant things by default, but which can also be customized in cases where you _do_ want to take the time to spruce up the output.
+
+**`Simulation.plot2D(ax=None, output_plane=None fields=None, labels=False, eps_parameters=None,boundary_parameters=None, source_parameters=None,monitor_parameters=None, field_parameters=None)`**
+—
+Plots a 2D cross section of the simulation domain using `matplotlib`. The plot includes the simulation geometry, boundary layers, sources, and monitors. Can also superimpose fields on top of 2D slice. Requires [matplotlib](https://matplotlib.org).
+
+* `ax`: A `matplotlib` axis object. `plot2D()` will add plot objects, like lines, patches, and scatter plots, to this object. If no `ax` is supplied, then the routine will create a new figure and grab its axis.
+* `output_plane`: A `Volume` object that specifies the plane over which to plot. Must be a 2D volume and a subset of the entire simulation volume (i.e. it should not span outside the domain at all).
+* `fields`: The field component (`mp.Ex`, `mp.Ey`, `mp.Ez`, `mp.Hx`, `mp.Hy`, `mp.Hz`) to superimpose over the simulation geometry. Default is `None`, where no fields are superimposed.
+* `labels`: If `True`, then labels will appear over each of the simulation elements.
+* `eps_parameters`: A `dict` of optional plotting parameters that override the default plotting parameters for the geometry.
+ * `interpolation='spline36'`: interpolation algorithm used to upsample the pixels.
+ * `cmap='binary'`: the color map of the geometry
+ * `alpha=1.0`: transparency of geometry
+* `boundary_parameters`:A `dict` of optional plotting parameters that override the default plotting parameters for the boundary layers. Available parameters include:
+ * `alpha=1.0`: transparency of boundary layers
+ * `facecolor='g'`: color of polygon face
+ * `edgecolor='g'`: color of outline stroke
+ * `linewidth=1`: line width of outline stroke
+ * `hatch='\'`: hatching pattern
+* `source_parameters`: A `dict` of optional plotting parameters that override the default plotting parameters for the sources.
+ * `color='r`: color of line and pt sources
+ * `alpha=1.0`: transparency of source
+ * `facecolor='none'`: color of polygon face for planar sources
+ * `edgecolor='r'`: color of outline stroke for planar sources
+ * `linewidth=1`: line width of outline stroke
+ * `hatch='\'`: hatching pattern
+ * `label_color='r'`: color of source labels
+ * `label_alpha=0.3`: transparency of source label box
+ * `offset=20`: distance from source center and label box
+* `monitor_parameters`: A `dict` of optional plotting parameters that override the default plotting parameters for the monitors.
+ * `color='g`: color of line and point monitors
+ * `alpha=1.0`: transparency of monitors
+ * `facecolor='none'`: color of polygon face for planar monitors
+ * `edgecolor='r'`: color of outline stroke for planar monitors
+ * `linewidth=1`: line width of outline stroke
+ * `hatch='\'`: hatching pattern
+ * `label_color='g'`: color of source labels
+ * `label_alpha=0.3`: transparency of monitor label box
+ * `offset=20`: distance from monitor center and label box
+* `field_parameters`: A `dict` of optional plotting parameters that override the default plotting parameters for the fields.
+ * `interpolation='spline36'`: interpolation function used to upsample field pixels
+ * `cmap='RdBu'`: color map for field pixels
+ * `alpha=0.6`: transparency of fields
+
+**`Simulation.plot3D()`**
+— Uses Mayavi to render a 3D simulation domain. The simulation object must be 3D. Can also be embedded in jupyter notebooks.
+
**`Simulation.visualize_chunks()`**
—
Displays an interactive image of how the cell is divided into chunks. Each rectangular region is a chunk, and each color represents a different processor. Requires [matplotlib](https://matplotlib.org).
+#### Animate2D
+A class used to record the fields while the simulation is timestepping (as a result of a `run` function). The user first intializes the object by specifying the simulation object and the field component of interest. The object can then be passed to any arbitrary step function modifier. For example, one can record the `E_z` component every 1 MEEP time units for 25 units with:
+```python
+animate = mp.Animate2D(sim,mp.Ez)
+sim.run(mp.at_every(1,animate),until=25)
+```
+
+By default, the object saves each frame as a png image into memory (not disk). This is typically more memory efficient than storing the actual fields. If the user sets the `normalize` argument, then the object will save the actual field information as a `numpy` array to be normalized for post processing. The user can also choose to update the fields of a figure in realtime by setting the `realtime` flag. This does not work for IPython or Jupyter notebooks, however.
+
+Once the simulation is run, the user can output the animation as an interactive JSHTML object, an mp4, or a GIF.
+
+Multiple Animate2D objects can be initialized and passed to the run function to track different volume locations (using `mp.in_volume`) or field components.
+
+Properties:
+
+**`sim`**
+— simulation object.
+
+**`fields`**
+— field component to record at each time step.
+
+**`f=None`**
+— optional `matplotlib` figure object that the routine will update on each call. If not supplied, then a new one will be created upon initialization.
+
+**`realtime=True`**
+— whether or not to update a figure window in realtime as the simulation progresses. Disabled by default. Not compatible with IPython/Jupyter notebooks.
+
+**`normalize=False`**
+— records fields at each time step in memory in a numpy array and then normalizes the result for future post-processing.
+
+**`plot_modifiers=None`**
+— A list of functions that can modify the figure's `axis` object. Each function modifier accepts a single argument, an `axis` object, and must return that same axis object. The following modifier changes the `xlabel`:
+```python
+def mod1(ax):
+ ax.set_xlabel('Testing')
+ return ax
+
+plot_modifiers = [mod1]
+```
+
+**`**customization_args`**
+— customization keyword arguments pass to `plot2D()` (i.e. `labels`, `eps_parameters`, `boundary_parameters`, etc.)
+
+Methods:
+
+**`Animate2D.to_jshtml(fps)`**
+— Outputs an interactable JSHTML animation object that is embeddable in Jupyter notebooks. The object is packaged with controls to manipulate the video's playback. User must specify a frame rate `fps` in frames per second.
+
+**`Animate2D.to_mp4(fps,filename)`**
+— Generates and outputs an mp4 video file of the animation with the filename, `filename`, and the frame rate, `fps`. Default encoding is h264 with yuv420p format. Requires `ffmpeg`.
+
+**`Animate2D.to_mp4(fps,filename)`**
+— Generates and outputs a GIF file of the animation with the filename, `filename` and the frame rate, `fps`. Requires `ffmpeg`.
+
Run and Step Functions
----------------------
diff --git a/python/Makefile.am b/python/Makefile.am
index 521f6b7da..95f843a06 100644
--- a/python/Makefile.am
+++ b/python/Makefile.am
@@ -68,6 +68,7 @@ TESTS = \
$(TEST_DIR)/simulation.py \
$(TEST_DIR)/source.py \
$(TEST_DIR)/user_defined_material.py \
+ $(TEST_DIR)/visualization.py \
$(WVG_SRC_TEST)
if WITH_COVERAGE
@@ -187,6 +188,7 @@ HL_IFACE = \
$(srcdir)/geom.py \
$(srcdir)/simulation.py \
$(srcdir)/source.py \
+ $(srcdir)/visualization.py \
$(srcdir)/materials.py
pkgpython_PYTHON = __init__.py $(HL_IFACE)
@@ -224,3 +226,5 @@ clean-local:
distclean-local:
rm -f *.h5
+ rm -f *.mp4
+ rm -f *.gif
diff --git a/python/examples/README.md b/python/examples/README.md
new file mode 100644
index 000000000..df668f7ba
--- /dev/null
+++ b/python/examples/README.md
@@ -0,0 +1,42 @@
+# Tutorials and Examples
+
+Meep simulations are Python scripts which involve specifying the device geometry, materials, current sources, monitor fields, and everything else necessary to set up a calculation. A Python script provides the flexibility to customize the simulation for practically any application particularly those involving parameter sweeps and optimization.
+
+Python libraries such as NumPy, SciPy, and Matplotlib can be used to augment the simulation functionality and will also be demonstrated. Much of the functionality of the low-level C++ interface has been abstracted in Python which means that you don't need to be an experienced programmer to set up simulations. Reasonable defaults are available where necessary.
+
+Several tutorials and examples are found here. These tutorials are meant to illustrate Meep's various features in an interactive and application-oriented manner.
+
+## iPython/Jupyter Notebooks
+
+Jupyter notebooks are interactive, browser based framework for displaying text, running python code, and visualizing results. There are several ways to read these notebooks online:
+
+### Local
+
+The recommended method to run the tutorial notebooks is by 1.) installing `meep` via `conda`, 2.) cloning this repo to your local drive, and 3.) launch a notebook server in this directory using `jupyter notebook`.
+
+### nbviewer
+
+`nbviewer` is a web platform that can render interactive features found inside Jupyter notebooks openly stored on web-servers. While `nbviewer` can't run python code, it can execute stored javascript code used to animate the simulations.
+
+### GitHub
+
+GitHub is able to render some of the smaller notebooks as plain text. However, they are not interactive and are often too large for GitHub.
+
+### Tutorials
+
+Below are summaries for each tutorial, along with the features the tutorials highlight. While there is no particular order to the tutorials, they progressively incorporate more complicated features.
+
+1. __`straight-waveguide.ipynb`__ -
+A simple 2D straight waveguide tutorial that explores basic meep features like `geometry`, `sources`, and `PML` layers. The tutorial also explores basic visualization and animation features.
+
+2. __`bent-waveguide.ipynb`__-
+A followup to the 2D straight waveguide tutorial by adding a bend.
+
+3. __`bend-flux.ipynb`__-
+Using the previous bent waveguide example, this tutorial calculates the loss, transmission, and reflection that the bent waveguide undergoes.
+
+4. __`ring.ipynb`__ -
+Computes the resonant mode frequencies of a 2D ring resonator using `harminv`.
+
+5. __`visualization.ipynb`__ -
+Demonstrates various visualization and animation features.
diff --git a/python/meep.i b/python/meep.i
index dca55ae30..efe060ada 100644
--- a/python/meep.i
+++ b/python/meep.i
@@ -1514,6 +1514,12 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where
SourceTime,
check_positive,
)
+ from .visualization import (
+ plot2D,
+ plot3D,
+ plot_fields,
+ Animate2D
+ )
if with_mpi():
try:
diff --git a/python/simulation.py b/python/simulation.py
index 2724c2b93..8aad4150b 100644
--- a/python/simulation.py
+++ b/python/simulation.py
@@ -18,6 +18,7 @@
import meep as mp
from meep.geom import Vector3, init_do_averaging
from meep.source import EigenModeSource, check_positive
+import meep.visualization as vis
try:
@@ -158,6 +159,48 @@ def __init__(self, center, size=Vector3(), dims=2, is_cylindrical=False):
vec2 = py_v3_to_vec(self.dims, v2, is_cylindrical)
self.swigobj = mp.volume(vec1, vec2)
+
+ def get_vertices(self):
+ xmin = self.center.x - self.size.x/2
+ xmax = self.center.x + self.size.x/2
+ ymin = self.center.y - self.size.y/2
+ ymax = self.center.y + self.size.y/2
+ zmin = self.center.z - self.size.z/2
+ zmax = self.center.z + self.size.z/2
+
+ # Iterate over and remove duplicates for collapsed dimensions (i.e. min=max))
+ return [Vector3(x,y,z) for x in list(set([xmin,xmax])) for y in list(set([ymin,ymax])) for z in list(set([zmin,zmax]))]
+
+
+ def get_edges(self):
+ vertices = self.get_vertices()
+ edges = []
+
+ # Useful for importing weird geometries and the sizes are slightly off
+ def nearly_equal(a,b,sig_fig=10):
+ return a==b or (abs(a-b) < 10**(-sig_fig))
+
+ for iter1 in range(len(vertices)):
+ for iter2 in range(iter1+1,len(vertices)):
+ if ((iter1 != iter2) and
+ nearly_equal((vertices[iter1]-vertices[iter2]).norm(),self.size.x) or
+ nearly_equal((vertices[iter1]-vertices[iter2]).norm(),self.size.y) or
+ nearly_equal((vertices[iter1]-vertices[iter2]).norm(),self.size.z)):
+ edges.append([vertices[iter1],vertices[iter2]])
+ return edges
+
+ def pt_in_volume(self,pt):
+ xmin = self.center.x - self.size.x/2
+ xmax = self.center.x + self.size.x/2
+ ymin = self.center.y - self.size.y/2
+ ymax = self.center.y + self.size.y/2
+ zmin = self.center.z - self.size.z/2
+ zmax = self.center.z + self.size.z/2
+
+ if (pt.x >= xmin and pt.x <= xmax and pt.y >= ymin and pt.y <= ymax and pt.z >= zmin and pt.z <= zmax):
+ return True
+ else:
+ return False
class FluxRegion(object):
@@ -2142,65 +2185,17 @@ def get_sfield_r(self):
def get_sfield_p(self):
return self.get_array(mp.Sp, cmplx=True)
- def visualize_chunks(self):
- if self.structure is None:
- self.init_sim()
- try:
- import matplotlib.pyplot as plt
- import matplotlib.cm
- import matplotlib.colors
- from mpl_toolkits.mplot3d import Axes3D
- from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection
- except ImportError:
- warnings.warn("matplotlib is required for visualize_chunks", ImportWarning)
- return
-
- vols = self.structure.get_chunk_volumes()
- owners = self.structure.get_chunk_owners()
-
- fig = plt.figure()
- ax = fig.add_subplot(111, projection='3d')
- chunk_colors = matplotlib.cm.rainbow(np.linspace(0, 1, mp.count_processors()))
-
- def plot_box(box, proc):
- low = mp.Vector3(box.low.x, box.low.y, box.low.z)
- high = mp.Vector3(box.high.x, box.high.y, box.high.z)
- points = [low, high]
-
- x_len = mp.Vector3(high.x) - mp.Vector3(low.x)
- y_len = mp.Vector3(y=high.y) - mp.Vector3(y=low.y)
- xy_len = mp.Vector3(high.x, high.y) - mp.Vector3(low.x, low.y)
-
- points += [low + x_len]
- points += [low + y_len]
- points += [low + xy_len]
- points += [high - x_len]
- points += [high - y_len]
- points += [high - xy_len]
- points = np.array([np.array(v) for v in points])
-
- edges = [
- [points[0], points[2], points[4], points[3]],
- [points[1], points[5], points[7], points[6]],
- [points[0], points[3], points[5], points[7]],
- [points[1], points[4], points[2], points[6]],
- [points[3], points[4], points[1], points[5]],
- [points[0], points[7], points[6], points[2]]
- ]
+ def plot2D(self,**kwargs):
+ return vis.plot2D(self,**kwargs)
+
+ def plot_fields(self,**kwargs):
+ return vis.plot_fields(self,**kwargs)
+
+ def plot3D(self):
+ return vis.plot3D(self)
- faces = Poly3DCollection(edges, linewidths=1, edgecolors='k')
- color_with_alpha = matplotlib.colors.to_rgba(chunk_colors[proc], alpha=0.2)
- faces.set_facecolor(color_with_alpha)
- ax.add_collection3d(faces)
-
- # Plot the points themselves to force the scaling of the axes
- ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0)
-
- if mp.am_master():
- for i, v in enumerate(vols):
- plot_box(mp.gv2box(v.surroundings()), owners[i])
- ax.set_aspect('auto')
- plt.show()
+ def visualize_chunks(self):
+ vis.visualize_chunks(self)
def _create_boundary_region_from_boundary_layers(boundary_layers, gv):
diff --git a/python/tests/visualization.py b/python/tests/visualization.py
new file mode 100644
index 000000000..46b859258
--- /dev/null
+++ b/python/tests/visualization.py
@@ -0,0 +1,128 @@
+
+# visualization.py - Tests the visualization module. Checks 2D
+# plotting of a waveguide with several sources, monitors, and
+# boundary conditions. Checks for subdomain plots.
+#
+# Also tests the animation run function, mp4 output, jshtml output, and git output.
+
+from __future__ import division
+
+import unittest
+import meep as mp
+import numpy as np
+
+# Make sure we have matplotlib installed
+import matplotlib
+matplotlib.use('agg') # Set backend for consistency and to pull pixels quickly
+from matplotlib import pyplot as plt
+import io
+
+def hash_figure(fig):
+ buf = io.BytesIO()
+ fig.savefig(buf, format='raw')
+ buf.seek(0)
+ data = np.frombuffer(buf.getvalue(), dtype=np.uint8)
+ return np.sum((data > np.mean(data)) + data)
+
+def setup_sim(zDim=0):
+ cell = mp.Vector3(16,8,zDim)
+
+ # A simple waveguide
+ geometry = [mp.Block(mp.Vector3(mp.inf,1,1),
+ center=mp.Vector3(),
+ material=mp.Medium(epsilon=12))]
+
+ # A point source
+ sources = [mp.Source(mp.ContinuousSource(frequency=0.15),
+ component=mp.Ez,
+ center=mp.Vector3(-5,0))]
+
+ # A line source
+ sources += [mp.Source(mp.ContinuousSource(frequency=0.15),
+ component=mp.Ez,
+ size=mp.Vector3(0,2,0),
+ center=mp.Vector3(-6,0))]
+
+ # A plane source
+ sources += [mp.Source(mp.ContinuousSource(frequency=0.15),
+ component=mp.Ez,
+ size=mp.Vector3(2,2,0),
+ center=mp.Vector3(-3,0))]
+
+ # Different pml layers
+ pml_layers = [mp.PML(2.0,mp.X),mp.PML(1.0,mp.Y,mp.Low),mp.PML(1.5,mp.Y,mp.High)]
+
+ resolution = 10
+
+ sim = mp.Simulation(cell_size=cell,
+ boundary_layers=pml_layers,
+ geometry=geometry,
+ sources=sources,
+ resolution=resolution)
+ # Line monitor
+ sim.add_flux(1,0,1,mp.FluxRegion(center=mp.Vector3(5,0,0),size=mp.Vector3(0,4), direction=mp.X))
+
+ # Plane monitor
+ sim.add_flux(1,0,1,mp.FluxRegion(center=mp.Vector3(2,0,0),size=mp.Vector3(4,4), direction=mp.X))
+
+ return sim
+
+class TestVisualization(unittest.TestCase):
+
+ def test_plot2D(self):
+
+ # Check plotting of geometry with several sources, monitors, and PMLs
+ f = plt.figure()
+ ax = f.gca()
+ sim = setup_sim()
+ ax = sim.plot2D(ax=ax)
+ if mp.am_master():
+ hash_figure(f)
+ #self.assertAlmostEqual(hash_figure(f),10231488)
+
+ # Check plotting of fields after timestepping
+ f = plt.figure()
+ ax = f.gca()
+ sim.run(until=200)
+ ax = sim.plot2D(ax=ax,fields=mp.Ez)
+ if mp.am_master():
+ hash_figure(f)
+ #self.assertAlmostEqual(hash_figure(f),79786722)
+
+ # Check output_plane feature
+ f = plt.figure()
+ ax = f.gca()
+ vol = mp.Volume(center=mp.Vector3(),size=mp.Vector3(2,2))
+ ax = sim.plot2D(ax=ax,fields=mp.Ez,output_plane=vol)
+ if mp.am_master():
+ hash_figure(f)
+ #self.assertAlmostEqual(hash_figure(f),68926258)
+
+ def test_animation_output(self):
+ # Check without normalization
+ sim = setup_sim()
+ Animate = mp.Animate2D(sim,fields=mp.Ez, realtime=False, normalize=False)
+ sim.run(mp.at_every(1,Animate),until=5)
+
+ # Check with normalization
+ animation = mp.Animate2D(sim,mp.Ez,realtime=False,normalize=True)
+ sim.run(mp.at_every(1),until=25)
+
+ # Check mp4 output
+ Animate.to_mp4(10,'test.mp4')
+
+ # Check gif output
+ Animate.to_gif(10,'test.gif')
+
+ # Check jshtml output
+ Animate.to_jshtml(10)
+ '''
+ Travis does not play well with Mayavi
+ def test_3D_mayavi(self):
+ sim = setup_sim(4)
+ sim.plot3D()
+ '''
+
+if __name__ == '__main__':
+ unittest.main()
+
diff --git a/python/visualization.py b/python/visualization.py
new file mode 100644
index 000000000..eb805ff8f
--- /dev/null
+++ b/python/visualization.py
@@ -0,0 +1,1107 @@
+from collections import namedtuple
+from collections import OrderedDict
+from collections import Sequence
+import warnings
+
+import numpy as np
+
+import meep as mp
+from meep.geom import Vector3, init_do_averaging
+from meep.source import EigenModeSource, check_positive
+
+
+# ------------------------------------------------------- #
+# Visualization
+# ------------------------------------------------------- #
+# Contains all necesarry visualation routines for use with
+# pymeep and pympb.
+
+# ------------------------------------------------------- #
+# Functions used to define the default plotting parameters
+# for the different plotting routines.
+
+default_source_parameters = {
+ 'color':'r',
+ 'edgecolor':'r',
+ 'facecolor':'none',
+ 'hatch':'/',
+ 'linewidth':2
+ }
+
+default_monitor_parameters = {
+ 'color':'b',
+ 'edgecolor':'b',
+ 'facecolor':'none',
+ 'hatch':'/',
+ 'linewidth':2
+ }
+
+default_field_parameters = {
+ 'interpolation':'spline36',
+ 'cmap':'RdBu',
+ 'alpha':0.6
+ }
+
+default_eps_parameters = {
+ 'interpolation':'spline36',
+ 'cmap':'binary',
+ 'alpha':1.0
+ }
+
+default_boundary_parameters = {
+ 'color':'g',
+ 'edgecolor':'g',
+ 'facecolor':'none',
+ 'hatch':'/'
+ }
+
+default_volume_parameters = {
+ 'alpha':1.0,
+ 'color':'k',
+ 'linestyle':'-',
+ 'linewidth':1,
+ 'marker':'.',
+ 'edgecolor':'k',
+ 'facecolor':'none',
+ 'hatch':'/'
+ }
+
+default_label_parameters = {
+ 'label_color':'r',
+ 'offset':20,
+ 'label_alpha':0.3
+}
+
+# ------------------------------------------------------- #
+# Routines to add legends to plot
+
+def place_label(ax,label_text,x,y,centerx,centery,label_parameters=None):
+
+ label_parameters = default_label_parameters if label_parameters is None else dict(default_label_parameters, **label_parameters)
+
+ offset = label_parameters['offset']
+ alpha = label_parameters['label_alpha']
+ color = label_parameters['label_color']
+
+ if x > centerx:
+ xtext = -offset
+ else:
+ xtext = offset
+ if y > centery:
+ ytext = -offset
+ else:
+ ytext = offset
+
+ ax.annotate(label_text, xy=(x, y), xytext=(xtext,ytext),
+ textcoords='offset points', ha='center', va='bottom',
+ bbox=dict(boxstyle='round,pad=0.2', fc=color, alpha=alpha),
+ arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0.5',
+ color=color))
+ return ax
+
+# ------------------------------------------------------- #
+# Helper functions used to plot volumes on a 2D plane
+
+def intersect_plane_line(plane_0,plane_n,line_0,line_1):
+ # Find the intersection point of a plane and line:
+ # http://www.ambrsoft.com/TrigoCalc/Plan3D/PlaneLineIntersection_.htm
+ # plane_0 ........ [Vector3] origin of plane
+ # plane_n ........ [Vector3] normal vector of plane
+ # line_0 ......... [Vector3] first point of line
+ # line_1 ......... [Vector3] second point of line
+ # If the line segment is parallel with the plane, then the first
+ # and last points (line_0 and line_1) are returned in a list.
+
+ # Plane coefficients
+ D = -plane_0.dot(plane_n)
+ A = plane_n.x
+ B = plane_n.y
+ C = plane_n.z
+
+ # Line coefficients
+ v = line_1 - line_0
+ a = v.x
+ b = v.y
+ c = v.z
+ x1 = line_0.x
+ y1 = line_0.y
+ z1 = line_0.z
+
+ den = A*a + B*b + C*c
+
+ # parallel case
+ if den == 0:
+ # coplanar
+ if (line_0-plane_0).dot(plane_n) == 0:
+ return [line_0,line_1]
+ # just parallel
+ else:
+ return None
+
+ pt = Vector3()
+ pt.x = x1 - a*(A*x1 + B*y1 + C*z1 + D) / den
+ pt.y = y1 - b*(A*x1 + B*y1 + C*z1 + D) / den
+ pt.z = z1 - c*(A*x1 + B*y1 + C*z1 + D) / den
+ return pt
+
+def intersect_volume_plane(volume,plane):
+ # returns the vertices that correspond to the polygon,
+ # line, or single point of the intersection of a volume
+ # object and a plane
+ # volume ......... [Volume] volume object
+ # plane .......... [Volume] volume object of the plane
+
+ # Get normal vector of plane
+ if plane.size.x == 0:
+ plane_n = Vector3(x=1)
+ elif plane.size.y == 0:
+ plane_n = Vector3(y=1)
+ elif plane.size.z == 0:
+ plane_n = Vector3(z=1)
+ else:
+ raise ValueError("plane volume must have a nonzero dimension")
+
+ # Get origin of plane
+ plane_0 = plane.center
+
+ intersection_vertices = []
+ edges = volume.get_edges()
+ for ce in edges:
+ pt = intersect_plane_line(plane_0,plane_n,ce[0],ce[1])
+ if isinstance(pt,(list,)):
+ for pt_iter in pt:
+ if (pt_iter is not None) and volume.pt_in_volume(pt_iter):
+ intersection_vertices.append(pt_iter)
+ else:
+ if (pt is not None) and volume.pt_in_volume(pt):
+ intersection_vertices.append(pt)
+ # For point sources, check if point lies on plane
+ if (volume.size.x == volume.size.y == volume.size.z == 0) and plane.pt_in_volume(volume.size):
+ intersection_vertices.append(volume.size)
+ return intersection_vertices
+
+# ------------------------------------------------------- #
+# actual plotting routines
+
+def plot_volume(sim,ax,volume,output_plane=None,plotting_parameters=None,label=None):
+ if not sim._is_initialized:
+ sim.init_sim()
+
+ import matplotlib.patches as patches
+ from matplotlib import pyplot as plt
+ from meep.simulation import Volume
+
+ # Set up the plotting parameters
+ plotting_parameters = default_volume_parameters if plotting_parameters is None else dict(default_volume_parameters, **plotting_parameters)
+
+ # Get domain measurements
+ if output_plane:
+ sim_center, sim_size = (output_plane.center, output_plane.size)
+ elif sim.output_volume:
+ sim_center, sim_size = mp.get_center_and_size(sim.output_volume)
+ else:
+ sim_center, sim_size = (sim.geometry_center, sim.cell_size)
+
+ plane = Volume(center=sim_center,size=sim_size)
+
+ # Pull volume parameters
+ size = volume.size
+ center = volume.center
+
+ xmax = center.x+size.x/2
+ xmin = center.x-size.x/2
+ ymax = center.y+size.y/2
+ ymin = center.y-size.y/2
+ zmax = center.z+size.z/2
+ zmin = center.z-size.z/2
+
+ # Add labels if requested
+ if label is not None and mp.am_master():
+ if sim_size.x == 0:
+ ax = place_label(ax,label,center.y,center.z,sim_center.y,sim_center.z,label_parameters=plotting_parameters)
+ elif sim_size.y == 0:
+ ax = place_label(ax,label,center.x,center.z,sim_center.x,sim_center.z,label_parameters=plotting_parameters)
+ elif sim_size.z == 0:
+ ax = place_label(ax,label,center.x,center.y,sim_center.x,sim_center.y,label_parameters=plotting_parameters)
+
+ # Intersect plane with volume
+ intersection = intersect_volume_plane(volume,plane)
+
+ # Sort the points in a counter clockwise manner to ensure convex polygon is formed
+ def sort_points(xy):
+ xy = np.squeeze(xy)
+ theta = np.arctan2(xy[:,1],xy[:,0])
+ return xy[np.argsort(theta,axis=0)]
+
+ if mp.am_master():
+ # Point volume
+ if len(intersection) == 1:
+ point_args = {key:value for key, value in plotting_parameters.items() if key in ['color','marker','alpha','linewidth']}
+ if sim_center.y == center.y:
+ ax.scatter(center.x,center.z, **point_args)
+ return ax
+ elif sim_center.x == center.x:
+ ax.scatter(center.y,center.z, **point_args)
+ return ax
+ elif sim_center.z == center.z:
+ ax.scatter(center.x,center.y, **point_args)
+ return ax
+ else:
+ return ax
+
+ # Line volume
+ elif len(intersection) == 2:
+ line_args = {key:value for key, value in plotting_parameters.items() if key in ['color','linestyle','linewidth','alpha']}
+ # Plot YZ
+ if sim_size.x == 0:
+ ax.plot([a.y for a in intersection],[a.z for a in intersection], **line_args)
+ return ax
+ #Plot XZ
+ elif sim_size.y == 0:
+ ax.plot([a.x for a in intersection],[a.z for a in intersection], **line_args)
+ return ax
+ # Plot XY
+ elif sim_size.z == 0:
+ ax.plot([a.x for a in intersection],[a.y for a in intersection], **line_args)
+ return ax
+ else:
+ return ax
+
+ # Planar volume
+ elif len(intersection) > 2:
+ planar_args = {key:value for key, value in plotting_parameters.items() if key in ['edgecolor','linewidth','facecolor','hatch','alpha']}
+ # Plot YZ
+ if sim_size.x == 0:
+ ax.add_patch(patches.Polygon(sort_points([[a.y,a.z] for a in intersection]), **planar_args))
+ return ax
+ #Plot XZ
+ elif sim_size.y == 0:
+ ax.add_patch(patches.Polygon(sort_points([[a.x,a.z] for a in intersection]), **planar_args))
+ return ax
+ # Plot XY
+ elif sim_size.z == 0:
+ ax.add_patch(patches.Polygon(sort_points([[a.x,a.y] for a in intersection]), **planar_args))
+ return ax
+ else:
+ return ax
+ else:
+ return ax
+ return ax
+
+def plot_eps(sim,ax,output_plane=None,eps_parameters=None):
+ if sim.structure is None:
+ sim.init_sim()
+
+
+ # consolidate plotting parameters
+ eps_parameters = default_eps_parameters if eps_parameters is None else dict(default_eps_parameters, **eps_parameters)
+
+ # Get domain measurements
+ if output_plane:
+ sim_center, sim_size = (output_plane.center, output_plane.size)
+ elif sim.output_volume:
+ sim_center, sim_size = mp.get_center_and_size(sim.output_volume)
+ else:
+ sim_center, sim_size = (sim.geometry_center, sim.cell_size)
+
+ xmin = sim_center.x - sim_size.x/2
+ xmax = sim_center.x + sim_size.x/2
+ ymin = sim_center.y - sim_size.y/2
+ ymax = sim_center.y + sim_size.y/2
+ zmin = sim_center.z - sim_size.z/2
+ zmax = sim_center.z + sim_size.z/2
+
+ center = Vector3(sim_center.x,sim_center.y,sim_center.z)
+ cell_size = Vector3(sim_size.x,sim_size.y,sim_size.z)
+
+ if sim_size.x == 0:
+ # Plot y on x axis, z on y axis (YZ plane)
+ extent = [ymin,ymax,zmin,zmax]
+ xlabel = 'Y'
+ ylabel = 'Z'
+ elif sim_size.y == 0:
+ # Plot x on x axis, z on y axis (XZ plane)
+ extent = [xmin,xmax,zmin,zmax]
+ xlabel = 'X'
+ ylabel = 'Z'
+ elif sim_size.z == 0:
+ # Plot x on x axis, y on y axis (XY plane)
+ extent = [xmin,xmax,ymin,ymax]
+ xlabel = 'X'
+ ylabel = 'Y'
+ else:
+ raise ValueError("A 2D plane has not been specified...")
+
+ eps_data = np.rot90(np.real(sim.get_array(center=center, size=cell_size, component=mp.Dielectric)))
+ if mp.am_master():
+ ax.imshow(eps_data, extent=extent, **eps_parameters)
+ ax.set_xlabel(xlabel)
+ ax.set_ylabel(ylabel)
+
+ return ax
+
+def plot_boundaries(sim,ax,output_plane=None,boundary_parameters=None):
+ if not sim._is_initialized:
+ sim.init_sim()
+
+ # consolidate plotting parameters
+ boundary_parameters = default_boundary_parameters if boundary_parameters is None else dict(default_boundary_parameters, **boundary_parameters)
+
+ def get_boundary_volumes(thickness,direction,side):
+ from meep.simulation import Volume
+
+ thickness = boundary.thickness
+
+ # Get domain measurements
+ sim_center, sim_size = (sim.geometry_center, sim.cell_size)
+
+ xmin = sim_center.x - sim_size.x/2
+ xmax = sim_center.x + sim_size.x/2
+ ymin = sim_center.y - sim_size.y/2
+ ymax = sim_center.y + sim_size.y/2
+ zmin = sim_center.z - sim_size.z/2
+ zmax = sim_center.z + sim_size.z/2
+
+ cell_x = sim.cell_size.x
+ cell_y = sim.cell_size.y
+ cell_z = sim.cell_size.z
+
+ if direction == mp.X and side == mp.Low:
+ return Volume(center=Vector3(xmin+thickness/2,sim.geometry_center.y,sim.geometry_center.z),
+ size=Vector3(thickness,cell_y,cell_z))
+ elif direction == mp.X and side == mp.High:
+ return Volume(center=Vector3(xmax-thickness/2,sim.geometry_center.y,sim.geometry_center.z),
+ size=Vector3(thickness,cell_y,cell_z))
+ elif direction == mp.Y and side == mp.Low:
+ return Volume(center=Vector3(sim.geometry_center.x,ymin+thickness/2,sim.geometry_center.z),
+ size=Vector3(cell_x,thickness,cell_z))
+ elif direction == mp.Y and side == mp.High:
+ return Volume(center=Vector3(sim.geometry_center.x,ymax-thickness/2,sim.geometry_center.z),
+ size=Vector3(cell_x,thickness,cell_z))
+ elif direction == mp.Z and side == mp.Low:
+ return Volume(center=Vector3(sim.geometry_center.x,sim.geometry_center.y,zmin+thickness/2),
+ size=Vector3(cell_x,cell_y,thickness))
+ elif direction == mp.Z and side == mp.High:
+ return Volume(center=Vector3(sim.geometry_center.x,sim.geometry_center.y,zmax-thickness/2),
+ size=Vector3(cell_x,cell_y,thickness))
+ else:
+ raise ValueError("Invalid boundary type")
+
+ import itertools
+ for boundary in sim.boundary_layers:
+ # All 4 side are the same
+ if boundary.direction == mp.ALL and boundary.side == mp.ALL:
+ if sim.dimensions == 1:
+ dims = [mp.X]
+ elif sim.dimensions == 2:
+ dims = [mp.X,mp.Y]
+ elif sim.dimensions == 3:
+ dims = [mp.X,mp.Y,mp.Z]
+ else:
+ raise ValueError("Invalid simulation dimensions")
+ for permutation in itertools.product(dims, [mp.Low, mp.High]):
+ vol = get_boundary_volumes(boundary.thickness,*permutation)
+ ax = plot_volume(sim,ax,vol,output_plane,plotting_parameters=boundary_parameters)
+ # 2 sides are the same
+ elif boundary.side == mp.ALL:
+ for side in [mp.Low, mp.High]:
+ vol = get_boundary_volumes(boundary.thickness,boundary.direction,side)
+ ax = plot_volume(sim,ax,vol,output_plane,plotting_parameters=boundary_parameters)
+ # only one side
+ else:
+ vol = get_boundary_volumes(boundary.thickness,boundary.direction,boundary.side)
+ ax = plot_volume(sim,ax,vol,output_plane,plotting_parameters=boundary_parameters)
+ return ax
+
+def plot_sources(sim,ax,output_plane=None,labels=False,source_parameters=None):
+ if not sim._is_initialized:
+ sim.init_sim()
+
+ from meep.simulation import Volume
+
+ # consolidate plotting parameters
+ source_parameters = default_source_parameters if source_parameters is None else dict(default_source_parameters, **source_parameters)
+
+ label = 'source' if labels else None
+
+ for src in sim.sources:
+ vol = Volume(center=src.center,size=src.size)
+ ax = plot_volume(sim,ax,vol,output_plane,plotting_parameters=source_parameters,label=label)
+ return ax
+
+def plot_monitors(sim,ax,output_plane=None,labels=False,monitor_parameters=None):
+ if not sim._is_initialized:
+ sim.init_sim()
+
+ from meep.simulation import Volume
+
+ # consolidate plotting parameters
+ monitor_parameters = default_monitor_parameters if monitor_parameters is None else dict(default_monitor_parameters, **monitor_parameters)
+
+ label = 'monitor' if labels else None
+
+ for mon in sim.dft_objects:
+ for reg in mon.regions:
+ vol = Volume(center=reg.center,size=reg.size)
+ ax = plot_volume(sim,ax,vol,output_plane,plotting_parameters=monitor_parameters,label=label)
+ return ax
+
+def plot_fields(sim,ax=None,fields=None,output_plane=None,field_parameters=None):
+ if not sim._is_initialized:
+ sim.init_sim()
+
+ if fields is None:
+ return ax
+
+ # user specifies a field component
+ if fields in [mp.Ex, mp.Ey, mp.Ez, mp.Hx, mp.Hy, mp.Hz]:
+ # Get domain measurements
+ if output_plane:
+ sim_center, sim_size = (output_plane.center, output_plane.size)
+ elif sim.output_volume:
+ sim_center, sim_size = mp.get_center_and_size(sim.output_volume)
+ else:
+ sim_center, sim_size = (sim.geometry_center, sim.cell_size)
+
+ xmin = sim_center.x - sim_size.x/2
+ xmax = sim_center.x + sim_size.x/2
+ ymin = sim_center.y - sim_size.y/2
+ ymax = sim_center.y + sim_size.y/2
+ zmin = sim_center.z - sim_size.z/2
+ zmax = sim_center.z + sim_size.z/2
+
+ center = Vector3(sim_center.x,sim_center.y,sim_center.z)
+ cell_size = Vector3(sim_size.x,sim_size.y,sim_size.z)
+
+ if sim_size.x == 0:
+ # Plot y on x axis, z on y axis (YZ plane)
+ extent = [ymin,ymax,zmin,zmax]
+ xlabel = 'Y'
+ ylabel = 'Z'
+ elif sim_size.y == 0:
+ # Plot x on x axis, z on y axis (XZ plane)
+ extent = [xmin,xmax,zmin,zmax]
+ xlabel = 'X'
+ ylabel = 'Z'
+ elif sim_size.z == 0:
+ # Plot x on x axis, y on y axis (XY plane)
+ extent = [xmin,xmax,ymin,ymax]
+ xlabel = 'X'
+ ylabel = 'Y'
+ fields = sim.get_array(center=center, size=cell_size, component=fields)
+ else:
+ raise ValueError('Please specify a valid field component (mp.Ex, mp.Ey, ...')
+
+ # Either plot the field, or return the array
+ if ax:
+ field_parameters = default_field_parameters if field_parameters is None else dict(default_field_parameters, **field_parameters)
+ if mp.am_master():
+ ax.imshow(np.rot90(fields), extent=extent, **field_parameters)
+ return ax
+ else:
+ return np.rot90(fields)
+ return ax
+
+def plot2D(sim,ax=None, output_plane=None, fields=None, labels=False,
+ eps_parameters=None,boundary_parameters=None,
+ source_parameters=None,monitor_parameters=None,
+ field_parameters=None):
+ if sim.structure is None:
+ sim.init_sim()
+
+ if ax is None and mp.am_master():
+ from matplotlib import pyplot as plt
+ ax = plt.gca()
+
+ # User incorrectly specified a 3D output plane
+ if output_plane and (output_plane.size.x != 0) and (output_plane.size.y != 0) and (output_plane.size.z != 0):
+ raise ValueError("output_plane must be a 2 dimensional volume (a plane).")
+ # User forgot to specify a 2D output plane for a 3D simulation
+ elif output_plane is None and (sim.cell_size.x != 0) and (sim.cell_size.y != 0) and (sim.cell_size.z != 0):
+ raise ValueError("For 3D simulations, you must specify an output_plane.")
+
+ # Plot geometry
+ ax = plot_eps(sim,ax,output_plane=output_plane,eps_parameters=eps_parameters)
+
+ # Plot boundaries
+ ax = plot_boundaries(sim,ax,output_plane=output_plane,boundary_parameters=boundary_parameters)
+
+ # Plot sources
+ ax = plot_sources(sim,ax,output_plane=output_plane,labels=labels,source_parameters=source_parameters)
+
+ # Plot monitors
+ ax = plot_monitors(sim,ax,output_plane=output_plane,labels=labels,monitor_parameters=monitor_parameters)
+
+ # Plot fields
+ ax = plot_fields(sim,ax,fields,output_plane=output_plane,field_parameters=field_parameters)
+
+ return ax
+
+def plot3D(sim):
+ from mayavi import mlab
+
+ if not sim._is_initialized:
+ sim.init_sim()
+
+ if sim.dimensions < 3:
+ raise ValueError("Simulation must have 3 dimensions to visualize 3D")
+
+ eps_data = sim.get_epsilon()
+ s = mlab.contour3d(eps_data, colormap="YlGnBu")
+ return s
+
+def visualize_chunks(sim):
+ if sim.structure is None:
+ sim.init_sim()
+
+ import matplotlib.pyplot as plt
+ import matplotlib.cm
+ import matplotlib.colors
+ from mpl_toolkits.mplot3d import Axes3D
+ from mpl_toolkits.mplot3d.art3d import Poly3DCollection, Line3DCollection
+
+ vols = sim.structure.get_chunk_volumes()
+ owners = sim.structure.get_chunk_owners()
+
+ def plot_box(box, proc, fig, ax):
+ low = mp.Vector3(box.low.x, box.low.y, box.low.z)
+ high = mp.Vector3(box.high.x, box.high.y, box.high.z)
+ points = [low, high]
+
+ x_len = mp.Vector3(high.x) - mp.Vector3(low.x)
+ y_len = mp.Vector3(y=high.y) - mp.Vector3(y=low.y)
+ xy_len = mp.Vector3(high.x, high.y) - mp.Vector3(low.x, low.y)
+
+ points += [low + x_len]
+ points += [low + y_len]
+ points += [low + xy_len]
+ points += [high - x_len]
+ points += [high - y_len]
+ points += [high - xy_len]
+ points = np.array([np.array(v) for v in points])
+
+ edges = [
+ [points[0], points[2], points[4], points[3]],
+ [points[1], points[5], points[7], points[6]],
+ [points[0], points[3], points[5], points[7]],
+ [points[1], points[4], points[2], points[6]],
+ [points[3], points[4], points[1], points[5]],
+ [points[0], points[7], points[6], points[2]]
+ ]
+
+ faces = Poly3DCollection(edges, linewidths=1, edgecolors='k')
+ color_with_alpha = matplotlib.colors.to_rgba(chunk_colors[proc], alpha=0.2)
+ faces.set_facecolor(color_with_alpha)
+ ax.add_collection3d(faces)
+
+ # Plot the points themselves to force the scaling of the axes
+ ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0)
+
+ if mp.am_master():
+ fig = plt.figure()
+ ax = fig.add_subplot(111, projection='3d')
+ chunk_colors = matplotlib.cm.rainbow(np.linspace(0, 1, mp.count_processors()))
+
+ for i, v in enumerate(vols):
+ plot_box(mp.gv2box(v.surroundings()), owners[i], fig, ax)
+ ax.set_aspect('auto')
+ plt.show()
+
+# ------------------------------------------------------- #
+# Animate2D
+# ------------------------------------------------------- #
+# An extensive run function used to visualize the fields
+# of a 2D simulation after every specified time step.
+# ------------------------------------------------------- #
+# Required arguments
+# sim ................. [Simulation object]
+# fields .............. [mp.Ex, mp.Ey, ..., mp. Hz]
+# ------------------------------------------------------- #
+# Optional arguments
+# f ................... [matplotlib figure object]
+# realtime ............ [bool] Update plot in each step
+# normalize ........... [bool] saves fields to normalize
+# after simulation ends.
+# plot_modifiers ...... [list] additional functions to
+# modify plot
+# customization_args .. [dict] other customization args
+# to pass to plot2D()
+
+class Animate2D(object):
+ def __init__(self,sim,fields,f=None,realtime=True,normalize=False,
+ plot_modifiers=None,**customization_args):
+ self.fields = fields
+
+ from matplotlib import pyplot as plt
+ from matplotlib import animation
+
+ if f:
+ self.f = f
+ self.ax = self.f.gca()
+ else:
+ self.f = None
+
+ self.ax = None
+
+ self.realtime = realtime
+ self.normalize = normalize
+ self.plot_modifiers = plot_modifiers
+ self.customization_args = customization_args
+
+
+ self.cumulative_fields = []
+ self._saved_frames = []
+
+ self.frame_format = 'png' # format in which each frame is saved in memory
+ self.codec = 'h264' # encoding of mp4 video
+ self.default_mode = 'loop' # html5 video control mode
+
+ self.init = False
+
+ # Needed for step functions
+ self.__code__ = namedtuple('gna_hack',['co_argcount'])
+ self.__code__.co_argcount=2
+
+ def __call__(self,sim,todo):
+ from matplotlib import pyplot as plt
+
+ if todo == 'step':
+ # Initialize the plot
+ if not self.init:
+ self.ax = sim.plot2D(ax=self.ax,fields=self.fields,**self.customization_args)
+ self.f=plt.gcf()
+ # Run the plot modifier functions
+ '''
+ if self.plot_modifiers:
+ for k in range(len(self.plot_modifiers)):
+ self.ax = self.plot_modifiers[k](self.ax)
+ '''
+ # Store the fields
+ self.w, self.h = self.f.get_size_inches()
+ if mp.am_master():
+ fields = self.ax.images[-1].get_array()
+ self.init = True
+
+ else:
+ # Update the plot
+ fields = sim.plot_fields(fields=self.fields)
+ if mp.am_master():
+ self.ax.images[-1].set_data(fields)
+ self.ax.images[-1].set_clim(vmin=0.8*np.min(fields), vmax=0.8*np.max(fields))
+
+ if self.realtime and mp.am_master():
+ # Redraw the current figure if requested
+ plt.pause(0.05)
+
+ if self.normalize and mp.am_master():
+ # Save fields as a numpy array to be normalized
+ # and saved later.
+ self.cumulative_fields.append(fields)
+ elif mp.am_master():
+ # Capture figure as a png, but store the png in memory
+ # to avoid writing to disk.
+ self.grab_frame()
+ return
+ elif todo == 'finish':
+
+ # Normalize the frames, if requested, and export
+ if self.normalize and mp.am_master():
+ print("Normalizing field data...")
+ fields = np.array(self.cumulative_fields) / np.max(np.abs(self.cumulative_fields),axis=(0,1,2))
+ for k in range(len(self.cumulative_fields)):
+ self.ax.images[-1].set_data(fields[k,:,:])
+ self.ax.images[-1].set_clim(vmin=-0.8, vmax=0.8)
+ self.grab_frame()
+
+ return
+
+ @property
+ def frame_size(self):
+ # A tuple ``(width, height)`` in pixels of a movie frame.
+ # modified from matplotlib library
+ w, h = self.f.get_size_inches()
+ return int(w * self.f.dpi), int(h * self.f.dpi)
+
+ def grab_frame(self):
+ # Saves the figures frame to memory.
+ # modified from matplotlib library
+ from io import BytesIO
+ bin_data = BytesIO()
+ self.f.savefig(bin_data, format=self.frame_format)
+ #imgdata64 = base64.encodebytes(bin_data.getvalue()).decode('ascii')
+ self._saved_frames.append(bin_data.getvalue())
+
+ def _embedded_frames(self, frame_list, frame_format):
+ # converts frame data stored in memory to html5 friendly format
+ # frame_list should be a list of base64-encoded png files
+ # modified from matplotlib
+ import base64
+ template = ' frames[{0}] = "data:image/{1};base64,{2}"\n'
+ return "\n" + "".join(
+ template.format(i, frame_format, base64.encodebytes(frame_data).decode('ascii').replace('\n', '\\\n'))
+ for i, frame_data in enumerate(frame_list))
+
+ def to_jshtml(self,fps):
+ # Exports a javascript enabled html object that is
+ # ready for jupyter notebook embedding.
+ # modified from matplotlib/animation.py code.
+ from uuid import uuid4
+ import sys
+ if sys.version_info[0] < 3:
+ warnings.warn('JSHTML output is not supported with python2 builds.')
+ return ""
+
+ # save the frames to an html file
+ fill_frames = self._embedded_frames(self._saved_frames, self.frame_format)
+ Nframes = len(self._saved_frames)
+ mode_dict = dict(once_checked='',
+ loop_checked='',
+ reflect_checked='')
+ mode_dict[self.default_mode + '_checked'] = 'checked'
+
+ interval = 1000 // fps
+
+ html_string = ""
+ html_string += JS_INCLUDE
+ html_string += STYLE_INCLUDE
+ html_string += DISPLAY_TEMPLATE.format(id=uuid4().hex,
+ Nframes=Nframes,
+ fill_frames=fill_frames,
+ interval=interval,
+ **mode_dict)
+ return JS_Animation(html_string)
+
+ def to_gif(self,fps,filename):
+ # Exports a gif of the recorded animation
+ # requires ffmpeg to be installed
+ # modified from the matplotlib library
+ from subprocess import Popen, PIPE
+ from io import TextIOWrapper, BytesIO
+ FFMPEG_BIN = 'ffmpeg'
+ command = [FFMPEG_BIN,
+ '-f', 'image2pipe', # force piping of rawvideo
+ '-vcodec', self.frame_format, # raw input codec
+ '-s', '%dx%d' % (self.frame_size),
+ '-r', str(fps), # frame rate in frames per second
+ '-i', 'pipe:', # The input comes from a pipe
+ '-vcodec', 'gif', # output gif format
+ '-r', str(fps), # frame rate in frames per second
+ '-y',
+ '-vf', 'pad=width=ceil(iw/2)*2:height=ceil(ih/2)*2',
+ '-an', filename # output filename
+ ]
+
+ print("Generating GIF...")
+ if mp.am_master():
+ proc = Popen(command, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+ for i in range(len(self._saved_frames)):
+ proc.stdin.write(self._saved_frames[i])
+ out, err = proc.communicate() # pipe in images
+ proc.stdin.close()
+ proc.wait()
+ return
+
+ def to_mp4(self,fps,filename):
+ # Exports an mp4 of the recorded animation
+ # requires ffmpeg to be installed
+ # modified from the matplotlib library
+ from subprocess import Popen, PIPE
+ from io import TextIOWrapper, BytesIO
+ FFMPEG_BIN = 'ffmpeg'
+ command = [FFMPEG_BIN,
+ '-f', 'image2pipe', # force piping of rawvideo
+ '-vcodec', self.frame_format, # raw input codec
+ '-s', '%dx%d' % (self.frame_size),
+ #'-pix_fmt', self.frame_format,
+ '-r', str(fps), # frame rate in frames per second
+ '-i', 'pipe:', # The input comes from a pipe
+ '-vcodec', self.codec, # output mp4 format
+ '-pix_fmt','yuv420p',
+ '-r', str(fps), # frame rate in frames per second
+ '-y',
+ '-vf', 'pad=width=ceil(iw/2)*2:height=ceil(ih/2)*2',
+ '-an', filename # output filename
+ ]
+ print("Generating MP4...")
+ if mp.am_master():
+ proc = Popen(command, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+ for i in range(len(self._saved_frames)):
+ proc.stdin.write(self._saved_frames[i])
+ out, err = proc.communicate() # pipe in images
+ proc.stdin.close()
+ proc.wait()
+ return
+
+ def reset(self):
+ self.cumulative_fields = []
+ self.ax = None
+ self.f = None
+
+ def set_figure(self,f):
+ self.f = f
+
+# ------------------------------------------------------- #
+# JS_Animation
+# ------------------------------------------------------- #
+# A helper class used to make jshtml animations embed
+# seamlessly within Jupyter notebooks.
+
+class JS_Animation():
+ def __init__(self,jshtml):
+ self.jshtml = jshtml
+ def _repr_html_(self):
+ return self.jshtml
+ def get_jshtml(self):
+ return self.jshtml
+
+
+# ------------------------------------------------------- #
+# JSHTML Templates
+# ------------------------------------------------------- #
+# Template functions grabbed from matplotlib/_animation_data.py
+# https://github.com/matplotlib/matplotlib/blob/aac57266803ff09c8c9dc10e1bb2bd389d79e42d/lib/matplotlib/_animation_data.py
+# Included here since python2 only supports matplotlib <3.0
+
+
+# Javascript template for HTMLWriter
+JS_INCLUDE = """
+
+
+"""
+
+
+# Style definitions for the HTML template
+STYLE_INCLUDE = """
+
+"""
+
+
+# HTML template for HTMLWriter
+DISPLAY_TEMPLATE = """
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+"""
+
+
+INCLUDED_FRAMES = """
+ for (var i=0; i<{Nframes}; i++){{
+ frames[i] = "{frame_dir}/frame" + ("0000000" + i).slice(-7) +
+ ".{frame_format}";
+ }}
+"""
\ No newline at end of file