Change lmoments3 behaviour with fitkwargs

CHANGELOG.rst

@@ -4,7 +4,7 @@ Changelog
 
 v0.55.0 (unreleased)
 --------------------
-Contributors to this version: Juliette Lavoie (:user:`juliettelavoie`), Trevor James Smith (:user:`Zeitsperre`).
+Contributors to this version: Juliette Lavoie (:user:`juliettelavoie`), Trevor James Smith (:user:`Zeitsperre`), Éric Dupuis (:user:`coxipi`).
 
 New indicators
 ^^^^^^^^^^^^^^
@@ -13,7 +13,7 @@ New indicators
 
 New features and enhancements
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-* New function ``ensemble.partition.general_partition`` (:pull:`2035`)
+* New function ``ensemble.partition.general_partition``. (:pull:`2035`).
 * Added a new ``xclim.indices.generic.bivariate_count_occurrences`` function to count instances where operations and performed and validated for two variables. (:pull:`2030`).
 * `xclim` now tracks energy usage and carbon emissions ("last run", "average", and "total") during CI workflows using the `eco-ci-energy-estimation` GitHub Action. (:pull:`2046`).
 
@@ -22,6 +22,8 @@ Internal changes
 * `sphinx-codeautolink` and `pygments` have been temporarily pinned due to breaking API changes. (:pull:`2030`).
 * Adjusted the ``TestOfficialYaml`` test to use a dynamic method for finding the installed location of `xclim`. (:pull:`2028`).
 * Adjusted two tests for better handling when running in Windows environments. (:pull:`2057`).
+* There is now a warning stating that `fitkwargs` are not employed when using the `lmoments3` distribution. One exception is the use of `'floc'` which is allowed with the gamma distribution. `'floc'` is used to shift the distribution before computing fitting parameters with the `lmoments3` distribution since ``loc=0`` is always assumed in the library. (:issue:`2043`, :pull:`2045`).
+* The `PWM` method (from `lmoments3`) is now available to be used with the `gamma` distribution in ``xclim.indices.standardized_precipitation_index`` and ``xclim.indices.standardized_precipitation_evapotranspiration_index``. (:issue:`2043`, :pull:`2045`).
 
 v0.54.0 (2024-12-16)
 --------------------

src/xclim/indices/_agro.py

