Use autodiff in templated network rate evaluation (#1614)

Note: this fixes an issue with powerlaw, which previously did not have any derivatives computed

Author: yut23

interfaces/rhs_type.H

@@ -95,12 +95,13 @@ struct rate_tab_t
     }
 };
 
-// number_t is currently only used in the screening routines
-template <typename number_t = amrex::Real>
+// number_t is currently used in the screening routines and templated network
+// rate evaluation
+template <typename number_t>
 struct rhs_state_t
 {
     amrex::Real rho;
-    tf_t tf;
+    tf_t<number_t> tf;
     rate_tab_t tab;
 #ifdef SCREENING
     plasma_state_t<number_t> pstate;
@@ -110,12 +111,11 @@ struct rhs_state_t
     amrex::Array1D<amrex::Real, 1, NumSpec> y;
 };
 
+template <typename number_t>
 struct rate_t
 {
-    amrex::Real fr;
-    amrex::Real rr;
-    amrex::Real frdt;
-    amrex::Real rrdt;
+    number_t fr;
+    number_t rr;
 };
 
 } // namespace RHS

interfaces/tfactors.H

@@ -2,76 +2,80 @@
 #define TFACTORS_H
 
 #include <AMReX.H>
+#include <AMReX_REAL.H>
+#include <microphysics_autodiff.H>
 #include <cmath>
 
 using namespace amrex::literals;
 
+template <typename number_t>
 struct tf_t {
-    amrex::Real temp;
-    amrex::Real t9;
-    amrex::Real t92;
-    amrex::Real t93;
-    // amrex::Real t94;
-    amrex::Real t95;
-    // amrex::Real t96;
-    amrex::Real t912;
-    amrex::Real t932;
-    amrex::Real t952;
-    amrex::Real t972;
-    amrex::Real t913;
-    amrex::Real t923;
-    amrex::Real t943;
-    amrex::Real t953;
-    // amrex::Real t973;
-    // amrex::Real t9113;
-    // amrex::Real t914;
-    // amrex::Real t934;
-    // amrex::Real t954;
-    // amrex::Real t974;
-    // amrex::Real t915;
-    // amrex::Real t935;
-    // amrex::Real t945;
-    // amrex::Real t965;
-    // amrex::Real t917;
-    // amrex::Real t927;
-    // amrex::Real t947;
-    // amrex::Real t918;
-    // amrex::Real t938;
-    // amrex::Real t958;
-    amrex::Real t9i;
-    amrex::Real t9i2;
-    // amrex::Real t9i3;
-    amrex::Real t9i12;
-    amrex::Real t9i32;
-    // amrex::Real t9i52;
-    // amrex::Real t9i72;
-    amrex::Real t9i13;
-    amrex::Real t9i23;
-    amrex::Real t9i43;
-    amrex::Real t9i53;
-    // amrex::Real t9i14;
-    // amrex::Real t9i34;
-    // amrex::Real t9i54;
-    // amrex::Real t9i15;
-    // amrex::Real t9i35;
-    // amrex::Real t9i45;
-    // amrex::Real t9i65;
-    // amrex::Real t9i17;
-    // amrex::Real t9i27;
-    // amrex::Real t9i47;
-    // amrex::Real t9i18;
-    // amrex::Real t9i38;
-    // amrex::Real t9i58;
-    // amrex::Real t916;
-    // amrex::Real t976;
-    // amrex::Real t9i76;
-    amrex::Real lnt9;
+    number_t temp;
+    number_t t9;
+    number_t t92;
+    number_t t93;
+    // number_t t94;
+    number_t t95;
+    // number_t t96;
+    number_t t912;
+    number_t t932;
+    number_t t952;
+    number_t t972;
+    number_t t913;
+    number_t t923;
+    number_t t943;
+    number_t t953;
+    // number_t t973;
+    // number_t t9113;
+    // number_t t914;
+    // number_t t934;
+    // number_t t954;
+    // number_t t974;
+    // number_t t915;
+    // number_t t935;
+    // number_t t945;
+    // number_t t965;
+    // number_t t917;
+    // number_t t927;
+    // number_t t947;
+    // number_t t918;
+    // number_t t938;
+    // number_t t958;
+    number_t t9i;
+    number_t t9i2;
+    // number_t t9i3;
+    number_t t9i12;
+    number_t t9i32;
+    // number_t t9i52;
+    // number_t t9i72;
+    number_t t9i13;
+    number_t t9i23;
+    number_t t9i43;
+    number_t t9i53;
+    // number_t t9i14;
+    // number_t t9i34;
+    // number_t t9i54;
+    // number_t t9i15;
+    // number_t t9i35;
+    // number_t t9i45;
+    // number_t t9i65;
+    // number_t t9i17;
+    // number_t t9i27;
+    // number_t t9i47;
+    // number_t t9i18;
+    // number_t t9i38;
+    // number_t t9i58;
+    // number_t t916;
+    // number_t t976;
+    // number_t t9i76;
+    number_t lnt9;
 };
 
