solprop.backends.numba_cpu.kernels#

The fused stage kernels: seven numbered passes in a handful of traversals.

docs/dev/architecture.md section 8 names seven passes and the reference backend performs each of them as its own traversal of the state. Every one of those traversals reads a field-sized array from memory and writes another, and at the sizes this package targets none of them fits a cache. The kernels do the same arithmetic in the same order with the intermediate arrays never formed: the stage combination, the lab frame, the noise and the pad are one traversal; the whole nonlinear term at one node – the rate lookup, the plasma column’s two running sums, the loss-current guard, the Kerr cubic and the Raman drive – is one traversal of t; and the crop, the prefactor, the windows, the return to the interaction picture and both accumulations are one more.

Specialisation is by closure constant, never by argument. A branch in the innermost loop is not free: the Raman flag was an argument first and the runtime test cost the RDW step 38 per cent (3.79 ms against 5.24). Every flag a kernel needs is bound when the factory is called, so numba folds the branch away and a problem without a term pays nothing for it.

Reductions are two-stage. A kernel writes one partial per column and the driver sums them; the order of either sum is free (SEP 25: a norm is held to its summation bound, not to bits). The column loop is a prange in the parallel variants, which a state of many columns takes (ACCEPT_ELEMENTS). Reassociating the masked sums within a column was measured and gained nothing: the mask’s select keeps the loop scalar, and multiplying by the mask instead would turn a non-finite value outside the band into a NaN norm.

The floating-point licence is SEP 25’s, applied by solprop._numba_cache.njit: every kernel contracts multiply-adds; the protected bodies (define_bodies’s protect) are compiled strict and called, never inlined. Nothing in here allocates, logs, raises or touches a Python object. numba.prange falls back to range when a function is compiled without parallel=True, so the serial and parallel variants of docs/dev/backends.md section 3.2 come from one body rather than from a copy of it.

Every kernel carries @no_type_check: the test session’s jaxtyping hook appends its checker innermost, and njit would otherwise compile the checker’s wrapper.

Module Attributes

`HOST`	The compiled body family every kernel inlines, and the protected bodies it calls.
`ACCEPT_ELEMENTS`	State size from which the accept pass and the weighted residual thread their columns.
`ACCEPT_COLUMNS`	Columns the parallel accept pass needs: with fewer the threads have too little.
`accept_and_partials_parallel`(state, ...)	`accept_and_partials()` with its columns threaded (`ACCEPT_ELEMENTS`).
`weighted_residual_parallel`(propagator, dz, ...)	`weighted_residual()` with its columns threaded (`ACCEPT_ELEMENTS`).
`BAND_DIAGNOSTICS`	Rows the band statistics column appends after the channels' two each: the largest exposure dip and rise, the largest `\|paid - freed\|` and `freed`, and the largest `\|f E\|` in the first `2K` samples and anywhere, per node (`solprop.backends.numpy_ref.passes.plasma_integration_diagnostics()`).
`BAND_STATS_ROWS`	Rows of `n_to` doubles the band statistics column needs per node: the raw rate and the filtered rate's exposure.

Functions

