Working pyharp binding (#107)

RadiationBand can act independently now

cmake/ktable_earth.cmake

@@ -7,16 +7,17 @@ macro(SET_IF_EMPTY _variable)
 endmacro()
 
 # athena variables
-set_if_empty(NUMBER_GHOST_CELLS 2)
-set_if_empty(NVAPOR 0)
+set_if_empty(NUMBER_GHOST_CELLS 0)
+set_if_empty(NVAPOR 1)
 
 # canoe variables
 set_if_empty(NCLOUD 0)
 set_if_empty(NTRACER 2)
 
 # canoe configure
 set(CMAKE_INSTALL_PREFIX "$ENV{HOME}/opt/")
-set(HYDROSTATIC ON)
+# set(HYDROSTATIC OFF)
 set(RFM ON)
 set(NETCDF ON)
+set(DISORT ON)
 set(PYTHON_BINDINGS ON)

python/pyharp.cpp

@@ -1,4 +1,5 @@
 // pybind
+#include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
@@ -20,13 +21,18 @@
 
 namespace py = pybind11;
 
-void init_index_map(ParameterInput *pin) { IndexMap::InitFromAthenaInput(pin); }
+void subscribe_species(std::map<std::string, std::vector<std::string>> smap) {
+  IndexMap::InitFromSpeciesMap(smap);
+}
 
 PYBIND11_MODULE(pyharp, m) {
   m.attr("__name__") = "pyharp";
   m.doc() = "Python bindings for harp module";
 
-  m.def("init_index_map", &init_index_map);
+  m.def("subscribe_species", &subscribe_species);
+
+  // MeshBlock
+  py::class_<MeshBlock>(m, "meshblock");
 
   // Radiation
   py::class_<Radiation>(m, "radiation")
@@ -52,14 +58,92 @@ PYBIND11_MODULE(pyharp, m) {
       .def(py::init<std::string, YAML::Node>())
 
       .def("resize", &RadiationBand::Resize, py::arg("nc1"), py::arg("nc2") = 1,
-           py::arg("nc3") = 1, py::arg("nstr") = 8)
+           py::arg("nc3") = 1, py::arg("nstr") = 4)
       .def("resize_solver", &RadiationBand::ResizeSolver, py::arg("nlyr"),
-           py::arg("nstr") = 1, py::arg("nuphi") = 1, py::arg("numu") = 1)
+           py::arg("nstr") = 4, py::arg("nuphi") = 1, py::arg("numu") = 1)
       .def("get_num_spec_grids", &RadiationBand::GetNumSpecGrids)
       .def("get_num_absorbers", &RadiationBand::GetNumAbsorbers)
+      .def("absorbers", &RadiationBand::Absorbers)
       .def("get_absorber", &RadiationBand::GetAbsorber)
       .def("get_absorber_by_name", &RadiationBand::GetAbsorberByName)
       .def("get_range", &RadiationBand::GetRange)
+      .def("cal_radiance", &RadiationBand::CalBandRadiance,
+           py::arg("pmb") = nullptr, py::arg("k") = 0, py::arg("j") = 0)
+
+      .def("get_toa",
+           [](RadiationBand &band) {
+             py::array_t<Real> ndarray(
+                 {band.GetNumSpecGrids(), band.GetNumOutgoingRays()},
+                 band._GetToaPtr());
+             return ndarray;
+           })
+
+      .def("get_tau",
+           [](RadiationBand &band) {
+             py::array_t<Real> ndarray(
+                 {band.GetNumSpecGrids(), band.GetNumLayers()},
+                 band._GetTauPtr());
+             return ndarray;
+           })
+
+      .def("set_spectral_properties",
+           [](RadiationBand &band,
+              std::map<std::string, std::vector<double>> &atm) {
+             auto pindex = IndexMap::GetInstance();
+
+             int nlayer = atm["HGT"].size();
+
+             AirColumn ac(nlayer);
+             std::vector<Real> x1v(nlayer), x1f(nlayer + 1);
+
+             for (int i = 0; i < nlayer; ++i) {
+               ac[i].SetType(AirParcel::Type::MoleFrac);
+               x1v[i] = atm["HGT"][i] * 1e3;  // km -> m
+               ac[i].w[IDN] = atm["TEM"][i];
+               ac[i].w[IPR] = atm["PRE"][i] * 100.;  // mbar -> Pa
+             }
+
+             for (int i = 1; i < nlayer; ++i) {
+               x1f[i] = 0.5 * (x1v[i - 1] + x1v[i]);
+             }
+
+             x1f[0] = x1v[0] - (x1v[1] - x1v[0]) / 2.;
+             x1f[nlayer] =
+                 x1v[nlayer - 1] + (x1v[nlayer - 1] - x1v[nlayer - 2]) / 2.;
+
+             for (auto const &pair : atm) {
+               if (pair.first == "HGT" || pair.first == "TEM" ||
+                   pair.first == "PRE") {
+                 continue;
+               }
+
+               if (pindex->HasVapor(pair.first))
+                 for (int i = 0; i < nlayer; ++i) {
+                   ac[i].w[pindex->GetVaporId(pair.first)] =
+                       pair.second[i] * 1e-6;
+                 }
+
+               if (pindex->HasCloud(pair.first))
+                 for (int i = 0; i < nlayer; ++i) {
+                   ac[i].c[pindex->GetCloudId(pair.first)] =
+                       pair.second[i] * 1e-6;
+                 }
+
+               if (pindex->HasChemistry(pair.first))
+                 for (int i = 0; i < nlayer; ++i) {
+                   ac[i].q[pindex->GetChemistryId(pair.first)] =
+                       pair.second[i] * 1e-6;
+                 }
+
+               if (pindex->HasTracer(pair.first))
+                 for (int i = 0; i < nlayer; ++i) {
+                   ac[i].x[pindex->GetTracerId(pair.first)] =
+                       pair.second[i] * 1e-6;
+                 }
+             }
+
+             band.SetSpectralProperties(ac, x1f.data());
+           })
 
       .def("__str__", [](RadiationBand &band) {
         std::stringstream ss;

src/exchanger/linear_exchanger.hpp

@@ -17,8 +17,13 @@
 
 template <typename T>
 LinearExchanger<T>::LinearExchanger() {
+#ifdef MPI_PARALLEL
   color_.resize(Globals::nranks);
   brank_.resize(Globals::nranks);
+#else
+  color_.resize(1);
+  brank_.resize(1);
+#endif  // MPI_PARALLEL
 }
 
 template <typename T>
@@ -35,6 +40,7 @@ int LinearExchanger<T>::GetRankInGroup() const {
 template <typename T>
 void LinearExchanger<T>::Regroup(MeshBlock const *pmb,
                                  CoordinateDirection dir) {
+#ifdef MPI_PARALLEL
   NeighborBlock bblock, tblock;
   ExchangerHelper::find_neighbors(pmb, dir, &bblock, &tblock);
 
@@ -67,7 +73,6 @@ void LinearExchanger<T>::Regroup(MeshBlock const *pmb,
     }
   }
 
-#ifdef MPI_PARALLEL
   MPI_Allgather(&bblock.snb.rank, 1, MPI_INT, brank_.data(), 1, MPI_INT,
                 MPI_COMM_WORLD);
 #else

src/harp/radiation.cpp

@@ -81,6 +81,8 @@ Radiation::Radiation(MeshBlock *pmb, ParameterInput *pin) {
     // set band toa
     int n = 0;
     for (auto &p : bands_) {
+      //! Delete the old array and initialize with a shallow slice
+      p->btoa.DeleteAthenaArray();
       p->btoa.InitWithShallowSlice(radiance, 3, n, p->GetNumOutgoingRays());
       n += p->GetNumOutgoingRays();
     }
@@ -91,8 +93,13 @@ Radiation::Radiation(MeshBlock *pmb, ParameterInput *pin) {
   flxdn.NewAthenaArray(bands_.size(), ncells3, ncells2, ncells1 + 1);
 
   for (int n = 0; n < bands_.size(); ++n) {
+    //! Delete the old array and initialize with a shallow slice
+    bands_[n]->bflxup.DeleteAthenaArray();
     bands_[n]->bflxup.InitWithShallowSlice(flxup, 4, n, 1);
-    bands_[n]->bflxup.InitWithShallowSlice(flxdn, 4, n, 1);
+
+    //! Delete the old array and initialize with a shallow slice
+    bands_[n]->bflxdn.DeleteAthenaArray();
+    bands_[n]->bflxdn.InitWithShallowSlice(flxdn, 4, n, 1);
   }
 
   // time control
@@ -128,10 +135,11 @@ void Radiation::CalRadiativeFlux(MeshBlock const *pmb, int k, int j, int il,
   AirColumn &&ac = AirParcelHelper::gather_from_primitive(pmb, k, j);
 
   Real grav = -pmb->phydro->hsrc.GetG1();
+  Real H0 = pmb->pcoord->GetPressureScaleHeight();
 
   for (auto &p : bands_) {
     // iu ~= ie + 1
-    p->SetSpectralProperties(ac, pcoord, grav, k, j);
+    p->SetSpectralProperties(ac, pcoord->x1f.data(), grav * H0, k, j);
     p->CalBandFlux(pmb, k, j, il, iu);
   }
 }
@@ -145,10 +153,11 @@ void Radiation::CalRadiance(MeshBlock const *pmb, int k, int j) {
   AirColumn &&ac = AirParcelHelper::gather_from_primitive(pmb, k, j);
 
   Real grav = -pmb->phydro->hsrc.GetG1();
+  Real H0 = pmb->pcoord->GetPressureScaleHeight();
 
   for (auto &p : bands_) {
     // iu ~= ie + 1
-    p->SetSpectralProperties(ac, pcoord, grav, k, j);
+    p->SetSpectralProperties(ac, pcoord->x1f.data(), grav * H0, k, j);
     p->CalBandRadiance(pmb, k, j);
   }
 }
@@ -192,19 +201,26 @@ size_t Radiation::LoadRestartData(char *psrc) {
 }
 
 namespace RadiationHelper {
+Direction parse_radiation_direction(std::string_view str) {
+  Direction ray;
+  ray.phi = 0.;
+
+  sscanf(str.data(), "(%lf,%lf)", &ray.mu, &ray.phi);
+  ray.mu = cos(deg2rad(ray.mu));
+  ray.phi = deg2rad(ray.phi);
+
+  return ray;
+}
+
 std::vector<Direction> parse_radiation_directions(std::string str) {
   std::vector<std::string> dstr = Vectorize<std::string>(str.c_str());
   int nray = dstr.size();
 
   std::vector<Direction> ray(nray);
 
   auto jt = dstr.begin();
-  for (auto it = ray.begin(); it != ray.end(); ++it, ++jt) {
-    it->phi = 0.;
-    sscanf(jt->c_str(), "(%lf,%lf)", &it->mu, &it->phi);
-    it->mu = cos(deg2rad(it->mu));
-    it->phi = deg2rad(it->phi);
-  }
+  for (auto it = ray.begin(); it != ray.end(); ++it, ++jt)
+    *it = parse_radiation_direction(*jt);
 
   return ray;
 }
@@ -238,7 +254,7 @@ uint64_t parse_radiation_flags(std::string str) {
     } else if (dstr[i] == "write_bin_radiance") {
       flags |= RadiationFlags::WriteBinRadiance;
     } else {
-      msg << "flag:" << dstr[i] << "unrecognized" << std::endl;
+      msg << "flag: '" << dstr[i] << "' unrecognized" << std::endl;
       throw RuntimeError("parse_radiation_flags", msg.str());
     }
 

src/harp/radiation.hpp

@@ -4,6 +4,7 @@
 // C/C++ headers
 #include <memory>
 #include <string>
+#include <string_view>
 #include <vector>
 
 // athena
@@ -100,6 +101,7 @@ const uint64_t WriteBinRadiance = 1LL << 8;
 }  // namespace RadiationFlags
 
 namespace RadiationHelper {
+Direction parse_radiation_direction(std::string_view str);
 std::vector<Direction> parse_radiation_directions(std::string str);
 uint64_t parse_radiation_flags(std::string str);
 void get_phase_momentum(Real *pmom, int iphas, Real gg, int npmom);

src/harp/radiation_band.cpp

@@ -54,16 +54,38 @@ RadiationBand::RadiationBand(std::string myname, YAML::Node const &rad)
   wrange_ = pgrid_->ReadRangeFrom(my);
 
   if (my["outdir"]) {
-    std::string dirstr = my["outdir"].as<std::string>();
-    rayOutput_ = RadiationHelper::parse_radiation_directions(dirstr);
+    if (!my["outdir"].IsSequence()) {
+      throw RuntimeError("RadiationBand", "outdir must be a sequence");
+    }
+
+    for (const auto &item : my["outdir"]) {
+      rayOutput_.push_back(
+          RadiationHelper::parse_radiation_direction(item.as<std::string>()));
+    }
   }
 
   // set absorbers
   if (my["opacity"]) {
+    if (!my["opacity"].IsSequence()) {
+      throw RuntimeError("RadiationBand", "opacity must be a sequence");
+    }
+
     auto names = my["opacity"].as<std::vector<std::string>>();
     absorbers_ = AbsorberFactory::CreateFrom(names, GetName(), rad);
   }
 
+  // set flags
+  if (my["flags"]) {
+    if (!my["flags"].IsSequence()) {
+      throw RuntimeError("RadiationBand", "flags must be a sequence");
+    }
+
+    auto flag_strs = my["flags"].as<std::vector<std::string>>();
+    for (auto flag : flag_strs) {
+      SetFlag(RadiationHelper::parse_radiation_flags(flag));
+    }
+  }
+
   // set rt solver
   if (my["rt-solver"]) {
     psolver_ = CreateRTSolverFrom(my["rt-solver"].as<std::string>(), rad);
@@ -110,7 +132,9 @@ void RadiationBand::Resize(int nc1, int nc2, int nc3, int nstr) {
   bssa.NewAthenaArray(nc3, nc2, nc1);
   bpmom.NewAthenaArray(nstr + 1, nc3, nc2, nc1);
 
-  //! \note btoa, bflxup, bflxdn are shallow slices to Radiation variables
+  btoa.NewAthenaArray(rayOutput_.size(), nc3, nc2);
+  bflxup.NewAthenaArray(nc3, nc2, nc1 + 1);
+  bflxdn.NewAthenaArray(nc3, nc2, nc1 + 1);
 }
 
 void RadiationBand::ResizeSolver(int nlyr, int nstr, int nuphi, int numu) {
@@ -129,8 +153,8 @@ AbsorberPtr RadiationBand::GetAbsorberByName(std::string const &name) {
   return nullptr;
 }
 
-void RadiationBand::CalBandFlux(MeshBlock const *pmb, int k, int j, int il,
-                                int iu) {
+RadiationBand const *RadiationBand::CalBandFlux(MeshBlock const *pmb, int k,
+                                                int j, int il, int iu) {
   // reset flux of this column
   for (int i = il; i <= iu; ++i) {
     bflxup(k, j, i) = 0.;
@@ -139,9 +163,12 @@ void RadiationBand::CalBandFlux(MeshBlock const *pmb, int k, int j, int il,
 
   psolver_->Prepare(pmb, k, j);
   psolver_->CalBandFlux(pmb, k, j, il, iu);
+
+  return this;
 }
 
-void RadiationBand::CalBandRadiance(MeshBlock const *pmb, int k, int j) {
+RadiationBand const *RadiationBand::CalBandRadiance(MeshBlock const *pmb, int k,
+                                                    int j) {
   // reset radiance of this column
   for (int n = 0; n < GetNumOutgoingRays(); ++n) {
     btoa(n, k, j) = 0.;
@@ -150,6 +177,8 @@ void RadiationBand::CalBandRadiance(MeshBlock const *pmb, int k, int j) {
 
   psolver_->Prepare(pmb, k, j);
   psolver_->CalBandRadiance(pmb, k, j);
+
+  return this;
 }
 
 void RadiationBand::WriteAsciiHeader(OutputParameters const *pout) const {