Use QLKNN_7_11 by default.

PiperOrigin-RevId: 732238534
google-deepmind · Mar 3, 2025 · 24ccdbb · 24ccdbb
1 parent dcab336
commit 24ccdbb
Show file tree

Hide file tree

Showing 33 changed files with 95 additions and 56 deletions.
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
@@ -63,10 +63,12 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
-    - run: git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git
-    - run: echo "TORAX_QLKNN_MODEL_PATH=$PWD/qlknn-hyper" >> "$GITHUB_ENV"
-
     - run: pip --version
+
+    - run: git clone https://github.com/google-deepmind/fusion_transport_surrogates.git
+    - run: pip install -e ./fusion_transport_surrogates
+    - run: echo "TORAX_QLKNN_MODEL_PATH=$PWD/fusion_transport_surrogates/fusion_transport_surrogates/models/qlknn_7_11.qlknn" >> "$GITHUB_ENV"
+
     - run: pip install -e .[dev]
     - run: pip freeze
 

diff --git a/README.md b/README.md
@@ -104,29 +104,7 @@ Activate the virtual env:
 source toraxvenv/bin/activate
 ```
 
-#### Install QLKNN-hyper
-
-Download QLKNN-hyper dependencies:
-
-```shell
-git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git
-```
-
-```shell
-export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper
-```
-
-It is recommended to automate the environment variable export. For example, if
-using bash, run:
-
-```shell
-echo export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper >> ~/.bashrc
-```
-The above command only needs to be run once on a given system.
-
-#### (Optional) Install QLKNN_7_11
-
-Optionally, you can instead use QLKNN_7_11, a more recent surrogate model:
+#### Install QLKNN_7_11
 
 ```shell
 git clone https://github.com/google-deepmind/fusion_transport_surrogates.git
@@ -140,6 +118,10 @@ We recommend automating the variable export. If using bash, run:
 echo export TORAX_QLKNN_MODEL_PATH="$PWD"/fusion_transport_surrogates/fusion_transport_surrogates/models/qlknn_7_11.qlknn >> ~/.bashrc
 ```
 
+The above command only needs to be run once on a given system.
+
+An alternative to QLKNN_7_11 is [QLKNN-hyper](#optional-install-qlknn-hyper)
+
 #### Install TORAX
 
 Download and install the TORAX codebase via http:
@@ -306,6 +288,34 @@ You can get out of the Python virtual env by deactivating it:
 deactivate
 ```
 
+#### (Optional) Install QLKNN-hyper
+
+An alternative to QLKNN_7_11 is to use QLKNN-hyper (also known as QLKNN10D)
+instead. QLKNN_7_11 was trained on more recent data that includes more
+information about impurities, and better coverage of the edge region compared
+to QLKNN-hyper. However, it is still widely used in other simulators, so it can
+be useful for comparative studies for instance.
+
+Download QLKNN-hyper dependencies:
+
+```shell
+git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git
+```
+
+Make QLKNN-hyper the default QLKNN model.
+
+```shell
+export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper
+```
+
+It is recommended to automate the environment variable export. For example, if
+using bash, run:
+
+```shell
+echo export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper >> ~/.bashrc
+```
+The above command only needs to be run once on a given system.
+
 ## Simulation tutorials
 
 Under construction

diff --git a/docs/installation.rst b/docs/installation.rst
@@ -54,36 +54,11 @@ Activate the virtual env:
 
 It is convenient to set up an alias for the above command.
 
-.. _install_qlknn_hyper:
-
-Install QLKNN-hyper
--------------------
-
-Download QLKNN dependencies:
-
-.. code-block:: console
-
-  git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git
-
-.. code-block:: console
-
-  export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper
-
-It is recommended to automate the environment variable export. For example, if using bash, run:
-
-.. code-block:: console
-
-  echo export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper >> ~/.bashrc
-
-The above command only needs to be run once on a given system.
-
 .. _install_qlknn_7_11:
 
-(Optional) Install QLKNN_7_11
+Install QLKNN_7_11
 -----------------------------
 
-Optionally, you can instead use QLKNN_7_11, a more recent surrogate model:
-
 .. code-block:: console
 
   git clone https://github.com/google-deepmind/fusion_transport_surrogates.git
@@ -96,6 +71,10 @@ We recommend automating the variable export. If using bash, run:
 
   echo export TORAX_QLKNN_MODEL_PATH="$PWD"/fusion_transport_surrogates/fusion_transport_surrogates/models/qlknn_7_11.qlknn >> ~/.bashrc
 
+The above command only needs to be run once on a given system.
+
+An alternative to QLKNN_7_11 is QLKNN-hyper (See :ref:`install_qlknn_hyper`).
+
 .. install_torax:
 
 Install TORAX
@@ -148,6 +127,11 @@ using bash, run:
 
 The above command only needs to be run once on a given system.
 
+.. _dev_install:
+
+(Optional) Install TORAX in development mode
+--------------------------------------------
+
 **Recommended** for developers. Instead of the above, install optional dependencies
 for (parallel) pytest and documentation generation. Also install in editable mode to
 not require reinstallation for every change.
@@ -156,6 +140,42 @@ not require reinstallation for every change.
 
   cd torax; pip install -e .[dev]
 