`accept_and_partials`(state, propagator, dz, ...)	Take the accept pass and the error norm's partials in one traversal.
`accept_and_partials_parallel`(state, ...)	`accept_and_partials()` with its columns threaded (`ACCEPT_ELEMENTS`).
`accept_kernels`(threads, n_columns, n_omega)	Choose the serial or the column-threaded accept pass and weighted residual.
`add_temporal_absorber`(samples, rate, output)	Add the temporal rate to a precomputed spatial absorber term in place.
`apply_hard_mask`(state, mask, out)	Zero every bin outside the active band, into `out`.
`band_work_rows`(spec)	Return the rows of `n_to` doubles the multipass plasma column needs per node.
`column_spec`(responses, kind, n_pol, *, parallel)	Read the specialisation off a compiled response set.
`exponentiate`(exponent, forward, inverse)	Form `exp(+dPhi)` and `exp(-dPhi)` from one real `exp` and one `sincos`.
`make_band_plasma_column`(spec)	Compile the plasma column as whole-row passes (SEP 18's rule, SEP 25's form).
`make_band_stats_column`(spec)	Compile the plasma statistics of a plane under a running-integral rule.
`make_combine_and_pad`(*, contiguous, noise[, ...])	Compile pass 1: stage combine, lab frame, noise and pad, in one traversal.
`make_finish_stage`(*, contiguous[, absorber])	Compile passes 6 and 7: crop, prefactor, window, frame, both accumulations.
`make_raman_finish`(spec)	Compile the delayed Raman term and the temporal collar, in one traversal.
`make_response_column`(spec, *[, ...])	Compile passes 4 and 5: the whole nonlinear term, one traversal of `t`.
`make_stats_column`(spec)	Compile the plasma statistics of a plane: one traversal for both of them.
`peak_reduce`(columns, partials)	Reduce `max_t sum_c \|x_c(t)\|^2` over a set of columns.
`scale_samples`(samples, factor)	Multiply the coarse time samples by the absorber's collar, in place.
`scale_spectrum`(spectrum, kernel, scale)	Multiply the Raman drive's spectrum by the damped kernel, in place.
`scale_tensor_spectrum`(spectrum, kernel, scale)	Multiply interleaved tensor drives by trace or traceless kernel spectra.
`weighted_residual`(propagator, dz, error, ...)	Sum the error norm's residual with a per-bin scale, one partial per column.
`weighted_residual_parallel`(propagator, dz, ...)	`weighted_residual()` with its columns threaded (`ACCEPT_ELEMENTS`).
`write_temporal_absorber`(samples, rate, output)	Write the temporal absorber's time-domain term, `-rate*samples/2`.

Classes

ColumnSpec(family, plasma, raman, raman_thg, ...)

Which response column a problem compiles to, and how it is executed.

solprop.backends.numba_cpu.kernels.ACCEPT_COLUMNS = 8#: Columns the parallel accept pass needs: with fewer the threads have too little.

solprop.backends.numba_cpu.kernels.ACCEPT_ELEMENTS = 65536#

State size from which the accept pass and the weighted residual thread their columns.

Both run a column per prange iteration in their parallel variants, which a state of at least this many elements in at least ACCEPT_COLUMNS columns takes when the kernel pool has more than one thread. Measured on six threads of an Apple-silicon laptop under load: 8 x 8192 elements 104 -> 80 us, 64 x 2048 224 -> 107 us, 2048 x 1024 3.5 -> 0.96 ms; 4 x 8192 slower (50 -> 63 us). Each column’s sum is the serial one, so the variants agree exactly.

solprop.backends.numba_cpu.kernels.BAND_DIAGNOSTICS = 6#: Rows the band statistics column appends after the channels’ two each: the largest exposure dip and rise, the largest |paid - freed| and freed, and the largest |f E| in the first 2K samples and anywhere, per node (solprop.backends.numpy_ref.passes.plasma_integration_diagnostics()).

solprop.backends.numba_cpu.kernels.BAND_STATS_ROWS = 2#: Rows of n_to doubles the band statistics column needs per node: the raw rate and the filtered rate’s exposure.

class solprop.backends.numba_cpu.kernels.ColumnSpec(family, plasma, raman, raman_thg, harmonic, n_pol, parallel, stash=False, n_channels=1, n_raman=1, raman_drives=1, band=False, preloaded=False)[source]#

Bases: NamedTuple

Which response column a problem compiles to, and how it is executed.

Every field is a closure constant of the kernel the factory returns, so two problems that differ in any of them get two compiled kernels and neither pays a branch for the other.

Attributes:

