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")