from collections import namedtuple import warnings from time import sleep import matplotlib.pyplot as plt 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 ## Typing imports from matplotlib.axes import Axes from matplotlib.figure import Figure from typing import Callable, Union, Any, Tuple, List, Optional # ------------------------------------------------------- # # Visualization # ------------------------------------------------------- # # Contains all necessary visualization 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.8, "post_process": np.real, "colorbar": False, } default_eps_parameters = { "interpolation": "spline36", "cmap": "binary", "alpha": 1.0, "contour": False, "contour_linewidth": 1, "frequency": None, "resolution": None, "colorbar": False, } default_colorbar_parameters = { "label": None, "orientation": "vertical", "extend": None, "position": "right", "size": "5%", "pad": "2%", } 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} # Used to remove the elements of a dictionary (dict_to_filter) that # 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: dict, func_with_kwargs: Callable) -> dict: import inspect filter_keys = [] try: # Python3 ... sig = inspect.signature(func_with_kwargs) filter_keys = [param.name for param in sig.parameters.values()] except: # Python2 ... filter_keys = inspect.getargspec(func_with_kwargs)[0] 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: Axes, label_text: str, x: float, y: float, centerx: float, centery: float, label_parameters: Optional[dict] = None, ) -> Axes: if label_parameters is None: label_parameters = default_label_parameters else: label_parameters = 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 # 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: Volume, volume2: Volume) -> List[Vector3]: # volume1 ............... [volume] # volume2 ............... [volume] # Represent the volumes by an "upper" and "lower" coordinate U1 = [ volume1.center.x + volume1.size.x / 2, volume1.center.y + volume1.size.y / 2, volume1.center.z + volume1.size.z / 2, ] L1 = [ volume1.center.x - volume1.size.x / 2, volume1.center.y - volume1.size.y / 2, volume1.center.z - volume1.size.z / 2, ] U2 = [ volume2.center.x + volume2.size.x / 2, volume2.center.y + volume2.size.y / 2, volume2.center.z + volume2.size.z / 2, ] L2 = [ volume2.center.x - volume2.size.x / 2, volume2.center.y - volume2.size.y / 2, volume2.center.z - volume2.size.z / 2, ] # Evaluate intersection U = np.min([U1, U2], axis=0) L = np.max([L1, L2], axis=0) # For single points we have to check manually 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): return [] # Pull all possible vertices vertices = [] for x_vals in [L[0], U[0]]: for y_vals in [L[1], U[1]]: 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 :] ] return vertices # All of the 2D plotting routines need an output plane over which to plot. # The user has many options to specify this output plane. They can pass # the output_plane parameter, which is a 2D volume object. They can specify # a volume using in_volume, which stores the volume as a C volume, not a Python # volume. They can also do nothing and plot the XY plane through Z=0. # # 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: Simulation, output_plane: Volume) -> Tuple[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) else: 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 + 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) if len(vertices) == 0: raise ValueError( "The specified user volume is completely outside of the simulation domain." ) intersection_vol = Volume(vertices=vertices) if (intersection_vol.size != plane_volume.size) or ( intersection_vol.center != plane_volume.center ): warnings.warn( "The specified user volume is larger than the simulation domain and has been truncated." ) sim_center, sim_size = (intersection_vol.center, intersection_vol.size) return sim_center, sim_size def box_vertices( box_center: Vector3, box_size: Vector3, is_cylindrical: bool = False ) -> Tuple[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. if is_cylindrical: xmin = 0 xmax = box_size.x else: xmin = box_center.x - 0.5 * box_size.x xmax = box_center.x + 0.5 * box_size.x ymin = box_center.y - 0.5 * box_size.y ymax = box_center.y + 0.5 * box_size.y zmin = box_center.z - 0.5 * box_size.z zmax = box_center.z + 0.5 * box_size.z return xmin, xmax, ymin, ymax, zmin, zmax # ------------------------------------------------------- # # actual plotting routines def plot_volume( sim: Simulation, ax: Axes, volume: Volume, output_plane: Optional[Volume] = None, plotting_parameters: Optional[dict] = None, label: Optional[str] = None, ) -> Axes: import matplotlib.patches as patches from matplotlib import pyplot as plt # Set up the plotting parameters if plotting_parameters is None: plotting_parameters = default_volume_parameters else: plotting_parameters = dict(default_volume_parameters, **plotting_parameters) # Get domain measurements sim_center, sim_size = get_2D_dimensions(sim, output_plane) plane = Volume(center=sim_center, size=sim_size) size = volume.size center = volume.center xmin, xmax, ymin, ymax, zmin, zmax = box_vertices(center, size, sim.is_cylindrical) # 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_volume(volume, plane) # Sort the points in a counter clockwise manner to ensure convex polygon is formed def sort_points(xy): xy = np.squeeze(xy) xy_mean = np.mean(xy, axis=0) theta = np.arctan2(xy[:, 1] - xy_mean[1], xy[:, 0] - xy_mean[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_size.y == 0: ax.scatter(center.x, center.z, **point_args) return ax elif sim_size.x == 0: ax.scatter(center.y, center.z, **point_args) return ax elif sim_size.z == 0: 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 _add_colorbar( ax: Axes, cmap: str, vmin: float, vmax: float, default_label: Optional[str] = None, colorbar_parameters: Optional[dict] = None, ) -> None: """Add a colorbar to the parent Figure of 'ax' by creating an additional Axes.""" import matplotlib as mpl from mpl_toolkits.axes_grid1 import make_axes_locatable if colorbar_parameters is None: colorbar_parameters = default_colorbar_parameters else: colorbar_parameters = dict(default_colorbar_parameters, **colorbar_parameters) # Use default label (specified by plot_eps or plot_fields) if no user-specified label if colorbar_parameters["label"] is None: colorbar_parameters["label"] = default_label # Create a map between field/eps values and colors in the colormap. # Note: cm.get_cmap() is deprecated for matplotlib>=3.6, use mpl.colormaps[cmap] instead if necessary. sm = mpl.cm.ScalarMappable( norm=mpl.colors.Normalize(vmin, vmax), cmap=mpl.cm.get_cmap(cmap), ) # Pop specific values out of colorbar params so user can add any kwargs to plt.colorbar cax = make_axes_locatable(ax).append_axes( pad=colorbar_parameters.pop("pad"), size=colorbar_parameters.pop("size"), position=colorbar_parameters.pop("position"), ) plt.colorbar(mappable=sm, cax=cax, **colorbar_parameters) def plot_eps( sim: Simulation, ax: Optional[Axes] = None, output_plane: Optional[Volume] = None, eps_parameters: Optional[dict] = None, colorbar_parameters: Optional[dict] = None, frequency: Optional[float] = None, ) -> Union[Axes, Any]: # consolidate plotting parameters if eps_parameters is None: eps_parameters = default_eps_parameters else: eps_parameters = dict(default_eps_parameters, **eps_parameters) # Determine a frequency to plot all epsilon if frequency is not None: warnings.warn( "The frequency parameter of plot2D has been deprecated. " "Use the frequency key of the eps_parameters dictionary instead." ) eps_parameters["frequency"] = frequency if eps_parameters["frequency"] is None: try: # Continuous sources eps_parameters["frequency"] = sim.sources[0].frequency except: try: # Gaussian sources eps_parameters["frequency"] = sim.sources[0].src.frequency except: try: # Custom sources eps_parameters["frequency"] = sim.sources[0].src.center_frequency except: # No sources eps_parameters["frequency"] = 0 # 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 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) 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" xtics = np.array([sim_center.x]) ytics = np.linspace(ymin, ymax, Ny) ztics = np.linspace(zmin, zmax, Nz) 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" xtics = np.linspace(xmin, xmax, Nx) ytics = np.array([sim_center.y]) ztics = np.linspace(zmin, zmax, Nz) 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" xtics = np.linspace(xmin, xmax, Nx) ytics = np.linspace(ymin, ymax, Ny) ztics = np.array([sim_center.z]) else: raise ValueError("A 2D plane has not been specified...") eps_data = np.rot90( np.real(sim.get_epsilon_grid(xtics, ytics, ztics, eps_parameters["frequency"])) ) if mp.am_master(): # If Axes was not provided, just return the eps_data, otherwise plot if not ax: return eps_data if eps_parameters["contour"]: ax.contour( eps_data, 0, levels=np.unique(eps_data), colors="black", origin="upper", extent=extent, linewidths=eps_parameters["contour_linewidth"], ) else: ax.imshow(eps_data, extent=extent, **filter_dict(eps_parameters, ax.imshow)) if eps_parameters["colorbar"]: _add_colorbar( ax=ax, cmap=eps_parameters["cmap"], vmin=np.amin(eps_data), vmax=np.amax(eps_data), default_label=r"$\epsilon_r$", colorbar_parameters=colorbar_parameters, ) ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) return ax def plot_boundaries( sim: Simulation, ax: Axes, output_plane: Optional[Volume] = None, boundary_parameters: Optional[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: float, direction: float, side) -> Volume: thickness = boundary.thickness xmin, xmax, ymin, ymax, zmin, zmax = box_vertices( sim.geometry_center, sim.cell_size, sim.is_cylindrical ) if direction == mp.X and side == mp.Low: return Volume( center=Vector3( xmin + 0.5 * thickness, sim.geometry_center.y, sim.geometry_center.z ), size=Vector3(thickness, sim.cell_size.y, sim.cell_size.z), ) elif (direction == mp.X and side == mp.High) or direction == mp.R: return Volume( center=Vector3( xmax - 0.5 * thickness, sim.geometry_center.y, sim.geometry_center.z ), size=Vector3(thickness, sim.cell_size.y, sim.cell_size.z), ) elif direction == mp.Y and side == mp.Low: return Volume( center=Vector3( sim.geometry_center.x, ymin + 0.5 * thickness, sim.geometry_center.z ), size=Vector3(sim.cell_size.x, thickness, sim.cell_size.z), ) elif direction == mp.Y and side == mp.High: return Volume( center=Vector3( sim.geometry_center.x, ymax - 0.5 * thickness, sim.geometry_center.z ), size=Vector3(sim.cell_size.x, thickness, sim.cell_size.z), ) elif direction == mp.Z and side == mp.Low: xcen = sim.geometry_center.x if sim.is_cylindrical: xcen += 0.5 * sim.cell_size.x return Volume( center=Vector3(xcen, sim.geometry_center.y, zmin + 0.5 * thickness), size=Vector3(sim.cell_size.x, sim.cell_size.y, thickness), ) elif direction == mp.Z and side == mp.High: xcen = sim.geometry_center.x if sim.is_cylindrical: xcen += 0.5 * sim.cell_size.x return Volume( center=Vector3(xcen, sim.geometry_center.y, zmax - 0.5 * thickness), size=Vector3(sim.cell_size.x, sim.cell_size.y, thickness), ) else: raise ValueError("Invalid boundary type") import itertools for boundary in sim.boundary_layers: # boundary on all four sides if boundary.direction == mp.ALL and boundary.side == mp.ALL: if sim.dimensions == 1: dims = [mp.X] elif sim.dimensions == mp.CYLINDRICAL or sim.is_cylindrical: dims = [mp.X, mp.Z] 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]): if ((permutation[0] == mp.X) and (permutation[1] == mp.Low)) and ( sim.dimensions == mp.CYLINDRICAL or sim.is_cylindrical ): continue vol = get_boundary_volumes(boundary.thickness, *permutation) ax = plot_volume( sim, ax, vol, output_plane, plotting_parameters=boundary_parameters ) # boundary on only two of four sides elif boundary.side == mp.ALL: for side in [mp.Low, mp.High]: if ((boundary.direction == mp.X) and (side == mp.Low)) and ( sim.dimensions == mp.CYLINDRICAL or sim.is_cylindrical ): continue vol = get_boundary_volumes(boundary.thickness, boundary.direction, side) ax = plot_volume( sim, ax, vol, output_plane, plotting_parameters=boundary_parameters ) # boundary on just one side else: if ((boundary.direction == mp.X) and (boundary.side == mp.Low)) and ( sim.dimensions == mp.CYLINDRICAL or sim.is_cylindrical ): continue 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: Simulation, ax: Axes, output_plane: Optional[Volume] = None, labels: bool = False, source_parameters: Optional[dict] = None, ) -> Axes: # consolidate plotting parameters if source_parameters is None: source_parameters = default_source_parameters else: source_parameters = 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: Simulation, ax: Axes, output_plane: Optional[Volume] = None, labels: bool = False, monitor_parameters: Optional[dict] = None, ) -> Axes: # consolidate plotting parameters if monitor_parameters is None: monitor_parameters = default_monitor_parameters else: monitor_parameters = 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: Simulation, ax: Optional[Axes] = None, fields: Optional = None, output_plane: Optional[Volume] = None, field_parameters: Optional[dict] = None, colorbar_parameters: Optional[dict] = None, ) -> Union[Axes, Any]: components = { mp.Ex, mp.Ey, mp.Ez, mp.Er, mp.Ep, mp.Dx, mp.Dy, mp.Dz, mp.Dr, mp.Dp, mp.Hx, mp.Hy, mp.Hz, mp.Hr, mp.Hp, mp.Bx, mp.By, mp.Bz, mp.Br, mp.Bp, mp.Sx, mp.Sy, mp.Sz, mp.Sr, mp.Sp, } if not sim._is_initialized: sim.init_sim() if fields is None: return ax if field_parameters is None: field_parameters = default_field_parameters else: field_parameters = dict(default_field_parameters, **field_parameters) # user specifies a field component if fields not in components: 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) 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] 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" field_data = sim.get_array(center=sim_center, size=sim_size, component=fields) field_data = field_parameters["post_process"](field_data) if (sim.dimensions == mp.CYLINDRICAL) or sim.is_cylindrical: field_data = np.flipud(field_data) else: field_data = np.rot90(field_data) # Either plot the field, or return the array if not ax: return field_data if mp.am_master(): ax.imshow(field_data, extent=extent, **filter_dict(field_parameters, ax.imshow)) if field_parameters["colorbar"]: _add_colorbar( ax=ax, cmap=field_parameters["cmap"], vmin=np.amin(field_data), vmax=np.amax(field_data), default_label="field value", colorbar_parameters=colorbar_parameters, ) return ax def plot2D( sim: Simulation, ax: Optional[Axes] = None, output_plane: Optional[Volume] = None, fields: Optional = None, labels: Optional[bool] = False, eps_parameters: Optional[dict] = None, boundary_parameters: Optional[dict] = None, source_parameters: Optional[dict] = None, monitor_parameters: Optional[dict] = None, field_parameters: Optional[dict] = None, colorbar_parameters: Optional[dict] = None, frequency: Optional[float] = None, plot_eps_flag: bool = True, plot_sources_flag: bool = True, plot_monitors_flag: bool = True, plot_boundaries_flag: bool = True, nb: bool = False, ) -> Axes: # Ensure a figure axis exists if ax is None and mp.am_master(): from matplotlib import pyplot as plt # nb = plt.get_backend() == 'module://ipympl.backend_nbagg' ax = plt.gca() # validate the output plane to ensure proper 2D coordinates sim_center, sim_size = get_2D_dimensions(sim, output_plane) output_plane = Volume(center=sim_center, size=sim_size) if eps_parameters is not None and field_parameters is not None: if field_parameters.get("colorbar", False) and eps_parameters.get( "colorbar", False ): raise ValueError( "'colorbar' parameter can only be specified for epsilon or fields, but not both." ) # Plot geometry if plot_eps_flag: ax = plot_eps( sim, ax, output_plane=output_plane, eps_parameters=eps_parameters, colorbar_parameters=colorbar_parameters, frequency=frequency, ) # Plot boundaries if plot_boundaries_flag: ax = plot_boundaries( sim, ax, output_plane=output_plane, boundary_parameters=boundary_parameters, ) # Plot sources if plot_sources_flag: ax = plot_sources( sim, ax, output_plane=output_plane, labels=labels, source_parameters=source_parameters, ) # Plot monitors if plot_monitors_flag: ax = plot_monitors( sim, ax, output_plane=output_plane, labels=labels, monitor_parameters=monitor_parameters, ) # Plot fields if fields is not None: ax = plot_fields( sim, ax, fields, output_plane=output_plane, field_parameters=field_parameters, colorbar_parameters=colorbar_parameters, ) # If using %matplotlib ipympl magic, we need to force the figure to be displayed immediately if mp.am_master() and nb: display_figure_immediately(ax.figure) sleep(0.05) return ax def plot3D(sim, save_to_image: bool = False, image_name: str = "sim.png", **kwargs): from vispy.scene.visuals import Box, Mesh from vispy.scene import SceneCanvas, transforms try: from skimage.measure import marching_cubes except: from skimage.measure import marching_cubes_lewiner as marching_cubes from vispy.visuals.filters import ShadingFilter # Set canvas canvas = SceneCanvas(keys="interactive", bgcolor="white") view = canvas.central_widget.add_view() view.camera = "turntable" # Get domain measurements sim_center, sim_size = sim.geometry_center, sim.cell_size xmin, xmax, ymin, ymax, zmin, zmax = mp.visualization.box_vertices( sim_center, sim_size, sim.is_cylindrical ) 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) xtics = np.linspace(xmin, xmax, Nx) ytics = np.linspace(ymin, ymax, Ny) ztics = np.linspace(zmin, zmax, Nz) # Get eps for geometry eps_data = np.round(np.real(sim.get_epsilon_grid(xtics, ytics, ztics)), 2) unique = np.unique(np.abs(eps_data)).tolist() # Remove background material unique.remove(np.round(np.abs(np.asarray(sim.default_material.epsilon_diag)), 2)[0]) mesh_midpoint = (sim_size[0] / 2, sim_size[1] / 2, sim_size[2] / 2) light_dir = (0, 0, -1, 0) # Build geometry for i, eps in enumerate(unique): eps_ = np.array(eps_data.flatten() == eps).astype(int).reshape(eps_data.shape) marching_cube = marching_cubes( eps_, 0.99, spacing=(sim.cell_size.x / Nx, sim.cell_size.y / Ny, sim.cell_size.z / Nz), ) vertices, faces = marching_cube[0], marching_cube[1] mesh = Mesh( vertices, faces, color=( 1 - ((i + 1) / len(unique)), 1 - ((i + 1) / len(unique)), 1 - ((i + 1) / len(unique)), 0.8, ), ) mesh.transform = transforms.MatrixTransform() mesh.transform.translate(np.asarray(sim.geometry_center)) shading_filter = ShadingFilter(shininess=100) shading_filter.light_dir = light_dir[:3] mesh.attach(shading_filter) view.add(mesh) # Build source thickness = ( sim.boundary_layers[0].thickness if not len(sim.boundary_layers) < 1 else 0 ) for source in sim.sources: size = tuple(source.size) source_box = Box( *size, color=(1, 0, 0, 1), # red ) center = list(source.center) source_box.transform = transforms.MatrixTransform() source_box.transform.translate(np.asarray(mesh_midpoint)) source_box.transform.translate(center) source_box.transform.translate(tuple(sim.geometry_center)) view.add(source_box) # Build monitors for mon in sim.dft_objects: for reg in mon.regions: size = list(reg.size) monitor_box = Box( *size, color=(0, 0, 1, 1), # blue ) center = list(reg.center) monitor_box.transform = transforms.MatrixTransform() vector = [0, 0, 0] vector[reg.direction] = 1 vector = mp.Vector3(*vector) monitor_box.transform.translate(tuple(mesh_midpoint)) monitor_box.transform.translate(center) monitor_box.transform.translate(tuple(sim.geometry_center)) view.add(monitor_box) # Build boundaries for box_center_top in [ np.add(mesh_midpoint, (0, 0, sim_size[2] / 2 - thickness / 2)), np.subtract(mesh_midpoint, (0, 0, sim_size[2] / 2 - thickness / 2)), ]: box = _build_3d_pml(sim_size[0], sim_size[1], thickness, box_center_top) view.add(box) for box_center_right in [ np.add(mesh_midpoint, (sim_size[0] / 2 - thickness / 2, 0, 0)), np.subtract(mesh_midpoint, (sim_size[0] / 2 - thickness / 2, 0, 0)), ]: box = _build_3d_pml(thickness, sim_size[1], sim_size[2], box_center_right) view.add(box) for box_center_front in [ np.add(mesh_midpoint, (0, sim_size[1] / 2 - thickness / 2, 0)), np.subtract(mesh_midpoint, (0, sim_size[1] / 2 - thickness / 2, 0)), ]: box = _build_3d_pml(sim_size[0], thickness, sim_size[2], box_center_front) view.add(box) # Camera options view.camera.center = mesh_midpoint view.camera.scale_factor = getattr( kwargs, "scale_factor", 2 * np.linalg.norm(sim_size) ) view.camera.elevation = getattr(kwargs, "elevation", 10) view.camera.azimuth = getattr(kwargs, "azimuth", 45) view.camera.transform.imap(light_dir) # Plot or save if save_to_image: image = canvas.render() import imageio imageio.imwrite(image_name, image) return canvas.show(run=True) def _build_3d_pml(x: float, y: float, thickness: float, translate: tuple): from vispy.scene.visuals import Box from vispy.scene import transforms from vispy.visuals.filters import WireframeFilter box = Box( x, y, thickness, color=(0, 1, 0, 0.2), # green but transparent # color=None, ) box.transform = transforms.MatrixTransform() box.transform.rotate(90, (1, 0, 0)) box.transform.translate(translate) wireframe_filter = WireframeFilter(width=2) box.mesh.attach(wireframe_filter) return box def visualize_chunks(sim: Simulation): if sim.structure is None: sim.init_sim() import matplotlib.pyplot as plt import matplotlib.cm import matplotlib.colors if sim.structure.gv.dim == 2: from mpl_toolkits.mplot3d import Axes3D from mpl_toolkits.mplot3d.art3d import Poly3DCollection else: from matplotlib.collections import PolyCollection vols = sim.structure.get_chunk_volumes() owners = sim.structure.get_chunk_owners() def plot_box(box, proc, fig, ax: Axes): if sim.structure.gv.dim == 2: 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 = 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] 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) else: low = Vector3(box.low.x, box.low.y) high = Vector3(box.high.x, box.high.y) points = [low, high] 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] points = np.array([np.array(v)[:-1] for v in points]) edges = [[points[0], points[2], points[1], points[3]]] faces = PolyCollection(edges, linewidths=1, edgecolors="k") color_with_alpha = matplotlib.colors.to_rgba(chunk_colors[proc]) faces.set_facecolor(color_with_alpha) ax.add_collection(faces) # Plot the points themselves to force the scaling of the axes ax.scatter(points[:, 0], points[:, 1], s=0) if mp.am_master(): fig = plt.figure() ax = fig.add_subplot( 111, projection="3d" if sim.structure.gv.dim == 2 else None ) 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_xlabel("x") ax.set_ylabel("y") ax.set_aspect("equal") cell_box = mp.gv2box(sim.structure.gv.surroundings()) if sim.structure.gv.dim == 2: ax.set_xlim3d(left=cell_box.low.x, right=cell_box.high.x) ax.set_ylim3d(bottom=cell_box.low.y, top=cell_box.high.y) ax.set_zlim3d(bottom=cell_box.low.z, top=cell_box.high.z) ax.set_zlabel("z") else: ax.set_xlim(left=cell_box.low.x, right=cell_box.high.x) ax.set_ylim(bottom=cell_box.low.y, top=cell_box.high.y) plt.tight_layout() plt.show() def display_figure_immediately(fig: Figure) -> None: """ Trigger the specified figure to display immediately, rather than waiting on the cell execution to end. Due to limitations in ipympl: https://github.com/matplotlib/ipympl/issues/290, which might be fixed at some point in the future. """ from IPython.display import display canvas = fig.canvas display(canvas) canvas._handle_message(canvas, {"type": "send_image_mode"}, []) canvas._handle_message(canvas, {"type": "refresh"}, []) canvas._handle_message(canvas, {"type": "initialized"}, []) canvas._handle_message(canvas, {"type": "draw"}, []) # ------------------------------------------------------- # # JS_Animation # ------------------------------------------------------- # # A helper class used to make jshtml animations embed # seamlessly within Jupyter notebooks. class JS_Animation: def __init__(self, jshtml: str): self.jshtml = jshtml def _repr_html_(self) -> str: return self.jshtml def get_jshtml(self) -> str: return self.jshtml # ------------------------------------------------------- # # 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: """ A class used to record the fields during timestepping (i.e., a [`run`](#run-functions) function). The object is initialized prior to timestepping by specifying the field component. The object can then be passed to any [step-function modifier](#step-function-modifiers). For example, one can record the $E_z$ fields at every one time unit using: ```py animate = mp.Animate2D(fields=mp.Ez, realtime=True, field_parameters={'alpha':0.8, 'cmap':'RdBu', 'interpolation':'none'}, boundary_parameters={'hatch':'o', 'linewidth':1.5, 'facecolor':'y', 'edgecolor':'b', 'alpha':0.3}) 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 fields of a figure can also be updated in realtime by setting the `realtime` flag. This does not work for IPython/Jupyter notebooks, however. Once the simulation is run, the animation can be output 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. """ def __init__( self, sim: Optional[Simulation] = None, fields: Optional = None, f: Optional[Figure] = None, realtime: bool = False, normalize: bool = False, plot_modifiers: Optional[list] = None, update_epsilon: bool = False, nb: bool = False, **customization_args ): """ Construct an `Animate2D` object. + **`sim=None`** — Optional Simulation object (this has no effect, and is included for backwards compatibility). + **`fields=None`** — Optional Field component to record at each time instant. + **`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=False`** — Whether or not to update a figure window in realtime as the simulation progresses. Disabled by default. + **`normalize=False`** — Records fields at each time step in memory in a NumPy array and then normalizes the result by dividing by the maximum field value at a single point in the cell over all the time snapshots. + **`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`: ```py def mod1(ax): ax.set_xlabel('Testing') return ax plot_modifiers = [mod1] ``` + **`update_epsilon=False`** — Redraw epsilon on each call. (Useful for topology optimization) + **`nb=False`** — For the animation work in a Jupyter notebook, set to True and use the cell magic: `%matplotlib ipympl` + **`**customization_args`** — Customization keyword arguments passed to `plot2D()` (i.e. `labels`, `eps_parameters`, `boundary_parameters`, etc.) """ if sim is not None: warnings.warn( "Warning: The 'sim' argument in Animate2D is deprecated and has no effect. It will be removed " "in a future release." ) self.fields = fields self.update_epsilon = update_epsilon self.nb = nb if f: self.f: Figure = f self.ax: Axes = self.f.gca() elif mp.am_master(): from matplotlib import pyplot as plt # To prevent 2 figures from being created in a notebook, interactive must be turned off and back on here # https://matplotlib.org/ipympl/examples/full-example.html#fixing-the-double-display-with-ioff if self.nb: plt.ioff() self.f: Figure = plt.figure() if self.nb: plt.ion() self.ax: Axes = self.f.gca() # This is another option for enabling notebook plotting # self.nb = plt.get_backend() == 'module://ipympl.backend_nbagg' 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: Simulation, todo: str) -> None: from matplotlib import pyplot as plt if todo == "step": # Initialize the plot if not self.init: filtered_plot2D = filter_dict(self.customization_args, plot2D) ax = sim.plot2D( ax=self.ax, fields=self.fields, nb=self.nb, **filtered_plot2D ) # Run the plot modifier functions if self.plot_modifiers: for k in range(len(self.plot_modifiers)): ax = self.plot_modifiers[k](self.ax) # Store the fields if mp.am_master(): fields = ax.images[-1].get_array() self.ax = ax self.w, self.h = self.f.get_size_inches() self.init = True else: if self.update_epsilon: # Update epsilon filtered_plot_eps = filter_dict(self.customization_args, plot_eps) # when calling with no 'ax', returns array of epsilon data eps = plot_eps(sim=sim, **filtered_plot_eps) if mp.am_master(): eps_idx = -1 if not self.fields else -2 self.ax.images[eps_idx].set_data(eps) # Need to check if None because mp.Ex == 0 if self.fields is not None: # Update fields filtered_plot_fields = filter_dict( self.customization_args, plot_fields ) # when calling with no 'ax', returns array of fields data fields = sim.plot_fields(fields=self.fields, **filtered_plot_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 in a Jupyter notebook, we need to redraw the canvas if self.nb and mp.am_master(): self.f.canvas.draw() if self.realtime and mp.am_master(): # Redraw the current figure if requested # For some reason, plt.pause() causes ipympl to redraw the same figure, and we end up with # a new copy of the figure every time this class is called. plt.pause(0.05) if not self.nb else sleep(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(): if mp.verbosity.meep > 0: 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) -> Tuple[int, int]: # 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) -> None: # 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: list, frame_format: str) -> str: # 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: int) -> JS_Animation: """ 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. """ # Exports a javascript enabled html object that is # ready for jupyter notebook embedding. # modified from matplotlib/animation.py code. # Only works with Python3 and matplotlib > 3.1.0 from distutils.version import LooseVersion import matplotlib if LooseVersion(matplotlib.__version__) < LooseVersion("3.1.0"): print("-------------------------------") print( "Warning: JSHTML output is not supported with your current matplotlib build. Consider upgrading to 3.1.0+" ) print("-------------------------------") return if mp.am_master(): from uuid import uuid4 from matplotlib._animation_data import ( DISPLAY_TEMPLATE, INCLUDED_FRAMES, JS_INCLUDE, STYLE_INCLUDE, ) # 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: int, filename: str) -> None: """ Generates and outputs a GIF file of the animation with the filename, `filename`, and the frame rate, `fps`. Note that GIFs are significantly larger than mp4 videos since they don't use any compression. Artifacts are also common because the GIF format only supports 256 colors from a _predefined_ color palette. Requires `ffmpeg`. """ # Exports a gif of the recorded animation # requires ffmpeg to be installed # modified from the matplotlib library if mp.am_master(): 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 ] if mp.verbosity.meep > 0: print("Generating GIF...") 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: int, filename: str) -> None: """ 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`. """ # Exports an mp4 of the recorded animation # requires ffmpeg to be installed # modified from the matplotlib library if mp.am_master(): 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 ] if mp.verbosity.meep > 0: print("Generating MP4...") 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) -> None: self.cumulative_fields = [] self.ax = None self.f = None def set_figure(self, f: Figure) -> None: self.f = f