Crate treams_core

Expand description

The numerical core of treams-rs: special functions, lattice sums, wave expansions, T-matrices and S-matrices, with analytic derivatives, in Rust and without Python.

A pullback turns the gradient of a real loss with respect to a result into the gradients with respect to the inputs; the glossary defines it with the related terms.

treams-core is an internal crate. The Python package treams_rs calls it through the bindings crate treams-py, which Python imports as treams_rs._native; neither crate is published on crates.io. The user documentation, with examples and the design rationale, is at https://yaugenst.github.io/treams-rs/latest/.

Upstream: treams 0.4.7 at commit 1f5d0d6. The kernels follow its formulas, conventions and mode order. Each module names the treams functions it mirrors, and the module map pairs every module with its treams_rs and treams namespaces.

§Layers

The modules form six layers:

Layer	Modules	Contents
L0	`numerics`, `linalg`, `fpenv`, `threads`	Numerical support
L1	`special`	Special functions
L2	`lattice`	Lattice sums and lattice geometry
L3	`basis`, `sw`, `cw`, `pw`, `rotation`, `channels`, `vectorwaves`, `fields`	Wave families
L4	`coeffs`, `tmatrix`, `ebcm`, `cluster`	Scattering by particles
L5	`smatrix`	Planar and periodic S-matrices

Production code should use its own layer and the layers below it, never a layer above. Modules of one layer may use each other: cw::to_sw() builds spherical modes, and tmatrix::sphere() takes the Mie coefficients of coeffs. Any module may use the root items Complex, Error, Result and MAX_DEGREE. Tests may use any module.

rotation and channels share one residual or gradient type across both multipole families (rotation::RotationResidual, channels::ChannelGradient), so they stay at the top level instead of inside sw and cw.

A directory module keeps its files private and re-exports their items, so callers write sw::expansion and smatrix::tr. Three modules keep their own path: special::coordinates, numerics::broadcast and numerics::parallel.

§Conventions

The contributor-facing implementation rules are the common home for ownership, validation and evidence requirements. The conventions below define the core’s terms.

§Numbers and labels

Complex is num_complex::Complex64, which NumPy calls complex128. Every kernel takes and returns double precision.
The degree is l and the order is m; lmax and mmax bound them. “Order” also names the order of a Bessel function, never a degree.
MAX_DEGREE = 128 bounds the degree l and the size of cylindrical orders m. Translations couple degrees up to special::MAX_ORDER = 256, and Wigner and Kambe labels reach special::MAX_LABEL = 260. treams has no such limits.
pol: u8 is the polarization index, 0 or 1, and helicity: bool selects the convention that treams calls poltype. With helicity true, pol 1 is positive and pol 0 negative helicity; with helicity false, pol 1 is TM (the N multipoles) and pol 0 TE (the M multipoles). Bases list pol 1 before pol 0, as treams does.
special::Radial is the treams modetype: Regular is 'regular', and Singular is 'singular', the outgoing Hankel function of the first kind. The bindings pass it as singular: bool.
A basis entry (pidx, mode) holds the position index (treams pidx) and the mode label, sw::Mode or cw::Mode. positions are Cartesian expansion centres.
A matrix between two bases maps coefficients of the source basis to the destination basis. Translations take the displacement destination - source; lattice sums run over source - destination, and each lattice path says where it negates the displacement.
Lattice sums take the Bloch vector kpar, the Ewald split parameter eta, the lattice vectors (treams a) and the shift (treams r).
Planar inputs: ks holds the two helicity wavenumbers of each medium, kzs the normal wavenumbers, zs the impedances and q the transverse wavevectors. smatrix::chirality_density and smatrix::oriented_chirality use the argument names of their native functions: normal for the normal wavenumbers and, in smatrix::oriented_chirality, transverse for the transverse wavevectors. The smatrix module docs give the order of the sides and of the four S-matrix blocks.

§Names

