diff --git a/src/parcels/interpolators/__init__.py b/src/parcels/interpolators/__init__.py new file mode 100644 index 000000000..bf8a4f4d9 --- /dev/null +++ b/src/parcels/interpolators/__init__.py @@ -0,0 +1,37 @@ +from ._uxinterpolators import ( + Ux_Velocity, + UxPiecewiseConstantFace, + UxPiecewiseLinearNode, +) +from ._xinterpolators import ( + CGrid_Tracer, + CGrid_Velocity, + XConstantField, + XFreeslip, + XLinear, + XLinear_Velocity, + XLinearInvdistLandTracer, + XNearest, + XPartialslip, + ZeroInterpolator, + ZeroInterpolator_Vector, +) + +__all__ = [ # noqa: RUF022 + # xinterpolators + "CGrid_Tracer", + "CGrid_Velocity", + "XConstantField", + "XFreeslip", + "XLinear", + "XLinearInvdistLandTracer", + "XLinear_Velocity", + "XNearest", + "XPartialslip", + "ZeroInterpolator", + "ZeroInterpolator_Vector", + # uxinterpolators + "UxPiecewiseConstantFace", + "UxPiecewiseLinearNode", + "Ux_Velocity", +] diff --git a/src/parcels/interpolators/_uxinterpolators.py b/src/parcels/interpolators/_uxinterpolators.py new file mode 100644 index 000000000..1d6899a14 --- /dev/null +++ b/src/parcels/interpolators/_uxinterpolators.py @@ -0,0 +1,73 @@ +"""Collection of pre-built interpolation kernels for unstructured grids.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from parcels._core.field import Field, VectorField + from parcels._core.uxgrid import _UXGRID_AXES + + +def UxPiecewiseConstantFace( + particle_positions: dict[str, float | np.ndarray], + grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], + field: Field, +): + """ + Piecewise constant interpolation kernel for face registered data. + This interpolation method is appropriate for fields that are + face registered, such as u,v in FESOM. + """ + return field.data.values[ + grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + ] + + +def UxPiecewiseLinearNode( + particle_positions: dict[str, float | np.ndarray], + grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], + field: Field, +): + """ + Piecewise linear interpolation kernel for node registered data located at vertical interface levels. + This interpolation method is appropriate for fields that are node registered such as the vertical + velocity W in FESOM2. Effectively, it applies barycentric interpolation in the lateral direction + and piecewise linear interpolation in the vertical direction. + """ + ti = grid_positions["T"]["index"] + zi, fi = grid_positions["Z"]["index"], grid_positions["FACE"]["index"] + z = particle_positions["z"] + bcoords = grid_positions["FACE"]["bcoord"] + node_ids = field.grid.uxgrid.face_node_connectivity[fi, :].values + # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels. + # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1. + # First, do barycentric interpolation in the lateral direction for each interface level + fzk = np.sum(field.data.values[ti[:, None], zi[:, None], node_ids] * bcoords, axis=-1) + fzkp1 = np.sum(field.data.values[ti[:, None], zi[:, None] + 1, node_ids] * bcoords, axis=-1) + + # Then, do piecewise linear interpolation in the vertical direction + zk = field.grid.z.values[zi] + zkp1 = field.grid.z.values[zi + 1] + return (fzk * (zkp1 - z) + fzkp1 * (z - zk)) / (zkp1 - zk) # Linear interpolation in the vertical direction + + +def Ux_Velocity( + particle_positions: dict[str, float | np.ndarray], + grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], + vectorfield: VectorField, +): + """Interpolation kernel for Vectorfields of velocity on a UxGrid.""" + u = vectorfield.U._interp_method(particle_positions, grid_positions, vectorfield.U) + v = vectorfield.V._interp_method(particle_positions, grid_positions, vectorfield.V) + if vectorfield.grid._mesh == "spherical": + u /= 1852 * 60 * np.cos(np.deg2rad(particle_positions["lat"])) + v /= 1852 * 60 + + if "3D" in vectorfield.vector_type: + w = vectorfield.W._interp_method(particle_positions, grid_positions, vectorfield.W) + else: + w = 0.0 + return u, v, w diff --git a/src/parcels/interpolators.py b/src/parcels/interpolators/_xinterpolators.py similarity index 87% rename from src/parcels/interpolators.py rename to src/parcels/interpolators/_xinterpolators.py index 52052ff85..14f92cb00 100644 --- a/src/parcels/interpolators.py +++ b/src/parcels/interpolators/_xinterpolators.py @@ -1,4 +1,4 @@ -"""Collection of pre-built interpolation kernels.""" +"""Collection of pre-built interpolation kernels for structured grids.""" from __future__ import annotations @@ -12,26 +12,8 @@ if TYPE_CHECKING: from parcels._core.field import Field, VectorField - from parcels._core.uxgrid import _UXGRID_AXES from parcels._core.xgrid import _XGRID_AXES -__all__ = [ - "CGrid_Tracer", - "CGrid_Velocity", - "UxPiecewiseConstantFace", - "UxPiecewiseLinearNode", - "Ux_Velocity", - "XConstantField", - "XFreeslip", - "XLinear", - "XLinearInvdistLandTracer", - "XLinear_Velocity", - "XNearest", - "XPartialslip", - "ZeroInterpolator", - "ZeroInterpolator_Vector", -] - def ZeroInterpolator( particle_positions: dict[str, float | np.ndarray], @@ -649,65 +631,3 @@ def XLinearInvdistLandTracer( values[exact_particles] = exact_vals[exact_particles] return values.compute() if is_dask_collection(values) else values - - -def UxPiecewiseConstantFace( - particle_positions: dict[str, float | np.ndarray], - grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], - field: Field, -): - """ - Piecewise constant interpolation kernel for face registered data. - This interpolation method is appropriate for fields that are - face registered, such as u,v in FESOM. - """ - return field.data.values[ - grid_positions["T"]["index"], grid_positions["Z"]["index"], grid_positions["FACE"]["index"] - ] - - -def UxPiecewiseLinearNode( - particle_positions: dict[str, float | np.ndarray], - grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], - field: Field, -): - """ - Piecewise linear interpolation kernel for node registered data located at vertical interface levels. - This interpolation method is appropriate for fields that are node registered such as the vertical - velocity W in FESOM2. Effectively, it applies barycentric interpolation in the lateral direction - and piecewise linear interpolation in the vertical direction. - """ - ti = grid_positions["T"]["index"] - zi, fi = grid_positions["Z"]["index"], grid_positions["FACE"]["index"] - z = particle_positions["z"] - bcoords = grid_positions["FACE"]["bcoord"] - node_ids = field.grid.uxgrid.face_node_connectivity[fi, :].values - # The zi refers to the vertical layer index. The field in this routine are assumed to be defined at the vertical interface levels. - # For interface zi, the interface indices are [zi, zi+1], so we need to use the values at zi and zi+1. - # First, do barycentric interpolation in the lateral direction for each interface level - fzk = np.sum(field.data.values[ti[:, None], zi[:, None], node_ids] * bcoords, axis=-1) - fzkp1 = np.sum(field.data.values[ti[:, None], zi[:, None] + 1, node_ids] * bcoords, axis=-1) - - # Then, do piecewise linear interpolation in the vertical direction - zk = field.grid.z.values[zi] - zkp1 = field.grid.z.values[zi + 1] - return (fzk * (zkp1 - z) + fzkp1 * (z - zk)) / (zkp1 - zk) # Linear interpolation in the vertical direction - - -def Ux_Velocity( - particle_positions: dict[str, float | np.ndarray], - grid_positions: dict[_UXGRID_AXES, dict[str, int | float | np.ndarray]], - vectorfield: VectorField, -): - """Interpolation kernel for Vectorfields of velocity on a UxGrid.""" - u = vectorfield.U._interp_method(particle_positions, grid_positions, vectorfield.U) - v = vectorfield.V._interp_method(particle_positions, grid_positions, vectorfield.V) - if vectorfield.grid._mesh == "spherical": - u /= 1852 * 60 * np.cos(np.deg2rad(particle_positions["lat"])) - v /= 1852 * 60 - - if "3D" in vectorfield.vector_type: - w = vectorfield.W._interp_method(particle_positions, grid_positions, vectorfield.W) - else: - w = 0.0 - return u, v, w diff --git a/src/parcels/kernels/__init__.py b/src/parcels/kernels/__init__.py index b4635869a..de252b003 100644 --- a/src/parcels/kernels/__init__.py +++ b/src/parcels/kernels/__init__.py @@ -1,4 +1,4 @@ -from .advection import ( +from ._advection import ( AdvectionAnalytical, AdvectionEE, AdvectionRK2, @@ -8,7 +8,7 @@ AdvectionRK4_3D_CROCO, AdvectionRK45, ) -from .advectiondiffusion import ( +from ._advectiondiffusion import ( AdvectionDiffusionEM, AdvectionDiffusionM1, DiffusionUniformKh, diff --git a/src/parcels/kernels/advection.py b/src/parcels/kernels/_advection.py similarity index 100% rename from src/parcels/kernels/advection.py rename to src/parcels/kernels/_advection.py diff --git a/src/parcels/kernels/advectiondiffusion.py b/src/parcels/kernels/_advectiondiffusion.py similarity index 100% rename from src/parcels/kernels/advectiondiffusion.py rename to src/parcels/kernels/_advectiondiffusion.py