Utility Functions

Utility helpers shared across microJAX modules.

microjax.utils.first_nonzero(x: Array, axis: int = 0) → Array

Return the index of the first non-zero entry along axis.

Parameters

  • x (Array) – Input array.

  • axis (int, optional) – Axis along which to search. Defaults to 0.

Returns

Index of the first non-zero entry. If no entry is non-zero the result is 0.

Return type

Array

microjax.utils.first_zero(x: Array, axis: int = 0) → Array

Return the index of the first zero entry along axis.

Parameters

  • x (Array) – Input array.

  • axis (int, optional) – Axis along which to search. Defaults to 0.

Returns

Index of the first zero entry. If no entry is zero the result is 0.

Return type

Array

microjax.utils.last_nonzero(x: Array, axis: int = 0) → Array

Return the index of the last non-zero entry along axis.

Parameters

  • x (Array) – Input array.

  • axis (int, optional) – Axis along which to search. Defaults to 0.

Returns

Index of the last non-zero entry. If all entries are zero the result is 0.

Return type

Array

microjax.utils.match_points(a: Array, b: Array) → Array

Greedy nearest-neighbour matching between two point sets.

The algorithm iterates over a and, for each element, picks the closest unused element of b. It does not solve the optimal assignment problem but provides a fast, deterministic ordering that is sufficient for tracking lens images across sampling steps.

Parameters

  • a (Array) – One-dimensional complex array of source points.

  • b (Array) – One-dimensional complex array of candidate points to match against.

Returns

Indices that permute b into the greedy match order of a.

Return type

Array

microjax.utils.max_zero_avoiding(x: Array) → Array

Maximum value of x while skipping zeros.

Parameters

x (Array) – One-dimensional input array.

Returns

Maximum non-zero value if present, otherwise 0.

Return type

Array

microjax.utils.mean_zero_avoiding(x: Array) → Array

Mean of non-zero entries in x.

Parameters

x (Array) – Input array.

Returns

Mean over the non-zero entries or 0 if all entries are zero.

Return type

Array

microjax.utils.min_zero_avoiding(x: Array) → Array

Minimum value of x while skipping zeros.

Parameters

x (Array) – One-dimensional input array.

Returns

Minimum non-zero value if present, otherwise 0.

Return type

Array

microjax.utils.sparse_argsort(a: Array) → Array

Argsort that treats zeros as missing values.

Parameters

a (Array) – Input array.

Returns

Indices that sort a after replacing zeros with NaN.

Return type

Array

microjax.utils.trapz_zero_avoiding(y: Array, x: Array, tail_idx: int) → Array

Trapezoidal integration omitting the final segment beyond tail_idx.

Equivalent to jax.scipy.integrate.trapezoid(y[:tail_idx+1], x[:tail_idx+1]) but avoids the last segment if tail_idx does not point to the final sample.

Parameters

  • y (Array) – Function samples.

  • x (Array) – Sample positions.

  • tail_idx (int) – Index of the last reliable sample to include in the integral.

Returns

Trapezoidal integral up to tail_idx.

Return type

Array