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)
path2 = bout_xyt_example_files(tmpdir_factory, nxpe=3, nype=4, nt=1)
bd2 = open_boutdataset(datapath=path2, inputfilepath=None)
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_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_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_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_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_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_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_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 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_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_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 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_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_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()
def test_storm_dataset(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)
print(ds.storm.normalisation)
assert False
def test_load_options_in_dataset(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=EXAMPLE_OPTIONS_FILE_PATH)
assert isinstance(ds.options, BoutOptions)
def test_to_field_aligned_dask(self, bout_xyt_example_files):
nz = 6
dataset_list = bout_xyt_example_files(
None, lengths=(3, 3, 4, nz), 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(nz):
n[t, x, y, z] = 1000.0 * t + 100.0 * x + 10.0 * y + z
# The above loop required the call to .load(), but that turned the data into a
# numpy array. Now convert back to dask
n = n.chunk({"t": 1})
assert isinstance(n.data, dask.array.Array)
n.attrs["direction_y"] = "Standard"
n_al = n.bout.to_field_aligned()
assert n_al.direction_y == "Aligned"
for t in range(ds.sizes["t"]):
for z in range(nz):
npt.assert_allclose(
n_al[t, 0, 0, z].values,
1000.0 * t + z % nz,
rtol=1.0e-15,
atol=5.0e-16,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_al[t, 0, 1, z].values,
1000.0 * t + 10.0 * 1.0 + (z + 1) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_al[t, 0, 2, z].values,
1000.0 * t + 10.0 * 2.0 + (z + 2) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_al[t, 0, 3, z].values,
1000.0 * t + 10.0 * 3.0 + (z + 3) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_al[t, 1, 0, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 0.0 + (z + 4) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_al[t, 1, 1, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 1.0 + (z + 5) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_al[t, 1, 2, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 2.0 + (z + 6) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_al[t, 1, 3, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 3.0 + (z + 7) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
def test_from_field_aligned(self, bout_xyt_example_files, nz):
dataset_list = bout_xyt_example_files(
None, lengths=(3, 3, 4, nz), 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"
n_nal = n.bout.from_field_aligned()
assert n_nal.direction_y == "Standard"
for t in range(ds.sizes["t"]):
for z in range(nz):
npt.assert_allclose(
n_nal[t, 0, 0, z].values,
1000.0 * t + z % nz,
rtol=1.0e-15,
atol=5.0e-16,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_nal[t, 0, 1, z].values,
1000.0 * t + 10.0 * 1.0 + (z - 1) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_nal[t, 0, 2, z].values,
1000.0 * t + 10.0 * 2.0 + (z - 2) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_nal[t, 0, 3, z].values,
1000.0 * t + 10.0 * 3.0 + (z - 3) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_nal[t, 1, 0, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 0.0 + (z - 4) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_nal[t, 1, 1, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 1.0 + (z - 5) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_nal[t, 1, 2, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 2.0 + (z - 6) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
for z in range(nz):
npt.assert_allclose(
n_nal[t, 1, 3, z].values,
1000.0 * t + 100.0 * 1 + 10.0 * 3.0 + (z - 7) % nz,
rtol=1.0e-15,
atol=0.0,
) # noqa: E501
def test_toFieldAligned(self, tmpdir_factory, bout_xyt_example_files, nz):
path = bout_xyt_example_files(tmpdir_factory,
lengths=(3, 3, 4, nz),
nxpe=1,
nype=1,
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')
# 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. * np.pi / ds.sizes['zeta']
for z in range(nz):
n[t, x, y, z] = 1000. * t + 100. * x + 10. * y + z
n.attrs['direction_y'] = 'Standard'
n_al = n.bout.toFieldAligned()
for t in range(ds.sizes['t']):
for z in range(nz):
assert_allclose(n_al[t, 0, 0, z].values,
1000. * t + z % nz,
rtol=1.e-15,
atol=5.e-16) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 0, 1, z].values,
1000. * t + 10. * 1. + (z + 1) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 0, 2, z].values,
1000. * t + 10. * 2. + (z + 2) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 0, 3, z].values,
1000. * t + 10. * 3. + (z + 3) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 0, z].values,
1000. * t + 100. * 1 + 10. * 0. + (z + 4) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 1, z].values,
1000. * t + 100. * 1 + 10. * 1. + (z + 5) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 2, z].values,
1000. * t + 100. * 1 + 10. * 2. + (z + 6) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 3, z].values,
1000. * t + 100. * 1 + 10. * 3. + (z + 7) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
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)
def test_toFieldAligned_dask(self, tmpdir_factory, bout_xyt_example_files):
nz = 6
path = bout_xyt_example_files(tmpdir_factory,
lengths=(3, 3, 4, nz),
nxpe=1,
nype=1,
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')
# 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. * np.pi / ds.sizes['zeta']
for z in range(nz):
n[t, x, y, z] = 1000. * t + 100. * x + 10. * y + z
# The above loop required the call to .load(), but that turned the data into a
# numpy array. Now convert back to dask
n = n.chunk({'t': 1})
assert isinstance(n.data, dask.array.Array)
n.attrs['direction_y'] = 'Standard'
n_al = n.bout.toFieldAligned()
for t in range(ds.sizes['t']):
for z in range(nz):
assert_allclose(n_al[t, 0, 0, z].values,
1000. * t + z % nz,
rtol=1.e-15,
atol=5.e-16) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 0, 1, z].values,
1000. * t + 10. * 1. + (z + 1) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 0, 2, z].values,
1000. * t + 10. * 2. + (z + 2) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 0, 3, z].values,
1000. * t + 10. * 3. + (z + 3) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 0, z].values,
1000. * t + 100. * 1 + 10. * 0. + (z + 4) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 1, z].values,
1000. * t + 100. * 1 + 10. * 1. + (z + 5) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 2, z].values,
1000. * t + 100. * 1 + 10. * 2. + (z + 6) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
for z in range(nz):
assert_allclose(n_al[t, 1, 3, z].values,
1000. * t + 100. * 1 + 10. * 3. + (z + 7) % nz,
rtol=1.e-15,
atol=0.) # noqa: E501
