BUG: fix ds.all_data() for spherical AMReX datasets

yt/frontends/amrex/data_structures.py

@@ -375,14 +375,9 @@ def _parse_index(self):
                 default_ybounds = (0.0, 1.0)
                 default_zbounds = (0.0, 1.0)
             case Geometry.CYLINDRICAL:
-                self.level_dds[:, 2] = 2 * np.pi
                 default_ybounds = (0.0, 1.0)
                 default_zbounds = (0.0, 2 * np.pi)
             case Geometry.SPHERICAL:
-                # BoxLib only supports 1D spherical, so ensure
-                # the other dimensions have the right extent.
-                self.level_dds[:, 1] = np.pi
-                self.level_dds[:, 2] = 2 * np.pi
                 default_ybounds = (0.0, np.pi)
                 default_zbounds = (0.0, 2 * np.pi)
             case _:
@@ -909,6 +904,12 @@ def _parse_header_file(self):
             dre = self.domain_right_edge.copy()
             dre[2] = 2.0 * np.pi
             self.domain_right_edge = dre
+        if self.geometry is Geometry.SPHERICAL and self.dimensionality < 3:
+            dre = self.domain_right_edge.copy()
+            dre[2] = 2.0 * np.pi
+            if self.dimensionality < 2:
+                dre[1] = np.pi
+            self.domain_right_edge = dre
 
         header_file.close()
 

yt/frontends/amrex/tests/test_outputs.py

@@ -443,3 +443,60 @@ def test_maestro_parameters():
     # Check an int parameter
     assert ds.parameters["s0_interp_type"] == 3
     assert type(ds.parameters["s0_interp_type"]) is int  # noqa: E721
+
+
+# test loading non-Cartesian coordinate systems in different dimensionalities
+
+
+def check_coordsys_data(ds):
+    # check that level_dds is consistent with domain_width
+    assert_allclose(
+        ds.index.level_dds[0] * ds.domain_dimensions,
+        ds.domain_width.to_value("code_length"),
+        rtol=1e-12,
+        atol=0.0,
+    )
+
+    # check that we get the expected number of data points when selecting the
+    # entire domain
+    expected_size = sum(np.count_nonzero(g.child_mask) for g in ds.index.grids)
+    ad = ds.all_data()
+    assert ad["boxlib", "Temp"].size == expected_size
+
+
+cyl_1d = "castro_sedov_1d_cyl_plt00150"
+cyl_2d = "castro_sedov_2d_sph_in_cyl_plt00130"
+sph_1d = "sedov_1d_sph_plt00120"
+sph_2d = "xrb_spherical_smallplt00010"
+
+
+@requires_file(cyl_1d)
+def test_coordsys_1d_cylindrical():
+    ds = data_dir_load(cyl_1d)
+    assert ds.geometry == "cylindrical"
+    assert ds.dimensionality == 1
+    check_coordsys_data(ds)
+
+
+@requires_file(cyl_2d)
+def test_coordsys_2d_cylindrical():
+    ds = data_dir_load(cyl_2d)
+    assert ds.geometry == "cylindrical"
+    assert ds.dimensionality == 2
+    check_coordsys_data(ds)
+
+
+@requires_file(sph_1d)
+def test_coordsys_1d_spherical():
+    ds = data_dir_load(sph_1d)
+    assert ds.geometry == "spherical"
+    assert ds.dimensionality == 1
+    check_coordsys_data(ds)
+
+
+@requires_file(sph_2d)
+def test_coordsys_2d_spherical():
+    ds = data_dir_load(sph_2d)
+    assert ds.geometry == "spherical"
+    assert ds.dimensionality == 2
+    check_coordsys_data(ds)

yt/sample_data_registry.json

@@ -646,6 +646,12 @@
     "load_name": null,
     "url": "https://yt-project.org/data/castro_sedov_2d_cyl_in_cart_plt00150.tar.gz"
   },
+  "castro_sedov_2d_sph_in_cyl_plt00130.tar.gz": {
+    "hash": "ef4d081a2a2f8e10afe132768725c573631b82021e91f07782f1c1fbe043e2b5",
+    "load_kwargs": {},
+    "load_name": null,
+    "url": "https://yt-project.org/data/castro_sedov_2d_sph_in_cyl_plt00130.tar.gz"
+  },
   "castro_sod_x_plt00036.tar.gz": {
     "hash": "3f0a586b41e7b54fa2b3cddd50f9384feb2efe1fe1a815e7348965ae7bf88f78",
     "load_kwargs": {},
@@ -946,6 +952,12 @@
     "load_name": "DD0045/DD0045.0.h5",
     "url": "https://yt-project.org/data/tiny_fof_halos.tar.gz"
   },
+  "xrb_spherical_smallplt00010.tar.gz": {
+    "hash": "27ede5ed03f7c89b2afac03a368beb56d5f25f0c7c95b81f14f071a54a795783",
+    "load_kwargs": {},
+    "load_name": null,
+    "url": "https://yt-project.org/data/xrb_spherical_smallplt00010.tar.gz"
+  },
   "ytdata_test.tar.gz": {
     "hash": "cafb2b06ab3190ba17909585b58a4724e25f27ac72f11d6dff1a482146eb8958",
     "load_kwargs": {},