familystr: "field_real", "field_analytic" or "envelope"; see solprop.backends.numba_cpu.fft.ColumnFamily.
plasmabool: Whether the ionisation rate, the plasma column and the loss current are evaluated. Only a field grid has them.
ramanbool: Whether the column writes the Raman drive out onto the doubled grid.
raman_thgbool: Whether the drive keeps the third harmonic (E^2 rather than |E_analytic|^2 / 2). Meaningless on an envelope grid, which cannot carry it in a quadratic drive.
harmonicbool: Whether the envelope’s third-harmonic Kerr term is added with the carrier.
n_polint: Transverse components at one node, 1 or 2. The Kerr sums run over them.
parallelbool: Whether the node loop is a prange.
stashbool: Whether the plasma branch leaves the final ionised fraction and the largest rate of each node behind for the between-planes statistics channel. A closure constant like every other flag, so a run without the channel compiles the two stores away and pays nothing.
raman_drivesint: One scalar or isotropic trace history, or three trace/difference/cross histories for a vector rank-two response.
bandbool: Whether the plasma column integrates by the band rule (SEP 18). The plasma then runs in make_band_plasma_column()’s passes, and the fused column is compiled without plasma and writes non-plasma polarisation separately.
preloadedbool: Whether the fused column starts each sample’s total from what out already holds: an existing non-plasma contribution rather than zero.

band: bool#: Alias for field number 11

family: str#: Alias for field number 0

harmonic: bool#: Alias for field number 4

n_channels: int#: Alias for field number 8

n_pol: int#: Alias for field number 5

n_raman: int#: Alias for field number 9

parallel: bool#: Alias for field number 6

plasma: bool#: Alias for field number 1

preloaded: bool#: Alias for field number 12

raman: bool#: Alias for field number 2

raman_drives: int#: Alias for field number 10

raman_thg: bool#: Alias for field number 3

stash: bool#: Alias for field number 7

solprop.backends.numba_cpu.kernels.HOST = Bodies(spatial_element=CPUDispatcher(<function define_bodies.<locals>.spatial_element>), slab_element=CPUDispatcher(<function define_bodies.<locals>.slab_element>), longitudinal_rate=CPUDispatcher(<function define_bodies.<locals>.longitudinal_rate>), longitudinal_root=CPUDispatcher(<function define_bodies.<locals>.longitudinal_root>), mask_plane_radius=CPUDispatcher(<function define_bodies.<locals>.mask_plane_radius>), aperture_weight=CPUDispatcher(<function define_bodies.<locals>.aperture_weight>), region_member=CPUDispatcher(<function define_bodies.<locals>.region_member>), propagator_pair=CPUDispatcher(<function define_bodies.<locals>.propagator_pair>), rate_at=CPUDispatcher(<function define_bodies.<locals>.rate_at>), magnitude_pair=CPUDispatcher(<function define_bodies.<locals>.magnitude_pair>), kerr_real=CPUDispatcher(<function define_bodies.<locals>.kerr_real>), kerr_analytic=CPUDispatcher(<function define_bodies.<locals>.kerr_analytic>), dispersive_isotropic=CPUDispatcher(<function define_bodies.<locals>.dispersive_isotropic>), dispersive_conserving=CPUDispatcher(<function define_bodies.<locals>.dispersive_conserving>), dispersive_harmonic=CPUDispatcher(<function define_bodies.<locals>.dispersive_harmonic>), dispersive_isotropic3=CPUDispatcher(<function define_bodies.<locals>.dispersive_isotropic3>), dispersive_conserving3=CPUDispatcher(<function define_bodies.<locals>.dispersive_conserving3>), dispersive_real_self=CPUDispatcher(<function define_bodies.<locals>.dispersive_real_self>), dispersive_real_pair=CPUDispatcher(<function define_bodies.<locals>.dispersive_real_pair>), dispersive_real_cross=CPUDispatcher(<function define_bodies.<locals>.dispersive_real_cross>), dispersive_conserving_self=CPUDispatcher(<function define_bodies.<locals>.dispersive_conserving_self>), dispersive_conserving_pair=CPUDispatcher(<function define_bodies.<locals>.dispersive_conserving_pair>), dispersive_conserving_cross=CPUDispatcher(<function define_bodies.<locals>.dispersive_conserving_cross>), dispersive_harmonic_self=CPUDispatcher(<function define_bodies.<locals>.dispersive_harmonic_self>), dispersive_harmonic_pair=CPUDispatcher(<function define_bodies.<locals>.dispersive_harmonic_pair>), dispersive_harmonic_cross=CPUDispatcher(<function define_bodies.<locals>.dispersive_harmonic_cross>), held_plasma_source=CPUDispatcher(<function define_bodies.<locals>.held_plasma_source>), loss_scale=CPUDispatcher(<function define_bodies.<locals>.loss_scale>), trapezoid=CPUDispatcher(<function define_bodies.<locals>.trapezoid>), fraction_at=CPUDispatcher(<function define_bodies.<locals>.fraction_at>), accumulate=CPUDispatcher(<function define_bodies.<locals>.accumulate>), low_storage_fold=CPUDispatcher(<function define_bodies.<locals>.low_storage_fold>), raman_drive=CPUDispatcher(<function define_bodies.<locals>.raman_drive>), raman_tensor_drive=CPUDispatcher(<function define_bodies.<locals>.raman_tensor_drive>), raman_damped=CPUDispatcher(<function define_bodies.<locals>.raman_damped>), raman_finish=CPUDispatcher(<function define_bodies.<locals>.raman_finish>), raman_tensor_finish=CPUDispatcher(<function define_bodies.<locals>.raman_tensor_finish>), absorber_sample=CPUDispatcher(<function define_bodies.<locals>.absorber_sample>), envelope_population=CPUDispatcher(<function define_bodies.<locals>.envelope_population>), envelope_axes=CPUDispatcher(<function define_bodies.<locals>.envelope_axes>), envelope_moments=CPUDispatcher(<function define_bodies.<locals>.envelope_moments>), envelope_loss=CPUDispatcher(<function define_bodies.<locals>.envelope_loss>), general_kerr_component=CPUDispatcher(<function define_bodies.<locals>.general_kerr_component>), general_kerr=CPUDispatcher(<function define_bodies.<locals>.general_kerr>), general_symmetric_row=CPUDispatcher(<function define_bodies.<locals>.general_symmetric_row>), general_symmetric_sum=CPUDispatcher(<function define_bodies.<locals>.general_symmetric_sum>), general_real_kerr=CPUDispatcher(<function define_bodies.<locals>.general_real_kerr>), general_pair_drive=CPUDispatcher(<function define_bodies.<locals>.general_pair_drive>), general_delayed=CPUDispatcher(<function define_bodies.<locals>.general_delayed>))#

