How to override velocities passed to the advection kernels

[willirath]

(My example is for adding Stokes drift to the ocean velocities used for advecting particles. The general idea of combining velocity fields in a non-invasive way is very interesting for various applications, though.)

We were working on the inclusion of Stokes drift. Ideally, we'd just make sure that the velocities used for advection are U_ocean + U_stokes. The following is a first stab at sub-classing FieldSet to make FieldSet.UV return a sum of (U, V) and the Stokes-drift velocities:

from parcels import FieldSet, Field
import numpy as np

class StokesFieldSet(FieldSet):
    """FieldSet with additional stokes velocities.

    :param U: :class:`parcels.field.Field` object for zonal velocity component
    :param V: :class:`parcels.field.Field` object for meridional velocity component
    :param US: :class:`parcels.field.Field` object for zonal stokes velocity component
    :param VS: :class:`parcels.field.Field` object for meridional stokes velocity component
    :param fields: Dictionary of additional :class:`parcels.field.Field` objects

    """
    def __init__(self, U, V, US=None, VS=None, fields={}):
        super(StokesFieldSet, self).__init__(U, V, fields=fields)

        if US:
            self.add_field(US)
        if VS:
            self.add_field(VS)

        self._UV = self.UV

        UV = Field('UV', None)
        UV.fieldset = self
        self.UV = UV
        self.UV.getUV = self.getUV_with_stokes

    def getUV_with_stokes(self, time, x, y, z):
        _U, _V = self._UV[time, x, y, z]
        if "US" in [f.name for f in self.fields]:
            _U = _U + self.US[time, x, y, z]
        if "VS" in [f.name for f in self.fields]:
            _V = _V + self.VS[time, x, y, z]
        return (_U, _V)

def test_stokes_fieldset():
    np.random.seed(1)

    xdim = 2
    ydim = 2

    lon = np.linspace(-20, 20, xdim, dtype=np.float32)
    lat = np.linspace(-20, 20, ydim, dtype=np.float32)

    dimensions = {'lon': lon, 'lat': lat}

    U = Field('U', np.random.randn(ydim, xdim), lon=lon, lat=lat)
    V = Field('V', np.random.randn(ydim, xdim), lon=lon, lat=lat)
    US = Field('US', np.random.randn(ydim, xdim), lon=lon, lat=lat)
    VS = Field('VS', np.random.randn(ydim, xdim), lon=lon, lat=lat)

    fs = FieldSet(U, V, fields={"US": US, "VS": VS})
    sfs = StokesFieldSet(U, V, US=US, VS=VS)

    assert fs.U[1, 2, 3, 4] == sfs.U[1, 2, 3, 4]
    assert fs.V[1, 2, 3, 4] == sfs.V[1, 2, 3, 4]

    assert fs.US[1, 2, 3, 4] == sfs.US[1, 2, 3, 4]
    assert fs.VS[1, 2, 3, 4] == sfs.VS[1, 2, 3, 4]

    assert fs.UV[1, 2, 3, 4][0] + fs.US[1, 2, 3, 4] == sfs.UV[1, 2, 3, 4][0]
    assert fs.UV[1, 2, 3, 4][1] + fs.VS[1, 2, 3, 4] == sfs.UV[1, 2, 3, 4][1]

test_stokes_fieldset()

The test is passing, which means that for the new class StokesFieldSet, the returned velocities are a sum of the ocean velocities and Stokes drift.

Using this class to create a fieldset and integrating trajectories fails if we use JIT. If I understand http://www.oceanparcels.org/#writing-parcels-kernels correctly, we cannot expect the code generator for the compiled parts of Parcels to handle too general Python code.

I did not dive into the JIT-details of Parcels. Before doing so: Do you think overriding Field or FieldSet is a viable option at all?

[willirath]

cc: @sruehs

[erikvansebille]

A very interesting approach, @willirath. However, I'm not surprised it doesn't work in JIT, as the compiler indeed really covers only a very small subset of all python commands.

A completely different approach is to rewrite the AdvectionRK4 kernel to use both fieldset.UV and fieldset.UVstokes. E.g.

(u1o, v1o) = fieldset.UV[time, particle.lon, particle.lat, particle.depth]
(u1s, v1s) = fieldset.UVstokes[time, particle.lon, particle.lat, particle.depth]
(u1, v1) = (u1o+u1s, v1o+v1s)

The big advantage of this is that Stokes and ocean currents don't need to be on the same grid then. In many cases, the Stokes flow is computed on a completely different, often coarser, grid (also in your case)?

If you want to take this approach, the trick is to create a new FieldSet fieldset.UVstokes. This may not be trivial either, however.

