[Refactor to use xarray] Cra refactor by tennlee · Pull Request #942 · nci/scores

tennlee · 2025-12-03T11:39:44Z

Demonstrates a possible approach to moving CRA towards using an xarray return type
Only applied to cra_2d (now cra_image) for the time being
Feedback welcome

Pre-commit checks not passing yet One test still failing Demonstrates approach to move cra_image to use an xarray data structure

Update some tests and some handling in the implementation

esteban-abellan · 2025-12-10T23:03:55Z

src/scores/spatial/cra_impl.py

 def cra(
-    fcst: xr.DataArray,
-    obs: xr.DataArray,
+    fcst: XarrayLike,


When testing these changes, the XarrayLike part failed. Did you mean ArrayLike instead? I tried ArrayLike and then adding the importing from numpy.typing import ArrayLike and it worked.

I don't have the full context in this PR, but just to add:

numpy's ArrayLike is a bit of a catch-all, we want to avoid using that if we have a more specific type-interface for that particular argument.

src/scores/spatial/cra_impl.py

Co-authored-by: esteban-abellan <esteban.abellan@gmail.com> Signed-off-by: Tennessee Leeuwenburg <134973832+tennlee@users.noreply.github.com>

Delete unnecessary code Continue refactoring

Various test fixes

Delete test case for exception no longer generated

One failing unit test, think the code needs some more tweaking for 2D image slices

Tests passing

…tation

Update CRA method to take max_distance_approx and add notes about distance calculations

reza-armuei

Thanks @esteban-abellan and @tennlee
I have added several comments, but the main points are:

1- There is an issue in how the optimal shift (that minimises the MSE) is computed. The shift_range is currently hard‑coded, which can lead to cases where the shift that yields the minimum MSE lies outside the allowed maximum shift provided by the user. In such cases, the code discards the optimal shift and returns the unshifted fcst.
However, there may still exist a valid shift within the user-specified maximum range that would reduce the MSE compared to the unshifted forecast. Defining shift_range dynamically based on the user-provided maximum allowed shift would resolve this problem.

2- Several unit tests are duplicated or test essentially the same behaviour. These could be consolidated to avoid redundancy.

3- The current unit tests mostly verify code mechanics rather than validating scientific correctness. We should add a couple of tests that check whether the function returns the expected total MSE and its decomposed components.

I will work on addressing these issues in a separate PR. Depending on my availability, I may be able to get to it soon.

src/scores/spatial/cra_impl.py

reza-armuei · 2026-02-03T00:25:00Z

src/scores/spatial/cra_impl.py

        obs (xr.DataArray): 2-D observation field.
-        threshold (float): Minimum value that a grid point must meet or exceed to be considered
+        minimum_intensity (float): Minimum value that a grid point must meet or exceed to be considered


To keep the consistency with the rest of scores, no need to include typehints in docstrings.

src/scores/spatial/cra_impl.py

reza-armuei · 2026-02-03T00:39:07Z

src/scores/spatial/cra_impl.py

    # Assign a unique label to each connected component. For instance, if there are 3 separate
    # blobs in our array, each blob will be assigned a different label (e.g., 1, 2, 3)
    labeled_array_obs, num_features_obs = scipy.ndimage.label(
        ~np.isnan(masked_obs), structure=structure
    )  # labels the connected components in the masked array where the values are not NaN
    if num_features_obs > 1:
        # Find the largest blob
        largest_blob_label_obs = np.argmax(np.bincount(labeled_array_obs.flat)[1:]) + 1

        # Create a new masked array with only the largest blob
        obs = masked_obs.where(labeled_array_obs == largest_blob_label_obs)
    else:
        obs = masked_obs

    labeled_array_fcst, num_features_fcst = scipy.ndimage.label(
        ~np.isnan(masked_fcst), structure=structure
    )  # labels the connected components in the masked array where the values are not NaN
    if num_features_fcst > 1:
        # Find the largest blob
        largest_blob_label_fcst = np.argmax(np.bincount(labeled_array_fcst.flat)[1:]) + 1

        # Create a new masked array with only the largest blob
        fcst = masked_fcst.where(labeled_array_fcst == largest_blob_label_fcst)
    else:
        fcst = masked_fcst


Since the logic for finding the largest blob is duplicated (once for observations and once for forecasts), would it make sense to extract it into a separate helper function and call it for both cases?

reza-armuei · 2026-02-03T00:43:49Z

src/scores/spatial/cra_impl.py

        fcst (xr.DataArray): Forecast field.
        obs (xr.DataArray): Observation field.
        x_name (str): Name of the zonal spatial dimension (e.g., 'x' or 'longitude').
        y_name (str): Name of the meridional spatial dimension (e.g., 'y' or 'latitude').
        max_distance (float) : Maximum distance in km allowed for the shifted blob.
        coord_units (str) : coordinates units, 'degrees' or 'metres'


Same here. No need for including typehints in docstrings.

reza-armuei · 2026-02-04T00:43:11Z

tests/spatial/test_cra.py

+def test_cra_image_shape_mismatch_raises_valueerror():
+    """Shape mismatches should raise ValueError (parity with cra_image behavior)."""
    fcst = create_array(shape=(10, 10))
    obs = create_array(shape=(8, 10))  # mismatched shape

    with pytest.raises(ValueError):
-        cra_core_2d(fcst, obs, threshold=5.0, y_name="y", x_name="x")
-
-
-def test_cra_core_2d_invalid_input_types_raise_typeerror():
-    """Non-xarray inputs should raise TypeError (parity with cra_2d behavior)."""
-    obs = create_array()
-    with pytest.raises(TypeError):
-        cra_core_2d("invalid", obs, threshold=5.0, y_name="y", x_name="x")
-
-    fcst = create_array()
-    with pytest.raises(TypeError):
-        cra_core_2d(fcst, "invalid", threshold=5.0, y_name="y", x_name="x")
+        _cra_image(fcst, obs, minimum_intensity=5.0, y_name="y", x_name="x")


I beleive this has already been tested.

reza-armuei · 2026-02-04T00:54:28Z

tests/spatial/test_cra.py

    Cover branch:
        if rmse_shifted > rmse_original or corr_shifted < corr_original or mse_shifted > original_mse:
            return None, None, None
    by monkeypatching metric functions to force the condition true.
    """


I am not sure about return None, None, None as I think the _translate_forecast_region should return un-shifted forecast with 0 values for shift in x an y for a case when shifting forecats can result in an increase in error (here based on MSE).

Suggested change

Cover branch:

if mse_shifted > mse_original or corr_shifted < corr_original or mse_shifted > original_mse:

return None, None, None

by monkeypatching metric functions to force the condition true.

"""

reza-armuei · 2026-02-04T00:58:30Z

tests/spatial/test_cra.py

+    assert np.isnan(result.mse_total)


 def test_cra_time_val_conversion_int_and_str(monkeypatch):


I am not sure if this test is required.

reza-armuei · 2026-02-04T00:59:59Z

tests/spatial/test_cra.py

    # Monkeypatch .sel to avoid KeyError when datetime64 is passed
    original_sel = xr.DataArray.sel

    def safe_sel(self, indexers=None, drop=False):


Is this used anywhere in this test?

reza-armuei · 2026-02-04T01:36:01Z

tests/spatial/test_cra.py



-def test_cra_core_2d_returns_none_when_shifted_fcst_is_none(monkeypatch):
+def test_cra_image_returns_none_when_shifted_fcst_is_none():


I think here _cra_image returns NaN because the fcst and obs fields contain fewer points than the default min_points value of 10. So I'm not sure whether this is the intended to test this behaviour?

Co-authored-by: reza-armuei <144857501+reza-armuei@users.noreply.github.com> Signed-off-by: Tennessee Leeuwenburg <134973832+tennlee@users.noreply.github.com>

esteban-abellan and others added 4 commits December 1, 2025 13:59

add refs and fix cra_2d docstring issues

d6c22b8

WIP refactor

ab76953

Pre-commit checks not passing yet One test still failing Demonstrates approach to move cra_image to use an xarray data structure

Refactoring major work done, now chasing test cases and coverage

019a974

Update notebook to use new API

558288c

Update some tests and some handling in the implementation

esteban-abellan reviewed Dec 10, 2025

View reviewed changes

src/scores/spatial/cra_impl.py Outdated Show resolved Hide resolved

tennlee and others added 4 commits December 11, 2025 10:43

Update src/scores/spatial/cra_impl.py

c8b39ed

Co-authored-by: esteban-abellan <esteban.abellan@gmail.com> Signed-off-by: Tennessee Leeuwenburg <134973832+tennlee@users.noreply.github.com>

Address pylint issues

57bf92d

Delete unnecessary code Continue refactoring

Update cra2d to expand_dims when required

5febd5f

Various test fixes

Various fixed

69c6fdb

Delete test case for exception no longer generated

tennlee changed the title ~~Cra refactor~~ [Dictionary Version] Cra refactor Jan 27, 2026

tennlee changed the title ~~[Dictionary Version] Cra refactor~~ [Refactor to use xarray] Cra refactor Jan 27, 2026

tennlee added 7 commits January 27, 2026 17:07

WIP refactoring

823ef24

Resolution of various test cases and associated refactoring

9f1779b

Further WIP.

b9f1a99

One failing unit test, think the code needs some more tweaking for 2D image slices

Refactored 2D to not take a time dimension

ba393a1

Tests passing

Update correlation coefficient calculation to use the scores implemen…

f8010da

…tation

Address pylint issues

534e7a8

Update CRA method to take max_distance_approx and add notes about distance calculations

Remove low-correlation penalty following code review

6922268

reza-armuei self-requested a review February 3, 2026 00:13

reza-armuei reviewed Feb 4, 2026

View reviewed changes

tennlee and others added 2 commits February 4, 2026 12:56

Update src/scores/spatial/cra_impl.py

b9d76c2

Co-authored-by: reza-armuei <144857501+reza-armuei@users.noreply.github.com> Signed-off-by: Tennessee Leeuwenburg <134973832+tennlee@users.noreply.github.com>

Update src/scores/spatial/cra_impl.py

0263f2f

Co-authored-by: reza-armuei <144857501+reza-armuei@users.noreply.github.com> Signed-off-by: Tennessee Leeuwenburg <134973832+tennlee@users.noreply.github.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Refactor to use xarray] Cra refactor#942

[Refactor to use xarray] Cra refactor#942
tennlee wants to merge 17 commits intocra_stagingfrom
cra_refactor

tennlee commented Dec 3, 2025

Uh oh!

esteban-abellan Dec 10, 2025

Uh oh!

nikeethr Jan 12, 2026

Uh oh!

Uh oh!

reza-armuei left a comment

Uh oh!

Uh oh!

reza-armuei Feb 3, 2026

Uh oh!

Uh oh!

reza-armuei Feb 3, 2026

Uh oh!

reza-armuei Feb 3, 2026

Uh oh!

reza-armuei Feb 4, 2026

Uh oh!

reza-armuei Feb 4, 2026

Uh oh!

reza-armuei Feb 4, 2026

Uh oh!

reza-armuei Feb 4, 2026

Uh oh!

reza-armuei Feb 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		assert np.isnan(result.mse_total)


		def test_cra_time_val_conversion_int_and_str(monkeypatch):



		def test_cra_core_2d_returns_none_when_shifted_fcst_is_none(monkeypatch):
		def test_cra_image_returns_none_when_shifted_fcst_is_none():

Conversation

tennlee commented Dec 3, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

reza-armuei left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants