From 6aa349481c04171d2404994864961cecba508286 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Wed, 25 Nov 2020 19:23:11 +0000 Subject: [PATCH 01/18] Function to create a single region for the whole grid Useful for slab or FCI cases. --- xbout/region.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/xbout/region.py b/xbout/region.py index f4ad4a86..533e6d6c 100644 --- a/xbout/region.py +++ b/xbout/region.py @@ -985,6 +985,37 @@ def _create_regions_toroidal(ds): return ds +def _create_single_region(ds, periodic_y=True): + nx = ds.metadata["nx"] + ny = ds.metadata["ny"] + + mxg = ds.metadata["MXG"] + myg = ds.metadata["MYG"] + # keep_yboundaries is 1 if there are y-boundaries and 0 if there are not + ybndry = ds.metadata["keep_yboundaries"] * myg + + connection = "all" if periodic_y else None + + regions = { + "all": Region( + name="all", + ds=ds, + xouter_ind=0, + xinner_ind=nx, + ylower_ind=0, + yupper_ind=ny, + connection_lower_y=connection, + connection_upper_y=connection, + ) + } + + _check_connections(regions) + + ds = _set_attrs_on_all_vars(ds, "regions", regions) + + return ds + + def _concat_inner_guards(da, da_global, mxg): """ Concatenate inner x-guard cells to da, which is in a single region, getting the guard From 5259dba249daec01d0721a4a80e78fa8f6a33924 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 12 Aug 2021 19:42:10 +0200 Subject: [PATCH 02/18] Fix integration of Field2D variables --- xbout/boutdataset.py | 10 ++-- xbout/tests/test_boutdataset.py | 83 +++++++++++++++++++++++++++++++++ xbout/tests/test_load.py | 8 +++- 3 files changed, 96 insertions(+), 5 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 602526f9..fcb28204 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -437,19 +437,21 @@ def integrate_midpoints(self, variable, *, dims=None, cumulative_t=False): spatial_dims = set(dims) - set([tcoord]) + integrand = variable * spatial_volume_element + # Need to check if the variable being integrated is a Field2D, which does not # have a z-dimension to sum over. Other variables are OK because metric # coefficients, dx and dy all have both x- and y-dimensions so variable would be # broadcast to include them if necessary missing_z_sum = zcoord in dims and zcoord not in variable.dims - integrand = variable * spatial_volume_element - - integral = integrand.sum(dim=spatial_dims) - # If integrand is a Field2D, need to multiply by nz if integrating over z if missing_z_sum: + spatial_dims -= set(zcoord) + integral = integrand.sum(dim=spatial_dims) integral = integral * ds.sizes[zcoord] + else: + integral = integrand.sum(dim=spatial_dims) if tcoord in dims: if cumulative_t: diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py index 0e2282f6..c2cd305c 100644 --- a/xbout/tests/test_boutdataset.py +++ b/xbout/tests/test_boutdataset.py @@ -1304,6 +1304,89 @@ def test_integrate_midpoints_slab(self, bout_xyt_example_files): rtol=1.4e-4, ) + # Create and test Field2D + ds["S"][...] = (tfunc * xfunc * yfunc).squeeze() + + # S is 'axisymmetric' so z-integral is just the length of the + # z-dimension + zintegral = 36.0 + + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims="t"), + (tintegral * xfunc * yfunc).squeeze(), + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims="x"), + (tfunc * xintegral * yfunc).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims="y"), + (tfunc * xfunc * yintegral).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims="z"), + (tfunc * xfunc * yfunc * zintegral).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["t", "x"]), + (tintegral * xintegral * yfunc).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["t", "y"]), + (tintegral * xfunc * yintegral).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["t", "z"]), + (tintegral * xfunc * yfunc * zintegral).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["x", "y"]), + (tfunc * xintegral * yintegral).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["x", "z"]), + (tfunc * xintegral * yfunc * zintegral).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["y", "z"]), + (tfunc * xfunc * yintegral * zintegral).squeeze(), + rtol=1.2e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["t", "x", "y"]), + (tintegral * xintegral * yintegral), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["t", "x", "z"]), + (tintegral * xintegral * yfunc * zintegral).squeeze(), + rtol=1.0e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["t", "y", "z"]), + (tintegral * xfunc * yintegral * zintegral).squeeze(), + rtol=1.2e-4, + ) + # default dims + npt.assert_allclose( + ds.bout.integrate_midpoints("S"), + (tfunc * xintegral * yintegral * zintegral).squeeze(), + rtol=1.4e-4, + ) + npt.assert_allclose( + ds.bout.integrate_midpoints("S", dims=["t", "x", "y", "z"]), + (tintegral * xintegral * yintegral * zintegral), + rtol=1.4e-4, + ) + @pytest.mark.parametrize( "location", ["CELL_CENTRE", "CELL_XLOW", "CELL_YLOW", "CELL_ZLOW"] ) diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 9c22c695..d0da64d3 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -464,10 +464,16 @@ def create_bout_ds( T = DataArray(data, dims=["t", "x", "y", "z"]) n = DataArray(data, dims=["t", "x", "y", "z"]) + S = DataArray(data[:, :, :, 0], dims=["t", "x", "y"]) for v in [n, T]: v.attrs["direction_y"] = "Standard" v.attrs["cell_location"] = "CELL_CENTRE" - ds = Dataset({"n": n, "T": T}) + v.attrs["direction_z"] = "Standard" + for v in [S]: + v.attrs["direction_y"] = "Standard" + v.attrs["cell_location"] = "CELL_CENTRE" + v.attrs["direction_z"] = "Average" + ds = Dataset({"n": n, "T": T, "S": S}) # BOUT_VERSION needed so that we know that number of points in z is MZ, not MZ-1 (as # it was in BOUT++ before v4.0 From 3e9294e759bfcc3b2cf87d219f9228c764add6a3 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 12 Aug 2021 22:32:31 +0200 Subject: [PATCH 03/18] Fix interpolation test for new Field2D variable "S" in test data --- xbout/tests/test_boutdataset.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py index c2cd305c..0b788815 100644 --- a/xbout/tests/test_boutdataset.py +++ b/xbout/tests/test_boutdataset.py @@ -1138,6 +1138,7 @@ def test_interpolate_parallel_all_variables_arg(self, bout_xyt_example_files): ( "n", "T", + "S", "g11", "g22", "g33", From 98c710f628b29cbfcf2f07cd8c417a8f917dda56 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Thu, 12 Aug 2021 22:33:02 +0200 Subject: [PATCH 04/18] Use 'coord' instead of 'dim' argument to xarray.DataArray.integrate() 'dim' was deprecated and has been removed in xarray-0.19.0. --- xbout/boutdataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index fcb28204..a040fd7d 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -457,7 +457,7 @@ def integrate_midpoints(self, variable, *, dims=None, cumulative_t=False): if cumulative_t: integral = integral.cumulative_integrate(coord=tcoord) else: - integral = integral.integrate(dim=tcoord) + integral = integral.integrate(coord=tcoord) return integral From 6bda7c3dcbeb1a2adc1e1082594b98ac1c2ff7ff Mon Sep 17 00:00:00 2001 From: John Omotani Date: Sun, 3 Oct 2021 20:51:42 +0200 Subject: [PATCH 05/18] Fix dimension names before writing in to_restart() --- xbout/utils.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/xbout/utils.py b/xbout/utils.py index de98d5eb..a8a1f939 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -315,6 +315,16 @@ def _pad_y_boundaries(ds): def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, overwrite): + # Rename dimension coordinates to BOUT++ standard names + ds = ds.copy() + ds = ds.rename( + { + ds.metadata["bout_tdim"]: "t", + ds.metadata["bout_xdim"]: "x", + ds.metadata["bout_ydim"]: "y", + ds.metadata["bout_zdim"]: "z", + } + ) _check_new_nxpe(ds, nxpe) _check_new_nype(ds, nype) @@ -338,10 +348,6 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over ny_inner = ds.metadata["ny_inner"] - tcoord = ds.metadata.get("bout_tdim", "t") - xcoord = ds.metadata.get("bout_xdim", "x") - ycoord = ds.metadata.get("bout_ydim", "y") - # These variables need to be saved to restart files in addition to evolving ones restart_metadata_vars = [ "zperiod", @@ -368,7 +374,7 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over if variables is None: # If variables to be saved were not specified, add all time-evolving variables - variables = [v for v in ds if tcoord in ds[v].dims] + variables = [v for v in ds if "t" in ds[v].dims] # Add extra variables always needed for v in [ @@ -400,7 +406,7 @@ 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) - nt = ds.sizes[tcoord] + nt = ds.sizes["t"] hist_hi = iteration - (nt - tind) if hist_hi < 0: hist_hi = -1 @@ -408,7 +414,7 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over has_second_divertor = ds.metadata["jyseps2_1"] != ds.metadata["jyseps1_2"] # select desired time-index for the restart files - ds = ds.isel({tcoord: tind}).persist() + ds = ds.isel({"t": tind}).persist() ds = _pad_x_boundaries(ds) ds = _pad_y_boundaries(ds) @@ -423,7 +429,7 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over else: yslice = slice(yproc * mysub, (yproc + 1) * mysub + 2 * myg) - ds_slice = ds.isel({xcoord: xslice, ycoord: yslice}) + ds_slice = ds.isel({"x": xslice, "y": yslice}) restart_ds = xr.Dataset() for v in variables: @@ -442,7 +448,7 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over restart_ds["hist_hi"] = hist_hi # tt is the simulation time where the restart happens - restart_ds["tt"] = ds[tcoord].values.flatten()[0] + restart_ds["tt"] = ds["t"].values.flatten()[0] restart_datasets.append(restart_ds) From 0e8325f373f93d920ba2a5b3077672cf44846816 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 4 Oct 2021 14:56:05 +0200 Subject: [PATCH 06/18] Fix exception message f-string --- xbout/boutdataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/xbout/boutdataset.py b/xbout/boutdataset.py index 3ca306e1..639db2b0 100644 --- a/xbout/boutdataset.py +++ b/xbout/boutdataset.py @@ -238,8 +238,8 @@ def find_with_dims(first_var, dims): if first_var is None: raise ValueError( f"Could not find variable to interpolate with both " - f"{ds.metadata.get('bout_xdim', 'x')} and " - f"{ds.metadata.get('bout_ydim', 'y')} dimensions" + f"{self.data.metadata.get('bout_xdim', 'x')} and " + f"{self.data.metadata.get('bout_ydim', 'y')} dimensions" ) variables.remove(first_var) ds = self.data[first_var].bout.interpolate_parallel( From 76b80d1482210e3dfc3cdd650f5afbfdb256ae3e Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 4 Oct 2021 14:56:35 +0200 Subject: [PATCH 07/18] Use default value for BOUT_VERSION if it is not found in metadata BOUT_VERSION was not added to the restart files in older versions of BOUT++, so default to BOUT_VERSION=4.3 to enable loading those. --- xbout/geometries.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index 77cc0870..9fb027d8 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -183,7 +183,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 + bout_v5 = updated_ds.metadata.get("BOUT_VERSION", 4.3) >= 5.0 use_metric_3d = updated_ds.metadata.get("use_metric_3d", False) can_use_1d_z_coord = (nz == 1) or use_metric_3d From 1c7a6a71a888431ee8261d1651a5af27c876d9ba Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 4 Oct 2021 14:59:06 +0200 Subject: [PATCH 08/18] Allow loading restart files A few small work-arounds are needed to load restart files and to save new restart files from a Dataset loaded from restart files, because restart files have no t-dimension and may be missing some variables that are present in dump files. By default, we detect restart files by checking if `restart` is in the file name, but the work-arounds can be forced on or off by passing `True` or `False` to the `is_restart` argument to `open_boutdataset()`. --- xbout/geometries.py | 8 ++- xbout/load.py | 80 +++++++++++++++------ xbout/tests/data/restart/BOUT.restart.0.nc | Bin 0 -> 43539 bytes xbout/tests/data/restart/BOUT.restart.1.nc | Bin 0 -> 43539 bytes xbout/tests/data/restart/README.md | 10 +++ xbout/tests/test_boutdataset.py | 9 +++ xbout/tests/test_load.py | 8 +++ xbout/utils.py | 49 +++++++------ 8 files changed, 122 insertions(+), 42 deletions(-) create mode 100644 xbout/tests/data/restart/BOUT.restart.0.nc create mode 100644 xbout/tests/data/restart/BOUT.restart.1.nc create mode 100644 xbout/tests/data/restart/README.md diff --git a/xbout/geometries.py b/xbout/geometries.py index 9fb027d8..f75650d4 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -270,7 +270,10 @@ def _set_default_toroidal_coordinates(coordinates, ds): coordinates = {} # Replace any values that have not been passed in with defaults - coordinates["t"] = coordinates.get("t", ds.metadata["bout_tdim"]) + if ds.metadata["is_restart"] == 0: + # Don't need "t" coordinate for restart files which have no time dimension, and + # adding it breaks the check for reloading in open_boutdataset + coordinates["t"] = coordinates.get("t", ds.metadata["bout_tdim"]) default_x = ( ds.metadata["bout_xdim"] if ds.metadata["bout_xdim"] != "x" else "psi_poloidal" @@ -341,7 +344,8 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): ds[coordinates["x"]].attrs["units"] = "Wb" # Record which dimensions 't', 'x', and 'y' were renamed to. - ds.metadata["bout_tdim"] = coordinates["t"] + if ds.metadata["is_restart"] == 0: + ds.metadata["bout_tdim"] = coordinates["t"] # x dimension not renamed, so this is still 'x' ds.metadata["bout_xdim"] = "x" ds.metadata["bout_ydim"] = coordinates["y"] diff --git a/xbout/load.py b/xbout/load.py index aa1a4c99..802c643c 100644 --- a/xbout/load.py +++ b/xbout/load.py @@ -24,12 +24,11 @@ "wtime_per_rhs", "wtime_per_rhs_e", "wtime_per_rhs_i", - "hist_hi", - "tt", "PE_XIND", "PE_YIND", "MYPE", ] +_BOUT_PER_PROC_VARIABLES_REQUIRED_FROM_RESTARTS = ["hist_hi", "tt"] _BOUT_TIME_DEPENDENT_META_VARS = ["iteration"] @@ -67,11 +66,12 @@ def open_boutdataset( keep_yboundaries=False, run_name=None, info=True, + is_restart=None, **kwargs, ): """ Load a dataset from a set of BOUT output files, including the input options - file. Can also load from a grid file. + file. Can also load from a grid file or from restart files. Note that when reloading a Dataset that was saved by xBOUT, the state of the saved Dataset is restored, and the values of `keep_xboundaries`, `keep_yboundaries`, and @@ -137,6 +137,12 @@ def open_boutdataset( Useful if you are going to open multiple simulations and compare the results. info : bool or "terse", optional + is_restart : bool, optional + Restart files require some special handling (e.g. working around variables that + are not present in restart files). By default, this special handling is enabled + if the files do not have a time dimension and `restart` is present in the file + name in `datapath`. This option can be set to True or False to explicitly enable + or disable the restart file handling. kwargs : optional Keyword arguments are passed down to `xarray.open_mfdataset`, which in turn passes extra kwargs down to `xarray.open_dataset`. @@ -151,6 +157,11 @@ def open_boutdataset( 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 "reload" in input_type: if input_type == "reload": if isinstance(datapath, Path): @@ -232,13 +243,14 @@ def attrs_remove_section(obj, section): # Determine if file is a grid file or data dump files remove_yboundaries = False - if "dump" in input_type: + if "dump" in input_type or "restart" in input_type: # Gather pointers to all numerical data from BOUT++ output files ds, remove_yboundaries = _auto_open_mfboutdataset( datapath=datapath, chunks=chunks, keep_xboundaries=keep_xboundaries, keep_yboundaries=keep_yboundaries, + is_restart=is_restart, **kwargs, ) elif "grid" in input_type: @@ -257,6 +269,7 @@ def attrs_remove_section(obj, section): # bool attributes metadata["keep_xboundaries"] = int(keep_xboundaries) metadata["keep_yboundaries"] = int(keep_yboundaries) + metadata["is_restart"] = int(is_restart) ds = _set_attrs_on_all_vars(ds, "metadata", metadata) if remove_yboundaries: @@ -264,13 +277,14 @@ def attrs_remove_section(obj, section): # grid file, as they will be removed from the full Dataset below keep_yboundaries = True - 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) + 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) @@ -450,6 +464,8 @@ def _check_dataset_type(datapath): - only one file, and no time dimension (iii) produced by BOUT++ - one or several files + (iv) restart files produced by BOUT++ + - one or several files, no time dimension, filenames include `restart` """ if not _is_path(datapath): @@ -483,12 +499,18 @@ def _check_dataset_type(datapath): if "metadata:keep_yboundaries" in ds.attrs: # (i) return "reload" - elif len(filepaths) > 1 or "t" in ds.dims: + elif "t" in ds.dims: # (iii) return "dump" - else: + elif all(["restart" in Path(p).name for p in filepaths]): + # (iv) + return "restart" + elif len(filepaths) == 1: # (ii) return "grid" + else: + # fall back to opening as dump files + return "dump" def _auto_open_mfboutdataset( @@ -497,11 +519,17 @@ def _auto_open_mfboutdataset( info=True, keep_xboundaries=False, keep_yboundaries=False, + is_restart=False, **kwargs, ): if chunks is None: chunks = {} + if is_restart: + data_vars = "minimal" + else: + data_vars = _BOUT_TIME_DEPENDENT_META_VARS + if _is_path(datapath): filepaths, filetype = _expand_filepaths(datapath) @@ -527,6 +555,7 @@ def _auto_open_mfboutdataset( keep_boundaries={"x": keep_xboundaries, "y": keep_yboundaries}, nxpe=nxpe, nype=nype, + is_restart=is_restart, ) paths_grid, concat_dims = _arrange_for_concatenation(filepaths, nxpe, nype) @@ -535,7 +564,7 @@ def _auto_open_mfboutdataset( paths_grid, concat_dim=concat_dims, combine="nested", - data_vars=_BOUT_TIME_DEPENDENT_META_VARS, + data_vars=data_vars, preprocess=_preprocess, engine=filetype, chunks=chunks, @@ -573,6 +602,7 @@ def _auto_open_mfboutdataset( keep_boundaries={"x": keep_xboundaries, "y": keep_yboundaries}, nxpe=nxpe, nype=nype, + is_restart=is_restart, ) datapath = [_preprocess(x) for x in datapath] @@ -582,15 +612,17 @@ def _auto_open_mfboutdataset( ds = xr.combine_nested( ds_grid, concat_dim=concat_dims, - data_vars=_BOUT_TIME_DEPENDENT_META_VARS, + data_vars=data_vars, join="exact", combine_attrs="no_conflicts", ) - # Remove any duplicate time values from concatenation - _, unique_indices = unique(ds["t_array"], return_index=True) + if not is_restart: + # Remove any duplicate time values from concatenation + _, unique_indices = unique(ds["t_array"], return_index=True) + ds = ds.isel(t=unique_indices) - return ds.isel(t=unique_indices), remove_yboundaries + return ds, remove_yboundaries def _expand_filepaths(datapath): @@ -749,7 +781,7 @@ def _arrange_for_concatenation(filepaths, nxpe=1, nype=1): return paths_grid, concat_dims -def _trim(ds, *, guards, keep_boundaries, nxpe, nype): +def _trim(ds, *, guards, keep_boundaries, nxpe, nype, is_restart): """ Trims all guard (and optionally boundary) cells off a single dataset read from a single BOUT dump file, to prepare for concatenation. @@ -767,6 +799,8 @@ def _trim(ds, *, guards, keep_boundaries, nxpe, nype): Number of processors in x direction nype : int Number of processors in y direction + is_restart : bool + Is data being loaded from restart files? """ if any(keep_boundaries.values()): @@ -791,7 +825,13 @@ def _trim(ds, *, guards, keep_boundaries, nxpe, nype): ): trimmed_ds = trimmed_ds.drop_vars(name) - return trimmed_ds.drop_vars(_BOUT_PER_PROC_VARIABLES, errors="ignore") + 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") def _infer_contains_boundaries(ds, nxpe, nype): diff --git a/xbout/tests/data/restart/BOUT.restart.0.nc b/xbout/tests/data/restart/BOUT.restart.0.nc new file mode 100644 index 0000000000000000000000000000000000000000..a6d15105b7172c8cec535e476669afced47901d1 GIT binary patch literal 43539 zcmeHQ349bq)~}xF9KgsSA&G!wLLwkX$VCzaY{(r*LLd==1;-&Xfl+ce8A1{l!A0co zLILFxkBG1yfV!@ztY}a`L?9k33IZyitE}q}K@r^bd#|hB%w!}I{bqDVyC75jURAw* z_3u~JU9VoxEX>W!zNl^2whU8fD3e)7^^iV(5UKNHoI^B~UrPrq1cYyON28gfvk6mq#BeP{C1*<|Al6(hkn0*YRBgD8*wa@ zVxs(l5W7vpdNzF$Z~Z0DocIO>9==H_zYt%vE5?*hk{ujjxW~bi?V70d2As3qr|}?_ zgSRwn3x#uJ^(IGD=*8GxFTFPpbDI9);9-^Y`;R|OxnzO2u{aQ*SVT?O9{&HoVBSXa zaDbvYF0}EH?A*L`Mo-EJZKIt7+*l>@3ZB!H{@T&>lboY|dK=460u;;dH=a080_-aulj@^IUXz_&46@i=LLa;jVxpW=IJ*|H zT4n?ipAr~bX;CvKL^|^g(!K@hCHdJbNPkm9Fe``NgtMDwSwq<6&>)N*Uvr`zGY7## z23x%3IRkqZPSk~sic~tVERY<|7HxVaj2(t%5$v^#@*-ImP(On0{o;;jb{*7>V98HB z8Ouhfa*wT=+MDeHxshzrOUHV%m1?tf_H}*P%}_UrZT@0te>O}#@gcT9jytn!n#4T81_}g))Sy@z$Eqhll3EV*foHKCvFGO zbDsFsfqwEd2C^`*l!HSn{9OW;F49eqcZKDtNXu0cN2Z$iC8gmf45JN|ehdeZj5{mf z@GSgwh2bMo80gm(mX;#L0U5z)L<`ewPo@0IcOdcxs2uK^8s-Dpg>|A*ZV(p}$!$q} z^hVZpwQ?Ah9#GGVh^dJizw_4EA5e>ceCJl+u) z7d$;o2mYP!`f$Z--p1mK0L5aD?GUVDYNu}BS2Eb1P=^{^9*W`S(6!cj}PUfsz@Cph@KSx_6y zaHk(BCI=y@9?~BF5P=A_rwA2Y%MlbsC`y8zua4mLcZICw=M|UWmI4oHLsvTP7qk>` zhe}M(hYv`2jDdBqmFmDx#tM1HASUfcY3#U^4}F7$$q%r=-5ZWY#~Hlq2MK}mU@^a{ zTeTC5=20I!pwXy^P1E5lxgAUX0d|9eIwIg`6VCk{wkAhbV?x|HWFd0}D$cTdNFo2q zUth6oudlJVEkLjsS{Kzfy(F)Y1><;B(@e0e?krLrW-v@M<2XBN-HxUlQANvsPYCe6 z;;z)CCF=F>M&w47>`vqxDtRf9p90y9^&nCQ)5ayR=`RxBJ2w(}C6sny(M0l%{+>kg z(1I8udGr9?ZA93>WmL*z7J3oM_x5`e$)gGS5ZOh>&0*a1R#QKf>`Nq%W9UaDj}4%h z3K2aqfJ%8_fEW@Gf&ro;*gz`f@e1)o?o>+?h~yC!iA3_ygd`#_3c_7?h_L1<8pv=q zh{#kRyEsFnMX-Z2;++R`q)}xa$WS@^9B%AiIQe=mToNF2;kd8gz0%iMtO^h;QWEFy zaAKhjB#B2X^M^t(Dhk)fXEIDgi@4nLJWueam5Gyi{abj(T7I5r5Lj&zr=Ta!pNQEN zPn=;8MG^@wmGBo7GAb+W7F&hcp}tzuJi--~peIk5%G5+T5_szTl?9jk28)#@^@Qxg zSQIbdHGgX?W}n#mu5Yl|eNAhz&@b{2EMl%7H`+H?9BTC#dFD3%z@p#b(^GtdMSpW^ z?V|J!|G?sSyQPu7!J@~tt;M3V2qyil=ZkeGw^@9H#lx+}V%uW>u#1?e*zUf;VpFTf zNVQm#|JHVKMfQQ!zQJOjrL|)uW~qPJMM~9IpZErgdrDi2#oNpL1B+vdmHQ)tmK}U@ zArrm`;n3>w&XrutuG4$yR`Ema^LK{;e|O>TC2$Ypi{kG#oL|Vyh&!eh$uV8#xf8&y^Op-1X1{jC&&1Rn2@{M1RL_ zi++-E-^SFkT^kq=^KKiS{1oHy;P(3ao@PAu{ezBa&oJ&5$+>a<-xv>6kNDr|=NS*3 z@A!kfRRp&Oe|2^@`It;igAZZ;mZro ziopAI^PV^>;`cwWJbRXLPe8-F$1{F+7hZJ8H*;ZltM3bQ@8gR6*7wMoW7}Qv+Fo6Q zCr{x$yzsg+X@!4ak*2(IyKn4bbE~oVWTk&#QFFt6t{7#n+Qs5&t+k6KtNjCu(F51| z2n*fR)?$&d#y_wK?(^p7z8NF8wi=7CAMg(>iaw}}_6-(oTRlb^9`X+?!tz$Q;;Fnk zUzAu|J4T8d`~!=IPdnH6#xC}^8Vmhe|G?q}d(p+d!D3QbYwhBhNBskfCEJhWb1cFn zN!>`3;LEEp@lV(O2GNw$z6a44+W!c%@g583!%ysgl&6*QF zj=jSDkjF}=>T)th=c~OGp#jw6P|@z4V19u0F?sSyKJiBh0Wz=dtWSK{@>CPme^Lk* zr}ar)F!gGjEa_D7d;*%T! z_mBEZuQyYDl0)w~6QAVJ3(ubXFb~{F4!r{HS)bm8_T_g$LLX?^ojQ1lboI{@MApqBq!E`PjW76 zf#1u6PjV>y!BhV}9`#8+1w>GNl0%UZ#HaccBH_uWSP9~j915BsKFOiz2~R$SQ4pWx z#J8Y71y)dfl1~v9#3wlvYC(LepWK3c3cjHFB!^~v;*%T-$MED+Tn6z;&fpf#dqoTU zAs&2^Gt`4m`V4Dd_@qyU2cPO^dhk2Y4MFf1(?l(3 zEz=Ot5YQ0N5YQ0N5YQ0N5cqWv7@d(jir)30Z{8E3HO#ORzn>1@^>M)d=Gd(4tg%^z z8CfR&u1`3u+;J6aF+0rFR>uX^#g!yJEjz!TTC|`i=#~Vye|IU?X(_TZtWR?ZMp_C9 z_TP`rio@RWh}9D_IJ;4l8m>jJ(9eCZ2_R4t<>>sN_nPp5*PTYs$DCy+-*oAaSgST> zzS9A{74ABqmtyR7|E3f2)BE(()!V9>{6ute?<0;r9{1j&Z)a^wvz(h$C}myL4DQ7d7yu4(M-m*8%;E zsgJCPb?Gqn7SzC(I-sx8T?h10)^Mr8(lh# zyxG?}px4`72lRG_9M+w6>0t1=)uPATT?h1(-!jwF+ZMn5$%(PkpRk?4kT#l1#T7xo|>J)_M`LA<5x{apA;_hfFq4i&z>*~}triA@HsSKCOn4D#aM`Jig zlDqK(DoZpad=*P&xTL@&+9jG2zsD6YBuU2Q1*Xu7I+LxU!fFo{0UMGmy2pa;Gb&6r z3%MsHS#*$f!~(mu+F`akOl9`Ua`IS8vQZTZQCbd1e=!gwS#*-g6AMO+PZJ|Sl12ZJ zomZGC>{XIQ*O(!F^i?IMaam*Wmj#Q2-Ac0P9n)EA>E6Gl*B=-~rw^=De{V{yB-K**fF#p2BGpmofF#pYBGpsafF$eHmVmkA zh9G_@$uwaGmrgU=lTs*FO@eQ_;wc1v9MNH#W=lcm86-&#<_9dK9)6h8R%uCZ!x@qs zB1(b^#us0ez8IZ}kRhcW`ni z=3FZVh$N%S3X`s)cDfiFlH7^&fLT{728ARG{efA3ix>@(EL3N@f{Ba8Fpy-SHIoa+ z7iF!c5g=sp4DW7+;m~LJhd(PB$}HCEQoF6n0l!PBRzb1M3F_vXIYsL!g#udJ{dXPT>xV(eKhMjU-{z3W?% zdK=5F=89@nE}R8_d6WOdrV!TOsP@5Xtg>5cY?U+MEKtj0v#8aknQN>@GqPYth&vN| z_yO$O`M;CP0Cx@k!<>j;sdf5v)AY(x6%^9_<^Hl5RA<4_M9uO$n zHPEaGM~LUQ>uDqTN3)7H7iqJpc1P3h$*{W9o+`8_JMGz28wc77M_Z%8y}2>n?R@W;Zdu+Z%-05(+h>l`fHo~DZIXz#%aj%bcu?%wx)5IL*;9sf8^G@{A>Cl0@{5_vjfcze!X$PtxfyX5YQ0N z5YQ0No5CL!70Y_E))DX}R z&=B||BXIkv*u1jAv$j@0)InD}qjalv@1k}$%_vn0Gxk>BJb0Eeu20h$cHIJ{ac93t z-Fx4yO#a|@UG(MmD*p?6^RqF^DrK_5Ui>iSK_%s`KOYG@_^@JY9J0%B_gZC8@KZx3 zT)kdN+cIzLr1-~`K96pFJLA#~O2&J6<~et5QZ@#EG&82#)5;4a8}he3^^B5o!@my9 zD1KHs)#>wX{n&HL_nkh@+4S6IrC-FDd)C*zpj5Wqv+$m*ElR_>eo?P>eM#9eFw!{v z%vPo7bH{|&zkFE<3+feh$I(}mdka6_clcnV^7Yj_LfXH#Oq>CyQGM*HT}s9LvNn0Q zyrDcgru)_{WxJK<^HaYoAMvJg>%f#T`JMMDp}odz8+Blh(y%J*T-oxy%Bx3m6PjZ8 zDXVOz2ZCqrQ%dh$6E~%CpYp@;UOnrBo0L^S1!w-9-lQzrey;u0nN7;K@D*yig--pz{ literal 0 HcmV?d00001 diff --git a/xbout/tests/data/restart/BOUT.restart.1.nc b/xbout/tests/data/restart/BOUT.restart.1.nc new file mode 100644 index 0000000000000000000000000000000000000000..512672f928776875f970869ec0ffe2c8573db06b GIT binary patch literal 43539 zcmeHQ31AdO)~=rE9LR8oBnBiC5)OgHT!f%t=RP0_SrQ|#=rGAl!YDbMgpkCaLq$Qn zP+*0Cpd7+_f#Lx!3ku6&L?EsQ3IZyitE|f+Cs8g8nT$lD|BTLP7i6m6tE$(p zzJ68R_3HJ^y;*4)ovsMEfC4Ks z!$q~vb^M^A6VD$P`4t)>0jFngstF-liQoyHbKPjDCy~y?pZ+uCrRG~JtIXCKi@BQe zd(gjdR!%=?sx@1yEtM4`V~13jYf4O|BjclnMa2)9Y$_c(GA1e}HY$1u8OE#rqR-%u ztbJKmI~7S%Fo=*O))C$vpoLs*033D;3uPT7L{EbJNzhDM>f;Y38jXfk?MkmhAX#VO z&H89~j<^K3jY~Atl!!0GB{e6xpn%eV2V7Vtn&xTQx#@WY*<M4q;X&HKxrQBQr!0L;q*61rLYwSg}MythGJjEQP z#}q?!@w9(bT4$nT*Nk0?_%uto#YE52)t}ffp``7sa7@Vh&dU~F8sP+9fnfgg5Z-kw zZ~i+#cxe1DM3`US$4c(4xHQ6e{W1SJ2orL=oFjZhTqm?)`442neE4eK{^(7P@QqHh zgn-6|$Yl}M@>#+k1~$=ZG&nuNC4!j~LDt$tVhzl79y?J)|3Wj?(tdw%cK4=B8*v0J zg@pNKA$FLEwPeaf-uf$^+3@x8K78X6ej&bKSA-~^#@BF!={=54+4hNAtAjb)aT@nw z9Jr-HTUt0vW;Z#c{JTSYt@QrgkmK}E4^}eL<3IjfeC1qMV{yn!u?QW%EAan+!CZ~z zQ7=Wa$bb9ljI8Vwf=|i-?RF~%czcz|D|n7m`p)AiXE=v_x*E$*y%fuf>rWlZ7@L(( z{3t9M#{~`o*B~CRcZp>PeLCgIw{?!o7r|{aqw|}Ph_j=f&VI`}@&*!Fnmy`)l)|hu z`n*l+b);pcf$xmKQasC%x;dY*TE0-L3#*kj#JSe0?Kfv+6i`{@A)yc42O(k3DVSaJ zNgdHs5}pzWX*98n2$8n@IqAUMl){`0;-j6h%9oVW-UN~v)6IV5YT6)xoLqjYBQg5W zhYT`*!7DoQ5cuHvj58m!pJSOZV-uo z=Gh36%;cU}S~h?*QMtin;?@%bNF!^u*1C2Oxr5dXC7Zq+F@%g{Cq6=YyK|2&aA&xGw={y`dfo@&;U#C1E9O>Uy9jMW`Py4_5^VG_QUzcr91cH;B>uDzAq$W+f<{r7T` z%7jfcEUP26R6+nbXnLuE+{K#hDXF@X^rFqe$?1x1r>M4Gli2k~>j&hZYZMlqxE;`* z^Te+X?I%xvqFI1g%F#nJ{kxJ{x=1r<9zvI=A}wYlj7*mJCB@+<45I^15dqm={gQ^Ooe_9We~lp92MB)KiI zM{i*5VwD51^bk8QAjT3me%q~a2t_T3{u7Ct5AzrRhnr^Nvs)g-3cpeA>U+m$ck_7yk(U5&+;UW&z*d9Rkc8jJsUDHbOxBR_UE7GHZQ7QwUD&2%*u-+Cz) zC6)RcT#dyAFU4YB--CXx#^MJr#p0%`D_dNRh4vnP@AS5NYWeFkn_P{BpO<1WZgN_X ztFh?lrC5xA?8Hs3#-fXtV&Q9kL2P%q&db`vOR?DZsdT%m?IOrav6y`(^suY3=;ftY zyf^INN>^hM>ZMrxbH<|`T#ZGzmtyhE`lxGMjl}>j#UkkZ%`Tri2YV?Nwcq?T*VT4$ zjhAB4FR@wgYAoWs6pPh+*IesrEQWb07WGfxbJT`Kt3b90;i$!2%Xad?2~Pdp#HRyf zu+t9|k$r$^ z4wRUZLqBMOV+gFHTd6hhNk|^g2#rbmQ5sv6@R574Fn9nKuzSO?=sJ~m{V*bM2`uJR z^{TdE(K_mb2Q=yxv1yt_;yaS~ALwq7PgejOY{IsmLzZXCEGEQ}gBEg^K*d(}2rA@R z`5Q|X?RPg8_jw5xBkDs3r4(l85nmXOEKQ@9)tdyfVMd2(S`=r8uHD_7DXOU1?+F3E zR~&*}TEMPmU1*VEFTBn(Nu(cc$I z9$FBNB#$1zyNw7N=!c~|W}!cld~bgMl02GVAd)?0*c^sUZvNZZd7s{~dX#3}HJ^Cx0< z#S>=$jUoxAFO~2Y6jCcIttLx_v4*``(mKKwl)xuXI+d}6IS_E_oW|U%-GfD=ft`?C z9*csxyykC>#mrOt-*pcbdv9(p7TS59fkpVOMPuB9#gTT8k&X9x1{Q;lo-1|_7DJ5f zwTqGmJOhi99Tx_>2a7(F+KWXG5ls49&lhXY>@c|piXu#51}h~DnOVne&f zNVQm#|JHVKZN{Nx?!n@Ksl8()e4%I9MMBl-Pu+vXBPH#{;-8B=1B(-infoJr79D=} zULt%E!lBjYo$I-lAyfKjmhwaO^LK{;e|O>G1@s=q7scN_cWyr63*r%R{U0Gbx;vtH z(O(D;KM$?myM*v%{p|}v8wqzq{7o0RoN#Z(?qz{12#*m@YkFb@;hu=KRnwjp(ccj} z!k#7Ew^3Hw^c>+~-d6<1KTmi(xV7QY7YL7i|FCP~M#9}9nYYh*k?=tEp#MF$nefp0 zu0P1zL~y(B>6X2O`yWnNBHt&U(5n?p_Kj*fNO-vWl5_LU5bh{wikWbhaED6XYjayf z;QiXWpJ@^C`yZNKY9ZVcuhSG#^UQ0o`FUFhm~RO!QzT`kC9bRcm@^$*-Py4R4$z_ z3eD{uBL%BG1B+Fk^{91^T^wvT7TPtQfyJxV{O<0-Vq$4~?PBATo`J=JUB_}b76FpP zHqs>eC_ zUcr9IQzd1Z%+xVCtd}A*0DJ5&+P&k;51>9IXFkeD{%9dU=GATWkx#cgS;G3y3c=zu zKdSHNR3GJFd(2S|zLt!9tdH*}JM-~HW#pqAd}A5;CBiZz%my*CV^Mz=f2lOAz%q}aDL!>P4K|$NQ38N)}k|S*>#AmV~vUX zbO3LKqYmJu7<@r~~+} z%ulAJ*maO+f(Gu?0VaD#9bmE#I`B5SAK~OKO!l5_{U0vYedvP6(y?cHN+*7E?>4&* zv!=UShrUiafR{6S)5+}_1qJB@zmQj=^XobA|4a0?at2hE=%hFg?sFdpEfcW>lpuD^ zFS>nvfZx~=J@6MVT(jzGyAJ#A1fASzo8U8a)B${jLC0pzwCgbcF3`Z8I)J~?Q3vod z4qLq>!mh*E*`R?tbpT(bqYmJstPLCZAe{B5lU>WXvv^JSjW_UIJ8B6Y>$C^U*V}a% zb%(ok0I#>B4&d$fJF021>!5SH)q=;|Q3vppKQzYN>^OdRxqV)T{Rl@bU?1Xde(}@o zNCezXSv2k78PQcii=X@C4)Ycy%pK2}Livt#suK;q?T`T}7e5~-)I zop2Wr^NF~|ORJ?xs%n(m$1njV-| zKpL$BX|xWc(K?Vu>p&W<18KAlq|rK%M(aQttpjPY`LEwy$~tUw$0nGUJaz8CUSN*B z({SS;z8%lkW@M!~{7MGx*W#&z-jQVV&}hqIV{Jwuqn}1w78hqT4H-AAwPi7}HUp92 z_5g`;<7{Rk<7PFL8)q{X8MmtCGDA#^&1Ph{-$SD8*jStC$hc8W*>Q0;3zDOGWvCq! zYcnJnx2b_0XR{|6r%{j{XEQ4qXH#jE&9r2kQMD;`^)>^OabDGA#@MV(#>rKa8EZ2& z8E05cW}MC5WSnNDLcA(M{a-n2ef_CrymTa|za-=BM$Q|bpT16nY)En-Vw^W2KmA2m z|Ij4Ku+WF(XQ`_;S5?Qnre`EA3FZe_61z7HOK|-J zF`;`>l0^qeS17QWt80wb8bhhIvK&2@lB{P!X_Qt?%@8pVBw2Kl$P;pt$0dr9AjzVC zz|PJ~6ZR^}qH9E#GUkRtLs9xz_+`O-VYiYjdPg*-I=uNyvgj6N*4yTI86#6PrUu(Y zFN+4S$`R5KbvfrCxm83chACP36Mx=TS9gt+4N~8u18<1ozYe_M$ zxy^?kN-|EEz9o~5*0=s9UGd~ce;l#KFxip-&NG@M*_R&>zXtlll$J_UN(au6 zWIs{jlRK{9hLri>Or%MoKaRzdPwsfiKirX1k}Mh$pWF$QAKV@Jq9O6iPdALCS0nc?&ny&!M3TGM zxDzrbi2)+X;Ie|Gsi>PG#)c$!<2=f&uM>kpl7;@1*)Us-21ypG6HV@f`C=GIve25y zdE@fam*EHyGI<7fH(e6#GyMlYE9p;6=IRowrK*PhE+wl1l0i$n{5RZ?ou)6XwCZc% z_bjU|1yr_!^~S|*E$Rhw+AHS3MQqBBC& z`G}({pl_G{PA(n2Yw#cD1pG?OQ>GZERF*JMXiot@uqVVXH9aTCkeZ%XI5r)AL=YYy z1n?HCxdf0j)YIy`A(%8aET}P7n2c7Fj@|%av&XqrgRK<|5Nwr9f4tRLX5MsN{35jn z2#R_QR4alJ;{5G;>PY_4tfJ0E>TIgs(bRh~UEQfq73!0n`fRF>19gR?uF>edz~k=+ zbanmH>*PzVt{b0vYK}?V>0r4Cj$&LkLXO8a&llx6)`xN^gB-&&AqP3|S0eK0;>5(* zgp&BknBl{RN5)4d7DpzUjPa4BLyZZ>=){DQ=+dF>@NzVqVVOJ)pxp&~;>EmD34-ioAQ>q=PcJS+s18!~Vr;31z zfQo>MfcpHwZ9v6A#Q_jd?ck4X2fX*{r;31zfQo>MfOZgi; zihzp19~lAZMtR|gv&~BU;+Sn&-#06g?`waD+O>0(~j_v&J-zm+?SN;AI*5KQ$kdHe=7H>bGoH+LVO5bS*lxXeN z!Oal|l+#ZS=uozJzY_b!?2V%j?NheS8E4Gtu}@jP=-YplC%>g!dGU$8TT1sTJEU{p zXU%?7*;rcADPn1pGSTnte8V#|Ug1g0i4L3hD1AGvOpe^XTWJaVs_D%=yOi56+&BHm zTRW8<>+k=j;O&1XAwg$J=l6Fgx5XvBcl7XfCFSK;dOdLbb>+k2;}hQa>NRC^_Bj2N z^V^ihqFGn{BV?;$y72nEN7A<_w;qUCxUT+H<*oWH$s1nTr0gsY+dr7Rq9mMr`OT>X zFDb!+-Hbb)->4jz@xY|;UN0y^jiC**9@?OEpZwkvsaHLxG%g-@=fvozmH$Sr?K1wx zbxP88wg1&Uyhb_m${kw+4zE;-il2G^qlCwmm%sXa^PiNZ%227}fUv9oq7>bjo6>v0 y!%ENk6@k(%bCtU;o-Mg!_;h8RWyg`}Q%jUvuFu#wZEDH3FE%GXZXP~;+y4RU-`(#3 literal 0 HcmV?d00001 diff --git a/xbout/tests/data/restart/README.md b/xbout/tests/data/restart/README.md new file mode 100644 index 00000000..e2b0b71f --- /dev/null +++ b/xbout/tests/data/restart/README.md @@ -0,0 +1,10 @@ +These restart files were created by running the 'conduction' example from the +BOUT-dev repo (at commit 696735c37ce31491c730aaf547795ff7cd593ae5). +``` +$ git clone git@github.com:boutproject/BOUT-dev.git +$ cd BOUT-dev +$ ./configure; make +$ cd examples/conduction +$ mpirun -np 2 ./conduction +``` +The files are created in the `data` subdirectory: `data/BOUT.restart.{0,1}.nc`. diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py index 12943b70..b4c07bbc 100644 --- a/xbout/tests/test_boutdataset.py +++ b/xbout/tests/test_boutdataset.py @@ -6,6 +6,7 @@ import dask.array import numpy as np +from pathlib import Path from scipy.integrate import quad_vec from xbout.tests.test_load import bout_xyt_example_files, create_bout_ds @@ -2254,3 +2255,11 @@ def test_to_restart_change_npe_doublenull_expect_fail( ) with pytest.raises(ValueError): ds.bout.to_restart(savepath=savepath, nxpe=nxpe, nype=nype) + + def test_from_restart_to_restart(self, tmp_path): + datapath = Path(__file__).parent.joinpath( + "data", "restart", "BOUT.restart.*.nc" + ) + ds = open_boutdataset(datapath, keep_xboundaries=True, keep_yboundaries=True) + + ds.bout.to_restart(savepath=tmp_path, nxpe=1, nype=4) diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 21207d0a..93156c37 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -1161,6 +1161,14 @@ def test_drop_vars(self, tmp_path_factory, bout_xyt_example_files): def test_combine_along_tx(self): ... + def test_restarts(self): + datapath = Path(__file__).parent.joinpath( + "data", "restart", "BOUT.restart.*.nc" + ) + ds = open_boutdataset(datapath, keep_xboundaries=True, keep_yboundaries=True) + + assert "T" in ds + _test_processor_layouts_list = [ # No parallelization diff --git a/xbout/utils.py b/xbout/utils.py index a8a1f939..2477ac77 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -317,14 +317,10 @@ def _pad_y_boundaries(ds): def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, overwrite): # Rename dimension coordinates to BOUT++ standard names ds = ds.copy() - ds = ds.rename( - { - ds.metadata["bout_tdim"]: "t", - ds.metadata["bout_xdim"]: "x", - ds.metadata["bout_ydim"]: "y", - ds.metadata["bout_zdim"]: "z", - } - ) + for d in ["t", "x", "y", "z"]: + key = f"bout_{d}dim" + if key in ds.metadata and ds.metadata[key] in ds.dims: + ds = ds.rename({ds.metadata[key]: d}) _check_new_nxpe(ds, nxpe) _check_new_nype(ds, nype) @@ -373,8 +369,13 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over ] if variables is None: - # If variables to be saved were not specified, add all time-evolving variables - variables = [v for v in ds if "t" in ds[v].dims] + if "t" in ds.dims: + # If variables to be saved were not specified, add all time-evolving + # variables + variables = [v for v in ds if "t" in ds[v].dims] + else: + # No time dimension, so just save all variables + variables = [v for v in ds] # Add extra variables always needed for v in [ @@ -403,18 +404,26 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over mxsub = (ds.metadata["nx"] - 2 * mxg) // nxpe mysub = ds.metadata["ny"] // nype - # hist_hi represents the number of iterations before the restart. Attempt to - # reconstruct here - iteration = ds.metadata.get("iteration", -1) - nt = ds.sizes["t"] - hist_hi = iteration - (nt - tind) - if hist_hi < 0: - hist_hi = -1 + if "hist_hi" in ds.metadata: + hist_hi = ds.metadata["hist_hi"] + else: + # hist_hi represents the number of iterations before the restart. Attempt to + # reconstruct here + iteration = ds.metadata.get("iteration", -1) + nt = ds.sizes["t"] + hist_hi = iteration - (nt - tind) + if hist_hi < 0: + hist_hi = -1 has_second_divertor = ds.metadata["jyseps2_1"] != ds.metadata["jyseps1_2"] - # select desired time-index for the restart files - ds = ds.isel({"t": tind}).persist() + if "t" in ds.dims: + # select desired time-index for the restart files + ds = ds.isel({"t": tind}).persist() + tt = ds["t"].values.flatten()[0] + else: + # If loaded from restart files, "tt" should be a scalar in metadata + tt = ds.metadata["tt"] ds = _pad_x_boundaries(ds) ds = _pad_y_boundaries(ds) @@ -448,7 +457,7 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over restart_ds["hist_hi"] = hist_hi # tt is the simulation time where the restart happens - restart_ds["tt"] = ds["t"].values.flatten()[0] + restart_ds["tt"] = tt restart_datasets.append(restart_ds) From c0730eddde86afdcc3a2e0900a9e1edf95a80107 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Mon, 4 Oct 2021 18:09:47 +0100 Subject: [PATCH 09/18] Pass is_restart in _trim() tests, test both False and True --- xbout/tests/test_load.py | 56 ++++++++++++++++++++++++++++++++-------- 1 file changed, 45 insertions(+), 11 deletions(-) diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 93156c37..924db2e4 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -1239,18 +1239,23 @@ def test_restarts(self): class TestTrim: - def test_no_trim(self): + @pytest.mark.parametrize("is_restart", [False, True]) + def test_no_trim(self, is_restart): ds = create_test_data(0) # Manually add filename - encoding normally added by xr.open_dataset ds.encoding["source"] = "folder0/BOUT.dmp.0.nc" - actual = _trim(ds, guards={}, keep_boundaries={}, nxpe=1, nype=1) + actual = _trim( + ds, guards={}, keep_boundaries={}, nxpe=1, nype=1, is_restart=is_restart + ) xrt.assert_equal(actual, ds) def test_trim_guards(self): ds = create_test_data(0) # Manually add filename - encoding normally added by xr.open_dataset ds.encoding["source"] = "folder0/BOUT.dmp.0.nc" - actual = _trim(ds, guards={"time": 2}, keep_boundaries={}, nxpe=1, nype=1) + actual = _trim( + ds, guards={"time": 2}, keep_boundaries={}, nxpe=1, nype=1, is_restart=False + ) selection = {"time": slice(2, -2)} expected = ds.isel(**selection) xrt.assert_equal(expected, actual) @@ -1377,7 +1382,8 @@ def test_infer_boundaries_2d_parallelization_doublenull_by_filenum( assert actual_lower_boundaries == lower_boundaries assert actual_upper_boundaries == upper_boundaries - def test_keep_xboundaries(self): + @pytest.mark.parametrize("is_restart", [False, True]) + def test_keep_xboundaries(self, is_restart): ds = create_test_data(0) ds = ds.rename({"dim2": "x"}) @@ -1387,11 +1393,19 @@ def test_keep_xboundaries(self): ds["jyseps2_1"] = 8 ds["jyseps1_2"] = 8 - actual = _trim(ds, guards={"x": 2}, keep_boundaries={"x": True}, nxpe=1, nype=1) + actual = _trim( + ds, + guards={"x": 2}, + keep_boundaries={"x": True}, + nxpe=1, + nype=1, + is_restart=is_restart, + ) expected = ds # Should be unchanged xrt.assert_equal(expected, actual) - def test_keep_yboundaries(self): + @pytest.mark.parametrize("is_restart", [False, True]) + def test_keep_yboundaries(self, is_restart): ds = create_test_data(0) ds = ds.rename({"dim2": "y"}) @@ -1401,7 +1415,14 @@ def test_keep_yboundaries(self): ds["jyseps2_1"] = 8 ds["jyseps1_2"] = 8 - actual = _trim(ds, guards={"y": 2}, keep_boundaries={"y": True}, nxpe=1, nype=1) + actual = _trim( + ds, + guards={"y": 2}, + keep_boundaries={"y": True}, + nxpe=1, + nype=1, + is_restart=is_restart, + ) expected = ds # Should be unchanged xrt.assert_equal(expected, actual) @@ -1409,7 +1430,10 @@ def test_keep_yboundaries(self): "filenum, lower, upper", [(0, True, False), (1, False, True), (2, True, False), (3, False, True)], ) - def test_keep_yboundaries_doublenull_by_filenum(self, filenum, lower, upper): + @pytest.mark.parametrize("is_restart", [False, True]) + def test_keep_yboundaries_doublenull_by_filenum( + self, filenum, lower, upper, is_restart + ): ds = create_test_data(0) ds = ds.rename({"dim2": "y"}) @@ -1421,7 +1445,14 @@ def test_keep_yboundaries_doublenull_by_filenum(self, filenum, lower, upper): ds["ny_inner"] = 8 ds["MYSUB"] = 4 - actual = _trim(ds, guards={"y": 2}, keep_boundaries={"y": True}, nxpe=1, nype=4) + actual = _trim( + ds, + guards={"y": 2}, + keep_boundaries={"y": True}, + nxpe=1, + nype=4, + is_restart=is_restart, + ) expected = ds # Should be unchanged if not lower: expected = expected.isel(y=slice(2, None, None)) @@ -1429,7 +1460,8 @@ def test_keep_yboundaries_doublenull_by_filenum(self, filenum, lower, upper): expected = expected.isel(y=slice(None, -2, None)) xrt.assert_equal(expected, actual) - def test_trim_timing_info(self): + @pytest.mark.parametrize("is_restart", [False, True]) + def test_trim_timing_info(self, is_restart): ds = create_test_data(0) from xbout.load import _BOUT_PER_PROC_VARIABLES @@ -1439,7 +1471,9 @@ def test_trim_timing_info(self): for v in _BOUT_PER_PROC_VARIABLES: ds[v] = 42.0 - ds = _trim(ds, guards={}, keep_boundaries={}, nxpe=1, nype=1) + ds = _trim( + ds, guards={}, keep_boundaries={}, nxpe=1, nype=1, is_restart=is_restart + ) expected = create_test_data(0) xrt.assert_equal(ds, expected) From f4932e0c3813a0d935bedc3b9def2be21cf2acc0 Mon Sep 17 00:00:00 2001 From: John Omotani Date: Tue, 5 Oct 2021 11:44:26 +0200 Subject: [PATCH 10/18] Write PE_XIND and PE_YIND in to_restart() Makes re-loading the created restart files simpler, and makes the created restart files more consistent with ones written directly by BOUT++. --- xbout/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/xbout/utils.py b/xbout/utils.py index 2477ac77..6346a7e4 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -454,6 +454,8 @@ def _split_into_restarts(ds, variables, savepath, nxpe, nype, tind, prefix, over restart_ds["MYSUB"] = mysub restart_ds["NXPE"] = nxpe restart_ds["NYPE"] = nype + restart_ds["PE_XIND"] = xproc + restart_ds["PE_YIND"] = yproc restart_ds["hist_hi"] = hist_hi # tt is the simulation time where the restart happens From 9545faff7d4058486cf1ec8c3a3b07c002831222 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20Schw=C3=B6rer?= Date: Mon, 11 Oct 2021 14:41:52 +0200 Subject: [PATCH 11/18] Add fci geometry --- xbout/geometries.py | 49 ++++++++++++++++++++++++++++----------------- 1 file changed, 31 insertions(+), 18 deletions(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index 6101d29a..68d5052e 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -288,6 +288,28 @@ def _set_default_toroidal_coordinates(coordinates, ds): return coordinates +def _add_vars_from_grid(ds, grid, variables): + # Get extra geometry information from grid file if it's not in the dump files + for v in variables: + if v not in ds: + if grid is None: + raise ValueError( + f"Grid file is required to provide {v}. Pass the grid " + f"file name as the 'gridfilepath' argument to " + f"open_boutdataset()." + ) + # ds[v] = grid[v] + # Work around issue where xarray drops attributes on coordinates when a new + # DataArray is assigned to the Dataset, see + # https://github.com/pydata/xarray/issues/4415 + # https://github.com/pydata/xarray/issues/4393 + # This way adds as a 'Variable' instead of as a 'DataArray' + ds[v] = (grid[v].dims, grid[v].values) + + _add_attrs_to_var(ds, v) + return ds + + @register_geometry("toroidal") def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): @@ -311,24 +333,7 @@ def add_toroidal_geometry_coords(ds, *, coordinates=None, grid=None): ) # Get extra geometry information from grid file if it's not in the dump files - needed_variables = ["psixy", "Rxy", "Zxy"] - for v in needed_variables: - if v not in ds: - if grid is None: - raise ValueError( - f"Grid file is required to provide {v}. Pass the grid " - f"file name as the 'gridfilepath' argument to " - f"open_boutdataset()." - ) - # ds[v] = grid[v] - # Work around issue where xarray drops attributes on coordinates when a new - # DataArray is assigned to the Dataset, see - # https://github.com/pydata/xarray/issues/4415 - # https://github.com/pydata/xarray/issues/4393 - # This way adds as a 'Variable' instead of as a 'DataArray' - ds[v] = (grid[v].dims, grid[v].values) - - _add_attrs_to_var(ds, v) + ds = _add_vars_from_grid(ds, grid, ["psixy", "Rxy", "Zxy"]) if "t" in ds.dims: # Rename 't' if user requested it @@ -418,3 +423,11 @@ def add_s_alpha_geometry_coords(ds, *, coordinates=None, grid=None): del ds["hthe"] return ds + + +@register_geometry("fci") +def add_fci_geometry_coords(ds, *, coordinates=None, grid=None): + assert coordinates is None, "Not implemented" + ds = _add_vars_from_grid(ds, grid, ["R", "Z"]) + ds = ds.set_coords(("R", "Z")) + return ds From 41372bda7c2da33bb4b4f37c4914ca79660efd83 Mon Sep 17 00:00:00 2001 From: dschwoerer Date: Mon, 11 Oct 2021 15:34:56 +0200 Subject: [PATCH 12/18] Apply pep8 recommendation --- xbout/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/utils.py b/xbout/utils.py index de98d5eb..36c182a9 100644 --- a/xbout/utils.py +++ b/xbout/utils.py @@ -64,7 +64,7 @@ def _separate_metadata(ds): var for var in variables if not any(dim in ["t", "x", "y", "z"] for dim in ds[var].dims) - and not var in exclude + and var not in exclude ] # Save metadata as a dictionary From d6fbeaa0d441a425eda2c948073875def8cdf667 Mon Sep 17 00:00:00 2001 From: johnomotani Date: Tue, 19 Oct 2021 13:34:10 +0100 Subject: [PATCH 13/18] Create single region for FCI geometry --- xbout/geometries.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xbout/geometries.py b/xbout/geometries.py index 68d5052e..035d5592 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -430,4 +430,5 @@ def add_fci_geometry_coords(ds, *, coordinates=None, grid=None): assert coordinates is None, "Not implemented" ds = _add_vars_from_grid(ds, grid, ["R", "Z"]) ds = ds.set_coords(("R", "Z")) + ds = _create_single_region(ds, periodic_y=True) return ds From da5f2ace36e367f9e0248619cd1da6f59ec8446d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20Schw=C3=B6rer?= Date: Mon, 25 Oct 2021 11:19:15 +0200 Subject: [PATCH 14/18] Add missing import --- xbout/geometries.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index 035d5592..c1acabd0 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, From a72c8e952584e76b3b65f39192d8fc7d55eb0e82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20Schw=C3=B6rer?= Date: Mon, 25 Oct 2021 11:20:41 +0200 Subject: [PATCH 15/18] Add basic FCI test --- xbout/tests/test_load.py | 55 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 21207d0a..48cb7ff9 100644 --- a/xbout/tests/test_load.py +++ b/xbout/tests/test_load.py @@ -50,6 +50,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 +1444,49 @@ def test_trim_timing_info(self): expected = create_test_data(0) xrt.assert_equal(ds, expected) + + +@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 + arr = np.arange(2 * 3 * 4).reshape(2, 3, 4) + grid = DataArray(data=arr, name="R", dims=["x", "y", "z"]).to_dataset() + grid["Z"] = DataArray(np.random.random((2, 3, 4)), dims=["x", "y", "z"]) + grid["dy"] = DataArray(np.ones((2, 3, 4)), 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=(2, 2, 3, 4), + nxpe=1, + nype=1, + # nt=1, + # guards=None, + syn_data_type="random", + grid=None, + squashed=False, + # topology="core", + write_to_disk=False, + bout_v5=True, + metric_3D=True, + ) From 00e53888934d2a0be7073728e9d4fa2a49327d4c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20Schw=C3=B6rer?= Date: Mon, 25 Oct 2021 11:39:44 +0200 Subject: [PATCH 16/18] FCI shape for test semi-flexible Make it a bit easier to change the shape consistently for FCI. Also include guards. --- xbout/tests/test_load.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/xbout/tests/test_load.py b/xbout/tests/test_load.py index 48cb7ff9..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 @@ -1446,6 +1448,10 @@ def test_trim_timing_info(self): 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): """ @@ -1456,10 +1462,11 @@ def create_example_grid_file_fci(tmp_path_factory): """ # Create grid dataset - arr = np.arange(2 * 3 * 4).reshape(2, 3, 4) + 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((2, 3, 4)), dims=["x", "y", "z"]) - grid["dy"] = DataArray(np.ones((2, 3, 4)), dims=["x", "y", "z"]) + 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 @@ -1477,11 +1484,11 @@ def create_example_files_fci(tmp_path_factory): return _bout_xyt_example_files( tmp_path_factory, - lengths=(2, 2, 3, 4), + lengths=fci_shape, nxpe=1, nype=1, # nt=1, - # guards=None, + guards={a: b for a, b in zip("xyz", fci_guards)}, syn_data_type="random", grid=None, squashed=False, From e1b1e1c6e1abcac20266db83965dde49fb9af164 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20Schw=C3=B6rer?= Date: Mon, 25 Oct 2021 12:31:16 +0200 Subject: [PATCH 17/18] Fix "all" region for FCI --- xbout/geometries.py | 14 ++++++++++---- xbout/region.py | 4 ++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/xbout/geometries.py b/xbout/geometries.py index c1acabd0..375e677f 100644 --- a/xbout/geometries.py +++ b/xbout/geometries.py @@ -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, From 007d7bc4610412175cf5bc0e512e6a265690705d Mon Sep 17 00:00:00 2001 From: John Omotani Date: Fri, 24 Dec 2021 20:01:05 +0000 Subject: [PATCH 18/18] In tests of restart file saving, do not check rank-dependent variables ... they cannot be consistent with `check_ds` since each restart file has a different value. --- xbout/tests/test_boutdataset.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/xbout/tests/test_boutdataset.py b/xbout/tests/test_boutdataset.py index 0390f065..877a8d22 100644 --- a/xbout/tests/test_boutdataset.py +++ b/xbout/tests/test_boutdataset.py @@ -2162,6 +2162,9 @@ def test_to_restart(self, tmp_path_factory, bout_xyt_example_files, tind): # No coordinates saved in restart files, so unset them in check_ds check_ds = check_ds.reset_coords() + # ignore variables that depend on the rank of the BOUT++ process - they + # cannot be consistent with check_ds + rank_dependent_vars = ["PE_XIND", "PE_YIND", "MYPE"] for v in restart_ds: if v in check_ds: xrt.assert_equal(restart_ds[v], check_ds[v]) @@ -2170,7 +2173,7 @@ def test_to_restart(self, tmp_path_factory, bout_xyt_example_files, tind): assert restart_ds[v].values == -1 elif v == "tt": assert restart_ds[v].values == t_array - else: + elif v not in rank_dependent_vars: assert restart_ds[v].values == check_ds.metadata[v] def test_to_restart_change_npe(self, tmp_path_factory, bout_xyt_example_files): @@ -2224,6 +2227,9 @@ def test_to_restart_change_npe(self, tmp_path_factory, bout_xyt_example_files): # No coordinates saved in restart files, so unset them in check_ds check_ds = check_ds.reset_coords() + # ignore variables that depend on the rank of the BOUT++ process - they + # cannot be consistent with check_ds + rank_dependent_vars = ["PE_XIND", "PE_YIND", "MYPE"] for v in restart_ds: if v in check_ds: xrt.assert_equal(restart_ds[v], check_ds[v]) @@ -2234,7 +2240,7 @@ def test_to_restart_change_npe(self, tmp_path_factory, bout_xyt_example_files): assert restart_ds[v].values == -1 elif v == "tt": assert restart_ds[v].values == t_array - else: + elif v not in rank_dependent_vars: assert restart_ds[v].values == check_ds.metadata[v] @pytest.mark.long @@ -2292,6 +2298,9 @@ def test_to_restart_change_npe_doublenull( # No coordinates saved in restart files, so unset them in check_ds check_ds = check_ds.reset_coords() + # ignore variables that depend on the rank of the BOUT++ process - they + # cannot be consistent with check_ds + rank_dependent_vars = ["PE_XIND", "PE_YIND", "MYPE"] for v in restart_ds: if v in check_ds: xrt.assert_equal(restart_ds[v], check_ds[v]) @@ -2302,7 +2311,7 @@ def test_to_restart_change_npe_doublenull( assert restart_ds[v].values == -1 elif v == "tt": assert restart_ds[v].values == t_array - else: + elif v not in rank_dependent_vars: assert restart_ds[v].values == check_ds.metadata[v] @pytest.mark.long