PyroScan: how fields, eigenfunctions and time are reduced when loading a scan

.gitignore

@@ -97,6 +97,11 @@ ipython_config.py
 #   install all needed dependencies.
 #Pipfile.lock
 
+# uv
+#  for similar reasons for pipenv uv.lock my be ignored
+
+uv.lock
+
 # PEP 582; used by e.g. github.com/David-OConnor/pyflow
 __pypackages__/
 

CHANGELOG.md

@@ -10,9 +10,23 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ### Added
 
 ### Changed
+  - `add_flags` has one implementation per input format: `GKInput.add_flags` for grouped inputs and the new `GKInputFlat.add_flags` for flat `KEY = value` inputs. Per-code overrides are removed and `GKInput.add_flags` is no longer abstract
   - Changed Pyro kwarg from `gk_type` to  `gk_code`
-
+  - `PyroScan.load_gk_output` uses one time setting for linear and nonlinear
+    runs: `time_mode` (`"average"` default, `"last"`, `"trace"`), with
+    `tolerance_time_range` (default 0.8) setting the averaging window. With
+    `"average"`, fields and eigenfunctions are `|field|**2` averaged in time,
+    stored as `phi_squared`, `apar_squared`, `bpar_squared` and
+    `eigenfunctions_squared`, since a time mean of a complex field is
+    meaningless. Linear growth rates and fluxes are now averaged rather than
+    taken at the last time.
+  - `phi`, `apar`, `bpar` and `eigenfunctions` are no longer reduced to `ky[0]`
+    and the smallest `|kx|`; they keep `kx` and `ky` unless summed with
+    `sum_kx` / `sum_ky`. `sum_ky` now defaults to False, for fluxes too.
+
 ### Fixed
+  - `add_flags` for TGLF, CGYRO and NEO matches keys case-insensitively, so e.g. `NBASIS_MAX` overwrites the existing TGLF `nbasis_max` instead of writing it twice. New keys take the code's stored case (`flag_key_case`: lowercase for TGLF, uppercase for CGYRO/NEO)
+  - `add_flags` raises `TypeError` when flags do not match the input structure (a dict of groups for namelist/TOML codes, a flat dict for TGLF/CGYRO/NEO) instead of crashing or writing an invalid file
 
 ## [0.0.1] - 2021-01-26  
 

src/pyrokinetics/diagnostics/saturation_rules.py

@@ -71,7 +71,18 @@ def mg_saturation(
         shat = pyro.local_geometry.shat
 
         theta = data["theta"].data
-        eigenfunctions = data["eigenfunctions"]
+        # The rule needs |eigenfunction|**2: loaded with time_mode="average" it
+        # is already stored, averaged in time; otherwise take the final time
+        if "eigenfunctions_squared" in data:
+            eigenfunctions_squared = data["eigenfunctions_squared"]
+        else:
+            eigenfunctions_squared = np.abs(data["eigenfunctions"]) ** 2
+            if "time" in eigenfunctions_squared.dims:
+                eigenfunctions_squared = eigenfunctions_squared.isel(time=-1, drop=True)
+        # The scan keeps kx; the rule uses one kx
+        if "kx" in eigenfunctions_squared.dims:
+            kx_min = abs(eigenfunctions_squared["kx"]).argmin().item()
+            eigenfunctions_squared = eigenfunctions_squared.isel(kx=kx_min, drop=True)
         growth_rate_tolerance = data["growth_rate_tolerance"]
 
         growth_rate = data["growth_rate"].where(
@@ -92,9 +103,8 @@ def mg_saturation(
             / heat_tot
         )
 
-        field_squared = (
-            np.abs(eigenfunctions.where(growth_rate_tolerance < gamma_tolerance, 0.0))
-            ** 2
+        field_squared = eigenfunctions_squared.where(
+            growth_rate_tolerance < gamma_tolerance, 0.0
         )
 
         # Set up Jacobian and k_perp

src/pyrokinetics/gk_code/__init__.py

@@ -3,6 +3,7 @@
 # Import built-in input and output readers to register them with GkInput and GkOutput.
 from .cgyro import GKInputCGYRO, GKOutputReaderCGYRO  # noqa
 from .gene import GKInputGENE, GKOutputReaderGENE  # noqa
+from .gftm import GKInputGFTM, GKOutputReaderGFTM  # noqa
 from .gk_input import GKInput, read_gk_input, supported_gk_input_types
 from .gk_output import GKOutput, read_gk_output, supported_gk_output_types
 from .gkw import GKInputGKW, GKOutputReaderGKW  # noqa

src/pyrokinetics/gk_code/cgyro.py

@@ -25,7 +25,7 @@
 from ..numerics import Numerics
 from ..templates import gk_templates
 from ..typing import PathLike
-from .gk_input import GKInput
+from .gk_input import GKInput, GKInputFlat
 from .gk_output import (
     Coords,
     Eigenfunctions,
@@ -37,13 +37,14 @@
 )
 
 
-class GKInputCGYRO(GKInput, FileReader, file_type="CGYRO", reads=GKInput):
+class GKInputCGYRO(GKInputFlat, FileReader, file_type="CGYRO", reads=GKInput):
     """
     Class that can read CGYRO input files, and produce
     Numerics, LocalSpecies, and LocalGeometry objects
     """
 
     code_name = "CGYRO"
+    flag_key_case = staticmethod(str.upper)
     default_file_name = "input.cgyro"
     norm_convention = "cgyro"
     _convention_dict = {}
@@ -265,13 +266,6 @@ def write(
     def is_nonlinear(self) -> bool:
         return bool(self.data.get("NONLINEAR_FLAG", 0))
 
-    def add_flags(self, flags) -> None:
-        """
-        Add extra flags to CGYRO input file
-        """
-        for key, value in flags.items():
-            self.data[key] = value
-
     def get_local_geometry(self) -> LocalGeometry:
         """
         Returns local geometry. Delegates to more specific functions

src/pyrokinetics/gk_code/gene.py

@@ -304,13 +304,6 @@ def write(
     def is_nonlinear(self) -> bool:
         return bool(self.data["general"].get("nonlinear", False))
 
-    def add_flags(self, flags) -> None:
-        """
-        Add extra flags to GENE input file
-        Uses default, which assumes a Fortan90 namelist
-        """
-        super().add_flags(flags)
-
     def get_local_geometry(self) -> LocalGeometry:
         """
         Returns local geometry. Delegates to more specific functions