Merge branch 'feature/IMAS_coupling' into subfeature/core_profiles_to_imas

MateoBell · web-flow · commit 5532c13f9fdc · 2025-06-12T11:57:53.000+02:00
diff --git a/README.md b/README.md
@@ -126,7 +126,9 @@ The following command will run TORAX using the default configuration file
 run_torax --config='examples/basic_config.py'
 ```
 
-Simulation progress is shown by a terminal progress bar indicating the current time and percentage completed.
+
+Simulation progress is shown by a terminal progress bar indicating the current
+time and percentage completed.
 
 To run more involved, ITER-inspired simulations, run:
 
diff --git a/docs/conf.py b/docs/conf.py
@@ -310,6 +310,9 @@
 .. |QLKNN| replace:: `QLKNN <QLKNN_target_>`_
 .. _QLKNN_target: {github_base_url}/torax/_src/transport_model/qlknn_transport_model.py
 
+.. |QuaLiKiz| replace:: `QuaLiKiz <qualikiz_target_>`_
+.. _qualikiz_target: {github_base_url}/torax/_src/transport_model/qualikiz_transport_model.py
+
 .. |transport_model| replace:: `transport_model <torax_src_transport_model_target_>`_
 .. _torax_src_transport_model_target: {github_base_url}/torax/_src/transport_model
 
diff --git a/docs/interfacing_with_surrogates.rst b/docs/interfacing_with_surrogates.rst
@@ -1,12 +1,17 @@
 .. _interfacing_with_surrogates:
 
-JAX-compatible interfaces with ML surrogates of physics models
+JAX-compatible interfaces with ML-surrogates of physics models
 ##############################################################
 
-This section discusses a variety of options for building JAX-friendly interfaces to surrogate models.
+This section discusses a variety of options for building JAX-friendly interfaces
+to surrogate models.
+
+As an illustrative example, suppose we have a new neural network surrogate
+transport model that we would like to use in TORAX. Assume that all the
+boilerplate described in the previous sections has been taken care of, as well
+as the definition of some functions to convert between TORAX structures and
+tensors for the neural network.
 
-As an illustrative example, suppose we have a new neural network surrogate transport model that we would like to use in TORAX.
-Assume that all the boilerplate described in the previous sections has been taken care of, as well as the definition of some functions to convert between TORAX structures and tensors for the neural network.
 
 .. code-block:: python
 
@@ -31,23 +36,33 @@ Assume that all the boilerplate described in the previous sections has been take
             v_e=v_e,
         )
 
-In this guide, we explore a few options for how you could make the ``_call_surrogate_model`` function for an existing surrogate, while maintaining the full power of JAX:
+In this guide, we explore a few options for how you could make the
+``_call_surrogate_model`` function for an existing surrogate, while maintaining
+the full power of JAX:
 
-1. **Manually reimplementing the model in JAX**,
-2. **Converting a Pytorch model to a JAX model**,
+1. **Manually reimplementing the model in JAX**.
+2. **Converting a Pytorch model to a JAX model**.
 3. **Using an ONNX model**.
 
 .. note::
-    These conversion methods are necessary in order to make an external model compatible with JAX's autodiff and JIT functionality, which is required for using TORAX's gradient-driven nonlinear solvers (e.g. Newton-Raphson).
-    Interfacing with non-differentiable, non-JITtable models is possible (for an example, see the `QuaLiKiz transport model`_) if the linear solver is used. However, note that if the model is called within the step function JIT will need to be disabled with ``TORAX_COMPILATION_ENABLED=0``.
+    These conversion methods are necessary in order to make an external model
+    compatible with JAX's autodiff and JIT functionality, which is required for
+    using TORAX's gradient-driven nonlinear solvers (e.g. Newton-Raphson).
+    Interfacing with non-differentiable, non-JITtable models is possible
+    (for an example, see the |QuaLiKiz| transport model implementation) if the
+    linear solver is used. However, note that if the model is called within the
+    step function, JIT will need to be disabled with
+    ``TORAX_COMPILATION_ENABLED=0``.
 
 
 Option 1: manually reimplementing the model in JAX
 ==================================================
 
-If the architecture of the surrogate is sufficiently simple, you might consider reimplementing the model in JAX.
-The surrogates in TORAX are mostly implemented using `Flax Linen`_, and can be found in the |fusion_surrogates|_ repository.
-If you're not familiar with Flax, you can check out the `Flax documentation`_ on how to define your own models.
+If the architecture of the surrogate is sufficiently simple, you might consider
+reimplementing the model in JAX. The surrogates in TORAX are mostly implemented
+using `Flax Linen`_, and can be found in the |fusion_surrogates|_ repository.
+If you're not familiar with Flax, you can check out the `Flax documentation`_
+on how to define your own models.
 
 Consider a PyTorch neural network,
 
@@ -97,8 +112,9 @@ This model can be replicated in Flax as follows:
 
     flax_model = FlaxMLP(hidden_dim, n_hidden, output_dim, input_dim)
 
-As this is only the model architecture, we need to load the trained weights separately.
-This can be a bit fiddly as you have to map from the parameter names in the weights checkpoint file to the parameter names in the Flax model.
+As this is only the model architecture, we need to load the trained weights
+separately. This can be a bit fiddly as you have to map from the parameter names
+in the weights checkpoint file to the parameter names in the Flax model.
 
 For loading weights from a PyTorch checkpoint, you might do something like:
 
@@ -130,24 +146,30 @@ The model can then be called like any Flax model,
 
 
 .. warning::
-    You need to be very careful when loading from a PyTorch state dict, as Flax and PyTorch may have slightly different representations of the weights (for example, one could be the transpose of the other). It's worth validating the output of your PyTorch model against your JAX model to make sure.
-
+    You need to be very careful when loading from a PyTorch state dict, as
+    Flax and PyTorch may have slightly different representations of the weights
+    (for example, one could be the transpose of the other). It's worth
+    validating the output of your PyTorch model against your JAX model to make
+    sure.
 
 
 Option 2: converting a PyTorch model to a JAX model
 ===================================================
 
 .. warning::
-    The `torch_xla2`_ package is still evolving, which means there may be unexpected breaking changes. Some of the methods described in this section may become deprecated with little warning.
+    The `torch_xla2`_ package is still evolving, which means there may be
+    unexpected breaking changes. Some of the methods described in this section
+    may become deprecated with little warning.
 
-If your model is in PyTorch, you could also consider using the `torch_xla2`_ package to do the conversion to JAX automatically.
+If your model is in PyTorch, you could also consider using the `torch_xla2`_
+package to do the conversion to JAX automatically.
 
 .. code-block:: python
 
     import torch
     import torch_xla2 as tx
 
-    trained_model = torch.load(PYTORCH_MODEL_PATH, weights_only=False) # Use weights_only=False if you want to load the full model
+    trained_model = torch.load(PYTORCH_MODEL_PATH, weights_only=False)  # Use weights_only=False if you want to load the full model
     params, jax_model_from_torch = tx.extract_jax(model)
 
 The model can then be called as a pure JAX function:
@@ -156,7 +178,9 @@ The model can then be called as a pure JAX function:
 
     output_tensor = jax.jit(jax_model_from_torch)(params, input_tensor)
 
-To remove the need for performing the conversion every time the model is loaded, you might want to save a JAX-compatible version of the weights and model to disk:
+To remove the need for performing the conversion every time the model is loaded,
+you might want to save a JAX-compatible version of the weights and model to
+disk:
 
 .. code-block:: python
 
@@ -181,16 +205,19 @@ The model can then be loaded and run as follows:
       model = jax.export.deserialize(model_as_bytes)
 
     # Load the weights
-    weights_as_npz = jnp.load('weights.npz')
+    weights_as_npz = np.load('weights.npz')
     weights = [jnp.array(v) for v in weights_as_npz.values()]
 
 
 Option 3: using an ONNX model
 =============================
 
-The `Open Neural Network Exchange`_ format (ONNX) is a highly interoperable format for sharing neural network models. ONNX files include the model architecture and weights bundled together.
+The `Open Neural Network Exchange`_ format (ONNX) is a highly interoperable
+format for sharing neural network models. ONNX files include the model
+architecture and weights bundled together.
 
-An ONNX model can be loaded and called as follows, making sure to specify the correct input and output node names for your specific model:
+An ONNX model can be loaded and called as follows, making sure to specify the
+correct input and output node names for your specific model:
 
 .. code-block:: python
 
@@ -207,7 +234,8 @@ An ONNX model can be loaded and called as follows, making sure to specify the co
     )
 
 However, JAX will not be able to differentiate through the InferenceSession.
-To convert the ONNX model to a JAX representation, you can use the `jaxonnxruntime`_ package:
+To convert the ONNX model to a JAX representation, you can use the
+`jaxonnxruntime`_ package:
 
 .. code-block:: python
 
@@ -225,10 +253,12 @@ To convert the ONNX model to a JAX representation, you can use the `jaxonnxrunti
 Best practices
 ==============
 
