From c5c8719718b6cdc45c61e8e3a203faafdec5828d Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 22 Dec 2021 17:31:21 -0800 Subject: [PATCH 1/9] Add restart.change_grid function Only works for nz=1 for now, as testing for axisymmetric simulations. Asserts that nz=1, and could be extended in future, perhaps by interpolating one Z slice at a time. --- boutdata/restart.py | 207 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 207 insertions(+) diff --git a/boutdata/restart.py b/boutdata/restart.py index 7729bf83e..be1aa42d0 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -956,3 +956,210 @@ def addvar(var, value, path="."): # Set the variable in the NetCDF file df.write(var, data) + + +def change_grid( + from_grid_file, + to_grid_file, + path="data", + output=".", + interpolator="nearest", + show=False, +): + """ + Convert a set of restart files from one grid to another + + Notes: + - Only working for 2D (axisymmetric) simulations with nz = 1 + + from_grid_file : str + File containing the input grid + to_grid_file : str + File containing the output grid + path : str, optional + Directory containing input restart files + output : str, optional + Directory where output restart files will be written + interpolator : str, optional + Interpolation method to use. Options are 'nearest', 'CloughTocher', 'RBF' + show : bool, optional + Display the interpolated fields using Matplotlib + + """ + + # Read in grid files + with DataFile(from_grid_file) as g: + from_Rxy = g["Rxy"] + from_Zxy = g["Zxy"] + + with DataFile(to_grid_file) as g: + to_Rxy = g["Rxy"] + to_Zxy = g["Zxy"] + + file_list = glob.glob(os.path.join(path, "BOUT.restart.*.nc")) + if len(file_list) == 0: + raise ValueError("ERROR: No restart files found") + + copy_vars = [ + "NXPE", + "NYPE", + "hist_hi", + "tt", + "MXG", + "MYG", + "MZG", + "nz", + "run_id", + "run_restart_from", + ] + copy_data = {} + interp_vars = [] + + # Read information from a restart file + with DataFile(file_list[0]) as f: + for var in copy_vars: + copy_data[var] = f[var] + + # Get a list of variables + varnames = f.list() + + for var in varnames: + dimensions = f.dimensions(var) + if dimensions == ("x", "y", "z"): + # Could be an evolving variable [x,y,z]x + interp_vars.append(var) + + # Only tested for nz = 1 + assert copy_data["nz"] == 1 + + def extrapolate_yguards(data2d, myg): + nx, ny = data2d.shape + result = np.zeros((nx, ny + 2 * myg)) + print(result.shape) + result[:, myg:-myg] = data2d + + dy = result[:, myg + 1] - result[:, myg] + for i in range(1, myg + 1): + result[:, myg - i] = result[:, myg] - i * dy + dy = result[:, -myg - 1] - result[:, -myg - 2] + for i in range(1, myg + 1): + result[:, -myg - 1 + i] = result[:, -myg - 1] + i * dy + return result + + from_Rxy = extrapolate_yguards(from_Rxy, copy_data["MYG"]) + from_Zxy = extrapolate_yguards(from_Zxy, copy_data["MYG"]) + to_Rxy = extrapolate_yguards(to_Rxy, copy_data["MYG"]) + to_Zxy = extrapolate_yguards(to_Zxy, copy_data["MYG"]) + + interp_data = {} + for var in interp_vars: + print("Interpolating " + var) + + from_data = collect( + var, + path=path, + xguards=True, + yguards=True, + prefix="BOUT.restart", + info=False, + ) + + if interpolator == "CloughTocher": + # Triangulate + from scipy.interpolate import CloughTocher2DInterpolator + + interp = CloughTocher2DInterpolator( + list(zip(from_Rxy.flatten(), from_Zxy.flatten())), from_data.flatten() + ) + elif interpolator == "RBF": + # Radial Basis Functions + from scipy.interpolate import RBFInterpolator + + interp = RBFInterpolator( + list(zip(from_Rxy.flatten(), from_Zxy.flatten())), + from_data.flatten(), + neighbors=50, + ) + + elif interpolator == "nearest": + # Nearest neighbour. Tends to be robust + from scipy.interpolate import NearestNDInterpolator + + interp = NearestNDInterpolator( + list(zip(from_Rxy.flatten(), from_Zxy.flatten())), from_data.flatten() + ) + else: + raise ValueError("Invalid interpolator") + + to_data = interp(list(zip(to_Rxy.flatten(), to_Zxy.flatten()))).reshape( + to_Rxy.shape + ) + + print( + "\tData ranges: {}:{} -> {}:{}".format( + np.amin(from_data), + np.amax(from_data), + np.amin(to_data), + np.amax(to_data), + ) + ) + if show: + import matplotlib.pyplot as plt + + plt.pcolormesh(to_Rxy, to_Zxy, to_data, shading="auto") + plt.plot(from_Rxy, from_Zxy, "ok") + plt.colorbar() + plt.axis("equal") + plt.show() + + interp_data[var] = to_data + + # Now have copy_data and interp_data dictionaries to write to the + # new restart files. Now need to partition the interpolated arrays + # with similar logic to redistribute() + + nxpe = copy_data["NXPE"] + nype = copy_data["NYPE"] + npes = nxpe * nype + + mxg = copy_data["MXG"] + myg = copy_data["MYG"] + + new_nx, new_ny = to_Rxy.shape + + if (new_nx - 2 * mxg) % nxpe != 0: + # Can't split grid in this way + raise ValueError("nxpe={} not compatible with nx = {}".format(nxpe, new_nx)) + if (new_ny - 2 * myg) % nxpe != 0: + # Can't split grid in this waymxsub = (new_nx - 2*mxg) // nxpe + raise ValueError("nype={} not compatible with ny = {}".format(nype, new_ny)) + + mxsub = (new_nx - 2 * mxg) // nxpe + mysub = (new_ny - 2 * myg) // nype + + copy_data["MXSUB"] = mxsub + copy_data["MYSUB"] = mysub + copy_data["nx"] = new_nx + copy_data["ny"] = new_ny - 2 * myg + + for i in range(npes): + ix = i % nxpe + iy = int(i / nxpe) + + def get_block(data): + return data[ + ix * mxsub : (ix + 1) * mxsub + 2 * mxg, + iy * mysub : (iy + 1) * mysub + 2 * myg, + ] + + outpath = os.path.join(output, "BOUT.restart." + str(i) + ".nc") + with DataFile(outpath, create=True) as f: + print("Creating " + outpath) + + # Write the scalars + for k in copy_data: + f.write(k, copy_data[k]) + + # Write fields + for k in interp_data: + f.write(k, get_block(interp_data[k])) From 14ed79c2814b1cd502a6a5ddaae01bce846091ea Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 23 Dec 2021 16:43:01 -0800 Subject: [PATCH 2/9] Add PE_XIND and PE_YIND to output restarts Apparently needed by xBOUT --- boutdata/restart.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/boutdata/restart.py b/boutdata/restart.py index be1aa42d0..327fdd559 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -1156,6 +1156,9 @@ def get_block(data): with DataFile(outpath, create=True) as f: print("Creating " + outpath) + f.write("PE_XIND", ix) + f.write("PE_YIND", iy) + # Write the scalars for k in copy_data: f.write(k, copy_data[k]) From 6b49fa9babfa14b104a60a83684a48e3e6e5c67b Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Thu, 23 Dec 2021 16:45:18 -0800 Subject: [PATCH 3/9] Small tidy of accidental paste --- boutdata/restart.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/boutdata/restart.py b/boutdata/restart.py index 327fdd559..26d03b4ac 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -1131,7 +1131,7 @@ def extrapolate_yguards(data2d, myg): # Can't split grid in this way raise ValueError("nxpe={} not compatible with nx = {}".format(nxpe, new_nx)) if (new_ny - 2 * myg) % nxpe != 0: - # Can't split grid in this waymxsub = (new_nx - 2*mxg) // nxpe + # Can't split grid in this way raise ValueError("nype={} not compatible with ny = {}".format(nype, new_ny)) mxsub = (new_nx - 2 * mxg) // nxpe From e1a94ced3bd9937c10c54a934524444ca63289e0 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 5 Jan 2022 11:38:01 -0800 Subject: [PATCH 4/9] Remove stray 'x' Co-authored-by: johnomotani --- boutdata/restart.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/boutdata/restart.py b/boutdata/restart.py index 26d03b4ac..71e6e4964 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -1026,7 +1026,7 @@ def change_grid( for var in varnames: dimensions = f.dimensions(var) if dimensions == ("x", "y", "z"): - # Could be an evolving variable [x,y,z]x + # Could be an evolving variable [x,y,z] interp_vars.append(var) # Only tested for nz = 1 From a8ab74bbf4270eb638d37064da1b4b7308854236 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 5 Jan 2022 11:38:28 -0800 Subject: [PATCH 5/9] Document missing support for Field2D & FieldPerp Co-authored-by: johnomotani --- boutdata/restart.py | 1 + 1 file changed, 1 insertion(+) diff --git a/boutdata/restart.py b/boutdata/restart.py index 71e6e4964..905f9879b 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -971,6 +971,7 @@ def change_grid( Notes: - Only working for 2D (axisymmetric) simulations with nz = 1 + - Does not support evolving Field2D or FieldPerp variables from_grid_file : str File containing the input grid From 187342714ad3bb84bfad2a51f280a28d83be7023 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 5 Jan 2022 11:39:10 -0800 Subject: [PATCH 6/9] Use integer division rather than conversion Co-authored-by: johnomotani --- boutdata/restart.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/boutdata/restart.py b/boutdata/restart.py index 905f9879b..0ff5ffdea 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -1145,7 +1145,7 @@ def extrapolate_yguards(data2d, myg): for i in range(npes): ix = i % nxpe - iy = int(i / nxpe) + iy = i // nxpe def get_block(data): return data[ From 7189c57a789748fe793eba5896f497fac17db69b Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 5 Jan 2022 13:53:14 -0800 Subject: [PATCH 7/9] Check for y_boundary_guards in grid files Not currently supported, so raise error if encountered. --- boutdata/restart.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/boutdata/restart.py b/boutdata/restart.py index 0ff5ffdea..35b1b463f 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -972,6 +972,7 @@ def change_grid( Notes: - Only working for 2D (axisymmetric) simulations with nz = 1 - Does not support evolving Field2D or FieldPerp variables + - Does not support grids with y boundary cells from_grid_file : str File containing the input grid @@ -990,10 +991,22 @@ def change_grid( # Read in grid files with DataFile(from_grid_file) as g: + # Check for y boundary cells + try: + if g["y_boundary_guards"] != 0: + raise ValueError("Support for grid files with y-boundary cells not implemented yet") + except KeyError: + pass # No y_boundary_guards key from_Rxy = g["Rxy"] from_Zxy = g["Zxy"] with DataFile(to_grid_file) as g: + # Check for y boundary cells + try: + if g["y_boundary_guards"] != 0: + raise ValueError("Support for grid files with y-boundary cells not implemented yet") + except KeyError: + pass # No y_boundary_guards key to_Rxy = g["Rxy"] to_Zxy = g["Zxy"] From 06c325d7a3007fc55be90690203e5ba6927a6e22 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 5 Jan 2022 13:54:01 -0800 Subject: [PATCH 8/9] Add outputs in form which can be collected Ensure that 3D inputs are 3D in the output, and that the BOUT_VERSION and MZ variables are in the outputs. This allows the output to be collected, and used as input to another interpolation if needed. --- boutdata/restart.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/boutdata/restart.py b/boutdata/restart.py index 35b1b463f..16b76f948 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -1015,6 +1015,7 @@ def change_grid( raise ValueError("ERROR: No restart files found") copy_vars = [ + "BOUT_VERSION", "NXPE", "NYPE", "hist_hi", @@ -1023,6 +1024,7 @@ def change_grid( "MYG", "MZG", "nz", + "MZ", "run_id", "run_restart_from", ] @@ -1161,10 +1163,11 @@ def extrapolate_yguards(data2d, myg): iy = i // nxpe def get_block(data): - return data[ + sliced = data[ ix * mxsub : (ix + 1) * mxsub + 2 * mxg, iy * mysub : (iy + 1) * mysub + 2 * myg, ] + return sliced.reshape(sliced.shape + (1,)) # make 3D outpath = os.path.join(output, "BOUT.restart." + str(i) + ".nc") with DataFile(outpath, create=True) as f: From c5be49206fc417d6fffcbc90429a774be90597b9 Mon Sep 17 00:00:00 2001 From: Ben Dudson Date: Wed, 5 Jan 2022 13:57:19 -0800 Subject: [PATCH 9/9] Apply Black formatting --- boutdata/restart.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/boutdata/restart.py b/boutdata/restart.py index 16b76f948..791bc07a4 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -994,9 +994,11 @@ def change_grid( # Check for y boundary cells try: if g["y_boundary_guards"] != 0: - raise ValueError("Support for grid files with y-boundary cells not implemented yet") + raise ValueError( + "Support for grid files with y-boundary cells not implemented yet" + ) except KeyError: - pass # No y_boundary_guards key + pass # No y_boundary_guards key from_Rxy = g["Rxy"] from_Zxy = g["Zxy"] @@ -1004,9 +1006,11 @@ def change_grid( # Check for y boundary cells try: if g["y_boundary_guards"] != 0: - raise ValueError("Support for grid files with y-boundary cells not implemented yet") + raise ValueError( + "Support for grid files with y-boundary cells not implemented yet" + ) except KeyError: - pass # No y_boundary_guards key + pass # No y_boundary_guards key to_Rxy = g["Rxy"] to_Zxy = g["Zxy"] @@ -1167,7 +1171,7 @@ def get_block(data): ix * mxsub : (ix + 1) * mxsub + 2 * mxg, iy * mysub : (iy + 1) * mysub + 2 * myg, ] - return sliced.reshape(sliced.shape + (1,)) # make 3D + return sliced.reshape(sliced.shape + (1,)) # make 3D outpath = os.path.join(output, "BOUT.restart." + str(i) + ".nc") with DataFile(outpath, create=True) as f: