Merge pull request #176 from darioizzo/thermonets

Thermonets
bluescarni · Jun 7, 2024 · 4bdd627 · 4bdd627
2 parents ee0d8df + 0a71007
commit 4bdd627
Show file tree

Hide file tree

Showing 22 changed files with 965 additions and 79 deletions.
diff --git a/.github/workflows/gh_actions_ci.yml b/.github/workflows/gh_actions_ci.yml
@@ -115,18 +115,17 @@ jobs:
     runs-on: windows-2022
     steps:
       - uses: actions/checkout@v4
-      - name: Add msbuild to PATH
-        uses: microsoft/[email protected]
-      - uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
           auto-update-conda: true
           python-version: "3.10"
           channels: conda-forge
           channel-priority: strict
+      - uses: ilammy/msvc-dev-cmd@v1
       - name: Build
         shell: pwsh
         run: |
-          conda install -y python=3.10 git pybind11 numpy cmake 'llvmdev=14.*' tbb-devel tbb astroquery libboost-devel fmt spdlog sleef sympy cloudpickle zlib libzlib 'mppp=1.*' numba
+          conda install -y python=3.10 git pybind11 numpy cmake llvmdev tbb-devel tbb astroquery libboost-devel fmt spdlog sleef sympy cloudpickle zlib libzlib 'mppp=1.*' numba
           git clone --depth 1 https://github.com/bluescarni/heyoka.git heyoka_cpp
           cd heyoka_cpp
           mkdir build

diff --git a/doc/api_lagham.rst b/doc/api_lagham.rst
@@ -5,6 +5,9 @@ Lagrangian and Hamiltonian mechanics
 
 .. currentmodule:: heyoka
 
+Functions
+---------
+
 .. autosummary::
    :toctree: autosummary_generated
 

diff --git a/doc/api_model.rst b/doc/api_model.rst
@@ -0,0 +1,16 @@
+.. _api_model:
+
+Models
+======
+
+.. currentmodule:: heyoka.model
+
+Functions
+---------
+
+.. autosummary::
+   :toctree: autosummary_generated
+
+   cart2geo
+   nrlmsise00_tn
+   jb08_tn
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -199,10 +199,10 @@ New
 - Implement an in-memory cache for ``llvm_state``. The cache is used
   to avoid re-optimising and re-compiling LLVM code which has
   already been optimised and compiled during the program execution
-  (`#132 <https://github.com/bluescarni/heyoka.py/pull/132>`__).
+  (`#134 <https://github.com/bluescarni/heyoka.py/pull/134>`__).
 - It is now possible to get the LLVM bitcode of
   an ``llvm_state``
-  (`#132 <https://github.com/bluescarni/heyoka.py/pull/132>`__).
+  (`#134 <https://github.com/bluescarni/heyoka.py/pull/134>`__).
 
 1.0.0 (2023-08-11)
 ------------------

diff --git a/doc/conf.py b/doc/conf.py
@@ -101,7 +101,8 @@
     "compiled_functions.ipynb",
     "Pseudo arc-length continuation*",
     "torch_and_heyoka*",
-    "differentiable_atm*"
+    "differentiable_atm*",
+    "thermoNETs*"
     ]
 
 latex_engine = "xelatex"

diff --git a/doc/examples_ml.rst b/doc/examples_ml.rst
@@ -30,4 +30,5 @@ Machine Learning
   notebooks/torch_and_heyoka
   notebooks/NeuralODEs
   notebooks/NeuralHamiltonianODEs
+  notebooks/thermoNETs
   notebooks/differentiable_atmosphere
diff --git a/doc/index.md b/doc/index.md
@@ -108,4 +108,5 @@ examples_others
 api_exsys
 api_integrators
 api_lagham
