Performance issues of domain cropping

[observingClouds]

This issue collects a few performance issues with the cropping implementation of #45 that I experience when wanting to crop a 11M cell grid. The time for such a cropping is > 2 days.

Main function

For a quick overview the main function is in mllam_data_prep/ops/cropping.py create_convex_hull_mask:

    da_lon, da_lat = _get_latlon_coords(ds)
    da_lon_ref, da_lat_ref = _get_latlon_coords(ds_reference)

    assert da_lat.dims == da_lon.dims
    assert da_lat_ref.dims == da_lon_ref.dims

    # latlon to (x, y, z) on unit sphere
    da_ref_xyz = _latlon_to_unit_sphere_xyz(da_lat=da_lat_ref, da_lon=da_lon_ref)

    chull_lam = SphericalPolygon.convex_hull(da_ref_xyz.values)

    # call .load() to avoid using dask arrays in the following apply_ufunc
    da_interior_mask = xr.apply_ufunc(
        chull_lam.contains_lonlat, da_lon.load(), da_lat.load(), vectorize=True
    ).astype(bool)
    da_interior_mask.attrs[
        "long_name"
    ] = "contained in convex hull of source dataset (da_ref)"

Redudant calculations

The slow part of create_convex_hull_mask is the test of whether or not a point is part of the hull, i.e. chull_lam.contains_lonlat. While apply_ufunc might parallelize this computation, contains_lonlat is still quite inefficient.
E.g. vector.lonlat_to_vector(lon, lat, degrees=degrees) is called on each call of contains_lonlat, even though it does not change.

contains_lonlat:

    def contains_lonlat(self, lon, lat, degrees=True):
        point = vector.lonlat_to_vector(lon, lat, degrees=degrees)  # is executed for each point call, though it does not change
        return self._contains_point(point, self._points, self._inside)

[observingClouds]

Special case: original grid contains many points outside region of interest

In cases where one has a very large grid (blue) and one likes to create an mllam dataset of a relatively small region (red), one can provide this selection in the coord_ranges of the config file (e.g. datastore.interior.yaml). To extend this "interior" domain by a boundary, one provides additional information in a datastore.boundary.yaml:

  domain_cropping:
    margin_width_degrees: 0.06
    interior_dataset_config_path: datastore.interior.yaml
    include_interior_points: false

The issue now is though that each cell of the big domain (red) is checked whether or not it is part of the convex hull. It would be great if there would be a method to make a "coarse check" to remove the obvious outliers before making a more fine-grained check.

[observingClouds]

Minimum distance calculation arcs

In case the number of arcs becomes substantial (in the above case it is about 1000), the calculation in its current implementation takes over 7 hours and slows continuously down.

Solution

Vectorizing the function shortest_distance_to_arc allows to apply the function to all arcs at once resulting in a compute time of 1h.

[Jayant-kernel]

Hi @observingClouds @joeloskarsson sir,

A 2-day runtime for cropping is a legendary level of patience, but I think we can bring that down to a coffee break. I’d love to tackle this performance overhaul.

My Plan:

The "Coarse-to-Fine" Filter: Implementing a vectorized lat/lon bounding-box heuristic to instantly 'yeet' millions of outliers before they ever touch the expensive spherical hull logic.

Pre-conversion Zen: Moving the unit-sphere XYZ transformations outside the loops to treat the entire 11M grid as a single tensor pass.
Broadcasting the Arcs: Refactoring the shortest_distance_to_arc logic into an $(N, \text{arcs}, 3)$ broadcasted operation to kill the iterative overhead.

Dask-Native Scaling: Ensuring the vectorization plays nice with xr.apply_ufunc for parallelized execution.

I've got a local prototype of the vectorization logic showing promising results. Mind assigning this to me so I can finalize the PR?

[observingClouds]

Hi @Jayant-kernel,
Thanks for showing interest in this issue. Your plan sounds great. Looking forward to your PR and discuss your approach more.

[Jayant-kernel]

Hi @observingClouds, thank you for the quick response and for assigning this to me!

I have start working on the optimization in cropping.py and will focus first on the two main bottlenecks from the issue:

• reducing memory usage in the arc-distance computation
• removing the per-point overhead in convex-hull membership checks

I’ll open a PR soon with benchmarks + correctness checks against the current implementation so we can compare speedup and verify behavior stays the same.