INNER CODE UNIT · Python
time_step_source
MatthewFilipovich/pycharge · src/pycharge/simulate.py:52
def time_step_source(source_idx):
ode_func = sources[source_idx].ode_func
t = ts[time_idx]
u = source_states[source_idx][time_idx]
dt = dts[time_idx]
other_charges_flat = [c for i, c_tuple in enumerate(charges) if i != source_idx for c in c_tuple]
u_step = rk4_step(ode_func, t, u, dt, other_charges_flat)
return source_states[source_idx].at[time_idx + 1].set(u_step)
source_states = tuple(time_step_source(source_idx) for source_idx in range(len(sources)))
return source_states
def print_timestep(time_idx: int):
if print_every:
jax.lax.cond(
time_idx % print_every == 0,