Add to build_sim

Author: theo-brown

torax/config/build_sim.py

@@ -41,6 +41,7 @@
 from torax.transport_model import constant as constant_transport
 from torax.transport_model import critical_gradient as critical_gradient_transport
 from torax.transport_model import qlknn_transport_model
+from torax.transport_model import tglfnn_transport_model
 # pylint: disable=g-import-not-at-top
 try:
   from torax.transport_model import qualikiz_transport_model
@@ -583,6 +584,16 @@ def build_transport_model_builder_from_config(
             **qualikiz_params,
         )
     )
+  elif transport_model == 'tglfnn':
+    tglfnn_params = dict(transport_config.pop('tglfnn_params', {}))
+    tglfnn_params.update(transport_config)
+    return tglfnn_transport_model.TGLFNNTransportModelBuilder(
+        runtime_params=config_args.recursive_replace(
+            tglfnn_transport_model.RuntimeParams(),
+            **tglfnn_params,
+        )
+    )
+
   # pylint: enable=undefined-variable
   raise ValueError(f'Unknown transport model: {transport_model}')
 

torax/examples/torax_tglfnn.py

@@ -1,7 +1,7 @@
 import json
 from copy import deepcopy
 from pathlib import Path
-
+import dataclasses
 import chex
 import jax.numpy as jnp
 from flax import linen as nn
@@ -12,7 +12,9 @@
 from torax.pedestal_model import pedestal_model as pedestal_model_lib
 from torax.transport_model import tglf_based_transport_model
 from warnings import warn
-from torax.transport_model.tglf_based_transport_model import TGLFInputs
+from torax.transport_model.tglf_based_transport_model import TGLFInputs, RuntimeParams
+from torax.transport_model import transport_model
+from typing import Callable
 
 
 class TGLFNNSurrogate(nn.Module):
@@ -118,9 +120,7 @@ def get_params_from_pytorch_state_dict(self, pytorch_state_dict: dict):
     return params
 
 
-class TGLFSurrogateTransportModel(
-    tglf_based_transport_model.TGLFBasedTransportModel
-):
+class TGLFNNTransportModel(tglf_based_transport_model.TGLFBasedTransportModel):
   """Calculate turbulent transport coefficients using a TGLF surrogate model."""
 
   def __init__(
@@ -239,3 +239,24 @@ def _call_implementation(
         gradient_reference_length=geo.Rmin,  # Device minor radius at LCFS
         gyrobohm_flux_reference_length=geo.Rmin,  # TODO: Check
     )
+
+
+@dataclasses.dataclass(kw_only=True)
+class TGLFNNTransportModelBuilder(transport_model.TransportModelBuilder):
+  """When called, instantiates a TGLFSurrogateTransportModel."""
+
+  runtime_params: RuntimeParams = dataclasses.field(
+      default_factory=RuntimeParams
+  )
+  weights_path: str = (
+      "/home/theo/documents/ukaea/torax/tglfnn/1.0.0/tglfnn_checkpoint.json"
+  )
+  scaling_path: str = "/home/theo/documents/ukaea/torax/tglfnn/1.0.0/stats.json"
+
+  def __call__(
+      self,
+  ) -> TGLFNNTransportModel:
+    return TGLFNNTransportModel(
+        path_to_model_weights_json=(self.weights_path),
+        path_to_model_scaling_json=(self.scaling_path),
+    )