def test_multiple_files(tmp_path):
"""Test creation of output file from multiple input files."""
os.chdir(tmp_path)
pu = "test_u.nc"
pv = "test_v.nc"
coords = {"lon": np.arange(5), "lat": np.arange(4), "time": np.arange(3)}
du = xr.Dataset({"U": (["time", "lat", "lon"], np.empty((3, 4, 5)))})
du.to_netcdf(pu)
dv = xr.Dataset({"V": (["time", "lat", "lon"], np.empty((3, 4, 5)))})
dv.to_netcdf(pv)
f = filtering.LagrangeFilter(
"multiple_files",
{"U": pu, "V": pv},
{"U": "U", "V": "V"},
{k: k for k in ["lon", "lat", "time"]},
sample_variables=["U"],
)
f.create_out().close()
out = Path("multiple_files.nc")
assert out.exists()
d = xr.open_dataset(out)
assert d.dims == {"lon": 5, "lat": 4, "time": 0}
assert "var_U" in d.variables
assert "var_V" not in d.variables
def test_filtering_output_times(tmp_chdir):
"""Test that input times are copied to the output file."""
nt = 37
w = 1 / 6
d, t, _ = velocity_dataset(nt, w)
f = filtering.LagrangeFilter(
"output_times",
d,
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=9 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=30 * 60,
)
f()
out = Path("output_times.nc")
assert out.exists()
d_out = xr.open_dataset(out)
xr.testing.assert_allclose(d_out.time, d.time[9:-9])
def test_absolute_times():
"""Test decoding of absolute times"""
nt = 37
w = 1 / 6
d, t, _ = velocity_dataset(nt, w)
t = t.copy()
# offset absolute and relative times
d.assign_coords(time=d["time"] + 1800)
f = filtering.LagrangeFilter(
"absolute_times",
d,
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=[],
window_size=0,
)
assert np.all(f._window_times(d["time"], True) == t)
def test_other_data(tmp_path):
"""Test creation of output file where a non-velocity variable is sampled."""
os.chdir(tmp_path)
p = "test.nc"
coords = {"lon": np.arange(5), "lat": np.arange(4), "time": np.arange(3)}
d = xr.Dataset(
{
"U": (["time", "lat", "lon"], np.empty((3, 4, 5))),
"V": (["time", "lat", "lon"], np.empty((3, 4, 5))),
"P": (["time", "lat", "lon"], np.empty((3, 4, 5))),
},
coords=coords,
)
d.to_netcdf(p)
f = filtering.LagrangeFilter(
"other_data",
{"U": p, "V": p, "P": p},
{"U": "U", "V": "V", "P": "P"},
{k: k for k in ["lon", "lat", "time"]},
sample_variables=["P"],
)
f.create_out().close()
out = Path("other_data.nc")
assert out.exists()
d = xr.open_dataset(out)
assert "var_P" in d.variables
def test_full_filtering(tmp_chdir):
"""Test running the full filtering workflow by calling the filter object."""
nt = 37
w = 1 / 6
d, t, _ = velocity_dataset(nt, w)
f = filtering.LagrangeFilter(
"full_filtering",
d,
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=18 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=30 * 60,
)
f(t[nt // 2])
out = Path("full_filtering.nc")
assert out.exists()
def test_curvilinear_advection():
"""Sanity check of advection on a curvilinear grid."""
nt = 37
w = 1 / 5
d, t, u = velocity_dataset(nt, w, curvilinear=True)
f = filtering.LagrangeFilter(
"curvilinear_test",
d,
{
"U": "u",
"V": "v"
},
{
"lon": "x_curv",
"lat": "y_curv",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=18 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=60,
)
transformed = f.advection_step(t[nt // 2], output_time=True)
t_trans = transformed["time"]
u_trans = transformed["var_U"][1][:, 4].compute()
assert np.allclose(u, u_trans, rtol=1e-1)
assert np.array_equal(t, t_trans)
def test_curvilinear(tmp_path):
"""Test creation of output file where grid is curvilinear."""
os.chdir(tmp_path)
p = "test.nc"
coords = {"xi": np.arange(5), "eta": np.arange(4), "time": np.arange(3)}
d = xr.Dataset(
{
"U": (["time", "eta", "xi"], np.empty((3, 4, 5))),
"V": (["time", "eta", "xi"], np.empty((3, 4, 5))),
"lat": (["eta", "xi"], np.ones((4, 5))),
"lon": (["eta", "xi"], 2 * np.ones((4, 5))),
},
coords=coords,
)
d.to_netcdf(p)
f = filtering.LagrangeFilter(
"curvilinear",
{"U": p, "V": p},
{"U": "U", "V": "V"},
{k: k for k in ["lon", "lat", "time"]},
sample_variables=["U"],
)
f.create_out().close()
out = Path("curvilinear.nc")
assert out.exists()
d = xr.open_dataset(out)
assert "var_U" in d.variables
assert d["var_U"].dims == ("time", "eta", "xi")
assert "lat" in d.variables
assert d["lat"].dims == ("eta", "xi")
def test_frequency_filter(leewave_data):
"""Test creation and application of frequency-space step filter."""
f = filtering.LagrangeFilter(
"frequency_filter",
leewave_data,
{
"U": "U",
"V": "V"
},
{
"lon": "x",
"lat": "y",
"time": "t"
},
["U"],
window_size=3 * 24 * 3600,
)
f.make_zonally_periodic()
f.make_meridionally_periodic()
# attach filter to object
filt = filtering.filter.FrequencySpaceFilter(1e-4, 3600)
f.inertial_filter = filt
adv = f.advection_step(7 * 24 * 3600)
U_filt = f.filter_step(adv)["var_U"].reshape(leewave_data.y.size, -1)
assert np.all((leewave_data.U_orig.data - U_filt[0, :])**2 < 3e-8)
def test_power_spectrum(leewave_data):
"""Test computation of the power spectrum for velocity in
the leewave dataset."""
f = filtering.LagrangeFilter(
"power_spectrum",
leewave_data,
{
"U": "U",
"V": "V"
},
{
"lon": "x",
"lat": "y",
"time": "t"
},
["U"],
window_size=3 * 24 * 3600,
)
f.make_zonally_periodic()
f.make_meridionally_periodic()
spectrum = filtering.analysis.power_spectrum(f, 7 * 24 * 3600)
assert "var_U" in spectrum
assert np.all(np.isreal(spectrum["var_U"]))
def test_masked_filtering(tmp_chdir):
"""Test running the full filtering workflow, seeding only on a subdomain."""
nt = 37
w = 1 / 6
d, t, _ = velocity_dataset(nt, w)
f = filtering.LagrangeFilter(
"masked_filtering",
d,
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=18 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=30 * 60,
)
f.seed_subdomain(500, 500, 500, 500) # seed only the middle poin
f(t[nt // 2])
out = Path("masked_filtering.nc")
assert out.exists()
def test_dims_indices_dicts(tmp_path):
"""Test creation of output file where dimensions and indices are specified in
per-variable dictionaries, instead of globally."""
os.chdir(tmp_path)
p = "test.nc"
coords = {
"X": np.arange(5),
"Xp": np.arange(5) + 0.5,
"Y": np.arange(4),
"Yp": np.arange(4) + 0.5,
"Z": np.arange(3),
"Zm": np.arange(1),
"time": np.arange(3),
}
d = xr.Dataset(
{
"UVEL": (["time", "Z", "Y", "X"], np.empty((3, 3, 4, 5))),
"VVEL": (["time", "Z", "Y", "X"], np.empty((3, 3, 4, 5))),
"UBAR": (["time", "Zm", "Y", "Xp"], np.empty((3, 1, 4, 5))),
"VBAR": (["time", "Zm", "Yp", "X"], np.empty((3, 1, 4, 5))),
},
coords=coords,
)
d.to_netcdf(p)
dims = {
"U": {"lon": "X", "lat": "Y", "time": "time", "depth": "Z"},
"V": {"lon": "X", "lat": "Y", "time": "time", "depth": "Z"},
"UBAR": {"lon": "Xp", "lat": "Y", "time": "time", "depth": "Zm"},
"VBAR": {"lon": "X", "lat": "Yp", "time": "time", "depth": "Zm"},
}
indices = {
"U": {"depth": [2]},
"V": {"depth": [2]},
"UBAR": {"depth": [0]},
"VBAR": {"depth": [0]},
}
f = filtering.LagrangeFilter(
"dims_indices_dicts",
{v: p for v in ["U", "V", "UBAR", "VBAR"]},
{"U": "UVEL", "V": "VVEL", "UBAR": "UBAR", "VBAR": "VBAR"},
dims,
sample_variables=["UBAR", "VBAR"],
indices=indices,
)
f.create_out().close()
out = Path("dims_indices_dicts.nc")
assert out.exists()
def test_filtering_output_times_with_calendar(tmp_chdir):
"""Test that input times from a file with a calendar
are correctly copied to the output file."""
nt = 37
w = 1 / 6
d, t, _ = velocity_dataset(nt, w)
# modify the dataset to have a "days since ..." calendar
d["time"] = d["time"] / 3600 / 24
d.time.attrs["units"] = "days since 1900-01-01 00:00:00"
d.time.attrs["calendar"] = "NOLEAP"
d.to_netcdf("data.nc")
f = filtering.LagrangeFilter(
"output_calendar",
{
"U": "data.nc",
"V": "data.nc"
},
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=9 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=30 * 60,
)
f()
out = Path("output_calendar.nc")
assert out.exists()
# open the datasets without decoding the time axis so we can compare
d_out = xr.open_dataset(out) # , decode_times=False)
d_in = xr.open_dataset("data.nc") # , decode_times=False)
# check the calendar attributes were propagated
assert d_out.time.attrs == d_in.time.attrs
# check the time values themselves
xr.testing.assert_allclose(d_out.time, d_in.time[9:-9])
def test_staggered(tmp_path):
"""Test creation of output file where velocity is staggered."""
os.chdir(tmp_path)
p = "test.nc"
coords = {
"xu": np.arange(5) + 0.5,
"xt": np.arange(5),
"yv": np.arange(4) + 0.5,
"yt": np.arange(4),
"time": np.arange(3),
}
d = xr.Dataset(
{
"U": (["time", "yt", "xu"], np.empty((3, 4, 5))),
"V": (["time", "yv", "xt"], np.empty((3, 4, 5))),
"P": (["time", "yt", "xt"], np.empty((3, 4, 5))),
},
coords=coords,
)
d.to_netcdf(p)
# variables
v = ["U", "V", "P"]
f = filtering.LagrangeFilter(
"staggered",
{k: p for k in v},
{k: k for k in v},
{
"U": {"lon": "xu", "lat": "yt", "time": "time"},
"V": {"lon": "xt", "lat": "yv", "time": "time"},
"P": {"lon": "xt", "lat": "yt", "time": "time"},
},
sample_variables=v,
)
f.create_out().close()
out = Path("staggered.nc")
assert out.exists()
d = xr.open_dataset(out)
for n in v:
assert f"var_{n}" in d.variables
def test_single_file(tmp_path):
"""Test creation of output file from a single input file."""
# because filtering puts files in the current directory, we need to change
# to the test directory
os.chdir(tmp_path)
# set up path for input file
p = "test.nc"
coords = {"lon": np.arange(5), "lat": np.arange(4), "time": np.arange(3)}
d = xr.Dataset(
{
"U": (["time", "lat", "lon"], np.empty((3, 4, 5))),
"V": (["time", "lat", "lon"], np.empty((3, 4, 5))),
},
coords=coords,
)
d.to_netcdf(p)
# create class
f = filtering.LagrangeFilter(
"single_file",
{"U": p, "V": p},
{"U": "U", "V": "V"},
{k: k for k in ["lon", "lat", "time"]},
sample_variables=["U"],
)
# create output file
f.create_out().close()
# check that we actually made the right file
out = Path("single_file.nc")
assert out.exists()
# check dimensions and sizes, and variables
d = xr.open_dataset(out)
assert d.dims == {"lon": 5, "lat": 4, "time": 0}
assert "var_U" in d.variables
assert "var_V" not in d.variables
def test_clobber(tmp_path):
"""Test whether existing output files are clobbered."""
os.chdir(tmp_path)
out_path = tmp_path / "clobbering.nc"
# write an input file
p = "test.nc"
coords = {"lon": np.arange(5), "lat": np.arange(4), "time": np.arange(3)}
d = xr.Dataset(
{
"U": (["time", "lat", "lon"], np.empty((3, 4, 5))),
"V": (["time", "lat", "lon"], np.empty((3, 4, 5))),
},
coords=coords,
)
d.to_netcdf(p)
# create filter
f = filtering.LagrangeFilter(
"clobbering",
{"U": p, "V": p},
{"U": "U", "V": "V"},
{k: k for k in ["lon", "lat", "time"]},
sample_variables=["U"],
)
# first time, file should create correctly
with f.create_out() as d:
pass
assert out_path.exists()
# second time, we should fail on clobbering
with pytest.raises(OSError):
f.create_out()
# but, we should be able to open the file with the clobber flag
assert out_path.exists()
with f.create_out(clobber=True) as d:
pass
assert out_path.exists()
def test_sanity_filtering_from_dataset():
"""Sanity check of filtering using the library.
As with the :func:`~test_sanity` test, this sets up a mean
velocity field (in 2D) with an oscillating component. Because the
velocity field is uniform in time, the Lagrangian timeseries
should be the same as the 1D timeseries.
"""
nt = 37
w = 1 / 6
d, t, _ = velocity_dataset(nt, w)
f = filtering.LagrangeFilter(
"sanity_test",
d,
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=18 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=30 * 60,
)
# filter from the middle of the series
filtered = f.filter_step(f.advection_step(t[nt // 2]))["var_U"]
# we expect a lot of parcels to hit the edge and die
# but some should stay alive
filtered = filtered[~np.isnan(filtered)]
assert filtered.size > 0
value = filtered.item(0)
assert value == pytest.approx(0.0, abs=1e-3)
def test_sanity_advection_from_file(tmp_path):
"""Sanity check of advection, with data loaded from a file."""
nt = 37
w = 1 / 6
d, t, u = velocity_dataset(nt, w)
p = tmp_path / "data.nc"
d.to_netcdf(p)
f = filtering.LagrangeFilter(
"advection_test",
{
"U": str(p),
"V": str(p)
},
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=18 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=60,
)
transformed = f.advection_step(t[nt // 2], output_time=True)
t_trans = transformed["time"]
u_trans = transformed["var_U"][1][:, 4].compute()
assert np.allclose(u, u_trans, rtol=1e-1)
assert np.array_equal(t, t_trans)
def test_doubly_periodic_advection():
"""Sanity check of advection in a doubly periodic domain.
The flow in this test is diagonal, but we set up a doubly-periodic
domain, so we expect that all particles remain alive.
"""
nt = 37
w = 1 / 6
d, t, u = velocity_dataset(nt, w)
f = filtering.LagrangeFilter(
"periodic_test",
d,
{
"U": "u",
"V": "u"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["V"],
mesh="flat",
window_size=18 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=60,
)
f.make_zonally_periodic(width=3)
f.make_meridionally_periodic(width=3)
transformed = f.advection_step(t[nt // 2])
v_trans = transformed["var_V"][1].compute()
assert not np.any(np.isnan(v_trans))
def test_sanity_advection():
"""Sanity check of advection.
Using a uniform velocity field, the particles should have the same
value regardless of where in the domain they go.
"""
nt = 37
w = 1 / 6
d, t, u = velocity_dataset(nt, w)
f = filtering.LagrangeFilter(
"advection_test",
d,
{
"U": "u",
"V": "v"
},
{
"lon": "x",
"lat": "y",
"time": "time"
},
sample_variables=["U"],
mesh="flat",
window_size=18 * 3600,
highpass_frequency=(w / 2) / 3600,
advection_dt=60,
)
transformed = f.advection_step(t[nt // 2], output_time=True)
t_trans = transformed["time"]
u_trans = transformed["var_U"][1][:, 4].compute()
assert np.allclose(u, u_trans, rtol=1e-1)
assert np.array_equal(t, t_trans)
"lat": "geolat_c",
"time": "time"
},
"V": {
"lon": "geolon_c",
"lat": "geolat_c",
"time": "time"
},
}
# construct filter object
ff = filtering.LagrangeFilter(
ncfile_out,
filenames,
variables,
dimensions,
sample_variables=["U", "V"],
mesh="spherical",
c_grid=False,
window_size=timedelta(days=2.0).total_seconds(),
)
# Filter only region below than 60N
ff.seed_subdomain(max_lat=60)
# change advection timestep (600 is the default)
ff.advection_dt = 600
# Construct a variable cut-off filter.
omega = 2 * np.pi / 24 / 60**2
lat = ff._grid_lat
mask = lat > 10
dimensions = {
"U": {"lon": "Xp1", "lat": "Y", "time": "T", "depth": "Zmd000200"},
"V": {"lon": "X", "lat": "Yp1", "time": "T", "depth": "Zmd000200"},
"RHO": {"lon": "X", "lat": "Y", "time": "T", "depth": "Zmd000200"}
}
# specify what depth level we should filter
zi=10
indices = {"depth": [zi]}
# construct filter object
ff = filtering.LagrangeFilter(
ncfile_out, filenames, variables, dimensions,
sample_variables=["U","V","RHO"], mesh="flat",
window_size=timedelta(days=2).total_seconds(),
indices=indices,
c_grid=True,
highpass_frequency=1.2060e-4
)
# To implement an Eulerian filter, turn off the advection at this point
# ff.kernel = ff.sample_kernel
# for density, pressure, temperature, tracers it is recommended to use the following interpolation method:
ff.fieldset.RHO.interp_method = "linear_invdist_land_tracer"
# Choose an output grid for the filtered fields
ff.set_particle_grid("RHO")
# Filter only a central subdomain
ff.seed_subdomain(min_lon=50000,
def test_dimension_files(tmp_path):
"""Test creation of output file where input variables pull dimensions
from different files.
"""
os.chdir(tmp_path)
pvel = "test_vel.nc"
pdata = "test_data.nc"
# create a velocity dataset with depth data
coords = {
"lon": np.arange(5),
"lat": np.arange(4),
"depth": np.arange(3),
"time": np.arange(3),
}
d = xr.Dataset(
{
"U": (["time", "depth", "lat", "lon"], np.empty((3, 3, 4, 5))),
"V": (["time", "depth", "lat", "lon"], np.empty((3, 3, 4, 5))),
},
coords=coords,
)
d.to_netcdf(pvel)
# create a data dataset without depth data
del coords["depth"]
d = xr.Dataset(
{
"UBAR": (["time", "lat", "lon"], np.empty((3, 4, 5))),
"VBAR": (["time", "lat", "lon"], np.empty((3, 4, 5))),
},
coords=coords,
)
d.to_netcdf(pdata)
# filename dicts
fd = {d: pvel for d in ["lon", "lat", "depth", "data"]}
dd = fd.copy()
dd["data"] = pdata
f = filtering.LagrangeFilter(
"dimension_files",
{"U": fd, "V": fd, "UBAR": dd, "VBAR": dd},
{"U": "U", "V": "V", "UBAR": "UBAR", "VBAR": "VBAR"},
{k: k for k in ["lon", "lat", "depth", "time"]},
sample_variables=["UBAR", "VBAR"],
indices={"depth": [0]},
)
f.create_out().close()
out = Path("dimension_files.nc")
assert out.exists()
d = xr.open_dataset(out)
assert "var_UBAR" in d.variables
assert "var_VBAR" in d.variables
assert d["var_UBAR"].dims == ("time", "lat", "lon")
