API

impact.Impact

Impact(*args, group=None, always_autophase=False, **kwargs)

Bases: CommandWrapper

Files will be written into a temporary directory within workdir. If workdir=None, a location will be determined by the system.

Source code in impact/impact.py

def __init__(self, *args, group=None, always_autophase=False, **kwargs):
    super().__init__(*args, **kwargs)
    # Save init
    self.original_input_file = self.input_file

    self.input = {"header": {}, "lattice": []}
    self.output = {}
    self._units = {}
    self._units.update(STATIC_UNITS)
    self.group = {}

    # MPI
    self._nnode = 1

    # Convenience lookup of elements in lattice by name
    self.ele = {}

    # Autophase settings to be applied.
    # This will be cleared when actually autophasing
    self._autophase_settings = {}
    self.always_autophase = always_autophase

    # Call configure
    if self.input_file:
        infile = tools.full_path(self.input_file)
        assert os.path.exists(infile), f"Impact input file does not exist: {infile}"
        self.load_input(self.input_file)
        self.configure()

        # Add groups, if any.
        if group:
            for k, v in group.items():
                self.add_group(k, **v)

    else:
        self.vprint("Using default input: 1 m drift lattice")
        self.input = deepcopy(DEFAULT_INPUT)
        self.configure()

Attributes

impact.Impact.cathode_start property writable

cathode_start

Returns a bool if cathode_start is requested. Can also be set as a bool.

impact.Impact.fieldmaps property

fieldmaps

Convenience pointer to .input['fieldmaps']

impact.Impact.header property

header

Convenience pointer to .input['header']

impact.Impact.lattice property

lattice

Convenience pointer to .input['lattice']

impact.Impact.nnode property writable

nnode

Number of MPI nodes to use

impact.Impact.numprocs property writable

numprocs

Number of MPI processors = Npcol*Nprow

impact.Impact.particles property

particles

Convenience pointer to .input['lattice']

impact.Impact.total_charge property writable

total_charge

Returns the total bunch charge in C. Can be set.

Methods:

impact.Impact.add_ele

add_ele(ele)

Adds an element to .lattice

Source code in impact/impact.py

def add_ele(self, ele):
    """
    Adds an element to .lattice
    """
    name = ele["name"]
    assert name not in self.lattice, "Element already exists"
    insert_ele_by_s(ele, self.lattice, verbose=self.verbose)
    # Add to the ele dict
    self.ele[name] = ele

impact.Impact.add_group

add_group(name, **kwargs)

Add a control group. See control.py

Source code in impact/impact.py

def add_group(self, name, **kwargs):
    """
    Add a control group. See control.py
    """
    assert name not in self.ele
    if name in self.group:
        self.vprint(f"Warning: group {name} already exists, overwriting.")

    g = ControlGroup(**kwargs, name=name)
    g.link(self.ele)
    self.group[name] = g

    return self.group[name]

impact.Impact.archive

archive(h5=None)

Archive all data to an h5 handle or filename.

If no file is given, a file based on the fingerprint will be created.

Source code in impact/impact.py

def archive(self, h5=None):
    """
    Archive all data to an h5 handle or filename.

    If no file is given, a file based on the fingerprint will be created.

    """
    if not h5:
        h5 = "impact_" + self.fingerprint() + ".h5"

    if isinstance(h5, str):
        new_h5file = True
        fname = os.path.expandvars(h5)
        g = h5py.File(fname, "w")
        self.vprint(f"Archiving to file {fname}")
    else:
        new_h5file = False
        g = h5

    # Write basic attributes
    archive.impact_init(g)

    # Initial particles
    if self.initial_particles:
        self.initial_particles.write(g, name="initial_particles")

    # All input
    archive.write_input_h5(g, self.input, name="input")

    # All output
    archive.write_output_h5(g, self.output, name="output", units=self._units)

    # Control groups
    if self.group:
        archive.write_control_groups_h5(g, self.group, name="control_groups")

    # Close file if created here.
    if new_h5file:
        g.close()

    return h5

