Add a vector wave equation example

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Current develop
 
 ### Added (new features/APIs/variables/...)
+- [[PR 1091]](https://github.com/parthenon-hpc-lab/parthenon/pull/1091) Add vector wave equation example.
 - [[PR 991]](https://github.com/parthenon-hpc-lab/parthenon/pull/991) Add fine fields.
 - [[PR 1084]](https://github.com/parthenon-hpc-lab/parthenon/pull/1084) Properly free swarm boundary MPI requests 
 - [[PR 1020]](https://github.com/parthenon-hpc-lab/parthenon/pull/1020) Add bi- and trilinear interpolation routines

example/fine_advection/advection_driver.cpp

@@ -53,6 +53,19 @@ AdvectionDriver::AdvectionDriver(ParameterInput *pin, ApplicationInput *app_in,
   pin->CheckDesired("Advection", "derefine_tol");
 }
 
+template <class pack_desc_t>
+TaskID AddUpdateTasks(TaskID dep, TaskList &tl, parthenon::CellLevel cl,
+                      parthenon::TopologicalType tt, Real beta, Real dt,
+                      pack_desc_t &desc, MeshData<Real> *mbase, MeshData<Real> *mc0,
+                      MeshData<Real> *mdudt, MeshData<Real> *mc1) {
+  const int ndim = mc0->GetParentPointer()->ndim;
+  auto flux_div = tl.AddTask(dep, Stokes<pack_desc_t>, cl, tt, desc, ndim, mc0, mdudt);
+  auto avg_data = tl.AddTask(flux_div, WeightedSumData<pack_desc_t>, cl, tt, desc, mc0,
+                             mbase, beta, 1.0 - beta, mc0);
+  return tl.AddTask(avg_data, WeightedSumData<pack_desc_t>, cl, tt, desc, mc0, mdudt, 1.0,
+                    beta * dt, mc1);
+}
+
 // See the advection.hpp declaration for a description of how this function gets called.
 TaskCollection AdvectionDriver::MakeTaskCollection(BlockList_t &blocks, const int stage) {
   using namespace parthenon::Update;
@@ -89,18 +102,22 @@ TaskCollection AdvectionDriver::MakeTaskCollection(BlockList_t &blocks, const in
     auto start_send = tl.AddTask(none, parthenon::StartReceiveBoundaryBuffers, mc1);
     auto start_flxcor = tl.AddTask(none, parthenon::StartReceiveFluxCorrections, mc0);
 
-    static auto desc =
-        parthenon::MakePackDescriptor<advection_package::Conserved::scalar>(
-            pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes},
-            {parthenon::PDOpt::WithFluxes});
+    using namespace advection_package::Conserved;
+    static auto desc = parthenon::MakePackDescriptor<phi>(
+        pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes},
+        {parthenon::PDOpt::WithFluxes});
     using pack_desc_t = decltype(desc);
 
-    static auto desc_fine =
-        parthenon::MakePackDescriptor<advection_package::Conserved::scalar_fine>(
-            pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes},
-            {parthenon::PDOpt::WithFluxes});
+    static auto desc_fine = parthenon::MakePackDescriptor<phi_fine>(
+        pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes},
+        {parthenon::PDOpt::WithFluxes});
     using pack_desc_fine_t = decltype(desc_fine);
 
+    static auto desc_vec = parthenon::MakePackDescriptor<C, D>(
+        pmesh->resolved_packages.get(), {parthenon::Metadata::WithFluxes},
+        {parthenon::PDOpt::WithFluxes});
+
+    using TT = parthenon::TopologicalType;
     using TE = parthenon::TopologicalElement;
     std::vector<TE> faces{TE::F1};
     if (pmesh->ndim > 1) faces.push_back(TE::F2);
@@ -115,37 +132,31 @@ TaskCollection AdvectionDriver::MakeTaskCollection(BlockList_t &blocks, const in
                             desc_fine, face, parthenon::CellLevel::fine, mc0.get());
     }
 
-    auto set_flx = parthenon::AddFluxCorrectionTasks(start_flxcor | flx | flx_fine, tl,
-                                                     mc0, pmesh->multilevel);
-
-    auto flux_div = tl.AddTask(set_flx, Stokes<pack_desc_t>, parthenon::CellLevel::same,
-                               parthenon::TopologicalType::Cell, desc, pmesh->ndim,
-                               mc0.get(), mdudt.get());
-
-    auto flux_div_fine = tl.AddTask(
-        set_flx, Stokes<pack_desc_fine_t>, parthenon::CellLevel::fine,
-        parthenon::TopologicalType::Cell, desc_fine, pmesh->ndim, mc0.get(), mdudt.get());
+    auto vf_dep = none;
+    for (auto edge : std::vector<TE>{TE::E1, TE::E2, TE::E3}) {
+      vf_dep = tl.AddTask(vf_dep, advection_package::CalculateVectorFluxes<C, D>, edge,
+                          parthenon::CellLevel::same, 1.0, mc0.get());
+      vf_dep = tl.AddTask(vf_dep, advection_package::CalculateVectorFluxes<D, C>, edge,
+                          parthenon::CellLevel::same, -1.0, mc0.get());
+    }
 
-    auto avg_data =
-        tl.AddTask(flux_div, WeightedSumData<pack_desc_t>, parthenon::CellLevel::same,
-                   parthenon::TopologicalElement::CC, desc, mc0.get(), mbase.get(), beta,
-                   1.0 - beta, mc0.get());
-    auto avg_data_fine =
-        tl.AddTask(flux_div_fine, WeightedSumData<pack_desc_fine_t>,
-                   parthenon::CellLevel::fine, parthenon::TopologicalElement::CC,
-                   desc_fine, mc0.get(), mbase.get(), beta, 1.0 - beta, mc0.get());
+    auto set_flx = parthenon::AddFluxCorrectionTasks(
+        start_flxcor | flx | flx_fine | vf_dep, tl, mc0, pmesh->multilevel);
 
     auto update =
-        tl.AddTask(avg_data, WeightedSumData<pack_desc_t>, parthenon::CellLevel::same,
-                   parthenon::TopologicalElement::CC, desc, mc0.get(), mdudt.get(), 1.0,
-                   beta * dt, mc1.get());
+        AddUpdateTasks(set_flx, tl, parthenon::CellLevel::same, TT::Cell, beta, dt, desc,
+                       mbase.get(), mc0.get(), mdudt.get(), mc1.get());
+
     auto update_fine =
-        tl.AddTask(avg_data_fine, WeightedSumData<pack_desc_fine_t>,
-                   parthenon::CellLevel::fine, parthenon::TopologicalElement::CC,
-                   desc_fine, mc0.get(), mdudt.get(), 1.0, beta * dt, mc1.get());
+        AddUpdateTasks(set_flx, tl, parthenon::CellLevel::fine, TT::Cell, beta, dt,
+                       desc_fine, mbase.get(), mc0.get(), mdudt.get(), mc1.get());
+
+    auto update_vec =
+        AddUpdateTasks(set_flx, tl, parthenon::CellLevel::same, TT::Face, beta, dt,
+                       desc_vec, mbase.get(), mc0.get(), mdudt.get(), mc1.get());
 
     auto boundaries = parthenon::AddBoundaryExchangeTasks(
-        update | update_fine | start_send, tl, mc1, pmesh->multilevel);
+        update | update_vec | update_fine | start_send, tl, mc1, pmesh->multilevel);
 
     auto fill_derived =
         tl.AddTask(boundaries, parthenon::Update::FillDerived<MeshData<Real>>, mc1.get());

