WIP: Add support for reading OpenPMD files

NEWS.md

@@ -6,6 +6,10 @@
 - Python 3.12 is now supported, Python 3.8.1 is required for development.
 - `total_filesize` now errors out when directories are passed.
 
+#### Features
+- Added experimental support for reading OpenPMD files (with most of the
+  heavy-lifting done by @krishivbhatia)
+
 ## Version 1.4.0 (2 May 2023)
 
 #### General

pyproject.toml

@@ -29,6 +29,7 @@ motionpicture = ">=0.2.0, ^0"
 argcomplete = "^2.0.0"
 tikzplotlib = ">=0.9.13, ^0"
 py-expression-eval = ">=0.3.14, ^0"
+openpmd-api = ">=0.14.0"
 
 [tool.poetry.dev-dependencies]
 black = "^23.1.0"

tests/grid_functions/batman.it00000000.bp4/profiling.json

@@ -0,0 +1,3 @@
+[
+{ "rank": 0, "start": "Mon_Nov_27_11:15:33_2023", "threads": 1, "bytes": 2432122, "mkdir_mus": 336, "aggregation_mus": 0, "meta_sort_merge_mus": 94, "minmax_mus": 0, "memcpy_mus": 250, "buffering_mus": 3653, "transport_0": { "type": "File_POSIX", "close_mus": 148, "write_mus": 2606, "open_mus": 3693}, "transport_1": { "type": "File_POSIX", "close_mus": 0, "write_mus": 172, "open_mus": 3241} }
+]

tests/grid_functions/batman.it00000001.bp4/profiling.json

@@ -0,0 +1,3 @@
+[
+{ "rank": 0, "start": "Mon_Nov_27_11:15:33_2023", "threads": 1, "bytes": 2432122, "mkdir_mus": 333, "aggregation_mus": 0, "meta_sort_merge_mus": 66, "minmax_mus": 0, "memcpy_mus": 215, "buffering_mus": 1018, "transport_0": { "type": "File_POSIX", "close_mus": 163, "write_mus": 2213, "open_mus": 1891}, "transport_1": { "type": "File_POSIX", "close_mus": 0, "write_mus": 209, "open_mus": 6161} }
+]

tests/test_cactus_grid_functions.py

@@ -23,6 +23,7 @@
 import numpy as np
 
 from kuibit import cactus_grid_functions as cg
+from kuibit import cactus_ascii_utils as ca
 from kuibit import grid_data
 from kuibit import simdir as sd
 
@@ -58,6 +59,10 @@ def test_GridFunctionsDir_string_or_tuple(self):
     def test_contains(self):
         self.assertIn("xyz", self.gd)
 
+    def test_contains_OpenPMDVars(self):
+        vars3D = self.gd.xyz
+        self.assertIn("wavetoyx_u", vars3D)
+
     def test__getitem(self):
         self.assertIs(self.gd["xy"], self.gd._all_griddata[(0, 1)])
 
@@ -90,6 +95,7 @@ def test_total_filesize(self):
 class TestAllGridFunctions(unittest.TestCase):
     def setUp(self):
         self.gf = sd.SimDir("tests/grid_functions").gf.xy
+        self.gf_xyz = sd.SimDir("tests/grid_functions").gf.xyz
 
     def test__init(self):
         self.assertCountEqual(self.gf.dimension, (0, 1))
@@ -136,6 +142,8 @@ def test__init(self):
             ],
         )
 
+        self.assertCountEqual(list(self.gf._vars_openpmd_files.keys()), [])
+
         # Here we are not testing that files are correctly organized...
 
     def test_keys(self):
@@ -183,7 +191,7 @@ def test__getitem(self):
         with self.assertRaises(KeyError):
             self.gf["hey"]
 
-        # Test a variable from ASCII and one from HDF5
+        # Test a variable from ASCII, one from HDF5, and one from OpenPMD
 
         # We don't test the details, we just test that object is initialized
         # with the correct data
@@ -197,6 +205,20 @@ def test__getitem(self):
             filename, "illinoisgrmhd-grmhd_primitives_allbutbi.xy.h5"
         )
 
+        # OpenPMD
+        self.assertTrue(
+            isinstance(self.gf_xyz["wavetoyx_u"], cg.OneGridFunctionOpenPMD)
+        )
+        self.assertEqual(self.gf_xyz["wavetoyx_u"].var_name, "wavetoyx_u")
+        # self.gf_xyz['wavetoy-u']allfiles is a set with two elements, the
+        # two iterations
+        filenames = {
+            os.path.split(f)[-1] for f in self.gf_xyz["wavetoyx_u"].allfiles
+        }
+        self.assertCountEqual(
+            filenames, {"batman.it00000000.bp4", "batman.it00000001.bp4"}
+        )
+
         # ASCII
 
         # First, we test that a warning is emitted when we don't have ghost
@@ -254,6 +276,10 @@ def setUp(self):
         # There's only one file
         self.rho_star_file = self.rho_star.allfiles[0]
 
+        # OpenPMD
+        self.wavetoy = sd.SimDir("tests/grid_functions").gf.xyz["wavetoyx_u"]
+        self.wavetoy_file_it0 = sorted(self.wavetoy.allfiles)[0]
+
         # We are going to test all the methods of the baseclass with
         # HDF5 files
 
@@ -287,6 +313,28 @@ def test__properties_in_file(self):
             self.P._components_in_file(self.P_file, 2, 0), [0, 1]
         )
 
