diff --git a/boutdata/restart.py b/boutdata/restart.py index 7729bf83e..791bc07a4 100644 --- a/boutdata/restart.py +++ b/boutdata/restart.py @@ -956,3 +956,234 @@ 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 + - 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 + 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: + # 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"] + + 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 = [ + "BOUT_VERSION", + "NXPE", + "NYPE", + "hist_hi", + "tt", + "MXG", + "MYG", + "MZG", + "nz", + "MZ", + "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] + 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 way + 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 = i // nxpe + + def get_block(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: + 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]) + + # Write fields + for k in interp_data: + f.write(k, get_block(interp_data[k]))