-**Caching and lazy loading**: Ideally, the model should be constructed and weights loaded once only, on the first call to the function.
-The loaded model should be cached and reused for subsequent calls.
+**Caching and lazy loading**: Ideally, the model should be constructed and
+weights loaded once only, on the first call to the function. The loaded model
+should be cached and reused for subsequent calls.
 
-For example, in the ``_combined`` function of the QLKNN transport model (the function that actually evaluates this model), we have:
+For example, in the ``_combined`` function of the QLKNN transport model (the
+function that actually evaluates this model), we have:
 
 .. code-block:: python
 
@@ -242,18 +272,21 @@ where
 
     @functools.lru_cache(maxsize=1)
     def get_model(path: str) -> base_qlknn_model.BaseQLKNNModel:
-    """Load the model."""
-    ...
-    return qlknn_10d.QLKNN10D(path)
+      """Load the model."""
+      ...
+      return qlknn_10d.QLKNN10D(path)
 
-By decorating with ``functools.lru_cache(maxsize=1)``, the result of this function - the loaded model - is stored in the cache and is only re-loaded if the function is called with a different ``path``.
+By decorating with ``functools.lru_cache(maxsize=1)``, the result of this
+function - the loaded model - is stored in the cache and is only re-loaded if
+the function is called with a different ``path``.
 
