def adjust_chunk_layout(self, sim: mp.Simulation, **kwargs) -> None:
"""Computes a new chunk layout, applies it to sim, and resets/re-inits sim.
This method also calls self.should_adjust_chunk_layout(sim). If the current
chunk layout is sufficiently well-balanced, no changes will be made.
NOTE: This is a state-changing method which may reset the simulation.
Args:
sim: the meep.Simulation object with the chunk layout to be adjusted
**kwargs: extra args to be passed to self.compute_new_chunk_layout
Raises:
ValueError: if sim.chunk_layout includes nodes with duplicate proc_ids
ValueError: if sim.chunk_layout has proc_ids not included in
sim.structure.get_chunk_owners() (this could occur if number of
processes exceeds the number of physical cores)
"""
self._validate_sim(sim)
if self.should_adjust_chunk_layout(sim):
old_chunk_layout = sim.chunk_layout
chunk_volumes = sim.structure.get_chunk_volumes()
chunk_owners = sim.structure.get_chunk_owners()
timing_measurements = MeepTimingMeasurements.new_from_simulation(sim, -1)
new_chunk_layout = self.compute_new_chunk_layout(timing_measurements,
old_chunk_layout,
chunk_volumes,
chunk_owners, **kwargs)
sim.reset_meep()
sim.chunk_layout = new_chunk_layout
sim.init_sim()
def output(self, simulation: mp.Simulation, each: Union[int, float],
until: Union[int, float]):
simulation.use_output_directory(self.dir_out)
simulation.run(mp.at_every(each,
mp.output_png(mp.Ez,
"-Zc" + self.colormap)),
until=until)
def before_sim(self, sim: mp.Simulation) -> None:
if not self._wg_mode:
self._wg_mode = _get_waveguide_mode(sim, self._src)
# Construct source time object on
for source in _create_waveguide_mode_source(
src=self._src,
wg_mode=self._wg_mode,
src_time=self._get_source_time()):
sim.add_source(source)
def before_sim(self, sim: mp.Simulation) -> None:
"""Creates the DFT field object used later to get the DFT fields."""
self._dft_field = sim.add_dft_fields([mp.Ex, mp.Ey, mp.Ez],
1 / self._wlen,
1 / self._wlen,
1,
center=self._sim_region.center,
size=self._sim_region.extents)
# TODO(logansu): Figure out smart way to handle coordinates.
self._xyzw = sim.get_array_metadata(center=self._sim_region.center,
size=self._sim_region.extents)
def eval(self, sim: mp.Simulation) -> goos.NumericFlow:
eps = sim.get_epsilon(omega=2 * np.pi / self._wlen)
if len(eps.shape) < 3:
eps = np.expand_dims(eps, axis=self._expand_axis)
return goos.NumericFlow(eps)
def gather_design_region_fields(
simulation: mp.Simulation,
design_region_monitors: List[mp.DftFields],
frequencies: List[float],
) -> List[List[onp.ndarray]]:
"""Collects the design region DFT fields from the simulation.
Args:
simulation: the simulation object.
design_region_monitors: the installed design region monitors.
frequencies: the frequencies to monitor.
Returns:
A list of lists. Each entry (list) in the overall list corresponds one-to-
one with a declared design region. For each such contained list, the
entries correspond to the field components that are monitored. The entries
are ndarrays of rank 4 with dimensions (freq, x, y, (z-or-pad)).
The design region fields are sampled on the *Yee grid*. This makes them
fairly awkward to inspect directly. Their primary use case is supporting
gradient calculations.
"""
design_region_fields = []
for monitor in design_region_monitors:
fields_by_component = []
for component in _compute_components(simulation):
fields_by_freq = []
for freq_idx, _ in enumerate(frequencies):
fields = simulation.get_dft_array(monitor, component, freq_idx)
fields_by_freq.append(_make_at_least_nd(fields))
fields_by_component.append(onp.stack(fields_by_freq))
design_region_fields.append(fields_by_component)
return design_region_fields
def before_adjoint_sim(self, adjoint_sim: mp.Simulation,
grad_val: goos.NumericFlow.Grad) -> None:
# It turns out that the gradient of the modal overlap turns out to be
# a TFSF source propagating in the direction opposite of overlap
# direction (i.e. if the overlap computes power flowing in +x direction
# then the mode source propagates in -x direction).
# TODO(logansu): Figure out how to set the appropriate bandwidth
# for the mode propagating backwards.
freq = 1 / self._overlap.wavelength
for src in _create_waveguide_mode_source(
src=self._overlap,
wg_mode=self._wg_mode,
src_time=mp.GaussianSource(frequency=freq, fwidth=0.2 * freq),
amp_factor=grad_val.array_grad / self._mode_norm * 0.5 *
self._amp_factor,
adjoint=True,
):
adjoint_sim.add_source(src)
def eval(self, sim: mp.Simulation) -> goos.NumericFlow:
"""Computes frequency-domain fields.
Frequency domain fields are computed for all three components and
stacked into one numpy array.
Returns:
A `NumericFlow` where the ith entry of the array corresponds to
electric field component i (e.g. 0th entry is Ex).
"""
fields = np.stack([
sim.get_dft_array(self._dft_field, mp.Ex, 0),
sim.get_dft_array(self._dft_field, mp.Ey, 0),
sim.get_dft_array(self._dft_field, mp.Ez, 0)
],
axis=0)
if len(fields.shape) < 4:
fields = np.expand_dims(fields, axis=self._expand_axis[0] + 1)
return goos.NumericFlow(fields)
def new_from_simulation(
cls,
sim: mp.Simulation,
elapsed_time: Optional[float] = -1,
time_per_step: Optional[List[float]] = None,
dft_relative_change: Optional[List[float]] = None,
overlap_relative_change: Optional[List[float]] = None,
relative_energy: Optional[List[float]] = None,
) -> 'MeepTimingMeasurements':
"""Creates a new `MeepTimingMeasurements` from a Meep simulation object.
Usage example:
```
sim = mp.Simulation(...)
sim.init_sim()
start_time = time.time()
sim.run(...)
measurements = meep_timing.MeepTimingMeasurements.new_from_simulation(
sim=sim,
elapsed_time=time.time() - start_time,
)
```
Args:
sim: a Meep simulation object that has previously been run.
elapsed_time: the wall time that has elapsed while running the simulation.
time_per_step: the wall time taken per step. Each element of this list is
expected to correspond to a unique time step, but could also be averaged
over several time steps.
dft_relative_change: the relative change in the DFT fields, measured at
each time step.
overlap_relative_change: the relative change in the mode overlaps,
measured at each time step.
relative_energy: the relative energy in the simulation domain, measured
at each time step.
Returns:
the resulting `MeepTimingMeasurements`
"""
measurements = {}
for name, timing_id in TIMING_MEASUREMENT_IDS.items():
measurements[name] = sim.time_spent_on(timing_id).tolist()
return cls(
measurements=measurements,
elapsed_time=elapsed_time,
num_time_steps=get_current_time_step(sim),
time_per_step=time_per_step if time_per_step is not None else [],
dft_relative_change=dft_relative_change
if dft_relative_change is not None else [],
overlap_relative_change=overlap_relative_change
if overlap_relative_change is not None else [],
relative_energy=relative_energy
if relative_energy is not None else [],
)
def eval(self, sim: mp.Simulation) -> goos.NumericFlow:
# Retrieve the relevant tangential fields.
fields = []
for comp in self._wg_mode.field_comps:
fields.append(
np.reshape(sim.get_dft_array(self._mon, comp, 0),
self._wg_mode.xyzw[3].shape))
norms = _get_overlap_integral(self._wg_mode.mode_fields, fields,
self._wg_mode.xyzw)
if gridlock.axisvec2polarity(self._overlap.normal) > 0:
val = 0.5 * (norms[0] + norms[1]) / self._mode_norm
else:
val = 0.5 * (norms[0] - norms[1]) / self._mode_norm
return goos.NumericFlow(val * self._amp_factor)
def install_design_region_monitors(
simulation: mp.Simulation,
design_regions: List[DesignRegion],
frequencies: List[float],
decimation_factor: int = 0,
) -> List[mp.DftFields]:
"""Installs DFT field monitors at the design regions of the simulation."""
design_region_monitors = [[
simulation.add_dft_fields([comp],
frequencies,
where=design_region.volume,
yee_grid=True,
decimation_factor=decimation_factor,
persist=True)
for comp in _compute_components(simulation)
] for design_region in design_regions]
return design_region_monitors
def before_sim(self, sim: mp.Simulation) -> None:
# Set extents to zero in the direction of the normal as MEEP
# can automatically deduce the normal this way.
extents = np.array(self._overlap.extents) * 1.0
extents[gridlock.axisvec2axis(self._overlap.normal)] = 0
# Add a flux region to monitor the power.
self._mon = sim.add_mode_monitor(
1 / self._overlap.wavelength, 0, 1,
mp.FluxRegion(center=self._overlap.center, size=extents))
# Calculate the eigenmode if we have not already.
if not self._wg_mode:
self._wg_mode = _get_waveguide_mode(sim, self._overlap)
norms = _get_overlap_integral(self._wg_mode.mode_fields,
self._wg_mode.mode_fields,
self._wg_mode.xyzw)
self._mode_norm = np.sqrt(0.5 * np.abs(norms[0] + norms[1]))
def _get_waveguide_mode(sim: mp.Simulation,
src: WaveguideModeSource) -> MeepWaveguideMode:
"""Computes the waveguide mode in Meep (via MPB).
Args:
sim: Meep simulation object.
src: Waveguide mode source properties.
Returns:
A tuple `(mode_fields, mode_comps, xyzw)`. `mode_fields`
is the sampled transverse fields of the mode. `mode_fields` is an array
with four elements, one for each transverse element as determined by
`mode_comps`. `xyzw` is the Meep array metadata for the region: It is
a 4-element tuple `(x, y, z, w)` where the `x`, `y`, and `z` are arrays
indicating the cell coordinates where the fields are sampled and `w`
is the volume weight factor used when calculating integrals involving
the fields (see Meep documentation for details).
"""
wlen = src.wavelength
fcen = 1 / wlen
normal_axis = gridlock.axisvec2axis(src.normal)
# Extract the eigenmode from Meep.
# `xyzw` is a tuple `(x_coords, y_coords, z_coords, weights)` where
# `x_coords`, `y_coords`, and `z_coords` is a list of the coordinates of
# the Yee cells within the source region.
# TODO(logansu): Remove this hack when Meep fixes its bugs with
# `get_array_metadata`. For now, we ensure that the slices are 3D, otherwise
# the values returned for `w` can be wrong and undeterministic.
dx = 1 / sim.resolution
# Count number of dimensions are effectively "zero". This is used to
# determine the dimensionality of the source (1D or 2D).
overlap_dims = 3 - sum(val <= dx for val in src.extents)
extents = [val if val >= dx else dx for val in src.extents]
xyzw = sim.get_array_metadata(center=src.center, size=extents)
# TODO(logansu): Understand how to guess a k-vector. Is it necessary?
k_guess = [0, 0, 0]
k_guess[normal_axis] = fcen
k_guess = mp.Vector3(*k_guess)
mode_dirs = [mp.X, mp.Y, mp.Z]
mode_data = sim.get_eigenmode(
fcen, mode_dirs[normal_axis],
mp.Volume(center=src.center, size=src.extents), src.mode_num + 1,
k_guess)
# Determine which field components are relevant for TFSF source (i.e. the
# field components tangential to the propagation direction.
# For simplicity, we order them in circular permutation order (x->y->z) with
# E fields followed by H fields.
if normal_axis == 0:
field_comps = [mp.Ey, mp.Ez, mp.Hy, mp.Hz]
elif normal_axis == 1:
field_comps = [mp.Ez, mp.Ex, mp.Hz, mp.Hx]
else:
field_comps = [mp.Ex, mp.Ey, mp.Hx, mp.Hy]
# Extract the actual field values for the relevant field components.
# Note that passing in `mode_data.amplitude` into `amp_func` parameter of
# `mp.Source` seems to give a bad source.
mode_fields = []
for c in field_comps:
field_slice = []
for x, y, z in itertools.product(xyzw[0], xyzw[1], xyzw[2]):
field_slice.append(mode_data.amplitude(mp.Vector3(x, y, z), c))
field_slice = np.reshape(field_slice,
(len(xyzw[0]), len(xyzw[1]), len(xyzw[2])))
mode_fields.append(field_slice)
# Sometimes Meep returns an extra layer of fields when one of the dimensions
# of the overlap region is zero. In that case, only use the first slice.
field_slicer = [slice(None), slice(None), slice(None)]
field_slicer[normal_axis] = slice(0, 1)
field_slicer = tuple(field_slicer)
for i in range(4):
mode_fields[i] = mode_fields[i][field_slicer]
xyzw[3] = np.reshape(xyzw[3], (len(xyzw[0]), len(xyzw[1]), len(xyzw[2])))
xyzw[3] = xyzw[3][field_slicer]
# TODO(logansu): See above TODO about hacking `get_array_metadata`.
# For now, just guess the correct value. The error introduced by this
# occurs at the edges of the source/overlap where the fields are small
# anyway.
xyzw[3][:] = dx**overlap_dims
# Fix the phase of the mode by normalizing the phase by the phase where
# the electric field as largest magnitude.
arr = np.hstack([mode_fields[0].flatten(), mode_fields[1].flatten()])
phase_norm = np.angle(arr[np.argmax(np.abs(arr))])
phase_norm = np.exp(1j * phase_norm)
mode_fields = [f / phase_norm for f in mode_fields]
return MeepWaveguideMode(wlen, field_comps, mode_fields, xyzw)