Merge pull request #5276 from Yurlungur/jmm/parthenon-frontend-updates

Fix small bugs/edge cases in the parthenon frontend

Author: chrishavlin

yt/frontends/parthenon/data_structures.py

@@ -53,10 +53,6 @@ def _setup_dx(self):
         id = self.id - self._id_offset
         LE, RE = self.index.grid_left_edge[id, :], self.index.grid_right_edge[id, :]
         self.dds = self.ds.arr((RE - LE) / self.ActiveDimensions, "code_length")
-        if self.ds.dimensionality < 2:
-            self.dds[1] = 1.0
-        if self.ds.dimensionality < 3:
-            self.dds[2] = 1.0
         self.field_data["dx"], self.field_data["dy"], self.field_data["dz"] = self.dds
 
     def retrieve_ghost_zones(self, n_zones, fields, all_levels=False, smoothed=False):
@@ -164,19 +160,13 @@ def __init__(
 
         self.geometry = _geom_map[self._handle["Info"].attrs["Coordinates"]]
 
-        if self.geometry is Geometry.CYLINDRICAL:
-            axis_order = ("r", "theta", "z")
-        else:
-            axis_order = None
-
         Dataset.__init__(
             self,
             filename,
             dataset_type,
             units_override=units_override,
             unit_system=unit_system,
             default_species_fields=default_species_fields,
-            axis_order=axis_order,
         )
         if storage_filename is None:
             storage_filename = self.basename + ".yt"

yt/frontends/parthenon/io.py

@@ -52,7 +52,10 @@ def _read_fluid_selection(self, chunks, selector, fields, size):
                 for gs in grid_sequences(chunk.objs):
                     start = gs[0].id - gs[0]._id_offset
                     end = gs[-1].id - gs[-1]._id_offset + 1
-                    data = ds[start:end, fdi, :, :, :].transpose()
+                    if len(ds.shape) == 4:
+                        data = ds[start:end, :, :, :].transpose()
+                    else:
+                        data = ds[start:end, fdi, :, :, :].transpose()
                     for i, g in enumerate(gs):
                         ind += g.select(selector, data[..., i], rv[field], ind)
             last_dname = dname
@@ -72,7 +75,10 @@ def _read_chunk_data(self, chunk, fields):
             for gs in grid_sequences(chunk.objs):
                 start = gs[0].id - gs[0]._id_offset
                 end = gs[-1].id - gs[-1]._id_offset + 1
-                data = ds[start:end, fdi, :, :, :].transpose()
+                if len(ds.shape) == 4:
+                    data = ds[start:end, :, :, :].transpose()
+                else:
+                    data = ds[start:end, fdi, :, :, :].transpose()
                 for i, g in enumerate(gs):
                     rv[g.id][field] = np.asarray(data[..., i], "=f8")
         return rv

yt/frontends/parthenon/tests/test_outputs.py

@@ -179,3 +179,25 @@ def test_load_cylindrical():
     # Check that the domain edges match r in [0.5,2.0], theta in [0, 2pi]
     assert_equal(ds.domain_left_edge.in_units("code_length").v[:2], (0.5, 0))
     assert_equal(ds.domain_right_edge.in_units("code_length").v[:2], (2.0, 2 * np.pi))
+
+
+# Sedov blast wave with curvlinear coords run with RIOT
+riot_sedov_curvlinear = "riot_sedov_curvlinear/sedov.out1.final.phdf"
+
+
+@requires_file(riot_sedov_curvlinear)
+def test_load_riot_curvilinear():
+    # Load a cylindrical dataset of a full disk
+    ds = data_dir_load(riot_sedov_curvlinear)
+
+    assert ("parthenon", "c.c.bulk.pressure") in ds.field_list
+
+    ad = ds.all_data()
+    dth = ad["index", "dtheta"]
+    vol = ad["index", "cell_volume"]
+    rho = ad["parthenon", "c.c.bulk.rho"]
+
+    assert np.all(np.abs(dth - 2 * np.pi) <= 1e-12)
+
+    total_mass = (rho * vol).sum()
+    assert np.all(np.abs(total_mass.value - 0.169646) <= 1e-8)