Use QLKNN_7_11 by default.

PiperOrigin-RevId: 732238534

.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
 

README.md

@@ -72,6 +72,8 @@ sudo apt-get install python3-tk
 
 ### How to install
 
+#### Prepare a virtual Environment
+
 Install virtualenv (if not already installed):
 
 ```shell
@@ -102,24 +104,26 @@ Activate the virtual env:
 source toraxvenv/bin/activate
 ```
 
-Download QLKNN dependencies:
+#### Install QLKNN_7_11
 
 ```shell
-git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git
+git clone https://github.com/google-deepmind/fusion_transport_surrogates.git
+pip install -e ./fusion_transport_surrogates
+export TORAX_QLKNN_MODEL_PATH="$PWD"/fusion_transport_surrogates/fusion_transport_surrogates/models/qlknn_7_11.qlknn
 ```
 
-```shell
-export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper
-```
-
-It is recommended to automate the environment variable export. For example, if
-using bash, run:
+We recommend automating the variable export. If using bash, run:
 
 ```shell
-echo export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper >> ~/.bashrc
+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:
 
 ```shell
@@ -177,8 +181,7 @@ completed.
 To run more involved, ITER-inspired simulations, run:
 
 ```shell
-python3 run_simulation_main.py
-   --config='torax.examples.iterhybrid_rampup'
+python3 run_simulation_main.py --config='torax.examples.iterhybrid_rampup'
 ```
 
 and
@@ -192,7 +195,7 @@ run command, and environment variables. For example, for increased output
 verbosity, can run with the `--log_progress` flag.
 
 ```shell
-python3 run_simulation_main.py
+python3 run_simulation_main.py \
    --config='torax.examples.iterhybrid_rampup' --log_progress
 ```
 
@@ -285,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

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 = '~/fusion_transport_surrogates/fusion_transport_surrogates/models/qlknn_7_11.qlknn'
 
 
 def get_default_model_path() -> str:

torax/transport_model/tests/qlknn_model_wrapper.py

@@ -0,0 +1,81 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for qlknn_model_wrapper."""
+
+import tempfile
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax.numpy as jnp
+from torax.transport_model import qlknn_model_wrapper
+# pylint: disable=g-import-not-at-top
+try:
+  from fusion_transport_surrogates import qlknn_model_test_utils
+except ImportError:
+  qlknn_model_test_utils = None
+# pylint: enable=g-import-not-at-top
+
+
+def _get_test_flux_name_map():
+  if qlknn_model_test_utils is None:
+    return {}
+  return dict(
+      (flux_name, f'torax_{flux_name}')
+      for flux_name in qlknn_model_test_utils.get_test_flux_map().keys()
+  )
+
+
+class QlknnModelWrapperTest(parameterized.TestCase):
+  """Tests for qlknn_model_wrapper."""
+
+  def setUp(self):
+    super().setUp()
+    if qlknn_model_test_utils is None:
+      self.skipTest('fusion_transport_surrogates is not available.')
+    # Create a test model on disk to be loaded by the wrapper.
+    self._config = qlknn_model_test_utils.get_test_model_config()
+    self._batch_dim = 10
+    batch_dims = (1, self._batch_dim)
+    model = qlknn_model_test_utils.init_model(self._config, batch_dims)
+    self._model_file = tempfile.NamedTemporaryFile(
+        'wb', suffix='.pkl', delete=False
+    )
+    self._flux_name_map = _get_test_flux_name_map()
+    model.export_model(self._model_file.name)
+    self._qlknn_model_wrapper = qlknn_model_wrapper.QLKNNModelWrapper(
+        path=self._model_file.name,
+        flux_name_map=self._flux_name_map
+    )
+
+  def test_predict_shape(self):
+    """Tests model output shape."""
+    inputs = jnp.empty((self._batch_dim, len(self._config.input_names)))
+    outputs = self._qlknn_model_wrapper.predict(inputs)
+    self.assertLen(outputs, len(self._flux_name_map))
+    for output in outputs.values():
+      self.assertEqual(output.shape, (self._batch_dim, 1))
+
+  def test_predict_names(self):
+    """Tests model output names are the TORAX flux names."""
+    inputs = jnp.empty((self._batch_dim, len(self._config.input_names)))
+    outputs = self._qlknn_model_wrapper.predict(inputs)
+    for flux_name in self._flux_name_map.values():
+      self.assertIn(flux_name, outputs)
+
+  # TODO(b/381134347): Add tests for get_model_inputs_from_qualikiz_inputs
+  # and inputs_and_ranges.
+
+
+if __name__ == '__main__':
+  absltest.main()