tsadar.core.modules.distribution_functions.spherical_harmonics

Classes

ArbitraryVr(nvr)	ArbitraryVr is a model for generating numerical radial functions for the spherical harmonics to produce a 2D distribution function.
FLM_MY(vr, dt)	Compute the first-order Legendre moment (FLM) of a distribution function.
FLM_NN(vr)	A neural network module for modeling spherical harmonics coefficients (FLM) as a function of the radial velocity vr.
SphericalHarmonics(dist_cfg)	Represents a 2D velocity distribution function using spherical harmonics expansion.

class tsadar.core.modules.distribution_functions.spherical_harmonics.FLM_NN(vr)Source

A neural network module for modeling spherical harmonics coefficients (FLM) as a function of the radial velocity vr. This module uses two separate MLPs to predict the magnitude and sign of the FLM coefficients, combining them to produce the final output.

flm_mag

MLP that predicts the (logarithmic) magnitude of the FLM coefficients.

Type:

eqx.nn.MLP

flm_sign

MLP that predicts the sign of the FLM coefficients.

Type:

eqx.nn.MLP

vr

Radial velocity array over which the FLM coefficients are evaluated.

Type:

Array

Parameters:

vr (Array) – The radial velocity array.

__call__(**kwargs)Source

Computes the FLM coefficients for the given input. :param f00: The normalization factor for the magnitude. :type f00: float or Array

Returns:

flm (Array) – The computed FLM coefficients as a function of vr.

flm_mag: MLP

flm_sign: MLP

vr: Array

__call__(**kwargs)Source

Call self as a function.

class tsadar.core.modules.distribution_functions.spherical_harmonics.FLM_MY(vr: Array, dt: float)Source

Compute the first-order Legendre moment (FLM) of a distribution function.

This module uses the Mora & Yahi (1982) model for thermal heat-flux reduction in laser-produced plasmas.

vr

Array of velocity values (normalized to thermal velocity).

Type:

Array

dt

Time step or scaling factor applied to the FLM coefficient.

Type:

float

References

Mora, P. & Yahi, H. (1982). Thermal heat-flux reduction in laser-produced plasmas. Phys. Rev. A 26, 2259–2261.

vr: Array

dt: float

__call__(**kwargs)Source

Call self as a function.

class tsadar.core.modules.distribution_functions.spherical_harmonics.ArbitraryVr(nvr)Source

ArbitraryVr is a model for generating numerical radial functions for the spherical harmonics to produce a 2D distribution function.

smooth

A function that applies 1D smoothing to input arrays, parameterized by window size.

Type:

Callable

flm_sign

Learnable parameters representing the sign component of the function, initialized to zeros.

Type:

Array

flm_mag

Learnable parameters representing the magnitude component of the function, initialized to zeros.

Type:

Array

Parameters:

nvr (int) – The number of radial velocity grid points, used to set the size of parameters and smoothing window.

__call__(**kwargs)Source

Computes the radial function by applying smoothing, nonlinearities (tanh and sigmoid), and scaling. Returns the resulting array representing the function values.

smooth: Callable

flm_sign: Array

flm_mag: Array

__call__(**kwargs)Source

Call self as a function.

class tsadar.core.modules.distribution_functions.spherical_harmonics.SphericalHarmonics(dist_cfg)Source

Represents a 2D velocity distribution function using spherical harmonics expansion. This class models the distribution function in velocity space using a truncated spherical harmonics expansion, with coefficients parameterized by various models (neural network, Mora-Yahi, or arbitrary). It supports normalization and parameterization of the distribution, and provides methods to evaluate the distribution on a velocity grid. .. attribute:: vr

Radial velocity grid.

type:

Array

th

Angular grid (theta) in velocity space.

Type:

Array

phi

Angular grid (phi) in velocity space.

Type:

Array

sph_harm_y

Vectorized spherical harmonics function.

Type:

Callable

vr_vxvy

Radial grid in (vx, vy) coordinates.

Type:

Array

Nl

Maximum order of spherical harmonics expansion.

Type:

int

flm

Dictionary of spherical harmonics coefficients.

Type:

Dict[str, Dict[str, Callable]]

m_scale

Scaling factor for the ‘m’ parameter.

Type:

float

m_shift

Shift for the ‘m’ parameter.

Type:

float

act_fun

Activation function for ‘m’ parameter normalization.

Type:

Callable

normed_m

Normalized ‘m’ parameter, defining the super-gaussian order for the f0 term.

Type:

Array

flm_type

Type of parameterization for spherical harmonics coefficients.

Type:

str

Parameters:

dist_cfg (dict) – Configuration dictionary containing distribution parameters.

get_unnormed_params()Source

Returns the unnormalized spherical harmonics coefficients and parameters.

get_unnormed_m()Source

Returns the unnormalized ‘m’ parameter (Super-gaussian order) for the distribution.

get_f00()Source

Computes the isotropic (l=0, m=0) part of the distribution function.

__call__()Source

Evaluates the full distribution function on the (vx, vy) grid using the spherical harmonics expansion and current parameters.

Raises:

NotImplementedError – If an unsupported ‘flm_type’ or spherical harmonics index is requested.

vr: Array

th: Array

phi: Array

vr_vxvy: Array

Nl: int

sph_harm_y: Callable

flm: Dict[str, Dict[str, Callable]]

m_scale: float

m_shift: float

act_fun: Callable

normed_m: Array

flm_type: str

get_unnormed_params()Source

Computes and returns the unnormalized parameters for the spherical harmonics distribution. This method constructs a dictionary of spherical harmonics coefficients (flm_dict) up to order self.Nl. The zeroth order coefficient (f00) is obtained from self.get_f00(). For higher orders, the coefficients are computed using the corresponding functions in self.flm, with keyword arguments including the unnormalized moment (m_f0) and the zeroth order coefficient (f00). :returns: dict –

A dictionary with a single key “flm”, whose value is a nested dictionary of spherical harmonics

coefficients indexed by their order and degree.

get_unnormed_m()Source

get_f00()Source

Computes the normalized isotropic (l=0, m=0) component of the distribution function in velocity space. Which is modeled as a super-gaussian distribution. :returns: jnp.ndarray – The normalized f00 distribution as a function of the radial velocity grid self.vr.

Notes

• The function uses a super-gaussian distribution parameterized by unnormed_m.

• The normalization ensures that the integral of f00 over all velocities equals 1.

• Requires self.get_unnormed_m() to provide the distribution shape parameter, and self.vr to be the velocity grid.

__call__()Source

Evaluates the distribution function on a velocity grid using spherical harmonics expansion. This method reconstructs the distribution function f(vx, vy) by interpolating the zeroth-order coefficient and adding higher-order spherical harmonics contributions. The result is normalized and ensured to be non-negative. :returns: jnp.ndarray – The normalized, non-negative distribution function evaluated on the (vx, vy) grid.

Notes

• FLM coefficients are computed on the radial velocity grid self.vr and interpolated to the (vx, vy) grid defined by self.vr_vxvy. This decoples the radial and cartesian grids.