jaxdem.analysis.kernels#

Kernel typing + small helpers for JAX analysis.

In the JAX engine, kernels are pure functions that operate on arrays directly:

kernel(arrays, t0, t1, **kwargs) -> pytree of arrays

Functions

`isf_angular_kernel_2d`(arrays, t0, t1, *, theta_0)	Angular ISF: <cos(θ₀ · Δθ)>
`isf_self_isotropic_kernel`(arrays, t0, t1, *, k)	Self intermediate scattering function (isotropic average).
`isf_self_kvecs_kernel`(arrays, t0, t1, *, kvecs)	Self ISF for explicit k-vectors: Fs({k}, t) = <cos(k·dr)>.
`msad_kernel_2d`(arrays, t0, t1)	Mean-squared angular displacement on unwrapped cumulative angle.
`msd_kernel`(arrays, t0, t1)	Mean-squared displacement.
`unwrap_angles_2d`(q_w, q_xyz)	Convert (T, N, 1) and (T, N, 3) quaternion trajectory to unwrapped cumulative angle (T, N).

Classes

KernelFn(*args, **kwargs)

class jaxdem.analysis.kernels.KernelFn(*args, **kwargs)[source][source]#: Bases: Protocol

jaxdem.analysis.kernels.msd_kernel(arrays: Mapping[str, Array], t0: Any, t1: Any) → Array[source][source]#

Mean-squared displacement.

Works for both: - pos[t]: (N,d) -> returns () scalar - pos[t]: (S,N,d) -> returns (S,) vector

jaxdem.analysis.kernels.isf_self_isotropic_kernel(arrays: Mapping[str, Array], t0: Any, t1: Any, *, k: Any) → Array[source][source]#

Self intermediate scattering function (isotropic average).

For isotropic averaging: - 2D: Fs(k, t) = < J0(k * |dr|) > - 3D: Fs(k, t) = < sin(k|dr|)/(k|dr|) >

Supports: - pos[t]: (N,d) -> returns () if k scalar else (K,) - pos[t]: (S,N,d)-> returns (S,) if k scalar else (S,K)

jaxdem.analysis.kernels.isf_self_kvecs_kernel(arrays: Mapping[str, Array], t0: Any, t1: Any, *, kvecs: Array) → Array[source][source]#: Self ISF for explicit k-vectors: Fs({k}, t) = <cos(k·dr)>.

jaxdem.analysis.kernels.unwrap_angles_2d(q_w: Array, q_xyz: Array) → Array[source][source]#: Convert (T, N, 1) and (T, N, 3) quaternion trajectory to unwrapped cumulative angle (T, N).

jaxdem.analysis.kernels.msad_kernel_2d(arrays: Mapping[str, Array], t0: Any, t1: Any) → Array[source][source]#: Mean-squared angular displacement on unwrapped cumulative angle.

jaxdem.analysis.kernels.isf_angular_kernel_2d(arrays: Mapping[str, Array], t0: Any, t1: Any, *, theta_0: Any) → Array[source][source]#: Angular ISF: <cos(θ₀ · Δθ)>