Fast MBCn (a la groupies)

CHANGES.rst

@@ -10,6 +10,11 @@ Announcements
 ^^^^^^^^^^^^^
 * `xclim` has migrated its development branch name from `master` to `main`. (:issue:`1667`, :pull:`1669`).
 
+
+New features and enhancements
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+* New `TrainAdjust` and class `MBCn` giving a faster and more accurate implementation of the 'MBCn' algorithm (:issue:`1551`, :pull:`1580`).
+
 New indicators
 ^^^^^^^^^^^^^^
 * New ``snw_season_length`` and ``snd_season_length`` computing the duration between the start and the end of the snow season, both defined as the first day of a continuous period with snow above/under a threshold. Previous versions of these indicators were renamed ``snw_days_above`` and ``snd_days_above`` to better reflect what they computed : the number of days with snow above a given threshold (with no notion of continuity). (:issue:`1703`, :pull:`1708`).
@@ -83,6 +88,7 @@ New features and enhancements
 * Added a new ``xclim.indices.generic.select_rolling_resample_op`` function to allow for computing rolling statistics. (:issue:`1480`, :pull:`1643`).
 * Add the possibility to use a group with a window in ``xc.sdba.processing.reordering``. (:pull:`1566`).
 
+
 Breaking changes
 ^^^^^^^^^^^^^^^^
 * `xclim` base Python version has been raised to Python3.9. Python3.9+ coding conventions are now supported. (:issue:`1268`, :pull:`1565`).

docs/notebooks/sdba.ipynb

