Merge pull request #66 from epochpic/multiple_probes

Fixed mutiple probes not working in the same file

Author: JoelLucaAdams

docs/key_functionality.rst

@@ -51,6 +51,12 @@ done via the ``preprocess`` parameter.
 Reading particle data
 ~~~~~~~~~~~~~~~~~~~~~
 
+.. warning::
+   It is **not recommended** to use :func:`xarray.open_mfdataset` or :func:`open_mfdataset` to load particle data from multiple SDF outputs. The number of particles often varies between outputs, which can lead to inconsistent array shapes that these functions cannot handle. Instead, consider loading each file individually and then concatenating them manually.
+
+.. note::
+   When loading multiple probes from a single SDF file, you **must** use the `probe_names` parameter to assign a unique name to each. For example, use `probe_names=["Front_Electron_Probe", "Back_Electron_Probe"]`. Failing to do so will result in dimension name conflicts.
+
 By default, particle data isn't kept as it takes up a lot of space.
 Pass ``keep_particles=True`` as a keyword argument to
 :func:`xarray.open_dataset` (for single files) or :func:`xarray.open_mfdataset` (for

src/sdf_xarray/__init__.py

@@ -69,6 +69,7 @@ def open_mfdataset(
     *,
     separate_times: bool = False,
     keep_particles: bool = False,
+    probe_names: list[str] | None = None,
 ) -> xr.Dataset:
     """Open a set of EPOCH SDF files as one `xarray.Dataset`
 
@@ -98,6 +99,8 @@ def open_mfdataset(
         different output frequencies
     keep_particles :
         If ``True``, also load particle data (this may use a lot of memory!)
+    probe_names :
+        List of EPOCH probe names
     """
 
     # TODO: This is not very robust, look at how xarray.open_mfdataset does it
@@ -108,10 +111,15 @@ def open_mfdataset(
     path_glob = sorted(list(path_glob))  # noqa: C414
 
     if not separate_times:
-        return combine_datasets(path_glob, keep_particles=keep_particles)
+        return combine_datasets(
+            path_glob, keep_particles=keep_particles, probe_names=probe_names
+        )
 
     time_dims, var_times_map = make_time_dims(path_glob)
-    all_dfs = [xr.open_dataset(f, keep_particles=keep_particles) for f in path_glob]
+    all_dfs = [
+        xr.open_dataset(f, keep_particles=keep_particles, probe_names=probe_names)
+        for f in path_glob
+    ]
 
     for df in all_dfs:
         for da in df:
@@ -211,14 +219,23 @@ class SDFDataStore(AbstractDataStore):
         "drop_variables",
         "keep_particles",
         "lock",
+        "probe_names",
     )
 
-    def __init__(self, manager, drop_variables=None, keep_particles=False, lock=None):
+    def __init__(
+        self,
+        manager,
+        drop_variables=None,
+        keep_particles=False,
+        lock=None,
+        probe_names=None,
+    ):
         self._manager = manager
         self._filename = self.ds.filename
         self.drop_variables = drop_variables
         self.keep_particles = keep_particles
         self.lock = ensure_lock(lock)
+        self.probe_names = probe_names
 
     @classmethod
     def open(
@@ -227,6 +244,7 @@ def open(
         lock=None,
         drop_variables=None,
         keep_particles=False,
+        probe_names=None,
     ):
         if isinstance(filename, os.PathLike):
             filename = os.fspath(filename)
@@ -237,6 +255,7 @@ def open(
             lock=lock,
             drop_variables=drop_variables,
             keep_particles=keep_particles,
+            probe_names=probe_names,
         )
 
     def _acquire(self, needs_lock=True):
@@ -347,7 +366,28 @@ def _process_grid_name(grid_name: str, transform_func) -> str:
 
             if value.is_point_data:
                 # Point (particle) variables are 1D
-                var_coords = (f"ID_{_process_grid_name(key, _grid_species_name)}",)
+
+                # Particle data does not maintain a fixed dimension size
+                # throughout the simulation. An example of a particle name comes
+                # in the form of `Particles/Px/Ion_H` which is then modified
+                # using `_process_grid_name()` into `Ion_H`. This is fine as the
+                # other components of the momentum (`Py`, `Pz`) will have the same
+                # size as they represent the same bunch of particles.
+
+                # Probes however have names in the form of `Electron_Front_Probe/Px`
+                # which are changed to just `Px`; this is fine when there is only one
+                # probe in the system but when there are multiple they will have
+                # conflicting sizes so we can't keep the names as simply `Px` so we
+                # instead set their dimension as the full name `Electron_Front_Probe_Px`.
+                is_probe_name_match = self.probe_names is not None and any(
+                    name in key for name in self.probe_names
+                )
+                name_processor = (
+                    _rename_with_underscore
+                    if is_probe_name_match
+                    else _grid_species_name
+                )
+                var_coords = (f"ID_{_process_grid_name(key, name_processor)}",)
             else:
                 # These are DataArrays
 
@@ -414,6 +454,7 @@ def open_dataset(
         *,
         drop_variables=None,
         keep_particles=False,
+        probe_names=None,
     ):
         if isinstance(filename_or_obj, Path):
             # sdf library takes a filename only
@@ -424,6 +465,7 @@ def open_dataset(
             filename_or_obj,
             drop_variables=drop_variables,
             keep_particles=keep_particles,
+            probe_names=probe_names,
         )
         with close_on_error(store):
             return store.load()
@@ -432,6 +474,7 @@ def open_dataset(
         "filename_or_obj",
         "drop_variables",
         "keep_particles",
+        "probe_names",
     ]
 
     def guess_can_open(self, filename_or_obj):

tests/example_two_probes_2D/0000.sdf

@@ -0,0 +1,188 @@
+
+### USER DEFINED CONSTANTS ###
+begin:constant
+  # User input constants
+  n_elec = 6e+29 #average electron number density
+  I_peak = 1e+28 #peak intensity of the laser
+
+  scale_length = 0.5e-6 #scale length of the target
+  L_target_x = 1.5e-6 #length of the target
+  L_target_y = 10e-6 #width of the target
+
+  lambda_L = 800e-9 #wavelength of the laser
+  y_fwhm_L = 2e-6 #intensity fwhm in y
+  tau_fwhm_L = 16e-15 #intensity fwhm in t
+
+  mppc = 100 #macro particles per cell
+end:constant
+
+begin:control
+  nx = 150
+  ny = 200
+  nsteps = -1
+  t_end = 200e-15
+  x_min = -10e-06
+  x_max =  5e-06
+  y_min = -10e-06
+  y_max = -y_min
+  nparticles = nint(((L_target_x * L_target_y) / ((x_max - x_min) * (y_max - y_min))) * nx * ny * mppc)
+  dt_multiplier = 0.8
+  smooth_currents = T
+  dlb_threshold = 0.8
+  stdout_frequency = 100 # Print ETA
+  particle_tstart = 30e-15 # Time before the laser hits the target
+end:control
+
+### CALCULATD CONSTANTS ###
+begin:constant
+  w_t = tau_fwhm_L / (2*sqrt(loge(2)))           # Beam waist in t
+  w_y = y_fwhm_L / (2*sqrt(loge(2)))             # Beam waist in y
+  t_hw01m = 0.5 * tau_fwhm_L * sqrt(loge(10)/loge(2))
+
+  x_foc = -x_min                                 # Boundary to focal point distance
+  k_L = 2.0 * pi / lambda_L                      # Laser wave number
+  x_Rr = pi * w_y^2 / lambda_L                   # Rayleigh range
+  omega_L = 2 * pi * c / lambda_L                # Angular frequency of the laser
+
+  w_boundary = w_y * sqrt(1 + (x_foc/x_Rr)^2)    # Waist on boundary
+  I_boundary = I_peak * (w_y / w_boundary)^2     # Intens. on boundary
+  r_curve = x_foc * (1 + (x_Rr/x_foc)^2)         # Boundary curv. rad.
+  phase_Gouy = atan(-x_foc/r_curve)              # Boundary Gouy shift
+end:constant
+
+begin:laser
+  boundary = x_min
+  intensity = I_boundary
+  omega = omega_L
+  phase = k_L * y^2 / (2 * r_curve) - phase_Gouy
+  profile = gauss(y, 0, w_boundary)
+  t_profile = gauss(time, t_hw01m, w_t) 
+  t_start = 0.0
+  t_end = end
+end:laser
+
+begin:qed
+  use_qed = F
+  qed_start_time = 0
+  produce_photons = T
+  photon_energy_min = 1 * kev
+  produce_pairs = T
+  photon_dynamics = T
+end:qed
+
+begin:collisions
+  use_collisions = T
+  coulomb_log = auto
+  collide = all
+end:collisions
+
+begin:boundaries
+  bc_x_min_particle = simple_outflow
+  bc_x_max_particle = simple_outflow
+  bc_y_min_particle = reflect
+  bc_y_max_particle = reflect
+  bc_x_min_field = simple_laser
+  bc_x_max_field = conduct
+  bc_y_min_field = conduct
+  bc_y_max_field = conduct
+end:boundaries
+
+begin:species
+  name = Electron
+  frac = 0.5
+  dump = T
+  temperature = 300
+  number_density  = if((abs(y) lt L_target_y/2), if((x lt 0) or (x gt L_target_x), 0, if((x gt 0) and (x lt scale_length), n_elec * (1-exp(x/scale_length))/(1-exp(1)), n_elec)), 0)
+  number_density_min = 1
+  identify:electron
+end:species
+
+begin:species
+  name = Ion_C
+  charge = 6
+  atomic_number = 12
+  mass = 1836.2 * 12
+  frac = 0.25
+  dump = T
+  number_density = number_density(Electron) * (6/7)
+  temperature = temperature_x(Electron)
+  number_density_min = 1
+end:species
+
+begin:species
+  name = Ion_H
+  charge = 1
+  atomic_number = 1
+  mass = 1836.2
+  fraction = 0.25
+  dump = T
+  number_density = number_density(Electron) / 7
+  temperature = temperature_x(Electron)
+  number_density_min = 1
+  identify:proton
+end:species
+
+begin:species
+  name = Photon
+  nparticles = 0
+  dump = T
+  identify:photon
+end:species
+
+begin:species
+  name = Positron
+  nparticles = 0
+  dump = T
+  identify:positron
+end:species
+
+### THE MANY PARTICLE PROBES ###
+
+#Electron#
+begin:probe
+  name = Electron_Front_Probe
+  point = (x_min, 0)
+  normal = (-1.0, 0.0)
+  ek_min = 0.0
+  ek_max = -1.0
+  include_species : Electron
+  dumpmask = always
+end:probe
+begin:probe
+  name = Electron_Back_Probe
+  point = (x_max, 0)
+  normal = (1.0, 0.0)
+  ek_min = 0.0
+  ek_max = -1.0
+  include_species : Electron
+  dumpmask = always
+end:probe
+
+################################
+
+begin:output_global
+  force_final_to_be_restartable = F
+end:output_global
+
+begin:output
+  name = normal
+  grid = always
+  particle_probes = always
+  dt_snapshot = 15e-14
+  ex = always + single
+  ey = always + single
+  ez = always + single
+  bx = always + single
+  by = always + single
+  bz = always + single
+  particles = always
+  particle_weight = always
+  px = always + single
+  py = always + single
+  pz = always + single
+  average_particle_energy = always + species + no_sum + single
+  number_density = always + species + no_sum + single
+  absorption = always 
+  total_energy_sum = always + species
+end:output
+

tests/example_two_probes_2D/0001.sdf

@@ -11,6 +11,7 @@
 )
 EXAMPLE_ARRAYS_DIR = pathlib.Path(__file__).parent / "example_array_no_grids"
 EXAMPLE_3D_DIST_FN = pathlib.Path(__file__).parent / "example_dist_fn"
+EXAMPLE_2D_PARTICLE_DATA = pathlib.Path(__file__).parent / "example_two_probes_2D"
 
 
 def test_basic():
@@ -226,3 +227,32 @@ def test_erroring_drop_variables():
         xr.open_dataset(
             EXAMPLE_FILES_DIR / "0000.sdf", drop_variables=["Electric_Field/E"]
         )
+
+
+def test_loading_multiple_probes():
+    with xr.open_dataset(
+        EXAMPLE_2D_PARTICLE_DATA / "0002.sdf",
+        keep_particles=True,
+        probe_names=["Electron_Front_Probe", "Electron_Back_Probe"],
+    ) as df:
+        assert "X_Probe_Electron_Front_Probe" in df.coords
+        assert "X_Probe_Electron_Back_Probe" in df.coords
+        assert "ID_Electron_Front_Probe_Px" in df.dims
+        assert "ID_Electron_Back_Probe_Px" in df.dims
+
+
+def test_loading_one_probe_drop_second_probe():
+    with xr.open_dataset(
+        EXAMPLE_2D_PARTICLE_DATA / "0002.sdf",
+        keep_particles=True,
+        drop_variables=[
+            "Electron_Back_Probe_Px",
+            "Electron_Back_Probe_Py",
+            "Electron_Back_Probe_Pz",
+            "Electron_Back_Probe_weight",
+        ],
+        probe_names=["Electron_Front_Probe"],
+    ) as df:
+        assert "X_Probe_Electron_Front_Probe" in df.coords
+        assert "ID_Electron_Front_Probe_Px" in df.dims
+        assert "ID_Electron_Back_Probe_Px" not in df.dims