If you're working on a Rectilinear grid (i.e. longitudes and latitudes only 1D vectors, so that no rotation is required), then you can simply do (u1s, v1s) = (fieldset.Ustokes[time, particle.lon, particle.lat, particle.depth], fieldset.Vstokes[time, particle.lon, particle.lat, particle.depth]) where fieldset.Ustokes and fieldset.Vstokes are fields with the Stokes flow. But on a Curvilinear grid, you will need to rotate, and I'll have to think about how best to do that. @delandmeterp, do you have any ideas?

By the way, if the ocean and stokes data are on the same grid, can't you simply add them up before doing pset.execute()? I.e. something like fieldset.U.data += fieldset.Ustokes.data and similar for the V fields?

[delandmeterp]

To summarise:

• On curvilinear grids, UV interpolation requires the existence of fields called cosU, sinU, cosV, sinV, that rotate the velocity to the correct zonal/meridional directions. cosU and sinU are defined on the same grid as U, and so on for field V

• We followed this approach because usually, curvi-grid based models store the U, V velocities on still aligned to their grid

• For the moment, only UV can be rotated this way. We could think about also rotating UVstokes the same way, but using which rotation fields?

To think about what to do, we should have a better idea of what are the Stokes drift data. Are the velocities alignes in zonal/meridional or not? How big are the data? (If they are small, it would be easier to rotate correctly the data in some prepro, before any interpolation.) Is it a feature we want to deal within Parcels, or is it something we want to treat with a kernel such as `

def stokesKernel(...):

        cU = fieldset.cosUstokes[time, x, y, z]
        sU = fieldset.sinUstokes[time, x, y, z]
        cV = fieldset.cosVstokes[time, x, y, z]
        sV = fieldset.sinVstokes[time, x, y, z]
        zonal = U * cU - V * sV
        meridional = U * sU + V * cV
        p.lon += zonal * dt
        p.lat += meridional * dt

` It is important to have in mind that the code hereabove is more or less what is done with this fieldset.UV. The UV interpolation is not faster, it only enables to write much shorter the advectionRK4 kernel itself.

So to conclude, depending on the data you have (size, grid equal or not to UV, curvi or recti) and how much it should be interpolated within parcels or within the scripts calling parcels, we have to design the solution! Could you give us more information about those different points? Thanks!

[willirath]

Thanks for these comments! And thanks for pointing out that the rotation may have to be performed manually.

First of all: For the moment, we'll probably just directly rewrite the RK4 kernel. So there's no real urgency to come up with a general solution. Having said that, it would be very helpful to have an easy way of superposing different velocity fields on different grids directly in Parcels. Here's why I think this is the case:

• It allows for much faster and cheaper setup of experiments. I also have things like U_mean_flow + U_meso_scale, or U_from_ocean_model + U_inertial_oscillation etc. in mind.

• As far as I understand, fieldset.UV[...] just triggers IO for the very data necessary to interpolate the fields to the desired location. This could make things very efficient at least for a moderate number of particles.

• Having fieldset.UV return a superposition of velocity fields rather than adapting kernels limits the work necessary to maintain and debug the code: Testing that fieldset.UV in fact returns the correct data is much easier than making sure than checking for subtle bugs in, e.g., a modified version of AdvectionRK45.

[erikvansebille]

I like these ideas yes.

I guess we would then need to have a way to label a field as U-like or something. Because note that there are also a lot of scalar fields (temperature, ssh, mixed_layer_depth) that shouldn't be rotated. And what if one grid is curvilinear and the other rectilinear? Then it becomes really tricky how to handle this.

Nevertheless, I agree that simply being able to add Fields, irrespective of their underlying grid, would be nice. Not in the next version, but certainly down the line before v2.0?

[delandmeterp]

For the moment, the special treatment for UV is simply hardcoded in Parcels.

I think the solution for this new (great!) idea will be having a meta-Field UV, that is simply a list of different vector fields to superpose. (We don't always want to simply superpose them. For example a particle located at the surface can typically follow the flow field + 2% of wind velocity, but this correction factor should be taken into account in field.set_scaling_factor, in the same way as for unit conversion)

Every vector field is also a list of two horizontal fields, plus potentially 4 rotation fields.

This should not be technically too difficult to do. When we come to that development, it would be good to work on it together, since it is important to have a practical application to test it!

[willirath]

This should not be technically too difficult to do. When we come to that development, it would be good to work on it together, since it is important to have a practical application to test it!

Then, let's leave it at this collection of ideas for now. Count me in!

[erikvansebille]

This has been implemented in #393 and is explained in the tutorial at http://nbviewer.jupyter.org/github/OceanParcels/parcels/blob/master/parcels/examples/tutorial_FieldLists.ipynb