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_register_new_geometry(self):
@register_geometry(name="Schwarzschild")
def add_schwarzschild_coords(ds, coordinates=None):
ds["event_horizon"] = 4.0
return ds
assert "Schwarzschild" in REGISTERED_GEOMETRIES.keys()
original = Dataset()
original["dy"] = DataArray(np.ones((3, 4)), dims=("x", "y"))
metadata = {
"bout_tdim": "t",
"bout_xdim": "x",
"bout_ydim": "y",
"bout_zdim": "z",
"keep_xboundaries": True,
"keep_yboundaries": True,
}
original.attrs["metadata"] = metadata
original["dy"].attrs["metadata"] = metadata
updated = apply_geometry(ds=original, geometry_name="Schwarzschild")
assert_equal(updated["event_horizon"], DataArray(4.0))
# clean up
del REGISTERED_GEOMETRIES["Schwarzschild"]
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_register_new_geometry(self):
@register_geometry(name="Schwarzschild")
def add_schwarzschild_coords(ds, coordinates=None):
ds['event_horizon'] = 4.0
return ds
assert "Schwarzschild" in REGISTERED_GEOMETRIES.keys()
original = Dataset()
updated = apply_geometry(ds=original, geometry_name="Schwarzschild")
assert_equal(updated['event_horizon'], DataArray(4.0))
# clean up
del REGISTERED_GEOMETRIES["Schwarzschild"]
ds2 = open_boutdataset(args.gridfile2,
keep_xboundaries=True,
keep_yboundaries=True,
info=False).load()
ds2.attrs["name"] = args.gridfile2
if args.poloidal_plot:
from xbout.geometries import apply_geometry
ds1.metadata["MXG"] = 2
ds1.metadata["MYG"] = ds1.metadata["y_boundary_guards"]
ds2.metadata["MXG"] = 2
ds2.metadata["MYG"] = ds2.metadata["y_boundary_guards"]
coordinates = {"x": "psi_poloidal", "y": "theta_coord"}
ds1 = apply_geometry(ds1, "toroidal", coordinates=coordinates)
ds2 = apply_geometry(ds2, "toroidal", coordinates=coordinates)
common_variables, common_scalar_variables = check_missing_variables(
ds1, ds2, ignore_ylow=args.ignore_ylow)
check_scalars(ds1, ds2, common_scalar_variables)
plot_grid_points(ds1,
ds2,
poloidal_plot=args.poloidal_plot,
show_all=args.show_all)
plot_arrays(
ds1,
ds2,
common_variables,
atol=args.atol,
poloidal_plot=args.poloidal_plot,
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_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
def test_unregistered_geometry(self):
with pytest.raises(UnregisteredGeometryError,
match="tesseract is not a registered geometry"):
apply_geometry(ds=Dataset(), geometry_name="tesseract")
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_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