INNER CODE UNIT · Python
time_step_body
MatthewFilipovich/pycharge · src/pycharge/simulate.py:48
def time_step_body(time_idx: int, source_states: tuple[Array, ...]) -> tuple[Array, ...]:
print_timestep(time_idx)
charges = create_charges(time_idx, source_states)
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