INNER CODE UNIT · Python

print_timestep

MatthewFilipovich/pycharge · src/pycharge/simulate.py:66

    def print_timestep(time_idx: int):
        if print_every:
            jax.lax.cond(
                time_idx % print_every == 0,
                lambda: jax.debug.print("Timestep {x}", x=time_idx),
                lambda: None,
            )

    def create_charges(time_idx: int, source_states: tuple[Array, ...]):
        def create_charge(source_idx: int, charge_idx: int):
            charge_0 = sources[source_idx].charges_0[charge_idx]
            position_0_fn = charge_0.position_fn
            source_state = source_states[source_idx]

            position_array = source_state[:, charge_idx, 0]
            velocity_array = source_state[:, charge_idx, 1]
            t_end = ts[time_idx]

View source record →