+.. _dev_install:
 
-Optional: Install additional GPU support for JAX if your machine has a GPU:
+(Optional) GPU support
+-------------------
+
+Install additional GPU support for JAX if your machine has a GPU:
 https://jax.readthedocs.io/en/latest/installation.html#supported-platforms
+
+
+.. _install_qlknn_hyper:
+
+(Optional) Install QLKNN-hyper
+-------------------
+
+An alternative to QLKNN_7_11 is to use QLKNN-hyper (also known as QLKNN10D)
+instead. QLKNN_7_11 was trained on more recent data that includes more
+information about impurities, and better coverage of the edge region compared
+to QLKNN-hyper. However, it is still widely used in other simulators, so it can
+be useful for comparative studies for instance.
+
+Download QLKNN dependencies:
+
+.. code-block:: console
+
+  git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git
+
+.. code-block:: console
+
+  export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper
+
+It is recommended to automate the environment variable export. For example, if
+using bash, run:
+
+.. code-block:: console
+
+  echo export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper >> ~/.bashrc
+
+The above command only needs to be run once on a given system.
+
diff --git a/torax/tests/test_data/test_all_transport_crank_nicolson.nc b/torax/tests/test_data/test_all_transport_crank_nicolson.nc
diff --git a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc
diff --git a/torax/tests/test_data/test_changing_config_after.nc b/torax/tests/test_data/test_changing_config_after.nc
diff --git a/torax/tests/test_data/test_changing_config_before.nc b/torax/tests/test_data/test_changing_config_before.nc
diff --git a/torax/tests/test_data/test_eqdsk.nc b/torax/tests/test_data/test_eqdsk.nc
diff --git a/torax/tests/test_data/test_fixed_dt.nc b/torax/tests/test_data/test_fixed_dt.nc
diff --git a/torax/tests/test_data/test_iterbaseline_mockup.nc b/torax/tests/test_data/test_iterbaseline_mockup.nc
diff --git a/torax/tests/test_data/test_iterhybrid_mockup.nc b/torax/tests/test_data/test_iterhybrid_mockup.nc
diff --git a/torax/tests/test_data/test_iterhybrid_newton.nc b/torax/tests/test_data/test_iterhybrid_newton.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc
diff --git a/...sts/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc b/...sts/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc
diff --git a/torax/tests/test_data/test_iterhybrid_rampup.nc b/torax/tests/test_data/test_iterhybrid_rampup.nc
diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc
diff --git a/torax/tests/test_data/test_ne_qlknn_defromchie.nc b/torax/tests/test_data/test_ne_qlknn_defromchie.nc
diff --git a/torax/tests/test_data/test_pc_method_ne.nc b/torax/tests/test_data/test_pc_method_ne.nc
diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.nc b/torax/tests/test_data/test_prescribed_timedependent_ne.nc
diff --git a/torax/tests/test_data/test_psi_and_heat.nc b/torax/tests/test_data/test_psi_and_heat.nc
diff --git a/torax/tests/test_data/test_qlknnheat.nc b/torax/tests/test_data/test_qlknnheat.nc
diff --git a/torax/tests/test_data/test_timedependence.nc b/torax/tests/test_data/test_timedependence.nc
diff --git a/torax/transport_model/qlknn_model_wrapper.py b/torax/transport_model/qlknn_model_wrapper.py
@@ -17,16 +17,23 @@
 
 from collections.abc import Mapping
 from typing import Final
+
+from absl import logging
 import immutabledict
 import jax
 import jax.numpy as jnp
 from torax.transport_model import base_qlknn_model
 from torax.transport_model import qualikiz_based_transport_model
+
+
 # pylint: disable=g-import-not-at-top
 try:
   from fusion_transport_surrogates import qlknn_model
   _FUSION_TRANSPORT_SURROGATES_AVAILABLE = True
-except ImportError:
+except ImportError as e:
+  logging.error(
+      'Error importing fusion_transport_surrogates: %s', e, exc_info=True
+  )
   _FUSION_TRANSPORT_SURROGATES_AVAILABLE = False
 # pylint: enable=g-import-not-at-top
 
@@ -53,7 +60,7 @@ def __init__(
     if not _FUSION_TRANSPORT_SURROGATES_AVAILABLE:
       raise ImportError(
           'QLKNNModelWrapper requires fusion_transport_surrogates to be '
-          'installed.'
+          'installed. There was an error importing it.'
       )
     if flux_name_map is None:
       flux_name_map = _FLUX_NAME_MAP

diff --git a/torax/transport_model/qlknn_transport_model.py b/torax/transport_model/qlknn_transport_model.py
@@ -42,7 +42,7 @@
 MODEL_PATH_ENV_VAR: Final[str] = 'TORAX_QLKNN_MODEL_PATH'
 # If no path is set in either the config or the environment variable, use
 # this path.
-DEFAULT_MODEL_PATH = '~/qlknn_hyper'
+DEFAULT_MODEL_PATH: Final[str] = '~/fusion_transport_surrogates/fusion_transport_surrogates/models/qlknn_7_11.qlknn'  # pylint: disable=line-too-long
 
 
 def get_default_model_path() -> str: