INNER CODE UNIT · Python
create_initial_state
MatthewFilipovich/pycharge · src/pycharge/simulate.py:40
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
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]