INNER CODE UNIT · Python

calculate_total_sources

MatthewFilipovich/pycharge · src/pycharge/potentials_and_fields.py:114

    def calculate_total_sources(r: Array, t: Array) -> Quantities:
        t_srcs = jnp.stack([emission_time(r, t, charge) for charge in charges])
        # Evaluate source properties at the original emission times
        r_srcs = jnp.stack([position(t_src, charge) for t_src, charge in zip(t_srcs, charges)])
        v_srcs = jnp.stack([velocity(t_src, charge) for t_src, charge in zip(t_srcs, charges)])
        a_srcs = jnp.stack([acceleration(t_src, charge) for t_src, charge in zip(t_srcs, charges)])

        # Compute individual contributions
        calculate_individual_source_vmap = jax.vmap(calculate_individual_source, in_axes=(0, 0, 0, 0, None))
        individual_quantities = calculate_individual_source_vmap(r_srcs, v_srcs, a_srcs, qs, r)
        # Sum contributions from all charges
        summed_quantities = Quantities(*(jnp.sum(value, axis=0) for value in individual_quantities))

        return summed_quantities

    def calculate_individual_source(
        r_src: Array, v_src: Array, a_src: Array, q: Array, r: Array
    ) -> Quantities:

View source record →