diff --git a/.github/workflows/pythonpackage.yml b/.github/workflows/pythonpackage.yml index c6e81bf7..e3b89c87 100644 --- a/.github/workflows/pythonpackage.yml +++ b/.github/workflows/pythonpackage.yml @@ -55,7 +55,7 @@ jobs: matrix: python-version: [3.7, 3.8] pip-packages: - - "setuptools pip pytest pytest-cov coverage codecov boutdata==0.1.4 xarray==0.18.0 dask==2.10.0 numpy==1.18.0 natsort==5.5.0 matplotlib==3.1.1 animatplot==0.4.2 netcdf4==1.4.2 Pillow==6.1.0" # test with oldest supported version of packages. Note, using numpy==1.18.0 as a workaround because numpy==1.17.0 is not supported on Python-3.7, even though we should currently support numpy==1.17.0. + - "setuptools pip pytest pytest-cov coverage codecov boutdata==0.1.4 xarray==0.18.0 dask==2.10.0 numpy==1.18.0 natsort==5.5.0 matplotlib==3.1.1 animatplot==0.4.2 netcdf4==1.4.2 Pillow==7.2.0" # test with oldest supported version of packages. Note, using numpy==1.18.0 as a workaround because numpy==1.17.0 is not supported on Python-3.7, even though we should currently support numpy==1.17.0. fail-fast: false steps: diff --git a/.github/workflows/pythonpublish.yml b/.github/workflows/pythonpublish.yml index 4fa26fe6..2979cc61 100644 --- a/.github/workflows/pythonpublish.yml +++ b/.github/workflows/pythonpublish.yml @@ -49,7 +49,7 @@ jobs: matrix: python-version: [3.7, 3.8] pip-packages: - - "setuptools pip pytest pytest-cov coverage codecov boutdata==0.1.4 xarray==0.18.0 dask==2.10.0 numpy==1.18.0 natsort==5.5.0 matplotlib==3.1.1 animatplot==0.4.2 netcdf4==1.4.2 Pillow==6.1.0" # test with oldest supported version of packages. Note, using numpy==1.18.0 as a workaround because numpy==1.17.0 is not supported on Python-3.7, even though we should currently support numpy==1.17.0. + - "setuptools pip pytest pytest-cov coverage codecov boutdata==0.1.4 xarray==0.18.0 dask==2.10.0 numpy==1.18.0 natsort==5.5.0 matplotlib==3.1.1 animatplot==0.4.2 netcdf4==1.4.2 Pillow==7.2.0" # test with oldest supported version of packages. Note, using numpy==1.18.0 as a workaround because numpy==1.17.0 is not supported on Python-3.7, even though we should currently support numpy==1.17.0. fail-fast: true steps: diff --git a/README.md b/README.md index ac1b7b10..ec283403 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # xBOUT [![Build Status](https://github.com/boutproject/xBOUT/workflows/master/badge.svg)](https://github.com/boutproject/xBOUT/actions) -[![codecov](https://codecov.io/gh/boutproject/xBOUT/branch/master/graph/badge.svg)](https://codecov.io/gh/boutproject/xBOUT) +[![codecov](https://codecov.io/gh/boutproject/xBOUT/branch/master/graph/badge.svg)](https://codecov.io/gh/boutproject/xBOUT) [![Documentation Status](https://readthedocs.org/projects/xbout/badge/?version=latest)](https://xbout.readthedocs.io/en/latest/?badge=latest) [![DOI](https://zenodo.org/badge/160846663.svg)](https://zenodo.org/badge/latestdoi/160846663) Documentation: https://xbout.readthedocs.io diff --git a/xbout/boutdataarray.py b/xbout/boutdataarray.py index 2ee9f134..cc09c3f1 100644 --- a/xbout/boutdataarray.py +++ b/xbout/boutdataarray.py @@ -17,6 +17,7 @@ from .plotting.utils import _create_norm from .region import _from_region from .utils import ( + _add_cartesian_coordinates, _make_1d_xcoord, _update_metadata_increased_x_resolution, _update_metadata_increased_y_resolution, @@ -823,6 +824,17 @@ def interpolate_to_new_grid( result = apply_geometry(result, self.data.geometry) return result[self.data.name] + def add_cartesian_coordinates(self): + """ + Add Cartesian (X,Y,Z) coordinates. + + Returns + ------- + DataArray with new coordinates added, which are named 'X_cartesian', + 'Y_cartesian', and 'Z_cartesian' + """ + return _add_cartesian_coordinates(self.data) + def remove_yboundaries(self, return_dataset=False, remove_extra_upper=False): """ Remove y-boundary points, if present, from the DataArray @@ -1443,6 +1455,42 @@ def interpolate_from_unstructured( return result + def interpolate_to_cartesian(self, *args, **kwargs): + """ + Interpolate the DataArray to a regular Cartesian grid. + + This method is intended to be used to produce data for visualisation, which + normally does not require double-precision values, so by default the data is + converted to `np.float32`. Pass `use_float32=False` to retain the original + precision. + + Parameters + ---------- + nX : int (default 300) + Number of grid points in the X direction + nY : int (default 300) + Number of grid points in the Y direction + nZ : int (default 100) + Number of grid points in the Z direction + use_float32 : bool (default True) + Downgrade precision to `np.float32`? + fill_value : float (default np.nan) + Value to use for points outside the interpolation domain (passed to + `scipy.RegularGridInterpolator`) + + See Also + -------- + BoutDataset.interpolate_to_cartesian + """ + da = self.data + name = da.name + ds = da.to_dataset() + # Dataset needs geometry and metadata attributes, but these are not copied from + # the DataArray by default + ds.attrs["geometry"] = da.geometry + ds.attrs["metadata"] = da.metadata + return ds.bout.interpolate_to_cartesian(*args, **kwargs)[name] + # BOUT-specific plotting functionality: methods that plot on a poloidal (R-Z) plane def contour(self, ax=None, **kwargs): """ diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index c128e061..8be1518f 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -28,7 +28,11 @@ _parse_coord_option, ) from .region import _from_region -from .utils import _get_bounding_surfaces, _split_into_restarts +from .utils import ( + _add_cartesian_coordinates, + _get_bounding_surfaces, + _split_into_restarts, +) @xr.register_dataset_accessor("bout") @@ -746,6 +750,171 @@ def interpolate_from_unstructured( return ds + def interpolate_to_cartesian( + self, nX=300, nY=300, nZ=100, *, use_float32=True, fill_value=np.nan + ): + """ + Interpolate the Dataset to a regular Cartesian grid. + + This method is intended to be used to produce data for visualisation, which + normally does not require double-precision values, so by default the data is + converted to `np.float32`. Pass `use_float32=False` to retain the original + precision. + + Parameters + ---------- + nX : int (default 300) + Number of grid points in the X direction + nY : int (default 300) + Number of grid points in the Y direction + nZ : int (default 100) + Number of grid points in the Z direction + use_float32 : bool (default True) + Downgrade precision to `np.float32`? + fill_value : float (default np.nan) + Value to use for points outside the interpolation domain (passed to + `scipy.RegularGridInterpolator`) + + See Also + -------- + BoutDataArray.interpolate_to_cartesian + """ + ds = self.data + ds = ds.bout.add_cartesian_coordinates() + + if not isinstance(use_float32, bool): + raise ValueError(f"use_float32 must be a bool, got '{use_float32}'") + if use_float32: + float_type = np.float32 + ds = ds.astype(float_type) + for coord in ds.coords: + # Coordinates are not converted by Dataset.astype, so convert explicitly + ds[coord] = ds[coord].astype(float_type) + fill_value = float_type(fill_value) + else: + float_type = ds[ds.data_vars[0]].dtype + + tdim = ds.metadata["bout_tdim"] + zdim = ds.metadata["bout_zdim"] + if tdim in ds.dims: + nt = ds.sizes[tdim] + n_toroidal = ds.sizes[zdim] + + # Create Cartesian grid to interpolate to + Xmin = ds["X_cartesian"].min() + Xmax = ds["X_cartesian"].max() + Ymin = ds["Y_cartesian"].min() + Ymax = ds["Y_cartesian"].max() + Zmin = ds["Z_cartesian"].min() + Zmax = ds["Z_cartesian"].max() + newX_1d = xr.DataArray(np.linspace(Xmin, Xmax, nX), dims="X") + newX = newX_1d.expand_dims({"Y": nY, "Z": nZ}, axis=[1, 2]) + newY_1d = xr.DataArray(np.linspace(Ymin, Ymax, nY), dims="Y") + newY = newY_1d.expand_dims({"X": nX, "Z": nZ}, axis=[0, 2]) + newZ_1d = xr.DataArray(np.linspace(Zmin, Zmax, nZ), dims="Z") + newZ = newZ_1d.expand_dims({"X": nX, "Y": nY}, axis=[0, 1]) + newR = np.sqrt(newX**2 + newY**2) + newzeta = np.arctan2(newY, newX) + # Define newzeta in range 0->2*pi + newzeta = np.where(newzeta < 0.0, newzeta + 2.0 * np.pi, newzeta) + + from scipy.interpolate import ( + RegularGridInterpolator, + griddata, + ) + + # Create Cylindrical coordinates for intermediate grid + Rcyl_min = float_type(ds["R"].min()) + Rcyl_max = float_type(ds["R"].max()) + Zcyl_min = float_type(ds["Z"].min()) + Zcyl_max = float_type(ds["Z"].max()) + n_Rcyl = int(round(nZ * (Rcyl_max - Rcyl_min) / (Zcyl_max - Zcyl_min))) + Rcyl = xr.DataArray(np.linspace(Rcyl_min, Rcyl_max, 2 * n_Rcyl), dims="r") + Zcyl = xr.DataArray(np.linspace(Zcyl_min, Zcyl_max, 2 * nZ), dims="z") + + # Create Dataset for result + result = xr.Dataset() + result.attrs["metadata"] = ds.metadata + + # Interpolate in two stages for efficiency. Unstructured 3d interpolation is + # very slow. Unstructured 2d interpolation onto Cartesian (R, Z) grids, followed + # by structured 3d interpolation onto the (X, Y, Z) grid, is much faster. + # Structured 3d interpolation straight from (psi, theta, zeta) to (X, Y, Z) + # leaves artifacts in the output, because theta does not vary continuously + # everywhere (has branch cuts). + + zeta_out = np.zeros(n_toroidal + 1) + zeta_out[:-1] = ds[zdim].values + zeta_out[-1] = zeta_out[-2] + ds["dz"].mean() + + def interp_single_time(da): + print(" interpolate poloidal planes") + + da_cyl = da.bout.interpolate_from_unstructured(R=Rcyl, Z=Zcyl).transpose( + "R", "Z", zdim, missing_dims="ignore" + ) + + if zdim not in da_cyl.dims: + da_cyl = da_cyl.expand_dims({zdim: n_toroidal + 1}, axis=-1) + else: + # Impose toroidal periodicity by appending zdim=0 to end of array + da_cyl = xr.concat((da_cyl, da_cyl.isel({zdim: 0})), zdim) + + print(" build 3d interpolator") + interp = RegularGridInterpolator( + (Rcyl.values, Zcyl.values, zeta_out), + da_cyl.values, + bounds_error=False, + fill_value=fill_value, + ) + + print(" do 3d interpolation") + return interp( + (newR, newZ, newzeta), + method="linear", + ) + + for name, da in ds.data_vars.items(): + print(f"\ninterpolating {name}") + # order of dimensions does not really matter here - output only depends on + # shape of newR, newZ, newzeta. Possibly more efficient to assign the 2d + # results in the loop to the last two dimensions, so put zeta first. Can't + # just use da.min().item() here (to get a scalar value instead of a + # zero-size array) because .item() doesn't work for dask arrays (yet!). + + datamin = float_type(da.min().values) + datamax = float_type(da.max().values) + + if tdim in da.dims: + data_cartesian = np.zeros((nt, nX, nY, nZ), dtype=float_type) + for tind in range(nt): + print(f" tind={tind}") + data_cartesian[tind, :, :, :] = interp_single_time( + da.isel({tdim: tind}) + ) + result[name] = xr.DataArray(data_cartesian, dims=[tdim, "X", "Y", "Z"]) + else: + data_cartesian = interp_single_time(da) + result[name] = xr.DataArray(data_cartesian, dims=["X", "Y", "Z"]) + + # Copy metadata to data variables, in case it is needed + result[name].attrs["metadata"] = ds.metadata + + result = result.assign_coords(X=newX_1d, Y=newY_1d, Z=newZ_1d) + + return result + + def add_cartesian_coordinates(self): + """ + Add Cartesian (X,Y,Z) coordinates. + + Returns + ------- + Dataset with new coordinates added, which are named 'X_cartesian', + 'Y_cartesian', and 'Z_cartesian' + """ + return _add_cartesian_coordinates(self.data) + def remove_yboundaries(self, **kwargs): """ Remove y-boundary points, if present, from the Dataset @@ -1006,7 +1175,13 @@ def to_restart( Number of processors in the y-direction. If not given, keep the number used for the original simulation tind : int, default -1 - Time-index of the slice to write to the restart files + Time-index of the slice to write to the restart files. Note, when creating + restart files from 'dump' files it is recommended to open the Dataset using + the full time range and use the `tind` argument here, rather than selecting + a time point manually, so that the calculation of `hist_hi` in the output + can be correct (which requires knowing the existing value of `hist_hi` + (output step count at the end of the simulation), `tind` and the total + number of time points in the current output data). prefix : str, default "BOUT.restart" Prefix to use for names of restart files overwrite : bool, default False @@ -1266,6 +1441,8 @@ def is_list(variable): animate=False, axis_coords=this_axis_coords, aspect=this_aspect, + vmin=this_vmin, + vmax=this_vmax, **this_kwargs, ) ) @@ -1279,6 +1456,8 @@ def is_list(variable): animate=False, axis_coords=this_axis_coords, aspect=this_aspect, + vmin=this_vmin, + vmax=this_vmax, label=w.name, **this_kwargs, ) diff --git a/xbout/calc/tests/test_turbulence.py b/xbout/calc/tests/test_turbulence.py index bc3d7557..a4e105da 100644 --- a/xbout/calc/tests/test_turbulence.py +++ b/xbout/calc/tests/test_turbulence.py @@ -19,7 +19,7 @@ def test_1d(self): dat = np.array([5, 7, 3.2, -1, -4.4]) orig = DataArray(dat, dims=["x"]) - sum_squares = np.sum(dat ** 2) + sum_squares = np.sum(dat**2) mean_squares = sum_squares / dat.size rootmeansquare = np.sqrt(mean_squares) @@ -31,7 +31,7 @@ def test_1d(self): def test_reduce_2d(self, dim, axis): dat = np.array([[5, 7, 3.2, -1, -4.4], [-1, -2.5, 0, 8, 3.0]]) orig = DataArray(dat, dims=["x", "t"]) - sum_squares = np.sum(dat ** 2, axis=axis) + sum_squares = np.sum(dat**2, axis=axis) mean_squares = sum_squares / dat.shape[axis] rootmeansquare = np.sqrt(mean_squares) @@ -44,7 +44,7 @@ def test_reduce_2d_dask(self): orig = DataArray(dat, dims=["x", "t"]) chunked = orig.chunk({"x": 1}) axis = 1 - sum_squares = np.sum(dat ** 2, axis=axis) + sum_squares = np.sum(dat**2, axis=axis) mean_squares = sum_squares / dat.shape[axis] rootmeansquare = np.sqrt(mean_squares) diff --git a/xbout/geometries.py b/xbout/geometries.py index 1a93e38c..78dd5ef2 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -181,7 +181,14 @@ def apply_geometry(ds, geometry_name, *, coordinates=None, grid=None): can_use_1d_z_coord = (nz == 1) or use_metric_3d if can_use_1d_z_coord: - z = _1d_coord_from_spacing(updated_ds["dz"], zcoord, updated_ds) + if updated_ds.geometry == "fci": + # dz is varying. just set to a linspace + z = xr.DataArray( + np.linspace(start=0, stop=2 * np.pi, num=nz, endpoint=False), + dims=zcoord, + ) + else: + z = _1d_coord_from_spacing(updated_ds["dz"], zcoord, updated_ds) else: if bout_v5: if not np.all(updated_ds["dz"].min() == updated_ds["dz"].max()): diff --git a/xbout/load.py b/xbout/load.py index e4ea57bb..f2aaca31 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -73,11 +73,11 @@ def open_boutdataset( keyword argument `drop_vars` to ignore the variables with conflicts, e.g. if `"S1"` and `"S2"` have conflicts ``` - ds = open_boutdataset("data*/boutdata.nc", drop_vars=["S1", "S2"]) + ds = open_boutdataset("data*/boutdata.nc", drop_variables=["S1", "S2"]) ``` will open a Dataset which is missing `"S1"` and `"S2"`.\ - [`drop_vars` is an argument of `xarray.open_dataset()` that is passed down through - `kwargs`.] + [`drop_variables` is an argument of `xarray.open_dataset()` that is passed down + through `kwargs`.] Parameters ---------- @@ -172,7 +172,10 @@ def attrs_to_dict(obj, section): sectionlength = len(section) for key in list(obj.attrs): if key[:sectionlength] == section: - result[key[sectionlength:]] = obj.attrs.pop(key) + val = obj.attrs.pop(key) + if isinstance(val, bytes): + val = val.decode() + result[key[sectionlength:]] = val return result def attrs_remove_section(obj, section): @@ -187,6 +190,10 @@ def attrs_remove_section(obj, section): # Restore metadata from attrs metadata = attrs_to_dict(ds, "metadata") + if "is_restart" not in metadata: + # Loading data that was saved with a version of xbout from before + # "is_restart" was added, so need to add it to the metadata. + metadata["is_restart"] = int(is_restart) ds.attrs["metadata"] = metadata # Must do this for all variables and coordinates in dataset too for da in chain(ds.data_vars.values(), ds.coords.values()): @@ -251,6 +258,15 @@ def attrs_remove_section(obj, section): else: raise ValueError(f"internal error: unexpected input_type={input_type}") + if not is_restart: + for var in _BOUT_TIME_DEPENDENT_META_VARS: + if var in ds: + # Assume different processors in x & y have same iteration etc. + latest_top_left = {dim: 0 for dim in ds[var].dims} + if "t" in ds[var].dims: + latest_top_left["t"] = -1 + ds[var] = ds[var].isel(latest_top_left).squeeze(drop=True) + ds, metadata = _separate_metadata(ds) # Store as ints because netCDF doesn't support bools, so we can't save # bool attributes @@ -264,15 +280,6 @@ def attrs_remove_section(obj, section): # grid file, as they will be removed from the full Dataset below keep_yboundaries = True - if not is_restart: - for var in _BOUT_TIME_DEPENDENT_META_VARS: - if var in ds: - # Assume different processors in x & y have same iteration etc. - latest_top_left = {dim: 0 for dim in ds[var].dims} - if "t" in ds[var].dims: - latest_top_left["t"] = -1 - ds[var] = ds[var].isel(latest_top_left).squeeze(drop=True) - ds = _add_options(ds, inputfilepath) if geometry is None: @@ -522,7 +529,7 @@ def _auto_open_mfboutdataset( # Open just one file to read processor splitting nxpe, nype, mxg, myg, mxsub, mysub, is_squashed_doublenull = _read_splitting( - filepaths[0], info + filepaths[0], info, keep_yboundaries ) if is_squashed_doublenull: @@ -655,7 +662,7 @@ def _expand_wildcards(path): return natsorted(filepaths, key=lambda filepath: str(filepath)) -def _read_splitting(filepath, info=True): +def _read_splitting(filepath, info, keep_yboundaries): ds = xr.open_dataset(str(filepath)) # Account for case of no parallelisation, when nxpe etc won't be in dataset @@ -719,6 +726,23 @@ def get_nonnegative_scalar(ds, key, default=1, info=True): nxpe = 1 nype = 1 is_squashed_doublenull = (ds["jyseps2_1"] != ds["jyseps1_2"]).values + elif ny_file == ny + 2 * myg: + # Older squashed file from double-null grid but containing only lower + # target boundary cells. + if keep_yboundaries: + raise ValueError( + "Cannot keep y-boundary points: squashed file is missing upper " + "target boundary points." + ) + has_yboundaries = not (ny_file == ny) + if not has_yboundaries: + myg = 0 + + nxpe = 1 + nype = 1 + # For this case, do not need the special handling enabled by + # is_squashed_doublenull=True, as keeping y-boundaries is not allowed + is_squashed_doublenull = False # Avoid trying to open this file twice ds.close() @@ -813,10 +837,6 @@ def _trim(ds, *, guards, keep_boundaries, nxpe, nype, is_restart): trimmed_ds = trimmed_ds.drop_vars(name) to_drop = _BOUT_PER_PROC_VARIABLES - if not is_restart: - # These variables are required to be consistent when loading restart files, so - # that they can be written out again in to_restart() - to_drop = to_drop + _BOUT_PER_PROC_VARIABLES_REQUIRED_FROM_RESTARTS return trimmed_ds.drop_vars(to_drop, errors="ignore") diff --git a/xbout/plotting/animate.py b/xbout/plotting/animate.py index e8135282..862aa266 100644 --- a/xbout/plotting/animate.py +++ b/xbout/plotting/animate.py @@ -81,7 +81,7 @@ def _normalise_time_coord(time_values): tmax = time_values.max() if tmax < 1.0e-2 or tmax > 1.0e6: scale_pow = int(np.floor(np.log10(tmax))) - scale_factor = 10 ** scale_pow + scale_factor = 10**scale_pow time_values = time_values / scale_factor suffix = f"e{scale_pow}" else: diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index a7fa7eda..0c09226e 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -64,6 +64,8 @@ def plot_separatrix(da, sep_pos, ax, radial_coord="x"): def _decompose_regions(da): + if da.geometry == "fci": + return {region: da for region in da.bout._regions} return { region: da.bout.from_region(region, with_guards=1) for region in da.bout._regions diff --git a/xbout/tests/test_boutdataarray.py b/xbout/tests/test_boutdataarray.py index fcb246fe..01f6c138 100644 --- a/xbout/tests/test_boutdataarray.py +++ b/xbout/tests/test_boutdataarray.py @@ -602,7 +602,7 @@ def test_interpolate_parallel_region_core(self, bout_xyt_example_files): def f(t): t = np.sin(t) - return t ** 3 - t ** 2 + t - 1.0 + return t**3 - t**2 + t - 1.0 n.data = f(theta).broadcast_like(n) @@ -664,7 +664,7 @@ def test_interpolate_parallel_region_core_change_n( def f(t): t = np.sin(t) - return t ** 3 - t ** 2 + t - 1.0 + return t**3 - t**2 + t - 1.0 n.data = f(theta).broadcast_like(n) @@ -712,7 +712,7 @@ def test_interpolate_parallel_region_sol(self, bout_xyt_example_files): def f(t): t = np.sin(t) - return t ** 3 - t ** 2 + t - 1.0 + return t**3 - t**2 + t - 1.0 n.data = f(theta).broadcast_like(n) @@ -761,7 +761,7 @@ def test_interpolate_parallel_region_singlenull(self, bout_xyt_example_files): def f(t): t = np.sin(3.0 * t) - return t ** 3 - t ** 2 + t - 1.0 + return t**3 - t**2 + t - 1.0 n.data = f(theta).broadcast_like(n) @@ -834,7 +834,7 @@ def test_interpolate_parallel(self, bout_xyt_example_files): def f_y(t): t = np.sin(3.0 * t) - return t ** 3 - t ** 2 + t - 1.0 + return t**3 - t**2 + t - 1.0 f = f_y(theta) * (x + 1.0) @@ -886,7 +886,7 @@ def test_interpolate_parallel_sol(self, bout_xyt_example_files): def f_y(t): t = np.sin(t) - return t ** 3 - t ** 2 + t - 1.0 + return t**3 - t**2 + t - 1.0 f = f_y(theta) * (x + 1.0) @@ -954,6 +954,85 @@ def test_interpolate_parallel_toroidal_points_list(self, bout_xyt_example_files) xrt.assert_identical(n_highres_truncated, n_highres.isel(zeta=points_list)) + def test_interpolate_to_cartesian(self, bout_xyt_example_files): + dataset_list = bout_xyt_example_files( + None, lengths=(2, 16, 17, 18), nxpe=1, nype=1, nt=1 + ) + with pytest.warns(UserWarning): + ds = open_boutdataset( + datapath=dataset_list, inputfilepath=None, keep_xboundaries=False + ) + + ds["psixy"] = ds["g11"].copy(deep=True) + ds["Rxy"] = ds["g11"].copy(deep=True) + ds["Zxy"] = ds["g11"].copy(deep=True) + + r = np.linspace(1.0, 2.0, ds.metadata["nx"]) + theta = np.linspace(0.0, 2.0 * np.pi, ds.metadata["ny"]) + R = r[:, np.newaxis] * np.cos(theta[np.newaxis, :]) + Z = r[:, np.newaxis] * np.sin(theta[np.newaxis, :]) + ds["Rxy"].values[:] = R + ds["Zxy"].values[:] = Z + + ds = apply_geometry(ds, "toroidal") + + da = ds["n"] + da.values[:] = 1.0 + + nX = 30 + nY = 30 + nZ = 10 + da_cartesian = da.bout.interpolate_to_cartesian(nX, nY, nZ) + + # Check a point inside the original grid + npt.assert_allclose( + da_cartesian.isel(t=0, X=round(nX * 4 / 5), Y=nY // 2, Z=nZ // 2).item(), + 1.0, + rtol=1.0e-15, + atol=1.0e-15, + ) + # Check a point outside the original grid + assert np.isnan(da_cartesian.isel(t=0, X=0, Y=0, Z=0).item()) + # Check output is float32 + assert da_cartesian.dtype == np.float32 + + def test_add_cartesian_coordinates(self, bout_xyt_example_files): + dataset_list = bout_xyt_example_files(None, nxpe=1, nype=1, nt=1) + with pytest.warns(UserWarning): + ds = open_boutdataset( + datapath=dataset_list, inputfilepath=None, keep_xboundaries=False + ) + + ds["psixy"] = ds["g11"].copy(deep=True) + ds["Rxy"] = ds["g11"].copy(deep=True) + ds["Zxy"] = ds["g11"].copy(deep=True) + + r = np.linspace(1.0, 2.0, ds.metadata["nx"]) + theta = np.linspace(0.0, 2.0 * np.pi, ds.metadata["ny"]) + R = r[:, np.newaxis] * np.cos(theta[np.newaxis, :]) + Z = r[:, np.newaxis] * np.sin(theta[np.newaxis, :]) + ds["Rxy"].values[:] = R + ds["Zxy"].values[:] = Z + + ds = apply_geometry(ds, "toroidal") + + zeta = ds["zeta"].values + + da = ds["n"].bout.add_cartesian_coordinates() + + npt.assert_allclose( + da["X_cartesian"], + R[:, :, np.newaxis] * np.cos(zeta[np.newaxis, np.newaxis, :]), + ) + npt.assert_allclose( + da["Y_cartesian"], + R[:, :, np.newaxis] * np.sin(zeta[np.newaxis, np.newaxis, :]), + ) + npt.assert_allclose( + da["Z_cartesian"], + Z[:, :, np.newaxis] * np.ones(ds.metadata["nz"])[np.newaxis, np.newaxis, :], + ) + def test_ddx(self, bout_xyt_example_files): nx = 64 diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py index 877a8d22..18e212fd 100644 --- a/xbout/tests/test_boutdataset.py +++ b/xbout/tests/test_boutdataset.py @@ -1204,14 +1204,14 @@ def test_integrate_midpoints_slab(self, bout_xyt_example_files): ds["dy"].data[...] = 0.2 ds["dz"] = 0.3 tfunc = 1.5 * t - xfunc = x ** 2 - yfunc = 10.0 * y - y ** 2 - zfunc = 2.0 * z ** 2 - 30.0 * z + xfunc = x**2 + yfunc = 10.0 * y - y**2 + zfunc = 2.0 * z**2 - 30.0 * z ds["n"].data[...] = tfunc * xfunc * yfunc * zfunc tintegral = 48.0 xintegral = 1000.0 / 3.0 - yintegral = 5.0 * 22.0 ** 2 - 22.0 ** 3 / 3.0 - zintegral = 2.0 * 36.0 ** 3 / 3.0 - 15.0 * 36.0 ** 2 + yintegral = 5.0 * 22.0**2 - 22.0**3 / 3.0 + zintegral = 2.0 * 36.0**3 / 3.0 - 15.0 * 36.0**2 npt.assert_allclose( ds.bout.integrate_midpoints("n", dims="t"), (tintegral * xfunc * yfunc * zfunc).squeeze(), @@ -1429,28 +1429,28 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files, location): # Default is to integrate over all spatial dimensions npt.assert_allclose( ds.bout.integrate_midpoints("n"), - 2.0 * np.pi * R * np.pi * (router ** 2 - rinner ** 2), + 2.0 * np.pi * R * np.pi * (router**2 - rinner**2), rtol=1.0e-5, atol=0.0, ) # Pass all spatial dims explicitly npt.assert_allclose( ds.bout.integrate_midpoints("n", dims=["x", "theta", "zeta"]), - 2.0 * np.pi * R * np.pi * (router ** 2 - rinner ** 2), + 2.0 * np.pi * R * np.pi * (router**2 - rinner**2), rtol=1.0e-5, atol=0.0, ) # Integrate in time too npt.assert_allclose( ds.bout.integrate_midpoints("n", dims=["t", "x", "theta", "zeta"]), - T_total * 2.0 * np.pi * R * np.pi * (router ** 2 - rinner ** 2), + T_total * 2.0 * np.pi * R * np.pi * (router**2 - rinner**2), rtol=1.0e-5, atol=0.0, ) # Integrate in time using dims=... npt.assert_allclose( ds.bout.integrate_midpoints("n", dims=...), - T_total * 2.0 * np.pi * R * np.pi * (router ** 2 - rinner ** 2), + T_total * 2.0 * np.pi * R * np.pi * (router**2 - rinner**2), rtol=1.0e-5, atol=0.0, ) @@ -1459,7 +1459,7 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files, location): ds.bout.integrate_midpoints( "n", dims=["t", "x", "theta", "zeta"], cumulative_t=True ), - T_cumulative * 2.0 * np.pi * R * np.pi * (router ** 2 - rinner ** 2), + T_cumulative * 2.0 * np.pi * R * np.pi * (router**2 - rinner**2), rtol=1.0e-5, atol=0.0, ) @@ -1496,14 +1496,14 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files, location): # router and circle with radius rinner, pi*(router**2 - rinner**2) npt.assert_allclose( ds.bout.integrate_midpoints("n", dims=["x", "theta"]), - np.pi * (router ** 2 - rinner ** 2), + np.pi * (router**2 - rinner**2), rtol=1.0e-5, atol=0.0, ) # Integrate in time too npt.assert_allclose( ds.bout.integrate_midpoints("n", dims=["t", "x", "theta"]), - T_total * np.pi * (router ** 2 - rinner ** 2), + T_total * np.pi * (router**2 - rinner**2), rtol=1.0e-5, atol=0.0, ) @@ -1514,7 +1514,7 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files, location): ), T_cumulative[:, np.newaxis] * np.pi - * (router ** 2 - rinner ** 2) + * (router**2 - rinner**2) * np.ones(ds.sizes["zeta"])[np.newaxis, :], rtol=1.0e-5, atol=0.0, @@ -1651,9 +1651,9 @@ def test_integrate_midpoints_salpha(self, bout_xyt_example_files, location): # / (cos(theta/2)**2 + (1-r/R0)/(1+r/R0)*sin(theta/2)**2)**2 # ) # field line length is int_{0}^{2 pi} dl - a = r ** 2 + a = r**2 c = (1 - r / R) / (1 + r / R) - b = q ** 2 * c + b = q**2 * c def func(theta): return np.sqrt( @@ -1838,6 +1838,87 @@ def test_interpolate_from_unstructured_unstructured_output( n_unstruct = n_unstruct.isel(t=0, zeta=0) npt.assert_allclose(n_unstruct.values, n_check, atol=1.0e-7) + def test_interpolate_to_cartesian(self, bout_xyt_example_files): + dataset_list = bout_xyt_example_files( + None, lengths=(2, 16, 17, 18), nxpe=1, nype=1, nt=1 + ) + with pytest.warns(UserWarning): + ds = open_boutdataset( + datapath=dataset_list, inputfilepath=None, keep_xboundaries=False + ) + + ds["psixy"] = ds["g11"].copy(deep=True) + ds["Rxy"] = ds["g11"].copy(deep=True) + ds["Zxy"] = ds["g11"].copy(deep=True) + + r = np.linspace(1.0, 2.0, ds.metadata["nx"]) + theta = np.linspace(0.0, 2.0 * np.pi, ds.metadata["ny"]) + R = r[:, np.newaxis] * np.cos(theta[np.newaxis, :]) + Z = r[:, np.newaxis] * np.sin(theta[np.newaxis, :]) + ds["Rxy"].values[:] = R + ds["Zxy"].values[:] = Z + + ds = apply_geometry(ds, "toroidal") + + ds["n"].values[:] = 1.0 + + nX = 30 + nY = 30 + nZ = 10 + ds_cartesian = ds.bout.interpolate_to_cartesian(nX, nY, nZ) + + # Check a point inside the original grid + npt.assert_allclose( + ds_cartesian["n"] + .isel(t=0, X=round(nX * 4 / 5), Y=nY // 2, Z=nZ // 2) + .item(), + 1.0, + rtol=1.0e-15, + atol=1.0e-15, + ) + # Check a point outside the original grid + assert np.isnan(ds_cartesian["n"].isel(t=0, X=0, Y=0, Z=0).item()) + # Check output is float32 + assert ds_cartesian["n"].dtype == np.float32 + + def test_add_cartesian_coordinates(self, bout_xyt_example_files): + dataset_list = bout_xyt_example_files(None, nxpe=1, nype=1, nt=1) + with pytest.warns(UserWarning): + ds = open_boutdataset( + datapath=dataset_list, inputfilepath=None, keep_xboundaries=False + ) + + ds["psixy"] = ds["g11"].copy(deep=True) + ds["Rxy"] = ds["g11"].copy(deep=True) + ds["Zxy"] = ds["g11"].copy(deep=True) + + R0 = 3.0 + r = np.linspace(1.0, 2.0, ds.metadata["nx"]) + theta = np.linspace(0.0, 2.0 * np.pi, ds.metadata["ny"]) + R = R0 + r[:, np.newaxis] * np.cos(theta[np.newaxis, :]) + Z = r[:, np.newaxis] * np.sin(theta[np.newaxis, :]) + ds["Rxy"].values[:] = R + ds["Zxy"].values[:] = Z + + ds = apply_geometry(ds, "toroidal") + + zeta = ds["zeta"].values + + ds = ds.bout.add_cartesian_coordinates() + + npt.assert_allclose( + ds["X_cartesian"], + R[:, :, np.newaxis] * np.cos(zeta[np.newaxis, np.newaxis, :]), + ) + npt.assert_allclose( + ds["Y_cartesian"], + R[:, :, np.newaxis] * np.sin(zeta[np.newaxis, np.newaxis, :]), + ) + npt.assert_allclose( + ds["Z_cartesian"], + Z[:, :, np.newaxis] * np.ones(ds.metadata["nz"])[np.newaxis, np.newaxis, :], + ) + class TestLoadInputFile: @pytest.mark.skip @@ -2110,12 +2191,14 @@ def test_to_restart(self, tmp_path_factory, bout_xyt_example_files, tind): nxpe = 3 nype = 2 + nt = 6 + path = bout_xyt_example_files( tmp_path_factory, nxpe=nxpe, nype=nype, nt=1, - lengths=[6, 4, 4, 7], + lengths=[nt, 4, 4, 7], guards={"x": 2, "y": 2}, write_to_disk=True, ) @@ -2170,7 +2253,10 @@ def test_to_restart(self, tmp_path_factory, bout_xyt_example_files, tind): xrt.assert_equal(restart_ds[v], check_ds[v]) else: if v == "hist_hi": - assert restart_ds[v].values == -1 + if t >= 0: + assert restart_ds[v].values == t + else: + assert restart_ds[v].values == nt - 1 elif v == "tt": assert restart_ds[v].values == t_array elif v not in rank_dependent_vars: @@ -2183,12 +2269,14 @@ def test_to_restart_change_npe(self, tmp_path_factory, bout_xyt_example_files): nxpe = 2 nype = 4 + nt = 6 + path = bout_xyt_example_files( tmp_path_factory, nxpe=nxpe_in, nype=nype_in, nt=1, - lengths=[6, 4, 4, 7], + lengths=[nt, 4, 4, 7], guards={"x": 2, "y": 2}, write_to_disk=True, ) @@ -2237,7 +2325,7 @@ def test_to_restart_change_npe(self, tmp_path_factory, bout_xyt_example_files): if v in ["NXPE", "NYPE", "MXSUB", "MYSUB"]: pass elif v == "hist_hi": - assert restart_ds[v].values == -1 + assert restart_ds[v].values == nt - 1 elif v == "tt": assert restart_ds[v].values == t_array elif v not in rank_dependent_vars: @@ -2253,13 +2341,15 @@ def test_to_restart_change_npe_doublenull( nxpe = 1 nype = 12 + nt = 6 + path = bout_xyt_example_files( tmp_path_factory, nxpe=nxpe_in, nype=nype_in, nt=1, guards={"x": 2, "y": 2}, - lengths=(6, 5, 4, 7), + lengths=(nt, 5, 4, 7), topology="upper-disconnected-double-null", write_to_disk=True, ) @@ -2308,7 +2398,7 @@ def test_to_restart_change_npe_doublenull( if v in ["NXPE", "NYPE", "MXSUB", "MYSUB"]: pass elif v == "hist_hi": - assert restart_ds[v].values == -1 + assert restart_ds[v].values == nt - 1 elif v == "tt": assert restart_ds[v].values == t_array elif v not in rank_dependent_vars: diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 3b7c7f75..df7eef76 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -713,8 +713,10 @@ def create_bout_ds( else: ds["dz"] = 2.0 * np.pi / nz - ds["iteration"] = t_length + ds["iteration"] = t_length - 1 + ds["hist_hi"] = t_length - 1 ds["t_array"] = DataArray(np.arange(t_length, dtype=float) * 10.0, dims="t") + ds["tt"] = ds["t_array"][-1] # xarray adds this encoding when opening a file. Emulate here as it may be used to # get the file number @@ -793,6 +795,8 @@ def create_bout_grid_ds(xsize=2, ysize=4, guards={}, topology="core", ny_inner=0 "MXSUB", "MYSUB", "MZSUB", + "hist_hi", + "iteration", "ixseps1", "ixseps2", "jyseps1_1", @@ -800,6 +804,7 @@ def create_bout_grid_ds(xsize=2, ysize=4, guards={}, topology="core", ny_inner=0 "jyseps2_1", "jyseps2_2", "ny_inner", + "tt", "zperiod", "ZMIN", "ZMAX", diff --git a/xbout/utils.py b/xbout/utils.py index bd711da9..00d78259 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -71,6 +71,7 @@ def _separate_metadata(ds): # Save metadata as a dictionary metadata_vals = [ds[var].values.item() for var in scalar_vars] + metadata_vals = [x.decode() if isinstance(x, bytes) else x for x in metadata_vals] metadata = dict(zip(scalar_vars, metadata_vals)) # Add default values for dimensions to metadata. These may be modified later by @@ -284,6 +285,32 @@ def _make_1d_xcoord(ds_or_da): _add_attrs_to_var(ds_or_da, xcoord) +def _add_cartesian_coordinates(ds): + # Add Cartesian X and Y coordinates if they do not exist already + # Works on either BoutDataset or BoutDataArray + + if ds.geometry == "toroidal": + R = ds["R"] + Z = ds["Z"] + zeta = ds[ds.metadata["bout_zdim"]] + if "X_cartesian" not in ds.coords: + X = R * np.cos(zeta) + ds = ds.assign_coords(X_cartesian=X) + if "Y_cartesian" not in ds.coords: + Y = R * np.sin(zeta) + ds = ds.assign_coords(Y_cartesian=Y) + if "Z_cartesian" not in ds.coords: + zcoord = ds.metadata["bout_zdim"] + nz = len(ds[zcoord]) + ds = ds.assign_coords(Z_cartesian=Z.expand_dims({zcoord: nz}, axis=-1)) + else: + raise ValueError( + f"Adding Cartesian coordinates to geometry={ds.geometry} is not supported" + ) + + return ds + + def _check_new_nxpe(ds, nxpe): # Check nxpe is valid @@ -486,16 +513,26 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over # hist_hi represents the number of iterations before the restart. Attempt to # reconstruct here - iteration = ds.metadata.get("iteration", -1) + final_hist_hi = ds.metadata.get("hist_hi", -1) if "t" in ds.dims: nt = ds.sizes["t"] - hist_hi = iteration - (nt - tind) + if tind < 0: + absolute_tind = tind + nt + if absolute_tind < 0: + absolute_tind = 0 + elif tind >= nt: + absolute_tind = nt - 1 + else: + absolute_tind = tind + hist_hi = final_hist_hi - (nt - 1 - absolute_tind) if hist_hi < 0: hist_hi = -1 - elif "hist_hi" in ds.metadata: - hist_hi = ds.metadata["hist_hi"] else: - hist_hi = -1 + # Either input is a set of restart files, in which case we should just use + # `hist_hi` if it exists, or the user has already selected a single time-point, + # and the best guess we have left is the existing `hist_hi` (although it might + # not be right). + hist_hi = final_hist_hi has_second_divertor = ds.metadata["jyseps2_1"] != ds.metadata["jyseps1_2"]