Factorize boxlib particle handling even more

yt/frontends/boxlib/data_structures.py

@@ -21,7 +21,7 @@
     NyxFieldInfo,
     WarpXFieldInfo,
 )
-from .misc import BoxlibParticleReaderMixin
+from .misc import BoxlibReadParticleFileMixin
 
 # This is what we use to find scientific notation that might include d's
 # instead of e's.
@@ -922,7 +922,7 @@ def relative_refinement(self, l0, l1):
         return self.refine_by ** (l1 - l0 + offset)
 
 
-class OrionHierarchy(BoxlibHierarchy, BoxlibParticleReaderMixin):
+class OrionHierarchy(BoxlibHierarchy, BoxlibReadParticleFileMixin):
     def __init__(self, ds, dataset_type="orion_native"):
         BoxlibHierarchy.__init__(self, ds, dataset_type)
         self._read_particles()

yt/frontends/boxlib/io.py

@@ -4,10 +4,11 @@
 import numpy as np
 from misc import BoxlibParticleSelectionMixin
 
-from yt.frontends.chombo.io import parse_orion_sinks
 from yt.funcs import mylog
 from yt.utilities.io_handler import BaseIOHandler
 
+from .misc import IOHandlerParticlesBoxlibMixin
+
 
 def _remove_raw(all_fields, raw_fields):
     centered_fields = set(all_fields)
@@ -188,56 +189,12 @@ def _read_particle_fields(self, chunks, ptf, selector):
                                 yield (ptype, field), data[mask].flatten()
 
 
-class IOHandlerOrion(IOHandlerBoxlib):
+class IOHandlerOrion(IOHandlerBoxlib, IOHandlerParticlesBoxlibMixin):
     _dataset_type = "orion_native"
 
-    _particle_filename = None
-
     @property
     def particle_filename(self):
         fn = self.ds.output_dir + "/StarParticles"
         if not os.path.exists(fn):
             fn = self.ds.output_dir + "/SinkParticles"
-        self._particle_filename = fn
-        return self._particle_filename
-
-    _particle_field_index = None
-
-    @property
-    def particle_field_index(self):
-
-        index = parse_orion_sinks(self.particle_filename)
-
-        self._particle_field_index = index
-        return self._particle_field_index
-
-    def _read_particles(self, grid, field):
-        """
-        parses the Orion Star Particle text files
-
-        """
-
-        particles = []
-
-        if grid.NumberOfParticles == 0:
-            return np.array(particles)
-
-        def read(line, field):
-            entry = line.strip().split(" ")[self.particle_field_index[field]]
-            return float(entry)
-
-        try:
-            lines = self._cached_lines
-            for num in grid._particle_line_numbers:
-                line = lines[num]
-                particles.append(read(line, field))
-            return np.array(particles)
-        except AttributeError:
-            fn = self.particle_filename
-            with open(fn) as f:
-                lines = f.readlines()
-                self._cached_lines = lines
-                for num in grid._particle_line_numbers:
-                    line = lines[num]
-                    particles.append(read(line, field))
-            return np.array(particles)
+        return fn

yt/frontends/boxlib/misc.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 
+from yt.frontends.chombo.io import parse_orion_sinks
 from yt.geometry.selection_routines import GridSelector
 
 
@@ -31,7 +32,7 @@ def _read_particle_selection(self, chunks, selector, fields):
         return rv
 
 
-class BoxlibParticleReaderMixin:
+class BoxlibReadParticleFileMixin:
     def _read_particle_file(self, fn):
         """actually reads the orion particle data file itself."""
         if not os.path.exists(fn):
@@ -97,3 +98,48 @@ def velocity(field, data):
             )
 
         super().setup_particle_fields(ptype)
