Add a version of test_screening that works on any templated network (#…

…1524)

I've been using this to test out autodiff on the screening implementations.

unit_test/test_screening_templated/GNUmakefile

@@ -0,0 +1,43 @@
+PRECISION  = DOUBLE
+PROFILE    = FALSE
+
+DEBUG      = FALSE
+
+DIM        = 3
+
+COMP	   = gnu
+
+USE_MPI    = FALSE
+USE_OMP    = FALSE
+
+USE_REACT = TRUE
+USE_CONDUCTIVITY = FALSE
+
+EBASE = main
+
+BL_NO_FORT = TRUE
+
+# define the location of the Microphysics top directory
+MICROPHYSICS_HOME  := ../..
+
+# This sets the EOS directory
+EOS_DIR     := helmholtz
+
+# This sets the network directory
+NETWORK_DIR := aprox21
+
+# This isn't actually used but we need VODE to compile with CUDA
+INTEGRATOR_DIR := VODE
+
+CONDUCTIVITY_DIR := stellar
+
+SCREEN_METHOD := screen5
+
+EXTERN_SEARCH += .
+
+Bpack   := ./Make.package
+Blocs   := .
+
+include $(MICROPHYSICS_HOME)/unit_test/Make.unit_test
+
+

unit_test/test_screening_templated/Make.package

@@ -0,0 +1,7 @@
+CEXE_sources += main.cpp
+
+CEXE_headers += test_screen.H
+
+CEXE_sources += variables.cpp
+CEXE_sources += screening_util.cpp
+CEXE_headers += variables.H

unit_test/test_screening_templated/README

@@ -0,0 +1,5 @@
+# test_screening_templated
+
+This is a unit test for the screening routines, using the templated network
+machinery to exercise all the rates in a network.
+

unit_test/test_screening_templated/_parameters

@@ -0,0 +1,11 @@
+@namespace: unit_test
+
+dens_min      real       1.d6
+dens_max      real       1.d9
+temp_min      real       1.d6
+temp_max      real       1.d12
+
+metalicity_max  real     0.1d0
+
+small_temp    real        1.d4
+small_dens    real        1.d-4

unit_test/test_screening_templated/inputs

@@ -0,0 +1,9 @@
+n_cell = 16
+max_grid_size = 32
+
+unit_test.dens_min = 10.e0
+unit_test.dens_max = 5.e9
+unit_test.temp_min = 1.e6
+unit_test.temp_max = 1.e10
+
+unit_test.metalicity_max = 0.5e0

unit_test/test_screening_templated/main.cpp