The compiled body family every kernel inlines, and the protected bodies it calls.

Not disk-cached: a kernel inlines these bodies’ IR and caches the result, and a body that captures another would pickle differently in every process, so its own entry would never be read back. The protected bodies are compiled strict, on their own, and called (SEP 25 section 5).

solprop.backends.numba_cpu.kernels.accept_and_partials(state, propagator, dz, solution, error, mask, advanced, partials)[source]#

Take the accept pass and the error norm’s partials in one traversal.

partials holds three sums per column – of |y|^2, |y_new|^2 and |propagator (h e)|^2 over the in-band mask – and the driver sums them over the columns. The column loop is a prange: accept_and_partials_parallel() is the same body compiled parallel.

Parameters:

statenumpy.ndarray: The state at the start of the step, complex, shape (n_columns, n_omega).
propagatornumpy.ndarray: exp(dPhi) over the whole step, same shape.
dzfloat: The step, m.
solutionnumpy.ndarray: The weighted stage sum of the solution, same shape.
errornumpy.ndarray: The weighted stage sum of the embedded difference, same shape.
masknumpy.ndarray: The in-band mask the norm is taken over, boolean, shape (n_omega,).
advancednumpy.ndarray: Written with the state at the end of the step.
partialsnumpy.ndarray: Written with the three sums, shape (3, n_columns).

solprop.backends.numba_cpu.kernels.accept_and_partials_parallel(state, propagator, dz, solution, error, mask, advanced, partials)#

accept_and_partials() with its columns threaded (ACCEPT_ELEMENTS).

solprop.backends.numba_cpu.kernels.accept_kernels(threads, n_columns, n_omega)[source]#

Choose the serial or the column-threaded accept pass and weighted residual.

Parameters:

threadsint: Threads the kernel pool has.
n_columnsint: Columns of the state the pass traverses.
n_omegaint: Bins per column.

Returns:

tuple: accept_and_partials() and weighted_residual(), or their parallel variants when the state is large enough to share (ACCEPT_ELEMENTS).

Examples

