add PlotFileData bindings (#320)

* add PlotFileData bindings

* Update PlotFileUtil.cpp

* Update PlotFileUtil.cpp

* fix bugs

* fix bindings for overloaded function

* add unit test for simple single-level 3D plotfile

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix box, string asserts

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix lint error

* add script to plot 2d plotfiles

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Apply suggestions from code review

Co-authored-by: Axel Huebl <[email protected]>

* Tests: Import pytest

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Import amr

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* generate test data on-the-fly

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove extra import

* remove extra import (again)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Axel Huebl <[email protected]>

Author: 3 people

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/test_plotfiledata.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+import pytest
+
+import amrex.space3d as amr
+
+
+def write_test_plotfile(filename):
+    """Write single-level plotfile (in order to read it back in)."""
+    domain_box = amr.Box([0, 0, 0], [31, 31, 31])
+    real_box = amr.RealBox([-0.5, -0.5, -0.5], [0.5, 0.5, 0.5])
+    geom = amr.Geometry(domain_box, real_box, amr.CoordSys.cartesian, [0, 0, 0])
+
+    ba = amr.BoxArray(domain_box)
+    dm = amr.DistributionMapping(ba, 1)
+    mf = amr.MultiFab(ba, dm, 1, 0)
+    mf.set_val(np.pi)
+
+    time = 1.0
+    level_step = 200
+    var_names = amr.Vector_string(["density"])
+    amr.write_single_level_plotfile(filename, mf, var_names, geom, time, level_step)
+
+
+@pytest.mark.skipif(amr.Config.spacedim != 3, reason="Requires AMREX_SPACEDIM = 3")
+def test_plotfiledata_read():
+    """Generate and then read plotfile using PlotFileUtil bindings."""
+    plt_filename = "test_plt00200"
+    write_test_plotfile(plt_filename)
+    plt = amr.PlotFileData(plt_filename)
+
+    assert plt.spaceDim() == 3
+    assert plt.time() == 1.0
+    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.small_end == amr.IntVect(0, 0, 0)
+    assert probDomain.big_end == amr.IntVect(31, 31, 31)
+
+    assert probSize == [1.0, 1.0, 1.0]
+    assert probLo == [-0.5, -0.5, -0.5]
+    assert probHi == [0.5, 0.5, 0.5]
+    assert cellSize == [1.0 / 32.0, 1.0 / 32.0, 1.0 / 32.0]
+    assert varNames == amr.Vector_string(["density"])
+    assert nComp == 1
+    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 == (32, 32, 32, 1)
+            assert np.all(marr_xp[:, :, :, :] == np.pi)
+            nboxes += 1
+
+        assert nboxes == 1

tools/plot_plotfile_2d.py

@@ -0,0 +1,57 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+
+import amrex.space2d 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):
+    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
+
+    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)