pjz documentation

Skip to main content
`Ctrl`+`K`

# 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.