@@ -1143,13 +1143,14 @@ def standardized_precipitation_index(
     window : int
         Averaging window length relative to the resampling frequency. For example, if `freq="MS"`,
         i.e. a monthly resampling, the window is an integer number of months.
-    dist : {"gamma", "fisk"}
-        Name of the univariate distribution. (see :py:mod:`scipy.stats`).
-    method : {"APP", "ML"}
+    dist : {'gamma', 'fisk'}
+        Name of the univariate distribution. (see :py:mod:`scipy.stats`). All possible distributions are allowed with 'PWM'.
+    method : {"APP", "ML", "PWM"}
         Name of the fitting method, such as `ML` (maximum likelihood), `APP` (approximate). The approximate method
         uses a deterministic function that does not involve any optimization.
     fitkwargs : dict, optional
         Kwargs passed to ``xclim.indices.stats.fit`` used to impose values of certains parameters (`floc`, `fscale`).
+        If method is `PWM`, `fitkwargs` should be empty, except for `floc` with `dist`=`gamma` which is allowed.
     cal_start : DateStr, optional
         Start date of the calibration period. A `DateStr` is expected, that is a `str` in format `"YYYY-MM-DD"`.
         Default option `None` means that the calibration period begins at the start of the input dataset.
@@ -1218,13 +1219,15 @@ def standardized_precipitation_index(
     >>> spi_3_fitted = standardized_precipitation_index(pr, params=params)
     """
     fitkwargs = fitkwargs or {}
+
     dist_methods = {"gamma": ["ML", "APP"], "fisk": ["ML", "APP"]}
     if dist in dist_methods:
         if method not in dist_methods[dist]:
             raise NotImplementedError(
-                f"{method} method is not implemented for {dist} distribution."
+                f"{method} method is not implemented for {dist} distribution"
             )
-    else:
+    # Constraints on distributions except for PWM
+    elif method != "PWM":
         raise NotImplementedError(f"{dist} distribution is not yet implemented.")
 
     # Precipitation is expected to be zero-inflated
@@ -1280,12 +1283,13 @@ def standardized_precipitation_evapotranspiration_index(
         Averaging window length relative to the resampling frequency. For example, if `freq="MS"`, i.e. a monthly
         resampling, the window is an integer number of months.
     dist : {'gamma', 'fisk'}
-        Name of the univariate distribution. (see :py:mod:`scipy.stats`).
-    method : {'APP', 'ML'}
+        Name of the univariate distribution. (see :py:mod:`scipy.stats`). All possible distributions are allowed with 'PWM'.
+    method : {'APP', 'ML', 'PWM'}
         Name of the fitting method, such as `ML` (maximum likelihood), `APP` (approximate). The approximate method
         uses a deterministic function that doesn't involve any optimization.
     fitkwargs : dict, optional
         Kwargs passed to ``xclim.indices.stats.fit`` used to impose values of certains parameters (`floc`, `fscale`).
+        If method is `PWM`, `fitkwargs` should be empty, except for `floc` with `dist`=`gamma` which is allowed.
     cal_start : DateStr, optional
         Start date of the calibration period. A `DateStr` is expected, that is a `str` in format `"YYYY-MM-DD"`.
         Default option `None` means that the calibration period begins at the start of the input dataset.
@@ -1317,7 +1321,8 @@ def standardized_precipitation_evapotranspiration_index(
             raise NotImplementedError(
                 f"{method} method is not implemented for {dist} distribution"
             )
-    else:
+    # Constraints on distributions except for PWM
+    elif method != "PWM":
         raise NotImplementedError(f"{dist} distribution is not yet implemented.")
 
     # Water budget is not expected to be zero-inflated

src/xclim/indices/stats.py

@@ -59,7 +59,21 @@ def _fitfunc_1d(arr, *, dist, nparams, method, **fitkwargs):
         # lmoments3 will raise an error if only dist.numargs + 2 values are provided
         if len(x) <= dist.numargs + 2:
             return np.asarray([np.nan] * nparams)
-        params = list(dist.lmom_fit(x).values())
+        if (type(dist).__name__ != "GammaGen" and len(fitkwargs.keys()) != 0) or (
+            type(dist).__name__ == "GammaGen"
+            and set(fitkwargs.keys()) - {"floc"} != set()
+        ):
+            raise ValueError(
+                "Lmoments3 does not use `fitkwargs` arguments, except for `floc` with the Gamma distribution."
+            )
+        if "floc" in fitkwargs and type(dist).__name__ == "GammaGen":
+            # lmoments3 assumes `loc` is 0, so `x` may need to be shifted
+            # note that `floc` must already be in appropriate units for `x`
+            params = dist.lmom_fit(x - fitkwargs["floc"])
+            params["loc"] = fitkwargs["floc"]
+            params = list(params.values())
+        else:
+            params = list(dist.lmom_fit(x).values())
     elif method == "APP":
         args, kwargs = _fit_start(x, dist.name, **fitkwargs)
         kwargs.setdefault("loc", 0)
@@ -796,12 +810,15 @@ def standardized_index_fit_params(
 
     dist_and_methods = {"gamma": ["ML", "APP"], "fisk": ["ML", "APP"]}
     dist = get_dist(dist)
-    if dist.name not in dist_and_methods:
-        raise NotImplementedError(f"The distribution `{dist.name}` is not supported.")
-    if method not in dist_and_methods[dist.name]:
-        raise NotImplementedError(
-            f"The method `{method}` is not supported for distribution `{dist.name}`."
-        )
+    if method != "PWM":
+        if dist.name not in dist_and_methods:
+            raise NotImplementedError(
+                f"The distribution `{dist.name}` is not supported."
+            )
+        if method not in dist_and_methods[dist.name]:
+            raise NotImplementedError(
+                f"The method `{method}` is not supported for distribution `{dist.name}`."
+            )
     da, group = preprocess_standardized_index(da, freq, window, **indexer)
     if zero_inflated:
         prob_of_zero = da.groupby(group).map(
@@ -870,8 +887,9 @@ def standardized_index(
         The approximate method uses a deterministic function that doesn't involve any optimization.
     zero_inflated : bool
         If True, the zeroes of `da` are treated separately.
-    fitkwargs : dict
+    fitkwargs : dict, optional
         Kwargs passed to ``xclim.indices.stats.fit`` used to impose values of certains parameters (`floc`, `fscale`).
+        If method is `PWM`, `fitkwargs` should be empty, except for `floc` with `dist`=`gamma` which is allowed.
     cal_start : DateStr, optional
         Start date of the calibration period. A `DateStr` is expected, that is a `str` in format `"YYYY-MM-DD"`.
         Default option `None` means that the calibration period begins at the start of the input dataset.

tests/test_indices.py

@@ -661,6 +661,38 @@ class TestStandardizedIndices:
                 [-1.08627775, -0.46491398, -0.77806462, 0.31759127, 0.03794528],
                 2e-2,
             ),
+            (
+                "D",
+                1,
+                "gamma",
+                "PWM",
+                [-0.13002, 1.346689, 0.965731, 0.245408, -0.427896],
+                2e-2,
+            ),
+            (
+                "D",
+                12,
+                "gamma",
+                "PWM",
+                [-0.209411, -0.086357, 0.636851, 1.022608, 0.634409],
+                2e-2,
+            ),
+            (
+                "MS",
+                1,
+                "gamma",
+                "PWM",
+                [1.364243, 1.478565, 1.915559, -3.055828, 0.905304],
+                2e-2,
+            ),
+            (
+                "MS",
+                12,
+                "gamma",
+                "PWM",
+                [0.577214, 1.522867, 1.634222, 0.967847, 0.689001],
+                2e-2,
+            ),
         ],
     )
     def test_standardized_precipitation_index(
@@ -672,6 +704,13 @@ def test_standardized_precipitation_index(
             and Version(__numpy_version__) < Version("2.0.0")
         ):
             pytest.skip("Skipping SPI/ML/D on older numpy")
+
+        # change `dist` to a lmoments3 object if needed
+        if method == "PWM":
+            lmom = pytest.importorskip("lmoments3.distr")
+            scipy2lmom = {"gamma": "gam"}
+            dist = getattr(lmom, scipy2lmom[dist])
+
         ds = open_dataset("sdba/CanESM2_1950-2100.nc").isel(location=1)
         if freq == "D":
             # to compare with ``climate_indices``
@@ -923,6 +962,42 @@ def test_zero_inflated(self, open_dataset):
             np.all(np.not_equal(spid[False].values, spid[True].values)), True
         )
 
+    def test_PWM_and_fitkwargs(self, open_dataset):
+        pr = (
+            open_dataset("sdba/CanESM2_1950-2100.nc")
+            .isel(location=1)
+            .sel(time=slice("1950", "1980"))
+        ).pr
+
+        lmom = pytest.importorskip("lmoments3.distr")
+        # for now, only one function used
+        scipy2lmom = {"gamma": "gam"}
+        dist = getattr(lmom, scipy2lmom["gamma"])
+        fitkwargs = {"floc": 0}
+        input_params = dict(
+            freq=None,
+            window=1,
+            method="PWM",
+            dist=dist,
+            fitkwargs=fitkwargs,
+            # doy_bounds=(180, 180),
+        )
+        # this should not cause a problem
+        params_d0 = xci.stats.standardized_index_fit_params(pr, **input_params).isel(
+            dayofyear=0
+        )
+        np.testing.assert_allclose(
+            params_d0, np.array([5.63e-01, 0, 3.37e-05]), rtol=0, atol=2e-2
+        )
+        # this should cause a problem
+        fitkwargs["fscale"] = 1
+        input_params["fitkwargs"] = fitkwargs
+        with pytest.raises(
+            ValueError,
+            match="Lmoments3 does not use `fitkwargs` arguments, except for `floc` with the Gamma distribution.",
+        ):
+            xci.stats.standardized_index_fit_params(pr, **input_params)
+
 
 class TestDailyFreezeThawCycles:
     @pytest.mark.parametrize(