Form Factor API

This is the module that handles the loss function, gradient calculation, hessian calculation, and parameter generation

class tsadar.core.physics.form_factor.FormFactor(lambda_range, npts, lam_shift, scattering_angles, num_grad_points, ud_ang, va_ang, calc_gain)Source

FormFactor class for calculating the Thomson scattering structure factor or spectral density function. This class encapsulates all static values and methods required for repeated calculations of the Thomson scattering structure factor or spectral density function, supporting both 1D and 2D electron distribution functions (EDFs), multiple plasma conditions, and scattering angles. :param lambda_range: Starting and ending wavelengths over which to calculate the spectrum. :type lambda_range: list :param npts: Number of wavelength points to use in the calculation. :type npts: int :param lam_shift: Wavelength shift to apply. :type lam_shift: float :param scattering_angles: Dictionary containing scattering angles (in degrees). :type scattering_angles: dict :param num_grad_points: Number of gradient points for plasma parameter profiles. :type num_grad_points: int :param ud_ang: Angle between electron drift and x-axis (degrees). :type ud_ang: float :param va_ang: Angle between ion flow and x-axis (degrees). :type va_ang: float

C

Speed of light in cm/s.

Type:

float

Me

Electron mass in keV/C^2.

Type:

float

Mp

Proton mass in keV/C^2.

Type:

float

npts

Number of wavelength points.

Type:

int

h

Step size for velocity grid.

Type:

float

xi1, xi2

Grids for velocity integration.

Type:

jnp.ndarray

Zpi

Precomputed plasma dispersion function values.

Type:

jnp.ndarray

lam_shift

Wavelength shift.

Type:

float

scattering_angles

Scattering angles.

Type:

dict

num_grad_points

Number of gradient points.

Type:

int

vmap_calc_chi_vals

Vectorized susceptibility calculation.

Type:

callable

ud_angle, va_angle

Electron drift and ion flow angles.

Type:

float

calc_all_chi_vals

Method for calculating susceptibility, possibly parallelized.

Type:

callable

__call__(params)Source

Calculates the standard collisionless Thomson spectral density function S(k,omg) for 1D EDFs.

params (dict): Plasma and distribution function parameters. formfactor (jnp.ndarray): Calculated spectrum. lams (jnp.ndarray): Wavelength axis.

rotate(vx, df, angle, reshape=False)Source

Rotates a 2D array by a given angle in radians.

vx (jnp.ndarray): Velocity grid. df (jnp.ndarray): 2D distribution function. angle (float): Rotation angle in radians. reshape (bool): Whether to reshape the output. jnp.ndarray: Rotated/interpolated 2D array.

scan_calc_chi_vals(carry, xs)Source

Calculates susceptibility values at a given point in the distribution function using scan.

carry (tuple): (velocity grid, 2D distribution function). xs (tuple): (angle, xie_mag_at, klde_mag_at). tuple: Updated carry and (fe_vphi, chiEI, chiERrat).

calc_chi_vals(vx, DF, inputs)Source

Calculates susceptibility values at a given point in the distribution function.

vx (jnp.ndarray): Velocity grid. DF (jnp.ndarray): 2D distribution function. inputs (tuple): (angle, xie_mag_at, klde_mag_at). tuple: (fe_vphi, chiEI, chiERrat).

_calc_all_chi_vals_(vx, DF, beta, xie_mag, klde_mag)Source

Calculates susceptibility values for all desired points xie (batch or vectorized).

vx (jnp.ndarray): Velocity grid. DF (jnp.ndarray): 2D distribution function. beta (jnp.ndarray): Angles. xie_mag (jnp.ndarray): Magnitudes of normalized velocity points. klde_mag (jnp.ndarray): Magnitudes of wavevector times Debye length. tuple: (fe_vphi, chiEI, chiERrat).

parallel_calc_all_chi_vals(x, DF, beta, xie_mag, klde_mag)Source

Parallelized calculation of susceptibility values across devices.

x (jnp.ndarray): Velocity grid. DF (jnp.ndarray): 2D distribution function. beta, xie_mag, klde_mag (jnp.ndarray): Parameters for susceptibility calculation. tuple: (fe_vphi, chiEI, chiERrat).

calc_in_2D(params)Source

Calculates the collisionless Thomson spectral density function S(k,omg) for a 2D numerical EDF.

params (dict): Plasma and distribution function parameters. formfactor (jnp.ndarray): Calculated spectrum. lams (jnp.ndarray): Wavelength axis.

__call__(params)Source

Calculates the standard collisionless Thomson spectral density function S(k,omg) and is capable of handling multiple plasma conditions and scattering angles. Distribution functions can be arbitrary as calculations of the susceptibility is done on-the-fly. Calculations are done in 4 dimension with the following shape, [number of gradient-points, number of wavelength points, number of angles, number of ion-species].

In angular, fe is a Tuple, Distribution function (DF), normalized velocity (x), and angles from k_L to f1 in radians

Parameters:

params – ThomsonParams object, contains all the parameters from the input deck

Returns:

formfactor –

array of the calculated spectrum, has the shape [number of gradient-points, number of

wavelength points, number of angles]

lams: wavelength axis

rotate(vx, df, angle, reshape: bool = False) → ArraySource

Rotate a 2D array by a specified angle in radians. This method rotates the input 2D array df using a rotation matrix constructed from the given angle. The rotation is performed around the origin, and the rotated coordinates are interpolated back onto the original grid using cubic interpolation.

vx (jnp.ndarray): 1D array representing the grid points along each axis. df (jnp.ndarray): 2D array to be rotated. angle (float): Rotation angle in radians (counterclockwise). reshape (bool, optional): Whether to reshape the output array. Defaults to False.

Returns:

jnp.ndarray – The rotated and interpolated 2D array.

scan_calc_chi_vals(carry, xs)Source

Calculate the values of the susceptibility at a given point in the distribution function

Parameters:

• carry –

  container for

  x: 1D array DF: 2D array

• xs –

  container for

  element: angle in radians xie_mag_at: float klde_mag_at: float

Returns:

fe_vphi – float, value of the projected distribution function at the point xie chiEI: float, value of the imaginary part of the electron susceptibility at the point xie chiERrat: float, value of the real part of the electron susceptibility at the point xie

calc_chi_vals(vx, DF, inputs)Source

Calculate the values of the susceptibility at a given point in the distribution function

Parameters:

carry –

container for

x: 1D array DF: 2D array inputs: container for

element: angle in radians xie_mag_at: float klde_mag_at: float

Returns:

fe_vphi – float, value of the projected distribution function at the point xie chiEI: float, value of the imaginary part of the electron susceptibility at the point xie chiERrat: float, value of the real part of the electron susceptibility at the point xie

_calc_all_chi_vals_(vx, DF, beta, xie_mag, klde_mag)Source

Calculate the susceptibility values for all the desired points xie

Parameters:

• x – normalized velocity grid

• beta – angle of the k-vector form the x-axis

• DF – 2D array, distribution function

• xie_mag – magnitude of the normalized velocity points where the calculations need to be performed

• klde_mag – magnitude of the wavevector time debye length where the calculations need to be performed

Returns:

fe_vphi – projected distribution function chiEI: imaginary part of the electron susceptibility chiERrat: real part of the electron susceptibility

parallel_calc_all_chi_vals(x, DF, beta, xie_mag, klde_mag)Source

calc_in_2D(params)Source

Calculates the collisionless Thomson spectral density function S(k,omg) for a 2D numerical EDF, capable of handling multiple plasma conditions and scattering angles. Distribution functions can be arbitrary as calculations of the susceptibility are done on-the-fly. Calculations are done in 4 dimension with the following shape, [number of gradient-points, number of wavelength points, number of angles, number of ion-species].

In angular, fe is a Tuple, Distribution function (DF), normalized velocity (x), and angles from k_L to f1 in radians

Parameters:

params – ThomsonParams object, contains all the parameters from the input deck

Returns:

formfactor –

array of the calculated spectrum, has the shape [number of gradient-points, number of

wavelength points, number of angles]

lams: wavelength axis

tsadar.core.physics.form_factor.zprimeMaxw(xi)Source

Calculates the derivative of the plasma dispersion function (Z-prime) for an array of normalized phase velocities (xi) using a combination of tabulated values and asymptotic approximations. For values of xi between -10 and 10, the function uses interpolated data from precomputed tables. For values outside this range, it applies the asymptotic approximation as described in Eqn. 5.2.10 of the Thomson scattering reference. :param xi: Array of normalized phase velocities (must be in ascending order). :type xi: np.ndarray

Returns:

Zp (np.ndarray) – 2D array where the first row contains the real components and the second row contains the imaginary components of Z-prime evaluated at each value of xi.