Extend and refactor MEQ mat file loading and add tests.

Updated the MEQ file loading function to handle both "LY" and "L" structured array wrappers commonly found in MEQ outputs. Added unit tests for the loader function.

PiperOrigin-RevId: 848554269

Author: jcitrin

torax/_src/geometry/geometry.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 
 """Classes for representing the problem geometry."""
+
 from collections.abc import Sequence
 import dataclasses
 import enum
+from typing import TypeVar
 
 import chex
 import jax
@@ -385,7 +387,10 @@ def z_magnetic_axis(self) -> chex.Numeric:
       raise ValueError('Geometry does not have a z magnetic axis.')
 
 
-def stack_geometries(geometries: Sequence[Geometry]) -> Geometry:
+GeometryT = TypeVar('GeometryT', bound='Geometry')
+
+
+def stack_geometries(geometries: Sequence[GeometryT]) -> GeometryT:
   """Batch together a sequence of geometries.
 
   Args:

torax/_src/geometry/geometry_loader.py

@@ -59,24 +59,29 @@ def _load_CHEASE_data(  # pylint: disable=invalid-name
 
 
 def _load_fbt_data(file_path: str | IO[bytes]) -> dict[str, np.ndarray]:
-  """Loads data into a dictionary from an FBT-LY file or file path."""
+  """Loads data into a dictionary from an MEQ LY or L file or file path."""
 
   meq_data = scipy.io.loadmat(file_path)
 
-  if "LY" not in meq_data:
-    # L file or LY file not the output of MEQ meqlpack. Return as is.
-    return meq_data
-  else:
-    # LY bundle file. numpy structured array likely resulting from
-    # scipy.io.loadmat returning the output of MEQ meqlpack.m.
-    # Reformat to return a dict. Shapes of 2D arrays are (radius, time)
-    LY_bundle_dict = {}  # pylint: disable=invalid-name
-    LY_bundle_array = meq_data["LY"]  # pylint: disable=invalid-name
-    field_names = LY_bundle_array.dtype.names
-    for name in field_names:
-      LY_bundle_dict[name] = LY_bundle_array[name].item()
-
-    return LY_bundle_dict
+  # Check for nested structured arrays common in MEQ outputs from MEQ meqlpack.
+  # The actual data is often wrapped in a top-level "LY" or "L" key.
+  for key in ("LY", "L"):
+    if key in meq_data:
+      data_array = meq_data[key]
+      # Check if the array is a structured array (has field names).
+      if data_array.dtype.names is None:
+        raise ValueError(f"Provided MEQ {key} data not in expected format.")
+      extracted_dict = {}
+      for name in data_array.dtype.names:
+        # structured_array[name] returns an array; .item() retrieves the
+        # inner array object for that specific field.
+        item = data_array[name].item()
+        if not isinstance(item, np.ndarray):
+          raise ValueError(f"MEQ data field '{name}' is not a numpy array.")
+        extracted_dict[name] = item
+      return extracted_dict
+  # If no nested structured array is found, return the loaded mat dict as is.
+  return meq_data
 
 
 def _load_eqdsk_data(file_path: str, cocos: int) -> dict[str, np.ndarray]:

torax/_src/geometry/tests/geometry_loader_test.py

@@ -0,0 +1,88 @@
+# Copyright 2025 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 the geometry_loader module."""
+
+from unittest import mock
+from absl.testing import absltest
+import numpy as np
+from torax._src.geometry import geometry_loader
+
+# pylint: disable=invalid-name
+
+
+class GeometryLoaderTest(absltest.TestCase):
+
+  @mock.patch("scipy.io.loadmat")
+  def test_load_data_flat(self, mock_loadmat):
+    """Tests loading MEQ data that is already flat in the mat file."""
+    mock_data = {
+        "rBt": np.array([5.0]),
+        "aminor": np.linspace(0.1, 1.0, 10),
+    }
+    mock_loadmat.return_value = mock_data
+
+    result = geometry_loader._load_fbt_data("dummy_path.mat")
+    self.assertEqual(result, mock_data)
+
+  @mock.patch("scipy.io.loadmat")
+  def test_load_fbt_data_nested_LY(self, mock_loadmat):
+    """Tests loading MEQ data nested inside an 'LY' structured array."""
+    # Build a structured array to mimic scipy.io.loadmat's behavior for
+    # meqlpack outputs.
+    aminor_data = np.linspace(0.1, 1.0, 10)
+    dtype = [("rBt", "O"), ("aminor", "O")]
+    # Structured array with 1 element, fields contain arrays.
+    LY_array = np.array([(np.array([5.0]), aminor_data)], dtype=dtype)
+
+    mock_loadmat.return_value = {"LY": LY_array}
+
+    result = geometry_loader._load_fbt_data("dummy_path.mat")
+
+    self.assertIn("rBt", result)
+    self.assertIn("aminor", result)
+    np.testing.assert_array_equal(result["rBt"], np.array([5.0]))
+    np.testing.assert_array_equal(result["aminor"], aminor_data)
+
+  @mock.patch("scipy.io.loadmat")
+  def test_load_fbt_data_nested_L(self, mock_loadmat):
+    """Tests loading MEQ data nested inside an 'L' structured array."""
+    pQ_data = np.linspace(0.0, 1.0, 20)
+    dtype = [("pQ", "O")]
+    L_array = np.array([(pQ_data,)], dtype=dtype)
+
+    mock_loadmat.return_value = {"L": L_array}
+
+    result = geometry_loader._load_fbt_data("dummy_path.mat")
+
+    self.assertIn("pQ", result)
+    np.testing.assert_array_equal(result["pQ"], pQ_data)
+
+  @mock.patch("scipy.io.loadmat")
+  def test_load_fbt_data_invalid_item_type(self, mock_loadmat):
+    """Tests that a ValueError is raised if a nested item is not an array."""
+    dtype = [("bad_field", "O")]
+    # Create a structured array where the item is NOT a numpy array (e.g., int)
+    LY_array = np.array([(123,)], dtype=dtype)
+
+    mock_loadmat.return_value = {"LY": LY_array}
+
+    with self.assertRaisesRegex(
+        ValueError, "MEQ data field 'bad_field' is not a numpy array"
+    ):
+      geometry_loader._load_fbt_data("dummy_path.mat")
+
+
+if __name__ == "__main__":
+  absltest.main()