+
+
+class IOHandlerParticlesBoxlibMixin:
+    def _read_particles(self, grid, field):
+        """
+        parses the Orion Star Particle text files
+
+        """
+
+        particles = []
+
+        if grid.NumberOfParticles == 0:
+            return np.array(particles)
+
+        def read(line, field):
+            entry = line.strip().split(" ")[self.particle_field_index[field]]
+            return float(entry)
+
+        try:
+            lines = self._cached_lines
+            for num in grid._particle_line_numbers:
+                line = lines[num]
+                particles.append(read(line, field))
+            return np.array(particles)
+        except AttributeError:
+            fn = self.particle_filename
+            with open(fn) as f:
+                lines = f.readlines()
+                self._cached_lines = lines
+                for num in grid._particle_line_numbers:
+                    line = lines[num]
+                    particles.append(read(line, field))
+            return np.array(particles)
+
+    _particle_field_index = None
+
+    @property
+    def particle_field_index(self):
+        if self._particle_field_index:
+            return self._particle_field_index
+
+        index = parse_orion_sinks(self.particle_filename)
+
+        self._particle_field_index = index
+        return self._particle_field_index

yt/frontends/chombo/data_structures.py

@@ -6,7 +6,7 @@
 
 from yt.data_objects.index_subobjects.grid_patch import AMRGridPatch
 from yt.data_objects.static_output import Dataset
-from yt.frontends.boxlib.misc import BoxlibParticleReaderMixin
+from yt.frontends.boxlib.misc import BoxlibReadParticleFileMixin
 from yt.funcs import mylog, setdefaultattr
 from yt.geometry.grid_geometry_handler import GridIndex
 from yt.utilities.file_handler import HDF5FileHandler, warn_h5py
@@ -576,7 +576,7 @@ def _is_valid(cls, filename, *args, **kwargs):
         return False
 
 
-class Orion2Hierarchy(ChomboHierarchy, BoxlibParticleReaderMixin):
+class Orion2Hierarchy(ChomboHierarchy, BoxlibReadParticleFileMixin):
     def __init__(self, ds, dataset_type="orion_chombo_native"):
         ChomboHierarchy.__init__(self, ds, dataset_type)
 

yt/frontends/chombo/io.py

@@ -2,7 +2,10 @@
 
 import numpy as np
 
-from yt.frontends.boxlib.misc import BoxlibParticleSelectionMixin
+from yt.frontends.boxlib.misc import (
+    BoxlibParticleSelectionMixin,
+    IOHandlerParticlesBoxlibMixin,
+)
 from yt.geometry.selection_routines import GridSelector
 from yt.utilities.io_handler import BaseIOHandler
 from yt.utilities.logger import ytLogger as mylog
@@ -236,48 +239,11 @@ def parse_orion_sinks(fn):
     return index
 
 
-class IOHandlerOrion2HDF5(IOHandlerChomboHDF5):
+class IOHandlerOrion2HDF5(IOHandlerChomboHDF5, IOHandlerParticlesBoxlibMixin):
     _dataset_type = "orion_chombo_native"
 
-    _particle_field_index = None
+    _particle_filename = None
 
     @property
-    def particle_field_index(self):
-
-        fn = self.ds.fullplotdir[:-4] + "sink"
-
-        index = parse_orion_sinks(fn)
-
-        self._particle_field_index = index
-        return self._particle_field_index
-
-    def _read_particles(self, grid, field):
-        """
-        parses the Orion Star Particle text files
-
-        """
-
-        particles = []
-
-        if grid.NumberOfParticles == 0:
-            return np.array(particles)
-
-        def read(line, field):
-            entry = line.strip().split(" ")[self.particle_field_index[field]]
-            return float(entry)
-
-        try:
-            lines = self._cached_lines
-            for num in grid._particle_line_numbers:
-                line = lines[num]
-                particles.append(read(line, field))
-            return np.array(particles)
-        except AttributeError:
-            fn = grid.ds.fullplotdir[:-4] + "sink"
-            with open(fn) as f:
-                lines = f.readlines()
-                self._cached_lines = lines
-                for num in grid._particle_line_numbers:
-                    line = lines[num]
-                    particles.append(read(line, field))
-            return np.array(particles)
+    def particle_filename(self):
+        return self.ds.fullplotdir[:-4] + "sink"