pytc.kmat

JAX implementation of kinetic matrix elements.

Functions

_pad_axis	Pad arr with zeros by pad along axis.
_stream_l_panels	Iterate (U_panel, phi_r_panel, phi_s_panel) along axis-1 of U.
calc_K1	Calculate K1 matrix: K1_{pqrs} = sum_{i,j} w_i w_j phi_p(i) phi_q(i) nabla_i u(i, j) phi_r(j) phi_s(j)
calc_K1_kernel	Calculate K1 kernel: K1_{kl} = sum_{g,h} w_g w_h xi_{grad}(k,g) nabla u(g,h) xi_phi(l,h)
calc_K3	Calculate K3 matrix: K3_{pqrs} = sum_{i,j} w_i w_j phi_p(i) phi_q(i) (nabla_i u(i, j))^2 phi_r(j) phi_s(j)
calc_K3_kernel	Calculate K3 kernel: K3_{kl} = sum_{g,h} w_g w_h xi_phi(k,g) |nabla u(g,h)|^2 xi_phi(l,h)
contract_K1_antisym_pq_isdf_streaming	Streaming wrapper around contract_K1_antisym_pq_isdf_jit().
contract_K1_isdf	Contract K1 using ISDF decomposition.
contract_K1_isdf_streaming	Streaming-capable wrapper around contract_K1_isdf_jit() (symmetric, p == q case).
contract_K1_minus_K2_isdf	Compute (K1 - K2)[pqrs] in one pass, halving GPU peak vs separate calls.
contract_K1_minus_K2_isdf_streaming	Streaming-capable wrapper around contract_K1_minus_K2_isdf_jit().
contract_K3_isdf	Contract K3 using ISDF decomposition.
contract_K3_isdf_streaming	Streaming-capable wrapper around contract_K3_isdf_jit().