@@ -0,0 +1,164 @@
+#include <AMReX_PlotFileUtil.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_Print.H>
+
+#include <AMReX_Geometry.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_BCRec.H>
+
+
+using namespace amrex;
+
+#include <test_screen.H>
+#include <AMReX_buildInfo.H>
+
+#include <network.H>
+#include <eos.H>
+#include <screen.H>
+
+#include <variables.H>
+
+#include <cmath>
+#include <unit_test.H>
+
+using namespace unit_test_rp;
+
+int main (int argc, char* argv[])
+{
+    amrex::Initialize(argc, argv);
+
+    main_main();
+
+    amrex::Finalize();
+    return 0;
+}
+
+void main_main ()
+{
+
+    // AMREX_SPACEDIM: number of dimensions
+    int n_cell, max_grid_size;
+
+    // inputs parameters
+    {
+        // ParmParse is way of reading inputs from the inputs file
+        ParmParse pp;
+
+        // We need to get n_cell from the inputs file - this is the
+        // number of cells on each side of a square (or cubic) domain.
+        pp.get("n_cell", n_cell);
+
+        // The domain is broken into boxes of size max_grid_size
+        max_grid_size = 32;
+        pp.query("max_grid_size", max_grid_size);
+
+    }
+
+
+    Vector<int> is_periodic(AMREX_SPACEDIM,0);
+    for (int idim=0; idim < AMREX_SPACEDIM; ++idim) {
+      is_periodic[idim] = 1;
+    }
+
+    // make BoxArray and Geometry
+    BoxArray ba;
+    Geometry geom;
+    {
+        IntVect dom_lo(AMREX_D_DECL(       0,        0,        0));
+        IntVect dom_hi(AMREX_D_DECL(n_cell-1, n_cell-1, n_cell-1));
+        Box domain(dom_lo, dom_hi);
+
+        // Initialize the boxarray "ba" from the single box "bx"
+        ba.define(domain);
+
+        // Break up boxarray "ba" into chunks no larger than
+        // "max_grid_size" along a direction
+        ba.maxSize(max_grid_size);
+
+        // This defines the physical box, [0, 1] in each direction.
+        RealBox real_box({AMREX_D_DECL(0.0, 0.0, 0.0)},
+                         {AMREX_D_DECL(1.0, 1.0, 1.0)});
+
+        // This defines a Geometry object
+        geom.define(domain, &real_box,
+                    CoordSys::cartesian, is_periodic.data());
+    }
+
+    // Nghost = number of ghost cells for each array
+    int Nghost = 0;
+
+    // do the runtime parameter initializations and microphysics inits
+    if (ParallelDescriptor::IOProcessor()) {
+      std::cout << "reading extern runtime parameters ..." << std::endl;
+    }
+
+    ParmParse ppa("amr");
+
+    init_unit_test();
+
+    eos_init(small_temp, small_dens);
+
+    network_init();
+
+    screening_init();
+
+    amrex::Vector<std::string> names;
+    plot_t vars = init_variables(names);
+
+    // time = starting time in the simulation
+    Real time = 0.0_rt;
+
+    // How Boxes are distributed among MPI processes
+    DistributionMapping dm(ba);
+
+    // we allocate our main multifabs
+    MultiFab state(ba, dm, vars.n_plot_comps, Nghost);
+
+    // Initialize the state to zero; we will fill
+    // it in below in do_eos.
+    state.setVal(0.0_rt);
+
+    // What time is it now?  We'll use this to compute total run time.
+    Real strt_time = ParallelDescriptor::second();
+
+
+    Real dlogrho = 0.0e0_rt;
+    Real dlogT   = 0.0e0_rt;
+    Real dmetal  = 0.0e0_rt;
+
+    if (n_cell > 1) {
+        dlogrho = (std::log10(dens_max) - std::log10(dens_min)) / static_cast<Real>(n_cell - 1);
+        dlogT   = (std::log10(temp_max) - std::log10(temp_min))/ static_cast<Real>(n_cell - 1);
+        dmetal  = (metalicity_max  - 0.0_rt)/ static_cast<Real>(n_cell - 1);
+    }
+
+    // Initialize the state and compute the different thermodynamics
+    // by inverting the EOS
+    for ( MFIter mfi(state); mfi.isValid(); ++mfi )
+    {
+      const Box& bx = mfi.validbox();
+
+      Array4<Real> const sp = state.array(mfi);
+
+      screen_test_C(bx, dlogrho, dlogT, dmetal, vars, sp);
+
+    }
+
+    // Call the timer again and compute the maximum difference between
+    // the start time and stop time over all processors
+    Real stop_time = ParallelDescriptor::second() - strt_time;
+    const int IOProc = ParallelDescriptor::IOProcessorNumber();
+    ParallelDescriptor::ReduceRealMax(stop_time, IOProc);
+
+
+    std::string name = "test_screening." + screen_name;
+
+    // Write a plotfile
+    WriteSingleLevelPlotfile(name, state, names, geom, time, 0);
+
+    write_job_info(name);
+
+    // Tell the I/O Processor to write out the "run time"
+    amrex::Print() << "Run time = " << stop_time << std::endl;
+
+}

unit_test/test_screening_templated/screening_util.cpp