Items inside sw, cw and pw carry no family prefix; other modules name them by path, as in sw::Mode, cw::Basis and pw::expansion. A module that serves both families prefixes the family: rotation::sw_rotate, channels::cw_periodic_to_pw.
A forward function carries the treams name or the name of its operation: coeffs::mie, cluster::interaction, smatrix::tr. Suffixes mark the variants: _array for an elementwise forward over arrays that returns a residual (special::bessel_array), _matrix for a matrix that changes the wave family (cw::to_sw_matrix), and _value for the forward that returns only the value of a forward with a residual (smatrix::tr_value). Matrices between two bases of one family are nouns: sw::expansion, sw::lattice_expansion, rotation::sw_rotation.
A forward X that supports gradients returns (value, XResidual) when its pullback does not read the value. When the pullback reads it, the forward returns only the residual, which lends the value through an accessor such as value(): coeffs::mie, smatrix::tr and linalg::solve work this way. XResidual::pullback(&self, cotangent) returns an XGradient with one field per differentiable input, or an array of gradients in argument order. The residual name drops the suffixes _array, _matrix and _value: sw::polar_translation_array returns a sw::PolarTranslationResidual, and cw::to_sw_matrix a cw::ToSwResidual.
The record methods of cluster::InteractionFactor and cluster::IterativeSphereCluster return residuals named after their native context classes. cluster::IlluminateResidual matches IlluminateContext, which diff.illuminate also returns. cluster::IterativeResidual matches IterativeContext.
Family twins share one residual: rotation::sw_rotation and rotation::cw_rotation return rotation::RotationResidual, and cluster::particle_cluster and cluster::cylindrical_particle_cluster return cluster::ParticleClusterResidual. linalg::solve_owned, which reuses the operator’s buffer, returns the linalg::SolveResidual of linalg::solve.
The pullbacks of linalg::EigResidual and smatrix::BandsResidual take one cotangent per output and name each after its output: values and vectors, and wavenumbers and vectors.

§Gradients of complex values

