Fixing an inconsistency in vertical index search between JIT and Scipy

[erikvansebille]

While working on #1816, we found a small inconsistency between the JIT and Scipy implementation of the vertical index search. This meant that a particle that was exactly at the top of a cell in JIT (zi=k, zeta=1) would be located at the bottom of the overlying cell in Scipy (zi=k+1, zeta=0).

While this does not matter for linear or nearest interpolation (the particle will get the velocities at the cell interface anyways), it does matter for C-grid interpolation because the horizontal velocities are vertically uniform over a cell; so the horizontal velocities would be different in JIT and Scipy for a particle located exactly at a depth level.

This PR makes the Scipy index-search consistent with what was already the JIT version. I think(?) it's arbitrary whether the upper (k+1) or lower (k) cell is chosen when a particle is at a depth level, and because JIT is used much more widely and in production, I decided to use the JIT-implementation also for Scipy.

[VeckoTheGecko]

Would you mind elaborating on 762d911 ? Does the first search get the wrong zi sometimes?

[erikvansebille]

Would you mind elaborating on 762d911 ? Does the first search get the wrong zi sometimes?

In some very rare cases, the zeta<0 or zeta>1; and this catches them so that 0<=zeta<=1. This is the same as we do in JIT. So these while-loops are only very rarely entered

[erikvansebille]

There's a failing mypy check; should I worry about that? Do you want to have a look for a fix, or can I merge?

[VeckoTheGecko]

Feel free to merge. The mypy checks can be worried about another time

[VeckoTheGecko]

I was doing some testing in #1816, and it looks like we're still getting failures on test_scipy_vs_jit[cgrid_velocity] on the line assert not np.isclose(pset_scipy[i].depth, z.flatten()[pset_scipy.pid[i]], atol=tol) where it doesn't look like the depth is updating for some particles.

I have pushed a commit to scipy_vertical_indexsearch_inconsistency with this modified test case

def test_scipy_vs_jit(interp_method):
    """Test that the scipy and JIT versions of the interpolation are the same."""
    variables = {"U": "U", "V": "V", "W": "W"}
    dimensions = {"time": "time", "lon": "lon", "lat": "lat", "depth": "depth"}
    fieldset = FieldSet.from_xarray_dataset(create_interpolation_data_with_land(), variables, dimensions, mesh="flat")

    for field in [fieldset.U, fieldset.V, fieldset.W]:  # Set a land point (for testing freeslip)
        field.interp_method = interp_method

    x, y, z = np.meshgrid(np.linspace(0, 1, 7), np.linspace(0, 1, 13), np.linspace(0, 1, 5))

    TestP = ScipyParticle.add_variable("pid", dtype=np.int32, initial=0)
    pset_scipy = ParticleSet(fieldset, pclass=TestP, lon=x, lat=y, depth=z, pid=np.arange(x.size))
    pset_jit = ParticleSet(fieldset, pclass=JITParticle, lon=x, lat=y, depth=z)

    def DeleteParticle(particle, fieldset, time):
        if particle.state >= 50:
            particle.delete()

    for pset in [pset_scipy, pset_jit]:
        pset.execute([AdvectionRK4_3D, DeleteParticle], runtime=4e-3, dt=1e-3)

    tol = 1e-6
+   count = 0
    for i in range(len(pset_scipy)):
        # Check that the Scipy and JIT particles are at the same location
        assert np.isclose(pset_scipy[i].lon, pset_jit[i].lon, atol=tol)
        assert np.isclose(pset_scipy[i].lat, pset_jit[i].lat, atol=tol)
        assert np.isclose(pset_scipy[i].depth, pset_jit[i].depth, atol=tol)
        # Check that the Scipy and JIT particles have moved
        assert not np.isclose(pset_scipy[i].lon, x.flatten()[pset_scipy.pid[i]], atol=tol)
        assert not np.isclose(pset_scipy[i].lat, y.flatten()[pset_scipy.pid[i]], atol=tol)
-      assert not np.isclose(pset_scipy[i].depth, z.flatten()[pset_scipy.pid[i]], atol=tol)
+      if np.isclose(pset_scipy[i].depth, z.flatten()[pset_scipy.pid[i]], atol=tol):
+            count += 1
+    print(f"{count}/{len(pset_scipy)} particles are at the same depth as the initial condition.")

And the outcome is 5/289 particles are at the same location. Initially I thought it was due to the grid geometry, and that some particles were initialized at the edges of the FieldSet, but changing those initialized positions doesn't fix things

- x, y, z = np.meshgrid(np.linspace(0, 1, 7), np.linspace(0, 1, 13), np.linspace(0, 1, 5))
+ dx = 0.2
+ x, y, z = np.meshgrid(
+     np.linspace(0 + dx, 1 - dx, 7), np.linspace(0 + dx, 1 - dx, 13), np.linspace(0 + dx, 1 - dx, 5)
+ )

Any thoughts as to what is going on here @erikvansebille ?

[erikvansebille]

Hmm, not sure what is going on here. I'll investigate in the coming days!

[erikvansebille]

I found the bug, just fixed it in bb275c3. The issue was with the 'land point' we created in the unit test. That set fieldset.W=0 for a whole column, and hence particle didn't move vertically.

[VeckoTheGecko]

I found the bug, just fixed it in bb275c3. The issue was with the 'land point' we created in the unit test. That set fieldset.W=0 for a whole column, and hence particle didn't move vertically.

I see! I'm a bit surprised that that caused it (since we're not using "nearest" interpolation), but at the same time I don't have a feel like I have a great intuition for this scenario