Improve default behavior of openPMD particle reader (#72)

* Improve defaults for reading openpmd bunches

* Update test

* Increase verson number

* Require latest APtools

* Fix bug

* Add tests for new functionality

* Remove whitespace

setup.cfg

@@ -23,7 +23,7 @@ install_requires =
     numpy
     numba
     scipy
-    aptools>=0.1.23
+    aptools>=0.1.25
     openpmd-api~=0.14.0
 python_requires = >=3.7
 

tests/test_bunch_read_save.py

@@ -1,11 +1,13 @@
 
 import os
 
+import pytest
 from numpy.testing import assert_array_almost_equal
 
 from wake_t.utilities.bunch_generation import (
     get_from_file, get_gaussian_bunch_from_size)
 from wake_t.utilities.bunch_saving import save_bunch_to_file
+from wake_t.diagnostics import OpenPMDDiagnostics
 
 
 tests_output_folder = './tests_output'
@@ -49,11 +51,7 @@ def test_bunch_read_save():
         file_path = os.path.join(
             output_folder,
             'bunch_{}.{}'.format(data_format, file_format))
-        if data_format in ['astra', 'csrtrack']:
-            bunch_saved = get_from_file(file_path, data_format)
-        if data_format in ['openpmd']:
-            bunch_saved = get_from_file(file_path, data_format,
-                                        species_name='elec_bunch')
+        bunch_saved = get_from_file(file_path, data_format)
 
         # For astra and csrtrack, remove the generated reference particle.
         if data_format in ['astra', 'csrtrack']:
@@ -75,5 +73,79 @@ def test_bunch_read_save():
         assert_array_almost_equal(bunch_saved.q, bunch.q)
 
 
+def test_openpmd_reading():
+    """
+    Check that the `get_from_file` method behaves as expected when reading
+    openPMD data.
+    """
+    output_folder = os.path.join(tests_output_folder, 'bunch_openpmd_read')
+    if not os.path.exists(output_folder):
+        os.makedirs(output_folder)
+
+    # Beam parameters.
+    emitt_nx = emitt_ny = 1e-6  # m
+    s_x = s_y = 3e-6  # m
+    s_t = 3.  # fs
+    gamma_avg = 100 / 0.511
+    gamma_spread = 1.  # %
+    q_bunch = 30  # pC
+    xi_avg = 0.  # m
+    n_part = 1e4
+
+    # Create particle bunch.
+    bunch_1 = get_gaussian_bunch_from_size(
+        emitt_nx, emitt_ny, s_x, s_y, gamma_avg, gamma_spread, s_t, xi_avg,
+        q_bunch, n_part, name='elec_bunch')
+    bunch_2 = get_gaussian_bunch_from_size(
+        emitt_nx, emitt_ny, s_x, s_y, gamma_avg, gamma_spread, s_t, xi_avg,
+        q_bunch, n_part, name='elec_bunch_2')
+
+    # Create diagnostics in output folder
+    diags = OpenPMDDiagnostics(write_dir=output_folder)
+
+    # Test 1.
+    # -------
+
+    # Write diagnostics without any species.
+    diags.write_diagnostics(time=0., dt=0.)
+    # Trying to read a species from this file should fail.
+    with pytest.raises(ValueError):
+        file_path = os.path.join(output_folder, 'hdf5', 'data00000000.h5')
+        bunch_saved = get_from_file(file_path, 'openpmd')
+
+    # Test 2.
+    # -------
+
+    # Write diagnostics with one species.
+    diags.write_diagnostics(time=0., dt=0., species_list=[bunch_1])
+    # Reading this file should raise no exception even if no species name
+    # is specified.
+    file_path = os.path.join(output_folder, 'hdf5', 'data00000001.h5')
+    bunch_saved = get_from_file(file_path, 'openpmd')
+    # Since the `name` parameter has not been specified, it must be equal to
+    # the original name of the openPMD species.
+    assert bunch_saved.name == 'elec_bunch'
+    # If a new custom name is given, use this name.
+    custom_name = 'bunch'
+    bunch_saved = get_from_file(file_path, 'openpmd', name=custom_name)
+    assert bunch_saved.name == custom_name
+
+    # Test 3.
+    # -------
+
+    # Write diagnostics with two species.
+    diags.write_diagnostics(time=0., dt=0., species_list=[bunch_1, bunch_2])
+    # Reading a species without specifying a species name should fail.
+    with pytest.raises(ValueError):
+        file_path = os.path.join(output_folder, 'hdf5', 'data00000002.h5')
+        bunch_saved = get_from_file(file_path, 'openpmd')
+    # Reading each species indicating the correct species name should not fail
+    bunch_saved = get_from_file(
+        file_path, 'openpmd', species_name='elec_bunch')
+    bunch_saved = get_from_file(
+        file_path, 'openpmd', species_name='elec_bunch_2')
+
+
 if __name__ == '__main__':
     test_bunch_read_save()
+    test_openpmd_reading()

wake_t/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.5.0'
+__version__ = '0.5.1'
 
 
 from .beamline_elements import (PlasmaStage, PlasmaRamp, ActivePlasmaLens,

wake_t/utilities/bunch_generation.py

@@ -5,6 +5,7 @@
 from scipy.stats import truncnorm
 import aptools.plasma_accel.general_equations as ge
 import aptools.data_handling.reading as dr
+from aptools.data_handling.utilities import get_available_species
 
 from wake_t.particles.particle_bunch import ParticleBunch
 
@@ -249,11 +250,72 @@ def get_matched_bunch(
 
 
 def get_from_file(file_path, code_name, preserve_prop_dist=False, name=None,
-                  **kwargs):
+                  species_name=None, **kwargs):
+    """Get particle bunch from file.
+
+    Parameters
+    ----------
+    file_path : str
+        Path to the file containing the particle distribution.
+    code_name : str
+        Name of the code from which the particle distribution originates.
+        Possible values are 'astra', 'csrtrack' and 'openpmd'.
+    preserve_prop_dist : bool, optional
+        Whether to preserve the average longitudinal position of the
+        distribution in the `prop_distance` parameter of the `ParticleBunch`,
+        by default False.
+    name : str, optional
+        Name to assign to the ParticleBunch, by default None. If None, when
+        reading a particle species from an `openPMD` file, the `name` is set
+        to the original name of the species. For other data formats, the
+        default name `elec_bunch_i` will be assigned, where `i`>=0 is an
+        integer that is increased for each unnamed `ParticleBunch` that is
+        generated.
+    species_name : str, optional
+        Name of the particle species to be read from an `openpmd` file.
+        Required only when more than one particle species is present in the
+        file.
+
+    Other Parameters
+    ----------------
+    **kwargs : dict
+        Other parameters to be passed to
+        `aptools.data_handling.reading.read_beam`.
+
+    Returns
+    -------
+    ParticleBunch
+        A ParticleBunch with the distribution from the specified file.
+    """
+    # If reading from an openPMD file, use the right `name` and `species_name`.
+    if code_name == 'openpmd':
+        if species_name is None:
+            available_species = get_available_species(file_path)
+            n_species = len(available_species)
+            if n_species == 0:
+                raise ValueError(
+                    "No particle species found in '{}'".format(file_path)
+                )
+            elif n_species == 1:
+                species_name = available_species[0]
+            else:
+                raise ValueError(
+                    'More than one particle species is available in' +
+                    "'{}'. ".format(file_path) +
+                    'Please specify a `species_name`. ' +
+                    'Available species are: ' + str(available_species)
+                )
+        kwargs['species_name'] = species_name
+        if name is None:
+            name = species_name
+    # Read particle species.
     x, y, z, px, py, pz, q = dr.read_beam(code_name, file_path, **kwargs)
+    # Center in z.
     z_avg = np.average(z, weights=q)
     xi = z - z_avg
+    # Create ParticleBunch
     bunch = ParticleBunch(q, x, y, xi, px, py, pz, name=name)
+    # Preserve `z_avg` as the initial propagation distance of the bunch.
     if preserve_prop_dist:
         bunch.prop_distance = z_avg
     return bunch