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