INNER CODE UNIT · Python
quantities_fn
MatthewFilipovich/pycharge · src/pycharge/potentials_and_fields.py:97
def quantities_fn(x: ArrayLike, y: ArrayLike, z: ArrayLike, t: ArrayLike) -> Quantities:
x, y, z, t = jnp.asarray(x), jnp.asarray(y), jnp.asarray(z), jnp.asarray(t)
if not (x.shape == y.shape == z.shape == t.shape):
raise ValueError("x, y, z, and t must have the same shape.")
original_shape = x.shape
r = jnp.stack([x, y, z], axis=-1) # Stack into (..., 3)
r_flat, t_flat = r.reshape(-1, 3), t.ravel() # Flatten for vmap
quantities_flat = jax.vmap(calculate_total_sources)(r_flat, t_flat)
scalar_flat, other_quantities_flat = quantities_flat[0], quantities_flat[1:]
return Quantities(
scalar_flat.reshape(original_shape), # 1D scalar quantity
*(q.reshape(*original_shape, 3) for q in other_quantities_flat), # 3D vectors quantities
)
def calculate_total_sources(r: Array, t: Array) -> Quantities: