tsadar.core.modules.distribution_functions.base#
Functions
|
Calculates the moment of the distribtuion function specified by m |
|
Generates a filter for seperating trainable parameters in a distribution function from static parameters, based on the distribution type and parameters. |
|
Applies a second-order Butterworth filter to a signal using JAX. |
|
Smooths a 1D array using a Hanning window. |
|
Smooths a 2D array using a Hanning window of the specified size. |
|
JAX compatible trapizoidal intergration. |
|
Updates the filter_spec tree by replacing the weights and biases of each layer in the distribution function neural network. |
Classes
|
Represents a 1D arbitrary velocity distribution function. |
|
Arbitrary2V is a two-velocity distribution function class that allows for arbitrary initialization and parameterization. |
|
DLM1V is a 1D distribution function model based on a parameterized "m" shape parameter, with support for activation functions and interpolation over precomputed distributions. |
|
Base class for 1D velocity distribution functions. |
|
A base class for 2D velocity distribution functions. |
- tsadar.core.modules.distribution_functions.base.smooth1d(array, window_size)Source#
Smooths a 1D array using a Hanning window.
- Parameters:
array (jnp.ndarray) – Input 1D array to be smoothed.
window_size (int) – Size of the Hanning window to use for smoothing.
- Returns:
jnp.ndarray – Smoothed array of the same shape as the input.
Notes
The function uses a Hanning window for smoothing and applies convolution with ‘same’ mode.
Requires JAX’s numpy module (jnp).
- tsadar.core.modules.distribution_functions.base.second_order_butterworth(signal: Array, f_sampling: int = 100, f_cutoff: int = 15, method: str = 'forward_backward') ArraySource#
Applies a second-order Butterworth filter to a signal using JAX. This function implements a digital Butterworth filter, similar to using scipy.signal.butter and scipy.signal.filtfilt, but is compatible with JAX arrays. It supports forward, backward, and forward-backward (zero-phase) filtering.
- Parameters:
signal (Array) – The input signal to be filtered.
f_sampling (int, optional) – The sampling frequency of the signal. Default is 100.
f_cutoff (int, optional) – The cutoff frequency of the filter. Default is 15.
method (str, optional) – The filtering method to use. Can be “forward”, “backward”, or “forward_backward” (default). “forward_backward” applies zero-phase filtering by filtering forward and then backward.
- Returns:
Array – The filtered signal.
- Raises:
NotImplementedError – If an unsupported method is specified.
References
Adapted from jax-ml/jax#17540
- tsadar.core.modules.distribution_functions.base.smooth2d(array, window_size)Source#
Smooths a 2D array using a Hanning window of the specified size.
- Parameters:
array (jnp.ndarray) – The 2D input array to be smoothed.
window_size (int) – The size of the Hanning window to use for smoothing.
- Returns:
jnp.ndarray – The smoothed 2D array, with the same shape as the input.
Notes
This function applies a 2D Hanning window to the input array and performs convolution.
The convolution is performed with ‘same’ mode, so the output has the same shape as the input.
Requires the input array and window size to be compatible with JAX (jnp).
- class tsadar.core.modules.distribution_functions.base.DistributionFunction1V(dist_cfg: Dict)Source#
Base class for 1D velocity distribution functions. This class represents a distribution function defined over a 1D velocity grid. It initializes the velocity grid vx based on the configuration provided.
- vx#
1D array of velocity grid points.
- Type:
Array
- Parameters:
dist_cfg (Dict) –
Configuration dictionary containing:
”nvx” (int): Number of velocity grid points.
- Raises:
NotImplementedError – If the instance is called directly, as this is an abstract base class.
- vx: Array#
- class tsadar.core.modules.distribution_functions.base.Arbitrary1V(dist_cfg)Source#
Represents a 1D arbitrary velocity distribution function. This class allows for the initialization, smoothing, and evaluation of a custom 1D distribution function. The distribution is initialized using a Super-gaussian distribtuion parameterized by a parameter m. The distribution function is defined in a 1D velocity space and can be smoothed using a second-order Butterworth filter.
- fval#
The internal representation of the distribution function values.
- Type:
Array
- smooth#
A smoothing function (Butterworth filter) applied to the distribution.
- Type:
Callable
- __init__(dist_cfg)Source#
Initializes the distribution function with configuration parameters, sets up the initial distribution and smoothing filter.
- init_dlm(m)Source#
Initializes the distribution function using the provided shape parameter m. Returns the processed distribution array.
- get_unnormed_params()Source#
Returns a dictionary containing the current (unnormalized) distribution function.
- __call__()Source#
Applies smoothing and normalization to the distribution function and returns the normalized distribution array.
- fval: Array#
- smooth: Callable#
- class tsadar.core.modules.distribution_functions.base.DLM1V(dist_cfg, activate=False)Source#
DLM1V is a 1D distribution function model based on a parameterized “m” shape parameter, with support for activation functions and interpolation over precomputed distributions. .. attribute:: normed_m
The normalized “m” parameter, possibly transformed by an activation function.
- type:
Array
- m_scale#
Scaling factor for the “m” parameter normalization.
- Type:
float
- m_shift#
Shift applied during “m” parameter normalization.
- Type:
float
- act_fun#
Activation function applied to the normalized “m” parameter.
- Type:
Callable
- f_vx_m#
Precomputed distribution values over velocity and “m” axes.
- Type:
Array
- interpolate_f_in_m#
Interpolation function for evaluating the distribution at arbitrary “m”.
- Type:
Callable
- m_ax#
Array of “m” values corresponding to precomputed distributions.
- Type:
Array
- __init__(dist_cfg, activate=False)Source#
Initializes the DLM1V distribution with configuration, normalization, and activation options. Loads precomputed distributions and sets up interpolation.
- __call__()Source#
Evaluates the distribution function for the current “m” parameter, interpolating as needed, and returns the normalized distribution over the velocity axis.
- m_scale: float#
- m_shift: float#
- act_fun: Callable#
- normed_m: Array#
- m_ax: Array#
- f_vx_m: Array#
- interpolate_f_in_m: Callable#
- __call__()Source#
Computes the normalized distribution function for the current parameters. This method applies the activation function to the normalized parameter normed_m, scales and shifts it to obtain unnormed_m, and then interpolates the distribution function using interpolate_f_in_m. The resulting distribution is normalized such that its sum over the velocity axis vx is unity. :returns: jnp.ndarray – The normalized distribution function evaluated over the velocity grid.
- call_matte(unnormed_m)Source#
Computes the normalized distribution function for the current parameters. This method applies the activation function to the normalized parameter normed_m, scales and shifts it to obtain unnormed_m, and then interpolates the distribution function using interpolate_f_in_m. The resulting distribution is normalized such that its sum over the velocity axis vx is unity. :returns: jnp.ndarray – The normalized distribution function evaluated over the velocity grid.
- class tsadar.core.modules.distribution_functions.base.DistributionFunction2V(dist_cfg)Source#
A base class for 2D velocity distribution functions. This class initializes a velocity grid for use in distribution function calculations, centered around zero and spanning from -vmax to vmax, with a specified number of grid points. This velocity grid is symetric in both x and y directions. Parameters ———- dist_cfg : dict
Configuration dictionary containing:
“nvx”: int
Number of velocity grid points along the x-axis.
Attributes#
vx : Array
1D array of velocity grid points along the x-axis.
Methods#
Calls the parent class’s __call__ method.
- vx: Array#
- class tsadar.core.modules.distribution_functions.base.Arbitrary2V(dist_cfg)Source#
Arbitrary2V is a two-velocity distribution function class that allows for arbitrary initialization and parameterization.
- fval#
The current value of the distribution function parameters.
- Type:
Array
- learn_log#
If True, the logarithm (base 10) of the distribution is learned instead of the distribution itself.
- Type:
bool
- __init__(dist_cfg)Source#
Initializes the Arbitrary2V distribution with configuration parameters. :param dist_cfg: Configuration dictionary containing initialization parameters. :type dist_cfg: dict
- init_dlm(m)Source#
Initializes the distribution function using a generalized Super-gaussian form. :param m: Super-gaussian order for the distribution. :type m: float
- Returns:
Array – The initialized distribution function values.
- get_unnormed_params()Source#
Returns the current (unnormalized) distribution parameters. :returns: dict – Dictionary with the current distribution function.
- __call__()Source#
Computes the normalized distribution function based on current parameters. :returns: Array – The normalized distribution function.
- learn_log: bool#
- fval: Array#
- init_dlm(m)Source#
Initialize the distribution function using the Dum-Langdon-Matte (DLM) form. Parameters ———- m : float
The super-gaussian order parameter for the DLM, controlling the shape of the distribution.
Returns#
jax.numpy.ndarray
The square root of the (optionally log-transformed) normalized DLM distribution function evaluated on the velocity grid defined by self.vx.
Notes#
The function computes the DLM distribution on a 2D velocity grid using the parameter m.
The distribution is normalized such that its sum over the grid equals one.
If self.learn_log is True, the function returns the negative base-10 logarithm of the distribution before taking the square root.
- init_bidlm(m, masym, tasym, theta)Source#
Initialize the distribution function using a 2D version of the Dum-Langdon-Matte (DLM) form. This formulation can have different widths and super-gaussian orders in the two dimensions. Parameters ———- m : float
The super-gaussian order parameter for the DLM, controlling the shape of the distribution.
Returns#
- jax.numpy.ndarray
The square root of the (optionally log-transformed) normalized DLM distribution function evaluated on the velocity grid defined by self.vx.
Notes#
The function computes the DLM distribution on a 2D velocity grid using the parameter m.
The distribution is normalized such that its sum over the grid equals one.
If self.learn_log is True, the function returns the negative base-10 logarithm of the distribution before taking the square root.
The distribution is rotated by an angle theta and has different widths in the x and y directions, controlled by tasym.
The distribution does not reduce to Matte (i.e. using masym = 1 is not the same as using regular dlm) due to being in cartesian not spherical, however the moments are correct
- __call__()Source#
Evaluates the normalized distribution function. This method computes the squared values of self.fval, optionally applies a logarithmic transformation if self.learn_log is True, and then normalizes the result so that the sum over the grid defined by self.vx integrates to one. :returns: jnp.ndarray – The normalized distribution function evaluated on the grid.
- tsadar.core.modules.distribution_functions.base.get_distribution_filter_spec(filter_spec: Dict, dist_params: Dict, replace: str | bool = True) DictSource#
Generates a filter for seperating trainable parameters in a distribution function from static parameters, based on the distribution type and parameters. This function modifies the filter_spec dictionary to indicate which parameters of the electron distribution functions are trainable, depending on the type of distribution specified in dist_params. It supports several distribution types, including ‘dlm’, ‘mx’, ‘arbitrary’, ‘arbitrary-nn’, and ‘sphericalharmonic’. :param filter_spec: The filter specification dictionary, typically representing the structure of the model or distribution functions. :type filter_spec: Dict :param dist_params: Dictionary containing distribution parameters, including the ‘type’ key and, for some types, additional nested parameters. :type dist_params: Dict
- Returns:
Dict – The updated filter specification dictionary, with trainable parameters marked according to the distribution type.
- Raises:
Warning – If the distribution type is ‘mx’ (Maxwellian), indicating no trainable parameters.
NotImplementedError – If the distribution type or a specific configuration is not supported.
- tsadar.core.modules.distribution_functions.base.update_distribution_layers(filter_spec, df)Source#
Updates the filter_spec tree by replacing the weights and biases of each layer in the distribution function neural network. :param filter_spec: The filter specification tree to be updated, typically used with Equinox (eqx) models. :param df: An object containing the distribution function neural network (df.f_nn), which is expected to have a ‘layers’ attribute. Each layer should have ‘linear.weight’ and ‘linear.bias’ attributes.
- Returns:
The updated filter_spec tree with the weights and biases of each layer replaced as specified.
Note
This function assumes that each layer in df.f_nn.layers has ‘linear.weight’ and ‘linear.bias’ attributes, and that eqx.tree_at is used for functional updates.