-**JITting model calls**: In general, you should make sure that your forward call of the model is JITted:
+**JITting model calls**: In general, you should make sure that your forward call
+of the model is JITted:
 
 .. code-block:: python
 
-    output_tensor = jax.jit(flax_model.apply)(params, input_tensor) # Good
-    output_tensor = flax_model.apply(params, input_tensor) # Bad
+    output_tensor = jax.jit(flax_model.apply)(params, input_tensor)  # Good
+    output_tensor = flax_model.apply(params, input_tensor)  # Bad
 
 This is vital to ensure fast performance.
 
@@ -262,6 +295,5 @@ This is vital to ensure fast performance.
 .. _torch_xla2: https://pytorch.org/xla/master/features/stablehlo.html
 .. _Open Neural Network Exchange: https://onnx.ai/
 .. _jaxonnxruntime: https://github.com/google/jaxonnxruntime
-.. _QuaLiKiz transport model: https://github.com/google-deepmind/torax/blob/main/torax/transport_model/qualikiz_transport_model.py
 .. |fusion_surrogates| replace:: ``google-deepmind/fusion_surrogates``
 .. _fusion_surrogates: https://github.com/google-deepmind/fusion_surrogates
diff --git a/docs/model_integration.rst b/docs/model_integration.rst
@@ -12,6 +12,11 @@ exposed this as part of the TORAX API.
 
 If you would like to use this please reach out to us. We aim to expose this
 functionality as part of the TORAX API in the very near future to further
-facilitate the integration of custom models.
+facilitate the integration of custom models, and further expand the
+documentation.
 
