Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions docs/user_guide/examples/tutorial_schism.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
"1. The mesh coordinates are in a **projected coordinate system (meters)**, not longitude/latitude, so\n",
" we build the grid as a *flat* (Cartesian) mesh.\n",
"2. Velocities are stored at **mesh nodes** over a vertical column of layers, ordered **bottom → surface**.\n",
"3. SCHISM uses a **localized vertical grid (LSC2)**: the number of valid vertical levels varies from\n",
"3. SCHISM uses a **localised vertical grid (LSC2)**: the number of valid vertical levels varies from\n",
" node to node, and levels below the seabed are stored as `NaN`. We use this to flag particles that\n",
" end up below the local bathymetry and stop advecting them.\n",
"\n",
Expand Down Expand Up @@ -147,7 +147,7 @@
"\n",
"* **Layer ordering.** SCHISM stores levels **bottom → surface**, while Parcels expects depth increasing\n",
" downward from the surface, so we reverse the layer axis.\n",
"* **Vertical coordinate.** SCHISM's LSC2 vertical grid (*Localized Sigma Coordinates with Shaved\n",
"* **Vertical coordinate.** SCHISM's LSC2 vertical grid (*Localised Sigma Coordinates with Shaved\n",
" cells*) varies with horizontal position: each node has its own layer depths, and even its own\n",
" *number* of levels (the true depths are written to a separate `zCoordinates` output, not used here).\n",
" **Parcels does not currently support a vertical grid that varies with lateral position.** `UxGrid`\n",
Expand Down
2 changes: 1 addition & 1 deletion src/parcels/_core/fieldset.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,12 +70,12 @@ def __init__(self, models: list[ModelData]):
for model in models:
if not isinstance(model, ModelData):
raise ValueError(f"Expected `model` to be a ModelData object. Got {model}")
# assert_compatible_calendars(fields)

self.models = models
self.constant_model: StructuredModelData | None = None
self._fields: dict[str, Field | VectorField] | None = None
self.reconstruct_fields()
assert_compatible_calendars(self.fields.values())
self.context: dict[str, float] = {}
assert_models_have_same_mesh(self.models)

Expand Down
10 changes: 6 additions & 4 deletions src/parcels/_core/particlefile.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,18 +166,20 @@ def write(self, pset: ParticleSet | ParticleSetView, t, indices=None):
)

if isinstance(t, (np.timedelta64, np.datetime64)):
print(t)
if self._time_interval is None:
raise ValueError(
"The time interval is unknown, so datetime/timedelta times cannot be converted. "
"Pass the ParticleFile to ParticleSet.execute() or call set_time_interval() first."
)
t = timedelta_to_float(t - self._time_interval.left)
vars_to_write = _get_vars_to_write(pclass)
if indices is None:
indices_to_write = _to_write_particles(particle_data, t)
else:
if indices is not None:
indices_to_write = indices
elif isinstance(pset, ParticleSetView):
# A view already selects the particles to write, so `t` is not used to filter
indices_to_write = np.flatnonzero(pset._index)
else:
indices_to_write = _to_write_particles(particle_data, t)

self._writer.write_table(
pa.table({v.name: pa.array(particle_data[v.name][indices_to_write]) for v in vars_to_write}),
Expand Down
51 changes: 28 additions & 23 deletions src/parcels/_core/xgrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -256,9 +256,9 @@ def get_axis_dim(self, axis: ptyping.XgridAxis) -> int:
assert self.sgrid_metadata.vertical_dimensions is not None
return get_cell_count_along_dim(self._ds, self.sgrid_metadata.vertical_dimensions[0])

def localize(
self, position: dict[ptyping.XgridAxis, tuple[int, float]], dims: list[str]
) -> dict[str, tuple[int, float]]:
def localise(
self, position: dict[ptyping.XgridAxis, dict[str, np.ndarray]], dims: list[str]
) -> dict[str, dict[str, np.ndarray]]:
"""
Uses the grid context (i.e., the staggering of the grid) to convert a position relative
to the f-points in the grid to a position relative to the staggered grid the array
Expand All @@ -271,24 +271,28 @@ def localize(
Parameters
----------
position : dict
A mapping of the axis to a tuple of (index, barycentric coordinate) for the
f-points in the grid.
A mapping of the axis to a dict with the arrays ``"index"`` and ``"bcoord"`` (barycentric
coordinate) for the f-points in the grid, as returned by ``search``.
dims : list[str]
A list of dimension names that the DataArray is defined on. This is used to determine
the staggering of the grid and which axis each dimension corresponds to.

Returns
-------
dict[str, tuple[int, float]]
A mapping of the dimension names to a tuple of (index, barycentric coordinate) for
dict[str, dict[str, np.ndarray]]
A mapping of the dimension names to a dict with the arrays ``"index"`` and ``"bcoord"`` for
the staggered grid the DataArray is defined on.

Example
-------
>>> position = {"X": (5, 0.51), "Y": (10, 0.25), "Z": (3, 0.75)}
>>> position = {
... "X": {"index": np.array([5]), "bcoord": np.array([0.51])},
... "Y": {"index": np.array([10]), "bcoord": np.array([0.25])},
... "Z": {"index": np.array([3]), "bcoord": np.array([0.75])},
... }
>>> dims = ["time", "depth", "YC", "XC"]
>>> grid.localize(position, dims)
{'depth': (3, 0.75), 'YC': (9, 0.75), 'XC': (5, 0.01)}
>>> grid.localise(position, dims)
{'depth': {'index': array([3]), 'bcoord': array([0.75])}, 'YC': {'index': array([9]), 'bcoord': array([0.75])}, 'XC': {'index': array([5]), 'bcoord': array([0.01])}}
"""
dim_to_axis = _get_dim_to_axis_mapping(self.sgrid_metadata)
axis_to_var = {dim_to_axis[dim]: dim for dim in dims if dim in dim_to_axis}
Expand All @@ -297,12 +301,13 @@ def localize(
}
return {
axis_to_var[axis]: _convert_center_pos_to_fpoint(
index=index,
bcoord=bcoord,
index=pos["index"],
bcoord=pos["bcoord"],
position=var_positions[axis],
f_point_position=self._fpoint_info[axis],
)
for axis, (index, bcoord) in position.items()
for axis, pos in position.items()
if axis in axis_to_var
}

@property
Expand Down Expand Up @@ -591,29 +596,29 @@ def assert_valid_depth(da_depth):

def _convert_center_pos_to_fpoint(
*,
index: int,
bcoord: float,
index: np.ndarray,
bcoord: np.ndarray,
position: ptyping.GridPosition,
f_point_position: sgrid.Padding,
) -> tuple[int, float]:
) -> dict[str, np.ndarray]:
"""Converts a physical position relative to the cell edges defined in the grid to be relative to the face center.

This is used to "localize" a position to be relative to the staggered grid at which the field is defined, so that
This is used to "localise" a position to be relative to the staggered grid at which the field is defined, so that
it can be easily interpolated.

This also handles different model input cell edges and centers are staggered in different directions (e.g., with NEMO and MITgcm).
"""
if position != "face": # Data is already defined on the F nodes
return index, bcoord
return {"index": index, "bcoord": bcoord}

bcoord = bcoord - 0.5
if bcoord < 0:
bcoord += 1.0
index -= 1
wrapped = bcoord < 0
bcoord = np.where(wrapped, bcoord + 1.0, bcoord)
index = np.where(wrapped, index - 1, index)

# Correct relative to the f-point position
# Padding.BOTH was "inner", Padding.LOW was "right" in xgcm vocabulary
if f_point_position in (sgrid.Padding.BOTH, sgrid.Padding.LOW):
index += 1
index = index + 1

return index, bcoord
return {"index": index, "bcoord": bcoord}
2 changes: 1 addition & 1 deletion src/parcels/_repr_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def particleset_repr(pset: ParticleSet) -> str:

def particlesetview_repr(pview: Any) -> str:
"""Return a pretty repr for ParticleSetView"""
time_string = "not_yet_set" if pview.time is None or np.isnan(pview.time) else f"{pview.time:f}"
time_string = "not_yet_set" if pview.t is None or np.isnan(pview.t) else f"{pview.t:f}"
out = f"P[{pview.particle_id}]: time={time_string}, z={pview.z:f}, y={pview.y:f}, x={pview.x:f}"
vars = [v.name for v in pview._pclass.variables if v.to_write is True and v.name not in ["z", "y", "x", "time"]]
for var in vars:
Expand Down
24 changes: 11 additions & 13 deletions tests/test_field.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import numpy as np
import pytest
import xarray as xr
from re_assert import Matches

from parcels import Field, VectorField
Expand Down Expand Up @@ -82,7 +83,16 @@ def test_field_time_interval():

