Fix objective.compile in Benchmarks

.github/workflows/benchmark.yml

@@ -83,13 +83,20 @@ jobs:
           pip install -r devtools/dev-requirements.txt
           pip install matplotlib==3.9.2
 
-      - name: Benchmark with pytest-benchmark (PR)
+      - name: Action Details
         if: env.has_changes == 'true'
         run: |
           source .venv-${{ env.version }}/bin/activate
+          which python
+          python --version
           pwd
           lscpu
           pip list
+
+      - name: Benchmark with pytest-benchmark (PR)
+        if: env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
           cd tests/benchmarks
           python -m pytest benchmark_cpu_small.py -vv \
             --benchmark-save='Latest_Commit' \
@@ -114,9 +121,6 @@ jobs:
         if: env.has_changes == 'true'
         run: |
           source .venv-${{ env.version }}/bin/activate
-          pwd
-          lscpu
-          pip list
           cd tests/benchmarks
           python -m pytest benchmark_cpu_small.py -vv \
             --benchmark-save='master' \

.github/workflows/notebook_tests.yml

@@ -76,13 +76,20 @@ jobs:
           pip install -r devtools/dev-requirements.txt
           pip install matplotlib==3.9.2
 
-      - name: Test notebooks with pytest and nbmake
+      - name: Action Details
         if: env.has_changes == 'true'
         run: |
           source .venv-${{ env.version }}/bin/activate
+          which python
+          python --version
           pwd
           lscpu
           pip list
+
+      - name: Test notebooks with pytest and nbmake
+        if: env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
           export PYTHONPATH=$(pwd)
           pytest -v --nbmake "./docs/notebooks" \
             --nbmake-timeout=2000 \

.github/workflows/regression_tests.yml

@@ -81,13 +81,20 @@ jobs:
         with:
           swap-size-gb: 10
 
-      - name: Test with pytest
+      - name: Action Details
         if: env.has_changes == 'true'
         run: |
           source .venv-${{ env.version }}/bin/activate
+          which python
+          python --version
           pwd
           lscpu
           pip list
+
+      - name: Test with pytest
+        if: env.has_changes == 'true'
+        run: |
+          source .venv-${{ env.version }}/bin/activate
           python -m pytest -v -m regression\
             --durations=0 \
             --cov-report xml:cov.xml \

desc/objectives/objective_funs.py

@@ -891,12 +891,12 @@
                 timer.disp("Hessian compilation time")
         if mode in ["lsq", "all"]:
             timer.start("Objective compilation time")
-            _ = self.compute_scaled(x, self.constants).block_until_ready()
+            _ = self.compute_scaled_error(x, self.constants).block_until_ready()
             timer.stop("Objective compilation time")
             if verbose > 1:
                 timer.disp("Objective compilation time")
             timer.start("Jacobian compilation time")
-            _ = self.jac_scaled(x, self.constants).block_until_ready()
+            _ = self.jac_scaled_error(x, self.constants).block_until_ready()
             timer.stop("Jacobian compilation time")
             if verbose > 1:
                 timer.disp("Jacobian compilation time")

docs/index.rst

@@ -21,6 +21,7 @@
    :caption: Tutorials
 
    installation
+   performance_tips
    notebooks/tutorials/basic_equilibrium.ipynb
    notebooks/tutorials/advanced_equilibrium_continuation.ipynb
    notebooks/tutorials/continuation_step_by_step.ipynb
@@ -35,7 +36,6 @@
    notebooks/tutorials/coil_stage_two_optimization.ipynb
    notebooks/tutorials/QFM_surface.ipynb
    notebooks/tutorials/ideal_ballooning_stability.ipynb
-   memory_usage
 
 .. toctree::
    :maxdepth: 1

docs/notebooks/tutorials/nae_constraint.ipynb

@@ -49,7 +49,7 @@
     "import os\n",
     "\n",
     "sys.path.insert(0, os.path.abspath(\".\"))\n",
-    "sys.path.append(os.path.abspath(\"../../\"))"
+    "sys.path.append(os.path.abspath(\"../../../\"))"
    ]
   },
   {
@@ -812,7 +812,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.0"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

docs/memory_usage.rst → docs/performance_tips.rst

@@ -1,3 +1,24 @@
+================
+Performance Tips
+================
+
+Caching the Compiled (Jitted) Code
+----------------------------------
+Although the compiled code is fast, it still takes time to compile. If you are running the same optimization, or some similar optimization, multiple times, you can save time by caching the compiled code. This automatically happens for a single session (for example, until you restart your kernel in Jupyter Notebook) but once you start using another session, the code will need to be recompiled. Fortunately, there is a way to bypass this. First create a cache directory (i.e. ``jax-caches``), and put the following code at the beginning of your script:
+
+.. code-block:: python
+
+    import jax
+    import jax.numpy as jnp
+
+    jax.config.update("jax_compilation_cache_dir", "../jax-caches")
+    jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+    jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+
+This will use a directory called ``jax-caches`` in the parent directory of the script to store the compiled code. The ``jax_persistent_cache_min_entry_size_bytes`` and ``jax_persistent_cache_min_compile_time_secs`` parameters are set to -1 and 0, respectively, to ensure that all compiled code is cached. For more details on caching, refer to official JAX documentation [`here <https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html#persistent-compilation-cache>`__].
+
+Note: Updating JAX version might re-compile some previously cached code, and this might increase the cache size. Every once in a while, you might need to clear your cache directory.
+
 
 Reducing Memory Size of Objective Jacobian Calculation
 ------------------------------------------------------