pjz 😴👕👖💤: Photonics on JAX

pjz is JAX and fdtd-z, a set of tools for runnning photonic simulation and optimization workflows at scale.

Main API

class pjz.SimParams(omega_range: Tuple[float, float], tt: int, dt: float = 0.5, source_ramp: float = 4.0, source_delay: float = 4.0, absorption_padding: int = 50, absorption_coeff: float = 0.0001, pml_widths: Tuple[int, int] = (16, 16), pml_alpha_coeff: float = 0.0, pml_sigma_lnr: float = 0.5, pml_sigma_m: float = 1.3, use_reduced_precision: bool = True, launch_params: Any | None = None)

Simulation parameters for interfacing with the fdtdz package.

Parameterize fdtdz.fdtdz() for the specific case of extracting time-harmonic solutions.

omega_range

(omega_min, omega_max) range of possible values for the angular frequencies to be extracted from a single simulation.

Type:

Tuple[float, float]

tt

Number of FDTD updates to perform.

Type:

int

source_ramp

Number of periods over which to bring sources to steady state.

Type:

float

source_delay

Number of periods before starting the source ramp.

Type:

float

absorption_padding

Number of cells additional padding cells in the x-y plane to use for adiabatic absorbing conditions.

Type:

int

absorption_coeff

Strength of the adiabatic absorber.

Type:

float

pml_widths

(pml_lo, pml_hi) denoting the number of PML cells to use at the bottom and top of the simulation domain. See fdtdz.fdtdz() docstring for additional details.

Type:

Tuple[int, int]

pml_alpha_coeff

Strength of alpha parameter of the PML, see the fdtdz.fdtdz() docstring for more details.

Type:

float

pml_sigma_lnr

Controls the sigma parameter of the PML, see the fdtdz.fdtdz() docstring for more details.

Type:

float

pml_sigma_m

Controls the sigma parameter of the PML, see the fdtdz.fdtdz() docstring for more details.

Type:

float

use_reduced_precision

See the fdtdz.fdtdz() docstring for more details.

Type:

bool

launch_params

See the fdtdz.fdtdz() docstring for more details.

Type:

Any

pjz.field(epsilon: Array, source: Array, omega: Array, source_pos: int, sim_params: SimParams)

Time-harmonic solution of Maxwell’s equations.

Parameters:

• epsilon – (3, xx, yy, zz) array of permittivity values.

• source – Array of excitation values of shape (2, 1, yy, zz), (2, xx, 1, yy), or (2, xx, yy, 1).

• omega – (ww,) array of angular frequencies.

• source_pos – Position of source along axis of propagation.

• sim_params – Simulation parameters.

Returns:

(ww, 3, xx, yy, zz) array of complex-valued field values at the various omega.

pjz.scatter(epsilon: Array, omega: Array, modes: Tuple[Array], pos: Tuple[int], sim_params: SimParams)

Differentiable time-harmonic scattering values between modes.

Parameters:

• epsilon – (3, xx, yy, zz) array of permittivity values. Differentiable.

• omega – (ww,) array of angular frequencies.

• modes – (2, xx, yy, zz) arrays with exactly one singular spatial dimension identifying the “ports” to be used when computing scattering parameters.

• pos – Integers denoting the location of modes along their respective propagation axes.

• sim_params – Simulation parameters.

Returns:

Scattering values as svals[i][j] nested lists of (ww,) arrays containing the scattering values from mode i to mode j over angular frequencies omega.

pjz.mode(epsilon: Array, omega: Array, num_modes: int, init: Array | None = None, shift_iters: int = 10, max_iters: int = 100000, tol: float = 0.0001) → Tuple[Array, Array, Array, int]

Solve for waveguide modes.

Uses the subspace iteration method to obtain modes without leaving JAX. Allows for updating mode solutions via the init parameter (e.g. for small changes in epsilon and/or omega).

Parameters:

• epsilon – (3, xx, yy, zz) array of permittivity values with exactly one xx, yy, or zz equal to 1.

• omega – Real-valued scalar angular frequency.

• num_modes – Integer denoting number of modes to solve for.

• init – (2, xx, yy, zz, num_modes) of values to use as initial guess.

• shift_iters – Number of iterations used to determine the largest eigenvalue of the waveguide operator.

• max_iters – Maximum number of eigenvalue solver iterations to execute.

• tol – Error threshold for eigenvalue solver.

Returns:

(wavevector, excitation, err, iters) where iters is the number of executed solver iterations, and excitation.shape == (2, xx, yy, zz, num_modes) and wavevector.shape == err.shape == (num_modes,), with excitation[..., i], wavevector[i], and err[i] being ordered such that i == 0 corresponds to the fundamental mode.

pjz.epsilon(layers: Array, interface_positions: Array, magnification: int, zz: int) → Array

Render a three-dimensional vector array of permittivity values.

Produces a 3D vector array of permittivity values on the Yee cell based on a layered stack of 2D profiles at magnification 2 * m. Along the z-axis, both the layer boundaries and grid positions are allowed to vary continuously, while along the x- and y-axes the size of each (unmagnified) cell is assumed to be 1.

Attempts to follow [1] but only computes the on-diagonal elements of the projection matrix and is adapted to a situation where there are no explicit interfaces because the pixel values are allowed to vary continuously within each layer.

Instead, the diagonal elements of the projection matrix for a given subvolume are estimated by computing gradients across it where df(u)/du is computed as the integral of f(u) * u over the integral of u**2 where u is relative to the center of the cell.

[1]

Farjadpour, Ardavan, et al. “Improving accuracy by subpixel smoothing in the finite-difference time domain.” Optics letters 31.20 (2006): 2972-2974

Parameters:

• layers – (ll, 2 * m * xx, 2 * m * yy) array of magnified layer profiles.

• interface_positions – (ll - 1) array of interface positions between the ll layer. Assumed to be in monotonically increasing order.

• magnification – Denotes a 2 * m in-plane magnification factor of layer profiles.

• zz – Number of cells along z-axis.

Returns:

(3, xx, yy, zz) array of permittivity values with offsets and vector components according to the finite-difference Yee cell.

pjz.density(u, radius, alpha, c=1.0, eta=0.5, eta_lo=0.25, eta_hi=0.75)

Computes a per-pixel density from raw optimization variable u.

Implements the “three-field” scheme detailed in [2] in order to allow for a final density that is binary (with the exception of boundary values) and that conforms to a minimum feature size requirement.

[2]

Zhou, Mingdong, et al. “Minimum length scale in topology optimization by geometric constraints.” Computer Methods in Applied Mechanics and Engineering 293 (2015): 266-282.

Parameters:

• u – (xx, yy)` variable array with values within [0, 1].

• radius – Radius of the conical filter used to blur u.

• alpha – Float within [0, 1] controlling binarization of the density, where 0 denotes no binarization, and 1 denotes full binarization of all pixels except those on boundaries (as given by eta) which are left unchanged (equal to the alpha = 0 case).

• c – Controls the detection of inflection points.

• eta – Threshold value used to binarize the density.

• eta_lo – Controls minimum feature size of void-phase features.

• eta_hi – Controls minimum feature size of density = 1 features.

Returns:

(density, loss) arrays, both of shape (xx, yy), corresponding to the pixel density and the minimum feature size loss respectively.

Shape functions

pjz.rect(shape: Tuple[int, ...], center: Array, widths: Array) → Array

Rectangle.

pjz.circ(shape: Tuple[int, ...], center: Array, radius: Array) → Array

Circle.

pjz.invert(a)

pjz.union(a, b)

pjz.intersect(a, b)

pjz.dilate(a: Array, radius: float) → Array

Dilate

pjz.shift(a: Array, axis: int, distance: float) → Array

Shift.