diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index aaa1735b..3ced053d 100644 --- 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 index baf5fcef..b938a0f0 100644 --- 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 index 612235ed..fceb9664 100644 --- 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 index 918f7d61..2e44a97d 100644 Binary files a/torax/tests/test_data/test_all_transport_crank_nicolson.nc and b/torax/tests/test_data/test_all_transport_crank_nicolson.nc differ diff --git a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc index eafc1fc1..cfe26c97 100644 Binary files a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc and b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc differ diff --git a/torax/tests/test_data/test_changing_config_after.nc b/torax/tests/test_data/test_changing_config_after.nc index b85ddcef..20f1e6a9 100644 Binary files a/torax/tests/test_data/test_changing_config_after.nc and b/torax/tests/test_data/test_changing_config_after.nc differ diff --git a/torax/tests/test_data/test_changing_config_before.nc b/torax/tests/test_data/test_changing_config_before.nc index 468697f6..189a6eca 100644 Binary files a/torax/tests/test_data/test_changing_config_before.nc and b/torax/tests/test_data/test_changing_config_before.nc differ diff --git a/torax/tests/test_data/test_eqdsk.nc b/torax/tests/test_data/test_eqdsk.nc index ef6d91d4..0ef28330 100644 Binary files a/torax/tests/test_data/test_eqdsk.nc and b/torax/tests/test_data/test_eqdsk.nc differ diff --git a/torax/tests/test_data/test_fixed_dt.nc b/torax/tests/test_data/test_fixed_dt.nc index 52663ab6..3a859884 100644 Binary files a/torax/tests/test_data/test_fixed_dt.nc and b/torax/tests/test_data/test_fixed_dt.nc differ diff --git a/torax/tests/test_data/test_iterbaseline_mockup.nc b/torax/tests/test_data/test_iterbaseline_mockup.nc index d542d571..e6c79094 100644 Binary files a/torax/tests/test_data/test_iterbaseline_mockup.nc and b/torax/tests/test_data/test_iterbaseline_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_mockup.nc b/torax/tests/test_data/test_iterhybrid_mockup.nc index a5af4b4a..2a2fe4b1 100644 Binary files a/torax/tests/test_data/test_iterhybrid_mockup.nc and b/torax/tests/test_data/test_iterhybrid_mockup.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_newton.nc b/torax/tests/test_data/test_iterhybrid_newton.nc index 286904d2..eb6e3c65 100644 Binary files a/torax/tests/test_data/test_iterhybrid_newton.nc and b/torax/tests/test_data/test_iterhybrid_newton.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc index 593296fe..e40c5e54 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc differ 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 index 502d606c..e40c5e54 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc index b52ad595..23b1706d 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc index 434a537c..a49ee18e 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc differ 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 index 0313ff61..9190d0ce 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc index 8ec314da..e520b5e9 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc differ 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 index f18813c0..16679613 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc differ 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 index 2bc924ef..6a1f1105 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc index f1545a5c..d93c315a 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc index f3a2b6c9..f7d7d10b 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc index 5ee9e41f..d5b51808 100644 Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc differ diff --git a/torax/tests/test_data/test_iterhybrid_rampup.nc b/torax/tests/test_data/test_iterhybrid_rampup.nc index e2c706a0..73825c3f 100644 Binary files a/torax/tests/test_data/test_iterhybrid_rampup.nc and b/torax/tests/test_data/test_iterhybrid_rampup.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc index 7e290f50..d495ab7c 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc and b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc differ diff --git a/torax/tests/test_data/test_ne_qlknn_defromchie.nc b/torax/tests/test_data/test_ne_qlknn_defromchie.nc index d312eebd..ecbc24e2 100644 Binary files a/torax/tests/test_data/test_ne_qlknn_defromchie.nc and b/torax/tests/test_data/test_ne_qlknn_defromchie.nc differ diff --git a/torax/tests/test_data/test_pc_method_ne.nc b/torax/tests/test_data/test_pc_method_ne.nc index fdd75265..402d061f 100644 Binary files a/torax/tests/test_data/test_pc_method_ne.nc and b/torax/tests/test_data/test_pc_method_ne.nc differ diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.nc b/torax/tests/test_data/test_prescribed_timedependent_ne.nc index 351b1479..9a531a55 100644 Binary files a/torax/tests/test_data/test_prescribed_timedependent_ne.nc and b/torax/tests/test_data/test_prescribed_timedependent_ne.nc differ diff --git a/torax/tests/test_data/test_psi_and_heat.nc b/torax/tests/test_data/test_psi_and_heat.nc index f370934c..402e03c6 100644 Binary files a/torax/tests/test_data/test_psi_and_heat.nc and b/torax/tests/test_data/test_psi_and_heat.nc differ diff --git a/torax/tests/test_data/test_qlknnheat.nc b/torax/tests/test_data/test_qlknnheat.nc index 93e05e59..dd4ecfa9 100644 Binary files a/torax/tests/test_data/test_qlknnheat.nc and b/torax/tests/test_data/test_qlknnheat.nc differ diff --git a/torax/tests/test_data/test_timedependence.nc b/torax/tests/test_data/test_timedependence.nc index c70de7a3..b64025bb 100644 Binary files a/torax/tests/test_data/test_timedependence.nc and b/torax/tests/test_data/test_timedependence.nc differ diff --git a/torax/transport_model/qlknn_model_wrapper.py b/torax/transport_model/qlknn_model_wrapper.py index 79d54ce9..262c8419 100644 --- 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 index 80248d2e..31c815ff 100644 --- 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: