Resolve "release v1.4.0"

Closes #317 (closed)

mlair/helpers/statistics.py

@@ -257,11 +257,12 @@ class SkillScores:
     """
     models_default = ["cnn", "persi", "ols"]
 
-    def __init__(self, external_data: Data, models=None, observation_name="obs"):
+    def __init__(self, external_data: Union[Data, None], models=None, observation_name="obs", ahead_dim="ahead"):
         """Set internal data."""
         self.external_data = external_data
         self.models = self.set_model_names(models)
         self.observation_name = observation_name
+        self.ahead_dim = ahead_dim
 
     def set_model_names(self, models: List[str]) -> List[str]:
         """Either use given models or use defaults."""
@@ -283,19 +284,17 @@ class SkillScores:
         combination_strings = [f"{first}-{second}" for (first, second) in combinations]
         return combinations, combination_strings
 
-    def skill_scores(self, window_lead_time: int, ahead_dim="ahead") -> pd.DataFrame:
+    def skill_scores(self) -> pd.DataFrame:
         """
         Calculate skill scores for all combinations of model names.
 
-        :param window_lead_time: length of forecast steps
-
         :return: skill score for each comparison and forecast step
         """
-        ahead_names = list(self.external_data[ahead_dim].data)
+        ahead_names = list(self.external_data[self.ahead_dim].data)
         combinations, combination_strings = self.get_model_name_combinations()
         skill_score = pd.DataFrame(index=combination_strings)
         for iahead in ahead_names:
-            data = self.external_data.sel(ahead=iahead)
+            data = self.external_data.sel({self.ahead_dim: iahead})
             skill_score[iahead] = [self.general_skill_score(data,
                                                             forecast_name=first,
                                                             reference_name=second,
@@ -303,8 +302,7 @@ class SkillScores:
                                    for (first, second) in combinations]
         return skill_score
 
-    def climatological_skill_scores(self, internal_data: Data, window_lead_time: int,
-                                    forecast_name: str, ahead_dim="ahead") -> xr.DataArray:
+    def climatological_skill_scores(self, internal_data: Data, forecast_name: str) -> xr.DataArray:
         """
         Calculate climatological skill scores according to Murphy (1988).
 
@@ -312,20 +310,19 @@ class SkillScores:
         is part of parameters.
 
         :param internal_data: internal data
-        :param window_lead_time: interested time step of forecast horizon to select data
         :param forecast_name: name of the forecast to use for this calculation (must be available in `data`)
 
         :return: all CASES as well as all terms
         """
-        ahead_names = list(self.external_data[ahead_dim].data)
+        ahead_names = list(self.external_data[self.ahead_dim].data)
 
         all_terms = ['AI', 'AII', 'AIII', 'AIV', 'BI', 'BII', 'BIV', 'CI', 'CIV', 'CASE I', 'CASE II', 'CASE III',
                      'CASE IV']
         skill_score = xr.DataArray(np.full((len(all_terms), len(ahead_names)), np.nan), coords=[all_terms, ahead_names],
-                                   dims=['terms', 'ahead'])
+                                   dims=['terms', self.ahead_dim])
 
         for iahead in ahead_names:
-            data = internal_data.sel(ahead=iahead)
+            data = internal_data.sel({self.ahead_dim: iahead})
 
             skill_score.loc[["CASE I", "AI", "BI", "CI"], iahead] = np.stack(self._climatological_skill_score(
                 data, mu_type=1, forecast_name=forecast_name, observation_name=self.observation_name).values.flatten())
@@ -334,7 +331,7 @@ class SkillScores:
                 data, mu_type=2, forecast_name=forecast_name, observation_name=self.observation_name).values.flatten())
 
             if self.external_data is not None and self.observation_name in self.external_data.coords["type"]:
-                external_data = self.external_data.sel(ahead=iahead, type=[self.observation_name])
+                external_data = self.external_data.sel({self.ahead_dim: iahead, "type": [self.observation_name]})
                 skill_score.loc[["CASE III", "AIII"], iahead] = np.stack(self._climatological_skill_score(
                     data, mu_type=3, forecast_name=forecast_name, observation_name=self.observation_name,
                     external_data=external_data).values.flatten())
@@ -373,12 +370,12 @@ class SkillScores:
         skill_score = 1 - mse(observation, forecast) / mse(observation, reference)
         return skill_score.values
 
-    @staticmethod
-    def skill_score_pre_calculations(data: Data, observation_name: str, forecast_name: str) -> Tuple[np.ndarray,
-                                                                                                     np.ndarray,
-                                                                                                     np.ndarray,
-                                                                                                     Data,
-                                                                                                     Dict[str, Data]]:
+    def skill_score_pre_calculations(self, data: Data, observation_name: str, forecast_name: str) -> Tuple[np.ndarray,
+                                                                                                           np.ndarray,
+                                                                                                           np.ndarray,
+                                                                                                           Data,
+                                                                                                           Dict[
+                                                                                                               str, Data]]:
         """
         Calculate terms AI, BI, and CI, mean, variance and pearson's correlation and clean up data.
 
@@ -391,7 +388,7 @@ class SkillScores:
 
         :returns: Terms AI, BI, and CI, internal data without nans and mean, variance, correlation and its p-value
         """
-        data = data.sel(type=[observation_name, forecast_name]).drop("ahead")
+        data = data.sel(type=[observation_name, forecast_name]).drop(self.ahead_dim)
         data = data.dropna("index")
 
         mean = data.mean("index")