INNER CODE UNIT · Python
dts
MatthewFilipovich/pycharge · src/pycharge/simulate.py:32
dts = ts[1:] - ts[:-1]
def simulate_fn():
initial_source_states = tuple(create_initial_state(source) for source in sources)
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)