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

pjz.scatter(epsilon: jax.Array, omega: jax.Array, ports: Sequence[Tuple[jax.Array, int]], field_kwargs: Dict)

  Returns scattering between `ports`, differentiable w.r.t. `epsilon`.

  Parameters:

    - ports – Sequence of `(mode, pos)` tuples that define both the excitation and output overlap operations for modes.

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

    - omega – `(ww,)` array of angular frequencies.

  Returns:

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

pjz.field(epsilon: jax.Array, source: jax.Array, omega: jax.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.

    - omega_range – `(omega_min, omega_max)` range for `omega` values.

    - tt – See `fdtdz_jax.fdtdz()`.

    - dt – See `fdtdz_jax.fdtdz()`.

    - source_width – Number of periods for ramp-up of time-harmonic sources.

    - source_delay – Delay before ramping up source, in `source_width` units.

    - absorption_padding – Padding cells to add along both boundaries of the x- and y-axes for adiabatic absorption boundary conditions.

    - absorption_strength – Scaling coefficient for adiabatic absorption boundary.

    - pml_widths – See `fdtdz_jax.fdtdz()` documentation.

    - pml_alpha_coeff – Constant value for `pml_alpha` parameter of `fdtdz_jax.fdtdz()`.

    - pml_sigma_lnr – Natural logarithm of PML reflectivity.

    - pml_sigma_m – Exponent for spatial scaling of PML.

    - use_reduced_precision – See `fdtdz_jax.fdtdz()` documentation.

    - launch_params – See `fdtdz_jax.fdtdz()` documentation.

  Returns:

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

class pjz.SimParams(omega_range, tt, dt, source_width, source_delay, absorption_padding, absorption_coeff, pml_widths, pml_alpha_coeff, pml_sigma_lnr, pml_sigma_m, use_reduced_precision, launch_params)

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

  Solve for waveguide modes.

  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: jax.Array, interface_positions: jax.Array, magnification: int, zz: int) → jax.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)

## Shape functions

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

  Rectangle.

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

  Circle.

pjz.invert(a)

pjz.union(a, b)

pjz.intersect(a, b)

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

  Dilate

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

  Shift.