diff --git a/xbout/geometries.py b/xbout/geometries.py index 035d5592..375e677f 100644 --- a/xbout/geometries.py +++ b/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)}" diff --git a/xbout/region.py b/xbout/region.py index ce401d2a..12391b12 100644 --- a/xbout/region.py +++ b/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, diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 21207d0a..7d9608f3 100644 --- a/xbout/tests/test_load.py +++ b/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, + )