add PlotFileData bindings

src/Base/PlotFileUtil.cpp

@@ -5,29 +5,47 @@
 #include "pyAMReX.H"
 
 #include <AMReX_PlotFileUtil.H>
-#include <AMReX_Vector.H>
 #include <AMReX_Print.H>
+#include <AMReX_Vector.H>
 
 #include <sstream>
 #include <string>
 
 namespace py = pybind11;
 using namespace amrex;
 
-void init_PlotFileUtil(py::module& m)
-{
-    m.def("write_single_level_plotfile",
-          &amrex::WriteSingleLevelPlotfile,
-          "Writes single level plotfile",
-          py::arg("plotfilename"),
-          py::arg("mf"),
-          py::arg("varnames"),
-          py::arg("geom"),
-          py::arg("time"),
-          py::arg("level_step"),
-          py::arg("versionName")="HyperCLaw-V1.1",
-          py::arg("levelPrefix")="Level_",
-          py::arg("mfPrefix")="Cell",
-          py::arg_v("extra_dirs", Vector<std::string>(), "list[str]")
-    );
+void init_PlotFileUtil(py::module &m) {
+  m.def("write_single_level_plotfile", &amrex::WriteSingleLevelPlotfile,
+        "Writes single level plotfile", py::arg("plotfilename"), py::arg("mf"),
+        py::arg("varnames"), py::arg("geom"), py::arg("time"),
+        py::arg("level_step"), py::arg("versionName") = "HyperCLaw-V1.1",
+        py::arg("levelPrefix") = "Level_", py::arg("mfPrefix") = "Cell",
+        py::arg_v("extra_dirs", Vector<std::string>(), "list[str]"));
+
+  py::class_<PlotFileData>(m, "PlotFileData")
+      // explicitly provide constructor argument types
+      .def(py::init<std::string const&>())
+
+      .def("spaceDim", &PlotFileData::spaceDim)
+      .def("time", &PlotFileData::time)
+      .def("finestLevel", &PlotFileData::finestLevel)
+      .def("refRatio", &PlotFileData::refRatio)
+      .def("levelStep", &PlotFileData::levelStep)
+      .def("boxArray", &PlotFileData::boxArray)
+      .def("DistributionMap", &PlotFileData::DistributionMap)
+      .def("syncDistributionMap", py::overload_cast<PlotFileData const&>(&PlotFileData::syncDistributionMap))
+      .def("syncDistributionMap", py::overload_cast<int, PlotFileData const&>(&PlotFileData::syncDistributionMap))
+
+      .def("coordSys", &PlotFileData::coordSys)
+      .def("probDomain", &PlotFileData::probDomain)
+      .def("probSize", &PlotFileData::probSize)
+      .def("probLo", &PlotFileData::probLo)
+      .def("probHi", &PlotFileData::probHi)
+      .def("cellSize", &PlotFileData::cellSize)
+      .def("varNames", &PlotFileData::varNames)
+      .def("nComp", &PlotFileData::nComp)
+      .def("nGrowVect", &PlotFileData::nGrowVect)
+
+      .def("get", py::overload_cast<int>(&PlotFileData::get))
+      .def("get", py::overload_cast<int, std::string const&>(&PlotFileData::get));
 }

tests/projz04000/Header

@@ -0,0 +1,22 @@
+HyperCLaw-V1.1
+3
+nH_wind
+nH_cloud
+nH
+3
+56649980960510.086
+0
+0 0 0
+2.4688e+21 6.172e+20 6.172e+20
+
+((0,0,0) (1023,255,0) (0,0,0))
+4000
+2.4109375e+18 2.4109375e+18 6.172e+20
+0
+0
+0 1 56649980960510.086
+4000
+0 2.4688e+21
+0 6.172e+20
+0 6.172e+20
+Level_0/Cell

tests/projz04000/Level_0/Cell_H

@@ -0,0 +1,15 @@
+1
+1
+3
+0
+(1 0
+((0,0,0) (1023,255,0) (0,0,0))
+)
+1
+FabOnDisk: Cell_D_00000 0
+
+1,3
+2.0621599037663667e+18,0.0000000000000000e+00,2.0621599037663667e+18,
+
+1,3
+2.0681399821329768e+19,5.5521812980561425e+20,5.6820568202245806e+20,

tests/test_plotfiledata.py

@@ -0,0 +1,57 @@
+# -*- coding: utf-8 -*-
+
+import os
+from pathlib import Path
+
+import pytest
+
+import amrex.space3d as amr
+
+
+@pytest.mark.skipif(amr.Config.spacedim != 3, reason="Requires AMREX_SPACEDIM = 3")
+def test_plotfiledata_read():
+    import amrex.space3d as amr
+
+    parent_path = Path(os.path.abspath(__file__)).parent
+    plt = amr.PlotFileData(str(parent_path / "projz04000"))
+
+    assert plt.spaceDim() == 3
+    assert plt.time() == 56649980960510.086
+    assert plt.finestLevel() == 0
+    assert plt.refRatio(0) == 0
+    assert plt.coordSys() == amr.CoordSys.cartesian
+
+    probDomain = plt.probDomain(0)
+    probSize = plt.probSize()
+    probLo = plt.probLo()
+    probHi = plt.probHi()
+    cellSize = plt.cellSize(0)
+    varNames = plt.varNames()
+    nComp = plt.nComp()
+    nGrowVect = plt.nGrowVect(0)
+
+    # assert probDomain == amr.Box([0,0,0], [1023,255,0]) # doesn't work
+    assert probDomain.small_end == amr.IntVect(0, 0, 0)
+    assert probDomain.big_end == amr.IntVect(1023, 255, 0)
+
+    assert probSize == [2.4688e21, 6.172e20, 6.172e20]
+    assert probLo == [0.0, 0.0, 0.0]
+    assert probHi == [2.4688e21, 6.172e20, 6.172e20]
+    assert cellSize == [2.4109375e18, 2.4109375e18, 6.172e20]
+    assert varNames == amr.Vector_string(["nH_wind", "nH_cloud", "nH"])
+    assert nComp == 3
+    assert nGrowVect == amr.IntVect(0, 0, 0)
+
+    for compname in varNames:
+        mfab_comp = plt.get(0, compname)
+        nboxes = 0
+
+        for mfi in mfab_comp:
+            marr = mfab_comp.array(mfi)
+            # numpy/cupy representation: non-copying view, including the
+            # guard/ghost region
+            marr_xp = marr.to_xp()
+            assert marr_xp.shape == (1024, 256, 1, 1)
+            nboxes += 1
+
+        assert nboxes == 1

tools/plot_plotfile_2d.py

@@ -0,0 +1,61 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+import amrex.space3d as amr
+
+
+def plot_mf(arr, compname, plo, phi):
+    plt.plot()
+    im = plt.imshow(
+        arr.T,
+        origin="lower",
+        interpolation="none",
+        extent=[plo[0], phi[0], plo[1], phi[1]],
+        aspect="equal",
+    )
+    plt.colorbar(im)
+    plt.tight_layout()
+    plt.savefig(f"{compname}.png", dpi=150)
+    plt.close()
+
+
+@pytest.mark.skipif(amr.Config.spacedim != 2, reason="Requires AMREX_SPACEDIM = 2")
+def plot_plotfile_2d(filename, level=0):
+    import amrex.space2d as amr
+
+    plt = amr.PlotFileData(filename)
+    assert level <= plt.finestLevel()
+
+    probDomain = plt.probDomain(level)
+    probLo = plt.probLo()
+    probHi = plt.probHi()
+    varNames = plt.varNames()
+
+    for compname in varNames:
+        mfab_comp = plt.get(level, compname)
+        arr = np.zeros((probDomain.big_end - probDomain.small_end + 1))
+        for mfi in mfab_comp:
+            bx = mfi.tilebox()
+            marr = mfab_comp.array(mfi)
+            marr_xp = marr.to_xp()
+            i_s, j_s = tuple(bx.small_end)
+            i_e, j_e = tuple(bx.big_end)
+            arr[i_s : i_e + 1, j_s : j_e + 1] = marr_xp[:, :, 0, 0]
+        plot_mf(arr, compname, probLo, probHi)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    import amrex.space2d as amr
+
+    parser = argparse.ArgumentParser(
+        "Plots each variable in a 2D plotfile using matplotlib."
+    )
+    parser.add_argument("filename", help="AMReX 2D plotfile to read")
+    parser.add_argument("level", type=int, help="AMR level to plot (default: 0)")
+    args = parser.parse_args()
+
+    amr.initialize([])
+    plot_plotfile_2d(args.filename, level=args.level)