def test_vectorfield_init_different_time_intervals():
# Tests that a VectorField raises a ValueError if the component fields have different time domains.
...
dsU = simple_UV_dataset(mesh_str="flat")[["U", "grid"]]
dsV = simple_UV_dataset(mesh_str="flat")[["V"]]

dsV = dsV.rename({"time": "timeV"})
dsV["timeV"] = dsV["timeV"] + np.timedelta64(1, "D")
dsV["timeV"].attrs["axis"] = "T"

ds = xr.merge([dsU, dsV])
with pytest.raises(ValueError, match="Multiple time coordinates"):
FieldSet.from_sgrid_conventions(ds)


def test_field_invalid_interpolator():
Expand Down Expand Up @@ -242,15 +252,3 @@ def test_field_eval_out_of_bounds_unstructured(field_name, location, expected):
match="Some interpolated values are out-of-bounds. These values are set to 0. Treat carefully.",
):
np.testing.assert_allclose(field.eval(*location), expected)


def test_field_unstructured_grid_creation(): ...


def test_field_interpolation(): ...


def test_field_interpolation_out_of_spatial_bounds(): ...


def test_field_interpolation_out_of_time_bounds(): ...
31 changes: 26 additions & 5 deletions tests/test_fieldset.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,11 @@
import parcels.tutorial
import tests
from parcels import Particle, ParticleFile, ParticleSet, Variable, convert
from parcels._core.fieldset import FieldSet, IncompatibleMeshesException, _datetime_to_msg
from parcels._core.fieldset import CalendarError, FieldSet, IncompatibleMeshesException, _datetime_to_msg
from parcels._core.mesh import SphericalMesh
from parcels._core.model import _default_vector_field_components, _infer_mesh_from_sgrid_data
from parcels._core.utils.time import TimeInterval
from parcels._datasets.structured.generated import simple_UV_sigma_dataset
from parcels._datasets.structured.generated import simple_UV_dataset, simple_UV_sigma_dataset
from parcels._datasets.structured.generic import datasets as datasets_structured
from parcels._datasets.structured.generic import datasets_sgrid
from parcels._datasets.structured.strategies import reduced_dims_dataset, sgrid_datasets_varying_dims
Expand Down Expand Up @@ -272,7 +272,16 @@ def test_multi_model_nonoverlapping_time_interval():

def test_fieldset_add_incompatible_calendars():
# tests the adding of fieldsets that have incompatible calendars
...
ds1 = datasets_structured["ds_2d_left"][["U_A_grid", "V_A_grid", "grid"]]
fieldset = FieldSet.from_sgrid_conventions(ds1)

ds2 = ds1.copy().rename({"U_A_grid": "U2", "V_A_grid": "V2"})
# Replace the numpy datetime64 time axis by an array of cftime objects with a different calendar
times = [cftime.DatetimeNoLeap(2000, 1, 1 + i) for i in range(ds2["time"].size)]
ds2["time"] = (ds2["time"].dims, np.array(times, dtype=object), ds2["time"].attrs)

with pytest.raises(CalendarError):
fieldset + FieldSet.from_sgrid_conventions(ds2)