+template <typename number_t>
 AMREX_GPU_HOST_DEVICE inline
-tf_t get_tfactors(amrex::Real temp)
+tf_t<number_t> get_tfactors(number_t temp)
 {
-    tf_t tf;
+    tf_t<number_t> tf;
 
     tf.temp = temp;
 
@@ -83,38 +87,38 @@ tf_t get_tfactors(amrex::Real temp)
     tf.t95   = tf.t92*tf.t93;
     // tf.t96   = tf.t9*tf.t95;
 
-    tf.t912  = std::sqrt(tf.t9);
+    tf.t912  = admath::sqrt(tf.t9);
     tf.t932  = tf.t9*tf.t912;
     tf.t952  = tf.t9*tf.t932;
     // tf.t972  = tf.t9*tf.t952;
     tf.t972  = tf.t92*tf.t932;
 
-    tf.t913  = std::cbrt(tf.t9);
+    tf.t913  = admath::cbrt(tf.t9);
     tf.t923  = tf.t913*tf.t913;
     tf.t943  = tf.t9*tf.t913;
     tf.t953  = tf.t9*tf.t923;
     // tf.t973  = tf.t953*tf.t923;
     // tf.t9113 = tf.t973*tf.t943;
 
-    // tf.t914  = std::pow(tf.t9, 0.25e0_rt);
+    // tf.t914  = admath::pow(tf.t9, 0.25e0_rt);
     // tf.t934  = tf.t914*tf.t914*tf.t914;
     // tf.t954  = tf.t9*tf.t914;
     // tf.t974  = tf.t9*tf.t934;
 
-    // tf.t915  = std::pow(tf.t9, 0.2_rt);
+    // tf.t915  = admath::pow(tf.t9, 0.2_rt);
     // tf.t935  = tf.t915*tf.t915*tf.t915;
     // tf.t945  = tf.t915 * tf.t935;
     // tf.t965  = tf.t9 * tf.t915;
 
-    // tf.t916  = std::pow(tf.t9, 1.0_rt/6.0_rt);
+    // tf.t916  = admath::pow(tf.t9, 1.0_rt/6.0_rt);
     // tf.t976  = tf.t9 * tf.t916;
     // tf.t9i76 = 1.0e0_rt/tf.t976;
 
-    // tf.t917  = std::pow(tf.t9, 1.0_rt/7.0_rt);
+    // tf.t917  = admath::pow(tf.t9, 1.0_rt/7.0_rt);
     // tf.t927  = tf.t917*tf.t917;
     // tf.t947  = tf.t927*tf.t927;
 
-    // tf.t918  = std::sqrt(tf.t914);
+    // tf.t918  = admath::sqrt(tf.t914);
     // tf.t938  = tf.t918*tf.t918*tf.t918;
     // tf.t958  = tf.t938*tf.t918*tf.t918;
 
@@ -149,7 +153,7 @@ tf_t get_tfactors(amrex::Real temp)
     // tf.t9i38 = tf.t9i18*tf.t9i18*tf.t9i18;
     // tf.t9i58 = tf.t9i38*tf.t9i18*tf.t9i18;
 
-    tf.lnt9 = std::log(tf.t9);
+    tf.lnt9 = admath::log(tf.t9);
 
     return tf;
 }