impact.Impact.autophase

autophase(settings=None, full_output=False)

Calculate the relative phases of each rf element by tracking a single particle. This uses a fast method that operates outside of Impact

Parameters:

Name	Type	Description	Default
settings		dict of ele_name:rel_phase_deg	None
full_output		type of output to return (see Returns)	False

Returns:

Type	Description
if full_output = True retuns a dict of:	ele_name:info_dict
Otherwise returns a dict of:	ele_name:rel_phase_deg
which is the same format as settings.

Source code in impact/impact.py

def autophase(self, settings=None, full_output=False):
    """
    Calculate the relative phases of each rf element
    by tracking a single particle.
    This uses a fast method that operates outside of Impact

    Parameters
    ----------
    settings: dict, optional=None
        dict of ele_name:rel_phase_deg

    full_output: bool, optional = False
        type of output to return (see Returns)


    Returns
    -------
    if full_output = True retuns a dict of:
            ele_name:info_dict

    Otherwise returns a dict of:
        ele_name:rel_phase_deg
    which is the same format as settings.


    """

    if self.initial_particles:
        t0 = self.initial_particles["mean_t"]
        pz0 = self.initial_particles["mean_pz"]
    else:
        t0 = 0
        pz0 = 0

    return fast_autophase_impact(
        self,
        settings=settings,
        t0=t0,
        pz0=pz0,
        full_output=full_output,
        verbose=self.verbose,
    )

impact.Impact.autophase_bookkeeper

autophase_bookkeeper()

Searches for 'autophase_deg' attribute in all eles. If one is found, autophase is called.

If .always_autophase == True, calls autophase is called.

Returns:

Name	Type	Description
settings	dict	Autophase settings found

Source code in impact/impact.py

def autophase_bookkeeper(self):
    """
    Searches for `'autophase_deg'` attribute in all eles.
    If one is found, autophase is called.

    If .always_autophase == True, calls autophase is called.

    Returns
    -------
    settings: dict
        Autophase settings found
    """
    if self._autophase_settings or self.always_autophase:
        if self.verbose:
            print("Autophase bookkeeper found settings, applying them")

        # Actual found settings
        settings = self.autophase(settings=self._autophase_settings)

        # Clear
        self._autophase_settings = {}

    else:
        settings = {}

    return settings

impact.Impact.ele_bookkeeper

ele_bookkeeper()

Link .ele = dict to the lattice elements by their 'name' field

Source code in impact/impact.py

def ele_bookkeeper(self):
    """
    Link .ele = dict to the lattice elements by their 'name' field
    """
    self.ele = ele_dict_from(self.input["lattice"])

impact.Impact.field

field(z=0, t=0, x=0, y=0, component='Ez')

Return the field component at a longitudinal position z at time t.

Warking: This is based on the parsed fieldmaps, and not calculated directly from Impact. Not all elements/parameters are implemented. Currently x, y must be 0.

Source code in impact/impact.py

def field(self, z=0, t=0, x=0, y=0, component="Ez"):
    """
    Return the field component at a longitudinal
    position z at time t.

    Warking: This is based on the parsed fieldmaps,
    and not calculated directly from Impact. Not all elements/parameters
    are implemented. Currently x, y must be 0.
    """
    return lattice_field(
        self.lattice, x=x, y=y, z=z, t=t, component=component, fmaps=self.fieldmaps
    )

impact.Impact.from_tao classmethod

from_tao(tao, fieldmap_style='fourier', n_coef=30, **kwargs)

Create a complete Impact object from a running Pytao Tao instance.

Wraps impact.interfaces.bmad.impact_from_tao

Parameters:

Name	Type	Description	Default
tao	Tao	Initialized Tao object	required
fieldmap_style	str	Style of fieldmap to create. One of: ('fourier', 'derivatives'), by default = 'fourier'	'fourier'
n_coef	int	a number of fourier (?) coefficients, by default 30	30
emfield_cartesian_eles	str or list[str]	Matching string or list of element names to be converted to Impact-T emfield_cartesian elements, by default = "EM_FIELD::*"	required
solrf_eles	str or list[str]	Matching string or list of element names to be converted to Impact-T solrf elements, by default = 'E_GUN::,SOLENOID::,LCAVITY::*'	required
quadrupole_eles	str or list[str]	Matching string or list of element names to be converted to Impact-T quadrupole elements, by default = 'quad::*'	required
write_beam_eles	str or list[str]	Matching string or list of element names to be converted to Impact-T write beam elements. Note that there is a limit to the number of these that can be created. By default = 'monitor::*'	required

Returns:

Type	Description
Impact	Initialized Impact object

Source code in impact/impact.py

@classmethod
def from_tao(cls, tao, fieldmap_style="fourier", n_coef=30, **kwargs):
    """
    Create a complete Impact object from a running Pytao Tao instance.

    Wraps impact.interfaces.bmad.impact_from_tao

    Parameters
    ----------
    tao : Tao
        Initialized Tao object

    fieldmap_style : str, optional
        Style of fieldmap to create. One of: ('fourier', 'derivatives'), by default = 'fourier'

    n_coef : int, optional
        a number of fourier (?) coefficients, by default 30

    emfield_cartesian_eles : str or list[str], optional
        Matching string or list of element names to be converted to Impact-T emfield_cartesian elements, by default = "EM_FIELD::*"

    solrf_eles : str or list[str], optional
        Matching string or list of element names to be converted to Impact-T solrf elements, by default = 'E_GUN::*,SOLENOID::*,LCAVITY::*'

    quadrupole_eles : str or list[str], optional
        Matching string or list of element names to be converted to Impact-T quadrupole elements, by default = 'quad::*'

    write_beam_eles : str or list[str], optional
        Matching string or list of element names to be converted to Impact-T write beam elements.
        Note that there is a limit to the number of these that can be created.  By default = 'monitor::*'

    Returns
    -------
    Impact
        Initialized Impact object
    """
    return impact_from_tao(
        tao, fieldmap_style=fieldmap_style, n_coef=n_coef, **kwargs
    )

impact.Impact.get_executable

get_executable()

Gets the full path of the executable from .command, .command_mpi Will search environmental variables: Impact.command_env='IMPACTT_BIN' Impact.command_mpi_env='IMPACTT_MPI_BIN'

Source code in impact/impact.py

def get_executable(self):
    """
    Gets the full path of the executable from .command, .command_mpi
    Will search environmental variables:
            Impact.command_env='IMPACTT_BIN'
            Impact.command_mpi_env='IMPACTT_MPI_BIN'

    """
    if self.use_mpi:
        exe = tools.find_executable(
            exename=self.command_mpi, envname=self.command_mpi_env
        )
    else:
        exe = tools.find_executable(exename=self.command, envname=self.command_env)
    return exe

impact.Impact.get_run_script

get_run_script(write_to_path=False, path=None)

Assembles the run script using self.mpi_run string of the form: 'mpirun -n {n} {command_mpi}'

Optionally writes a file 'run' with this line to path.

Source code in impact/impact.py

def get_run_script(self, write_to_path=False, path=None):
    """
    Assembles the run script using self.mpi_run string of the form:
        'mpirun -n {n} {command_mpi}'

    Optionally writes a file 'run' with this line to path.
    """

    n_procs = self.numprocs

    exe = self.get_executable()

    if self.use_mpi:
        # mpi_exe could be a complicated string like:
        # 'srun -N 1 --cpu_bind=cores {n} {command_mpi}'
        # 'mpirun -n {n} {command_mpi}'

        runscript = self.mpi_run.format(
            nnode=self.nnode, nproc=n_procs, command_mpi=exe
        )

    else:
        if n_procs > 1:
            raise ValueError("Error: n_procs > 1 but use_mpi = False")
        runscript = exe

    if write_to_path:
        if path is None:
            path = self.path
        path = os.path.join(path, "run")
        with open(path, "w") as f:
            f.write(runscript)
        tools.make_executable(path)
    return runscript

impact.Impact.load_archive