+api_model
 ```
diff --git a/doc/notebooks/compiled_functions.ipynb b/doc/notebooks/compiled_functions.ipynb
@@ -640,7 +640,7 @@
     "\n",
     "### JAX\n",
     "\n",
-    "As a last benchmark, we will be performing the same evaluation with [JAX](https://jax.readthedocs.io/en/latest/index.html). Similarly to heyoka.py, JAX offers the possibility to [JIT compile Python functions](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html#using-jit-to-speed-up-functions), so we expect similar performance to heyoka.py. Note that, in order to perform a fair comparison, for the execution of this notebook we [enabled 64-bit floats in JAX](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision) and we used JAX's CPU backend [forcing a single thread of execution](https://github.com/google/jax/issues/1539).\n",
+    "As a last benchmark, we will be performing the same evaluation with [JAX](https://jax.readthedocs.io/en/latest/index.html). Similarly to heyoka.py, JAX offers the possibility to [JIT compile Python functions](https://jax.readthedocs.io/en/latest/quickstart.html#just-in-time-compilation-with-jax-jit), so we expect similar performance to heyoka.py. Note that, in order to perform a fair comparison, for the execution of this notebook we [enabled 64-bit floats in JAX](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision) and we used JAX's CPU backend [forcing a single thread of execution](https://github.com/google/jax/issues/1539).\n",
     "\n",
     "Let us see the jax code:"
    ]

diff --git a/doc/notebooks/differentiable_atmosphere.ipynb b/doc/notebooks/differentiable_atmosphere.ipynb
diff --git a/doc/notebooks/elp2000.ipynb b/doc/notebooks/elp2000.ipynb
@@ -91,7 +91,7 @@
    "id": "de4449ca-5f30-4905-8701-469ace9cc9ff",
    "metadata": {},
    "source": [
-    "It is possible to change the expression used to represent the time variable in the ELP2000 solution via the ``time`` keyword argument. This allows to rescale and change the origin of time in the ELP2000 formulae.\n",
+    "It is possible to change the expression used to represent the time variable in the ELP2000 solution via the ``time_expr`` keyword argument. This allows to rescale and change the origin of time in the ELP2000 formulae.\n",
     "\n",
     "For instance, if we want time to be measured in Julian days (rather than centuries) since J2000, we can write:"
    ]
@@ -114,7 +114,7 @@
    ],
    "source": [
     "print(\"x coordinate, threshold = 1e-2, time in days since J2000:\\n{}\\n\"\n",
-    "      .format(hy.model.elp2000_cartesian_fk5(time=hy.time/36525., thresh = 1e-2)[0]))"
+    "      .format(hy.model.elp2000_cartesian_fk5(time_expr=hy.time/36525., thresh = 1e-2)[0]))"
    ]
   },
   {
@@ -143,7 +143,7 @@
    ],
    "source": [
     "print(\"x coordinate, threshold = 1e-2, time in JD:\\n{}\\n\"\n",
-    "      .format(hy.model.elp2000_cartesian_fk5(time=(hy.time-2451545.0)/36525., thresh = 1e-2)[0]))"
+    "      .format(hy.model.elp2000_cartesian_fk5(time_expr=(hy.time-2451545.0)/36525., thresh = 1e-2)[0]))"
    ]
   },
   {
@@ -172,7 +172,7 @@
    ],
    "source": [
     "print(\"x coordinate, threshold = 1e-2, time represented as variable 'y':\\n{}\\n\"\n",
-    "      .format(hy.model.elp2000_cartesian_fk5(time=hy.expression(\"y\"), thresh = 1e-2)[0]))"
+    "      .format(hy.model.elp2000_cartesian_fk5(time_expr=hy.expression(\"y\"), thresh = 1e-2)[0]))"
    ]
   },
   {
@@ -230,7 +230,7 @@
     "    # Build the ELP2000 formulae for the geocentric Cartesian\n",
     "    # position of the Moon in the FK5 frame. Replace the\n",
     "    # default time variable and set a custom threshold level.\n",
-    "    moon_x, moon_y, moon_z = hy.model.elp2000_cartesian_fk5(time=(tm-2451545.0)/36525.,\n",
+    "    moon_x, moon_y, moon_z = hy.model.elp2000_cartesian_fk5(time_expr=(tm-2451545.0)/36525.,\n",
     "                                                            thresh=thr)\n",
     "\n",
     "    # Compile the function for the evaluation of moon_x/y/z.\n",

diff --git a/doc/notebooks/ex_system_internals.ipynb b/doc/notebooks/ex_system_internals.ipynb
@@ -7,9 +7,7 @@
    "source": [
     "(ex_system_internals)=\n",
     "\n",
-    "# Implementation details\n",
-    "\n",
-    "In this section, we will delve into specific implementation details of heyoka.py's expression system. Our main goal is to explain certain pitfalls that can occur in innocent-looking code and which can have deleterious consequence on performance, and to teach you, dear user, how to avoid these pitfalls.\n",
+    "# Supporting large computational graphs\n",
     "\n",
     "## Reference semantics and shared (sub)expressions\n",
     "\n",

diff --git a/doc/notebooks/ex_system_revisited.ipynb b/doc/notebooks/ex_system_revisited.ipynb
@@ -7,7 +7,7 @@
    "source": [
     "(ex_system_rev)=\n",
     "\n",
-    "# Tips & tricks\n",
+    "# Common pitfalls\n",
     "\n",
     "## Long sums and products\n",
     "\n",

diff --git a/doc/notebooks/thermoNETs.ipynb b/doc/notebooks/thermoNETs.ipynb
diff --git a/doc/notebooks/vsop2013.ipynb b/doc/notebooks/vsop2013.ipynb
@@ -104,7 +104,7 @@
    "id": "7ff04f3f-3e51-4e45-bb20-a43ebdf9b0e3",
    "metadata": {},
    "source": [
-    "It is possible to change the expression used to represent the time variable in the VSOP2013 solution via the ``time`` keyword argument. This allows to rescale and change the origin of time in the VSOP2013 formulae.\n",
+    "It is possible to change the expression used to represent the time variable in the VSOP2013 solution via the ``time_expr`` keyword argument. This allows to rescale and change the origin of time in the VSOP2013 formulae.\n",
     "\n",
     "For instance, if we want time to be measured in Julian days (rather than millenia) since J2000, we can write:"
    ]
@@ -127,7 +127,7 @@
    ],
    "source": [
     "print(\"Mars' semi-major axis solution, threshold = 6e-5, time in days since J2000:\\n{}\\n\"\n",
-    "      .format(hy.model.vsop2013_elliptic(4, 1, time=hy.time/365250., thresh = 6e-5)))"
+    "      .format(hy.model.vsop2013_elliptic(4, 1, time_expr=hy.time/365250., thresh = 6e-5)))"
    ]
   },
   {
@@ -156,7 +156,7 @@
    ],
    "source": [
     "print(\"Mars' semi-major axis solution, threshold = 6e-5, time in JD:\\n{}\\n\"\n",
-    "      .format(hy.model.vsop2013_elliptic(4, 1, time=(hy.time-2451545.0)/365250., thresh = 6e-5)))"
+    "      .format(hy.model.vsop2013_elliptic(4, 1, time_expr=(hy.time-2451545.0)/365250., thresh = 6e-5)))"
    ]
   },
   {
@@ -185,7 +185,7 @@
    ],
    "source": [
     "print(\"Mars' semi-major axis solution, threshold = 6e-5, time represented as the 'x' variable:\\n{}\\n\"\n",
-    "      .format(hy.model.vsop2013_elliptic(4, 1, time=hy.expression('x'), thresh = 6e-5)))"
+    "      .format(hy.model.vsop2013_elliptic(4, 1, time_expr=hy.expression('x'), thresh = 6e-5)))"
    ]
   },
   {
@@ -248,7 +248,7 @@
     "    venus_x, venus_y, venus_z = hy.model.vsop2013_cartesian_icrf(2,\n",
     "                                                                 # NOTE: measure time\n",
     "                                                                 # using the Julian date.\n",
-    "                                                                 time=(tm-2451545.0)/365250.,\n",
+    "                                                                 time_expr=(tm-2451545.0)/365250.,\n",
     "                                                                 thresh=thr)[:3]\n",
     "\n",
     "    # Compile the function for the evaluation of venus_x/y/z.\n",

diff --git a/heyoka/_test_batch_integrator.py b/heyoka/_test_batch_integrator.py
@@ -490,8 +490,7 @@ def test_basic(self):
             self.assertTrue(ta.with_events)
             self.assertFalse(ta.compact_mode)
             self.assertFalse(ta.high_accuracy)
-            self.assertEqual(ta.state_vars, [x, v])
-            self.assertEqual(ta.rhs, [v, -9.8 * sin(x)])
+            self.assertEqual(ta.sys, sys)
 
             ta = taylor_adaptive_batch(
                 sys=sys,

diff --git a/heyoka/_test_model.py b/heyoka/_test_model.py
@@ -270,7 +270,7 @@ def test_cr3bp(self):
         self.assertTrue("0.06250000000" in str(jac))
 
     def test_ffnn(self):
-        from . import model, make_vars, tanh, expression, par, sum as hysum
+        from . import model, make_vars, expression, par, sum as hysum
 
         x, y = make_vars("x", "y")
 
@@ -285,3 +285,138 @@ def test_ffnn(self):
 
         my_ffnn4 = model.ffnn([x], [], 1, [linear], [1.2, 1.3])
         self.assertEqual(my_ffnn4[0], expression(1.3) + (expression(1.2) * x))
+
+    def test_cart2geo(self):
+        from . import model, make_vars, cfunc
+
+        x, y, z = make_vars("x", "y", "z")
+        geodesic1 = model.cart2geo([x, y, z], ecc2=0.13, R_eq=60.0, n_iters=1)
+        geodesic2 = model.cart2geo([x, y, z])
+
+        # We test on all inputs (no defaults)
+        geodesic1_cf = cfunc(geodesic1, vars=[x, y, z])
+        self.assertTrue(
+            (geodesic1_cf([1, -1, 1])[0] + 59.791916138446254) ** 2 < 1e-12**2
+        )
+        self.assertTrue(
+            (geodesic1_cf([1, -1, 1])[1] + 0.2053312550471871) ** 2 < 1e-12**2
+        )
+        self.assertTrue(
+            (geodesic1_cf([1, -1, 1])[2] + 0.7853981633974483) ** 2 < 1e-12**2
+        )
+
+        # We test that the default values are correctly set
+        geodesic2_cf = cfunc(geodesic2, vars=[x, y, z])
+        self.assertTrue(
+            (geodesic2_cf([6000000, 6000000, 6000000])[0] - 4021307.660867557) ** 2
+            < 1e-12**2
+        )
+        self.assertTrue(
+            (geodesic2_cf([6000000, 6000000, 6000000])[1] - 0.6174213396277664) ** 2
+            < 1e-12**2
+        )
+        self.assertTrue(
+            (geodesic2_cf([6000000, 6000000, 6000000])[2] - 0.7853981633974483) ** 2
+            < 1e-12**2
+        )
+
+    def test_nrlmsise00(self):
+        from . import model, make_vars, cfunc, time
+
+        h, lat, lon, f107, f107a, ap = make_vars(
+            "h", "lat", "lon", "f107", "f107a", "ap"
+        )
+        nrlmsise00 = model.nrlmsise00_tn(
+            geodetic=[h, lat, lon], f107=f107, f107a=f107a, ap=ap, time_expr=time / 86400
+        )
+        nrlmsise00_cf = cfunc([nrlmsise00], vars=[h, lat, lon, f107, f107a, ap])
+
+        # We test on zero time
+        self.assertTrue(
+            (
+                nrlmsise00_cf([600, 1.2, 3.9, 21.2, 12.2, 22.0], time=0.0)[0]
+                - 9.599548606663777e-15
+            )
+            ** 2
+            < 1e-12**2
+        )
+        # We test some days later
+        self.assertTrue(
+            (
+                nrlmsise00_cf([234, 4.5, 1.02, 4, 3, 5], time=123.23 * 86400.0)[0]
+                - 3.549961466488851e-11
+            )
+            ** 2
+            < 1e-12**2
+        )
+
+    def test_jb08(self):
+        from . import model, make_vars, cfunc, time
+
+        h, lat, lon, f107a, f107, s107a, s107, m107a, m107, y107a, y107, dDstdT = (
+            make_vars(
+                "h",
+                "lat",
+                "lon",
+                "f107a",
+                "f107",
+                "s107a",
+                "s107",
+                "m107a",
+                "m107",
+                "y107a",
+                "y107",
+                "dDstdT",
+            )
+        )
+        jb08 = model.jb08_tn(
+            geodetic=[h, lat, lon],
+            f107=f107,
+            f107a=f107a,
+            s107=s107,
+            s107a=s107a,
+            m107=m107,
+            m107a=m107a,
+            y107=y107,
+            y107a=y107a,
+            dDstdT=dDstdT,
+            time_expr=time / 86400,
+        )
+        jb08_cf = cfunc(
+            [jb08],
+            vars=[
+                h,
+                lat,
+                lon,
+                f107a,
+                f107,
+                s107a,
+                s107,
+                m107a,
+                m107,
+                y107a,
+                y107,
+                dDstdT,
+            ],
+        )
+
+        # We test on zero time
+        self.assertTrue(
+            (
+                jb08_cf([600, 1.2, 3.9, 3, 4, 5, 6, 7, 8, 9, 10, 11], time=0.0)[0]
+                - 6.805408788157112e-15
+            )
+            ** 2
+            < 1e-12**2
+        )
+        # We test some days later
+        self.assertTrue(
+            (
+                jb08_cf(
+                    [234, 4.5, 1.02, 11, 10, 9, 8, 7, 6, 5, 4, 3], time=123.23 * 86400.0
+                )[0]
+                - 1.3364825974582714e-11
+            )
+            ** 2
+            < 1e-12**2
+        )
diff --git a/heyoka/_test_scalar_integrator.py b/heyoka/_test_scalar_integrator.py
@@ -152,8 +152,7 @@ def test_basic(self):
             self.assertTrue(ta.with_events)
             self.assertFalse(ta.compact_mode)
             self.assertFalse(ta.high_accuracy)
-            self.assertEqual(ta.state_vars, [x, v])
-            self.assertEqual(ta.rhs, [v, -9.8 * sin(x)])
+            self.assertEqual(ta.sys, sys)
 
             ta = taylor_adaptive(
                 sys=sys,