>>> from solprop.backends.numba_cpu import kernels
>>> kernels.accept_kernels(6, 4, 8192)[0] is kernels.accept_and_partials
True

solprop.backends.numba_cpu.kernels.add_temporal_absorber(samples, rate, output)[source]#

Add the temporal rate to a precomputed spatial absorber term in place.

Parameters:

samplesnumpy.ndarray: Fine temporal samples, shape (n_columns, n_to).
ratenumpy.ndarray: Power absorption rates in m^-1, shape (n_to,).
outputnumpy.ndarray: Existing derivative term; receives -rate*samples/2.

solprop.backends.numba_cpu.kernels.apply_hard_mask(state, mask, out)[source]#

Zero every bin outside the active band, into out.

The active-band mask, applied once per accepted step. It is deliberately not part of a stage’s pass 7: the reference applies it in its own finishing pass, and a fused pass 7 that added it would mask the stage derivatives as well as the state.

Parameters:

statenumpy.ndarray: The state, complex, shape (n_columns, n_omega).
masknumpy.ndarray: The active band, boolean, shape (n_omega,).
outnumpy.ndarray: Written with the masked state. May be state itself.

solprop.backends.numba_cpu.kernels.band_work_rows(spec)[source]#

Return the rows of n_to doubles the multipass plasma column needs per node.

The raw rate, the weight f q/m (summed over channels when there are several), the loss guard, the free-current integrand, and the filtered integrand a scan reads. The total current is written directly to the caller’s output. The components and the channels are processed one after another, so the count does not grow with them.

Parameters:

specColumnSpec: The specialisation.

Returns:

int: Five.

Examples

>>> from solprop.backends.numba_cpu.kernels import ColumnSpec, band_work_rows
>>> band_work_rows(ColumnSpec("field_real", True, False, False, False, 2, False))
5

solprop.backends.numba_cpu.kernels.column_spec(responses, kind, n_pol, *, parallel, stash=False)[source]#

Read the specialisation off a compiled response set.

One place decides which kernel a problem gets, so the backend and the tests cannot disagree about it.

Parameters:

responsessolprop.translate.responses.CompiledResponses: The compiled response set the contract carries.
kind{‘field’, ‘envelope’}: The spectral layout.
n_polint: Transverse components at one node, 1 or 2.
parallelbool: Whether the node loop is a prange, which the threading regime decides.
stashbool, optional: Whether the plasma branch leaves each node’s final ionised fraction and largest rate behind, which the between-planes statistics channel reads.

Returns:

ColumnSpec: The specialisation.

Examples

>>> from solprop.backends.numba_cpu.kernels import ColumnSpec
>>> ColumnSpec("envelope", False, False, False, True, 1, False).harmonic
True

solprop.backends.numba_cpu.kernels.exponentiate(exponent, forward, inverse)[source]#

Form exp(+dPhi) and exp(-dPhi) from one real exp and one sincos.

A complex exponential per sign, which is what the reference evaluates, costs the transcendental twice for numbers that share all their trigonometry.

Parameters:

exponentnumpy.ndarray: dPhi for one stage fraction, complex, shape (n_columns, n_omega).
forwardnumpy.ndarray: Written with exp(+dPhi), same shape and dtype.
inversenumpy.ndarray: Written with exp(-dPhi), same shape and dtype.

solprop.backends.numba_cpu.kernels.make_band_plasma_column(spec)[source]#

Compile the plasma column as whole-row passes (SEP 18’s rule, SEP 25’s form).

Each of the two running integrals is a filter of its integrand followed by a trapezoid scan, and each integrand is known only once the scan before it is done, so the column runs in passes along t, each a contiguous traversal of a row of scratch: the rate; the exposure (its filter and scan) with the fraction, the loss guard and the free-current weight; for each component, the free current (its integrand’s filter and scan) with the total current. The stencil reads the row the pass before it finished, so there is no ring to keep. A pass whose whole stencil reads zeros – before the rate first leaves zero, after it last does – writes exact zeros without reading. The exposure is floored at zero where the fraction is formed; the loss current keeps the raw rate. With the trapezoid’s one identity tap it is the trapezoid’s column.

The total current is written into out uncollared. The fused non-plasma column separately forms Kerr and Raman polarisation, with its usual collar. Source assembly consumes gJ and g’J; there is no third running integral.

Parameters:

specColumnSpec: The specialisation: field family, polarisation count, channels, stash and the node loop’s parallelism.

Returns:

numba.core.dispatcher.Dispatcher: The compiled kernel.

solprop.backends.numba_cpu.kernels.make_band_stats_column(spec)[source]#

Compile the plasma statistics of a plane under a running-integral rule.

The band rule’s first stage alone – the filtered exposure, floored, and the fraction – for the final fraction and the largest rate of every channel, with the six per-node parts of SEP 18’s run diagnostics beside them, in whole-row passes as make_band_plasma_column() runs it: the rate, its filter, the exposure’s scan with its increments’ extremes, then the fraction, the work paid and the integrand’s levels. Given the trapezoid’s one identity tap it is the plain column with the diagnostics, which is how a trapezoid run that asks for them is answered.

Parameters:

specColumnSpec: The specialisation: field family, polarisation and plasma channel counts.

Returns:

numba.core.dispatcher.Dispatcher: The compiled kernel.

solprop.backends.numba_cpu.kernels.make_combine_and_pad(*, contiguous, noise, parallel=False)[source]#

Compile pass 1: stage combine, lab frame, noise and pad, in one traversal.

Parameters:

contiguousbool: Whether the crop index is the leading block of the fine spectrum, which it is on a field grid and is not on an envelope grid, where the negative frequencies sit at the far end. The contiguous form writes the pad tail as one run.
noisebool: Whether the modified noise model’s realisation is added at every stage.
parallelbool, optional: Whether independent columns execute in parallel. The small modal path retains the serial default.

Returns:

numba.core.dispatcher.Dispatcher: The compiled kernel.

solprop.backends.numba_cpu.kernels.make_finish_stage(*, contiguous, absorber=False)[source]#

Compile passes 6 and 7: crop, prefactor, window, frame, both accumulations.

Parameters:

contiguousbool: Whether the crop index is the leading block of the fine spectrum.
absorberbool, optional: Whether the in-stage temporal absorber contributes a second term. It is a compile-time flag rather than a zero array so that a problem without a collar pays neither the read nor the add.

Returns:

numba.core.dispatcher.Dispatcher: The compiled kernel.

solprop.backends.numba_cpu.kernels.make_raman_finish(spec)[source]#

Compile the delayed Raman term and the temporal collar, in one traversal.

P_R = eps0 rho E(t) (h * D)(t), with the convolution already done on the doubled grid. The only pass the Raman term costs beyond its transform pair: the drive was written by the response column and the total is already in out.

Parameters:

specColumnSpec: Field family, physical component count and Raman history count. A field grid multiplies the delayed response matrix by the real field; an envelope grid multiplies it by the complex envelope.

Returns:

numba.core.dispatcher.Dispatcher: The compiled kernel.

solprop.backends.numba_cpu.kernels.make_response_column(spec, *, reuse_unchanged_fraction=False)[source]#

Compile passes 4 and 5: the whole nonlinear term, one traversal of t.

The kernel this backend exists for. At one node it reads the field once and forms the ionisation rate, the plasma column’s two running sums, the loss-current guard, the Kerr cubic and the Raman drive from the same registers, where the reference makes a separate pass over the chunk for each – two of them with scipy.integrate.cumulative_trapezoid(), which cannot be given a running sum.

Every flag is a closure constant. The Raman one was an argument first and the runtime test in the innermost loop cost the RDW step 38 per cent (3.79 ms against 5.24); numba folds a constant away and a problem without a term pays nothing.

Parameters:

specColumnSpec: The specialisation.
reuse_unchanged_fractionbool, optional: Reuse the one-channel fraction when the rounded exposure is unchanged. Sample zero always evaluates the original scalar body. The production caller selects the measured component and traversal domain. Multiple channels keep their independent original calculation.

Returns:

numba.core.dispatcher.Dispatcher: The compiled kernel.

solprop.backends.numba_cpu.kernels.make_stats_column(spec)[source]#

Compile the plasma statistics of a plane: one traversal for both of them.

