Update from naunet

CMakeLists.txt

@@ -1,15 +1,15 @@
 # To use CUDA by enabling language
 cmake_minimum_required(VERSION 3.18)
-project(NAUNET VERSION 0.2.0 LANGUAGES CUDA CXX)
+project(NAUNET VERSION 0.4.2 LANGUAGES CUDA CXX)
 
 include(CMakePrintHelpers)
 
-set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_CXX_STANDARD 14)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)
 
 if(NOT CMAKE_BUILD_TYPE)
-    set(CMAKE_BUILD_TYPE Debug)
+    set(CMAKE_BUILD_TYPE Release)
 endif()
 
 if(CMAKE_BUILD_TYPE STREQUAL Debug)
@@ -22,10 +22,13 @@ message(STATUS "CMAKE_CXX_COMPILER_VERSION = ${CMAKE_CXX_COMPILER_VERSION}")
 message(STATUS "CMAKE_COMPILER_IS_GNUCXX   = ${CMAKE_COMPILER_IS_GNUCXX}"  )
 message(STATUS "CMAKE_BUILD_TYPE           = ${CMAKE_BUILD_TYPE}"          )
 
-set(SUNDIALS_MIN_VERSION "5.6.1")
+set(SUNDIALS_MIN_VERSION "6.0.0")
 set(SUNDIALS_DIR "" CACHE PATH "Path to SUNDIALSConfig.cmake")
 find_package(SUNDIALS ${SUNDIALS_MIN_VERSION} REQUIRED CONFIG)
 
