Refactoring field interpolation and allow custom interpolation methods in Scipy mode #1816

VeckoTheGecko · 2025-01-08T09:32:57Z

This PR refactors many of the indexing methods and interpolation methods out of field.py, and moves indexing code to a separate file to make things more manageable and reduce the coupling with the Field class.

This PR also allows users to easily overwrite the behaviour of existing interpolation methods or (untested) define new interpolation methods in Scipy mode. This can be done via the register_2d_interpolator(...) and register_3d_interpolator(...) decorators. This behaviour is in beta and is subject to change.

This also makes the interpolation functions more easily testable by providing only the required data.

Fixes #1823

VeckoTheGecko

Thoughts on the below @erikvansebille ? Let me know if you want me to break this into separate PRs

parcels/_interpolation.py

parcels/tools/statuscodes.py

parcels/_interpolation.py

VeckoTheGecko · 2025-01-16T17:28:55Z

Note for review: de89311 includes testing that shows the refactor of the 3d interpolation is 100% equivalent to the previous version for all combinations of interp method, grid types, and cell locations.

VeckoTheGecko · 2025-01-16T17:47:34Z

I'm quite aware that this is getting to be a big PR. To be expected for such a big refactor, but I think it would be good to continue in other PRs (for vector interp and indexing) to keep this reviewable.

VeckoTheGecko · 2025-01-16T18:10:44Z

tests/test_interpolation.py

+    @pytest.mark.parametrize(
+        "func, eta, xsi, expected",
+        [
+            pytest.param(interpolation._nearest_2d, 0.49, 0.49, 3.0, id="nearest_2d-1"),
+            pytest.param(interpolation._nearest_2d, 0.49, 0.51, 4.0, id="nearest_2d-2"),
+            pytest.param(interpolation._nearest_2d, 0.51, 0.49, 5.0, id="nearest_2d-3"),
+            pytest.param(interpolation._nearest_2d, 0.51, 0.51, 6.0, id="nearest_2d-4"),
+            pytest.param(interpolation._tracer_2d, None, None, 6.0, id="tracer_2d"),
+            # pytest.param(interpolation._linear_2d, ...),
+            # pytest.param(interpolation._linear_invdist_land_tracer_2d, ...),
+        ],
+    )
+    def test_2d(self, data_2d, func, eta, xsi, expected):
+        ctx = interpolation.InterpolationContext2D(data_2d, eta, xsi, self.ti, self.yi, self.xi)
+        assert func(ctx) == expected
+
+    @pytest.mark.parametrize(
+        "func, eta, xsi, expected",
+        [
+            # pytest.param(interpolation._nearest_3d, ...),
+            # pytest.param(interpolation._cgrid_velocity_3d, ...),
+            # pytest.param(interpolation._linear_invdist_land_tracer_3d, ...),
+            # pytest.param(interpolation._linear_3d, ...),
+            # pytest.param(interpolation._tracer_3d, ...),
+        ],
+    )


Thoughts on test cases @erikvansebille ? I think it might be good to add some on the raw arrays

And yes, we absolutely also need 3D tests. Perhaps again comparing to JIT mode?

VeckoTheGecko · 2025-01-17T10:35:49Z

I added an extra commit in the history that shows 100% equivalence of the refactor. See updated #1816 (comment)

erikvansebille

First bit of reviews (of the three main new files). Rest to come after the weekend

parcels/_index_search.py

parcels/_interpolation.py

parcels/_index_search.py

erikvansebille

