Add ML options to flame problem

.github/workflows/good_defines.txt

@@ -2,9 +2,11 @@ AMREX_SPACEDIM
 AMREX_USE_CUDA
 AMREX_USE_GPU
 DO_PROBLEM_POST_TIMESTEP
+ML
 NAUX_NET
 AUX_THERMO
 ROTATION
 SDC
+_ML_LIB
 _OPENMP
 __cplusplus

Exec/Make.Maestro

@@ -24,6 +24,9 @@ BL_NO_FORT ?= TRUE
 # Require C++17
 CXXSTD := c++17
 
+# default is to not compile with ML (PyTorch)
+USE_ML ?= FALSE
+
 # Check to make sure that Microphysics actually exists,
 # using an arbitrary file that must exist to check.
 # Throw an error if we don't have Microphysics.
@@ -49,6 +52,10 @@ ifeq ($(USE_SDC), TRUE)
   USERSuffix = .SDC
 endif
 
+ifeq ($(USE_ML), TRUE)
+  USERSuffix = .ML
+endif
+
 include $(AMREX_HOME)/Tools/GNUMake/Make.defs
 
 MAESTROEX_AUTO_SOURCE_DIR := $(TmpBuildDir)/maestroex_sources/$(optionsSuffix).EXE
@@ -143,6 +150,37 @@ Blocs += $(foreach dir, $(EXTERN_CORE), $(dir))
 # external libraries
 #------------------------------------------------------------------------------
 
+#------------------------------------------------------------------------------
+# PyTorch
+#------------------------------------------------------------------------------
+
+# Torch directories -- note the Make.package for these adds these
+# directories into INCLUDE_LOCATIONS and LIBRARY_LOCATIONS for us,
+# so we don't need to do it here
+ifeq ($(USE_ML), TRUE)
+  PYTORCH_ROOT := $(TOP)/external/libtorch
+
+  TORCH_LIBPATH = $(PYTORCH_ROOT)/lib
+
+  ifeq ($(USE_CUDA),TRUE)
+    TORCH_LIBS = -ltorch -ltorch_cpu -lc10 -lc10_cuda -lcuda  # -lcaffe2 -lcaffe2_gpu
+  else
+    TORCH_LIBS = -ltorch -ltorch_cpu -lc10  # -lcaffe2
+  endif
+
+  INCLUDE_LOCATIONS += $(PYTORCH_ROOT)/include \
+                       $(PYTORCH_ROOT)/include/torch/csrc/api/include
+  LIBRARY_LOCATIONS += $(TORCH_LIBPATH)
+
+  DEFINES += -DML -D_GLIBCXX_USE_CXX11_ABI=1
+
+  ifeq ($(USE_CUDA), TRUE)
+    LDFLAGS += -Xlinker '--no-as-needed,-rpath $(TORCH_LIBPATH) $(TORCH_LIBS)'
+  else
+    LDFLAGS += -Wl,--no-as-needed,-rpath=$(TORCH_LIBPATH) $(TORCH_LIBS)
+  endif
+endif
+
 #------------------------------------------------------------------------------
 # include all of the necessary directories
 #------------------------------------------------------------------------------

Exec/science/README

@@ -17,6 +17,11 @@ flame/
    the SNe code show that SDC can overcome these problems.
 
 
+flame_ml/
+
+   The same combustion-mode problem as in flame/ but using ML model to
+   replace the reaction kernel of the solver.
+
 sub_chandra/
 
    Surpise! We're modelling convection.  This time in 3D spherical

Exec/science/flame_ml/GNUmakefile

@@ -0,0 +1,29 @@
+DEBUG      = FALSE
+DIM        = 2
+COMP	   = gnu
+USE_MPI    = TRUE
+USE_OMP    = FALSE
+USE_REACT  = TRUE
+USE_ML     = TRUE
+
+TINY_PROFILE = FALSE
+PROFILE      = FALSE # TRUE overrides TINY_PROFILE
+
+
+# define the location of the MAESTROEX home directory
+MAESTROEX_HOME  := ../../..
+
+
+# Set the EOS, conductivity, and network directories
+# We first check if these exist in $(MAESTROEX_HOME)/Microphysics/(EOS/conductivity/networks)
+# If not we use the version in $(MICROPHYSICS_HOME)/Microphysics/(EOS/conductivity/networks)
+EOS_DIR          := helmholtz
+CONDUCTIVITY_DIR := stellar
+NETWORK_DIR      := ignition_simple
+INTEGRATOR_DIR   := VODE
+
+Bpack   := ./Make.package
+Blocs   := .
+PROBIN_PARAMETER_DIRS := .
+
+include $(MAESTROEX_HOME)/Exec/Make.Maestro