+message(STATUS "SUNDIALS_FOUND             = ${SUNDIALS_FOUND}")
+
+
 list(APPEND SUNDIALSLIB
     SUNDIALS::cvode
     SUNDIALS::nveccuda
@@ -79,14 +82,17 @@ message(STATUS "NVIDIA_GPU_COUNT           = " ${NVIDIA_GPU_COUNT})
 set(MAKE_STATIC ON CACHE BOOL "")
 set(MAKE_SHARED ON CACHE BOOL "")
 set(MAKE_PYTHON OFF CACHE BOOL "")
+set(MAKE_TEST ON CACHE BOOL "")
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
 if(MAKE_PYTHON)
-    set(PYMODNAME "chemistry" CACHE STRING "")
-    find_package(PythonInterp REQUIRED)
+    set(PYMODNAME "pynaunet" CACHE STRING "")
+    find_package(Python REQUIRED COMPONENTS Interpreter Development)
 endif()
 
-enable_testing()
-
 add_subdirectory(src)
-add_subdirectory(test)
+
+if(MAKE_TEST)
+    enable_testing()
+    add_subdirectory(tests)
+endif()

demo.ipynb

@@ -0,0 +1,85 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import tomlkit\n",
+    "\n",
+    "from pathlib import Path\n",
+    "\n",
+    "plt.rcParams['figure.figsize'] = [8, 6]\n",
+    "plt.rcParams['figure.dpi'] = 100 # 200 e.g. is really fine, but slower"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"naunet_config.toml\") as config:\n",
+    "    content = tomlkit.loads(config.read())\n",
+    "\n",
+    "summary = content[\"summary\"]\n",
+    "species = summary[\"list_of_species\"]\n",
+    "nspec = summary[\"num_of_species\"]\n",
+    "\n",
+    "bindata = np.fromfile(f\"build/tests/evolution_singlegrid.bin\", dtype=\"float64\").reshape(-1, nspec+1)\n",
+    "time = bindata[:, 0]\n",
+    "\n",
+    "spec = \"H2\"\n",
+    "sidx = species.index(spec) + 1\n",
+    "\n",
+    "fig, ax = plt.subplots(facecolor=\"w\")\n",
+    "\n",
+    "ax.plot(time, bindata[:,sidx]);\n",
+    "\n",
+    "ax.set_xscale('log');\n",
+    "ax.set_yscale('log');\n",
+    "ax.set_xlabel('Time (yr)');\n",
+    "ax.set_ylabel(f'[{spec}]');"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "elapsed = np.loadtxt(\"build/tests/time_singlegrid.txt\")\n",
+    "\n",
+    "fig, ax = plt.subplots(facecolor=\"w\")\n",
+    "\n",
+    "ax.plot(np.arange(elapsed.size), elapsed);\n",
+    "\n",
+    "ax.set_yscale(\"log\");\n",
+    "ax.set_xlabel(\"Step\");\n",
+    "ax.set_ylabel(\"Elapsed Time (s)\");"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.6 ('.venv': poetry)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.8.6"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "314ff8c970d95e799925aa8f25686df887ef025d7910d0eca0a071c657e55e81"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

include/naunet.h

@@ -23,14 +23,31 @@ class Naunet {
    public:
     Naunet();
     ~Naunet();
-    int Init(int nsystem = MAX_NSYSTEMS, double atol = 1e-20, double rtol = 1e-5, int mxsteps=500);
-    int DebugInfo();
     int Finalize();
-    /* */
-    int Reset(int nsystem = MAX_NSYSTEMS, double atol = 1e-20, double rtol = 1e-5, int mxsteps=500);
-    /* */
+    int Init(int nsystem = MAX_NSYSTEMS, double atol = 1e-20,
+             double rtol = 1e-5, int mxsteps = 500);
+    int PrintDebugInfo();
+#ifdef IDX_ELEM_H
+    // Renormalize the species abundances accroding to `ab_ref_`(private).
+    // This conserves the element fractions but does not change the overall
+    // density/
+    int Renorm(realtype *ab);
+#endif
+    int Reset(int nsystem = MAX_NSYSTEMS, double atol = 1e-20,
+              double rtol = 1e-5, int mxsteps = 500);
+#ifdef IDX_ELEM_H
+    // Set the reference abundance `ab_ref_`. `opt == 0` assumes that the input
+    // is element abundances. `opt == 1` assumes the input is species
+    // abundances.
+    int SetReferenceAbund(realtype *ref, int opt = 0);
+#endif
     int Solve(realtype *ab, realtype dt, NaunetData *data);
 #ifdef PYMODULE
+#ifdef IDX_ELEM_H
+    py::array_t<realtype> PyWrapRenorm(py::array_t<realtype> arr);
+    py::array_t<realtype> PyWrapSetReferenceAbund(py::array_t<realtype> arr,
+                                                  int opt);
+#endif
     py::array_t<realtype> PyWrapSolve(py::array_t<realtype> arr, realtype dt,
                                       NaunetData *data);
 #endif
@@ -40,6 +57,8 @@ class Naunet {
     int mxsteps_;
     realtype atol_;
     realtype rtol_;
+    FILE *errfp_;
+    realtype ab_ref_[NELEMENTS];
 
     /* */
 
@@ -49,6 +68,7 @@ class Naunet {
     SUNMatrix cv_a_[NSTREAMS];
     void *cv_mem_[NSTREAMS];
     SUNLinearSolver cv_ls_[NSTREAMS];
+    SUNContext cv_sunctx_[NSTREAMS];
 
     cusparseHandle_t cusp_handle_[NSTREAMS];
     cusolverSpHandle_t cusol_handle_[NSTREAMS];
@@ -61,6 +81,13 @@ class Naunet {
     NaunetData *h_data;
 
     /*  */
+
+    int GetCVStates(void *cv_mem, long int &nst, long int &nfe,
+                    long int &nsetups, long int &nje, long int &netf,
+                    long int &nge, long int &nni, long int &ncfn);
+    int HandleError(int flag, realtype *ab, realtype dt, realtype t0);
+    static int CheckFlag(void *flagvalue, const char *funcname, int opt,
+                         FILE *errf);
 };
 
 #ifdef PYMODULE
@@ -72,32 +99,42 @@ PYBIND11_MODULE(PYMODNAME, m) {
              py::arg("atol") = 1e-20, py::arg("rtol") = 1e-5,
              py::arg("mxsteps") = 500)
         .def("Finalize", &Naunet::Finalize)
-#ifdef USE_CUDA
+#ifdef IDX_ELEM_H
+        .def("Renorm", &Naunet::PyWrapRenorm)
+#endif
         .def("Reset", &Naunet::Reset, py::arg("nsystem") = 1,
              py::arg("atol") = 1e-20, py::arg("rtol") = 1e-5,
              py::arg("mxsteps") = 500)
+#ifdef IDX_ELEM_H
+        .def("SetReferenceAbund", &Naunet::PyWrapSetReferenceAbund,
+             py::arg("ref"), py::arg("opt") = 0)
 #endif
         .def("Solve", &Naunet::PyWrapSolve);
 
     // clang-format off
     py::class_<NaunetData>(m, "NaunetData")
         .def(py::init())
+        .def_readwrite("nH", &NaunetData::nH)
+        .def_readwrite("Tgas", &NaunetData::Tgas)
+        .def_readwrite("Tdust", &NaunetData::Tdust)
+        .def_readwrite("zeta_cr", &NaunetData::zeta_cr)
+        .def_readwrite("Av", &NaunetData::Av)
+        .def_readwrite("omega", &NaunetData::omega)
+        .def_readwrite("zeta_xr", &NaunetData::zeta_xr)
+        .def_readwrite("G0", &NaunetData::G0)
         .def_readwrite("rG", &NaunetData::rG)
-        .def_readwrite("barr", &NaunetData::barr)
         .def_readwrite("sites", &NaunetData::sites)
+        .def_readwrite("barr", &NaunetData::barr)
         .def_readwrite("hop", &NaunetData::hop)
         .def_readwrite("nMono", &NaunetData::nMono)
+        .def_readwrite("opt_frz", &NaunetData::opt_frz)
+        .def_readwrite("opt_thd", &NaunetData::opt_thd)
+        .def_readwrite("opt_crd", &NaunetData::opt_crd)
         .def_readwrite("duty", &NaunetData::duty)
         .def_readwrite("Tcr", &NaunetData::Tcr)
+        .def_readwrite("opt_uvd", &NaunetData::opt_uvd)
+        .def_readwrite("opt_rcd", &NaunetData::opt_rcd)
         .def_readwrite("branch", &NaunetData::branch)
-        .def_readwrite("nH", &NaunetData::nH)
-        .def_readwrite("zeta_cr", &NaunetData::zeta_cr)
-        .def_readwrite("zeta_xr", &NaunetData::zeta_xr)
-        .def_readwrite("Tgas", &NaunetData::Tgas)
-        .def_readwrite("Tdust", &NaunetData::Tdust)
-        .def_readwrite("Av", &NaunetData::Av)
-        .def_readwrite("G0", &NaunetData::G0)
-        .def_readwrite("omega", &NaunetData::omega)
         ;
     // clang-format on
 }