Implement DataTree in auto_regression

mesmer/calibrate_mesmer/train_gv.py

@@ -11,7 +11,7 @@
 import xarray as xr
 
 from mesmer.io.save_mesmer_bundle import save_mesmer_data
-from mesmer.stats import _fit_auto_regression_scen_ens, _select_ar_order_scen_ens
+from mesmer.stats import fit_auto_regression_scen_ens, select_ar_order_scen_ens
 
 
 def train_gv(gv, targ, esm, cfg, save_params=True, **kwargs):
@@ -172,10 +172,10 @@ def train_gv_AR(params_gv, gv, max_lag, sel_crit):
     # create temporary DataArray objects
     data = [xr.DataArray(data, dims=["run", "time"]) for data in gv.values()]
 
-    AR_order = _select_ar_order_scen_ens(
+    AR_order = select_ar_order_scen_ens(
         *data, dim="time", ens_dim="run", maxlag=max_lag, ic=sel_crit
     )
-    params = _fit_auto_regression_scen_ens(
+    params = fit_auto_regression_scen_ens(
         *data, dim="time", ens_dim="run", lags=AR_order
     )
 

mesmer/calibrate_mesmer/train_lv.py

@@ -14,9 +14,9 @@
 )
 from mesmer.io.save_mesmer_bundle import save_mesmer_data
 from mesmer.stats import (
-    _fit_auto_regression_scen_ens,
     adjust_covariance_ar1,
     find_localized_empirical_covariance,
+    fit_auto_regression_scen_ens,
 )
 
 
@@ -235,7 +235,7 @@ def train_lv_AR1_sci(params_lv, targs, y, wgt_scen_eq, aux, cfg):
         dims = ("run", "time", "cell")
         data = [xr.DataArray(data, dims=dims) for data in targ.values()]
 
-        params = _fit_auto_regression_scen_ens(*data, dim="time", ens_dim="run", lags=1)
+        params = fit_auto_regression_scen_ens(*data, dim="time", ens_dim="run", lags=1)
 
         params_lv["AR1_int"][targ_name] = params.intercept.values
         params_lv["AR1_coef"][targ_name] = params.coeffs.values.squeeze()

mesmer/stats/__init__.py

@@ -1,12 +1,12 @@
 from mesmer.stats._auto_regression import (
-    _fit_auto_regression_scen_ens,
-    _select_ar_order_scen_ens,
     draw_auto_regression_correlated,
     draw_auto_regression_monthly,
     draw_auto_regression_uncorrelated,
     fit_auto_regression,
     fit_auto_regression_monthly,
+    fit_auto_regression_scen_ens,
     select_ar_order,
+    select_ar_order_scen_ens,
 )
 from mesmer.stats._gaspari_cohn import gaspari_cohn, gaspari_cohn_correlation_matrices
 from mesmer.stats._harmonic_model import fit_harmonic_model, predict_harmonic_model
@@ -27,8 +27,8 @@
 
 __all__ = [
     # auto regression
-    "_fit_auto_regression_scen_ens",
-    "_select_ar_order_scen_ens",
+    "fit_auto_regression_scen_ens",
+    "select_ar_order_scen_ens",
     "draw_auto_regression_correlated",
     "draw_auto_regression_uncorrelated",
     "fit_auto_regression",