Cargrid output

src/CMakeLists.txt

@@ -88,6 +88,7 @@ add_executable(
         outputs/io_wrapper.cpp
         outputs/outputs.cpp
         outputs/basetype_output.cpp
+        outputs/cartgrid.cpp
         outputs/derived_variables.cpp
         outputs/binary.cpp
         outputs/eventlog.cpp

src/outputs/cartgrid.cpp

@@ -0,0 +1,152 @@
+//========================================================================================
+// AthenaXXX astrophysical plasma code
+// Copyright(C) 2020 James M. Stone <[email protected]> and the Athena code team
+// Licensed under the 3-clause BSD License (the "LICENSE")
+//========================================================================================
+//! \file cartgrid.cpp
+//! \brief writes data on a Cartesian sub-grid in binary format
+
+#include <sys/stat.h>  // mkdir
+
+#include <cstdio> // snprintf
+#include <fstream>
+#include <string>
+#include <sstream>
+
+#include "athena.hpp"
+#include "globals.hpp"
+#include "mesh/mesh.hpp"
+#include "outputs.hpp"
+#include "parameter_input.hpp"
+#include "utils/cart_grid.hpp"
+
+CartesianGridOutput::CartesianGridOutput(ParameterInput *pin, Mesh *pm,
+                                         OutputParameters op)
+    : BaseTypeOutput(pin, pm, op) {
+  Real center[3], extent[3];
+  int numpoints[3];
+  bool is_cheb;
+
+  mkdir("cart", 0755);
+
+  center[0] = pin->GetOrAddReal(op.block_name, "center_x", 0.0);
+  center[1] = pin->GetOrAddReal(op.block_name, "center_y", 0.0);
+  center[2] = pin->GetOrAddReal(op.block_name, "center_z", 0.0);
+  extent[0] = pin->GetOrAddReal(op.block_name, "extent_x", 2.0);
+  extent[1] = pin->GetOrAddReal(op.block_name, "extent_y", 2.0);
+  extent[2] = pin->GetOrAddReal(op.block_name, "extent_z", 2.0);
+  numpoints[0] = pin->GetOrAddInteger(op.block_name, "numpoints_x", 32);
+  numpoints[1] = pin->GetOrAddInteger(op.block_name, "numpoints_y", 32);
+  numpoints[2] = pin->GetOrAddInteger(op.block_name, "numpoints_z", 32);
+  is_cheb = pin->GetOrAddBoolean(op.block_name, "chebyshev", false);
+
+  pcart = new CartesianGrid(pm->pmb_pack, center, extent, numpoints, is_cheb);
+
+  for (int d = 0; d < 3; ++d) {
+    md.center[d] = center[d];
+    md.extent[d] = extent[d];
+    md.numpoints[d] = numpoints[d];
+  }
+  md.is_cheb = is_cheb;
+}
+
+CartesianGridOutput::~CartesianGridOutput() { delete pcart; }
+
+void CartesianGridOutput::LoadOutputData(Mesh *pm) {
+  // If AMR is enabled we need to reset the CartesianGrid
+  if (pm->adaptive) {
+    pcart->SetInterpolationIndices();
+    pcart->SetInterpolationWeights();
+  }
+
+  int nout_vars = outvars.size();
+  Kokkos::realloc(outarray, nout_vars, 1, md.numpoints[0], md.numpoints[1],
+                  md.numpoints[2]);
+
+  // Calculate derived variables, if required
+  if (out_params.contains_derived) {
+    ComputeDerivedVariable(out_params.variable, pm);
+  }
+
+  for (int n = 0; n < nout_vars; ++n) {
+    pcart->InterpolateToGrid(outvars[n].data_index, *(outvars[n].data_ptr));
+    auto v_slice =
+        Kokkos::subview(outarray, n, 0, Kokkos::ALL, Kokkos::ALL, Kokkos::ALL);
+    Kokkos::deep_copy(v_slice, pcart->interp_vals.h_view);
+  }
+#if MPI_PARALLEL_ENABLED
+  // Note that InterpolateToGrid will set zero values for points not owned by
+  // current rank
+  int count = nout_vars * md.numpoints[0] * md.numpoints[1] * md.numpoints[2];
+  if (0 == global_variable::my_rank) {
+    MPI_Reduce(MPI_IN_PLACE, outarray.data(), count, MPI_ATHENA_REAL, MPI_SUM,
+               0, MPI_COMM_WORLD);
+  } else {
+    MPI_Reduce(outarray.data(), outarray.data(), count, MPI_ATHENA_REAL,
+               MPI_SUM, 0, MPI_COMM_WORLD);
+  }
+#endif
+}
+
+void CartesianGridOutput::WriteOutputFile(Mesh *pm, ParameterInput *pin) {
+#if MPI_PARALLEL_ENABLED
+  if (0 == global_variable::my_rank) {
+#endif
+
+    // Assemble filename
+    char fname[BUFSIZ];
+    std::snprintf(fname, BUFSIZ, "cart/%s.%s.%05d.bin",
+                  out_params.file_basename.c_str(), out_params.file_id.c_str(),
+                  out_params.file_number);
+
+    // Open file
+    std::ofstream ofile(fname, std::ios::binary);
+
+    // Write metadata
+    md.cycle = pm->ncycle;
+    md.time = pm->time;
+    md.noutvars = outvars.size();
+    ofile.write(reinterpret_cast<char *>(&md), sizeof(MetaData));
+
+    // Write list of variables
+    {
+      std::stringstream msg;
+      for (int n = 0; n < md.noutvars - 1; ++n) {
+        msg << outvars[n].label << " ";
+      }
+      msg << outvars[md.noutvars - 1].label;
+      std::string smsg = msg.str();
+      int len = smsg.size();
+      ofile.write(reinterpret_cast<char *>(&len), sizeof(int));
+      ofile.write(smsg.c_str(), len);
+    }
+
+    // Write actual data
+    for (int n = 0; n < md.noutvars; ++n) {
+      for (int k = 0; k < md.numpoints[2]; ++k) {
+        for (int j = 0; j < md.numpoints[1]; ++j) {
+          for (int i = 0; i < md.numpoints[0]; ++i) {
+            // Note that we are accessing the array with the convention of
+            // CartesianGrid which is opposite from the one used in the rest of
+            // the code, but we write the output as k, j, i
+            float var = outarray(n, 0, i, j, k);
+            ofile.write(reinterpret_cast<char *>(&var), sizeof(float));
+          }
+        }
+      }
+    }
+
+#if MPI_PARALLEL_ENABLED
+  }
+#endif
+
+  // increment counters
+  out_params.file_number++;
+  if (out_params.last_time < 0.0) {
+    out_params.last_time = pm->time;
+  } else {
+    out_params.last_time += out_params.dt;
+  }
+  pin->SetInteger(out_params.block_name, "file_number", out_params.file_number);
+  pin->SetReal(out_params.block_name, "last_time", out_params.last_time);
+}

src/outputs/outputs.cpp

@@ -272,6 +272,9 @@ Outputs::Outputs(ParameterInput *pin, Mesh *pm) {
       } else if (opar.file_type.compare("bin") == 0) {
         pnode = new MeshBinaryOutput(pin,pm,opar);
         pout_list.insert(pout_list.begin(),pnode);
+      } else if (opar.file_type.compare("cart") == 0) {
+        pnode = new CartesianGridOutput(pin,pm,opar);
+        pout_list.insert(pout_list.begin(),pnode);
       } else if (opar.file_type.compare("rst") == 0) {
       // Add restarts to the tail end of BaseTypeOutput list, so file counters for other
       // output types are up-to-date in restart file

src/outputs/outputs.hpp

@@ -384,6 +384,34 @@ class RestartOutput : public BaseTypeOutput {
   void WriteOutputFile(Mesh *pm, ParameterInput *pin) override;
 };
 
+// Forward declaration
+class CartesianGrid;
+
+//----------------------------------------------------------------------------------------
+//! \class CartesianGridOutput
+//  \brief derived BaseTypeOutput class for output on a Cartesian grid
+class CartesianGridOutput : public BaseTypeOutput {
+  struct MetaData {
+    int cycle;
+    float time;
+    float center[3];
+    float extent[3];
+    int numpoints[3];
+    bool is_cheb;
+    int noutvars;
+  };
+ public:
+  CartesianGridOutput(ParameterInput *pin, Mesh *pm, OutputParameters oparams);
+  ~CartesianGridOutput();
+  //! Interpolate the data on the Cartesian grid and handle MPI communication
+  void LoadOutputData(Mesh *pm) override;
+  //! Write the data to file
+  void WriteOutputFile(Mesh *pm, ParameterInput *pin) override;
+ private:
+  CartesianGrid *pcart;
+  MetaData md;
+};
+
 //----------------------------------------------------------------------------------------
 //! \class EventLogOutput
 //  \brief derived BaseTypeOutput class for event counter data

src/utils/cart_grid.cpp

@@ -74,8 +74,6 @@ CartesianGrid::CartesianGrid(MeshBlockPack *pmy_pack, Real center[3],
   return;
 }
 
-CartesianGrid::~CartesianGrid() {}
-
 void CartesianGrid::ResetCenter(Real center[3]) {
   // grid center
   center_x1 = center[0];
@@ -210,6 +208,11 @@ void CartesianGrid::SetInterpolationWeights() {
       Real x0 = min_x1 + nx * d_x1;
       Real y0 = min_x2 + ny * d_x2;
       Real z0 = min_x3 + nz * d_x3;
+      if (is_cheby) {
+        x0 = center_x1 + extent_x1*std::cos(nx*M_PI/(nx1-1));
+        y0 = center_x2 + extent_x2*std::cos(ny*M_PI/(nx2-1));
+        z0 = center_x3 + extent_x3*std::cos(nz*M_PI/(nx3-1));
+      }
       // extract MeshBlock bounds
       Real &x1min = size.h_view(ii0).x1min;
       Real &x1max = size.h_view(ii0).x1max;

src/utils/cart_grid.hpp

@@ -21,8 +21,7 @@ class CartesianGrid {
  public:
   // Creates a geodesic grid with refinement level nlev and radius rad
   CartesianGrid(MeshBlockPack *pmy_pack, Real center[3],
-                    Real extent[3], int numpoints[3], bool is_cheb = false);
-  ~CartesianGrid();
+                Real extend[3], int numpoints[3], bool is_cheb = false);
 
   // parameters for the grid
   Real center_x1, center_x2, center_x3;   // grid centers
@@ -39,14 +38,14 @@ class CartesianGrid {
   DualArray3D<Real> interp_vals;   // container for data interpolated to sphere
   void InterpolateToGrid(int nvars, DvceArray5D<Real> &val);  // interpolate to sphere
   void ResetCenter(Real center[3]);  // set indexing for interpolation
+  void SetInterpolationIndices();      // set indexing for interpolation
+  void SetInterpolationWeights();      // set weights for interpolation
   void ResetCenterAndExtent(Real center[3], Real extent[3]);
 
  private:
   MeshBlockPack* pmy_pack;  // ptr to MeshBlockPack containing this Hydro
   DualArray4D<int> interp_indcs;   // indices of MeshBlock and zones therein for interp
   DualArray5D<Real> interp_wghts;  // weights for interpolation
-  void SetInterpolationIndices();      // set indexing for interpolation
-  void SetInterpolationWeights();      // set weights for interpolation
 };
 
 #endif // UTILS_CART_GRID_HPP_

vis/python/cartgrid.py

@@ -0,0 +1,70 @@
+"""
+Reads CartesianGrid binary format
+"""
+
+import numpy as np
+from struct import Struct
+
+
+class CartesianGridData:
+    """Class representing a single CartesianGrid dump
+
+    Members are:
+    * cycle:int           simulation cycle
+    * time:float          simulation time
+    * center[3]:float     box center
+    * extent[3]:float     box extent
+    * numpoints[3]:int    number of grid points in each direction
+    * is_cheb:bool        Chebyshev grid
+    * variables           dictionary with all grid functions
+    """
+
+    def __init__(self, fname):
+        mdata = Struct("i7f3i?2i")
+        with open(fname, "rb") as f:
+            # Parse metadata
+            blob = f.read(mdata.size)
+            ncycle, time, cx, cy, cz, ex, ey, ez, nx, ny, nz, cheb, nout, nstr = (
+                mdata.unpack(blob)
+            )
+            self.cycle = ncycle
+            self.time = time
+            self.center = (cx, cy, cz)
+            self.extent = (ex, ey, ez)
+            self.numpoints = (nx, ny, nz)
+            self.is_cheb = cheb
+            self.variables = {}
+
+            # Read variable names
+            blob = f.read(nstr).decode("ascii")
+            names = blob.split()
+            if len(names) != nout:
+                raise IOError(f"Could not read {fname}")
+
+            # Read variables
+            for n in names:
+                self.variables[n] = (
+                    np.fromfile(f, dtype=np.float32, count=np.prod(self.numpoints))
+                    .reshape(self.numpoints)
+                    .transpose()
+                )
+
+    def coords(self, d=None):
+        """Returns the coordinates"""
+        if d is None:
+            return self.coords(0), self.coords(1), self.coords(2)
+        if self.is_cheb:
+            from math import pi
+
+            return self.center[d] + self.extent[d] * np.cos(
+                np.linspace(0, pi, self.numpoints[d])
+            )
+        else:
+            return self.center[d] + self.extent[d] * np.linspace(
+                -1, 1, self.numpoints[d]
+            )
+
+    def meshgrid(self):
+        """Returns a mesh grid with all the coordinates"""
+        x, y, z = self.coords()
+        return np.meshgrid(x, y, z, indexing="ij")