Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix loading of grid files #307

Open
wants to merge 8 commits into
base: master
Choose a base branch
from
2 changes: 1 addition & 1 deletion examples/plot_grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,6 @@
grid["psi_poloidal"].bout.pcolormesh()
grid["psi_poloidal"].bout.pcolormesh(shading="gouraud")

grid["psi_poloidal"].bout.regions()
grid["psi_poloidal"].bout.plot_regions()

plt.show()
2 changes: 1 addition & 1 deletion xbout/geometries.py
Original file line number Diff line number Diff line change
Expand Up @@ -401,7 +401,7 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None):
ds = _maybe_rename_dimension(ds, "y", coordinates["y"])

# TODO automatically make this coordinate 1D in simplified cases?
ds = ds.rename(psixy=coordinates["x"])
ds[coordinates["x"]] = ds["psixy"]
ds = ds.set_coords(coordinates["x"])
ds[coordinates["x"]].attrs["units"] = "Wb"

Expand Down
13 changes: 9 additions & 4 deletions xbout/load.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ def open_boutdataset(
run_name=None,
info=True,
is_restart=None,
is_mms_dump=False,
**kwargs,
):
"""Load a dataset from a set of BOUT output files, including the
Expand Down Expand Up @@ -179,12 +180,12 @@ def open_boutdataset(
chunks = {}

input_type = _check_dataset_type(datapath)

if is_restart is None:
is_restart = input_type == "restart"
elif is_restart is True:
input_type = "restart"

if is_mms_dump:
input_type = "dump"
if "reload" in input_type:
if input_type == "reload":
if isinstance(datapath, Path):
Expand Down Expand Up @@ -358,7 +359,12 @@ def attrs_remove_section(obj, section):
ds.metadata[v] = grid[v].values

# Update coordinates to match particular geometry of grid
ds = geometries.apply_geometry(ds, geometry, grid=grid)
if input_type == "grid":
# Specify coordinate names to avoid name clash
coordinates = {"x": "psi_poloidal", "y": "y", "z": "zeta"}
else:
coordinates = None
ds = geometries.apply_geometry(ds, geometry, grid=grid, coordinates=coordinates)

if remove_yboundaries:
ds = ds.bout.remove_yboundaries()
Expand Down Expand Up @@ -616,7 +622,6 @@ def _auto_open_mfboutdataset(
nxpe, nype, mxg, myg, mxsub, mysub, is_squashed_doublenull = _read_splitting(
filepaths[0], info, keep_yboundaries
)

if is_squashed_doublenull:
# Need to remove y-boundaries after loading: (i) in case
# we are loading a squashed data-set, in which case we
Expand Down
1 change: 1 addition & 0 deletions xbout/tests/test_boutdataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -1136,6 +1136,7 @@ def test_interpolate_parallel_all_variables_arg(self, bout_xyt_example_files):
"G3",
"J",
"Bxy",
"psixy",
)
)

Expand Down
Loading