solprop.backends.mlx_gpu.backend#

The mlx backend: the single-precision oracle’s step, on the Metal GPU.

MlxBackend keeps everything the single-precision oracle (NumpySingleBackend) does on the host – admission, the narrowed tables and their exponents, the state scale, the stage bundles, landing, statistics from the widened state, metadata – and moves the state and every per-step pass to the device: the stage combinations, the double-float propagators and prefactors (DoubleFloatTerms), the transforms, the synthesis and projection products, the responses, the finish and the error norm, reduced on the device to a few scalars per trial.

It admits the fibre route – mode-averaged and multimode capillaries, solid-core and GNLSE fibres on field and envelope grids, with Kerr, the envelope grid’s harmonic, scalar and tensor Raman, gas mixtures, the silica bundle, the field-grid plasma with its held current, loss, both temporal absorber schemes, the polarisation taper and device-ready gradients – and radial and Cartesian free space, including a line, with their transverse collars, stepped by Runge-Kutta. Everything else is refused by name.

Functions

refuse_unsupported(problem)

Refuse, by name, what the mlx backend does not run in this phase.

Classes

`DeviceSources`(polarisation, current)	A chunk's windowed polarisation and ungated current, device arrays.
`MlxBackend`(problem, *[, grid, modes, ...])	Propagate a single-precision problem on an Apple-silicon GPU.
`MlxLowStorage`(backend)	The low-storage RK4IP engine on the device (SEP 24 section 6.5).

class solprop.backends.mlx_gpu.backend.DeviceSources(polarisation, current)[source]#

Bases: object

A chunk’s windowed polarisation and ungated current, device arrays.

Attributes:

polarisationmlx.core.array: The windowed polarisation, divided by s.
currentmlx.core.array or None: The ungated plasma current, divided by s, or None without plasma.

class solprop.backends.mlx_gpu.backend.MlxBackend(problem, *, grid=None, modes=None, profiles=None, tdse_binding=None)[source]#

Bases: NumpySingleBackend

Propagate a single-precision problem on an Apple-silicon GPU.

Parameters:

problemsolprop.translate.problem.KernelProblem: A validated single-precision fibre, radial or Cartesian problem.
gridsolprop.grid.TimeGrid, optional: The grid, needed only by stats.
modescollections.abc.Sequence, optional: The propagated modes, needed only by the zdw statistic.
profilescollections.abc.Mapping, optional: The z-profiles keyed by statistic name.
tdse_bindingsolprop.translate.tdse.TDSEBinding or None, optional: Must be None.

Raises:

BackendError: Naming what is refused, or a protected operation the device does not establish.

memory_ledger()[source]#

Return the reference’s ledger with MLX’s device memory beside it.

Returns:

dict: The oracle’s rows, mlx_device_active and mlx_device_peak (MLX’s own counters, measured) and mlx_device_cache.

metadata()[source]#

Return the backend’s metadata with MLX’s version, device and GEMM probe.

Returns:

dict: The reference’s metadata and the mlx entries.

name = 'mlx'#

set_initial(field, noise, *, z=0.0)[source]#

Load a state onto the device at the launch’s scale.

Parameters:

fieldnumpy.ndarray: The physical state in the canonical layout.
noiseobject: Must be None.
zfloat, optional: The position, m.

class solprop.backends.mlx_gpu.backend.MlxLowStorage(backend)[source]#

Bases: object

The low-storage RK4IP engine on the device (SEP 24 section 6.5).

The host engine’s four register passes (HostLowStorage) with device arrays; the half-step flows are formed every step by the double-float kernel, which agrees with the host’s tabulated flows to rounding.

Parameters:

backendMlxBackend: The backend whose state is the register A.

atomic = False#

close()[source]#

Hold nothing between steps.

held()[source]#

Return the bytes held between steps: none, the registers are transients.

Returns:

dict: Zero rows, as the host engine names them.

reset()[source]#

Forget the last step’s indicator.

step(z, dz)[source]#

Take one step, leaving the raw accepted state in the backend.

Parameters:

zfloat: Where the step starts, m.
dzfloat: The step, m, positive.

Returns:

float: The accepted state’s squared norm over the norm mask.

solprop.backends.mlx_gpu.backend.refuse_unsupported(problem)[source]#

Refuse, by name, what the mlx backend does not run in this phase.

Parameters:

problemsolprop.translate.problem.KernelProblem: The problem.

Raises:

BackendError: Naming the refusal and the backend that runs it.