load_archive(h5, configure=True)

Loads input and output from archived h5 file.

See: Impact.archive

Source code in impact/impact.py

def load_archive(self, h5, configure=True):
    """
    Loads input and output from archived h5 file.

    See: Impact.archive
    """
    if isinstance(h5, str):
        fname = os.path.expandvars(h5)
        g = h5py.File(fname, "r")

        glist = archive.find_impact_archives(g)
        n = len(glist)
        if n == 0:
            # legacy: try top level
            message = "legacy"
        elif n == 1:
            gname = glist[0]
            message = f"group {gname} from"
            g = g[gname]
        else:
            raise ValueError(f"Multiple archives found in file {fname}: {glist}")

        self.vprint(f"Reading {message} archive file {h5}")
    else:
        g = h5

    self.input = archive.read_input_h5(g["input"], verbose=self.verbose)
    # Rewrite legacy header keys (e.g. "sigx(m)") to their canonical
    # form on load so the in-memory header is always canonical.
    _normalize_header_keys(self.input["header"])
    self.output, self._units = archive.read_output_h5(
        g["output"], verbose=self.verbose
    )
    self._units.update(STATIC_UNITS)

    if "initial_particles" in g:
        self.initial_particles = ParticleGroup(h5=g["initial_particles"])

    if "control_groups" in g:
        self.group = archive.read_control_groups_h5(
            g["control_groups"], verbose=self.verbose
        )
    self.vprint("Loaded from archive. Note: Must reconfigure to run again.")
    self.configured = False

    if configure:
        self.configure()

        # Re-link groups
        # TODO: cleaner logic
        for _, cg in self.group.items():
            cg.link(self.ele)

impact.Impact.load_output

load_output()

Loads stats, slice_info, and particles.

Source code in impact/impact.py

def load_output(self):
    """
    Loads stats, slice_info, and particles.
    """
    self.output["stats"], u = load_stats(
        self.path, species=self.species, types=FORT_STAT_TYPES, verbose=self.verbose
    )
    self._units.update(u)

    # This is not always present
    dipole_stats, u = load_stats(
        self.path,
        species=self.species,
        types=FORT_DIPOLE_STAT_TYPES,
        verbose=self.verbose,
    )
    if dipole_stats:
        self.output["dipole_stats"] = dipole_stats
        self._units.update(u)

    self.output["slice_info"], u = load_slice_info(self.path, self.verbose)
    self._units.update(u)

    self.load_particles()

impact.Impact.old_plot

old_plot(y='sigma_x', x='mean_z', nice=True, include_layout=True)

Simple stat plot

Source code in impact/impact.py

def old_plot(self, y="sigma_x", x="mean_z", nice=True, include_layout=True):
    """
    Simple stat plot
    """
    return plot_stat(self, y=y, x=x, nice=nice)

impact.Impact.plot

plot(y=('sigma_x', 'sigma_y'), x='mean_z', xlim=None, ylim=None, ylim2=None, y2=(), nice=True, include_layout=True, include_labels=False, include_markers=True, include_particles=True, include_field=True, field_t=None, include_legend=True, return_figure=False, tex=True, **kwargs)

Source code in impact/impact.py

def plot(
    self,
    y: str | Sequence[str] = ("sigma_x", "sigma_y"),
    x="mean_z",
    xlim=None,
    ylim=None,
    ylim2=None,
    y2: str | Sequence[str] = (),
    nice=True,
    include_layout=True,
    include_labels=False,
    include_markers=True,
    include_particles=True,
    include_field=True,
    field_t=None,
    include_legend=True,
    return_figure=False,
    tex=True,
    **kwargs,
):
    """ """
    if not isinstance(y, str):
        y = list(y)
    if not isinstance(y2, str):
        y2 = list(y2)

    # Just plot fieldmaps if there are no stats
    if "stats" not in self.output:
        return plot_layout(
            self,
            xlim=xlim,
            include_markers=include_markers,
            include_labels=include_labels,
            include_field=include_field,
            field_t=field_t,
            return_figure=return_figure,
            **kwargs,
        )

    return plot_stats_with_layout(
        self,
        ykeys=y,
        ykeys2=y2,
        xkey=x,
        xlim=xlim,
        ylim=ylim,
        ylim2=ylim2,
        nice=nice,
        tex=tex,
        include_layout=include_layout,
        include_labels=include_labels,
        include_field=include_field,
        field_t=field_t,
        include_markers=include_markers,
        include_particles=include_particles,
        include_legend=include_legend,
        return_figure=return_figure,
        **kwargs,
    )

