Speed up landice.mesh.preprocess_ais_data extrapolation

[trhille]

Speed up extrapolation step in landice.mesh.preprocess_ais_data by replacing scipy.interpolate.NearestNDInterpolator with smarter algorithm using scipy.ndimage.distance_transform_edt.

Checklist

• Document (in a comment titled Testing in this PR) any testing that was used to verify the changes

[trhille]

Testing

Using current implementation in main with antarctica_1km_2024_01_29.nc input file, we get this:

calling preprocess_ais_data
Beginning building interpolator for thk
Finished building interpolator in 9.447678919066675 seconds
Beginning interpolation for thk
Interpolation completed in 529.3699218650581 seconds
Beginning building interpolator for bheatflx
Finished building interpolator in 9.67620110407006 seconds
Beginning interpolation for bheatflx
Interpolation completed in 501.6405820630025 seconds
Beginning building interpolator for vx
Finished building interpolator in 9.535582176060416 seconds
Beginning interpolation for vx
Interpolation completed in 540.0124549659668 seconds
Beginning building interpolator for vy
Finished building interpolator in 9.551368530956097 seconds
Beginning interpolation for vy
Interpolation completed in 539.4295184250223 seconds
Beginning building interpolator for ex
Finished building interpolator in 9.488464000984095 seconds
Beginning interpolation for ex
Interpolation completed in 561.9257037510397 seconds
Beginning building interpolator for ey
Finished building interpolator in 9.515914903022349 seconds
Beginning interpolation for ey
Interpolation completed in 560.7867423400749 seconds
Beginning building interpolator for thkerr
Finished building interpolator in 9.296070468961261 seconds
Beginning interpolation for thkerr
Interpolation completed in 526.2319564099889 seconds
Beginning building interpolator for dhdt
Finished building interpolator in 8.225148725905456 seconds
Beginning interpolation for dhdt
Interpolation completed in 328.5181841789745 seconds
All interpolations completed in 4162.666724437033 seconds.

Using these changes, we get this:

calling preprocess_ais_data
Beginning nearest-fill preprocessing for thk
thk: 13301973 valid cells, 31146916 cells to fill
Nearest-fill preprocessing for thk completed in 3.007 seconds
Beginning nearest-fill preprocessing for bheatflx
bheatflx: 13313571 valid cells, 31135318 cells to fill
Nearest-fill preprocessing for bheatflx completed in 2.935 seconds
Beginning nearest-fill preprocessing for vx
vx: 13278353 valid cells, 31170536 cells to fill
Nearest-fill preprocessing for vx completed in 2.995 seconds
Beginning nearest-fill preprocessing for vy
vy: 13278353 valid cells, 31170536 cells to fill
Nearest-fill preprocessing for vy completed in 3.006 seconds
Beginning nearest-fill preprocessing for ex
ex: 13201865 valid cells, 31247024 cells to fill
Nearest-fill preprocessing for ex completed in 2.999 seconds
Beginning nearest-fill preprocessing for ey
ey: 13201861 valid cells, 31247028 cells to fill
Nearest-fill preprocessing for ey completed in 3.004 seconds
Beginning nearest-fill preprocessing for thkerr
thkerr: 13301973 valid cells, 31146916 cells to fill
Nearest-fill preprocessing for thkerr completed in 2.986 seconds
Beginning nearest-fill preprocessing for dhdt
dhdt: 11333700 valid cells, 33115189 cells to fill
Nearest-fill preprocessing for dhdt completed in 2.728 seconds
All nearest-fill preprocessing completed in 23.674 seconds.

So, a 200x speedup.
I still need to verify that these changes are producing the expected behavior.

[trhille]

I've confirmed that AIS 8–30km meshes look visually identical with these changes. Maximum differences are mostly minuscule:

thickness: 4.547473508864641e-13
bedTopography: 4.547473508864641e-13
iceMask: 8.881784197001252e-16
basalHeatFlux: 0.03157263926767409
observedSurfaceVelocityX: 3.3881317890172014e-21
observedSurfaceVelocityY: 1.3552527156068805e-20
observedSurfaceVelocityUncertainty: 1.0587911840678754e-22
observedThicknessTendency: 8.97833206677805e-08
thicknessUncertainty: 1.1368683772161603e-13

However, the maximum differences for basalHeatFlux and observedThicknessTendency are relatively much larger than for other fields. Differences appear to be along seams in the extrapolation, so I don't think this is an issue.

Left panel in each screenshot is the field on the mesh created using main; right panel is the absolute value of the difference between the two meshes.

[matthewhoffman]

Did you mean to switch to a coarser dataset? Is that accounting for any of the speedup you reported?

[trhille]

Yes, this was intentional, and no I don't think it accounts for the speedup I reported. This was inadvertently switched from 8km to 1km in this recent commit: 0b52238#diff-992c4799a18d79fe0e3d55859f4b5f326104e5af5e6b8a5b0459848aa90a84ca (and was missed because there were far too many changes in that PR). This just resets to the original 8km data set, but I have been careful to be consistent in which resolution I use between main and this branch.

[trhille]

In my most recent testing, in which I definitely used the 8km source dataset, I get All nearest-fill preprocessing completed in 0.288 seconds for this branch and All interpolations completed in 11.316709512029774 seconds. on main. So still a ~2 order of magnitude speedup when using a coarse data set, but not so impactful when just cutting of 11 seconds instead of 4100 seconds :).