solprop.backends.mlx_gpu.passes

The passes of one stage on the device, in binary32: the oracle’s, in MLX.

Each function follows the single-precision oracle’s pass of the same name (solprop.backends.numpy_ref.single and solprop.backends.numpy_ref.passes) and is compared with it, class R, by tests/backends/test_mlx.py. The running integrals use the oracle’s doubling order exactly; the transforms are MLX’s FFT, the matrix products MLX’s float32 GEMM (the backend has probed it), and the remaining element-wise work MLX’s own operations. Every array is float32 or complex64.

Functions

analytic(values)	Return the analytic signal of a real array along its last axis, in complex64.
kerr(field, coefficients, factor, carrier)	Evaluate the Kerr polarisation, the oracle's kerr_single.
plasma_current(field, tables, constants, ...)	Return the field-grid plasma's ungated current, the oracle's.
project(values, columns, taper, n_pol)	Project node values back onto the modes, P / tau.
project_many(values, columns, taper, n_pol)	Project several node arrays onto the modes in one matrix product.
raman(field, spectrum, factor, *, kind, thg)	Evaluate the delayed Raman polarisation, the oracle's raman_single.
running_integral(values, dt, rule, *[, evaluate])	Return a running integral, the oracle's running_integral.
synthesise(modal_t, rows, taper, n_pol)	Synthesise the field on a chunk of nodes, tau S a.
to_frequency(values, crop, *, kind)	Transform to frequency and gather the coarse bins.
to_time(spectrum, crop, n_to, *, kind)	Pad onto the fine grid and transform every modal column to time.

Classes

DeviceRateTable(coefficients, constant_tail, ...)

An ionisation rate table on the device.

class solprop.backends.mlx_gpu.passes.DeviceRateTable(coefficients, constant_tail, inverse_edges, inverse_tail, log_field_min, d_log_field, field_min, field_max, rate_max)

Bases: object

An ionisation rate table on the device.

The spline of solprop.translate.ionisation.RateTable in ln E, with its local coordinate formed as log1p(E / E_i - 1) against the narrowed left edge E_i of its interval: a binary32 ln E would carry an error of |ln E| u, which the rate’s steep slope multiplies. The polynomial’s constant term, ln W at the interval’s edge (of order 70), is held as a double-float pair and added to the small remainder by an exact two-sum, and the rate is exp(hi) (1 + lo): a binary32 constant term would cost |ln W| u relative in the rate.

Attributes:

coefficientsmlx.core.array

(4, n_intervals) float32; the last row is the constant term’s head.

constant_tailmlx.core.array

The constant term’s tail, float32.

inverse_edgesmlx.core.array

1 / E_i per interval, float32.

inverse_tailmlx.core.array

The tail of 1 / E_i as a double-float pair, float32.

log_field_minfloat

ln E_0.

d_log_fieldfloat

The interval width in ln E.

field_minfloat

Below it the rate is zero, V/m.

field_maxfloat

Above it the rate holds rate_max, V/m.

rate_maxfloat

s^-1.

coefficients: mlx.core.array

constant_tail: mlx.core.array

d_log_field: float

field_max: float

field_min: float

inverse_edges: mlx.core.array

inverse_tail: mlx.core.array

log_field_min: float

rate_max: float

classmethod upload(table)

Narrow a rate table onto the device.

Parameters:

tablesolprop.translate.ionisation.RateTable

The compiled table.

Returns:

DeviceRateTable

The device table.

solprop.backends.mlx_gpu.passes.analytic(values)

Return the analytic signal of a real array along its last axis, in complex64.

Parameters:

valuesmlx.core.array

Float32 values.

Returns:

mlx.core.array

Complex64.

solprop.backends.mlx_gpu.passes.kerr(field, coefficients, factor, carrier)

Evaluate the Kerr polarisation, the oracle’s kerr_single.

Parameters:

fieldmlx.core.array

The scaled field at the nodes.

coefficientssolprop.translate.responses.KerrCoefficients

The compiled coefficients.

factorfloat

coefficient * density * s**2, rounded once.

carriermlx.core.array or None

The harmonic carrier on an envelope grid that radiates it.

Returns:

mlx.core.array

The polarisation divided by s.

solprop.backends.mlx_gpu.passes.plasma_current(field, tables, constants, initials, dt, rule, *, evaluate=False)

Return the field-grid plasma’s ungated current, the oracle’s.

Parameters:

fieldmlx.core.array