@@ -0,0 +1,151 @@
+#include <AMReX_PlotFileUtil.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_Print.H>
+
+#include <AMReX_Geometry.H>
+#include <AMReX_MultiFab.H>
+#include <AMReX_BCRec.H>
+#include <AMReX_Loop.H>
+
+#include <variables.H>
+#include <network.H>
+#include <eos.H>
+#include <rhs_type.H>
+
+#include <screen.H>
+
+#include <cmath>
+
+using namespace amrex;
+using namespace unit_test_rp;
+
+void screen_test_C(const Box& bx,
+                   const Real dlogrho, const Real dlogT, const Real dmetal,
+                   const plot_t& vars,
+                   Array4<Real> const sp) {
+
+  const int ih1 = network_spec_index("hydrogen-1");
+  if (ih1 < 0) amrex::Error("Error: ih1 not found");
+
+  const int ihe4 = network_spec_index("helium-4");
+  if (ihe4 < 0) amrex::Error("Error: ihe4 not found");
+
+  amrex::ParallelFor(bx,
+  [=] AMREX_GPU_DEVICE (int i, int j, int k)
+  {
+
+    // set the composition -- approximately solar
+    Real metalicity = 0.0 + static_cast<Real> (k) * dmetal;
+
+    Real xn[NumSpec];
+
+    // for now... the screening using 1-based indexing
+    Array1D<Real, 1, NumSpec> ymass;
+
+    for (auto& x : xn) {
+      x = metalicity / static_cast<Real>(NumSpec - 2);
+    }
+    xn[ih1] = 0.75_rt - 0.5_rt * metalicity;
+    xn[ihe4] = 0.25_rt - 0.5_rt * metalicity;
+
+    for (int n = 0; n < NumSpec; n++) {
+      ymass(n+1) = xn[n] / aion[n];
+    }
+
+    Real temp_zone = std::pow(10.0, std::log10(temp_min) + static_cast<Real>(j)*dlogT);
+
+    Real dens_zone = std::pow(10.0, std::log10(dens_min) + static_cast<Real>(i)*dlogrho);
+
+    // store default state
+    sp(i, j, k, vars.irho) = dens_zone;
+    sp(i, j, k, vars.itemp) = temp_zone;
+    for (int n = 0; n < NumSpec; n++) {
+      sp(i, j, k, vars.ispec+n) = xn[n];
+    }
+
+    plasma_state_t pstate;
+    fill_plasma_state(pstate, temp_zone, dens_zone, ymass);
+
+    Real sc1a;
+    Real sc1adt;
+
+    constexpr_for<1, Rates::NumRates+1>([&] (auto n) {
+      constexpr int rate = n;
+      constexpr RHS::rhs_t data = RHS::rhs_data(rate);
+
+      if (i == 1 && j == 1 && k == 1) {
+        std::cout << "working on rate " << rate << "\n";
+      }
+
+      if constexpr (data.screen_forward_reaction == 0 && data.screen_reverse_reaction == 0) {
+        return;
+      }
+      if (vars.iscn(rate).value == -1) {
+        return;
+      }
+
+      if constexpr (data.exponent_A == 1 && data.exponent_B == 1 && data.exponent_C == 0) {
+        // Forward reaction is A + B, screen using these two species
+
+        constexpr amrex::Real Z1 = NetworkProperties::zion(data.species_A);
+        constexpr amrex::Real A1 = NetworkProperties::aion(data.species_A);
+
+        constexpr amrex::Real Z2 = NetworkProperties::zion(data.species_B);
+        constexpr amrex::Real A2 = NetworkProperties::aion(data.species_B);
+
+        constexpr auto scn_fac = scrn::calculate_screen_factor(Z1, A1, Z2, A2);
+
+        // Require scn_fac to be evaluated at compile time
+        static_assert(scn_fac.z1 == Z1);
+
+        actual_screen(pstate, scn_fac, sc1a, sc1adt);
+        sp(i, j, k, vars.iscn(rate).value) = sc1a;
+        sp(i, j, k, vars.iscn(rate).dt) = sc1adt;
+      }
+
+      if constexpr (data.exponent_A == 2 && data.exponent_B == 0 && data.exponent_C == 0) {
+        // Forward reaction is A + A, screen using just this species
+
+        constexpr amrex::Real Z1 = NetworkProperties::zion(data.species_A);
+        constexpr amrex::Real A1 = NetworkProperties::aion(data.species_A);
+
+        constexpr auto scn_fac = scrn::calculate_screen_factor(Z1, A1, Z1, A1);
+
+        static_assert(scn_fac.z1 == Z1);
+
+        actual_screen(pstate, scn_fac, sc1a, sc1adt);
+        sp(i, j, k, vars.iscn(rate).value) = sc1a;
+        sp(i, j, k, vars.iscn(rate).dt) = sc1adt;
+      }
+
+      if constexpr (data.exponent_A == 3 && data.exponent_B == 0 && data.exponent_C == 0) {
+        // Forward reaction is triple alpha or an equivalent, screen using A + A
+        // and then A + X where X has twice the number of protons and neutrons.
+
+        constexpr amrex::Real Z1 = NetworkProperties::zion(data.species_A);
+        constexpr amrex::Real A1 = NetworkProperties::aion(data.species_A);
+
+        constexpr auto scn_fac1 = scrn::calculate_screen_factor(Z1, A1, Z1, A1);
+
+        static_assert(scn_fac1.z1 == Z1);
+
+        actual_screen(pstate, scn_fac1, sc1a, sc1adt);
+        sp(i, j, k, vars.iscn(rate).value) = sc1a;
+        sp(i, j, k, vars.iscn(rate).dt) = sc1adt;
+
+        constexpr amrex::Real Z2 = 2.0_rt * Z1;
+        constexpr amrex::Real A2 = 2.0_rt * A1;
+
+        constexpr auto scn_fac2 = scrn::calculate_screen_factor(Z1, A1, Z2, A2);
+
+        static_assert(scn_fac2.z1 == Z1);
+
+        actual_screen(pstate, scn_fac2, sc1a, sc1adt);
+        sp(i, j, k, vars.iscn(rate).aux_value) = sc1a;
+        sp(i, j, k, vars.iscn(rate).aux_dt) = sc1adt;
+      }
+    });
+
+  });
+
+}