Skip to content

Investigate antimeridian improvements and pyproj updates #645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
djhoese wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

Investigate antimeridian improvements and pyproj updates #645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
cc4cb05
Initial antimeridian handling refactoring
djhoese Jan 30, 2025
bd26af0
Revert deprecation of modify_extents
djhoese Feb 27, 2025
b18d6c4
Remove valid input/output masks in kdtree resamplers
djhoese Feb 28, 2025
06057aa
Add test for lon_wrap=180 in nearest resampler
djhoese Mar 4, 2025
2c4004b
Add antimeridian kdtree tests with antimeridian input
djhoese Mar 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 24 additions & 49 deletions pyresample/future/resamplers/nearest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,24 +41,20 @@
dask = None


def query_no_distance(target_lons, target_lats, valid_output_index,
mask=None, valid_input_index=None,
def query_no_distance(target_lons, target_lats,
mask=None,
neighbours=None, epsilon=None, radius=None,
kdtree=None):
"""Query the kdtree. No distances are returned.

NOTE: Dask array arguments must always come before other keyword arguments
for `da.blockwise` arguments to work.
"""
voi = valid_output_index
shape = voi.shape + (neighbours,)
voir = voi.ravel()
if mask is not None:
mask = mask.ravel()[valid_input_index.ravel()]
target_lons_valid = target_lons.ravel()[voir]
target_lats_valid = target_lats.ravel()[voir]
mask = mask.ravel()

coords = lonlat2xyz(target_lons_valid, target_lats_valid)
# TODO: Convert between CRSes

Check notice on line 56 in pyresample/future/resamplers/nearest.py

View check run for this annotation

codefactor.io / CodeFactor

pyresample/future/resamplers/nearest.py#L56 

Unresolved comment '# TODO: Convert between CRSes'. (C100)
coords = lonlat2xyz(target_lons.ravel(), target_lats.ravel())
distance_array, index_array = kdtree.query(
coords,
k=neighbours,
Expand All @@ -72,28 +68,21 @@
# KDTree query returns out-of-bounds neighbors as `len(arr)`
# which is an invalid index, we mask those out so -1 represents
# invalid values
# voi is 2D (trows, tcols)
# index_array is 2D (valid output pixels, neighbors)
# there are as many Trues in voi as rows in index_array
good_pixels = index_array < kdtree.n
res_ia = np.empty(shape, dtype=int)
mask = np.zeros(shape, dtype=bool)
mask[voi, :] = good_pixels
res_ia[mask] = index_array[good_pixels]
res_ia[~mask] = -1
return res_ia


def _my_index(index_arr, vii, data_arr, vii_slices=None, ia_slices=None,
fill_value=np.nan):
res_ia = index_array.astype(int, copy=False) # usually a copy since pykdtree creates uint32/64 indexes
res_ia[res_ia >= kdtree.n] = -1
return res_ia.reshape(target_lons.shape + (neighbours,))


def _my_index(index_arr, data_arr, ia_slices=None, fill_value=np.nan):
"""Wrap index logic for 'get_sample_from_neighbour_info' to be used inside dask map_blocks."""
vii_slices = tuple(
x if x is not None else vii.ravel() for x in vii_slices)
# vii_slices = tuple(
# x if x is not None else vii.ravel() for x in vii_slices)
mask_slices = tuple(
x if x is not None else (index_arr == -1) for x in ia_slices)
ia_slices = tuple(
x if x is not None else index_arr for x in ia_slices)
res = data_arr[vii_slices][ia_slices]
res = data_arr[ia_slices]
res[mask_slices] = fill_value
return res

Expand Down Expand Up @@ -156,21 +145,17 @@
"""Set up kd tree on input."""
source_lons, source_lats = self.source_geo_def.get_lonlats(
chunks=chunks)
valid_input_idx = ((source_lons >= -180) & (source_lons <= 180) & (source_lats <= 90) & (source_lats >= -90))
input_coords = lonlat2xyz(source_lons, source_lats)
input_coords = input_coords[valid_input_idx.ravel(), :]

# Build kd-tree on input
input_coords = input_coords.astype(np.float64)
input_coords = input_coords.astype(np.float64, copy=False)
delayed_kdtree = dask.delayed(KDTree, pure=True)(input_coords)
return valid_input_idx, delayed_kdtree
return delayed_kdtree

def _query_resample_kdtree(self,
resample_kdtree,
tlons,
tlats,
valid_input_index,
valid_output_index,
mask,
neighbors,
radius_of_influence,
Expand All @@ -181,12 +166,12 @@
else:
ndims = self.source_geo_def.ndim
dims = 'mn'[:ndims]
args = (mask, dims, valid_input_index, dims)
args = (mask, dims, dims)
# res.shape = rows, cols, neighbors
# j=rows, i=cols, k=neighbors, m=source rows, n=source cols
res = da.blockwise(
query_no_distance, 'jik', tlons, 'ji', tlats, 'ji',
valid_output_index, 'ji', *args, kdtree=resample_kdtree,
*args, kdtree=resample_kdtree,
neighbours=neighbors, epsilon=epsilon,
radius=radius_of_influence, dtype=np.int64,
meta=np.array((), dtype=np.int64),
Expand All @@ -208,27 +193,25 @@

# Create kd-tree
chunks = mask.chunks if mask is not None else CHUNK_SIZE
valid_input_idx, resample_kdtree = self._create_resample_kdtree(chunks=chunks)
resample_kdtree = self._create_resample_kdtree(chunks=chunks)

# TODO: Add 'chunks' keyword argument to this method and use it
target_lons, target_lats = self.target_geo_def.get_lonlats(chunks=CHUNK_SIZE)
valid_output_idx = ((target_lons >= -180) & (target_lons <= 180) & (target_lats <= 90) & (target_lats >= -90))

if mask is not None:
if mask.shape != self.source_geo_def.shape:
raise ValueError("'mask' must be the same shape as the source geo definition")
mask = mask.data
index_arr = self._query_resample_kdtree(
resample_kdtree, target_lons, target_lats, valid_input_idx,
valid_output_idx, mask,
resample_kdtree, target_lons, target_lats,
mask,
neighbors, radius_of_influence, epsilon)

return valid_input_idx, index_arr
return index_arr

def get_sample_from_neighbor_info(
self,
data,
valid_input_index,
index_array,
neighbors=1,
fill_value=np.nan):
Expand Down Expand Up @@ -278,7 +261,6 @@
"handle more than 1 neighbor yet.")
# Convert from multiple neighbor shape to 1 neighbor
ia = index_array[:, :, 0]
vii = valid_input_index

src_geo_dims = self._get_src_geo_dims()
dst_geo_dims = ('y', 'x')
Expand All @@ -287,7 +269,6 @@
# shape of the source data after we flatten the geo dimensions
flat_src_shape = []
# slice objects to index in to the source data
vii_slices = []
ia_slices = []
# whether we have seen the geo dims in our analysis
geo_handled = False
Expand All @@ -302,15 +283,13 @@
src_dim_to_ind[dim] = i
if dim in src_geo_dims and not geo_handled:
flat_src_shape.append(-1)
vii_slices.append(None) # mark for replacement
ia_slices.append(None) # mark for replacement
flat_adim.append(i)
src_adims.append(i)
dst_dims.extend(dst_geo_dims)
geo_handled = True
elif dim not in src_geo_dims:
flat_src_shape.append(data.sizes[dim])
vii_slices.append(slice(None))
ia_slices.append(slice(None))
src_adims.append(i)
dst_dims.append(dim)
Expand All @@ -322,7 +301,6 @@
# neighbors_dim = i + 3

new_data = data.data.reshape(flat_src_shape)
vii = vii.ravel()
dst_adims = [dst_dim_to_ind[dim] for dim in dst_dims]
ia_adims = [dst_dim_to_ind[dim] for dim in dst_geo_dims]
# FUTURE: when we allow more than one neighbor add neighbors dimension
Expand All @@ -336,9 +314,8 @@
res = da.blockwise(
_my_index, dst_adims,
ia, ia_adims,
vii, flat_adim,
new_data, src_adims,
vii_slices=vii_slices, ia_slices=ia_slices,
ia_slices=ia_slices,
fill_value=fill_value,
meta=np.array((), dtype=new_data.dtype),
dtype=new_data.dtype, concatenate=True)
Expand Down Expand Up @@ -410,10 +387,9 @@
internal_cache_key = (mask_hash, neighbors, radius_of_influence, epsilon)
in_int_cache = internal_cache_key in self._internal_cache
if not in_int_cache:
valid_input_index, index_arr = self._get_neighbor_info(
index_arr = self._get_neighbor_info(
mask, neighbors, radius_of_influence, epsilon)
item_to_cache = {
"valid_input_index": valid_input_index,
"index_array": index_arr,
}
self._internal_cache[internal_cache_key] = item_to_cache
Expand Down Expand Up @@ -465,7 +441,6 @@
precompute_dict = self._internal_cache[cache_key]
result = self.get_sample_from_neighbor_info(
new_data,
precompute_dict["valid_input_index"],
precompute_dict["index_array"],
fill_value=fill_value)
return self._verify_result_object_type(result, data)
Expand Down
27 changes: 16 additions & 11 deletions pyresample/geometry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1220,7 +1220,7 @@ def _get_crs_area_of_use(self, projection):
return aou

def freeze(self, lonslats=None, resolution=None, shape=None, proj_info=None,
antimeridian_mode=None):
antimeridian_mode: str = "modify_extents") -> AreaDefinition:
"""Create an AreaDefinition from this area with help of some extra info.

Parameters
Expand Down Expand Up @@ -1248,13 +1248,15 @@ def freeze(self, lonslats=None, resolution=None, shape=None, proj_info=None,
right side. However, this means that some coordinates will be
outside the coordinate space of the projection. Although most
PROJ and pyresample functionality can handle this there may be
some edge cases.
* "modify_crs": Change the prime meridian of the projection
from 0 degrees longitude to 180 degrees longitude. This has
the effect of putting the data on a continuous coordinate
some edge cases. This is the default mode.
* "modify_crs": Modify the projection definition to wrap longitude
values so they are between 0 and 360 degrees (PROJ ``lon_wrap=180``).
This mode also includes the extent modifications of the
"modify_extents" mode.
This has the effect of putting the data on a continuous coordinate
system. However, this means that comparing data resampled to
this resulting area and an area not over the anti-meridian
would be more difficult.
this resulting area and an area not using this longitude
wrapping may require extra care.
* "global_extents": Ignore the bounds of the data and use -180/180
degrees as the west and east bounds of the data. This will
generate a large output area, but with the benefit of keeping
Expand All @@ -1264,6 +1266,12 @@ def freeze(self, lonslats=None, resolution=None, shape=None, proj_info=None,

Shape parameters are ignored if the instance is created
with the `optimize_projection` flag set to True.

.. versionchanged:: 1.33
The "modify_crs" mode was changed to use "+lon_wrap=180" instead of
"+pm=180" as the CRS modification. This means extents of this mode
are now equivalent to the extents in "modify_extents".

"""
with ignore_pyproj_proj_warnings():
proj_dict = self._get_proj_dict()
Expand Down Expand Up @@ -1311,7 +1319,7 @@ def _compute_bound_centers(self, proj_dict, lonslats, antimeridian_mode):
# cross anti-meridian of projection
xmin, xmax = self._compute_new_x_corners_for_antimeridian(xarr, antimeridian_mode)
if antimeridian_mode == "modify_crs":
proj_dict.update({"pm": 180.0})
proj_dict.update({"lon_wrap": 180})
return proj_dict, (xmin, ymin, xmax, ymax)

@staticmethod
Expand All @@ -1334,9 +1342,6 @@ def _compute_new_x_corners_for_antimeridian(self, xarr, antimeridian_mode):
xmax = np.nanmax(wrapped_array)
if hasattr(wrapped_array, "compute"):
xmin, xmax = da.compute(xmin, xmax)
if antimeridian_mode == "modify_crs":
xmin -= 180
xmax -= 180
return xmin, xmax


Expand Down
Loading
Loading