diff --git a/docs/gallery_code/general/plot_custom_aggregation.py b/docs/gallery_code/general/plot_custom_aggregation.py index 5fba3669b6..6ef6075fb3 100644 --- a/docs/gallery_code/general/plot_custom_aggregation.py +++ b/docs/gallery_code/general/plot_custom_aggregation.py @@ -72,7 +72,7 @@ def main(): # Make an aggregator from the user function. SPELL_COUNT = Aggregator( - "spell_count", count_spells, units_func=lambda units: 1 + "spell_count", count_spells, units_func=lambda units, **kwargs: 1 ) # Define the parameters of the test. diff --git a/docs/src/userguide/cube_statistics.rst b/docs/src/userguide/cube_statistics.rst index 08297c2a51..9dc21f91b5 100644 --- a/docs/src/userguide/cube_statistics.rst +++ b/docs/src/userguide/cube_statistics.rst @@ -14,7 +14,7 @@ Cube Statistics Collapsing Entire Data Dimensions --------------------------------- -.. testsetup:: +.. testsetup:: collapsing import iris filename = iris.sample_data_path('uk_hires.pp') @@ -125,7 +125,7 @@ in order to calculate the area of the grid boxes:: These areas can now be passed to the ``collapsed`` method as weights: -.. doctest:: +.. doctest:: collapsing >>> new_cube = cube.collapsed(['grid_longitude', 'grid_latitude'], iris.analysis.MEAN, weights=grid_areas) >>> print(new_cube) @@ -141,8 +141,8 @@ These areas can now be passed to the ``collapsed`` method as weights: altitude - x Scalar coordinates: forecast_reference_time 2009-11-19 04:00:00 - grid_latitude 1.5145501 degrees, bound=(0.14430022, 2.8848) degrees - grid_longitude 358.74948 degrees, bound=(357.494, 360.00497) degrees + grid_latitude 1.5145501 degrees, bound=(0.13755022, 2.89155) degrees + grid_longitude 358.74948 degrees, bound=(357.48724, 360.01172) degrees surface_altitude 399.625 m, bound=(-14.0, 813.25) m Cell methods: mean grid_longitude, grid_latitude @@ -155,6 +155,50 @@ Several examples of area averaging exist in the gallery which may be of interest including an example on taking a :ref:`global area-weighted mean `. +In addition to plain arrays, weights can also be given as cubes or (names of) +:meth:`~iris.cube.Cube.coords`, :meth:`~iris.cube.Cube.cell_measures`, or +:meth:`~iris.cube.Cube.ancillary_variables`. +This has the advantage of correct unit handling, e.g., for area-weighted sums +the units of the resulting cube are multiplied by an area unit: + +.. doctest:: collapsing + + >>> from iris.coords import CellMeasure + >>> cell_areas = CellMeasure( + ... grid_areas, + ... standard_name='cell_area', + ... units='m2', + ... measure='area', + ... ) + >>> cube.add_cell_measure(cell_areas, (0, 1, 2, 3)) + >>> area_weighted_sum = cube.collapsed( + ... ['grid_longitude', 'grid_latitude'], + ... iris.analysis.SUM, + ... weights='cell_area' + ... ) + >>> print(area_weighted_sum) + air_potential_temperature / (m2.K) (time: 3; model_level_number: 7) + Dimension coordinates: + time x - + model_level_number - x + Auxiliary coordinates: + forecast_period x - + level_height - x + sigma - x + Derived coordinates: + altitude - x + Scalar coordinates: + forecast_reference_time 2009-11-19 04:00:00 + grid_latitude 1.5145501 degrees, bound=(0.13755022, 2.89155) degrees + grid_longitude 358.74948 degrees, bound=(357.48724, 360.01172) degrees + surface_altitude 399.625 m, bound=(-14.0, 813.25) m + Cell methods: + sum grid_longitude, grid_latitude + Attributes: + STASH m01s00i004 + source 'Data from Met Office Unified Model' + um_version '7.3' + .. _cube-statistics-aggregated-by: Partially Reducing Data Dimensions @@ -338,3 +382,44 @@ from jja-2006 to jja-2010: mam 2010 jja 2010 +Moreover, :meth:`Cube.aggregated_by ` supports +weighted aggregation. +For example, this is helpful for an aggregation over a monthly time +coordinate that consists of months with different numbers of days. +Similar to :meth:`Cube.collapsed `, weights can be +given as arrays, cubes, or as (names of) :meth:`~iris.cube.Cube.coords`, +:meth:`~iris.cube.Cube.cell_measures`, or +:meth:`~iris.cube.Cube.ancillary_variables`. +When weights are not given as arrays, units are correctly handled for weighted +sums, i.e., the original unit of the cube is multiplied by the units of the +weights. +The following example shows a weighted sum (notice the change of the units): + +.. doctest:: aggregation + + >>> from iris.coords import AncillaryVariable + >>> time_weights = AncillaryVariable( + ... cube.coord("time").bounds[:, 1] - cube.coord("time").bounds[:, 0], + ... long_name="Time Weights", + ... units="hours", + ... ) + >>> cube.add_ancillary_variable(time_weights, 0) + >>> seasonal_sum = cube.aggregated_by("clim_season", iris.analysis.SUM, weights="Time Weights") + >>> print(seasonal_sum) + surface_temperature / (3600 s.K) (-- : 4; latitude: 18; longitude: 432) + Dimension coordinates: + latitude - x - + longitude - - x + Auxiliary coordinates: + clim_season x - - + forecast_reference_time x - - + season_year x - - + time x - - + Scalar coordinates: + forecast_period 0 hours + Cell methods: + mean month, year + sum clim_season + Attributes: + Conventions 'CF-1.5' + STASH m01s00i024 diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 592c713b4d..b34816e616 100644 --- a/docs/src/whatsnew/latest.rst +++ b/docs/src/whatsnew/latest.rst @@ -39,6 +39,12 @@ This document explains the changes made to Iris for this release #. `@rcomer`_ enabled lazy evaluation of :obj:`~iris.analysis.RMS` calcuations with weights. (:pull:`5017`) +#. `@schlunma`_ allowed the usage of cubes, coordinates, cell measures, or + ancillary variables as weights for cube aggregations + (:meth:`iris.cube.Cube.collapsed`, :meth:`iris.cube.Cube.aggregated_by`, and + :meth:`iris.cube.Cube.rolling_window`). This automatically adapts cube units + if necessary. (:pull:`5084`) + 🐛 Bugs Fixed ============= diff --git a/lib/iris/analysis/__init__.py b/lib/iris/analysis/__init__.py index c80469c3c7..173487cfb0 100644 --- a/lib/iris/analysis/__init__.py +++ b/lib/iris/analysis/__init__.py @@ -39,8 +39,10 @@ from collections.abc import Iterable import functools from functools import wraps +from inspect import getfullargspec import warnings +from cf_units import Unit import dask.array as da import numpy as np import numpy.ma as ma @@ -55,7 +57,9 @@ ) from iris.analysis._regrid import CurvilinearRegridder, RectilinearRegridder import iris.coords +from iris.coords import _DimensionalMetadata from iris.exceptions import LazyAggregatorError +import iris.util __all__ = ( "Aggregator", @@ -467,11 +471,13 @@ def __init__( Kwargs: * units_func (callable): - | *Call signature*: (units) + | *Call signature*: (units, \**kwargs) If provided, called to convert a cube's units. Returns an :class:`cf_units.Unit`, or a value that can be made into one. + To ensure backwards-compatibility, also accepts a callable with + call signature (units). * lazy_func (callable or None): An alternative to :data:`call_func` implementing a lazy @@ -479,7 +485,8 @@ def __init__( main operation, but should raise an error in unhandled cases. Additional kwargs:: - Passed through to :data:`call_func` and :data:`lazy_func`. + Passed through to :data:`call_func`, :data:`lazy_func`, and + :data:`units_func`. Aggregators are used by cube aggregation methods such as :meth:`~iris.cube.Cube.collapsed` and @@ -625,7 +632,11 @@ def update_metadata(self, cube, coords, **kwargs): """ # Update the units if required. if self.units_func is not None: - cube.units = self.units_func(cube.units) + argspec = getfullargspec(self.units_func) + if argspec.varkw is None: # old style + cube.units = self.units_func(cube.units) + else: # new style (preferred) + cube.units = self.units_func(cube.units, **kwargs) def post_process(self, collapsed_cube, data_result, coords, **kwargs): """ @@ -693,13 +704,13 @@ class PercentileAggregator(_Aggregator): """ def __init__(self, units_func=None, **kwargs): - """ + r""" Create a percentile aggregator. Kwargs: * units_func (callable): - | *Call signature*: (units) + | *Call signature*: (units, \**kwargs) If provided, called to convert a cube's units. Returns an :class:`cf_units.Unit`, or a @@ -934,13 +945,13 @@ class WeightedPercentileAggregator(PercentileAggregator): """ def __init__(self, units_func=None, lazy_func=None, **kwargs): - """ + r""" Create a weighted percentile aggregator. Kwargs: * units_func (callable): - | *Call signature*: (units) + | *Call signature*: (units, \**kwargs) If provided, called to convert a cube's units. Returns an :class:`cf_units.Unit`, or a @@ -1172,8 +1183,112 @@ def post_process(self, collapsed_cube, data_result, coords, **kwargs): return result +class _Weights(np.ndarray): + """Class for handling weights for weighted aggregation. + + This subclasses :class:`numpy.ndarray`; thus, all methods and properties of + :class:`numpy.ndarray` (e.g., `shape`, `ndim`, `view()`, etc.) are + available. + + Details on subclassing :class:`numpy.ndarray` are given here: + https://numpy.org/doc/stable/user/basics.subclassing.html + + """ + + def __new__(cls, weights, cube, units=None): + """Create class instance. + + Args: + + * weights (Cube, string, _DimensionalMetadata, array-like): + If given as a :class:`iris.cube.Cube`, use its data and units. If + given as a :obj:`str` or :class:`iris.coords._DimensionalMetadata`, + assume this is (the name of) a + :class:`iris.coords._DimensionalMetadata` object of the cube (i.e., + one of :meth:`iris.cube.Cube.coords`, + :meth:`iris.cube.Cube.cell_measures`, or + :meth:`iris.cube.Cube.ancillary_variables`). If given as an + array-like object, use this directly and assume units of `1`. If + `units` is given, ignore all units derived above and use the ones + given by `units`. + * cube (Cube): + Input cube for aggregation. If weights is given as :obj:`str` or + :class:`iris.coords._DimensionalMetadata`, try to extract the + :class:`iris.coords._DimensionalMetadata` object and corresponding + dimensional mappings from this cube. Otherwise, this argument is + ignored. + * units (string, Unit): + If ``None``, use units derived from `weights`. Otherwise, overwrite + the units derived from `weights` and use `units`. + + """ + # `weights` is a cube + # Note: to avoid circular imports of Cube we use duck typing using the + # "hasattr" syntax here + # --> Extract data and units from cube + if hasattr(weights, "add_aux_coord"): + obj = np.asarray(weights.data).view(cls) + obj.units = weights.units + + # `weights`` is a string or _DimensionalMetadata object + # --> Extract _DimensionalMetadata object from cube, broadcast it to + # correct shape using the corresponding dimensional mapping, and use + # its data and units + elif isinstance(weights, (str, _DimensionalMetadata)): + dim_metadata = cube._dimensional_metadata(weights) + arr = dim_metadata._values + if dim_metadata.shape != cube.shape: + arr = iris.util.broadcast_to_shape( + arr, + cube.shape, + dim_metadata.cube_dims(cube), + ) + obj = np.asarray(arr).view(cls) + obj.units = dim_metadata.units + + # Remaining types (e.g., np.ndarray): try to convert to ndarray. + else: + obj = np.asarray(weights).view(cls) + obj.units = Unit("1") + + # Overwrite units from units argument if necessary + if units is not None: + obj.units = units + + return obj + + def __array_finalize__(self, obj): + """See https://numpy.org/doc/stable/user/basics.subclassing.html. + + Note + ---- + `obj` cannot be `None` here since ``_Weights.__new__`` does not call + ``super().__new__`` explicitly. + + """ + self.units = getattr(obj, "units", Unit("1")) + + @classmethod + def update_kwargs(cls, kwargs, cube): + """Update ``weights`` keyword argument in-place. + + Args: + + * kwargs (dict): + Keyword arguments that will be updated in-place if a `weights` + keyword is present which is not ``None``. + * cube (Cube): + Input cube for aggregation. If weights is given as :obj:`str`, try + to extract a cell measure with the corresponding name from this + cube. Otherwise, this argument is ignored. + + """ + if kwargs.get("weights") is not None: + kwargs["weights"] = cls(kwargs["weights"], cube) + + def create_weighted_aggregator_fn(aggregator_fn, axis, **kwargs): - """Return an aggregator function that can explicitely handle weights. + """Return an aggregator function that can explicitly handle weights. Args: @@ -1398,7 +1513,7 @@ def _weighted_quantile_1D(data, weights, quantiles, **kwargs): array or float. Calculated quantile values (set to np.nan wherever sum of weights is zero or masked) """ - # Return np.nan if no useable points found + # Return np.nan if no usable points found if np.isclose(weights.sum(), 0.0) or ma.is_masked(weights.sum()): return np.resize(np.array(np.nan), len(quantiles)) # Sort the data @@ -1535,7 +1650,7 @@ def _proportion(array, function, axis, **kwargs): # Otherwise, it is possible for numpy to return a masked array that has # a dtype for its data that is different to the dtype of the fill-value, # which can cause issues outside this function. - # Reference - tests/unit/analyis/test_PROPORTION.py Test_masked.test_ma + # Reference - tests/unit/analysis/test_PROPORTION.py Test_masked.test_ma numerator = _count(array, axis=axis, function=function, **kwargs) result = ma.asarray(numerator / total_non_masked) @@ -1630,6 +1745,18 @@ def _sum(array, **kwargs): return rvalue +def _sum_units_func(units, **kwargs): + """Multiply original units with weight units if possible.""" + weights = kwargs.get("weights") + if weights is None: # no weights given or weights are None + result = units + elif hasattr(weights, "units"): # weights are _Weights + result = units * weights.units + else: # weights are regular np.ndarrays + result = units + return result + + def _peak(array, **kwargs): def column_segments(column): nan_indices = np.where(np.isnan(column))[0] @@ -1745,7 +1872,7 @@ def interp_order(length): COUNT = Aggregator( "count", _count, - units_func=lambda units: 1, + units_func=lambda units, **kwargs: 1, lazy_func=_build_dask_mdtol_function(_count), ) """ @@ -1777,7 +1904,7 @@ def interp_order(length): MAX_RUN = Aggregator( None, iris._lazy_data.non_lazy(_lazy_max_run), - units_func=lambda units: 1, + units_func=lambda units, **kwargs: 1, lazy_func=_build_dask_mdtol_function(_lazy_max_run), ) """ @@ -2021,7 +2148,11 @@ def interp_order(length): """ -PROPORTION = Aggregator("proportion", _proportion, units_func=lambda units: 1) +PROPORTION = Aggregator( + "proportion", + _proportion, + units_func=lambda units, **kwargs: 1, +) """ An :class:`~iris.analysis.Aggregator` instance that calculates the proportion, as a fraction, of :class:`~iris.cube.Cube` data occurrences @@ -2122,6 +2253,7 @@ def interp_order(length): SUM = WeightedAggregator( "sum", _sum, + units_func=_sum_units_func, lazy_func=_build_dask_mdtol_function(_sum), ) """ @@ -2159,7 +2291,7 @@ def interp_order(length): VARIANCE = Aggregator( "variance", ma.var, - units_func=lambda units: units * units, + units_func=lambda units, **kwargs: units * units, lazy_func=_build_dask_mdtol_function(da.var), ddof=1, ) @@ -2801,7 +2933,7 @@ def __init__(self, mdtol=1): Both sourge and target cubes must have an XY grid defined by separate X and Y dimensions with dimension coordinates. All of the XY dimension coordinates must also be bounded, and have - the same cooordinate system. + the same coordinate system. """ if not (0 <= mdtol <= 1): diff --git a/lib/iris/cube.py b/lib/iris/cube.py index abe37c35fb..fcd7b5b828 100644 --- a/lib/iris/cube.py +++ b/lib/iris/cube.py @@ -28,6 +28,7 @@ import iris._lazy_data as _lazy import iris._merge import iris.analysis +from iris.analysis import _Weights from iris.analysis.cartography import wrap_lons import iris.analysis.maths import iris.aux_factory @@ -3721,9 +3722,15 @@ def collapsed(self, coords, aggregator, **kwargs): sum :data:`~iris.analysis.SUM`. Weighted aggregations support an optional *weights* keyword argument. - If set, this should be supplied as an array of weights whose shape - matches the cube. Values for latitude-longitude area weights may be - calculated using :func:`iris.analysis.cartography.area_weights`. + If set, this can be supplied as an array, cube, or (names of) + :meth:`~iris.cube.Cube.coords`, :meth:`~iris.cube.Cube.cell_measures`, + or :meth:`~iris.cube.Cube.ancillary_variables`. In all cases, the + weights should be 1d (for collapsing over a 1d coordinate) or match the + shape of the cube. When weights are not given as arrays, units are + correctly handled for weighted sums, i.e., the original unit of the + cube is multiplied by the units of the weights. Values for + latitude-longitude area weights may be calculated using + :func:`iris.analysis.cartography.area_weights`. Some Iris aggregators support "lazy" evaluation, meaning that cubes resulting from this method may represent data arrays which are @@ -3802,6 +3809,10 @@ def collapsed(self, coords, aggregator, **kwargs): cube.collapsed(['latitude', 'longitude'], iris.analysis.VARIANCE) """ + # Update weights kwargs (if necessary) to handle different types of + # weights + _Weights.update_kwargs(kwargs, self) + # Convert any coordinate names to coordinates coords = self._as_list_of_coords(coords) @@ -3970,10 +3981,14 @@ def aggregated_by( also be supplied. These include :data:`~iris.analysis.MEAN` and :data:`~iris.analysis.SUM`. - Weighted aggregations support an optional *weights* keyword argument. If - set, this should be supplied as an array of weights whose shape matches - the cube or as 1D array whose length matches the dimension over which is - aggregated. + Weighted aggregations support an optional *weights* keyword argument. + If set, this can be supplied as an array, cube, or (names of) + :meth:`~iris.cube.Cube.coords`, :meth:`~iris.cube.Cube.cell_measures`, + or :meth:`~iris.cube.Cube.ancillary_variables`. In all cases, the + weights should be 1d or match the shape of the cube. When weights are + not given as arrays, units are correctly handled for weighted sums, + i.e., the original unit of the cube is multiplied by the units of the + weights. Parameters ---------- @@ -4032,6 +4047,10 @@ def aggregated_by( STASH m01s00i024 """ + # Update weights kwargs (if necessary) to handle different types of + # weights + _Weights.update_kwargs(kwargs, self) + groupby_coords = [] dimension_to_groupby = None @@ -4070,10 +4089,16 @@ def aggregated_by( f"that is aggregated, got {len(weights):d}, expected " f"{self.shape[dimension_to_groupby]:d}" ) - weights = iris.util.broadcast_to_shape( - weights, - self.shape, - (dimension_to_groupby,), + + # iris.util.broadcast_to_shape does not preserve _Weights type + weights = _Weights( + iris.util.broadcast_to_shape( + weights, + self.shape, + (dimension_to_groupby,), + ), + self, + units=weights.units, ) if weights.shape != self.shape: raise ValueError( @@ -4289,8 +4314,11 @@ def rolling_window(self, coord, aggregator, window, **kwargs): * kwargs: Aggregator and aggregation function keyword arguments. The weights - argument to the aggregator, if any, should be a 1d array with the - same length as the chosen window. + argument to the aggregator, if any, should be a 1d array, cube, or + (names of) :meth:`~iris.cube.Cube.coords`, + :meth:`~iris.cube.Cube.cell_measures`, or + :meth:`~iris.cube.Cube.ancillary_variables` with the same length as + the chosen window. Returns: :class:`iris.cube.Cube`. @@ -4358,6 +4386,10 @@ def rolling_window(self, coord, aggregator, window, **kwargs): possible windows of size 3 from the original cube. """ + # Update weights kwargs (if necessary) to handle different types of + # weights + _Weights.update_kwargs(kwargs, self) + coord = self._as_list_of_coords(coord)[0] if getattr(coord, "circular", False): @@ -4459,8 +4491,14 @@ def rolling_window(self, coord, aggregator, window, **kwargs): "as the window." ) kwargs = dict(kwargs) - kwargs["weights"] = iris.util.broadcast_to_shape( - weights, rolling_window_data.shape, (dimension + 1,) + + # iris.util.broadcast_to_shape does not preserve _Weights type + kwargs["weights"] = _Weights( + iris.util.broadcast_to_shape( + weights, rolling_window_data.shape, (dimension + 1,) + ), + self, + units=weights.units, ) data_result = aggregator.aggregate( rolling_window_data, axis=dimension + 1, **kwargs diff --git a/lib/iris/tests/test_aggregate_by.py b/lib/iris/tests/test_aggregate_by.py index 90bf0e5d4e..e5614f6b63 100644 --- a/lib/iris/tests/test_aggregate_by.py +++ b/lib/iris/tests/test_aggregate_by.py @@ -413,6 +413,30 @@ def test_single(self): aggregateby_cube.data, self.single_rms_expected ) + def test_str_aggregation_single_weights_none(self): + # mean group-by with single coordinate name. + aggregateby_cube = self.cube_single.aggregated_by( + "height", iris.analysis.MEAN, weights=None + ) + self.assertCML( + aggregateby_cube, ("analysis", "aggregated_by", "single.cml") + ) + np.testing.assert_almost_equal( + aggregateby_cube.data, self.single_expected + ) + + def test_coord_aggregation_single_weights_none(self): + # mean group-by with single coordinate. + aggregateby_cube = self.cube_single.aggregated_by( + self.coord_z_single, iris.analysis.MEAN, weights=None + ) + self.assertCML( + aggregateby_cube, ("analysis", "aggregated_by", "single.cml") + ) + np.testing.assert_almost_equal( + aggregateby_cube.data, self.single_expected + ) + def test_weighted_single(self): # weighted mean group-by with single coordinate name. aggregateby_cube = self.cube_single.aggregated_by( @@ -1328,5 +1352,153 @@ def test_weights_fail_with_non_weighted_aggregator(self): ) +# Simply redo the tests of TestAggregateBy with other cubes as weights +# Note: other weights types (e.g., coordinates, cell measures, etc.) are not +# tested this way here since this would require adding dimensional metadata +# objects to the cubes, which would change the CMLs of all resulting cubes of +# TestAggregateBy. + + +class TestAggregateByWeightedByCube(TestAggregateBy): + def setUp(self): + super().setUp() + + self.weights_single = self.cube_single[:, 0, 0].copy( + self.weights_single + ) + self.weights_single.units = "m2" + self.weights_multi = self.cube_multi[:, 0, 0].copy(self.weights_multi) + self.weights_multi.units = "m2" + + def test_str_aggregation_weighted_sum_single(self): + aggregateby_cube = self.cube_single.aggregated_by( + "height", + iris.analysis.SUM, + weights=self.weights_single, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_coord_aggregation_weighted_sum_single(self): + aggregateby_cube = self.cube_single.aggregated_by( + self.coord_z_single, + iris.analysis.SUM, + weights=self.weights_single, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_str_aggregation_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + ["height", "level"], + iris.analysis.SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_str_aggregation_rev_order_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + ["level", "height"], + iris.analysis.SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_coord_aggregation_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + [self.coord_z1_multi, self.coord_z2_multi], + iris.analysis.SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_coord_aggregation_rev_order_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + [self.coord_z2_multi, self.coord_z1_multi], + iris.analysis.SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + +class TestAggregateByWeightedByObj(tests.IrisTest): + def setUp(self): + self.dim_coord = iris.coords.DimCoord( + [0, 1, 2], standard_name="latitude", units="degrees" + ) + self.aux_coord = iris.coords.AuxCoord( + [0, 1, 1], long_name="auxcoord", units="kg" + ) + self.cell_measure = iris.coords.CellMeasure( + [0, 0, 0], standard_name="cell_area", units="m2" + ) + self.ancillary_variable = iris.coords.AncillaryVariable( + [1, 1, 1], var_name="ancvar", units="kg" + ) + self.cube = iris.cube.Cube( + [1, 2, 3], + standard_name="air_temperature", + units="K", + dim_coords_and_dims=[(self.dim_coord, 0)], + aux_coords_and_dims=[(self.aux_coord, 0)], + cell_measures_and_dims=[(self.cell_measure, 0)], + ancillary_variables_and_dims=[(self.ancillary_variable, 0)], + ) + + def test_weighting_with_str_dim_coord(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights="latitude" + ) + np.testing.assert_array_equal(res_cube.data, [0, 8]) + self.assertEqual(res_cube.units, "K degrees") + + def test_weighting_with_str_aux_coord(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights="auxcoord" + ) + np.testing.assert_array_equal(res_cube.data, [0, 5]) + self.assertEqual(res_cube.units, "K kg") + + def test_weighting_with_str_cell_measure(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights="cell_area" + ) + np.testing.assert_array_equal(res_cube.data, [0, 0]) + self.assertEqual(res_cube.units, "K m2") + + def test_weighting_with_str_ancillary_variable(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights="ancvar" + ) + np.testing.assert_array_equal(res_cube.data, [1, 5]) + self.assertEqual(res_cube.units, "K kg") + + def test_weighting_with_dim_coord(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights=self.dim_coord + ) + np.testing.assert_array_equal(res_cube.data, [0, 8]) + self.assertEqual(res_cube.units, "K degrees") + + def test_weighting_with_aux_coord(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights=self.aux_coord + ) + np.testing.assert_array_equal(res_cube.data, [0, 5]) + self.assertEqual(res_cube.units, "K kg") + + def test_weighting_with_cell_measure(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights=self.cell_measure + ) + np.testing.assert_array_equal(res_cube.data, [0, 0]) + self.assertEqual(res_cube.units, "K m2") + + def test_weighting_with_ancillary_variable(self): + res_cube = self.cube.aggregated_by( + "auxcoord", iris.analysis.SUM, weights=self.ancillary_variable + ) + np.testing.assert_array_equal(res_cube.data, [1, 5]) + self.assertEqual(res_cube.units, "K kg") + + if __name__ == "__main__": unittest.main() diff --git a/lib/iris/tests/test_analysis.py b/lib/iris/tests/test_analysis.py index e0a5d0971e..0717368d98 100644 --- a/lib/iris/tests/test_analysis.py +++ b/lib/iris/tests/test_analysis.py @@ -12,6 +12,7 @@ import dask.array as da import numpy as np import numpy.ma as ma +import pytest import iris import iris.analysis.cartography @@ -288,7 +289,7 @@ def test_weighted_mean(self): iris.analysis.MEAN, ) - # Test collpasing of non data coord + # Test collapsing of non data coord self.assertRaises( iris.exceptions.CoordinateCollapseError, e.collapsed, @@ -1702,5 +1703,237 @@ def test_weights_in_kwargs(self): self.assertEqual(kwargs, {"test_kwarg": "test", "weights": "ignored"}) +class TestWeights: + @pytest.fixture(autouse=True) + def setup_test_data(self): + self.lat = iris.coords.DimCoord( + [0, 1], standard_name="latitude", units="degrees" + ) + self.lon = iris.coords.DimCoord( + [0, 1, 2], standard_name="longitude", units="degrees" + ) + self.cell_measure = iris.coords.CellMeasure( + np.arange(6).reshape(2, 3), standard_name="cell_area", units="m2" + ) + self.aux_coord = iris.coords.AuxCoord( + [3, 4], long_name="auxcoord", units="s" + ) + self.ancillary_variable = iris.coords.AncillaryVariable( + [5, 6, 7], var_name="ancvar", units="kg" + ) + self.cube = iris.cube.Cube( + np.arange(6).reshape(2, 3), + standard_name="air_temperature", + units="K", + dim_coords_and_dims=[(self.lat, 0), (self.lon, 1)], + aux_coords_and_dims=[(self.aux_coord, 0)], + cell_measures_and_dims=[(self.cell_measure, (0, 1))], + ancillary_variables_and_dims=[(self.ancillary_variable, 1)], + ) + + def test_init_with_weights(self): + weights = iris.analysis._Weights([], self.cube) + new_weights = iris.analysis._Weights(weights, self.cube) + assert isinstance(new_weights, iris.analysis._Weights) + assert new_weights is not weights + np.testing.assert_array_equal(new_weights, []) + assert new_weights.units == "1" + assert weights.units == "1" + + def test_init_with_weights_and_units(self): + weights = iris.analysis._Weights([], self.cube) + new_weights = iris.analysis._Weights(weights, self.cube, units="J") + assert isinstance(new_weights, iris.analysis._Weights) + assert new_weights is not weights + np.testing.assert_array_equal(new_weights, []) + assert new_weights.units == "J" + assert weights.units == "1" + + def test_init_with_cube(self): + weights = iris.analysis._Weights(self.cube, self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.arange(6).reshape(2, 3)) + assert weights.units == "K" + + def test_init_with_cube_and_units(self): + weights = iris.analysis._Weights(self.cube, self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.arange(6).reshape(2, 3)) + assert weights.units == "J" + + def test_init_with_str_dim_coord(self): + weights = iris.analysis._Weights("latitude", self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[0, 0, 0], [1, 1, 1]]) + assert weights.units == "degrees" + + def test_init_with_str_dim_coord_and_units(self): + weights = iris.analysis._Weights("latitude", self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[0, 0, 0], [1, 1, 1]]) + assert weights.units == "J" + + def test_init_with_str_aux_coord(self): + weights = iris.analysis._Weights("auxcoord", self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[3, 3, 3], [4, 4, 4]]) + assert weights.units == "s" + + def test_init_with_str_aux_coord_and_units(self): + weights = iris.analysis._Weights("auxcoord", self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[3, 3, 3], [4, 4, 4]]) + assert weights.units == "J" + + def test_init_with_str_ancillary_variable(self): + weights = iris.analysis._Weights("ancvar", self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[5, 6, 7], [5, 6, 7]]) + assert weights.units == "kg" + + def test_init_with_str_ancillary_variable_and_units(self): + weights = iris.analysis._Weights("ancvar", self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[5, 6, 7], [5, 6, 7]]) + assert weights.units == "J" + + def test_init_with_str_cell_measure(self): + weights = iris.analysis._Weights("cell_area", self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.arange(6).reshape(2, 3)) + assert weights.units == "m2" + + def test_init_with_str_cell_measure_and_units(self): + weights = iris.analysis._Weights("cell_area", self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.arange(6).reshape(2, 3)) + assert weights.units == "J" + + def test_init_with_dim_coord(self): + weights = iris.analysis._Weights(self.lat, self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[0, 0, 0], [1, 1, 1]]) + assert weights.units == "degrees" + + def test_init_with_dim_coord_and_units(self): + weights = iris.analysis._Weights(self.lat, self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[0, 0, 0], [1, 1, 1]]) + assert weights.units == "J" + + def test_init_with_aux_coord(self): + weights = iris.analysis._Weights(self.aux_coord, self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[3, 3, 3], [4, 4, 4]]) + assert weights.units == "s" + + def test_init_with_aux_coord_and_units(self): + weights = iris.analysis._Weights(self.aux_coord, self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[3, 3, 3], [4, 4, 4]]) + assert weights.units == "J" + + def test_init_with_ancillary_variable(self): + weights = iris.analysis._Weights(self.ancillary_variable, self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[5, 6, 7], [5, 6, 7]]) + assert weights.units == "kg" + + def test_init_with_ancillary_variable_and_units(self): + weights = iris.analysis._Weights( + self.ancillary_variable, self.cube, units="J" + ) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [[5, 6, 7], [5, 6, 7]]) + assert weights.units == "J" + + def test_init_with_cell_measure(self): + weights = iris.analysis._Weights(self.cell_measure, self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.arange(6).reshape(2, 3)) + assert weights.units == "m2" + + def test_init_with_cell_measure_and_units(self): + weights = iris.analysis._Weights( + self.cell_measure, self.cube, units="J" + ) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.arange(6).reshape(2, 3)) + assert weights.units == "J" + + def test_init_with_list(self): + weights = iris.analysis._Weights([1, 2, 3], self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [1, 2, 3]) + assert weights.units == "1" + + def test_init_with_list_and_units(self): + weights = iris.analysis._Weights([1, 2, 3], self.cube, units="J") + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, [1, 2, 3]) + assert weights.units == "J" + + def test_init_with_ndarray(self): + weights = iris.analysis._Weights(np.zeros((5, 5)), self.cube) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.zeros((5, 5))) + assert weights.units == "1" + + def test_init_with_ndarray_and_units(self): + weights = iris.analysis._Weights( + np.zeros((5, 5)), self.cube, units="J" + ) + assert isinstance(weights, iris.analysis._Weights) + np.testing.assert_array_equal(weights, np.zeros((5, 5))) + assert weights.units == "J" + + def test_init_with_invalid_obj(self): + with pytest.raises(KeyError): + iris.analysis._Weights("invalid_obj", self.cube) + + def test_init_with_invalid_obj_and_units(self): + with pytest.raises(KeyError): + iris.analysis._Weights("invalid_obj", self.cube, units="J") + + def test_update_kwargs_no_weights(self): + kwargs = {"test": [1, 2, 3]} + iris.analysis._Weights.update_kwargs(kwargs, self.cube) + assert kwargs == {"test": [1, 2, 3]} + + def test_update_kwargs_weights_none(self): + kwargs = {"test": [1, 2, 3], "weights": None} + iris.analysis._Weights.update_kwargs(kwargs, self.cube) + assert kwargs == {"test": [1, 2, 3], "weights": None} + + def test_update_kwargs_weights(self): + kwargs = {"test": [1, 2, 3], "weights": [1, 2]} + iris.analysis._Weights.update_kwargs(kwargs, self.cube) + assert len(kwargs) == 2 + assert kwargs["test"] == [1, 2, 3] + assert isinstance(kwargs["weights"], iris.analysis._Weights) + np.testing.assert_array_equal(kwargs["weights"], [1, 2]) + assert kwargs["weights"].units == "1" + + +CUBE = iris.cube.Cube(0) + + +@pytest.mark.parametrize( + "kwargs,expected", + [ + ({}, "s"), + ({"test": "m"}, "s"), + ({"weights": None}, "s"), + ({"weights": [1, 2, 3]}, "s"), + ({"weights": iris.analysis._Weights([1], CUBE)}, "s"), + ({"weights": iris.analysis._Weights([1], CUBE, units="kg")}, "s kg"), + ], +) +def test_sum_units_func(kwargs, expected): + units = cf_units.Unit("s") + result = iris.analysis._sum_units_func(units, **kwargs) + assert result == expected + + if __name__ == "__main__": tests.main() diff --git a/lib/iris/tests/test_lazy_aggregate_by.py b/lib/iris/tests/test_lazy_aggregate_by.py index d1ebc9a36a..57b748e52f 100644 --- a/lib/iris/tests/test_lazy_aggregate_by.py +++ b/lib/iris/tests/test_lazy_aggregate_by.py @@ -6,6 +6,7 @@ import unittest from iris._lazy_data import as_lazy_data +from iris.analysis import SUM from iris.tests import test_aggregate_by @@ -44,5 +45,65 @@ def tearDown(self): assert self.cube_easy.has_lazy_data() +class TestLazyAggregateByWeightedByCube(TestLazyAggregateBy): + def setUp(self): + super().setUp() + + self.weights_single = self.cube_single[:, 0, 0].copy( + self.weights_single + ) + self.weights_single.units = "m2" + self.weights_multi = self.cube_multi[:, 0, 0].copy(self.weights_multi) + self.weights_multi.units = "m2" + + def test_str_aggregation_weighted_sum_single(self): + aggregateby_cube = self.cube_single.aggregated_by( + "height", + SUM, + weights=self.weights_single, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_coord_aggregation_weighted_sum_single(self): + aggregateby_cube = self.cube_single.aggregated_by( + self.coord_z_single, + SUM, + weights=self.weights_single, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_str_aggregation_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + ["height", "level"], + SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_str_aggregation_rev_order_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + ["level", "height"], + SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_coord_aggregation_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + [self.coord_z1_multi, self.coord_z2_multi], + SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + def test_coord_aggregation_rev_order_weighted_sum_multi(self): + aggregateby_cube = self.cube_multi.aggregated_by( + [self.coord_z2_multi, self.coord_z1_multi], + SUM, + weights=self.weights_multi, + ) + self.assertEqual(aggregateby_cube.units, "kelvin m2") + + if __name__ == "__main__": unittest.main() diff --git a/lib/iris/tests/unit/analysis/test_Aggregator.py b/lib/iris/tests/unit/analysis/test_Aggregator.py index ec837ea49a..45081ad07f 100644 --- a/lib/iris/tests/unit/analysis/test_Aggregator.py +++ b/lib/iris/tests/unit/analysis/test_Aggregator.py @@ -15,6 +15,7 @@ import numpy.ma as ma from iris.analysis import Aggregator +from iris.cube import Cube from iris.exceptions import LazyAggregatorError @@ -286,10 +287,30 @@ def test_units_change(self): units_func = mock.Mock(return_value=mock.sentinel.new_units) aggregator = Aggregator("", None, units_func) cube = mock.Mock(units=mock.sentinel.units) - aggregator.update_metadata(cube, []) - units_func.assert_called_once_with(mock.sentinel.units) + aggregator.update_metadata(cube, [], kw1=1, kw2=2) + units_func.assert_called_once_with(mock.sentinel.units, kw1=1, kw2=2) self.assertEqual(cube.units, mock.sentinel.new_units) + def test_units_func_no_kwargs(self): + # To ensure backwards-compatibility, Aggregator also supports + # units_func that accept the single argument `units` + def units_func(units): + return units**2 + + aggregator = Aggregator("", None, units_func) + cube = Cube(0, units="s") + aggregator.update_metadata(cube, [], kw1=1, kw2=2) + self.assertEqual(cube.units, "s2") + + def test_units_func_kwargs(self): + def units_func(units, **kwargs): + return units**2 + + aggregator = Aggregator("", None, units_func) + cube = Cube(0, units="s") + aggregator.update_metadata(cube, [], kw1=1, kw2=2) + self.assertEqual(cube.units, "s2") + class Test_lazy_aggregate(tests.IrisTest): def test_kwarg_pass_through_no_kwargs(self): diff --git a/lib/iris/tests/unit/cube/test_Cube.py b/lib/iris/tests/unit/cube/test_Cube.py index 8e9e00dce8..aa9e3b51b1 100644 --- a/lib/iris/tests/unit/cube/test_Cube.py +++ b/lib/iris/tests/unit/cube/test_Cube.py @@ -20,7 +20,7 @@ from iris._lazy_data import as_lazy_data import iris.analysis -from iris.analysis import MEAN, Aggregator, WeightedAggregator +from iris.analysis import MEAN, SUM, Aggregator, WeightedAggregator import iris.aux_factory from iris.aux_factory import HybridHeightFactory from iris.common.metadata import BaseMetadata @@ -320,18 +320,36 @@ def test_dim0_lazy(self): self.assertArrayAlmostEqual(cube_collapsed.data, [1.5, 2.5, 3.5]) self.assertFalse(cube_collapsed.has_lazy_data()) + def test_dim0_lazy_weights_none(self): + cube_collapsed = self.cube.collapsed("y", MEAN, weights=None) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAlmostEqual(cube_collapsed.data, [1.5, 2.5, 3.5]) + self.assertFalse(cube_collapsed.has_lazy_data()) + def test_dim1_lazy(self): cube_collapsed = self.cube.collapsed("x", MEAN) self.assertTrue(cube_collapsed.has_lazy_data()) self.assertArrayAlmostEqual(cube_collapsed.data, [1.0, 4.0]) self.assertFalse(cube_collapsed.has_lazy_data()) + def test_dim1_lazy_weights_none(self): + cube_collapsed = self.cube.collapsed("x", MEAN, weights=None) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAlmostEqual(cube_collapsed.data, [1.0, 4.0]) + self.assertFalse(cube_collapsed.has_lazy_data()) + def test_multidims(self): # Check that MEAN works with multiple dims. cube_collapsed = self.cube.collapsed(("x", "y"), MEAN) self.assertTrue(cube_collapsed.has_lazy_data()) self.assertArrayAllClose(cube_collapsed.data, 2.5) + def test_multidims_weights_none(self): + # Check that MEAN works with multiple dims. + cube_collapsed = self.cube.collapsed(("x", "y"), MEAN, weights=None) + self.assertTrue(cube_collapsed.has_lazy_data()) + self.assertArrayAllClose(cube_collapsed.data, 2.5) + def test_non_lazy_aggregator(self): # An aggregator which doesn't have a lazy function should still work. dummy_agg = Aggregator( @@ -342,18 +360,19 @@ def test_non_lazy_aggregator(self): self.assertArrayEqual(result.data, np.mean(self.data, axis=1)) -class Test_collapsed__multidim_weighted(tests.IrisTest): +class Test_collapsed__multidim_weighted_with_arr(tests.IrisTest): def setUp(self): self.data = np.arange(6.0).reshape((2, 3)) self.lazydata = as_lazy_data(self.data) - # Test cubes wth (same-valued) real and lazy data - cube_real = Cube(self.data) + # Test cubes with (same-valued) real and lazy data + cube_real = Cube(self.data, units="m") for i_dim, name in enumerate(("y", "x")): npts = cube_real.shape[i_dim] coord = DimCoord(np.arange(npts), long_name=name) cube_real.add_dim_coord(coord, i_dim) self.cube_real = cube_real self.cube_lazy = cube_real.copy(data=self.lazydata) + self.cube_lazy.units = "kg" # Test weights and expected result for a y-collapse self.y_weights = np.array([0.3, 0.5]) self.full_weights_y = np.broadcast_to( @@ -375,6 +394,7 @@ def test_weighted_fullweights_real_y(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_y ) + self.assertEqual(cube_collapsed.units, "m") def test_weighted_fullweights_lazy_y(self): # Full-shape weights, lazy data : Check lazy result, same values as real calc. @@ -385,6 +405,7 @@ def test_weighted_fullweights_lazy_y(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_y ) + self.assertEqual(cube_collapsed.units, "kg") def test_weighted_1dweights_real_y(self): # 1-D weights, real data : Check same results as full-shape. @@ -394,6 +415,7 @@ def test_weighted_1dweights_real_y(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_y ) + self.assertEqual(cube_collapsed.units, "m") def test_weighted_1dweights_lazy_y(self): # 1-D weights, lazy data : Check lazy result, same values as real calc. @@ -404,6 +426,7 @@ def test_weighted_1dweights_lazy_y(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_y ) + self.assertEqual(cube_collapsed.units, "kg") def test_weighted_fullweights_real_x(self): # Full weights, real data, ** collapse X ** : as for 'y' case above @@ -413,6 +436,7 @@ def test_weighted_fullweights_real_x(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_x ) + self.assertEqual(cube_collapsed.units, "m") def test_weighted_fullweights_lazy_x(self): # Full weights, lazy data, ** collapse X ** : as for 'y' case above @@ -423,6 +447,7 @@ def test_weighted_fullweights_lazy_x(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_x ) + self.assertEqual(cube_collapsed.units, "kg") def test_weighted_1dweights_real_x(self): # 1-D weights, real data, ** collapse X ** : as for 'y' case above @@ -432,6 +457,7 @@ def test_weighted_1dweights_real_x(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_x ) + self.assertEqual(cube_collapsed.units, "m") def test_weighted_1dweights_lazy_x(self): # 1-D weights, lazy data, ** collapse X ** : as for 'y' case above @@ -442,6 +468,148 @@ def test_weighted_1dweights_lazy_x(self): self.assertArrayAlmostEqual( cube_collapsed.data, self.expected_result_x ) + self.assertEqual(cube_collapsed.units, "kg") + + def test_weighted_sum_fullweights_adapt_units_real_y(self): + # Check that units are adapted correctly ('m' * '1' = 'm') + cube_collapsed = self.cube_real.collapsed( + "y", SUM, weights=self.full_weights_y + ) + self.assertEqual(cube_collapsed.units, "m") + + def test_weighted_sum_fullweights_adapt_units_lazy_y(self): + # Check that units are adapted correctly ('kg' * '1' = 'kg') + cube_collapsed = self.cube_lazy.collapsed( + "y", SUM, weights=self.full_weights_y + ) + self.assertEqual(cube_collapsed.units, "kg") + + def test_weighted_sum_1dweights_adapt_units_real_y(self): + # Check that units are adapted correctly ('m' * '1' = 'm') + # Note: the same test with lazy data fails: + # https://github.com/SciTools/iris/issues/5083 + cube_collapsed = self.cube_real.collapsed( + "y", SUM, weights=self.y_weights + ) + self.assertEqual(cube_collapsed.units, "m") + + def test_weighted_sum_with_unknown_units_real_y(self): + # Check that units are adapted correctly ('unknown' * '1' = 'unknown') + # Note: does not need to be adapted in subclasses since 'unknown' + # multiplied by any unit is 'unknown' + self.cube_real.units = "unknown" + cube_collapsed = self.cube_real.collapsed( + "y", + SUM, + weights=self.full_weights_y, + ) + self.assertEqual(cube_collapsed.units, "unknown") + + def test_weighted_sum_with_unknown_units_lazy_y(self): + # Check that units are adapted correctly ('unknown' * '1' = 'unknown') + # Note: does not need to be adapted in subclasses since 'unknown' + # multiplied by any unit is 'unknown' + self.cube_lazy.units = "unknown" + cube_collapsed = self.cube_lazy.collapsed( + "y", + SUM, + weights=self.full_weights_y, + ) + self.assertEqual(cube_collapsed.units, "unknown") + + +# Simply redo the tests of Test_collapsed__multidim_weighted_with_arr with +# other allowed objects for weights + + +class Test_collapsed__multidim_weighted_with_cube( + Test_collapsed__multidim_weighted_with_arr +): + def setUp(self): + super().setUp() + + self.y_weights_original = self.y_weights + self.full_weights_y_original = self.full_weights_y + self.x_weights_original = self.x_weights + self.full_weights_x_original = self.full_weights_x + + self.y_weights = self.cube_real[:, 0].copy(self.y_weights_original) + self.y_weights.units = "m2" + self.full_weights_y = self.cube_real.copy(self.full_weights_y_original) + self.full_weights_y.units = "m2" + self.x_weights = self.cube_real[0, :].copy(self.x_weights_original) + self.full_weights_x = self.cube_real.copy(self.full_weights_x_original) + + def test_weighted_sum_fullweights_adapt_units_real_y(self): + # Check that units are adapted correctly ('m' * 'm2' = 'm3') + cube_collapsed = self.cube_real.collapsed( + "y", SUM, weights=self.full_weights_y + ) + self.assertEqual(cube_collapsed.units, "m3") + + def test_weighted_sum_fullweights_adapt_units_lazy_y(self): + # Check that units are adapted correctly ('kg' * 'm2' = 'kg m2') + cube_collapsed = self.cube_lazy.collapsed( + "y", SUM, weights=self.full_weights_y + ) + self.assertEqual(cube_collapsed.units, "kg m2") + + def test_weighted_sum_1dweights_adapt_units_real_y(self): + # Check that units are adapted correctly ('m' * 'm2' = 'm3') + # Note: the same test with lazy data fails: + # https://github.com/SciTools/iris/issues/5083 + cube_collapsed = self.cube_real.collapsed( + "y", SUM, weights=self.y_weights + ) + self.assertEqual(cube_collapsed.units, "m3") + + +class Test_collapsed__multidim_weighted_with_str( + Test_collapsed__multidim_weighted_with_cube +): + def setUp(self): + super().setUp() + + self.full_weights_y = "full_y" + self.full_weights_x = "full_x" + self.y_weights = "y" + self.x_weights = "1d_x" + + self.dim_metadata_full_y = iris.coords.CellMeasure( + self.full_weights_y_original, + long_name=self.full_weights_y, + units="m2", + ) + self.dim_metadata_full_x = iris.coords.AuxCoord( + self.full_weights_x_original, + long_name=self.full_weights_x, + units="m2", + ) + self.dim_metadata_1d_y = iris.coords.DimCoord( + self.y_weights_original, long_name=self.y_weights, units="m2" + ) + self.dim_metadata_1d_x = iris.coords.AncillaryVariable( + self.x_weights_original, long_name=self.x_weights, units="m2" + ) + + for cube in (self.cube_real, self.cube_lazy): + cube.add_cell_measure(self.dim_metadata_full_y, (0, 1)) + cube.add_aux_coord(self.dim_metadata_full_x, (0, 1)) + cube.remove_coord("y") + cube.add_dim_coord(self.dim_metadata_1d_y, 0) + cube.add_ancillary_variable(self.dim_metadata_1d_x, 1) + + +class Test_collapsed__multidim_weighted_with_dim_metadata( + Test_collapsed__multidim_weighted_with_str +): + def setUp(self): + super().setUp() + + self.full_weights_y = self.dim_metadata_full_y + self.full_weights_x = self.dim_metadata_full_x + self.y_weights = self.dim_metadata_1d_y + self.x_weights = self.dim_metadata_1d_x class Test_collapsed__cellmeasure_ancils(tests.IrisTest): @@ -501,7 +669,7 @@ def _assert_warn_collapse_without_weight(self, coords, warn): self.assertIn(mock.call(msg.format(coord)), warn.call_args_list) def _assert_nowarn_collapse_without_weight(self, coords, warn): - # Ensure that warning is not rised. + # Ensure that warning is not raised. msg = "Collapsing spatial coordinate {!r} without weighting" for coord in coords: self.assertNotIn(mock.call(msg.format(coord)), warn.call_args_list) @@ -590,7 +758,7 @@ def _assert_warn_cannot_check_contiguity(self, warn): self.assertIn(mock.call(msg), warn.call_args_list) def _assert_cube_as_expected(self, cube): - """Ensure that cube data and coordiantes are as expected.""" + """Ensure that cube data and coordinates are as expected.""" self.assertArrayEqual(cube.data, np.array(3)) lat = cube.coord("latitude") @@ -604,16 +772,14 @@ def _assert_cube_as_expected(self, cube): def test_collapsed_lat_with_3_bounds(self): """Collapse latitude with 3 bounds.""" with mock.patch("warnings.warn") as warn: - collapsed_cube = self.cube.collapsed("latitude", iris.analysis.SUM) + collapsed_cube = self.cube.collapsed("latitude", SUM) self._assert_warn_cannot_check_contiguity(warn) self._assert_cube_as_expected(collapsed_cube) def test_collapsed_lon_with_3_bounds(self): """Collapse longitude with 3 bounds.""" with mock.patch("warnings.warn") as warn: - collapsed_cube = self.cube.collapsed( - "longitude", iris.analysis.SUM - ) + collapsed_cube = self.cube.collapsed("longitude", SUM) self._assert_warn_cannot_check_contiguity(warn) self._assert_cube_as_expected(collapsed_cube) @@ -621,7 +787,7 @@ def test_collapsed_lat_lon_with_3_bounds(self): """Collapse latitude and longitude with 3 bounds.""" with mock.patch("warnings.warn") as warn: collapsed_cube = self.cube.collapsed( - ["latitude", "longitude"], iris.analysis.SUM + ["latitude", "longitude"], SUM ) self._assert_warn_cannot_check_contiguity(warn) self._assert_cube_as_expected(collapsed_cube) @@ -741,9 +907,9 @@ def test_different_array_attrs_incompatible(self): class Test_rolling_window(tests.IrisTest): def setUp(self): - self.cube = Cube(np.arange(6)) + self.cube = Cube(np.arange(6), units="kg") self.multi_dim_cube = Cube(np.arange(36).reshape(6, 6)) - val_coord = DimCoord([0, 1, 2, 3, 4, 5], long_name="val") + val_coord = DimCoord([0, 1, 2, 3, 4, 5], long_name="val", units="s") month_coord = AuxCoord( ["jan", "feb", "mar", "apr", "may", "jun"], long_name="month" ) @@ -770,6 +936,7 @@ def test_string_coord(self): np.array([1, 2, 3, 4]), bounds=np.array([[0, 2], [1, 3], [2, 4], [3, 5]]), long_name="val", + units="s", ) month_coord = AuxCoord( np.array( @@ -818,6 +985,30 @@ def test_ancillary_variables_and_cell_measures_removed(self): self.assertEqual(res_cube.ancillary_variables(), []) self.assertEqual(res_cube.cell_measures(), []) + def test_weights_arr(self): + weights = [0, 0, 1, 0, 2] + res_cube = self.cube.rolling_window("val", SUM, 5, weights=weights) + np.testing.assert_array_equal(res_cube.data, [10, 13]) + self.assertEqual(res_cube.units, "kg") + + def test_weights_cube(self): + weights = Cube([0, 0, 1, 0, 2], units="m2") + res_cube = self.cube.rolling_window("val", SUM, 5, weights=weights) + np.testing.assert_array_equal(res_cube.data, [10, 13]) + self.assertEqual(res_cube.units, "kg m2") + + def test_weights_str(self): + weights = "val" + res_cube = self.cube.rolling_window("val", SUM, 6, weights=weights) + np.testing.assert_array_equal(res_cube.data, [55]) + self.assertEqual(res_cube.units, "kg s") + + def test_weights_dim_coord(self): + weights = self.cube.coord("val") + res_cube = self.cube.rolling_window("val", SUM, 6, weights=weights) + np.testing.assert_array_equal(res_cube.data, [55]) + self.assertEqual(res_cube.units, "kg s") + class Test_slices_dim_order(tests.IrisTest): """ @@ -905,7 +1096,7 @@ def setUp(self): len(self.cube.coord("model_level_number").points) ) self.exp_iter_2d = np.ndindex(6, 70, 1, 1) - # Define maximum number of interations for particularly long + # Define maximum number of interactions for particularly long # (and so time-consuming) iterators. self.long_iterator_max = 5