@@ -481,7 +481,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "##### Perform an initial univariate adjustment."
+    "##### Perform the multivariate adjustment (MBCn)."
    ]
   },
   {
@@ -490,132 +490,42 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# additive for tasmax\n",
-    "QDMtx = sdba.QuantileDeltaMapping.train(\n",
-    "    dref.tasmax, dhist.tasmax, nquantiles=20, kind=\"+\", group=\"time\"\n",
+    "ADJ = sdba.MBCn.train(\n",
+    "    dref,\n",
+    "    dhist,\n",
+    "    base_kws={\"nquantiles\": 20, \"group\": \"time\"},\n",
+    "    adj_kws={\"interp\": \"nearest\", \"extrapolation\": \"constant\"},\n",
+    "    n_iter=20,  # perform 20 iteration\n",
+    "    n_escore=1000,  # only send 1000 points to the escore metric\n",
     ")\n",
-    "# Adjust both hist and sim, we'll feed both to the Npdf transform.\n",
-    "scenh_tx = QDMtx.adjust(dhist.tasmax)\n",
-    "scens_tx = QDMtx.adjust(dsim.tasmax)\n",
-    "\n",
-    "# remove == 0 values in pr:\n",
-    "dref[\"pr\"] = sdba.processing.jitter_under_thresh(dref.pr, \"0.01 mm d-1\")\n",
-    "dhist[\"pr\"] = sdba.processing.jitter_under_thresh(dhist.pr, \"0.01 mm d-1\")\n",
-    "dsim[\"pr\"] = sdba.processing.jitter_under_thresh(dsim.pr, \"0.01 mm d-1\")\n",
-    "\n",
-    "# multiplicative for pr\n",
-    "QDMpr = sdba.QuantileDeltaMapping.train(\n",
-    "    dref.pr, dhist.pr, nquantiles=20, kind=\"*\", group=\"time\"\n",
-    ")\n",
-    "# Adjust both hist and sim, we'll feed both to the Npdf transform.\n",
-    "scenh_pr = QDMpr.adjust(dhist.pr)\n",
-    "scens_pr = QDMpr.adjust(dsim.pr)\n",
-    "\n",
-    "scenh = xr.Dataset(dict(tasmax=scenh_tx, pr=scenh_pr))\n",
-    "scens = xr.Dataset(dict(tasmax=scens_tx, pr=scens_pr))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "##### Stack the variables to multivariate arrays and standardize them\n",
-    "The standardization process ensure the mean and standard deviation of each column (variable) is 0 and 1 respectively.\n",
     "\n",
-    "`scenh` and `scens` are standardized together, so the two series are coherent. As we'll see further, we do not need to keep the mean and standard deviation, as we only keep the rank order information from the `NpdfTransform` output."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Stack the variables (tasmax and pr)\n",
-    "ref = sdba.processing.stack_variables(dref)\n",
-    "scenh = sdba.processing.stack_variables(scenh)\n",
-    "scens = sdba.processing.stack_variables(scens)\n",
-    "\n",
-    "# Standardize\n",
-    "ref, _, _ = sdba.processing.standardize(ref)\n",
-    "\n",
-    "allsim_std, _, _ = sdba.processing.standardize(xr.concat((scenh, scens), \"time\"))\n",
-    "scenh_std = allsim_std.sel(time=scenh.time)\n",
-    "scens_std = allsim_std.sel(time=scens.time)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "##### Perform the N-dimensional probability density function transform\n",
-    "\n",
-    "The NpdfTransform will iteratively randomly rotate our arrays in the \"variables\" space and apply the univariate adjustment before rotating it back. In Cannon (2018) and Pitié et al. (2005), it can be seen that the source array's joint distribution converges toward the target's joint distribution when a large number of iterations is done."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from xclim import set_options\n",
-    "\n",
-    "# See the advanced notebook for details on how this option work\n",
-    "with set_options(sdba_extra_output=True):\n",
-    "    out = sdba.adjustment.NpdfTransform.adjust(\n",
-    "        ref,\n",
-    "        scenh_std,\n",
-    "        scens_std,\n",
-    "        base=sdba.QuantileDeltaMapping,  # Use QDM as the univariate adjustment.\n",
-    "        base_kws={\"nquantiles\": 20, \"group\": \"time\"},\n",
-    "        n_iter=20,  # perform 20 iteration\n",
-    "        n_escore=1000,  # only send 1000 points to the escore metric (it is really slow)\n",
+    "scenh, scens = (\n",
+    "    ADJ.adjust(\n",
+    "        sim=ds,\n",
+    "        ref=dref,\n",
+    "        hist=dhist,\n",
+    "        base=sdba.QuantileDeltaMapping,\n",
+    "        base_kws_vars={\n",
+    "            \"pr\": {\n",
+    "                \"kind\": \"*\",\n",
+    "                \"jitter_under_thresh_value\": \"0.01 mm d-1\",\n",
+    "                \"adapt_freq_thresh\": \"0.1 mm d-1\",\n",
+    "            },\n",
+    "            \"tasmax\": {\"kind\": \"+\"},\n",
+    "        },\n",
+    "        adj_kws={\"interp\": \"nearest\", \"extrapolation\": \"constant\"},\n",
     "    )\n",
-    "\n",
-    "scenh_npdft = out.scenh.rename(time_hist=\"time\")  # Bias-adjusted historical period\n",
-    "scens_npdft = out.scen  # Bias-adjusted future period\n",
-    "extra = out.drop_vars([\"scenh\", \"scen\"])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "##### Restoring the trend\n",
-    "\n",
-    "The NpdfT has given us new \"hist\" and \"sim\" arrays with a correct rank structure. However, the trend is lost in this process. We reorder the result of the initial adjustment according to the rank structure of the NpdfT outputs to get our final bias-adjusted series.\n",
-    "\n",
-    "`sdba.processing.reordering`: 'ref' the argument that provides the order, 'sim' is the argument to reorder."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "scenh = sdba.processing.reordering(scenh_npdft, scenh, group=\"time\")\n",
-    "scens = sdba.processing.reordering(scens_npdft, scens, group=\"time\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "scenh = sdba.processing.unstack_variables(scenh)\n",
-    "scens = sdba.processing.unstack_variables(scens)"
+    "    for ds in [dhist, dsim]\n",
+    ")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "##### There we are!\n",
+    "##### Let's trigger all the computations.\n",
     "\n",
-    "Let's trigger all the computations. The use of `dask.compute` allows the three DataArrays to be computed at the same time, avoiding repeating the common steps."
+    "The use of `dask.compute` allows the three DataArrays to be computed at the same time, avoiding repeating the common steps."
    ]
   },
   {
@@ -629,7 +539,7 @@
     "\n",
     "with ProgressBar():\n",
     "    scenh, scens, escores = compute(\n",
-    "        scenh.isel(location=2), scens.isel(location=2), extra.escores.isel(location=2)\n",
+    "        scenh.isel(location=2), scens.isel(location=2), ADJ.ds.escores.isel(location=2)\n",
     "    )"
    ]
   },
@@ -686,7 +596,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.11.7"
   },
   "toc": {
    "base_numbering": 1,