example/fine_advection/advection_package.cpp

@@ -66,23 +66,44 @@ std::shared_ptr<StateDescriptor> Initialize(ParameterInput *pin) {
   }
   pkg->AddParam<>("profile", profile_str);
 
-  pkg->AddField<Conserved::scalar_fine>(
+  pkg->AddField<Conserved::phi_fine>(
       Metadata({Metadata::Cell, Metadata::Fine, Metadata::Independent,
                 Metadata::WithFluxes, Metadata::FillGhost}));
 
-  pkg->AddField<Conserved::scalar>(Metadata({Metadata::Cell, Metadata::Independent,
-                                             Metadata::WithFluxes, Metadata::FillGhost}));
-  pkg->AddField<Conserved::scalar_fine_restricted>(
+  pkg->AddField<Conserved::phi>(Metadata({Metadata::Cell, Metadata::Independent,
+                                          Metadata::WithFluxes, Metadata::FillGhost}));
+  pkg->AddField<Conserved::phi_fine_restricted>(
       Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}));
+
+  Metadata m(
+      {Metadata::Face, Metadata::Independent, Metadata::WithFluxes, Metadata::FillGhost});
+  m.RegisterRefinementOps<parthenon::refinement_ops::ProlongateSharedMinMod,
+                          parthenon::refinement_ops::RestrictAverage,
+                          parthenon::refinement_ops::ProlongateInternalTothAndRoe>();
+  pkg->AddField<Conserved::C>(m);
+  pkg->AddField<Conserved::D>(m);
+  pkg->AddField<Conserved::recon>(Metadata(
+      {Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector<int>{4}));
+  pkg->AddField<Conserved::recon_f>(Metadata(
+      {Metadata::Face, Metadata::Derived, Metadata::OneCopy}, std::vector<int>{2}));
+  pkg->AddField<Conserved::C_cc>(Metadata(
+      {Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector<int>{3}));
+  pkg->AddField<Conserved::D_cc>(Metadata(
+      {Metadata::Cell, Metadata::Derived, Metadata::OneCopy}, std::vector<int>{3}));
+  pkg->AddField<Conserved::divC>(
+      Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}));
+  pkg->AddField<Conserved::divD>(
+      Metadata({Metadata::Cell, Metadata::Derived, Metadata::OneCopy}));
+
   pkg->CheckRefinementBlock = CheckRefinement;
   pkg->EstimateTimestepMesh = EstimateTimestep;
-  pkg->FillDerivedMesh = RestrictScalarFine;
+  pkg->FillDerivedMesh = FillDerived;
   return pkg;
 }
 
 AmrTag CheckRefinement(MeshBlockData<Real> *rc) {
   // refine on advected, for example.  could also be a derived quantity
-  static auto desc = parthenon::MakePackDescriptor<Conserved::scalar>(rc);
+  static auto desc = parthenon::MakePackDescriptor<Conserved::phi>(rc);
   auto pack = desc.GetPack(rc);
 
   auto pmb = rc->GetBlockPointer();
@@ -121,7 +142,7 @@ Real EstimateTimestep(MeshData<Real> *md) {
   const auto &vy = pkg->Param<Real>("vy");
   const auto &vz = pkg->Param<Real>("vz");
 
-  static auto desc = parthenon::MakePackDescriptor<Conserved::scalar>(md);
+  static auto desc = parthenon::MakePackDescriptor<Conserved::phi>(md);
   auto pack = desc.GetPack(md);
 
   IndexRange ib = md->GetBoundsI(IndexDomain::interior);
@@ -147,9 +168,12 @@ Real EstimateTimestep(MeshData<Real> *md) {
   return cfl * min_dt / 2.0;
 }
 
-TaskStatus RestrictScalarFine(MeshData<Real> *md) {
-  static auto desc = parthenon::MakePackDescriptor<Conserved::scalar_fine,
-                                                   Conserved::scalar_fine_restricted>(md);
+TaskStatus FillDerived(MeshData<Real> *md) {
+  static auto desc =
+      parthenon::MakePackDescriptor<Conserved::phi_fine, Conserved::phi_fine_restricted,
+                                    Conserved::C, Conserved::C_cc, Conserved::D,
+                                    Conserved::D_cc, Conserved::divC, Conserved::divD>(
+          md);
   auto pack = desc.GetPack(md);
 
   IndexRange ib = md->GetBoundsI(IndexDomain::interior);
@@ -160,21 +184,66 @@ TaskStatus RestrictScalarFine(MeshData<Real> *md) {
   parthenon::par_for(
       PARTHENON_AUTO_LABEL, 0, pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
       KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
-        const int kf = ndim > 2 ? (k - nghost) * 2 + nghost : k;
-        const int jf = ndim > 1 ? (j - nghost) * 2 + nghost : j;
-        const int fi = ndim > 0 ? (i - nghost) * 2 + nghost : i;
-        pack(b, Conserved::scalar_fine_restricted(), k, j, i) = 0.0;
+        const int kf = (ndim > 2) ? (k - nghost) * 2 + nghost : k;
+        const int jf = (ndim > 1) ? (j - nghost) * 2 + nghost : j;
+        const int fi = (ndim > 0) ? (i - nghost) * 2 + nghost : i;
+        pack(b, Conserved::phi_fine_restricted(), k, j, i) = 0.0;
         Real ntot = 0.0;
         for (int ioff = 0; ioff <= (ndim > 0); ++ioff)
           for (int joff = 0; joff <= (ndim > 1); ++joff)
             for (int koff = 0; koff <= (ndim > 2); ++koff) {
               ntot += 1.0;
-              pack(b, Conserved::scalar_fine_restricted(), k, j, i) +=
-                  pack(b, Conserved::scalar_fine(), kf + koff, jf + joff, fi + ioff);
+              pack(b, Conserved::phi_fine_restricted(), k, j, i) +=
+                  pack(b, Conserved::phi_fine(), kf + koff, jf + joff, fi + ioff);
             }
-        pack(b, Conserved::scalar_fine_restricted(), k, j, i) /= ntot;
+        pack(b, Conserved::phi_fine_restricted(), k, j, i) /= ntot;
       });
 
+  using TE = parthenon::TopologicalElement;
+  parthenon::par_for(
+      PARTHENON_AUTO_LABEL, 0, pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
+      KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
+        pack(b, Conserved::C_cc(0), k, j, i) =
+            0.5 * (pack(b, TE::F1, Conserved::C(), k, j, i) +
+                   pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0)));
+        pack(b, Conserved::C_cc(1), k, j, i) =
+            0.5 * (pack(b, TE::F2, Conserved::C(), k, j, i) +
+                   pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i));
+        pack(b, Conserved::C_cc(2), k, j, i) =
+            0.5 * (pack(b, TE::F3, Conserved::C(), k, j, i) +
+                   pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i));
+        auto &coords = pack.GetCoordinates(b);
+        pack(b, Conserved::divC(), k, j, i) =
+            (pack(b, TE::F1, Conserved::C(), k, j, i + (ndim > 0)) -
+             pack(b, TE::F1, Conserved::C(), k, j, i)) /
+                coords.Dxc<X1DIR>(k, j, i) +
+            (pack(b, TE::F2, Conserved::C(), k, j + (ndim > 1), i) -
+             pack(b, TE::F2, Conserved::C(), k, j, i)) /
+                coords.Dxc<X2DIR>(k, j, i) +
+            (pack(b, TE::F3, Conserved::C(), k + (ndim > 2), j, i) -
+             pack(b, TE::F3, Conserved::C(), k, j, i)) /
+                coords.Dxc<X3DIR>(k, j, i);
+
+        pack(b, Conserved::D_cc(0), k, j, i) =
+            0.5 * (pack(b, TE::F1, Conserved::D(), k, j, i) +
+                   pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0)));
+        pack(b, Conserved::D_cc(1), k, j, i) =
+            0.5 * (pack(b, TE::F2, Conserved::D(), k, j, i) +
+                   pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i));
+        pack(b, Conserved::D_cc(2), k, j, i) =
+            0.5 * (pack(b, TE::F3, Conserved::D(), k, j, i) +
+                   pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i));
+        pack(b, Conserved::divD(), k, j, i) =
+            (pack(b, TE::F1, Conserved::D(), k, j, i + (ndim > 0)) -
+             pack(b, TE::F1, Conserved::D(), k, j, i)) /
+                coords.Dxc<X1DIR>(k, j, i) +
+            (pack(b, TE::F2, Conserved::D(), k, j + (ndim > 1), i) -
+             pack(b, TE::F2, Conserved::D(), k, j, i)) /
+                coords.Dxc<X2DIR>(k, j, i) +
+            (pack(b, TE::F3, Conserved::D(), k + (ndim > 2), j, i) -
+             pack(b, TE::F3, Conserved::D(), k, j, i)) /
+                coords.Dxc<X3DIR>(k, j, i);
+      });
   return TaskStatus::complete;
 }
 } // namespace advection_package