Merge pull request #228 from boutproject/fci-geometry-new

FCI geometry: basic test

xbout/geometries.py

@@ -4,7 +4,7 @@
 import xarray as xr
 import numpy as np
 
-from .region import Region, _create_regions_toroidal
+from .region import Region, _create_regions_toroidal, _create_single_region
 from .utils import (
     _add_attrs_to_var,
     _set_attrs_on_all_vars,
@@ -184,7 +184,7 @@ def apply_geometry(ds, geometry_name, *, coordinates=None, grid=None):
         # In BOUT++ v5, dz is either a Field2D or Field3D.
         # We can use it as a 1D coordinate if it's a Field3D, _or_ if nz == 1
         bout_v5 = updated_ds.metadata["BOUT_VERSION"] > 5.0 or (
-            updated_ds.metadata["BOUT_VERSION"] == 5.0 and updated_ds["dz"].ndim == 2
+            updated_ds.metadata["BOUT_VERSION"] == 5.0 and updated_ds["dz"].ndim >= 2
         )
         use_metric_3d = updated_ds.metadata.get("use_metric_3d", False)
         can_use_1d_z_coord = (nz == 1) or use_metric_3d
@@ -197,14 +197,20 @@ def apply_geometry(ds, geometry_name, *, coordinates=None, grid=None):
                     raise ValueError(
                         f"Spacing is not constant. Cannot create z coordinate"
                     )
-                dz = updated_ds["dz"][0, 0]
+
+                dz = updated_ds["dz"].min()
             else:
                 dz = updated_ds["dz"]
 
             z0 = 2 * np.pi * updated_ds.metadata["ZMIN"]
             z1 = z0 + nz * dz
-            if not np.isclose(
-                z1, 2.0 * np.pi * updated_ds.metadata["ZMAX"], rtol=1.0e-15, atol=0.0
+            if not np.all(
+                np.isclose(
+                    z1,
+                    2.0 * np.pi * updated_ds.metadata["ZMAX"],
+                    rtol=1.0e-15,
+                    atol=0.0,
+                )
             ):
                 warn(
                     f"Size of toroidal domain as calculated from nz*dz ({str(z1 - z0)}"

xbout/region.py

@@ -1239,8 +1239,8 @@ def _create_single_region(ds, periodic_y=True):
         "all": Region(
             name="all",
             ds=ds,
-            xouter_ind=0,
-            xinner_ind=nx,
+            xouter_ind=nx,
+            xinner_ind=0,
             ylower_ind=0,
             yupper_ind=ny,
             connection_lower_y=connection,

xbout/tests/test_load.py

@@ -3,6 +3,8 @@
 import inspect
 from pathlib import Path
 import re
+from functools import reduce
+import operator
 
 import pytest
 
@@ -50,6 +52,15 @@ def test_check_extensions(tmp_path):
         filetype = _check_filetype(example_invalid_file)
 
 
+def test_set_fci_coords(create_example_grid_file_fci, create_example_files_fci):
+    grid = create_example_grid_file_fci
+    data = create_example_files_fci
+
+    ds = open_boutdataset(data, gridfilepath=grid, geometry="fci")
+    assert "R" in ds
+    assert "Z" in ds
+
+
 class TestPathHandling:
     def test_glob_expansion_single(self, tmp_path):
         files_dir = tmp_path.joinpath("data")
@@ -1435,3 +1446,54 @@ def test_trim_timing_info(self):
 
         expected = create_test_data(0)
         xrt.assert_equal(ds, expected)
+
+
+fci_shape = (2, 2, 3, 4)
+fci_guards = (2, 2, 0)
+
+
+@pytest.fixture
+def create_example_grid_file_fci(tmp_path_factory):
+    """
+    Mocks up a FCI-like netCDF file, and return the temporary test
+    directory containing them.
+
+    Deletes the temporary directory once that test is done.
+    """
+
+    # Create grid dataset
+    shape = (fci_shape[1] + 2 * fci_guards[0], *fci_shape[2:])
+    arr = np.arange(reduce(operator.mul, shape, 1)).reshape(shape)
+    grid = DataArray(data=arr, name="R", dims=["x", "y", "z"]).to_dataset()
+    grid["Z"] = DataArray(np.random.random(shape), dims=["x", "y", "z"])
+    grid["dy"] = DataArray(np.ones(shape), dims=["x", "y", "z"])
+    grid = grid.set_coords(["dy"])
+
+    # Create temporary directory
+    save_dir = tmp_path_factory.mktemp("griddata")
+
+    # Save
+    filepath = save_dir.joinpath("fci.nc")
+    grid.to_netcdf(filepath, engine="netcdf4")
+
+    return filepath
+
+
+@pytest.fixture
+def create_example_files_fci(tmp_path_factory):
+
+    return _bout_xyt_example_files(
+        tmp_path_factory,
+        lengths=fci_shape,
+        nxpe=1,
+        nype=1,
+        # nt=1,
+        guards={a: b for a, b in zip("xyz", fci_guards)},
+        syn_data_type="random",
+        grid=None,
+        squashed=False,
+        # topology="core",
+        write_to_disk=False,
+        bout_v5=True,
+        metric_3D=True,
+    )