Exec/science/flame_ml/MaestroBCFill.cpp

@@ -0,0 +1,259 @@
+
+#include <Maestro.H>
+#include <Maestro_F.H>
+
+#include <MaestroInletBCs.H>
+
+using namespace amrex;
+using namespace InletBCs;
+
+void Maestro::ScalarFill(const Array4<Real>& scal, const Box& bx,
+                         const Box& domainBox, const Real* dx, const BCRec bcs,
+                         const Real* gridlo, const int comp) {
+    // timer for profiling
+    BL_PROFILE_VAR("Maestro::ScalarFill()", ScalarFill);
+
+    fab_filcc(bx, scal, 1, domainBox, dx, gridlo, &bcs);
+
+    FillExtBC(scal, bx, domainBox, dx, bcs, comp, false);
+}
+
+void Maestro::VelFill(const Array4<Real>& vel, const Box& bx,
+                      const Box& domainBox, const Real* dx, const BCRec bcs,
+                      const Real* gridlo, const int comp) {
+    // timer for profiling
+    BL_PROFILE_VAR("Maestro::VelFill()", VelFill);
+
+    fab_filcc(bx, vel, 1, domainBox, dx, gridlo, &bcs);
+
+    FillExtBC(vel, bx, domainBox, dx, bcs, comp, true);
+}
+
+void Maestro::FillExtBC(const Array4<Real>& q, const Box& bx,
+                        const Box& domainBox, const Real* dx, const BCRec bcs,
+                        const int comp, const bool is_vel) {
+    // timer for profiling
+    BL_PROFILE_VAR("Maestro::FillExtBC()", FillExtBC);
+
+    // do nothing if there are no exterior boundaries
+    bool found_ext_boundary = false;
+    for (auto i = 0; i < AMREX_SPACEDIM; ++i) {
+        if (bcs.lo(i) == BCType::ext_dir || bcs.hi(i) == BCType::ext_dir) {
+            found_ext_boundary = true;
+            break;
+        }
+    }
+    if (!found_ext_boundary) return;
+
+    // get parameters for EXT_DIR bcs
+    const Real INLET_RHO_l = INLET_RHO;
+    const Real INLET_RHOH_l = INLET_RHOH;
+    const Real INLET_TEMP_l = INLET_TEMP;
+    RealVector INLET_RHOX_l(NumSpec);
+    for (int i = 0; i < NumSpec; ++i) INLET_RHOX_l[i] = INLET_RHOX[i];
+    const Real INLET_VEL_l = INLET_VEL;
+
+    const auto domlo = domainBox.loVect3d();
+    const auto domhi = domainBox.hiVect3d();
+
+    const auto lo = bx.loVect3d();
+    const auto hi = bx.hiVect3d();
+
+    if (lo[0] < domlo[0]) {
+        auto imax = domlo[0];
+
+        if (bcs.lo(0) == BCType::ext_dir) {
+            ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                if (i < imax) {
+                    q(i, j, k) = 1.0;
+                }
+            });
+            Gpu::synchronize();
+        }
+    }
+
+    if (hi[0] > domhi[0]) {
+        auto imin = domhi[0];
+
+        if (bcs.hi(0) == BCType::ext_dir) {
+            ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                if (i > imin) {
+                    q(i, j, k) = 1.0;
+                }
+            });
+            Gpu::synchronize();
+        }
+    }
+
+    if (lo[1] < domlo[1]) {
+        auto jmax = domlo[1];
+
+        if (bcs.lo(1) == BCType::ext_dir) {
+            if (is_vel) {
+                if (comp == 0) {
+                    ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                        if (j < jmax) {
+                            q(i, j, k) = 0.0;
+                        }
+                    });
+                } else if (comp == 1) {
+                    ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                        if (j < jmax) {
+                            q(i, j, k) = INLET_VEL_l;
+                        }
+                    });
+                }
+                Gpu::synchronize();
+            } else {
+                if (INLET_VEL != 0.0) {
+                    if (comp == Rho) {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (j < jmax) {
+                                            q(i, j, k) = INLET_RHO_l;
+                                        }
+                                    });
+                    } else if (comp == RhoH) {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (j < jmax) {
+                                            q(i, j, k) = INLET_RHOH_l;
+                                        }
+                                    });
+                    } else if (comp >= FirstSpec &&
+                               comp < FirstSpec + NumSpec) {
+                        Real* AMREX_RESTRICT INLET_RHOX_p =
+                            INLET_RHOX.dataPtr();
+                        ParallelFor(
+                            bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                if (j < jmax) {
+                                    q(i, j, k) = INLET_RHOX_p[comp - FirstSpec];
+                                }
+                            });
+                    } else if (comp == Temp) {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (j < jmax) {
+                                            q(i, j, k) = INLET_TEMP_l;
+                                        }
+                                    });
+                    } else {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (j < jmax) {
+                                            q(i, j, k) = 0.0;
+                                        }
+                                    });
+                    }
+                } else {
+                    ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                        if (j < jmax) {
+                            q(i, j, k) = 0.0;
+                        }
+                    });
+                }
+                Gpu::synchronize();
+            }  // end if is_vel
+        }
+    }
+
+    if (hi[1] > domhi[1]) {
+        auto jmin = domhi[1];
+
+        if (bcs.hi(1) == BCType::ext_dir) {
+            ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                if (j > jmin) {
+                    q(i, j, k) = 1.0;
+                }
+            });
+        }
+        Gpu::synchronize();
+    }
+
+#if AMREX_SPACEDIM == 3
+
+    if (lo[2] < domlo[2]) {
+        auto kmax = domlo[2];
+
+        if (bcs.lo(2) == BCType::ext_dir) {
+            if (is_vel) {
+                if (comp == 0 || comp == 1) {
+                    ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                        if (k < kmax) {
+                            q(i, j, k) = 0.0;
+                        }
+                    });
+                } else if (comp == 2) {
+                    ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                        if (k < kmax) {
+                            q(i, j, k) = INLET_VEL_l;
+                        }
+                    });
+                }
+            } else {
+                if (INLET_VEL != 0.0) {
+                    if (comp == Rho) {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (k < kmax) {
+                                            q(i, j, k) = INLET_RHO_l;
+                                        }
+                                    });
+                    } else if (comp == RhoH) {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (k < kmax) {
+                                            q(i, j, k) = INLET_RHOH_l;
+                                        }
+                                    });
+                    } else if (comp >= FirstSpec &&
+                               comp < FirstSpec + NumSpec) {
+                        Real* AMREX_RESTRICT INLET_RHOX_p =
+                            INLET_RHOX.dataPtr();
+                        ParallelFor(
+                            bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                if (k < kmax) {
+                                    q(i, j, k) = INLET_RHOX_p[comp - FirstSpec];
+                                }
+                            });
+                    } else if (comp == Temp) {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (k < kmax) {
+                                            q(i, j, k) = INLET_TEMP_l;
+                                        }
+                                    });
+                    } else {
+                        ParallelFor(bx,
+                                    [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                                        if (k < kmax) {
+                                            q(i, j, k) = 0.0;
+                                        }
+                                    });
+                    }
+                } else {
+                    ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                        if (k < kmax) {
+                            q(i, j, k) = 0.0;
+                        }
+                    });
+                }
+            }  // end if is_vel
+            Gpu::synchronize();
+        }
+    }
+
+    if (hi[2] > domhi[2]) {
+        auto kmin = domhi[2];
+
+        if (bcs.hi(2) == BCType::ext_dir) {
+            ParallelFor(bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) {
+                if (k > kmin) {
+                    q(i, j, k) = 0.0;
+                }
+            });
+        }
+        Gpu::synchronize();
+    }
+#endif
+}