@pytest.mark.parametrize(
Expand All @@ -295,7 +304,11 @@ def test_datetime_to_msg(input_, expected):

def test_fieldset_samegrids_UV():
"""Test that if a simple fieldset with U and V is created, that only one grid object is defined."""
...
ds = simple_UV_dataset(mesh_str="flat")
fieldset = FieldSet.from_sgrid_conventions(ds)

assert fieldset.U.grid is fieldset.V.grid
assert len(fieldset.gridset) == 1


def test_fieldset_grid_deduplication():
Expand All @@ -305,7 +318,15 @@ def test_fieldset_grid_deduplication():
When grid deduplication is actually implemented, this might need to be refactored
into multiple tests (/more might be needed).
"""
...
ds = simple_UV_dataset(mesh_str="flat")
fieldset = FieldSet.from_sgrid_conventions(ds)

ds2 = simple_UV_dataset(mesh_str="flat").rename({"U": "U2", "V": "V2"})
fieldset += FieldSet.from_sgrid_conventions(ds2)
assert fieldset.U.grid is fieldset.V.grid
assert fieldset.U2.grid is fieldset.V2.grid
assert fieldset.U.grid is not fieldset.U2.grid
assert fieldset.U.grid is not fieldset.V2.grid


def test_fieldset_from_icon():
Expand Down
41 changes: 0 additions & 41 deletions tests/test_particlefile.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import pyarrow as pa
import pyarrow.parquet as pq
import pytest
import xarray as xr
from re_assert import Matches

import parcels.tutorial
Expand Down Expand Up @@ -156,46 +155,6 @@ def test_write_dtypes_pfile(fieldset, tmp_parquet):
assert tab[f"v_{d.__name__}"].type == pa.from_numpy_dtype(d)


@pytest.mark.skip(reason="Pending ParticleFile refactor; see issue #2386")
@pytest.mark.parametrize("dt", [-np.timedelta64(1, "s"), np.timedelta64(1, "s")])
@pytest.mark.parametrize("maxvar", [2, 4, 10])
def test_pset_repeated_release_delayed_adding_deleting(fieldset, tmp_parquet, dt, maxvar):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I spent quite some time on making this test work for v4, but then decided that it was too complex because the parquet format and the kernel loop have changed so much. So I suggest to remove it - we test/check output in tutorials and on other places

"""Tests that if particles are released and deleted based on age that resulting output file is correct."""
npart = 10
fieldset.add_context("maxvar", maxvar)

MyParticle = Particle.add_variable(
[Variable("sample_var", initial=0.0), Variable("v_once", dtype=np.float64, initial=0.0)]
)

pset = ParticleSet(
fieldset,
x=np.zeros(npart),
y=np.zeros(npart),
pclass=MyParticle,
t=fieldset.time_interval.left + [np.timedelta64(i + 1, "s") for i in range(npart)],
)
pfile = ParticleFile(tmp_parquet, outputdt=abs(dt))

def IncrLon(particles, fieldset): # pragma: no cover
particles.sample_var += 1.0
particles.state = np.where(
particles.sample_var > fieldset.maxvar,
StatusCode.Delete,
particles.state,
)

for _ in range(npart):
pset.execute(IncrLon, dt=dt, runtime=np.timedelta64(1, "s"), output_file=pfile)

ds = xr.open_zarr(tmp_parquet)
samplevar = ds["sample_var"][:]
assert samplevar.shape == (npart, min(maxvar, npart + 1))
# test whether samplevar[:, k] = k
for k in range(samplevar.shape[1]):
assert np.allclose([p for p in samplevar[:, k] if np.isfinite(p)], k + 1)


def test_file_warnings(fieldset, tmp_parquet):
pset = ParticleSet(fieldset, x=[0, 0], y=[0, 0], t=[np.timedelta64(0, "s"), np.timedelta64(1, "s")])
pfile = ParticleFile(tmp_parquet, outputdt=np.timedelta64(2, "s"))
Expand Down
21 changes: 7 additions & 14 deletions tests/test_particleset_execute.py
Original file line number Diff line number Diff line change
Expand Up @@ -390,12 +390,6 @@ def FieldAccessOutsideTime(particles, fieldset): # pragma: no cover
pset.execute(FieldAccessOutsideTime, runtime=np.timedelta64(1, "D"), dt=np.timedelta64(10, "D"))


def test_raise_grid_searching_error(): ...


def test_raise_general_error(): ...


def test_errorinterpolation(fieldset):
class NaNInterpolator(VectorInterpolator): # pragma: no cover
def interp(self, particle_positions, grid_positions, field):
Expand Down Expand Up @@ -571,24 +565,23 @@ def test_uxstommelgyre_multiparticle_pset_execute():
)


@pytest.mark.xfail(reason="Output file not implemented yet")
def test_uxstommelgyre_pset_execute_output():
def test_uxstommelgyre_pset_execute_output(tmp_parquet):
ds = datasets_unstructured["stommel_gyre_delaunay"]
fieldset = FieldSet.from_ugrid_conventions(ds)
pset = ParticleSet(
fieldset,
x=[30.0],
y=[5.0],
z=[50.0],
t=[0.0],
x=30.0,
y=5.0,
z=50.0,
t=np.timedelta64(0, "s"),
pclass=Particle,
)
output_file = ParticleFile(
name="stommel_uxarray_particles.zarr", # the file name
path=tmp_parquet, # the file name
outputdt=np.timedelta64(5, "m"), # the time step of the outputs
)
pset.execute(
runtime=np.timedelta64(10, "m"), dt=np.timedelta64(60, "s"), kernels=AdvectionEE, output_file=output_file
runtime=np.timedelta64(10, "m"), dt=np.timedelta64(150, "s"), kernels=AdvectionEE, output_file=output_file
)


Expand Down
Loading
Loading