More comments, now on all files except for the tests/* files

parcels/_interpolation.py

parcels/_index_search.py

erikvansebille · 2025-01-17T16:23:15Z

parcels/field.py

Should we not also move the _search_indices(), _search_indices_curvilinear() and _search_indices_rectilinear() methods to the _index_search.py file? Why are they still in field.py?

In general, I see that even in this PR, field.py still contains a lot of methods/functions that don't specifically need to be here. That would really clean up the field.py file?

E.g. VectorField.dist, VectorField._is_land2D() and VectorField.jacobian can go to an interpolation_utils.py file (mimicking what is in include folder for C)?

And the spatial interpolation for VectorFields can also go to _interpolation.py?

Agreed - I have pushed these changes. I haven't made any changes to VectorField, I will make those in another PR

parcels/tools/statuscodes.py

#1816 (comment)

Not needed anymore since show_time isn't in the codebase anymore #1816 (comment)

parcels/_index_search.py

parcels/_interpolation.py

parcels/field.py

erikvansebille · 2025-01-21T07:24:08Z

tests/test_interpolation.py

+            # pytest.param(interpolation._linear_2d, ...),
+            # pytest.param(interpolation._linear_invdist_land_tracer_2d, ...),


I think it's crucial that we test all interpolation methods with these simple unit tests. But I must say I don't really understand how this test-function is now set up. Why a class?

And indeed, I would not test just one point, but a few edge cases too.

I think it's crucial that we test all interpolation methods with these simple unit tests

Agreed

But I must say I don't really understand how this test-function is now set up. Why a class?

Admittedly classes aren't used much in Pytest beyond a way of grouping together tests (even then - not used often). I thought it would be a straight-forward way to group the tests that are testing on ti,zi,yi,xi = 0,1,1,1 on the data_2d and data_3d datasets.

And indeed, I would not test just one point, but a few edge cases too.

Agreed - those can be functions outside of this class. Also testing of the land interp method

erikvansebille · 2025-01-21T07:26:09Z

tests/test_interpolation.py

+    )
+    def test_2d(self, data_2d, func, eta, xsi, expected):
+        ctx = interpolation.InterpolationContext2D(data_2d, eta, xsi, self.ti, self.yi, self.xi)
+        assert func(ctx) == expected


If we don't want to hand-code all the expecteds, we could also compare to JIT mode? I know that perhaps we will remove JIT mode in the long term, but until then it would be good to check that they are consistent

Can do. I think that would require going back up to constructing the Field class, and then testing the interpolation on that.

I guess that would be the only way of doing it? I assume that the JIT functions don't follow the same structure as the new scipy interpolation functions?

Yes, I fear that the only way to test again JIT is going through the pset.execute(). So make a FieldSet, a large ParticleSet, and then call a Field-evaoluation (particle.u = fieldset.U[particle.time, particle.depth, particle.lat, particle.lon]) in a custom kernel and assert whether all particle.u are the same value as the new Scipy interpolation.

Not the cleanest code and I can imagine you're slightly disappointed to require all these extra classes in this test-function (defeats the purpose of unit tests somewhat), but by far the most robust validation that our new Scipy works. When/if we move away from JIT, we can remove all that code ;-)

Not the cleanest code and I can imagine you're slightly disappointed to require all these extra classes in this test-function

Yes, but at the end of the day its important to have proper validation and clean code can come gradually as we refactor parts of the codebase and improve :)

erikvansebille · 2025-01-21T07:26:35Z

tests/test_interpolation.py

+    @pytest.mark.parametrize(
+        "func, eta, xsi, expected",
+        [
+            pytest.param(interpolation._nearest_2d, 0.49, 0.49, 3.0, id="nearest_2d-1"),
+            pytest.param(interpolation._nearest_2d, 0.49, 0.51, 4.0, id="nearest_2d-2"),
+            pytest.param(interpolation._nearest_2d, 0.51, 0.49, 5.0, id="nearest_2d-3"),
+            pytest.param(interpolation._nearest_2d, 0.51, 0.51, 6.0, id="nearest_2d-4"),
+            pytest.param(interpolation._tracer_2d, None, None, 6.0, id="tracer_2d"),
+            # pytest.param(interpolation._linear_2d, ...),
+            # pytest.param(interpolation._linear_invdist_land_tracer_2d, ...),
+        ],
+    )
+    def test_2d(self, data_2d, func, eta, xsi, expected):
+        ctx = interpolation.InterpolationContext2D(data_2d, eta, xsi, self.ti, self.yi, self.xi)
+        assert func(ctx) == expected
+
+    @pytest.mark.parametrize(
+        "func, eta, xsi, expected",
+        [
+            # pytest.param(interpolation._nearest_3d, ...),
+            # pytest.param(interpolation._cgrid_velocity_3d, ...),
+            # pytest.param(interpolation._linear_invdist_land_tracer_3d, ...),
+            # pytest.param(interpolation._linear_3d, ...),
+            # pytest.param(interpolation._tracer_3d, ...),
+        ],
+    )


And yes, we absolutely also need 3D tests. Perhaps again comparing to JIT mode?

tests/test_interpolation.py

erikvansebille · 2025-01-29T07:32:48Z

parcels/_index_search.py

+            if gridindexingtype == "mom5" and z > 2 * grid.depth[0] - grid.depth[1]:
+                return (-1, z / grid.depth[0])
+            else:
+                _raise_field_out_of_bound_surface_error(z, 0, 0)


It's not ideal that this function passes y=0 and x=0. I understand we don't know what x and y are at this moment, but printing them as zeros can be misleading to users. Would it not be better to print not available or unknown or so? Or is there no way to somehow figure out what x and y are?

The previous code was also putting them to be 0, so this was just a continuation of that. Easy fix to set x and y to be None!

OK, and could we then print these Nones as 'unknown' in the actual warning message to the user? Would then make more sense to them?

parcels/_index_search.py

#1816 (comment)

fix

… from field.py to _index_search.py

#1816 (comment)

Not needed anymore since show_time isn't in the codebase anymore #1816 (comment)

Note that "cgrid_velocity" at the moment breaks; so there may be an issue there with the new scipy interpolation?

#1816 (comment)

VeckoTheGecko · 2025-01-29T15:07:51Z

Cleaned up now. We can either merge this first, or #1834. Whichever one comes first, I'll rebase the other

VeckoTheGecko · 2025-01-29T16:07:12Z

I encountered this

=========================== short test summary info ============================
FAILED tests/test_interpolation.py::test_scipy_vs_jit[nearest] - assert not np.True_
 +  where np.True_ = <function isclose at 0x7f9f38c37130>(np.float32(0.75000733), np.float64(0.75), atol=1e-08)
 +    where <function isclose at 0x7f9f38c37130> = np.isclose
 +    and   np.float32(0.75000733) = P[118388](lon=0.166994, lat=0.335120, depth=0.750007, pid=148.000000, time=0.003000).depth
= 1 failed, 1181 passed, 1 skipped, 8 xfailed, 4087 warnings in 687.10s (0:11:27) =

Perhaps due to #1834 due to it being in the depth dimension (though could not recreate locally)

erikvansebille · 2025-01-29T16:56:45Z

I encountered this

=========================== short test summary info ============================
FAILED tests/test_interpolation.py::test_scipy_vs_jit[nearest] - assert not np.True_
 +  where np.True_ = <function isclose at 0x7f9f38c37130>(np.float32(0.75000733), np.float64(0.75), atol=1e-08)
 +    where <function isclose at 0x7f9f38c37130> = np.isclose
 +    and   np.float32(0.75000733) = P[118388](lon=0.166994, lat=0.335120, depth=0.750007, pid=148.000000, time=0.003000).depth
= 1 failed, 1181 passed, 1 skipped, 8 xfailed, 4087 warnings in 687.10s (0:11:27) =

Perhaps due to #1834 due to it being in the depth dimension (though could not recreate locally)

Hmm, it's a 1e-6 error; it could also be because of accumulating round-off errors. I found the 1e-8 extremely tight already, but since it worked locally on my computer too I didn't change it. I think 1e-6 is still totally acceptable, so I'll see if that fixes the breaking assert

VeckoTheGecko commented Jan 8, 2025

View reviewed changes

parcels/_interpolation.py Outdated Show resolved Hide resolved

parcels/tools/statuscodes.py Show resolved Hide resolved

VeckoTheGecko commented Jan 9, 2025

View reviewed changes

parcels/_interpolation.py Outdated Show resolved Hide resolved

VeckoTheGecko requested a review from erikvansebille January 16, 2025 18:04

VeckoTheGecko changed the title ~~Refactoring indexing and interpolation~~ Refactoring field interpolation (and move indexing code) Jan 16, 2025

VeckoTheGecko commented Jan 16, 2025

View reviewed changes

VeckoTheGecko marked this pull request as ready for review January 16, 2025 18:18

VeckoTheGecko force-pushed the v/refactor-interp branch from 8abec45 to b38a17c Compare January 17, 2025 10:34

VeckoTheGecko changed the title ~~Refactoring field interpolation (and move indexing code)~~ Refactoring field interpolation and allow custom interpolation methods in Scipy mode Jan 17, 2025

erikvansebille reviewed Jan 17, 2025

View reviewed changes

VeckoTheGecko force-pushed the v/refactor-interp branch from bff8ca2 to 14b93c9 Compare January 17, 2025 15:41

VeckoTheGecko requested a review from erikvansebille January 17, 2025 16:11

erikvansebille reviewed Jan 17, 2025

View reviewed changes

parcels/_interpolation.py Outdated Show resolved Hide resolved

parcels/_interpolation.py Outdated Show resolved Hide resolved

parcels/_index_search.py Show resolved Hide resolved

erikvansebille reviewed Jan 17, 2025

View reviewed changes

VeckoTheGecko added a commit that referenced this pull request Jan 20, 2025

Remove casts to float32

2b46ba9

#1816 (comment)

VeckoTheGecko added a commit that referenced this pull request Jan 20, 2025

Remove msg from TimeExtrapolation constructor

579707c

Not needed anymore since show_time isn't in the codebase anymore #1816 (comment)

erikvansebille reviewed Jan 21, 2025

View reviewed changes

erikvansebille reviewed Jan 27, 2025

View reviewed changes

tests/test_interpolation.py Outdated Show resolved Hide resolved

erikvansebille mentioned this pull request Jan 29, 2025

Fixing an inconsistency in vertical index search between JIT and Scipy #1834

Merged

erikvansebille reviewed Jan 29, 2025

View reviewed changes

parcels/_index_search.py Show resolved Hide resolved

VeckoTheGecko added a commit that referenced this pull request Jan 29, 2025

Review feedback

6fdfc3a

#1816 (comment)

VeckoTheGecko force-pushed the v/refactor-interp branch from f4c6d83 to 973b293 Compare January 29, 2025 15:00

VeckoTheGecko added 5 commits January 29, 2025 16:05

refactor error classes removing custom inits

b4d9c2b

update to use _raise_out_of_bound_error

681f668

fix

update to use _raise_out_of_bound_error

6990c02

update to use _raise_out_of_bound_surface_error

44fc36a

Rename parse_particletime

08f560c

VeckoTheGecko and others added 21 commits January 29, 2025 16:05

Delete old 3d interpolator

78429eb

update comment

a235874

Review feedback

46fc3c3

patch indexerror

5ed3310

Rename file _indexing.py -> _index_search.py

bf79d18

Refactor test

353d220

Rename test data function

d604ae1

review feedback

b9d9242

Move calc_cell_edge_sizes, and cell_areas out of field.py

7052b26

Move methods _search_indices_curvilinear, _search_indices_rectilinear…

cb0ba68

… from field.py to _index_search.py

move tests

8a861c3

Move reconnect_bnd_indices to grid.py

77d8a8e

Remove casts to float32

76c6865

#1816 (comment)

Remove msg from TimeExtrapolation constructor

d88bce7

Not needed anymore since show_time isn't in the codebase anymore #1816 (comment)

Fix citations

7d2e762

review feedback

3854ea1

Adding a unit test to compare JIT and SciPy interpolation/integration

a0a0336

Note that "cgrid_velocity" at the moment breaks; so there may be an issue there with the new scipy interpolation?

Review edits

8ac1593

Review feedback

af8ffe8

#1816 (comment)

cleanup test_interpolation.py

d61563f

xfail cgrid_velocity on test_scipy_vs_jit

46a2853

VeckoTheGecko force-pushed the v/refactor-interp branch from fb5a593 to 46a2853 Compare January 29, 2025 15:05

VeckoTheGecko closed this Jan 29, 2025

VeckoTheGecko reopened this Jan 29, 2025

erikvansebille and others added 3 commits January 29, 2025 17:58

Relaxing jit-vs-scipy tolerance

09aec5e

Merge branch 'main' into v/refactor-interp

baaf04c

Copying #1834 changes into new _index_search functions

10a44f8

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactoring field interpolation and allow custom interpolation methods in Scipy mode #1816

Refactoring field interpolation and allow custom interpolation methods in Scipy mode #1816

VeckoTheGecko commented Jan 8, 2025 •

edited

Loading

VeckoTheGecko left a comment

VeckoTheGecko commented Jan 16, 2025 •

edited

Loading

VeckoTheGecko commented Jan 16, 2025

VeckoTheGecko Jan 16, 2025 •

edited

Loading

erikvansebille Jan 21, 2025

VeckoTheGecko commented Jan 17, 2025

erikvansebille left a comment

erikvansebille left a comment

erikvansebille Jan 17, 2025

VeckoTheGecko Jan 20, 2025

erikvansebille Jan 21, 2025

VeckoTheGecko Jan 21, 2025

erikvansebille Jan 21, 2025

VeckoTheGecko Jan 21, 2025 •

edited

Loading

erikvansebille Jan 22, 2025

VeckoTheGecko Jan 22, 2025

erikvansebille Jan 21, 2025

erikvansebille Jan 29, 2025

VeckoTheGecko Jan 29, 2025

erikvansebille Jan 29, 2025

VeckoTheGecko commented Jan 29, 2025

VeckoTheGecko commented Jan 29, 2025

erikvansebille commented Jan 29, 2025

		# pytest.param(interpolation._linear_2d, ...),
		# pytest.param(interpolation._linear_invdist_land_tracer_2d, ...),

Refactoring field interpolation and allow custom interpolation methods in Scipy mode #1816

Are you sure you want to change the base?

Refactoring field interpolation and allow custom interpolation methods in Scipy mode #1816

Conversation

VeckoTheGecko commented Jan 8, 2025 • edited Loading

VeckoTheGecko left a comment

Choose a reason for hiding this comment

VeckoTheGecko commented Jan 16, 2025 • edited Loading

VeckoTheGecko commented Jan 16, 2025

VeckoTheGecko Jan 16, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

VeckoTheGecko commented Jan 17, 2025

erikvansebille left a comment

Choose a reason for hiding this comment

erikvansebille left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

VeckoTheGecko Jan 21, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

VeckoTheGecko commented Jan 29, 2025

VeckoTheGecko commented Jan 29, 2025

erikvansebille commented Jan 29, 2025

VeckoTheGecko commented Jan 8, 2025 •

edited

Loading

VeckoTheGecko commented Jan 16, 2025 •

edited

Loading

VeckoTheGecko Jan 16, 2025 •

edited

Loading

VeckoTheGecko Jan 21, 2025 •

edited

Loading