def test_concat(self, tmpdir_factory, bout_xyt_example_files):
path1 = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=4, nt=1)
bd1 = open_boutdataset(datapath=path1, inputfilepath=None,
keep_xboundaries=False)
path2 = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=4, nt=1)
bd2 = open_boutdataset(datapath=path2, inputfilepath=None,
keep_xboundaries=False)
result = concat([bd1, bd2], dim='run')
assert result.dims == {**bd1.dims, 'run': 2}
def test_remove_yboundaries(
self, bout_xyt_example_files, mxg, myg, remove_extra_upper
):
dataset_list, grid_ds = bout_xyt_example_files(
None,
lengths=(2, 3, 4, 3),
nxpe=1,
nype=6,
nt=1,
grid="grid",
guards={"x": mxg, "y": myg},
topology="connected-double-null",
syn_data_type="linear",
)
ds = open_boutdataset(
datapath=dataset_list,
gridfilepath=grid_ds,
geometry="toroidal",
keep_yboundaries=True,
)
dataset_list_no_yboundaries, grid_ds_no_yboundaries = bout_xyt_example_files(
None,
lengths=(2, 3, 4, 3),
nxpe=1,
nype=6,
nt=1,
grid="grid",
guards={"x": mxg, "y": 0},
topology="connected-double-null",
syn_data_type="linear",
)
ds_no_yboundaries = open_boutdataset(
datapath=dataset_list_no_yboundaries,
gridfilepath=grid_ds_no_yboundaries,
geometry="toroidal",
keep_yboundaries=False,
)
if remove_extra_upper:
ds_no_yboundaries = xr.concat(
[
ds_no_yboundaries.isel(theta=slice(None, 11)),
ds_no_yboundaries.isel(theta=slice(12, -1)),
],
dim="theta",
)
n = ds["n"].bout.remove_yboundaries(remove_extra_upper=remove_extra_upper)
assert n.metadata["keep_yboundaries"] == 0
npt.assert_equal(n.values, ds_no_yboundaries["n"].values)
def test_getFieldAligned(self, tmpdir_factory, bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=4, nt=1)
ds = open_boutdataset(datapath=path, inputfilepath=None, keep_xboundaries=False)
ds['psixy'] = ds['x']
ds['Rxy'] = ds['x']
ds['Zxy'] = ds['y']
ds = apply_geometry(ds, 'toroidal')
n = ds['n']
n.attrs['direction_y'] = 'Standard'
n_aligned_from_array = n.bout.toFieldAligned()
# check n_aligned does not exist yet
try:
ds['n_aligned']
assert False
except KeyError:
pass
n_aligned_from_ds = ds.bout.getFieldAligned('n')
xrt.assert_allclose(n_aligned_from_ds, n_aligned_from_array)
xrt.assert_allclose(ds['n_aligned'], n_aligned_from_array)
# check getting the cached version
ds['n_aligned'] = ds['T']
xrt.assert_allclose(ds.bout.getFieldAligned('n'), ds['T'])
def test_single_file(self, tmp_path_factory):
# Create temp directory for files
test_dir = tmp_path_factory.mktemp("test_data")
# Generate some test data
generated_ds = create_bout_ds(syn_data_type="linear")
generated_ds.to_netcdf(test_dir.joinpath("BOUT.dmp.0.nc"))
var = "n"
expected = old_collect(var, path=test_dir, xguards=True, yguards=False)
# Test against new standard - open_boutdataset
with pytest.warns(UserWarning):
ds = open_boutdataset(test_dir.joinpath("BOUT.dmp.0.nc"))
actual = ds[var].values
assert expected.shape == actual.shape
npt.assert_equal(actual, expected)
# Test against backwards compatible collect function
actual = new_collect(var, path=test_dir, xguards=True, yguards=False)
assert expected.shape == actual.shape
npt.assert_equal(actual, expected)
def test_isel(self, tmpdir_factory, bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=1, nype=1, nt=1)
bd = open_boutdataset(datapath=path, inputfilepath=None,
keep_xboundaries=False)
actual = bd.isel(x=slice(None,None,2))
expected = bd.bout.data.isel(x=slice(None,None,2))
xrt.assert_equal(actual, expected)
def test_ddy(self, bout_xyt_example_files):
ny = 64
dataset_list, gridfilepath = bout_xyt_example_files(
None,
lengths=(2, 3, ny, 4),
nxpe=1,
nype=1,
grid="grid",
)
ds = open_boutdataset(
datapath=dataset_list, geometry="toroidal", gridfilepath=gridfilepath
)
n = ds["n"]
t = ds["t"].broadcast_like(n)
x = ds["x"].broadcast_like(n)
y = ds["theta"].broadcast_like(n)
z = ds["zeta"].broadcast_like(n)
n.values[:] = (np.sin(3.0 * y / ny) * (1.0 + t + x + z)).values
expected = 3.0 / ny * np.cos(3.0 * y / ny) * (1.0 + t + x + z)
npt.assert_allclose(
n.bout.ddy().isel(theta=slice(1, -1)).values,
expected.isel(theta=slice(1, -1)).values,
rtol=1.0e-2,
atol=1.0e-13,
)
def test_ddz(self, bout_xyt_example_files):
dataset_list = bout_xyt_example_files(
None,
lengths=(2, 3, 4, 64),
nxpe=1,
nype=1,
)
with pytest.warns(UserWarning):
ds = open_boutdataset(
datapath=dataset_list,
)
n = ds["n"]
t = ds["t"].broadcast_like(n)
x = ds["x"].broadcast_like(n)
y = ds["y"].broadcast_like(n)
z = ds["z"].broadcast_like(n)
n.values[:] = (np.sin(z) * (1.0 + t + x + y)).values
expected = np.cos(z) * (1.0 + t + x + y)
npt.assert_allclose(
n.bout.ddz().values, expected.values, rtol=1.0e-2, atol=1.0e-13
)
def test_ddx(self, bout_xyt_example_files):
nx = 64
dataset_list = bout_xyt_example_files(
None,
lengths=(2, nx, 4, 3),
nxpe=1,
nype=1,
)
with pytest.warns(UserWarning):
ds = open_boutdataset(
datapath=dataset_list,
)
n = ds["n"]
t = ds["t"].broadcast_like(n)
ds["x_1d"] = _1d_coord_from_spacing(ds["dx"], "x")
x = ds["x_1d"].broadcast_like(n)
y = ds["y"].broadcast_like(n)
z = ds["z"].broadcast_like(n)
n.values[:] = (np.sin(12.0 * x / nx) * (1.0 + t + y + z)).values
expected = 12.0 / nx * np.cos(12.0 * x / nx) * (1.0 + t + y + z)
npt.assert_allclose(
n.bout.ddx().isel(x=slice(1, -1)).values,
expected.isel(x=slice(1, -1)).values,
rtol=1.0e-2,
atol=1.0e-13,
)
def test_interpolate_parallel_toroidal_points_list(self, bout_xyt_example_files):
dataset_list, grid_ds = bout_xyt_example_files(
None,
lengths=(2, 3, 16, 3),
nxpe=1,
nype=3,
nt=1,
grid="grid",
guards={"y": 2},
topology="single-null",
)
ds = open_boutdataset(
datapath=dataset_list,
gridfilepath=grid_ds,
geometry="toroidal",
keep_yboundaries=True,
)
n_highres = ds["n"].bout.interpolate_parallel()
points_list = [1, 2]
n_highres_truncated = ds["n"].bout.interpolate_parallel(
toroidal_points=points_list
)
xrt.assert_identical(n_highres_truncated, n_highres.isel(zeta=points_list))
def test_multiple_files_along_y(self, tmp_path_factory):
# Create temp directory for files
test_dir = tmp_path_factory.mktemp("test_data")
# Generate some test data
ds_list, file_list = create_bout_ds_list("BOUT.dmp",
nxpe=1,
nype=3,
syn_data_type="linear")
for temp_ds, file_name in zip(ds_list, file_list):
temp_ds.to_netcdf(test_dir.joinpath(file_name))
var = "n"
expected = old_collect(var,
path=test_dir,
prefix="BOUT.dmp",
xguards=True)
# Test against new standard - .open_boutdataset
with pytest.warns(UserWarning):
ds = open_boutdataset(test_dir.joinpath("BOUT.dmp.*.nc"))
actual = ds[var].values
assert expected.shape == actual.shape
npt.assert_equal(actual, expected)
# Test against backwards compatible collect function
actual = new_collect(var,
path=test_dir,
prefix="BOUT.dmp",
xguards=True)
assert expected.shape == actual.shape
npt.assert_equal(actual, expected)
def test_storm_dataset_inheritance(self, tmpdir_factory,
bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=1, nype=1, nt=1)
ds = open_boutdataset(datapath=path, inputfilepath=None)
ds.storm.set_extra_data('options')
print(ds.storm.extra_data)
assert False
def test_object_permanence(self, tmpdir_factory, bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=1, nype=1, nt=1)
ds = open_boutdataset(datapath=path, inputfilepath=None)
ds.storm.extra_info = 'options'
new_ds = ds.isel(t=-1)
print(new_ds.storm.extra_info)
assert False
def test_to_dataset(self, tmpdir_factory, bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=4, nt=1)
ds = open_boutdataset(datapath=path,
inputfilepath=None,
keep_xboundaries=False)
da = ds['n']
new_ds = da.bout.to_dataset()
assert dict_equiv(ds.attrs, new_ds.attrs)
assert dict_equiv(ds.metadata, new_ds.metadata)
def test_to_dataset(self, bout_xyt_example_files):
dataset_list = bout_xyt_example_files(None, nxpe=3, nype=4, nt=1)
with pytest.warns(UserWarning):
ds = open_boutdataset(
datapath=dataset_list, inputfilepath=None, keep_xboundaries=False
)
da = ds["n"]
new_ds = da.bout.to_dataset()
assert dict_equiv(ds.attrs, new_ds.attrs)
assert dict_equiv(ds.metadata, new_ds.metadata)
def test_interpolate_parallel(self, bout_xyt_example_files):
dataset_list, grid_ds = bout_xyt_example_files(
None,
lengths=(2, 3, 16, 3),
nxpe=1,
nype=3,
nt=1,
grid="grid",
guards={"y": 2},
topology="single-null",
)
ds = open_boutdataset(
datapath=dataset_list,
gridfilepath=grid_ds,
geometry="toroidal",
keep_yboundaries=True,
)
n = ds["n"]
thetalength = 2.0 * np.pi
dtheta = thetalength / 48.0
theta = xr.DataArray(
np.linspace(0.0 - 1.5 * dtheta, thetalength + 1.5 * dtheta, 52),
dims="theta",
)
dtheta_fine = thetalength / 3.0 / 128.0
theta_fine = xr.DataArray(
np.linspace(
0.0 + 0.5 * dtheta_fine, thetalength - 0.5 * dtheta_fine, 3 * 128
),
dims="theta",
)
x = xr.DataArray(np.arange(3), dims="x")
def f_y(t):
t = np.sin(3.0 * t)
return t ** 3 - t ** 2 + t - 1.0
f = f_y(theta) * (x + 1.0)
n.data = f.broadcast_like(n)
f_fine = f_y(theta_fine) * (x + 1.0)
n_highres = n.bout.interpolate_parallel().isel(theta=slice(2, -2))
expected = f_fine.broadcast_like(n_highres)
npt.assert_allclose(n_highres.values, expected.values, rtol=0.0, atol=1.1e-2)
def test_interpolate_parallel_region_core_change_n(
self, bout_xyt_example_files, res_factor
):
dataset_list, grid_ds = bout_xyt_example_files(
None,
lengths=(2, 3, 16, 3),
nxpe=1,
nype=1,
nt=1,
grid="grid",
guards={"y": 2},
topology="core",
)
ds = open_boutdataset(
datapath=dataset_list,
gridfilepath=grid_ds,
geometry="toroidal",
keep_yboundaries=True,
)
n = ds["n"]
thetalength = 2.0 * np.pi
dtheta = thetalength / 16.0
theta = xr.DataArray(
np.linspace(0.0 - 1.5 * dtheta, thetalength + 1.5 * dtheta, 20),
dims="theta",
)
dtheta_fine = thetalength / res_factor / 16.0
theta_fine = xr.DataArray(
np.linspace(
0.0 + dtheta_fine / 2.0,
thetalength - dtheta_fine / 2.0,
res_factor * 16,
),
dims="theta",
)
def f(t):
t = np.sin(t)
return t ** 3 - t ** 2 + t - 1.0
n.data = f(theta).broadcast_like(n)
n_highres = n.bout.interpolate_parallel("core", n=res_factor)
expected = f(theta_fine).broadcast_like(n_highres)
npt.assert_allclose(n_highres.values, expected.values, rtol=0.0, atol=1.0e-2)
def test_from_field_aligned_staggered(self, bout_xyt_example_files, stag_location):
dataset_list = bout_xyt_example_files(
None, lengths=(3, 3, 4, 8), nxpe=1, nype=1, nt=1
)
with pytest.warns(UserWarning):
ds = open_boutdataset(
datapath=dataset_list, inputfilepath=None, keep_xboundaries=False
)
ds["psixy"] = ds["x"]
ds["Rxy"] = ds["x"]
ds["Zxy"] = ds["y"]
ds = apply_geometry(ds, "toroidal")
# set up test variable
n = ds["n"].load()
zShift = ds["zShift"].load()
for t in range(ds.sizes["t"]):
for x in range(ds.sizes["x"]):
for y in range(ds.sizes["theta"]):
zShift[x, y] = (
(x * ds.sizes["theta"] + y) * 2.0 * np.pi / ds.sizes["zeta"]
)
for z in range(ds.sizes["zeta"]):
n[t, x, y, z] = 1000.0 * t + 100.0 * x + 10.0 * y + z
n.attrs["direction_y"] = "Aligned"
ds["T"].attrs["direction_y"] = "Aligned"
n_nal = n.bout.from_field_aligned().copy(deep=True)
assert n_nal.direction_y == "Standard"
# make 'n' staggered
ds["n"].attrs["cell_location"] = stag_location
if stag_location != "CELL_ZLOW":
with pytest.raises(ValueError):
# Check exception raised when needed zShift_CELL_*LOW is not present
ds["n"].bout.from_field_aligned()
ds["zShift_" + stag_location] = zShift
ds["zShift_" + stag_location].attrs["cell_location"] = stag_location
ds = ds.set_coords("zShift_" + stag_location)
ds = ds.drop_vars("zShift")
with pytest.raises(ValueError):
# Check shifting non-staggered field fails without zShift
ds["T"].bout.from_field_aligned()
n_stag_nal = ds["n"].bout.from_field_aligned()
npt.assert_equal(n_stag_nal.values, n_nal.values)
def create_test_file(tmpdir_factory):
# Create temp dir for output of animate1D/2D
save_dir = tmpdir_factory.mktemp("test_data")
# Generate some test data
ds_list, file_list = create_bout_ds_list("BOUT.dmp", nxpe=3, nype=3,
syn_data_type="linear")
for ds, file_name in zip(ds_list, file_list):
ds.to_netcdf(str(save_dir.join(str(file_name))))
ds = open_boutdataset(save_dir.join("BOUT.dmp.*.nc")) # Open test data
return save_dir, ds
def test_test_method(self, tmpdir_factory, bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=1, nype=1, nt=1)
ds = open_boutdataset(datapath=path, inputfilepath=None)
#ds = collect(path=path)
#bd = BoutAccessor(ds)
print(ds)
#ds.bout.test_method()
#print(ds.bout.options)
#print(ds.bout.metadata)
print(ds.isel(t=-1))
#ds.bout.set_extra_data('stored')
ds.bout.extra_data = 'stored'
print(ds.bout.extra_data)
def test_save_dtype(self, tmpdir_factory, bout_xyt_example_files, save_dtype):
# Create data
path = bout_xyt_example_files(tmpdir_factory, nxpe=1, nype=1, nt=1)
# Load it as a boutdataset
original = open_boutdataset(datapath=path, inputfilepath=None)
# Save it to a netCDF file
savepath = str(Path(path).parent) + 'temp_boutdata.nc'
original.bout.save(savepath=savepath, save_dtype=np.dtype(save_dtype))
# Load it again using bare xarray
recovered = open_dataset(savepath)
assert recovered['n'].values.dtype == np.dtype(save_dtype)
def test_save_all(self, tmpdir_factory, bout_xyt_example_files):
# Create data
path = bout_xyt_example_files(tmpdir_factory, nxpe=4, nype=5, nt=1)
# Load it as a boutdataset
original = open_boutdataset(datapath=path, inputfilepath=None)
# Save it to a netCDF file
savepath = str(Path(path).parent) + 'temp_boutdata.nc'
original.bout.save(savepath=savepath)
# Load it again using bare xarray
recovered = open_dataset(savepath)
# Compare equal (not identical because attributes are changed when saving)
xrt.assert_equal(original, recovered)
def collectData(self, quant, simIndex=0, simType=''):
os.chdir('{}/{}/{}'.format(self.outDir, simIndex, simType))
if os.path.isfile('BOUT.dmp.nc') is False:
print('{}-{} not squashed. squashing now'.format(simIndex, simType))
squashoutput(quiet=True)
gridFile = fnmatch.filter(next(os.walk('./'))[2], '*profile*')[0]
print(gridFile)
ds = open_boutdataset(
'BOUT.dmp.nc',
gridfilepath=gridFile,
coordinates={'x': 'psi_pol', 'y': 'theta', 'z': 'zeta'},
geometry='toroidal')
try:
quant2 = ds[quant]
except(KeyError):
quant2 = ds['t_array'].metadata[quant]
return np.squeeze(quant2)
def test_save_separate_variables(self, tmpdir_factory, bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=4, nype=1, nt=1)
# Load it as a boutdataset
original = open_boutdataset(datapath=path, inputfilepath=None)
# Save it to a netCDF file
savepath = str(Path(path).parent) + '/temp_boutdata.nc'
original.bout.save(savepath=savepath, separate_vars=True)
for var in ['n', 'T']:
# Load it again using bare xarray
savepath = str(Path(path).parent) + '/temp_boutdata_' + var + '.nc'
recovered = open_dataset(savepath)
# Compare equal (not identical because attributes are changed when saving)
xrt.assert_equal(recovered[var], original[var])
def test_region_limiter(
self,
bout_xyt_example_files,
guards,
keep_xboundaries,
keep_yboundaries,
with_guards,
):
# Note using more than MXG x-direction points and MYG y-direction points per
# output file ensures tests for whether boundary cells are present do not fail
# when using minimal numbers of processors
dataset_list, grid_ds = bout_xyt_example_files(
None,
lengths=(2, 5, 4, 3),
nxpe=1,
nype=1,
nt=1,
guards=guards,
grid="grid",
topology="limiter",
)
ds = open_boutdataset(
datapath=dataset_list,
gridfilepath=grid_ds,
geometry="toroidal",
keep_xboundaries=keep_xboundaries,
keep_yboundaries=keep_yboundaries,
)
ds["R"] = ds["R"].copy(deep=True)
ds["R"].data[:, :] = np.linspace(0.0, 1.0, ds.sizes["x"])[:, np.newaxis]
ds["Z"] = ds["Z"].copy(deep=True)
ds["Z"].data[:, :] = np.linspace(0.0, 1.0, ds.sizes["theta"])[np.newaxis, :]
n = ds["n"].isel(t=-1, zeta=0)
n.bout.contour()
plt.close()
n.bout.contourf()
plt.close()
n.bout.pcolormesh()
plt.close()
def test_save_all(self, tmpdir_factory, bout_xyt_example_files, options):
# Create data
path = bout_xyt_example_files(tmpdir_factory, nxpe=4, nype=5, nt=1)
# Load it as a boutdataset
original = open_boutdataset(datapath=path, inputfilepath=None)
if not options:
original.attrs['options'] = {}
# Save it to a netCDF file
savepath = str(Path(path).parent) + 'temp_boutdata.nc'
original.bout.save(savepath=savepath)
# Load it again using bare xarray
recovered = open_dataset(savepath)
# Compare
xrt.assert_equal(original, recovered)
def test_metadata(self, tmpdir_factory):
# Create temp directory for files
test_dir = tmpdir_factory.mktemp("test_data")
# Generate some test data
generated_ds = create_bout_ds(syn_data_type="linear")
generated_ds.to_netcdf(str(test_dir.join("BOUT.dmp.0.nc")))
ds = open_boutdataset(test_dir.join('BOUT.dmp.*.nc'))
for v in METADATA_VARS:
expected = old_collect(v, path=test_dir)
# Check metadata against new standard - open_boutdataset
actual = ds.bout.metadata[v]
npt.assert_equal(actual, expected)
# Check against backwards compatible collect function
actual = new_collect(v, path=test_dir)
npt.assert_equal(actual, expected)
def test_reload_separate_variables(
self, tmpdir_factory, bout_xyt_example_files, geometry
):
if geometry is not None:
grid = "grid"
else:
grid = None
path = bout_xyt_example_files(tmpdir_factory, nxpe=4, nype=1, nt=1, grid=grid)
if grid is not None:
gridpath = str(Path(path).parent) + "/grid.nc"
else:
gridpath = None
# Load it as a boutdataset
original = open_boutdataset(
datapath=path,
inputfilepath=None,
geometry=geometry,
gridfilepath=gridpath,
)
# Save it to a netCDF file
savepath = str(Path(path).parent) + '/temp_boutdata.nc'
original.bout.save(savepath=savepath, separate_vars=True)
# Load it again
savepath = str(Path(path).parent) + '/temp_boutdata_*.nc'
recovered = reload_boutdataset(savepath, pre_squashed=True)
# Compare
for coord in original.coords.values():
# Get rid of the options if they exist, because options are not dealt with
# totally consistently: they exist if a coord was created from a variable
# loaded from the BOUT++ output, but not if the coord was calculated from
# some parameters or loaded from a grid file
try:
del coord.attrs["options"]
except KeyError:
pass
xrt.assert_identical(recovered, original)
def test_derivatives_doublenull(self, bout_xyt_example_files):
# Check function does not error on double-null topology
dataset_list, grid_ds = bout_xyt_example_files(
None,
lengths=(2, 3, 4, 3),
nxpe=1,
nype=6,
nt=1,
grid="grid",
guards={"x": 2, "y": 2},
topology="connected-double-null",
)
ds = open_boutdataset(
datapath=dataset_list,
gridfilepath=grid_ds,
geometry="toroidal",
keep_yboundaries=True,
)
test_ddx = ds["n"].bout.ddx()
test_ddy = ds["n"].bout.ddy()
test_ddz = ds["n"].bout.ddz()
from matplotlib import pyplot as plt from xbout import open_boutdataset # We do not distribute binary files with xBOUT, so you need to provide your own gridfile, # for example one created using Hypnotoad. gridfilepath = "grid.nc" grid = open_boutdataset(gridfilepath, geometry="toroidal") grid["psi_poloidal"].bout.contourf() grid["psi_poloidal"].bout.contour() grid["psi_poloidal"].bout.pcolormesh() grid["psi_poloidal"].bout.pcolormesh(shading="gouraud") grid["psi_poloidal"].bout.regions() plt.show()
def test_dataset_duck_typing(self, tmpdir_factory, bout_xyt_example_files):
path = bout_xyt_example_files(tmpdir_factory, nxpe=1, nype=1, nt=1)
ds = open_boutdataset(datapath=path, inputfilepath=None)
result = concat([ds.bout, ds.bout], dim='run')
print(result)