prysm.mathops#

A submodule which allows the user to swap out the backend for mathematics.

class prysm.mathops.BackendShim(src)#

Bases: object

A shim that allows a backend to be swapped at runtime.

prysm.mathops.set_backend_to_cupy()#: Convenience method to automatically configure prysm’s backend to cupy.

prysm.mathops.set_backend_to_defaults()#: Convenience method to restore prysm’s default backend options.

prysm.mathops.set_backend_to_pytorch()#: Convenience method to automatically configure prysm’s backend to PyTorch.

prysm.mathops.set_backend_to_mlx()#: Convenience method to automatically configure prysm’s backend to MLX.

prysm.mathops.set_fft_backend_to_mkl_fft()#: Convenience method to automatically configure prysm’s backend to MKL_FFT for FFTs.

prysm.mathops.array_to_true_numpy(*args)#

Convert one or more arrays from an alternate backend to numpy.

Needed for plotting, serialization, etc.

Does nothing if given an actual numpy array

Parameters:: args (any number of arrays, of any dimension and dtype)
Return type:: array, or list of bonefide numpy arrays

prysm.mathops.row_dot(a, b)#

Batched dot product along the last (row) axis: sum(a * b, axis=-1).

For inputs of shape (N, K) returns shape (N,). Implementation uses einsum because it is the fastest route on numpy and works unchanged on the cupy / torch backends. See the prior incarnation in prysm/x/raytracing/spencer_and_murty.py (_multi_dot) for benchmark notes — this task is memory-bandwidth limited.

Parameters:

a (ndarray) – shape (N, K)
b (ndarray) – shape (N, K)

Returns:

shape (N,)

Return type:

ndarray

prysm.mathops.jinc(r)#

Jinc.

The first zero of jinc occurs at r=pi

Parameters:: r (number) – radial distance
Returns:: the value of j1(x)/x for x != 0, 0.5 at 0
Return type:: float

prysm.mathops.is_odd(int)#

Determine if an interger is odd using binary operations.

Parameters:: int (int) – an integer
Returns:: true if odd, False if even
Return type:: bool

prysm.mathops.is_power_of_2(value)#

Check if a value is a power of 2 using binary operations.

Parameters:: value (number) – value to check
Returns:: true if the value is a power of two, False if the value is no
Return type:: bool

Notes

c++ inspired implementation, see SO: https://stackoverflow.com/questions/29480680/finding-if-a-number-is-a-power-of-2-using-recursion

prysm.mathops.sign(x)#: Sign of a number. Note only works for single values, not arrays.

prysm.mathops.kronecker(i, j)#: Kronecker delta function, 1 if i = j, otherwise 0.

prysm.mathops.gamma(n, m)#: Gamma function.