From 79351423317189c937ddd4491f3701f89205192c Mon Sep 17 00:00:00 2001 From: Thomas Date: Fri, 29 Jul 2022 14:26:06 -0700 Subject: [PATCH 1/2] add type hinting to visualization --- python/visualization.py | 290 +++++++++++++++++++++++----------------- 1 file changed, 170 insertions(+), 120 deletions(-) diff --git a/python/visualization.py b/python/visualization.py index 4f4879e8e..1fe46f0a6 100644 --- a/python/visualization.py +++ b/python/visualization.py @@ -1,16 +1,21 @@ -import warnings from collections import namedtuple +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 +from meep.simulation import Simulation, Volume -import meep as mp +## Typing imports +from matplotlib.axes import Axes +from typing import Callable, Union, Any, Iterable # ------------------------------------------------------- # # Visualization # ------------------------------------------------------- # -# Contains all necesarry visualation routines for use with +# Contains all necessary visualization routines for use with # pymeep and pympb. # ------------------------------------------------------- # @@ -74,7 +79,7 @@ # don't correspond to the keyword arguments of a particular # function (func_with_kwargs.) # Adapted from https://stackoverflow.com/questions/26515595/how-does-one-ignore-unexpected-keyword-arguments-passed-to-a-function/44052550 -def filter_dict(dict_to_filter, func_with_kwargs): +def filter_dict(dict_to_filter: dict, func_with_kwargs: Callable) -> dict: import inspect filter_keys = [] @@ -86,18 +91,21 @@ def filter_dict(dict_to_filter, func_with_kwargs): # Python2 ... filter_keys = inspect.getargspec(func_with_kwargs)[0] - return { + filtered_dict = { filter_key: dict_to_filter[filter_key] for filter_key in filter_keys if filter_key in dict_to_filter } + return filtered_dict # ------------------------------------------------------- # # Routines to add legends to plot -def place_label(ax, label_text, x, y, centerx, centery, label_parameters=None): +def place_label( + ax: Axes, label_text: str, x, y, centerx, centery, label_parameters: dict = None +) -> Axes: if label_parameters is None: label_parameters = default_label_parameters @@ -108,8 +116,15 @@ def place_label(ax, label_text, x, y, centerx, centery, label_parameters=None): alpha = label_parameters["label_alpha"] color = label_parameters["label_color"] - xtext = -offset if x > centerx else offset - ytext = -offset if y > centery else offset + if x > centerx: + xtext = -offset + else: + xtext = offset + if y > centery: + ytext = -offset + else: + ytext = offset + ax.annotate( label_text, xy=(x, y), @@ -129,7 +144,7 @@ def place_label(ax, label_text, x, y, centerx, centery, label_parameters=None): # Returns the intersection points of two Volumes. # Volumes must be a line, plane, or rectangular prism # (since they are volume objects) -def intersect_volume_volume(volume1, volume2): +def intersect_volume_volume(volume1: Volume, volume2: Volume) -> list: # volume1 ............... [volume] # volume2 ............... [volume] @@ -161,11 +176,11 @@ def intersect_volume_volume(volume1, volume2): L = np.max([L1, L2], axis=0) # For single points we have to check manually - if np.all(U - L == 0) and ( - (not volume1.pt_in_volume(Vector3(*U))) - or (not volume2.pt_in_volume(Vector3(*U))) - ): - return [] + if np.all(U - L == 0): + if (not volume1.pt_in_volume(Vector3(*U))) or ( + not volume2.pt_in_volume(Vector3(*U)) + ): + return [] # Check for two volumes that don't intersect if np.any(U - L < 0): @@ -175,7 +190,9 @@ def intersect_volume_volume(volume1, volume2): vertices = [] for x_vals in [L[0], U[0]]: for y_vals in [L[1], U[1]]: - vertices.extend(Vector3(x_vals, y_vals, z_vals) for z_vals in [L[2], U[2]]) + for z_vals in [L[2], U[2]]: + vertices.append(Vector3(x_vals, y_vals, z_vals)) + # Remove any duplicate points caused by coplanar lines vertices = [ vertices[i] for i, x in enumerate(vertices) if x not in vertices[i + 1 :] @@ -193,28 +210,27 @@ def intersect_volume_volume(volume1, volume2): # Not only do we need to check for all of these possibilities, but we also need # to check if the user accidentally specifies a plane that stretches beyond the # simulation domain. -def get_2D_dimensions(sim, output_plane): - from meep.simulation import Volume - +def get_2D_dimensions(sim: Simulation, output_plane: Volume) -> (Vector3, Vector3): # Pull correct plane from user if output_plane: plane_center, plane_size = (output_plane.center, output_plane.size) elif sim.output_volume: plane_center, plane_size = mp.get_center_and_size(sim.output_volume) - elif (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical: - plane_center, plane_size = ( - sim.geometry_center + mp.Vector3(sim.cell_size.x / 2), - sim.cell_size, - ) else: - plane_center, plane_size = (sim.geometry_center, sim.cell_size) + if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical: + plane_center, plane_size = ( + sim.geometry_center + Vector3(sim.cell_size.x / 2), + sim.cell_size, + ) + else: + plane_center, plane_size = (sim.geometry_center, sim.cell_size) plane_volume = Volume(center=plane_center, size=plane_size) if plane_size.x != 0 and plane_size.y != 0 and plane_size.z != 0: raise ValueError("Plane volume must be 2D (a plane).") if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical: - center = sim.geometry_center + mp.Vector3(sim.cell_size.x / 2) - check_volume = mp.Volume(center=center, size=sim.cell_size) + center = sim.geometry_center + Vector3(sim.cell_size.x / 2) + check_volume = Volume(center=center, size=sim.cell_size) else: check_volume = Volume(center=sim.geometry_center, size=sim.cell_size) vertices = intersect_volume_volume(check_volume, plane_volume) @@ -237,7 +253,9 @@ def get_2D_dimensions(sim, output_plane): return sim_center, sim_size -def box_vertices(box_center, box_size, is_cylindrical=False): +def box_vertices( + box_center: Vector3, box_size: Vector3, is_cylindrical: bool = False +) -> (float, float, float, float, float, float): # in cylindrical coordinates, radial (R) axis # is in the range (0,R) rather than (-R/2,+R/2) # as in Cartesian coordinates. @@ -257,14 +275,16 @@ def box_vertices(box_center, box_size, is_cylindrical=False): # ------------------------------------------------------- # # actual plotting routines - - def plot_volume( - sim, ax, volume, output_plane=None, plotting_parameters=None, label=None -): + sim: Simulation, + ax: Axes, + volume: Volume, + output_plane: Volume = None, + plotting_parameters: dict = None, + label: str = None, +) -> Axes: import matplotlib.patches as patches from matplotlib import pyplot as plt - from meep.simulation import Volume # Set up the plotting parameters if plotting_parameters is None: @@ -357,7 +377,7 @@ def sort_points(xy): ax.plot( [a.y for a in intersection], [a.z for a in intersection], - **line_args, + **line_args ) return ax # Plot XZ @@ -365,7 +385,7 @@ def sort_points(xy): ax.plot( [a.x for a in intersection], [a.z for a in intersection], - **line_args, + **line_args ) return ax # Plot XY @@ -373,7 +393,7 @@ def sort_points(xy): ax.plot( [a.x for a in intersection], [a.y for a in intersection], - **line_args, + **line_args ) return ax else: @@ -417,7 +437,13 @@ def sort_points(xy): return ax -def plot_eps(sim, ax, output_plane=None, eps_parameters=None, frequency=None): +def plot_eps( + sim: Simulation, + ax: Axes, + output_plane: Volume = None, + eps_parameters: dict = None, + frequency: float = None, +) -> Axes: # consolidate plotting parameters if eps_parameters is None: eps_parameters = default_eps_parameters @@ -454,7 +480,11 @@ def plot_eps(sim, ax, output_plane=None, eps_parameters=None, frequency=None): sim_center, sim_size, sim.is_cylindrical ) - grid_resolution = eps_parameters["resolution"] or sim.resolution + if eps_parameters["resolution"]: + grid_resolution = eps_parameters["resolution"] + else: + grid_resolution = sim.resolution + Nx = int((xmax - xmin) * grid_resolution + 1) Ny = int((ymax - ymin) * grid_resolution + 1) Nz = int((zmax - zmin) * grid_resolution + 1) @@ -470,10 +500,10 @@ def plot_eps(sim, ax, output_plane=None, eps_parameters=None, frequency=None): elif sim_size.y == 0: # Plot x on x axis, z on y axis (XZ plane) extent = [xmin, xmax, zmin, zmax] - xlabel = ( - "R" if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical else "X" - ) - + if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical: + xlabel = "R" + else: + xlabel = "X" ylabel = "Z" xtics = np.linspace(xmin, xmax, Nx) ytics = np.array([sim_center.y]) @@ -512,16 +542,19 @@ def plot_eps(sim, ax, output_plane=None, eps_parameters=None, frequency=None): return ax -def plot_boundaries(sim, ax, output_plane=None, boundary_parameters=None): +def plot_boundaries( + sim: Simulation, + ax: Axes, + output_plane: Volume = None, + boundary_parameters: dict = None, +) -> Axes: # consolidate plotting parameters if boundary_parameters is None: boundary_parameters = default_boundary_parameters else: boundary_parameters = dict(default_boundary_parameters, **boundary_parameters) - def get_boundary_volumes(thickness, direction, side): - from meep.simulation import Volume - + def get_boundary_volumes(thickness: float, direction: float, side) -> Volume: thickness = boundary.thickness xmin, xmax, ymin, ymax, zmin, zmax = box_vertices( @@ -625,8 +658,13 @@ def get_boundary_volumes(thickness, direction, side): return ax -def plot_sources(sim, ax, output_plane=None, labels=False, source_parameters=None): - from meep.simulation import Volume +def plot_sources( + sim: Simulation, + ax: Axes, + output_plane: Volume = None, + labels: bool = False, + source_parameters: dict = None, +): # consolidate plotting parameters if source_parameters is None: @@ -649,14 +687,18 @@ def plot_sources(sim, ax, output_plane=None, labels=False, source_parameters=Non return ax -def plot_monitors(sim, ax, output_plane=None, labels=False, monitor_parameters=None): - from meep.simulation import Volume - +def plot_monitors( + sim: Simulation, + ax: Axes, + output_plane: Volume = None, + labels: bool = False, + monitor_parameters: dict = None, +) -> Axes: # consolidate plotting parameters if monitor_parameters is None: monitor_parameters = default_monitor_parameters else: - monitor_parametesr = dict(default_monitor_parameters, **monitor_parameters) + monitor_parameters = dict(default_monitor_parameters, **monitor_parameters) label = "monitor" if labels else None @@ -674,7 +716,13 @@ def plot_monitors(sim, ax, output_plane=None, labels=False, monitor_parameters=N return ax -def plot_fields(sim, ax=None, fields=None, output_plane=None, field_parameters=None): +def plot_fields( + sim: Simulation, + ax: Axes = None, + fields=None, + output_plane: Volume = None, + field_parameters: dict = None, +) -> Union[Axes, Any]: if not sim._is_initialized: sim.init_sim() @@ -686,7 +734,8 @@ def plot_fields(sim, ax=None, fields=None, output_plane=None, field_parameters=N else: field_parameters = dict(default_field_parameters, **field_parameters) - if fields not in [ + # user specifies a field component + if fields in [ mp.Ex, mp.Ey, mp.Ez, @@ -699,64 +748,68 @@ def plot_fields(sim, ax=None, fields=None, output_plane=None, field_parameters=N mp.Hy, mp.Hz, ]: - raise ValueError("Please specify a valid field component (mp.Ex, mp.Ey, ...") - - # Get domain measurements - sim_center, sim_size = get_2D_dimensions(sim, output_plane) + # Get domain measurements + sim_center, sim_size = get_2D_dimensions(sim, output_plane) - xmin, xmax, ymin, ymax, zmin, zmax = box_vertices( - sim_center, sim_size, sim.is_cylindrical - ) - - 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 = ( - "R" if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical else "X" + xmin, xmax, ymin, ymax, zmin, zmax = box_vertices( + sim_center, sim_size, sim.is_cylindrical ) - 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=sim_center, size=sim_size, component=fields) + 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] + if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical: + xlabel = "R" + else: + 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=sim_center, size=sim_size, component=fields) + else: + raise ValueError("Please specify a valid field component (mp.Ex, mp.Ey, ...") + fields = field_parameters["post_process"](fields) if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical: fields = np.flipud(fields) else: fields = np.rot90(fields) - if not ax: + # Either plot the field, or return the array + if ax: + if mp.am_master(): + ax.imshow(fields, extent=extent, **filter_dict(field_parameters, ax.imshow)) + return ax + else: return fields - if mp.am_master(): - ax.imshow(fields, extent=extent, **filter_dict(field_parameters, ax.imshow)) return ax def plot2D( - sim, - ax=None, - output_plane=None, + sim: Simulation, + ax: Axes = None, + output_plane: Volume = None, fields=None, - labels=False, - eps_parameters=None, - boundary_parameters=None, - source_parameters=None, - monitor_parameters=None, - field_parameters=None, - frequency=None, - plot_eps_flag=True, - plot_sources_flag=True, - plot_monitors_flag=True, - plot_boundaries_flag=True, -): + labels: bool = False, + eps_parameters: dict = None, + boundary_parameters: dict = None, + source_parameters: dict = None, + monitor_parameters: dict = None, + field_parameters: dict = None, + frequency: float = None, + plot_eps_flag: bool = True, + plot_sources_flag: bool = True, + plot_monitors_flag: bool = True, + plot_boundaries_flag: bool = True, +) -> Axes: # Ensure a figure axis exists if ax is None and mp.am_master(): @@ -765,8 +818,6 @@ def plot2D( ax = plt.gca() # validate the output plane to ensure proper 2D coordinates - from meep.simulation import Volume - sim_center, sim_size = get_2D_dimensions(sim, output_plane) output_plane = Volume(center=sim_center, size=sim_size) @@ -819,7 +870,7 @@ def plot2D( return ax -def plot3D(sim): +def plot3D(sim: Simulation): from mayavi import mlab if sim.dimensions < 3: @@ -838,16 +889,17 @@ def plot3D(sim): ztics = np.linspace(zmin, zmax, Nz) eps_data = sim.get_epsilon_grid(xtics, ytics, ztics) - return mlab.contour3d(eps_data, colormap="YlGnBu") + s = mlab.contour3d(eps_data, colormap="YlGnBu") + return s -def visualize_chunks(sim): +def visualize_chunks(sim: Simulation): if sim.structure is None: sim.init_sim() + import matplotlib.pyplot as plt import matplotlib.cm import matplotlib.colors - import matplotlib.pyplot as plt if sim.structure.gv.dim == 2: from mpl_toolkits.mplot3d import Axes3D @@ -860,13 +912,13 @@ def visualize_chunks(sim): def plot_box(box, proc, fig, ax): if sim.structure.gv.dim == 2: - low = mp.Vector3(box.low.x, box.low.y, box.low.z) - high = mp.Vector3(box.high.x, box.high.y, box.high.z) + low = Vector3(box.low.x, box.low.y, box.low.z) + high = 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) + x_len = Vector3(high.x) - Vector3(low.x) + y_len = Vector3(y=high.y) - Vector3(y=low.y) + xy_len = Vector3(high.x, high.y) - Vector3(low.x, low.y) points += [low + x_len] points += [low + y_len] @@ -893,12 +945,12 @@ def plot_box(box, proc, fig, ax): # Plot the points themselves to force the scaling of the axes ax.scatter(points[:, 0], points[:, 1], points[:, 2], s=0) else: - low = mp.Vector3(box.low.x, box.low.y) - high = mp.Vector3(box.high.x, box.high.y) + low = Vector3(box.low.x, box.low.y) + high = Vector3(box.high.x, box.high.y) 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) + x_len = Vector3(high.x) - Vector3(low.x) + y_len = Vector3(y=high.y) - Vector3(y=low.y) points += [low + x_len] points += [low + y_len] @@ -1003,7 +1055,7 @@ def __init__( realtime=False, normalize=False, plot_modifiers=None, - **customization_args, + **customization_args ): """ Construct an `Animate2D` object. @@ -1175,7 +1227,6 @@ def to_jshtml(self, fps): # Only works with Python3 and matplotlib > 3.1.0 from distutils.version import LooseVersion - import matplotlib if LooseVersion(matplotlib.__version__) < LooseVersion("3.1.0"): @@ -1187,7 +1238,6 @@ def to_jshtml(self, fps): return if mp.am_master(): from uuid import uuid4 - from matplotlib._animation_data import ( DISPLAY_TEMPLATE, INCLUDED_FRAMES, @@ -1199,7 +1249,7 @@ def to_jshtml(self, fps): 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[f"{self.default_mode}_checked"] = "checked" + mode_dict[self.default_mode + "_checked"] = "checked" interval = 1000 // fps @@ -1211,7 +1261,7 @@ def to_jshtml(self, fps): Nframes=Nframes, fill_frames=fill_frames, interval=interval, - **mode_dict, + **mode_dict ) return JS_Animation(html_string) @@ -1227,8 +1277,8 @@ def to_gif(self, fps, filename): # requires ffmpeg to be installed # modified from the matplotlib library if mp.am_master(): - from io import BytesIO, TextIOWrapper - from subprocess import PIPE, Popen + from subprocess import Popen, PIPE + from io import TextIOWrapper, BytesIO FFMPEG_BIN = "ffmpeg" command = [ @@ -1273,8 +1323,8 @@ def to_mp4(self, fps, filename): # requires ffmpeg to be installed # modified from the matplotlib library if mp.am_master(): - from io import BytesIO, TextIOWrapper - from subprocess import PIPE, Popen + from subprocess import Popen, PIPE + from io import TextIOWrapper, BytesIO FFMPEG_BIN = "ffmpeg" command = [ From 5f0ad2edce665aab367d1906b2d22c8e2e075774 Mon Sep 17 00:00:00 2001 From: Thomas Date: Fri, 29 Jul 2022 14:29:02 -0700 Subject: [PATCH 2/2] import visualization on-the-fly in Simulation --- python/simulation.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/simulation.py b/python/simulation.py index 0a4453c0a..6a58cd601 100644 --- a/python/simulation.py +++ b/python/simulation.py @@ -15,7 +15,6 @@ except ImportError: from collections.abc import Sequence -import meep.visualization as vis import numpy as np from meep.geom import GeometricObject, Medium, Vector3, init_do_averaging from meep.source import ( @@ -4680,6 +4679,8 @@ def plot2D( - `post_process=np.real`: post processing function to apply to fields (must be a function object) """ + import meep.visualization as vis + return vis.plot2D( self, ax=ax, @@ -4700,6 +4701,8 @@ def plot2D( ) def plot_fields(self, **kwargs): + import meep.visualization as vis + return vis.plot_fields(self, **kwargs) def plot3D(self): @@ -4707,6 +4710,8 @@ def plot3D(self): Uses Mayavi to render a 3D simulation domain. The simulation object must be 3D. Can also be embedded in Jupyter notebooks. """ + import meep.visualization as vis + return vis.plot3D(self) def visualize_chunks(self): @@ -4715,6 +4720,8 @@ def visualize_chunks(self): rectangular region is a chunk, and each color represents a different processor. Requires [matplotlib](https://matplotlib.org). """ + import meep.visualization as vis + vis.visualize_chunks(self)