electron_density and peak_ionisation_rate are read at the node of largest ionised fraction, so both need the rate and the exposure at every node – which is the response column’s first half and nothing else. The reference re-synthesises every node at once and then runs the multipass scan, which was measured at two thirds of a whole CtC step on an x86-64 server; a run with many save planes cannot afford that.

Parameters:

specColumnSpec: The specialisation: field family, polarisation and plasma channel counts.

Returns:

numba.core.dispatcher.Dispatcher: The compiled kernel.

solprop.backends.numba_cpu.kernels.peak_reduce(columns, partials)[source]#

Reduce max_t sum_c |x_c(t)|^2 over a set of columns.

The reduction behind the on-axis peak intensity. The sum over components is inside the maximum, so the answer is the largest total at any one instant and not the sum of separately timed maxima.

Parameters:

columnsnumpy.ndarray: The analytic field at one point, complex, shape (n_components, n_t).
partialsnumpy.ndarray: Written with the running maximum, shape (1,).

solprop.backends.numba_cpu.kernels.scale_samples(samples, factor)[source]#

Multiply the coarse time samples by the absorber’s collar, in place.

Parameters:

samplesnumpy.ndarray: The state in time on the coarse grid, shape (n_columns, n_t).
factornumpy.ndarray: exp(-rate dz / 2) on the coarse time axis, shape (n_t,).

solprop.backends.numba_cpu.kernels.scale_spectrum(spectrum, kernel, scale)[source]#

Multiply the Raman drive’s spectrum by the damped kernel, in place.

Parameters:

spectrumnumpy.ndarray: The drive’s one-sided spectrum on the doubled grid, complex, shape (n_rows, n_to + 1).
kernelnumpy.ndarray: The damped kernel’s spectrum at this position, shape (n_to + 1,).
scalefloat: The 1 / (2 n_to) that numpy.fft.irfft() applies and pyfftw’s execute does not – the same fold as the modal pair’s, with its own factor.

solprop.backends.numba_cpu.kernels.scale_tensor_spectrum(spectrum, kernel, scale)[source]#

Multiply interleaved tensor drives by trace or traceless kernel spectra.

Parameters:

spectrumnumpy.ndarray: Owned FFT rows, one or three consecutive rows per physical node.
kernelnumpy.ndarray: Trace-only or trace/traceless spectra, including dt.
scalefloat: Explicit inverse-transform normalization 1 / (2 n_to).

solprop.backends.numba_cpu.kernels.weighted_residual(propagator, dz, error, mask, scale, partials)[source]#

Sum the error norm’s residual with a per-bin scale, one partial per column.

solprop.stepper.controller.PIController.scale() returns one number when the controller has no weights and one number per bin when it has, and the weighted form divides before the norm is taken: the weights redistribute the tolerance across the band rather than rescaling the whole of it. A single summed residual divided by a scalar cannot express that, so a weighted run pays this second traversal – and only a weighted run does.

Parameters:

propagatornumpy.ndarray: exp(dPhi) over the whole step.
dzfloat: The step, m.
errornumpy.ndarray: The weighted stage sum of the embedded difference.
masknumpy.ndarray: The in-band mask the norm is taken over.
scalenumpy.ndarray: The controller’s per-bin scale.
partialsnumpy.ndarray: Written with one sum of squares per column, shape (n_columns,).

solprop.backends.numba_cpu.kernels.weighted_residual_parallel(propagator, dz, error, mask, scale, partials)#

weighted_residual() with its columns threaded (ACCEPT_ELEMENTS).

solprop.backends.numba_cpu.kernels.write_temporal_absorber(samples, rate, output)[source]#

Write the temporal absorber’s time-domain term, -rate*samples/2.

The product is grouped as solprop.boundaries.nonlinear_absorber_term() groups it, (-0.5 * rate) * samples, so the samples agree with the reference’s bit for bit.

Parameters:

samplesnumpy.ndarray: Fine temporal samples, real for a field grid (an analytic signal’s real part may be a strided view), shape (n_columns, n_to).
ratenumpy.ndarray: Power absorption rates in m^-1, shape (n_to,).
outputnumpy.ndarray: The planned forward transform’s input, overwritten.