The physical field s E, float32, (n_nodes, n_pol, n_to).

tablestuple of DeviceRateTable

One per channel.

constantstuple of tuple of float

Per channel, density (e^2/m) n_c and density I_p.

initialstuple of float

Per channel, the pre-ionised fraction.

dtfloat

The sample spacing, s.

rulesolprop.nonlinear.ionisation.RunningIntegral

The running-integral rule.

evaluatebool, optional

Whether each running integral is evaluated as it is scanned.

Returns:

mlx.core.array

The current density, A/m^2.

solprop.backends.mlx_gpu.passes.project(values, columns, taper, n_pol)

Project node values back onto the modes, P / tau.

Parameters:

valuesmlx.core.array

(n_nodes_chunk, n_pol_node, n_to).

columnsmlx.core.array

The chunk’s columns of the narrowed projection matrix.

taperfloat

The taper scale.

n_polint

Components the transform resolves at each node.

Returns:

mlx.core.array

(n_modes, n_pol_state, n_to).

solprop.backends.mlx_gpu.passes.project_many(values, columns, taper, n_pol)

Project several node arrays onto the modes in one matrix product.

The arrays’ real parts (a complex array’s real and imaginary parts) are placed side by side along the time axis, so that the polarisation and the plasma’s two current streams take one launch rather than four; every column of the product is the one project() forms, by the same inner products.

Parameters:

valueslist of mlx.core.array

(n_nodes_chunk, n_pol_node, n_to) each, float32 or complex64.

columnsmlx.core.array

The chunk’s columns of the narrowed projection matrix.

taperfloat

The taper scale.

n_polint

Components the transform resolves at each node.

Returns:

list of mlx.core.array

Each array projected, (n_modes, n_pol_state, n_to), in its own dtype.

solprop.backends.mlx_gpu.passes.raman(field, spectrum, factor, *, kind, thg)

Evaluate the delayed Raman polarisation, the oracle’s raman_single.

Parameters:

fieldmlx.core.array

The scaled field, (n_nodes, n_pol, n_to).

spectrummlx.core.array

The narrowed kernel spectrum: (n_to + 1,) scalar, (1 | 2, n_to + 1) tensor.

factorfloat

The combined coefficient, rounded once.

kind{‘field’, ‘envelope’}

The spectral layout.

thgbool

Whether the drive keeps the harmonic.

Returns:

mlx.core.array

The polarisation divided by s.

solprop.backends.mlx_gpu.passes.running_integral(values, dt, rule, *, evaluate=False)

Return a running integral, the oracle’s running_integral.

Parameters:

valuesmlx.core.array

The integrand along the last axis, float32.

dtfloat

The sample spacing, s.

rulesolprop.nonlinear.ionisation.RunningIntegral

The problem’s rule.

evaluatebool, optional

Whether to evaluate the integral at once, releasing the scan’s levels.

Returns:

mlx.core.array

The running integral, zero at the first sample.

solprop.backends.mlx_gpu.passes.synthesise(modal_t, rows, taper, n_pol)

Synthesise the field on a chunk of nodes, tau S a.

Parameters:

modal_tmlx.core.array

(n_modes, n_pol_state, n_to).

rowsmlx.core.array

The chunk’s rows of the narrowed synthesis matrix.

taperfloat

The taper scale at this position.

n_polint

Components the transform resolves at each node.

Returns:

mlx.core.array

(n_nodes_chunk, n_pol * n_pol_state, n_to).

solprop.backends.mlx_gpu.passes.to_frequency(values, crop, *, kind)

Transform to frequency and gather the coarse bins.

Parameters:

valuesmlx.core.array

(n_modes, n_pol, n_to).

cropmlx.core.array

The coarse bins’ positions.

kind{‘field’, ‘envelope’}

The spectral layout.

Returns:

mlx.core.array

(n_modes, n_pol, n_omega) complex64.

solprop.backends.mlx_gpu.passes.to_time(spectrum, crop, n_to, *, kind)

Pad onto the fine grid and transform every modal column to time.

Parameters:

spectrummlx.core.array

(n_modes, n_pol, n_omega) complex64.

cropmlx.core.array or None

Where each coarse bin sits in the fine spectrum; None when the coarse bins are the fine spectrum’s first ones, which is padded rather than scattered.

n_toint

Fine time samples.

kind{‘field’, ‘envelope’}

The spectral layout.

Returns:

mlx.core.array

(n_modes, n_pol, n_to), float32 on a field grid, complex64 otherwise.