-For information on JAX-friendly interfacing with ML-surrogates of physics models, see :ref:`interfacing_with_surrogates`.
+.. toctree::
+   :maxdepth: 1
+   :caption: Model Integration Topics
+
+   interfacing_with_surrogates
diff --git a/torax/_src/orchestration/run_simulation.py b/torax/_src/orchestration/run_simulation.py
@@ -44,66 +44,66 @@ def prepare_simulation(
     transport_model = torax_config.transport.build_transport_model()
     pedestal_model = torax_config.pedestal.build_pedestal_model()
 
-  geometry_provider = torax_config.geometry.build_provider
-  source_models = source_models_lib.SourceModels(
-      torax_config.sources, neoclassical=torax_config.neoclassical
-  )
-
-  static_runtime_params_slice = (
-      build_runtime_params.build_static_params_from_config(torax_config)
-  )
-
-  solver = torax_config.solver.build_solver(
-      static_runtime_params_slice=static_runtime_params_slice,
-      transport_model=transport_model,
-      source_models=source_models,
-      pedestal_model=pedestal_model,
-  )
-
-  mhd_models = torax_config.mhd.build_mhd_models(
-      static_runtime_params_slice=static_runtime_params_slice,
-      transport_model=transport_model,
-      source_models=source_models,
-      pedestal_model=pedestal_model,
-  )
-
-  step_fn = step_function.SimulationStepFn(
-      solver=solver,
-      time_step_calculator=torax_config.time_step_calculator.time_step_calculator,
-      transport_model=transport_model,
-      pedestal_model=pedestal_model,
-      mhd_models=mhd_models,
-  )
-
-  dynamic_runtime_params_slice_provider = (
-      build_runtime_params.DynamicRuntimeParamsSliceProvider.from_config(
-          torax_config
-      )
-  )
-
-  if torax_config.restart and torax_config.restart.do_restart:
-    initial_state, post_processed_outputs = (
-        initial_state_lib.get_initial_state_and_post_processed_outputs_from_file(
-            t_initial=torax_config.numerics.t_initial,
-            file_restart=torax_config.restart,
-            static_runtime_params_slice=static_runtime_params_slice,
-            dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
-            geometry_provider=geometry_provider,
-            step_fn=step_fn,
-        )
+    geometry_provider = torax_config.geometry.build_provider
+    source_models = source_models_lib.SourceModels(
+        torax_config.sources, neoclassical=torax_config.neoclassical
+    )
+
+    static_runtime_params_slice = (
+        build_runtime_params.build_static_params_from_config(torax_config)
+    )
+
+    solver = torax_config.solver.build_solver(
+        static_runtime_params_slice=static_runtime_params_slice,
+        transport_model=transport_model,
+        source_models=source_models,
+        pedestal_model=pedestal_model,
+    )
+
+    mhd_models = torax_config.mhd.build_mhd_models(
+        static_runtime_params_slice=static_runtime_params_slice,
+        transport_model=transport_model,
+        source_models=source_models,
+        pedestal_model=pedestal_model,
     )
-    restart_case = True
-  else:
-    initial_state, post_processed_outputs = (
-        initial_state_lib.get_initial_state_and_post_processed_outputs(
-            t=torax_config.numerics.t_initial,
-            static_runtime_params_slice=static_runtime_params_slice,
-            dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
-            geometry_provider=geometry_provider,
-            step_fn=step_fn,
+
+    step_fn = step_function.SimulationStepFn(
+        solver=solver,
+        time_step_calculator=torax_config.time_step_calculator.time_step_calculator,
+        transport_model=transport_model,
+        pedestal_model=pedestal_model,
+        mhd_models=mhd_models,
+    )
+
+    dynamic_runtime_params_slice_provider = (
+        build_runtime_params.DynamicRuntimeParamsSliceProvider.from_config(
+            torax_config
         )
     )
-    restart_case = False
+
+    if torax_config.restart and torax_config.restart.do_restart:
+        initial_state, post_processed_outputs = (
+            initial_state_lib.get_initial_state_and_post_processed_outputs_from_file(
+                t_initial=torax_config.numerics.t_initial,
+                file_restart=torax_config.restart,
+                static_runtime_params_slice=static_runtime_params_slice,
+                dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
+                geometry_provider=geometry_provider,
+                step_fn=step_fn,
+            )
+        )
+        restart_case = True
+    else:
+        initial_state, post_processed_outputs = (
+            initial_state_lib.get_initial_state_and_post_processed_outputs(
+                t=torax_config.numerics.t_initial,
+                static_runtime_params_slice=static_runtime_params_slice,
+                dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
+                geometry_provider=geometry_provider,
+                step_fn=step_fn,
+            )
+        )
+        restart_case = False
 
     return (
         static_runtime_params_slice,
@@ -162,7 +162,7 @@ def run_simulation(
         torax_config=torax_config,
     )
 
-  return (
-      state_history.simulation_output_to_xr(torax_config.restart),
-      state_history,
-  )
+    return (
+        state_history.simulation_output_to_xr(torax_config.restart),
+        state_history,
+    )
diff --git a/torax/imas_tools/equilibrium.py b/torax/imas_tools/equilibrium.py
@@ -144,7 +144,7 @@ def geometry_from_IMAS(
     rhon = IMAS_data.profiles_1d.rho_tor_norm
     vpr = 4 * np.pi * Phi[-1] * rhon / (F * flux_surf_avg_1_over_R2)
     spr = vpr / (2 * np.pi * R_major)
-    Ip_profile_unscaled = scipy.integrate.cumulative_trapezoid(y=spr * jtor, x=rhon, initial=0.0)    # this Ip_profile by integration results in a discrepancy between this term and the total ip from IDSAdd commentMore actions
+    Ip_profile_unscaled = scipy.integrate.cumulative_trapezoid(y=spr * jtor, x=rhon, initial=0.0)    # this Ip_profile by integration results in a discrepancy between this term and the total ip from IDS
 
     # Because of the discrepancy between Ip_profile[-1] (computed by integration) and global_quantities.ip, here we will scale Ip_profile such that the total plasma current is equal
     Ip_total = -1 * IMAS_data.global_quantities.ip
diff --git a/torax/imas_tools/util.py b/torax/imas_tools/util.py
@@ -162,6 +162,7 @@ def update_dict(old_dict:dict, updates:dict) -> dict:
   return new_dict
 
 
+
 # todo check if we can copy form geometry without weird dependency loops
 def face_to_cell(face):
     """Infers cell values corresponding to a vector of face values.
diff --git a/torax/tests/sim_test.py b/torax/tests/sim_test.py