impact.Impact.print_lattice

print_lattice()

Pretty printing of the lattice

Source code in impact/impact.py

def print_lattice(self):
    """
    Pretty printing of the lattice
    """
    for ele in self.input["lattice"]:
        line = ele_str(ele)
        print(line)

impact.Impact.run_impact

run_impact(verbose=False, timeout=None)

Runs Impact-T

Source code in impact/impact.py

def run_impact(self, verbose=False, timeout=None):
    """
    Runs Impact-T

    """

    # Clear output
    self.output = {}
    self.finished = False

    # Autophase
    autophase_settings = self.autophase_bookkeeper()
    if autophase_settings:
        self.output["autophase_info"] = autophase_settings

    run_info = self.output["run_info"] = {"error": False}

    # Run script, gets executables
    runscript = self.get_run_script()
    run_info["run_script"] = runscript

    t1 = time()
    run_info["start_time"] = t1

    self.vprint("Running Impact-T in " + self.path)
    self.vprint(runscript)
    # Write input
    self.write_input()

    # Remove previous files
    for f in fort_files(self.path):
        os.remove(f)

    if timeout:
        res = tools.execute2(runscript.split(), timeout=timeout, cwd=self.path)
        log = res["log"]
        self.error = res["error"]
        run_info["error"] = self.error
        run_info["why_run_error"] = res["why_error"]
    else:
        # Interactive output, for Jupyter
        log = []
        counter = 0
        for path in tools.execute(runscript.split(), cwd=self.path):
            # Fancy clearing of old lines
            counter += 1
            if verbose:
                if counter < 15:
                    print(path, end="")
                else:
                    print(
                        "\r",
                        path.strip() + ", elapsed: " + str(time() - t1),
                        end="",
                    )
            log.append(path)
        self.vprint("Finished.")
    self.log = log

    # Load output
    self.load_output()

    run_info["run_time"] = time() - t1

    self.finished = True

impact.Impact.stat

stat(key)

Array from .output['stats'][key]

Additional keys are avalable: 'mean_energy': mean energy 'Ez': z component of the electric field at the centroid particle 'Bz' z component of the magnetic field at the centroid particle 'cov_{a}__{b}': any symmetric covariance matrix term

Source code in impact/impact.py

def stat(self, key):
    """
    Array from `.output['stats'][key]`

    Additional keys are avalable:
        'mean_energy': mean energy
        'Ez': z component of the electric field at the centroid particle
        'Bz'  z component of the magnetic field at the centroid particle
        'cov_{a}__{b}': any symmetric covariance matrix term

    """

    if key in ("Ez", "Bz"):
        return self.centroid_field(component=key[0:2])

    if key == "mean_energy":
        return self.stat("mean_kinetic_energy") + self.mc2

    # Allow flipping covariance keys
    if key.startswith("cov_") and key not in self.output["stats"]:
        k1, k2 = key[4:].split("__")
        key = f"cov_{k2}__{k1}"

    if key not in self.output["stats"]:
        raise ValueError(f"{key} is not available in the output data")

    return self.output["stats"][key]

impact.Impact.track

track(particles, s=None)

Track a ParticleGroup. An optional stopping s can be given.

Source code in impact/impact.py

def track(self, particles, s=None):
    """
    Track a ParticleGroup. An optional stopping s can be given.
    """
    if not s:
        s = self.stop
    return track_to_s(self, particles, s)

impact.Impact.track1

track1(x0=0, px0=0, y0=0, py0=0, z0=0, pz0=1e-15, t0=0, s=None, species=None)

