basis

Functions related to vectors and field expansion in the FMM scheme.

Copyright (c) Meta Platforms, Inc. and affiliates.

class fmmax.basis.Expansion(basis_coefficients: ndarray[Any, Any])

Stores the expansion.

The expansion consists of the integer coefficients of the reciprocal lattice vectors used in the Fourier expansion of fields, permittivities, etc. in the FMM scheme.

basis_coefficients

The integer coefficients of the primitive reciprocal lattice vectors, which generate the full set of reciprocal-space vectors in the expansion.

Type:

numpy.ndarray[Any, Any]

num_terms

The number of terms in the expansion.

class fmmax.basis.LatticeVectors(u: Array, v: Array)

Stores a pair of lattice vectors.

Note that this is just a pair of 2-dimensional vectors, which may be either for the real-space lattice or the reciprocal space lattice, depending on usage.

u

The first primitive lattice vector.

Type:

jax.Array

v

The second primitive lattice vector, with identical shape.

Type:

jax.Array

property reciprocal: LatticeVectors

Returns the corresponding vectors on the reciprocal lattice.

class fmmax.basis.Truncation(value)

Enumerates truncation modes.

fmmax.basis.brillouin_zone_in_plane_wavevector(brillouin_grid_shape: Tuple[int, int], primitive_lattice_vectors: LatticeVectors) → Array

Computes in-plane wavevectors on a regular grid in the first Brillouin zone.

The wavevectors are found by dividing the Brillouin zone into a grid with the specified shape; the wavevectors are at the centers of the grid voxels.

Parameters:

• brillouin_grid_shape – The shape of the wavevector grid.

• primitive_lattice_vectors – The primitive vectors for the real-space lattice.

Returns:

The in-plane wavevectors, with shape brillouin_grid_shape + (2,).

fmmax.basis.generate_expansion(primitive_lattice_vectors: LatticeVectors, approximate_num_terms: int, truncation: Truncation) → Expansion

Generates the expansion for the specified real-space basis.

Parameters:

• primitive_lattice_vectors – The primitive vectors for the real-space lattice.

• approximate_num_terms – The approximate number of terms in the expansion. To maintain a symmetric expansion, the total number of terms may differ from this value.

• truncation – The truncation to be used for the expansion.

Returns:

The Expansion. The basis coefficients of the expansion are sorted so that the zeroth-order term is first.

fmmax.basis.min_array_shape_for_expansion(expansion: Expansion) → Tuple[int, int]

Returns the minimum allowed shape for an array to be expanded.

fmmax.basis.plane_wave_in_plane_wavevector(wavelength: Array, polar_angle: Array, azimuthal_angle: Array, permittivity: Array) → Array

Computes the in-plane wavevector for a plane-wave excitation.

Parameters:

• wavelength – The free-space wavelength of the plane-wave excitation.

• polar_angle – Polar angle of the plane-wave excitation.

• azimuthal_angle – Azimuthal angle of plane-wave the excitation.

• permittivity – Scalar permittivity of the medium in which the polar angle and azimuthal angle are specified.

Returns:

The fundamental transverse wavevector, i.e. (kx0, ky0).

fmmax.basis.transverse_wavevectors(in_plane_wavevector: Array, primitive_lattice_vectors: LatticeVectors, expansion: Expansion) → Array

Obtains all the relevant transverse wavevectors given the expansion.

Parameters:

• in_plane_wavevector – The zeroth-order transverse wavevector.

• primitive_lattice_vectors – The primitive vectors for the real-space lattice.

• expansion – The expansion for which the set of transverse wavevectors are sought.

Returns:

The transverse wavevectors.

fmmax.basis.unit_cell_coordinates(primitive_lattice_vectors: LatticeVectors, shape: Tuple[int, int], num_unit_cells: Tuple[int, int]) → Tuple[Array, Array]

Compute spatial coordinates given the grid shape and number of unit cells.

fmmax.basis.validate_shape_for_expansion(shape: Tuple[int, ...], expansion: Expansion) → None

Validates that the shape is sufficient for the provided expansion.