Use QLKNN_7_11 in tests and make QLKNN_7_11 the default model in install instructions.

.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_surrogates.git
+    - run: pip install -e ./fusion_surrogates
+    - run: echo "TORAX_QLKNN_MODEL_PATH=$PWD/fusion_surrogates/fusion_surrogates/models/qlknn_7_11.qlknn" >> "$GITHUB_ENV"
+
     - run: pip install -e .[dev]
     - run: pip freeze
 

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_surrogates.git
@@ -140,6 +118,10 @@ We recommend automating the variable export. If using bash, run:
 echo export TORAX_QLKNN_MODEL_PATH="$PWD"/fusion_surrogates/fusion_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,36 @@ 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-10D, also known as QLKNN10D
+[K.L. van de Plassche PoP 2020](https://doi.org/10.1063/1.5134126). QLKNN_7_11
+is based on QuaLiKiz 2.8.1 which has an improved collision operator compared to
+the QLKNN10D training set. QLKNN_7_11 training data includes impurity density
+gradients as an input feature and has better coverage of the near-LCFS region
+compared to QLKNN-hyper-10D. 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

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_surrogates.git
@@ -96,6 +71,10 @@ We recommend automating the variable export. If using bash, run:
 
   echo export TORAX_QLKNN_MODEL_PATH="$PWD"/fusion_surrogates/fusion_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,44 @@ 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-10D, also known as QLKNN10D
+[K.L. van de Plassche PoP 2020](https://doi.org/10.1063/1.5134126). QLKNN_7_11
+is based on QuaLiKiz 2.8.1 which has an improved collision operator compared to
+the QLKNN10D training set. QLKNN_7_11 training data includes impurity density
+gradients as an input feature and has better coverage of the near-LCFS region
+compared to QLKNN-hyper-10D. 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.
+

torax/tests/sim_no_compile_test.py

@@ -36,7 +36,7 @@ class SimNoCompileTest(sim_test_case.SimTestCase):
           'test_qlknnheat.py',
           _ALL_PROFILES,
           0,
-          1e-11,
+          1e-10,
       ),
   )
   def test_run_simulation(

torax/tests/sim_test.py

@@ -99,7 +99,7 @@ class SimTest(sim_test_case.SimTestCase):
           'test_qlknnheat.py',
           _ALL_PROFILES,
           0,
-          1e-11,
+          1e-10,
       ),
       # Tests fixed_dt timestep
       (
@@ -188,6 +188,7 @@ class SimTest(sim_test_case.SimTestCase):
           'test_all_transport_fusion_qlknn.py',
           _ALL_PROFILES,
           0,
+          1e-8,
       ),
       # Tests CHEASE geometry. Implicit solver. Heat transport only.
       (
@@ -347,6 +348,7 @@ class SimTest(sim_test_case.SimTestCase):
           'test_iterhybrid_predictor_corrector_zeffprofile.py',
           _ALL_PROFILES,
           0,
+          1e-8,
       ),
       # Predictor-corrector solver with a time-dependent isotope mix.
       (
@@ -708,12 +710,12 @@ def test_ip_bc_vloop_bc_equivalence(self):
       np.testing.assert_allclose(
           sim_outputs_ip_bc.core_profiles[profile].value[middle_index, :],
           sim_outputs_vloop_bc.core_profiles[profile].value[middle_index, :],
-          rtol=1e-4,
+          rtol=1e-3,
       )
       np.testing.assert_allclose(
           sim_outputs_ip_bc.core_profiles[profile].value[-1, :],
           sim_outputs_vloop_bc.core_profiles[profile].value[-1, :],
-          rtol=1e-4,
+          rtol=1e-3,
       )
 
     # pylint: enable=invalid-name

torax/transport_model/pydantic_model.py

@@ -93,6 +93,9 @@ def _conform_data(cls, data: dict[str, Any]) -> dict[str, Any]:
         # The QLK version this specific QLKNN was trained on tends to
         # underpredict ITG electron heat flux in shaped, high-beta scenarios.
         data['ITG_flux_ratio_correction'] = 2.0
+    else:
+      if 'smoothing_sigma' not in data:
+        data['smoothing_sigma'] = 0.1
     return data
 
   def build_transport_model(self) -> qlknn_transport_model.QLKNNTransportModel:

torax/transport_model/qlknn_model_wrapper.py

@@ -14,18 +14,25 @@
 """A wrapper for QLKNN transport surrogate models."""
 from collections.abc import Mapping
 from typing import Final
+
+from absl import logging
 import immutabledict
 import jax
 import jax.numpy as jnp
 from torax import jax_utils
 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_surrogates import qlknn_model
 
   _FUSION_SURROGATES_AVAILABLE = True
-except ImportError:
+except ImportError as e:
+  logging.error(
+      'Error importing fusion_surrogates: %s', e, exc_info=True
+  )
   _FUSION_SURROGATES_AVAILABLE = False
 # pylint: enable=g-import-not-at-top
 
@@ -51,7 +58,8 @@ def __init__(
   ):
     if not _FUSION_SURROGATES_AVAILABLE:
       raise ImportError(
-          'QLKNNModelWrapper requires fusion_surrogates to be installed.'
+          'QLKNNModelWrapper requires fusion_surrogates to be installed. '
+          'There was an error importing it.'
       )
     if flux_name_map is None:
       flux_name_map = _FLUX_NAME_MAP

torax/transport_model/qlknn_transport_model.py

@@ -37,7 +37,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_surrogates/fusion_surrogates/models/qlknn_7_11.qlknn'  # pylint: disable=line-too-long
 
 
 def get_default_model_path() -> str:

torax/transport_model/tests/pydantic_model_test.py

@@ -80,16 +80,16 @@ def test_build_qualikiz_transport_model(self):
         transport_model, qualikiz_transport_model.QualikizTransportModel
     )
 
-  def test_qlknn_flux_ratio_correction_default(self):
-    """Tests that the flux ratio correction is set to 2.0 for QLKNN."""
+  def test_qlknn_defaults(self):
+    """Tests that correct default values are set for QLKNN."""
     transport = transport_pydantic_model.Transport.from_dict(
         {'transport_model': 'qlknn'}
     )
     transport_model_config = transport.transport_model_config
     self.assertIsInstance(
         transport_model_config, transport_pydantic_model.QLKNNTransportModel
     )
-    self.assertEqual(transport_model_config.ITG_flux_ratio_correction, 2.0)
+    self.assertEqual(transport_model_config.smoothing_sigma, 0.1)
     self.assertEqual(transport_model_config.ETG_correction_factor, 1.0 / 3.0)