Skip to content
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.

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

PI-2472: _regrid_area_weighted_array: Set order y_dim, x_dim axis to be last dimensions #3587

Closed
Changes from 1 commit
Commits
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
41 changes: 39 additions & 2 deletions lib/iris/experimental/regrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,31 @@ def _regrid_area_weighted_array(
cached_x_bounds.append(x_bounds)
cached_x_indices.append(x_indices)

# Move y_dim and x_dim to last dimensions
x_dim_orig = copy.copy(x_dim)
y_dim_orig = copy.copy(y_dim)
if x_dim is None and y_dim is None:
# e.g. a scalar point such as a vertical profile
pass
elif x_dim is not None and y_dim is None:
# test cross_section along line longitude
src_data = np.moveaxis(src_data, x_dim, -1)
x_dim = src_data.ndim - 1
elif y_dim is not None and x_dim is None:
# test cross_section along line latitude
src_data = np.moveaxis(src_data, y_dim, -1)
y_dim = src_data.ndim - 1
elif x_dim < y_dim:
src_data = np.moveaxis(src_data, x_dim, -1)
src_data = np.moveaxis(src_data, y_dim - 1, -2)
x_dim = src_data.ndim - 1
y_dim = src_data.ndim - 2
else:
src_data = np.moveaxis(src_data, x_dim, -1)
src_data = np.moveaxis(src_data, y_dim, -2)
x_dim = src_data.ndim - 1
y_dim = src_data.ndim - 2

# Create empty data array to match the new grid.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • test_regrid_transposed is the only test that has an input that is affected by np.moveaxis (the fourth loop). Would it be possible to add a test with an input that is affected by np.moveaxis for each of the other loops?
  • Since all the other tests don't require any axes to be moved, is it worth performing a check to determine whether an axis needs to be moved before moving it?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possibly. If we add a test for data ordered: [x, z, y], and like [y, x, z] I think this captures most of it. In the next PR we ensure we have x and y dims to might want to re-evaluate this next.

# Note that dtype is not preserved and that the array is
# masked to allow for regions that do not overlap.
Expand Down Expand Up @@ -577,8 +602,6 @@ def _regrid_area_weighted_array(
# Transpose weights to match dim ordering in data.
weights_shape_y = weights.shape[0]
weights_shape_x = weights.shape[1]
if x_dim is not None and y_dim is not None and x_dim < y_dim:
weights = weights.T
# Broadcast the weights array to allow numpy's ma.average
# to be called.
weights_padded_shape = [1] * data.ndim
Expand Down Expand Up @@ -608,6 +631,20 @@ def _regrid_area_weighted_array(
if not src_masked and not new_data.mask.any():
new_data = new_data.data

# Restore axis to original order
if x_dim_orig is None and y_dim_orig is None:
pass
elif x_dim_orig is not None and y_dim_orig is None:
new_data = np.moveaxis(new_data, -1, x_dim_orig)
elif y_dim_orig is not None and x_dim_orig is None:
new_data = np.moveaxis(new_data, -1, y_dim_orig)
elif x_dim_orig < y_dim_orig:
new_data = np.moveaxis(new_data, -1, x_dim_orig)
new_data = np.moveaxis(new_data, -1, y_dim_orig)
else:
new_data = np.moveaxis(new_data, -2, y_dim_orig)
new_data = np.moveaxis(new_data, -1, x_dim_orig)

return new_data


Expand Down