Repository: NVIDIA/earth2studio (language: python)

Issue:
🐛[BUG]: ACC calculation uses wrong mean
### Version

v0.1.0

### On which installation method(s) does this occur?

Pip, Source

### Describe the issue

The calculation of `x_bar` for ACC should subtract the mean of `x_hat` from `x_hat`, not the mean of `x`. Same for `y_bar`/`y_hat`/`y`. See [here](https://github.com/NVIDIA/earth2studio/blob/main/earth2studio/statistics/acc.py#L144) on lines 144 and 148.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index feb239ba..f0f1b5f8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,6 +22,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added the ability to reload zarr and netcdf backends
 - Added the ability to read from an IOBackend
 - Add spread/skill ratio
+- Added rank histogram
+- Added reduction_dimensions as required property in statistics
+  and metrics API.
 
 ### Changed
 
diff --git a/docs/modules/statistics.rst b/docs/modules/statistics.rst
index 6a5d8c26..4b71631f 100644
--- a/docs/modules/statistics.rst
+++ b/docs/modules/statistics.rst
@@ -17,6 +17,7 @@ Various statistic and metric calculations for analysing inference data.
     statistics.acc
     statistics.crps
     statistics.mean
+    statistics.rank_histogram
     statistics.rmse
     statistics.spread_skill_ratio
     statistics.std
diff --git a/earth2studio/statistics/__init__.py b/earth2studio/statistics/__init__.py
index 7678b423..244219bd 100644
--- a/earth2studio/statistics/__init__.py
+++ b/earth2studio/statistics/__init__.py
@@ -18,5 +18,6 @@
 from .base import Metric, Statistic
 from .crps import crps
 from .moments import mean, std, variance  # noqa
+from .rank import rank_histogram  # noqa
 from .rmse import rmse, spread_skill_ratio  # noqa
 from .weights import lat_weight  # noqa
diff --git a/earth2studio/statistics/acc.py b/earth2studio/statistics/acc.py
index 97e3bd26..97db1e45 100644
--- a/earth2studio/statistics/acc.py
+++ b/earth2studio/statistics/acc.py
@@ -56,13 +56,17 @@ def __init__(
                 raise ValueError(
                     "Error! Weights must be the same dimension as reduction_dimensions"
                 )
-        self.reduction_dimensions = reduction_dimensions
+        self._reduction_dimensions = reduction_dimensions
         self.weights = weights
 
         self.climatology = climatology
 
     def __str__(self) -> str:
-        return "acc"
+        return "_".join(self._reduction_dimensions + ["acc"])
+
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return self._reduction_dimensions
 
     def __call__(
         self,
@@ -78,16 +82,18 @@ def __call__(
         Parameters
         ----------
         x : torch.Tensor
-            Input tensor #1 intended to apply metric to.
+            Input tensor, typically the forecast or prediction tensor, but ACC is
+            symmetric with respect to `x` and `y`.
         x_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `x` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
             "time" and "variable" must be in x_coords.
         y : torch.Tensor
-            Input tensor #2 intended to be used as validation data.
+            Input tensor #2 intended to be used as validation data, but ACC is symmetric
+            with respect to `x` and `y`.
         y_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `y` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
             "time" and "variable" must be in y_coords.
             If "lead_time" is in x_coords, then "lead_time" must also be in y_coords. The
             intention, in this case, is that users will use `fetch_data` to make it easier
@@ -110,17 +116,17 @@ def __call__(
             handshake_dim(y_coords, c, i)
             handshake_coords(x_coords, y_coords, c)
 
-        dims = [list(x_coords).index(rd) for rd in self.reduction_dimensions]
+        dims = [list(x_coords).index(rd) for rd in self._reduction_dimensions]
         output_coords = CoordSystem(
             {
                 key: x_coords[key]
                 for key in x_coords
-                if key not in self.reduction_dimensions
+                if key not in self._reduction_dimensions
             }
         )
 
         weights = _broadcast_weights(
-            self.weights, self.reduction_dimensions, x_coords
+            self.weights, self._reduction_dimensions, x_coords
         ).to(x.device)
         weights_sum = torch.sum(weights)
 
@@ -141,11 +147,11 @@ def __call__(
             clim = torch.zeros_like(x)
 
         x_hat = x - clim
-        x_bar = torch.sum(weights * x, dim=dims, keepdim=True) / weights_sum
+        x_bar = torch.sum(weights * x_hat, dim=dims, keepdim=True) / weights_sum
         x_diff = x_hat - x_bar
 
         y_hat = y - clim
-        y_bar = torch.sum(weights * y, dim=dims, keepdim=True) / weights_sum
+        y_bar = torch.sum(weights * y_hat, dim=dims, keepdim=True) / weights_sum
         y_diff = y_hat - y_bar
 
         p1 = torch.sum(weights * x_diff * y_diff, dim=dims)
diff --git a/earth2studio/statistics/base.py b/earth2studio/statistics/base.py
index c5519bc8..a60aa177 100644
--- a/earth2studio/statistics/base.py
+++ b/earth2studio/statistics/base.py
@@ -25,11 +25,18 @@
 class Statistic(Protocol):
     """Statistic interface."""
 
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        """Gives the input dimensions of which the statistic performs a reduction
+        over. The is used to determine, a priori, the output dimensions of a statistic.
+        """
+        pass
+
     def __call__(
         self, x: torch.Tensor, coords: CoordSystem
     ) -> tuple[torch.Tensor, CoordSystem]:
         """Apply statistic to data `x`, with coordinates `coords` and reduce
-        over dimensions `reduction_dims`.
+        over dimensions `reduction_dimensions`.
 
         Parameters
         ----------
@@ -37,7 +44,7 @@ def __call__(
             Input tensor intended to apply statistic to.
         coords : CoordSystem
             Ordered dict representing coordinate system that describes the tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
         """
         pass
 
@@ -46,6 +53,10 @@ def __call__(
 class Metric(Protocol):
     """Metrics interface."""
 
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        pass
+
     def __call__(
         self,
         x: torch.Tensor,
@@ -54,19 +65,21 @@ def __call__(
         y_coords: CoordSystem,
     ) -> tuple[torch.Tensor, CoordSystem]:
         """Apply metric to data `x` and `y`, checking that their coordinates
-        are broadcastable. While reducing over `reduction_dims`.
+        are broadcastable. While reducing over `reduction_dimensions`.
 
         Parameters
         ----------
         x : torch.Tensor
-            Input tensor #1 intended to apply metric to.
+            Input tensor #1 intended to apply metric to. `x` is typically understood
+            to be the forecast or prediction tensor.
         x_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `x` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
         y : torch.Tensor
-            Input tensor #2 intended to apply statistic to.
+            Input tensor #2 intended to apply statistic to. `y` is typically the observation
+            or validation tensor.
         y_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `y` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
         """
         pass
diff --git a/earth2studio/statistics/crps.py b/earth2studio/statistics/crps.py
index 857aebce..5ee94876 100644
--- a/earth2studio/statistics/crps.py
+++ b/earth2studio/statistics/crps.py
@@ -16,6 +16,7 @@
 
 import torch
 
+from earth2studio.statistics.moments import mean
 from earth2studio.utils.coords import handshake_coords, handshake_dim
 from earth2studio.utils.type import CoordSystem
 
@@ -35,21 +36,47 @@ class crps:
 
     Parameters
     ----------
-    reduction_dimension: str
+    ensemble_dimension: str
         A name corresponding to a dimension to perform the
-        statistical reduction over. Example: 'ensemble'
+        ensemble reduction over. Example: 'ensemble'
+    reduction_dimensions: list[str]
+        A list of dimensions over which to average the crps over.
+        optional, by default none. If none, no additional reduction is done.
+    weights: torch.Tensor, optional
+        A tensor containing weights to assign to the reduction dimensions.
+        Note that these weights must have the same number of dimensions
+        as passed in reduction_dimensions.
+        Example: if reduction_dimensions = ['lat', 'lon'] then
+        assert weights.ndim == 2.
+        By default None.
     """
 
-    def __init__(self, reduction_dimension: str):
-        if not isinstance(reduction_dimension, str):
+    def __init__(
+        self,
+        ensemble_dimension: str,
+        reduction_dimensions: list[str] | None = None,
+        weights: torch.Tensor = None,
+    ):
+        if not isinstance(ensemble_dimension, str):
             raise ValueError(
                 "Error! CRPS currently assumes reduction over a single dimension."
             )
 
-        self.reduction_dimension = reduction_dimension
+        self.ensemble_dimension = ensemble_dimension
+        self._reduction_dimensions = reduction_dimensions
+        if reduction_dimensions is not None:
+            self.mean = mean(reduction_dimensions, weights=weights, batch_update=False)
 
     def __str__(self) -> str:
-        return "crps"
+        return "_".join(self.reduction_dimensions + ["crps"])
+
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return (
+            [self.ensemble_dimension]
+            if self._reduction_dimensions is None
+            else [self.ensemble_dimension] + self._reduction_dimensions
+        )
 
     def __call__(
         self,
@@ -62,37 +89,35 @@ def __call__(
         Apply metric to data `x` and `y`, checking that their coordinates
         are broadcastable. While reducing over `reduction_dims`.
 
-        If batch_update was passed True upon metric initialization then this method
-        returns the running sample RMSE over all seen batches.
-
         Parameters
         ----------
         x : torch.Tensor
-            Input tensor #1 intended to apply metric to.
+            Input tensor of ensemble forecast or prediction data. This is the tensor
+            over which the CRPS/CDF is calculated with respect to.
         x_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `x` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
         y : torch.Tensor
-            Input tensor #2 intended to apply statistic to.
+            Observation or validation tensor.
         y_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `y` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
 
         Returns
         -------
         tuple[torch.Tensor, CoordSystem]
-            Returns root mean squared error tensor with appropriate reduced coordinates.
+            Returns CRPS tensor with appropriate reduced coordinates.
         """
-        if self.reduction_dimension not in x_coords:
+        if not all([rd in x_coords for rd in self.reduction_dimensions]):
             raise ValueError(
                 "Initialized reduction dimension does not appear in passed coords"
             )
 
         # Do some coordinate checking
-        # Assume reduction_dim is in x_coords but not y_coords
-        if self.reduction_dimension in y_coords:
+        # Assume ensemble_dim is in x_coords but not y_coords
+        if self.ensemble_dimension in y_coords:
             raise ValueError(
-                f"{self.reduction_dimension} should not be in y_coords but is."
+                f"{self.ensemble_dimension} should not be in y_coords but is."
             )
         if x.ndim != y.ndim + 1:
             raise ValueError(
@@ -102,15 +127,19 @@ def __call__(
         # Input coordinate checking
         coord_count = 0
         for c in x_coords:
-            if c != self.reduction_dimension:
+            if c != self.ensemble_dimension:
                 handshake_dim(y_coords, c, coord_count)
                 handshake_coords(y_coords, x_coords, c)
                 coord_count += 1
 
-        dim = list(x_coords).index(self.reduction_dimension)
+        dim = list(x_coords).index(self.ensemble_dimension)
+
+        out = _crps_from_empirical_cdf(x, y, dim=dim)
+        out_coords = y_coords.copy()
 
-        _crps = _crps_from_empirical_cdf(x, y, dim=dim)
-        return _crps, y_coords.copy()
+        if self._reduction_dimensions is not None:
+            out, out_coords = self.mean(out, out_coords)
+        return out, out_coords
 
 
 def _crps_from_empirical_cdf(
diff --git a/earth2studio/statistics/moments.py b/earth2studio/statistics/moments.py
index 125a7e50..51d3cb76 100644
--- a/earth2studio/statistics/moments.py
+++ b/earth2studio/statistics/moments.py
@@ -54,7 +54,7 @@ def __init__(
                 raise ValueError(
                     "Error! Weights must be the same dimension as reduction_dimensions"
                 )
-        self.reduction_dimensions = reduction_dimensions
+        self._reduction_dimensions = reduction_dimensions
         self.weights = weights
 
         self.batch_update = batch_update
@@ -62,7 +62,11 @@ def __init__(
             self.n = 0
 
     def __str__(self) -> str:
-        return "_".join(self.reduction_dimensions + ["mean"])
+        return "_".join(self._reduction_dimensions + ["mean"])
+
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return self._reduction_dimensions
 
     def __call__(
         self, x: torch.Tensor, coords: CoordSystem
@@ -80,17 +84,21 @@ def __call__(
         coords: CoordSystem
             Coordinates referring to the input data, x.
         """
-        if not all([rd in coords for rd in self.reduction_dimensions]):
+        if not all([rd in coords for rd in self._reduction_dimensions]):
             raise ValueError(
-                "Initialized reduction dimensions do not appear in passed coords."
+                "Initialized reduction_dimensions do not appear in passed coords."
             )
 
-        dims = [list(coords).index(rd) for rd in self.reduction_dimensions]
+        dims = [list(coords).index(rd) for rd in self._reduction_dimensions]
         output_coords = CoordSystem(
-            {key: coords[key] for key in coords if key not in self.reduction_dimensions}
+            {
+                key: coords[key]
+                for key in coords
+                if key not in self._reduction_dimensions
+            }
         )
         weights = _broadcast_weights(
-            self.weights, self.reduction_dimensions, coords
+            self.weights, self._reduction_dimensions, coords
         ).to(x.device)
         weights_sum = torch.sum(weights)
 
@@ -143,7 +151,7 @@ def __init__(
                 raise ValueError(
                     "Error! Weights must be the same dimension as reduction_dimensions"
                 )
-        self.reduction_dimensions = reduction_dimensions
+        self._reduction_dimensions = reduction_dimensions
         self.weights = weights
 
         self.batch_update = batch_update
@@ -151,7 +159,11 @@ def __init__(
             self.n = 0
 
     def __str__(self) -> str:
-        return "_".join(self.reduction_dimensions + ["variance"])
+        return "_".join(self._reduction_dimensions + ["variance"])
+
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return self._reduction_dimensions
 
     def __call__(
         self, x: torch.Tensor, coords: CoordSystem
@@ -169,18 +181,22 @@ def __call__(
         coords: CoordSystem
             Coordinates referring to the input data, x.
         """
-        if not all([rd in coords for rd in self.reduction_dimensions]):
+        if not all([rd in coords for rd in self._reduction_dimensions]):
             raise ValueError(
-                "Initialized reduction dimensions do not appear in passed coords."
+                "Initialized reduction_dimensions do not appear in passed coords."
             )
 
-        dims = [list(coords).index(rd) for rd in self.reduction_dimensions]
+        dims = [list(coords).index(rd) for rd in self._reduction_dimensions]
         output_coords = CoordSystem(
-            {key: coords[key] for key in coords if key not in self.reduction_dimensions}
+            {
+                key: coords[key]
+                for key in coords
+                if key not in self._reduction_dimensions
+            }
         )
 
         weights = _broadcast_weights(
-            self.weights, self.reduction_dimensions, coords
+            self.weights, self._reduction_dimensions, coords
         ).to(x.device)
         weights_sum = torch.sum(weights)
 
@@ -248,11 +264,15 @@ def __init__(
         self.var = variance(
             reduction_dimensions, weights=weights, batch_update=batch_update
         )
-        self.reduction_dimensions = reduction_dimensions
+        self._reduction_dimensions = reduction_dimensions
         self.weights = weights
 
     def __str__(self) -> str:
-        return "_".join(self.reduction_dimensions + ["std"])
+        return "_".join(self._reduction_dimensions + ["std"])
+
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return self._reduction_dimensions
 
     def __call__(
         self, x: torch.Tensor, coords: CoordSystem
diff --git a/earth2studio/statistics/rank.py b/earth2studio/statistics/rank.py
new file mode 100644
index 00000000..014dad5c
--- /dev/null
+++ b/earth2studio/statistics/rank.py
@@ -0,0 +1,158 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import OrderedDict
+
+import numpy as np
+import torch
+from modulus.metrics.general.histogram import _count_bins, linspace
+
+from earth2studio.utils.coords import handshake_coords, handshake_dim
+from earth2studio.utils.type import CoordSystem
+
+
+class rank_histogram:
+    """
+    Compute the Rank Histogram for a given set of ensemble forecasts.
+
+
+
+    This statistic reduces over a single dimension, where the presumed ensemble dimension
+    does not appear in the truth/observation tensor.
+
+    Parameters
+    ----------
+    ensemble_dimension: str
+        A name corresponding to a dimension to perform the
+        ranking over. Example: 'ensemble'
+    reduction_dimensions: list[str]
+        A list of dimensions over which to bin the ranks.
+    number_of_bins: int
+        The number of bins to discretize the unit interval over.
+        by default, 10
+    """
+
+    def __init__(
+        self,
+        ensemble_dimension: str,
+        reduction_dimensions: list[str],
+        number_of_bins: int = 10,
+    ):
+        if not isinstance(ensemble_dimension, str):
+            raise ValueError(
+                "Error! Rank Histogram currently assumes reduction over a single dimension."
+            )
+
+        self.ensemble_dimension = ensemble_dimension
+        self._reduction_dimensions = reduction_dimensions
+        self.number_of_bins = number_of_bins
+
+    def __str__(self) -> str:
+        return "rank_histogram"
+
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return [self.ensemble_dimension] + self._reduction_dimensions
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        x_coords: CoordSystem,
+        y: torch.Tensor,
+        y_coords: CoordSystem,
+    ) -> tuple[torch.Tensor, CoordSystem]:
+        """
+        Apply metric to data `x` and `y`, checking that their coordinates
+        are broadcastable. While reducing over `reduction_dimensions`.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of ensemble data. The rank of observation input tensor `y`
+            is determined with respect to the ensemble dimension of `x`.
+        x_coords : CoordSystem
+            Ordered dict representing coordinate system that describes the `x` tensor.
+            `reduction_dimensions` must be in coords.
+        y : torch.Tensor
+            The observation input tensor.
+        y_coords : CoordSystem
+            Ordered dict representing coordinate system that describes the `y` tensor.
+            `reduction_dimensions` must be in coords.
+
+        Returns
+        -------
+        tuple[torch.Tensor, CoordSystem]
+            Returns rank histogram tensor with appropriate reduced coordinates.
+        """
+        if not all([rd in x_coords for rd in self.reduction_dimensions]):
+            raise ValueError(
+                "Initialized reduction dimension does not appear in passed coords"
+            )
+
+        # Do some coordinate checking
+        # Assume ensemble_dim is in x_coords but not y_coords
+        if self.ensemble_dimension in y_coords:
+            raise ValueError(
+                f"{self.ensemble_dimension} should not be in y_coords but is."
+            )
+        if x.ndim != y.ndim + 1:
+            raise ValueError(
+                "x and y must have broadcastable shapes but got"
+                + f"{x.shape} and {y.shape}"
+            )
+        # Input coordinate checking
+        coord_count = 0
+        for c in x_coords:
+            if c != self.ensemble_dimension:
+                handshake_dim(y_coords, c, coord_count)
+                handshake_coords(y_coords, x_coords, c)
+                coord_count += 1
+
+        dim = list(x_coords).index(self.ensemble_dimension)
+
+        _ranks = torch.mean(1.0 * torch.ge(y, x), dim=dim)
+
+        # Reshape ranks using reduction dimensions
+        dims = [list(y_coords).index(rd) for rd in self._reduction_dimensions]
+        new_dims = list(range(len(dims)))
+        _ranks = torch.flatten(
+            torch.movedim(_ranks, dims, new_dims),
+            start_dim=new_dims[0],
+            end_dim=new_dims[-1],
+        )
+
+        # Compute histogram
+        _bin_edges = linspace(
+            torch.zeros_like(_ranks[0]), torch.ones_like(_ranks[0]), self.number_of_bins
+        )
+        _rank_histogram = _count_bins(_ranks, _bin_edges)
+
+        # Create output tensor and coords
+        out = torch.stack((0.5 * (_bin_edges[:-1] + _bin_edges[1:]), _rank_histogram))
+        out_coords = y_coords.copy()
+        for rd in self._reduction_dimensions:
+            out_coords.pop(rd)
+
+        out_coords = (
+            OrderedDict(
+                {
+                    "histogram_data": np.array(["bin_centers", "bin_counts"]),
+                    "bin": np.arange(self.number_of_bins),
+                }
+            )
+            | out_coords
+        )
+        return out, out_coords
diff --git a/earth2studio/statistics/rmse.py b/earth2studio/statistics/rmse.py
index c832cd91..acfe8085 100644
--- a/earth2studio/statistics/rmse.py
+++ b/earth2studio/statistics/rmse.py
@@ -55,12 +55,16 @@ def __init__(
             reduction_dimensions, weights=weights, batch_update=batch_update
         )
         self.weights = weights
-        self.reduction_dimensions = reduction_dimensions
+        self._reduction_dimensions = reduction_dimensions
 
         self.batch_update = batch_update
 
     def __str__(self) -> str:
-        return "_".join(self.reduction_dimensions + ["rmse"])
+        return "_".join(self._reduction_dimensions + ["rmse"])
+
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return self._reduction_dimensions
 
     def __call__(
         self,
@@ -79,15 +83,17 @@ def __call__(
         Parameters
         ----------
         x : torch.Tensor
-            Input tensor #1 intended to apply metric to.
+            Input tensor, typically the forecast or prediction tensor, but RMSE is
+            symmetric with respect to `x` and `y`.
         x_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `x` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
         y : torch.Tensor
-            Input tensor #2 intended to apply statistic to.
+            Input tensor #2 intended to be used as validation data, but ACC is symmetric
+            with respect to `x` and `y`.
         y_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `y` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
 
         Returns
         -------
@@ -147,7 +153,7 @@ def __init__(
         reduction_batch_update: bool = False,
     ):
         self.ensemble_dimension = [ensemble_dimension]
-        self.reduction_dimensions = reduction_dimensions
+        self._reduction_dimensions = reduction_dimensions
         self.ensemble_mean = mean(
             reduction_dimensions=self.ensemble_dimension,
             weights=ensemble_weights,
@@ -171,9 +177,13 @@ def __init__(
 
     def __str__(self) -> str:
         return "_".join(
-            self.ensemble_dimension + self.reduction_dimensions + ["spread_skill"]
+            self.ensemble_dimension + self._reduction_dimensions + ["spread_skill"]
         )
 
+    @property
+    def reduction_dimensions(self) -> list[str]:
+        return self.ensemble_dimension + self._reduction_dimensions
+
     def __call__(
         self,
         x: torch.Tensor,
@@ -191,15 +201,16 @@ def __call__(
         Parameters
         ----------
         x : torch.Tensor
-            Input tensor #1 intended to apply metric to.
+            The ensemble forecast input tensor. This is the tensor over which the
+            ensemble mean and spread are calculated with respect to.
         x_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `x` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
         y : torch.Tensor
-            Input tensor #2 intended to apply statistic to.
+            The observation input tensor.
         y_coords : CoordSystem
             Ordered dict representing coordinate system that describes the `y` tensor.
-            'reduction_dims' must be in coords.
+            `reduction_dimensions` must be in coords.
 
         Returns
         -------


Test patch:
diff --git a/test/statistics/test_crps.py b/test/statistics/test_crps.py
index a1d65de9..3660921f 100644
--- a/test/statistics/test_crps.py
+++ b/test/statistics/test_crps.py
@@ -21,19 +21,34 @@
 import pytest
 import torch
 
-from earth2studio.statistics import crps
+from earth2studio.statistics import crps, lat_weight
 from earth2studio.statistics.crps import _crps_from_empirical_cdf
 
+lat_weights = lat_weight(torch.as_tensor(np.linspace(-90.0, 90.0, 361)))
+
 
 @pytest.mark.parametrize(
-    "reduction_dimension",
+    "ensemble_dimension",
     [
         "ensemble",
         "time",
     ],
 )
+@pytest.mark.parametrize(
+    "reduction_weights",
+    [
+        (None, None),
+        (["lat", "lon"], lat_weights.unsqueeze(1).repeat(1, 720)),
+        (["lat", "lon"], None),
+        (["lat"], lat_weights),
+    ],
+)
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
-def test_crps(reduction_dimension: str, device: str) -> None:
+def test_crps(
+    ensemble_dimension: str,
+    reduction_weights: tuple[list[str], np.ndarray],
+    device: str,
+) -> None:
 
     x = torch.randn((10, 1, 2, 361, 720), device=device)
 
@@ -48,16 +63,23 @@ def test_crps(reduction_dimension: str, device: str) -> None:
     )
 
     y_coords = copy.deepcopy(x_coords)
-    y_coords.pop(reduction_dimension)
+    y_coords.pop(ensemble_dimension)
     y_shape = [len(y_coords[c]) for c in y_coords]
     y = torch.randn(y_shape, device=device)
 
-    CRPS = crps(reduction_dimension)
+    reduction_dimensions, weights = reduction_weights
+    if weights is not None:
+        weights = weights.to(device)
+    CRPS = crps(
+        ensemble_dimension, reduction_dimensions=reduction_dimensions, weights=weights
+    )
 
     z, c = CRPS(x, x_coords, y, y_coords)
-    assert reduction_dimension not in c
+    assert ensemble_dimension not in c
+    if reduction_dimensions is not None:
+        print(list(c), reduction_dimensions)
+        assert all([rd not in c for rd in reduction_dimensions])
     assert list(z.shape) == [len(val) for val in c.values()]
-    assert z.shape == y.shape
 
 
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
diff --git a/test/statistics/test_ranks.py b/test/statistics/test_ranks.py
new file mode 100644
index 00000000..5b950116
--- /dev/null
+++ b/test/statistics/test_ranks.py
@@ -0,0 +1,152 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+import torch
+
+from earth2studio.statistics import rank_histogram
+
+
+@pytest.mark.parametrize(
+    "ensemble_dimension",
+    [
+        "ensemble",
+        "time",
+    ],
+)
+@pytest.mark.parametrize("device", ["cpu", "cuda:0"])
+def test_rank_histogram(ensemble_dimension: str, device: str) -> None:
+
+    x = torch.randn((10, 1, 2, 361, 720), device=device)
+
+    x_coords = OrderedDict(
+        {
+            "ensemble": np.arange(10),
+            "time": np.array([np.datetime64("1993-04-05T00:00")]),
+            "variable": ["t2m", "tcwv"],
+            "lat": np.linspace(-90.0, 90.0, 361),
+            "lon": np.linspace(0.0, 360.0, 720, endpoint=False),
+        }
+    )
+
+    y_coords = copy.deepcopy(x_coords)
+    y_coords.pop(ensemble_dimension)
+    y_shape = [len(y_coords[c]) for c in y_coords]
+    y = torch.randn(y_shape, device=device)
+
+    reduction_dimensions = ["lat", "lon"]
+    RH = rank_histogram(ensemble_dimension, reduction_dimensions)
+
+    z, c = RH(x, x_coords, y, y_coords)
+    for di in [ensemble_dimension] + reduction_dimensions:
+        assert di not in c
+    assert list(z.shape) == [len(val) for val in c.values()]
+
+
+@pytest.mark.parametrize("device", ["cpu", "cuda:0"])
+def test_rank_histogram_failures(device: str) -> None:
+    reduction_dimension = "ensemble"
+    x = torch.randn((10, 1, 2, 361, 720), device=device)
+
+    x_coords = OrderedDict(
+        {
+            "ensemble": np.arange(10),
+            "time": np.array([np.datetime64("1993-04-05T00:00")]),
+            "variable": ["t2m", "tcwv"],
+            "lat": np.linspace(-90.0, 90.0, 361),
+            "lon": np.linspace(0.0, 360.0, 720, endpoint=False),
+        }
+    )
+
+    # Raise error for training to pass a list
+    with pytest.raises(ValueError):
+        rank_histogram(["ensemble", "time"], ["lat"])
+
+    RH = rank_histogram(reduction_dimension, ["lat", "lon"])
+
+    # Test reduction_dimension in y error
+    with pytest.raises(ValueError):
+        y_coords = copy.deepcopy(x_coords)
+        y_shape = [len(y_coords[c]) for c in y_coords]
+        y = torch.randn(y_shape, device=device)
+        z, c = RH(x, x_coords, y, y_coords)
+
+    # Test x and y don't have broadcastable shapes
+    with pytest.raises(ValueError):
+        y_coords = OrderedDict({"phony": np.arange(1)})
+        for c in x_coords:
+            if c != reduction_dimension:
+                y_coords[c] = x_coords[c]
+
+        y_shape = [len(y_coords[c]) for c in y_coords]
+        y = torch.randn(y_shape, device=device)
+        z, c = RH(x, x_coords, y, y_coords)
+
+    # Test reduction_dimension not in x_coords
+    with pytest.raises(ValueError):
+        y_coords = OrderedDict({})
+        for c in x_coords:
+            if c != reduction_dimension:
+                y_coords[c] = x_coords[c]
+
+        y_shape = [len(y_coords[c]) for c in y_coords]
+        y = torch.randn(y_shape, device=device)
+
+        x_coords.pop("ensemble")
+        z, c = RH(x, x_coords, y, y_coords)
+
+
+@pytest.mark.parametrize("reduction_dimension", [20, 50, 80])
+@pytest.mark.parametrize("number_of_bins", [3, 5, 8])
+@pytest.mark.parametrize("device", ["cpu", "cuda:0"])
+def test_rank_histogram_accuracy(
+    reduction_dimension: int, number_of_bins: int, device: str
+) -> None:
+
+    x = (
+        1.0
+        + torch.randn((1000, reduction_dimension, reduction_dimension), device=device)
+        ** 2
+    )
+
+    x_coords = OrderedDict(
+        {
+            "ensemble": np.arange(1000),
+            "lat": np.arange(reduction_dimension),
+            "lon": np.arange(reduction_dimension),
+        }
+    )
+
+    y_coords = copy.deepcopy(x_coords)
+    y_coords.pop("ensemble")
+    y_shape = [len(y_coords[c]) for c in y_coords]
+    y = 1.0 + torch.randn(y_shape, device=device) ** 2
+
+    reduction_dimensions = ["lat", "lon"]
+    RH = rank_histogram("ensemble", reduction_dimensions, number_of_bins=number_of_bins)
+
+    z, _ = RH(x, x_coords, y, y_coords)
+
+    # This should be result in a near uniform distribution but is prone to statistical error
+    assert torch.allclose(
+        z[1, :],
+        reduction_dimension**2 / number_of_bins * torch.ones_like(z[1, :]),
+        rtol=2.0 * number_of_bins / reduction_dimension,
+    )