Tracks a single particle with starting coordinates: x0, y0, z0 in meters px0, py0, pz0 in eV/c t0 in seconds

to a position 's' in meters

Used for phasing and scaling elements.

If successful, returns a ParticleGroup with the final particle.

Otherwise, returns None

Source code in impact/impact.py

def track1(
    self,
    x0=0,
    px0=0,
    y0=0,
    py0=0,
    z0=0,
    pz0=1e-15,
    t0=0,
    s=None,  # final s
    species=None,
):
    """
    Tracks a single particle with starting coordinates:
    x0, y0, z0 in meters
    px0, py0, pz0 in eV/c
    t0 in seconds

    to a position 's' in meters

    Used for phasing and scaling elements.

    If successful, returns a ParticleGroup with the final particle.

    Otherwise, returns None

    """
    if not s:
        s = self.stop

    if not species:
        species = self.species

    # Change to serial exe just for this
    n_procs_save = self.numprocs
    self.numprocs = 1
    result = track1_to_s(
        self,
        s=s,
        x0=x0,
        px0=px0,
        y0=y0,
        py0=py0,
        z0=z0,
        pz0=pz0,
        t0=t0,
        species=species,
    )
    self.numprocs = n_procs_save
    return result

impact.Impact.units

units(key)

pmd_unit of a given key

Source code in impact/impact.py

def units(self, key):
    """pmd_unit of a given key"""

    # Allow flipping covariance keys
    if key.startswith("cov_") and key not in self._units:
        k1, k2 = key[4:].split("__")
        key = f"cov_{k2}__{k1}"

    if key not in self._units:
        raise ValueError(f"Unknown unit for {key}")

    return self._units[key]

impact.Impact.vprint

vprint(*args)

verbose print

Source code in impact/impact.py

def vprint(self, *args):
    """verbose print"""
    if self.verbose:
        print(*args)

impact.Impact.write_input

write_input(input_filename='ImpactT.in', path=None)

Write all input.

If .initial_particles are given,

Source code in impact/impact.py

def write_input(self, input_filename="ImpactT.in", path=None):
    """
    Write all input.

    If .initial_particles are given,
    """

    if path is None:
        path = self.path

    pathlib.Path(path).mkdir(exist_ok=True, parents=True)

    filePath = os.path.join(path, input_filename)

    # Write fieldmaps
    for name, fieldmap in self.input["fieldmaps"].items():
        file = os.path.join(path, name)
        fieldmaps.write_fieldmap(file, fieldmap)

    # Initial particles (ParticleGroup)
    if self.initial_particles:
        self.write_initial_particles(update_header=True, path=path)

        # Check consistency
        if (
            self.header["Flagimg"] == 1
            and self.header["Nemission"] < 1
            and self.total_charge > 0
        ):
            raise ValueError(
                f"Cathode start with space charge must "
                f"set header['Nemission'] > 0. "
                f"The current value is {self.header['Nemission']}."
            )

    # Symlink
    elif self.header["Flagdist"] == 16:
        src = self.input["input_particle_file"]
        dest = os.path.join(path, "partcl.data")

        # Don't worry about overwriting in temporary directories
        if self._tempdir and os.path.exists(dest):
            os.remove(dest)

        if not os.path.exists(dest):
            writers.write_input_particles_from_file(src, dest, self.header["Np"])
        else:
            self.vprint("partcl.data already exits, will not overwrite.")

    # Check for point-to-point spacechage processor criteria
    if self.numprocs > 1:
        for ele in self.lattice:
            if ele["type"] == "point_to_point_spacecharge":
                Np = self.header["Np"]
                numprocs = self.numprocs
                if Np % numprocs != 0:
                    raise ValueError(
                        f"The number of electrons ({Np}) divided by the number of processors ({numprocs}) must be an integer."
                    )

    # Write main input file. This should come last.
    writers.write_impact_input(filePath, self.header, self.lattice)

    # Write run script
    self.get_run_script(write_to_path=True, path=path)