OpenPMD to Ocelot format

[zhangli28]

Hi @AngelFP ,

It would be nice to have feature so that particle distribution can be converted from OpenPMD to Ocelot format. It would be very powerful for simulating beamline.

[AngelFP]

Hi @zhangli28, thanks for the suggestion. You could use the Wake-T adaptor with very minor modifications to convert an APtools particle distribution to ocelot. I think the code below should work:

from typing import Optional

import numpy as np
import scipy.constants as ct
from aptools.particle_distributions import ParticleDistribution

from ocelot.cpbd.beam import ParticleArray
from ocelot.common.globals import m_e_GeV

def particle_distribution_to_parray(
    particle_distribution: ParticleDistribution,
    gamma_ref: Optional[float] = None,
    z_ref: Optional[float] = None,
) -> ParticleArray:
    """
    Converts an APtools ParticleDistribution to an Ocelot ParticleArray.

    Parameters
    ------------
    particle_distribution : ParticleDistribution
        The original particle distribution from APtools.

    gamma_ref : float, optional
        Reference energy of the particle beam used for tracking in Ocelot. If
        not specified, the reference energy will be taken as the average
        energy of the input distribution.

    z_ref : float, optional
        Reference longitudinal position of the particle beam used for tracking
        in Ocelot. If not specified, the reference value will be taken as the
        average longitudinal position of the input distribution.

    Returns
    -------
    An Ocelot ParticleArray.

    """
    # Make sure that the distribution is an electron beam.
    is_electron_beam = (
        particle_distribution.q_species == -ct.e and particle_distribution.m_species == ct.m_e
    )
    assert is_electron_beam, 'Only electron beams are supported in Ocelot.'

    # Extract particle coordinates.
    x = particle_distribution.x  # [m]
    y = particle_distribution.y  # [m]
    z = particle_distribution.z  # [m]
    px = particle_distribution.px  # [m_e * c]
    py = particle_distribution.py  # [m_e * c]
    pz = particle_distribution.pz  # [m_e * c]
    q = np.abs(particle_distribution.q)  # [C]

    # Calculate gamma.
    gamma = np.sqrt(1 + px**2 + py**2 + pz**2)

    # Determine reference energy and longitudinal position.
    if gamma_ref is None:
        gamma_ref = np.average(gamma, weights=q)
    if z_ref is None:
        z_ref = np.average(z, weights=q)

    # Calculate momentum deviation (dp) and kinetic momentum (p_kin).
    b_ref = np.sqrt(1 - gamma_ref**(-2))
    dp = (gamma-gamma_ref)/(gamma_ref*b_ref)
    p_kin = np.sqrt(gamma_ref**2 - 1)

    # Create particle array
    p_array = ParticleArray(len(q))
    p_array.rparticles[0] = x
    p_array.rparticles[2] = y
    p_array.rparticles[4] = -(z - z_ref)
    p_array.rparticles[1] = px/p_kin
    p_array.rparticles[3] = py/p_kin
    p_array.rparticles[5] = dp
    p_array.q_array[:] = q
    p_array.s = z_ref
    p_array.E = gamma_ref * m_e_GeV

    return p_array

[zhangli28]

Hi @AngelFP ,

Thank you. You are very helpful and great personality. I very much appreciate your kind help.