A loss L changes by dL = Re Σ conj(g)·dx when a value changes by dx, where g is the cotangent of that value. For a complex value x = a + i b, the cotangent is ∂L/∂a + i ∂L/∂b, and real inputs get real gradients. A pullback through a holomorphic map y = f(x) therefore multiplies by the conjugate derivative: g_x = conj(f'(x)) g_y.

§Reusable derivatives and saved state

A residual represents fixed primal inputs. Its pushforward(&self, tangent) maps input directions to output directions; pullback(&self, cotangent) maps an output cotangent to input gradients. Both borrow the residual, so repeated calls can reuse recorded values and factorizations. linalg::SolveResidual is an example. These are first-order actions; their presence does not imply support for differentiating the derivative implementation itself.

Keep invariant-bearing residual fields private and retain only data needed by those actions. Returned gradient records may expose plain data. Prefer moving owned buffers when the caller has no further use for them, as linalg::solve_owned does, over cloning to satisfy a borrowed interface.

saved::SavedState carries numerical state through array-only callbacks. Its bytes belong to one invocation of the current extension, not a persistent file format. Keep shape-derived sizes, tags, writes and restores synchronized; validate lengths, domains and allocation sizes before numerical work. Tests should produce state through the public operation, then check repeated round trips, both derivative directions, deferred errors and zero directions. A layout-preserving cleanup retains every byte; removing payload or padding is a separate compatibility decision.

§Errors

Every fallible function returns Result, whose Error says why it failed: InvalidInput for input outside the domain, SpecialFunction, NonFinite and NotConverged for numerical failures, Singular for a singular linear system, and OutOfMemory when the system refuses the memory of a large output or workspace. The bindings raise OutOfMemory as Python MemoryError and every other variant as ValueError, with the displayed message. treams returns NaN or emits NumPy warnings in the numerical cases.

Reserve advertised large outputs and workspaces fallibly, through numerics::zeros, numerics::filled, numerics::reserve and the LU workspace of linalg, so a refused request returns OutOfMemory instead of aborting the process. This is not a process-wide guarantee: input copies, matrix products, gradient buffers and other allocations can still abort when refused. Branch on typed error identity, not display text, and propagate resource and programming errors through numerical fallback.

§Parallelism and reproducibility

Rayon runs independent items in parallel, on the pool of threads, once their number reaches a threshold; numerics::parallel lists every threshold. The elementwise evaluations and the helpers there keep every output in index order. A pullback that adds over items splits them into chunks fixed by the item count, or into blocks fixed by the shape, and adds the partial sums in order, so the thread count sets how many workers run the chunks but never the order of the additions. The docs of these residuals, and of numerics::broadcast for the elementwise ones, say how they add; the thread count changes the result of none of them:

Outputs in index order, or sums in a fixed order: the elementwise residuals of special and lattice::SumResidual, sw::PolarTranslationResidual, sw::ExpansionResidual, sw::LatticeExpansionFromTableResidual, cw::PolarTranslationResidual, cw::ExpansionResidual, cw::LatticeExpansionResidual, rotation::RotationResidual, pw::PhasesResidual, pw::PermutationResidual, vectorwaves::VectorWaveResidual, ebcm::QmatResidual, coeffs::MieResidual, coeffs::MieCylResidual, tmatrix::SphereResidual, tmatrix::CylinderResidual, tmatrix::MetricResidual, cluster::InteractionResidual, cluster::IlluminateResidual and cluster::SphereClusterResidual.
Partial sums over fixed chunks or blocks, added in order: sw::LatticeExpansionResidual, sw::PeriodicToCwResidual, cw::ToSwResidual, pw::ExpansionResidual, pw::FieldResidual, channels::SphericalChannelsResidual, channels::CylindricalChannelsResidual, fields::FieldResidual, fields::OperatorResidual, cluster::IterativeResidual and smatrix::LayerStackResidual.

The intended contract is identical bits at every thread budget, including one thread (threads::set_num_threads or TREAMS_RS_NUM_THREADS=1).

§Floating-point environment

The kernels assume IEEE 754 gradual underflow: subnormal numbers stay subnormal. JAX (through XLA) and torch.set_flush_denormal(True) make the calling thread flush them to zero, which breaks algorithms that scale through the subnormal range, such as faer’s complex reciprocal of an LU pivot. fpenv::ieee runs a closure with subnormals on the calling thread and restores the caller’s mode afterwards. Its name says what the closure gets, the IEEE 754 default arithmetic, and stays short because the bindings wrap the body of every Python entry point in it: ieee(|| ...). fpenv::flushing does the opposite, for tests. The fpenv module docs explain the control registers and why the guard is sound.

§Lints

The workspace keeps clippy’s default lints, adds pedantic and nursery, requires docs on every public item, and the checks treat every warning as an error. clippy::indexing_slicing stays on, so new code indexes through iterators or checked access. A numerical module that validates its indices on entry opts out once, with #![allow(clippy::indexing_slicing)] // <why the indices are in range>; the reason comment is required. Tests may index freely (clippy.toml).

§Glossary

A record is a function that returns a value and a reusable context. The context stores what is needed to compute derivatives later: context.pullback(g) takes the gradient g of a real-valued loss with respect to the value and returns the gradients with respect to the inputs, one for each differentiable input, in the order of the arguments. Gradients follow the convention dL = Re Σ conj(g)·dx.

In the Rust core, a function returns (value, XResidual), and XResidual::pullback(&self, cotangent) returns the input gradients as XGradient.

Term	Meaning
forward	A function that computes a value. One that supports gradients returns a residual, beside the value or holding it.
residual	What a forward saves for its pushforward and pullback: `XResidual` for the forward `X`. Not the `b - A x` of a linear system, which the GMRES solver of `linalg` calls the residual.
tangent	An input or output direction at the recorded primal inputs.
pushforward	`XResidual::pushforward`: maps input tangents to output tangents while borrowing the reusable residual.
pullback	`XResidual::pullback`: turns the gradient with respect to the value into the gradients with respect to the inputs. It borrows the reusable residual.
cotangent	The gradient of a real-valued loss with respect to one value: the `g` of the definition above.
gradient	`XGradient`: the input gradients that a pullback returns, in the order of the forward’s arguments.
record	The Python name of a forward with a residual: a function of `treams_rs.diff` that returns `(value, context)`.
context	The Python object that holds a residual, named after its record: `MieContext` for `diff.mie`.
jet	A value carried through the arithmetic with its first derivatives along `N` directions (`numerics::Jet<N>`).
`RadialJet`	A radial function with its first and second derivatives in its one complex argument.
split (`eta`)	The Ewald split parameter, which divides a lattice sum into a real-space and a reciprocal-space part; `eta = 0` selects it automatically.
Bloch vector (`kpar`)	The wavevector components along a lattice: the term of lattice point `R` carries the phase `exp(i kpar · R)`.
sheet	The branch of the complex logarithm on which a reciprocal-space term of a lattice sum is evaluated.

The lattice module docs define the other terms of the lattice sums.

§Tests

The tests have two tiers:

A module’s inline tests module checks how that module computes its results: fast paths against reference paths, tables and plans, dispatch thresholds, error paths, accuracy against reference tables, and private helpers.
src/properties/<domain>.rs checks what the public operations of one domain promise: physical laws, analytic identities, and pullbacks against finite differences and adjoint pairings. The domains are special, linalg, lattice, waves, plane, smatrix and tmatrix, as in the Python suite; lattice is the directory src/properties/lattice/.

A test of one module’s implementation goes inline; a physical, analytic or adjoint identity goes into its domain file. Both tiers use crate-private items, so the crate has no tests/ directory. cargo test -p treams-core runs both. Testing on the docs site gives the helpers and conventions of both tiers.

§Module map

Each module mirrors one treams namespace where one exists. A dash marks a module without a counterpart.

Rust module	treams_rs namespace	treams	Contents
`special`	`treams_rs.special`	`treams.special`	Bessel, Legendre and Wigner functions, coordinate transforms, incomplete gamma and Kambe integrals
`lattice`	`treams_rs.lattice`, `treams_rs.misc.firstbrillouin*`	`treams.lattice`, `treams.misc.firstbrillouin*`	Lattice sums and lattice geometry
`sw`	`treams_rs.sw`, `treams_rs.special.tl_vsw_*`	`treams.sw`, `treams.special.tl_vsw_*`	Spherical modes, translation coefficients, expansion matrices and the conversion of chains to cylindrical waves
`cw`	`treams_rs.cw`, `treams_rs.special.tl_vcw`	`treams.cw`, `treams.special.tl_vcw`	Cylindrical modes, translation coefficients, expansion matrices and the conversion to spherical waves
`pw`	`treams_rs.pw`, `treams_rs.misc.wave_vec_z`	`treams.pw`, `treams.misc.wave_vec_z`	Plane-wave polarizations, fields and multipole expansions
`basis`	`treams_rs.SphericalBasis`, `treams_rs.CylindricalBasis`	`treams.SphericalWaveBasis`, `treams.CylindricalWaveBasis`	Multipole bases and the expansion gradients that both families share
`rotation`	`treams_rs.sw.rotate`, `treams_rs.cw.rotate`, `treams_rs.operators.Rotate`	`treams.sw.rotate`, `treams.cw.rotate`, `treams.Rotate`	Rotation coefficients and rotation matrices of multipole bases
`channels`	`treams_rs.sw.periodic_to_pw`, `treams_rs.cw.periodic_to_pw`	`treams.sw.periodic_to_pw`, `treams.cw.periodic_to_pw`	Plane-wave channels of periodic arrays
`vectorwaves`	`treams_rs.special.vsh_`, `vsw_`, `vcw_`, `vpw_`, `sph_harm`	`treams.special.vsh_`, `vsw_`, `vcw_`, `vpw_`, `sph_harm`	Vector spherical harmonics and vector spherical, cylindrical and plane waves
`fields`	`treams_rs.operators.efield`	`treams.efield`	Electric fields of multipole expansions at sample points
`coeffs`	`treams_rs.coeffs`, `treams_rs.Material`, `treams_rs.misc.refractive_index`	`treams.coeffs`, `treams.Material`, `treams.misc.refractive_index`	Mie coefficients of multilayer spheres and cylinders, and materials
`tmatrix`	`treams_rs.TMatrix`, `treams_rs.CylindricalTMatrix`	`treams.TMatrix`, `treams.TMatrixC`	T-matrices of single spheres and cylinders, and the cd, db and chi metrics
`ebcm`	`treams_rs.ebcm`	`treams.ebcm`	Q-matrices of the extended boundary condition method
`cluster`	`treams_rs.Cluster`, `treams_rs.iterative`	`treams.TMatrix.cluster`, `treams.TMatrix.interaction`	Multiple scattering in particle clusters, dense and matrix-free
`smatrix`	`treams_rs.SMatrix`, `treams_rs.coeffs.fresnel`	`treams.SMatrices`, `treams.coeffs.fresnel`	Planar S-matrices: interfaces, layer stacks, composition, illumination, transmittance and bands
`linalg`	`treams_rs.diff.solve`, `svdvals`, `eig`	-	Dense LU, SVD and eigensystems with pullbacks, and restarted GMRES
`numerics`	-	-	Forward-mode jets, broadcasting and parallel thresholds
`fpenv`	-	-	A guard that keeps subnormal numbers when the caller flushes them to zero
`threads`	`treams_rs.set_num_threads`, `threads`, `thread_info`	-	The thread budget and the fork-safe pool that runs every parallel region

Modules§

basis: Multipole bases of modes at expansion positions, and the expansion gradients shared by the spherical and cylindrical families.
channels: Plane-wave channels of periodic arrays: how the multipoles of a 2D array of spheres or a 1D array of cylinders couple to the diffracted plane waves on either side.
cluster: Multiple scattering in clusters of particles: solves of (I - T C) X = T B.
coeffs: Mie coefficients of multilayer chiral spheres and cylinders, and the materials they take.
cw: Cylindrical vector waves: mode labels, translation coefficients of mode pairs, expansion matrices between bases and over lattices, and the conversion to spherical waves.
ebcm: Q-matrices of the extended boundary condition method (EBCM) for particles that are symmetric about the z axis, with analytic gradients for the surface and the media.
fields: Cartesian spherical and cylindrical vector waves, field sums and field operators.
fpenv: IEEE 754 subnormals for the numerical work, whatever the caller’s floating-point mode. treams-rs extension.
lattice: Lattice sums of outgoing spherical and cylindrical waves, and lattice geometry.
linalg: Dense linear solves and general complex eigensystems with analytic derivatives, and restarted GMRES. treams-rs extension.
numerics 🔒: Numerical support without physics: complex helpers, forward-mode jets, broadcasting, fallible allocation and the parallel policy. treams-rs extension.
pw: Plane waves: polarization vectors, fields, translation phases, cyclic permutations of the Cartesian axes, and the expansion of plane waves in regular spherical and cylindrical waves.
rotation: Multipole rotations in the z-y-z convention, with analytic Euler-angle pullbacks.
saved: Numeric derivative state carried by array-only framework callbacks.
smatrix: Planar scattering matrices of plane-wave modes and their analytic pullbacks.
special: Special functions: Bessel, Legendre, Ferrers and Wigner functions, solid harmonics, incomplete gamma and Kambe integrals, and coordinate transforms.
sw: Spherical vector waves: mode labels, translation coefficients of mode pairs, expansion matrices between bases and over lattices, and the cylindrical waves radiated by a z-periodic chain of spheres.
threads: The thread budget and the Rayon pool that runs every parallel region. treams-rs extension.
tmatrix: T-matrices of single spheres and cylinders, and three helicity measures of a T-matrix.
vectorwaves: Vector waves one at a time: the vector spherical harmonics and the spherical, cylindrical and plane vector waves in their local component frames, with pullbacks.

Enums§

DerivativeError: A stable reason why a requested derivative is unavailable.
Error: A numerical failure, invalid physical input or refused memory.

Constants§

MAX_DEGREE: The largest multipole degree l, and the largest |m| of a cylindrical mode.

Type Aliases§

Complex: Complex double precision used throughout the numerical core.
Result: Numerical result with a typed error.