INNER CODE UNIT · Python
r
MatthewFilipovich/pycharge · src/pycharge/potentials_and_fields.py:103
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:
t_srcs = jnp.stack([emission_time(r, t, charge) for charge in charges])
# Evaluate source properties at the original emission times
r_srcs = jnp.stack([position(t_src, charge) for t_src, charge in zip(t_srcs, charges)])
v_srcs = jnp.stack([velocity(t_src, charge) for t_src, charge in zip(t_srcs, charges)])
a_srcs = jnp.stack([acceleration(t_src, charge) for t_src, charge in zip(t_srcs, charges)])