+        self.assertCountEqual(
+            self.wavetoy._iterations_in_file(self.wavetoy_file_it0), [0]
+        )
+        self.assertEqual(
+            self.wavetoy._min_iteration_in_file(self.wavetoy_file_it0), 0
+        )
+        self.assertEqual(
+            self.wavetoy._max_iteration_in_file(self.wavetoy_file_it0), 0
+        )
+        self.assertCountEqual(
+            self.wavetoy._ref_levels_in_file(self.wavetoy_file_it0, 0),
+            [0, 1],
+        )
+        self.assertCountEqual(
+            self.wavetoy._components_in_file(
+                self.wavetoy_file_it0,
+                0,
+                0,
+            ),
+            [0, 1, 2, 3],
+        )
+
     def test_restarts(self):
         # TODO: This test is not robust when we are dealing with only one file...
         #       Add a second file and rewrite tests.
@@ -299,6 +347,11 @@ def test_iterations(self):
         self.assertCountEqual(self.P.available_iterations, [0, 1, 2])
         self.assertCountEqual(self.P.iterations, [0, 1, 2])
 
+        self.assertEqual(self.wavetoy.min_iteration, 0)
+        self.assertEqual(self.wavetoy.max_iteration, 1)
+        self.assertCountEqual(self.wavetoy.available_iterations, [0, 1])
+        self.assertCountEqual(self.wavetoy.iterations, [0, 1])
+
         # Iteration not available
         with self.assertRaises(ValueError):
             self.P._files_with_iteration(3)
@@ -311,9 +364,14 @@ def test_times(self):
         self.assertCountEqual(self.P.available_times, [0, 0.25, 0.5])
         self.assertCountEqual(self.P.times, [0, 0.25, 0.5])
 
+        self.assertCountEqual(self.wavetoy.available_times, [0, 0.0078125])
+        self.assertCountEqual(self.wavetoy.times, [0, 0.0078125])
+
     def test_iteration_at_time(self):
         self.assertEqual(self.P.iteration_at_time(0.5), 2)
 
+        self.assertEqual(self.wavetoy.iteration_at_time(0.0078125), 1)
+
         # Time not available
         with self.assertRaises(ValueError):
             self.P.iteration_at_time(5)
@@ -489,6 +547,44 @@ def test_init(self):
         with self.assertRaises(RuntimeError):
             self.rho_star._parse_file("/tmp/wrongname")
 
+    def test_openpmd(self):
+        iteration = 0
+        component = 0
+
+        # with cg.openpmd_series(self.wavetoy_file_it0) as series:
+        #     variable = series.iterations[iteration].meshes["wavetoyx_state_lev00"]["wavetoyx_u"]
+        #     chunk = variable.available_chunks()[component]
+
+        #     # Do the actual reading
+        #     data = variable.load_chunk(
+        #         chunk.offset, chunk.extent
+        #     )
+
+        #     print("YO", chunk.extent)
+
+        #     series.flush()
+
+
+        # expected_grid = grid_data.UniformGrid(
+        #     [9, 65, 65],
+        #     x0=[-1.0, -1.0, -1.0],
+        #     x1=[1.0, 1.0, 1.0],
+        #     ref_level=0,
+        #     num_ghost=[0, 0, 0],
+        #     time=0,
+        #     iteration=0,
+        #     component=0,
+        # )
+
+        # expected_grid_data = grid_data.UniformGridData(expected_grid, data)
+
+        # self.assertEqual(
+        #     self.wavetoy._read_component_as_uniform_grid_data(
+        #         self.wavetoy_file_it0, 0, 0, 0
+        #     ),
+        #     expected_grid_data,
+        # )
+
     def test_clear_cache(self):
         # Read something
         self.P[0]
@@ -501,3 +597,53 @@ def test_clear_cache(self):
 
         # Check that we are clear
         self.assertIsNone(self.P.alldata[self.P_file][0][0][0])
+
+
+class TestOneGridFunctionOpenPMD(unittest.TestCase):
+    def setUp(self):
+        self.gf = sd.SimDir("tests/grid_functions").gf.xyz
+
+    def test__init(self):
+        self.assertCountEqual(self.gf.dimension, (0, 1, 2))
+        # Here we check that we indexed the correct variables. We must check
+        # the OpenPMD file
+
+        # There are five files including one OpenPMD file in the test folder:
+        # 1. illinoisgrmhd-grmhd_primitives_allbutbi.xyz.asc (ASCII one group)
+        # 2. rho_star.xyz.asc (ASCII one var)
+        # 3. illinoisgrmhd-grmhd_primitives_allbutbi.xyz.file_0.h5 (HDF5 one group)
+        # 4. illinoisgrmhd-grmhd_primitives_allbutbi.xyz.file_1.h5 (HDF5 one group)
+        # 5. batman.it00000000.bp4 (OpenPMD bp4)
+
+        # Here we can we find all the variables
+
+        # assert variables from OpenPMD bp4 file
+        self.assertCountEqual(
+            list(self.gf._vars_openpmd_files.keys()),
+            [
+                "wavetoyx_rho",
+                "wavetoyx_u",
+            ],
+        )
+
+        # assert variables from h5 file
+        self.assertCountEqual(
+            list(self.gf._vars_h5_files.keys()),
+            ["P", "rho_b", "vx", "vy", "vz"],
+        )
+
+        # assert variables from ascii file
+        self.assertCountEqual(
+            list(self.gf._vars_ascii_files.keys()),
+            ["P", "rho_b", "vx", "vy", "vz", "rho_star"],
+        )
+
+    def test_allfiles(self):
+        # This is a weak test, we are just testing how many files we have...
+        # There should be 6 files including the OpenPMD bp4 files
+        self.assertEqual(len(self.gf.allfiles), 6)
+
+    def test_multiple_mesh_refinement_levels(self):
+        # wavetouyx_u variable in the OpenPMD file has 2 mesh refinement levels
+        wavetoyx_u = self.gf.fields.wavetoyx_u
+        self.assertEqual(wavetoyx_u[0].refinement_levels, [0, 1])