INNER CODE UNIT · Python
final_source_states
MatthewFilipovich/pycharge · src/pycharge/simulate.py:36
final_source_states = jax.lax.fori_loop(0, len(ts) - 1, time_step_body, initial_source_states)
return final_source_states
def create_initial_state(source: Source) -> Array:
source_state = jnp.full([len(ts), len(source.charges_0), 2, 3], jnp.nan)
for charge_idx, charge in enumerate(source.charges_0):
pos0, vel0 = position(ts[0], charge), velocity(ts[0], charge)
source_state = source_state.at[0, charge_idx, :, :].set([pos0, vel0])
return source_state
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