diff --git a/aptools/__init__.py b/aptools/__init__.py index 3319366..bc2df0c 100644 --- a/aptools/__init__.py +++ b/aptools/__init__.py @@ -1,4 +1,4 @@ from .data_analysis import beam_diagnostics -__version__ = "0.1.25" +__version__ = "0.2.0" __all__ = ['beam_diagnostics', '__version__'] diff --git a/aptools/data_handling/reading.py b/aptools/data_handling/reading.py index 0744151..d635703 100644 --- a/aptools/data_handling/reading.py +++ b/aptools/data_handling/reading.py @@ -4,12 +4,18 @@ import numpy as np import scipy.constants as ct from h5py import File as H5File +from deprecated import deprecated from aptools.helper_functions import join_infile_path, reposition_bunch from aptools.plasma_accel.general_equations import plasma_skin_depth from aptools.data_processing.beam_filtering import filter_beam +@deprecated( + version="0.2.0", + reason=("This method is replaced by those in the new " + "`particle_distributions.read` module.") +) def read_beam(code_name, file_path, reposition=False, avg_pos=[None, None, None], avg_mom=[None, None, None], filter_min=[None, None, None, None, None, None, None], diff --git a/aptools/data_handling/saving.py b/aptools/data_handling/saving.py index 7c58feb..977dcd5 100644 --- a/aptools/data_handling/saving.py +++ b/aptools/data_handling/saving.py @@ -7,6 +7,7 @@ import scipy.constants as ct from openpmd_api import (Series, Access, Dataset, Mesh_Record_Component, Unit_Dimension) +from deprecated import deprecated from aptools.helper_functions import reposition_bunch, get_particle_subset from aptools import __version__ @@ -15,6 +16,11 @@ SCALAR = Mesh_Record_Component.SCALAR +@deprecated( + version="0.2.0", + reason=("This method is replaced by those in the new " + "`particle_distributions.save` module.") +) def save_beam(code_name, beam_data, folder_path, file_name, reposition=False, avg_pos=[None, None, None], avg_mom=[None, None, None], n_part=None, **kwargs): diff --git a/aptools/particle_distributions/__init__.py b/aptools/particle_distributions/__init__.py new file mode 100644 index 0000000..f90ec5d --- /dev/null +++ b/aptools/particle_distributions/__init__.py @@ -0,0 +1,6 @@ +from .particle_distribution import ParticleDistribution +from .read import read_distribution +from .save import save_distribution + + +__all__ = ['ParticleDistribution', 'read_distribution', 'save_distribution'] diff --git a/aptools/particle_distributions/particle_distribution.py b/aptools/particle_distributions/particle_distribution.py new file mode 100644 index 0000000..d2cfa42 --- /dev/null +++ b/aptools/particle_distributions/particle_distribution.py @@ -0,0 +1,52 @@ +"""Defiles the ParticleDistribution class""" + +import numpy as np + + +class ParticleDistribution(): + """Class for storing a particle distribution. + + Parameters + ---------- + x, y, xi : ndarray + Position of the macropparticles in the x, y, and xi directions in + units of m. + px, py, pz : ndarray + Momentum of the macropparticles in the x, y, and z directions in + non-dimensional units (beta*gamma). + w : ndarray + Weight of the macroparticles, i.e., the number of real particles + represented by each macroparticle. In practice, :math:`w = q_m / q`, + where :math:`q_m` and :math:`q` are, respectively, the charge of the + macroparticle and of the real particle (e.g., an electron). + q_species, m_species : float + Charge and mass of a single particle of the species represented + by the macroparticles. For an electron bunch, + ``q_species=-e`` and ``m_species=m_e`` + """ + def __init__( + self, + x: np.ndarray, + y: np.ndarray, + z: np.ndarray, + px: np.ndarray, + py: np.ndarray, + pz: np.ndarray, + w: np.ndarray, + q_species: float, + m_species: float + ) -> None: + self.x = x + self.y = y + self.z = z + self.px = px + self.py = py + self.pz = pz + self.w = w + self.q_species = q_species + self.m_species = m_species + + @property + def q(self) -> np.ndarray: + """Return the total macroparticle charge.""" + return self.w * self.q_species diff --git a/aptools/particle_distributions/read.py b/aptools/particle_distributions/read.py new file mode 100644 index 0000000..7eae228 --- /dev/null +++ b/aptools/particle_distributions/read.py @@ -0,0 +1,218 @@ +"""Defines methods for reading particle distributions from file.""" + +from typing import Optional + +import numpy as np +import scipy.constants as ct +from h5py import File + +from aptools.helper_functions import join_infile_path +from .particle_distribution import ParticleDistribution + + +def read_distribution( + file_path: str, + data_format: str, + **kwargs +) -> ParticleDistribution: + """Read particle distribution from file. + + Parameters + ---------- + file_path : str + Path to the file with particle data + data_format : str + Internal format of the data. Possible values + are 'astra', 'csrtrack' and 'openpmd'. + + Other Parameters + ---------------- + **kwargs + Additional parameters to be passed to the particle readers. + + Returns + ------- + ParticleDistribution + The particle distribution. + """ + return _readers[data_format](file_path, **kwargs) + + +def read_from_astra( + file_path: str, + remove_non_standard: bool = True +) -> ParticleDistribution: + """Read particle distribution from ASTRA. + + Parameters + ---------- + file_path : str + Path to the file with particle data + remove_non_standard : bool + Determines whether non-standard particles (those with a status flag + other than 5) should be removed from the read data. + + Returns + ------- + ParticleDistribution + The particle distribution. + """ + # Read data. + data = np.genfromtxt(file_path) + + # Get status flag and remove non-standard particles, if needed. + status_flag = data[:, 9] + if remove_non_standard: + data = data[np.where(status_flag == 5)] + + # Extract phase space and particle index. + x = data[:, 0] + y = data[:, 1] + z = data[:, 2] + px = data[:, 3] + py = data[:, 4] + pz = data[:, 5] + q = data[:, 7] * 1e-9 + index = data[:, 8] + + # Apply reference particle. + z[1:] += z[0] + pz[1:] += pz[0] + + # Determine charge and mass of particle species. + assert index[0] in [1, 2, 3], ( + 'Only electrons, positrons and protons are supported when reading ' + 'ASTRA distributions') + q_species = ct.e * (-1 if index[0] in [1, 3] else 1) + m_species = ct.m_e if index[0] in [1, 2] else ct.m_p + + # Convert momentum to normalized units (beta * gamma). + px /= m_species * ct.c**2 / ct.e + py /= m_species * ct.c**2 / ct.e + pz /= m_species * ct.c**2 / ct.e + + # Get particle weights. + w = q / q_species + + # Return distribution. + return ParticleDistribution(x, y, z, px, py, pz, w, q_species, m_species) + + +def read_from_csrtrack( + file_path: str, +) -> ParticleDistribution: + """Read particle distribution from CSRtrack. + + Parameters + ---------- + file_path : str + Path to the file with particle data + + Returns + ------- + ParticleDistribution + The particle distribution. + """ + # Read data. + data = np.genfromtxt(file_path) + + # Extract phase space, + z = data[1:, 0] + x = data[1:, 1] + y = data[1:, 2] + pz = data[1:, 3] / (ct.m_e*ct.c**2/ct.e) + px = data[1:, 4] / (ct.m_e*ct.c**2/ct.e) + py = data[1:, 5] / (ct.m_e*ct.c**2/ct.e) + q = data[1:, 6] + + # Apply reference particle. + x[1:] += x[0] + y[1:] += y[0] + z[1:] += z[0] + px[1:] += px[0] + py[1:] += py[0] + pz[1:] += pz[0] + + # Determine charge and mass of particle species (only electrons in + # CSRtrack). + q_species = -ct.e + m_species = ct.m_e + + # Get particle weights. + w = q / q_species + + # Return distribution. + return ParticleDistribution(x, y, z, px, py, pz, w, q_species, m_species) + + +def read_from_openpmd( + file_path: str, + species_name: Optional[str] = None +) -> ParticleDistribution: + """Read particle distribution from ASTRA. + + Parameters + ---------- + file_path : str + Path to the file with particle data + species_name : str, Optional + Name of the particle species. Optional if only one particle species + is available in the openpmd file. + + Returns + ------- + ParticleDistribution + The particle distribution. + """ + # Open file. + file_content = File(file_path, mode='r') + + # Get base path in file. + iteration = list(file_content['/data'].keys())[0] + base_path = '/data/{}'.format(iteration) + + # Get path under which particle data is stored. + particles_path = file_content.attrs['particlesPath'].decode() + + # Get list of available species. + available_species = list( + file_content[join_infile_path(base_path, particles_path)]) + assert len(available_species) > 0, ( + "No particle species found in '{}'.".format(file_path)) + + # It not specified, read first species. + if species_name is None: + species_name = available_species[0] + + # Get species. + beam_species = file_content[ + join_infile_path(base_path, particles_path, species_name)] + + # Get phase space and attributes. + mass = beam_species['mass'] + charge = beam_species['charge'] + position = beam_species['position'] + position_off = beam_species['positionOffset'] + momentum = beam_species['momentum'] + m_species = mass.attrs['value'] * mass.attrs['unitSI'] + q_species = charge.attrs['value'] * charge.attrs['unitSI'] + x = (position['x'][:] * position['x'].attrs['unitSI'] + + position_off['x'].attrs['value'] * position_off['x'].attrs['unitSI']) + y = (position['y'][:] * position['y'].attrs['unitSI'] + + position_off['y'].attrs['value'] * position_off['y'].attrs['unitSI']) + z = (position['z'][:] * position['z'].attrs['unitSI'] + + position_off['z'].attrs['value'] * position_off['z'].attrs['unitSI']) + px = momentum['x'][:] * momentum['x'].attrs['unitSI'] / (m_species*ct.c) + py = momentum['y'][:] * momentum['y'].attrs['unitSI'] / (m_species*ct.c) + pz = momentum['z'][:] * momentum['z'].attrs['unitSI'] / (m_species*ct.c) + w = beam_species['weighting'][:] + + # Return distribution. + return ParticleDistribution(x, y, z, px, py, pz, w, q_species, m_species) + + +_readers = { + 'astra': read_from_astra, + 'csrtrack': read_from_csrtrack, + 'openpmd': read_from_openpmd +} diff --git a/aptools/particle_distributions/save.py b/aptools/particle_distributions/save.py new file mode 100644 index 0000000..bcd5b89 --- /dev/null +++ b/aptools/particle_distributions/save.py @@ -0,0 +1,307 @@ +"""Defines methods for saving particle distributions to file.""" + +from typing import Optional + +import numpy as np +import scipy.constants as ct +from openpmd_api import (Series, Access, Dataset, Mesh_Record_Component, + Unit_Dimension) + +from aptools import __version__ +from .particle_distribution import ParticleDistribution + + +SCALAR = Mesh_Record_Component.SCALAR + + +def save_distribution( + distribution: ParticleDistribution, + file_path: str, + data_format: str, + **kwargs +) -> None: + """Save particle distribution to file in the specified format. + + Parameters + ---------- + distribution : ParticleDistribution + The particle distribution to save. + file_path : str + Path to the file in which to save the data. + data_format : str + Internal format of the data. Possible values + are 'astra', 'csrtrack' and 'openpmd'. + + Other Parameters + ---------------- + **kwargs + Additional parameters to be passed to the particle savers. + """ + _savers[data_format](distribution, file_path, **kwargs) + + +def save_to_astra( + distribution: ParticleDistribution, + file_path: str +) -> None: + """Save particle distribution to text file in ASTRA format. + + Parameters + ---------- + distribution : ParticleDistribution + The particle distribution to save. + file_path : str + Path to the file in which to save the data. + """ + + # Get beam data + m_species = distribution.m_species + q_species = distribution.q_species + x_orig = distribution.x + y_orig = distribution.y + z_orig = distribution.z + px_orig = distribution.px * m_species * ct.c**2 / ct.e # eV/c + py_orig = distribution.py * m_species * ct.c**2 / ct.e # eV/c + pz_orig = distribution.pz * m_species * ct.c**2 / ct.e # eV/c + q_orig = distribution.q * 1e9 # nC + w = distribution.w + + # Determine particle index (type of species). + if m_species == ct.m_e: + index = 1 if q_species == -ct.e else 2 + if m_species == ct.m_p: + index == 3 + + # Create arrays + x = np.zeros(q_orig.size + 1) + y = np.zeros(q_orig.size + 1) + z = np.zeros(q_orig.size + 1) + px = np.zeros(q_orig.size + 1) + py = np.zeros(q_orig.size + 1) + pz = np.zeros(q_orig.size + 1) + q = np.zeros(q_orig.size + 1) + + # Reference particle + x[0] = np.average(x_orig, weights=w) + y[0] = np.average(y_orig, weights=w) + z[0] = np.average(z_orig, weights=w) + px[0] = np.average(px_orig, weights=w) + py[0] = np.average(py_orig, weights=w) + pz[0] = np.average(pz_orig, weights=w) + q[0] = np.sum(q_orig) / len(q_orig) + + # Put relative to reference particle + x[1::] = x_orig + y[1::] = y_orig + z[1::] = z_orig - z[0] + px[1::] = px_orig + py[1::] = py_orig + pz[1::] = pz_orig - pz[0] + q[1::] = q_orig + t = z / ct.c + + # Add flags and indices + ind = np.ones(q.size) * index + flag = np.ones(q.size) * 5 + + # Save to file + data = np.column_stack((x, y, z, px, py, pz, t, q, ind, flag)) + np.savetxt( + file_path, + data, + '%1.12e %1.12e %1.12e %1.12e %1.12e %1.12e %1.12e %1.12e %i %i' + ) + + +def save_to_csrtrack( + distribution: ParticleDistribution, + file_path: str +) -> None: + """Save particle distribution to text file using the CSRtrack fmt1 format. + + Parameters + ---------- + distribution : ParticleDistribution + The particle distribution to save. + file_path : str + Path to the file in which to save the data. + """ + + # Get beam data. + x_orig = distribution.x + y_orig = distribution.y + z_orig = distribution.z + px_orig = distribution.px * ct.m_e*ct.c**2/ct.e + py_orig = distribution.py * ct.m_e*ct.c**2/ct.e + pz_orig = distribution.pz * ct.m_e*ct.c**2/ct.e + q_orig = distribution.q + + # Create arrays. + x = np.zeros(q_orig.size+2) + y = np.zeros(q_orig.size+2) + z = np.zeros(q_orig.size+2) + px = np.zeros(q_orig.size+2) + py = np.zeros(q_orig.size+2) + pz = np.zeros(q_orig.size+2) + q = np.zeros(q_orig.size+2) + + # Reference particle. + x[1] = np.average(x_orig, weights=q_orig) + y[1] = np.average(y_orig, weights=q_orig) + z[1] = np.average(z_orig, weights=q_orig) + px[1] = np.average(px_orig, weights=q_orig) + py[1] = np.average(py_orig, weights=q_orig) + pz[1] = np.average(pz_orig, weights=q_orig) + q[1] = sum(q_orig)/len(q_orig) + + # Relative coordinates. + x[2::] = x_orig - x[1] + y[2::] = y_orig - y[1] + z[2::] = z_orig - z[1] + px[2::] = px_orig - px[1] + py[2::] = py_orig - py[1] + pz[2::] = pz_orig - pz[1] + q[2::] = q_orig + + # Save to file. + data = np.column_stack((z, x, y, pz, px, py, q)) + if not file_path.endswith('.fmt1'): + file_path += '.fmt1' + np.savetxt( + file_path, + data, + '%1.12e %1.12e %1.12e %1.12e %1.12e %1.12e %1.12e' + ) + + +def save_to_openpmd( + distribution: ParticleDistribution, + file_path: str, + species_name: Optional[str] = 'particle_distribution' +) -> None: + """ + Save particle distribution to an HDF5 file following the openPMD standard. + + Parameters + ---------- + distribution : ParticleDistribution + The particle distribution to save. + file_path : str + Path to the file in which to save the data. + species_name : str + Optional. Name under which the particle species should be stored. + + """ + # Get beam data + x = np.ascontiguousarray(distribution.x) + y = np.ascontiguousarray(distribution.y) + z = np.ascontiguousarray(distribution.z) + px = np.ascontiguousarray(distribution.px) + py = np.ascontiguousarray(distribution.py) + pz = np.ascontiguousarray(distribution.pz) + w = np.ascontiguousarray(distribution.w) + q_species = distribution.q_species + m_species = distribution.m_species + + # Save to file + if not file_path.endswith('.h5'): + file_path += '.h5' + opmd_series = Series(file_path, Access.create) + + # Set basic attributes. + opmd_series.set_software('APtools', __version__) + opmd_series.set_particles_path('particles') + + # Create iteration + it = opmd_series.iterations[0] + + # Create particles species. + particles = it.particles[species_name] + + # Create additional necessary arrays and constants. + px = px * m_species * ct.c + py = py * m_species * ct.c + pz = pz * m_species * ct.c + + # Generate datasets. + d_x = Dataset(x.dtype, extent=x.shape) + d_y = Dataset(y.dtype, extent=y.shape) + d_z = Dataset(z.dtype, extent=z.shape) + d_px = Dataset(px.dtype, extent=px.shape) + d_py = Dataset(py.dtype, extent=py.shape) + d_pz = Dataset(pz.dtype, extent=pz.shape) + d_w = Dataset(w.dtype, extent=w.shape) + d_q = Dataset(np.dtype('float64'), extent=[1]) + d_m = Dataset(np.dtype('float64'), extent=[1]) + d_xoff = Dataset(np.dtype('float64'), extent=[1]) + d_yoff = Dataset(np.dtype('float64'), extent=[1]) + d_zoff = Dataset(np.dtype('float64'), extent=[1]) + + # Record data. + particles['position']['x'].reset_dataset(d_x) + particles['position']['y'].reset_dataset(d_y) + particles['position']['z'].reset_dataset(d_z) + particles['positionOffset']['x'].reset_dataset(d_xoff) + particles['positionOffset']['y'].reset_dataset(d_yoff) + particles['positionOffset']['z'].reset_dataset(d_zoff) + particles['momentum']['x'].reset_dataset(d_px) + particles['momentum']['y'].reset_dataset(d_py) + particles['momentum']['z'].reset_dataset(d_pz) + particles['weighting'][SCALAR].reset_dataset(d_w) + particles['charge'][SCALAR].reset_dataset(d_q) + particles['mass'][SCALAR].reset_dataset(d_m) + + # Prepare for writting. + particles['position']['x'].store_chunk(x) + particles['position']['y'].store_chunk(y) + particles['position']['z'].store_chunk(z) + particles['positionOffset']['x'].make_constant(0.) + particles['positionOffset']['y'].make_constant(0.) + particles['positionOffset']['z'].make_constant(0.) + particles['momentum']['x'].store_chunk(px) + particles['momentum']['y'].store_chunk(py) + particles['momentum']['z'].store_chunk(pz) + particles['weighting'][SCALAR].store_chunk(w) + particles['charge'][SCALAR].make_constant(q_species) + particles['mass'][SCALAR].make_constant(m_species) + + # Set units. + particles['position'].unit_dimension = {Unit_Dimension.L: 1} + particles['positionOffset'].unit_dimension = {Unit_Dimension.L: 1} + particles['momentum'].unit_dimension = { + Unit_Dimension.L: 1, + Unit_Dimension.M: 1, + Unit_Dimension.T: -1, + } + particles['charge'].unit_dimension = { + Unit_Dimension.T: 1, + Unit_Dimension.I: 1, + } + particles['mass'].unit_dimension = {Unit_Dimension.M: 1} + + # Set weighting attributes. + particles['position'].set_attribute('macroWeighted', np.uint32(0)) + particles['positionOffset'].set_attribute( + 'macroWeighted', np.uint32(0)) + particles['momentum'].set_attribute('macroWeighted', np.uint32(0)) + particles['weighting'][SCALAR].set_attribute( + 'macroWeighted', np.uint32(1)) + particles['charge'][SCALAR].set_attribute( + 'macroWeighted', np.uint32(0)) + particles['mass'][SCALAR].set_attribute('macroWeighted', np.uint32(0)) + particles['position'].set_attribute('weightingPower', 0.) + particles['positionOffset'].set_attribute('weightingPower', 0.) + particles['momentum'].set_attribute('weightingPower', 1.) + particles['weighting'][SCALAR].set_attribute('weightingPower', 1.) + particles['charge'][SCALAR].set_attribute('weightingPower', 1.) + particles['mass'][SCALAR].set_attribute('weightingPower', 1.) + + # Flush data. + opmd_series.flush() + + +_savers = { + 'astra': save_to_astra, + 'csrtrack': save_to_csrtrack, + 'openpmd': save_to_openpmd +} diff --git a/setup.cfg b/setup.cfg index 7521279..a0bc9f7 100644 --- a/setup.cfg +++ b/setup.cfg @@ -24,3 +24,4 @@ install_requires = scipy h5py openpmd-api~=0.14.0 + deprecated