pythonCRO 0.1.8__py3-none-any.whl

pyCRO/__init__.py

@@ -0,0 +1,16 @@
+# __init__.py
+
+__version__ = "0.1.8"
+
+from .fitting import RO_fitting
+from .solver import RO_solver
+
+from .analytic import RO_analytic_std, RO_analytic_solver
+
+from .utils import RO_BWJ, func_mon_std
+from .par_load import par_load
+
+from .fit_LR import fit_LR
+from .fit_MLE import fit_MLE
+
+from .visual import plot_RO_par, plot_ens_RO_par

pyCRO/analytic.py

@@ -0,0 +1,244 @@
+import sys
+import numpy as np
+
+def RO_analytic_std(par):
+    """
+    Compute analytical standard deviation of T and h for the Recharge Oscillator (RO) model.
+    ONLY for linear RO model with white noise (annual mean parameters only)
+
+    Parameters
+    ----------
+    par : dict
+        Dictionary containing parameter arrays:
+        'R', 'F1', 'epsilon', 'F2', 'sigma_T', 'sigma_h'.
+
+    Returns
+    -------
+    T_std : float
+        Standard deviation of T.
+    h_std : float
+        Standard deviation of h.
+    """
+    
+    R_value = par['R'][0]
+    F1_value = par['F1'][0]
+    epsilon_value = par['epsilon'][0]
+    F2_value = par['F2'][0]
+    sigma_T_value = par['sigma_T'][0]
+    sigma_h_value = par['sigma_h'][0]
+
+    # Precompute useful terms
+    numerator_T = ((F1_value * F2_value - epsilon_value * R_value + epsilon_value**2) * sigma_T_value**2 +
+                   (F1_value**2) * sigma_h_value**2)
+    denominator = 2 * (-R_value + epsilon_value) * (F1_value * F2_value - R_value * epsilon_value)
+    T_std = np.sqrt(numerator_T / denominator)
+
+    numerator_h = ((F2_value**2) * sigma_T_value**2 +
+                   (F1_value * F2_value - epsilon_value * R_value + R_value**2) * sigma_h_value**2)
+    h_std = np.sqrt(numerator_h / denominator)
+
+    return T_std, h_std
+
+
+def RO_analytic_solver(par, IC, N, NE, dt=0.1, saveat=1.0, savemethod="sampling", noise_custom=None):
+    """
+    Analytical solution of the Recharge Oscillator (RO) model with deterministic and stochastic forcing.
+    ONLY for linear RO model with white noise (annual mean parameters only)
+
+    Parameters
+    ----------
+    par : dict
+        Dictionary of model parameters, each as at least a one-element array:
+        
+        - ``'R'`` : float
+            Damping parameter.
+        - ``'F1'`` : float
+            Feedback parameter relating thermocline depth to SST.
+        - ``'epsilon'`` : float
+            Thermocline damping parameter.
+        - ``'F2'`` : float
+            Feedback parameter relating SST to thermocline.
+        - ``'sigma_T'`` : float
+            Noise amplitude for SST.
+        - ``'sigma_h'`` : float
+            Noise amplitude for thermocline depth.
+    IC : tuple of float
+        Initial condition ``(To, ho)``:
+        
+        - ``To`` : initial SST anomaly.
+        - ``ho`` : initial thermocline depth anomaly.
+    N : float
+        Total simulation length (time units).
+    NE : int
+        Number of ensemble members.
+    dt : float, optional
+        Numerical integration time step. Default is ``0.1``.
+    saveat : float, optional
+        Output saving interval. Must be divisible by ``dt``. Default is ``1.0``.
+    savemethod : {'sampling', 'mean'}, optional
+        Method for saving results:
+        
+        - ``'sampling'`` : take samples every ``saveat``.
+        - ``'mean'`` : average values over each ``saveat`` interval.
+    noise_custom : {None, int, ndarray}, optional
+        Specification of stochastic noise:
+        
+        - ``None`` (default): new Gaussian noise is generated for each ensemble.
+        - ``int`` : seed for reproducible noise (same across ensembles).
+        - ``ndarray`` of shape ``(NT-1, 4)`` : user-provided noise, repeated across ensembles.
+        - ``ndarray`` of shape ``(NT-1, 4, NE)`` : user-provided noise for each ensemble.
+
+    Returns
+    -------
+    T_anal_out : ndarray of shape (N_out, NE)
+        SST anomalies including deterministic and stochastic forcing, 
+        after applying the saving scheme. The final row is NaN-padded 
+        for dimensional consistency.
+    h_anal_out : ndarray of shape (N_out, NE)
+        Thermocline anomalies including deterministic and stochastic forcing, 
+        after applying the saving scheme. The final row is NaN-padded.
+    noise_array : ndarray of shape (NT-1, 4, NE)
+        Realizations of Gaussian noise used in the simulation.
+
+    Notes
+    -----
+    The analytical solution consists of:
+    
+    - Deterministic part: exponential-sinusoidal functions of time.
+    - Stochastic part: weighted sums (discrete convolutions) of noise 
+      with exponential and sinusoidal kernels.
+    
+    Ensemble members are evaluated simultaneously using vectorized 
+    NumPy broadcasting, avoiding explicit Python loops.
+
+    Examples
+    --------
+    >>> par = {'R':[0.5], 'F1':[1.2], 'epsilon':[0.3], 'F2':[0.8],
+    ...        'sigma_T':[0.2], 'sigma_h':[0.1]}
+    >>> IC = (0.1, -0.2)
+    >>> T, h, noise = RO_solver_analytic(par, IC, N=50, NE=10, dt=0.1, saveat=1.0)
+    >>> T.shape, h.shape
+    ((51, 10), (51, 10))
+    """
+    
+    NT = int(round(N / dt))
+    step = int(round(saveat / dt))
+    ratio = saveat / dt
+
+    ### Checking simulation setups ###
+    ##################################
+    print("---------------------------------------------------------------------------------")
+    print("Welcome to the CRO Analytical Solver!")
+    print("Ensure that the same arguments are provided as for RO_solver,\n except that 'NM' and 'EF' are not required.")
+    print("---------------------------------------------------------------------------------")
+
+    ### Check timestep settings ###
+    if dt > saveat:
+        raise ValueError(f"dt={dt} and saveat={saveat}: The numerical time step (dt) must be smaller than or equal to the data saving interval (saveat).")
+    if not ratio.is_integer():
+        raise ValueError(f"dt = {dt}, saveat = {saveat}: saveat must be divisible by dt so that saveat/dt results in an integer.")
+    ###############################
+
+    ### extracting linear and noise amplitude parameters ###
+    ########################################################
+    R_value = par['R'][0]
+    F1_value = par['F1'][0]
+    epsilon_value = par['epsilon'][0]
+    F2_value = par['F2'][0]
+    sigma_T_value = par['sigma_T'][0]
+    sigma_h_value = par['sigma_h'][0]
+    ########################################################
+    ########################################################
+
+
+    ### initial setups ###
+    ######################
+    To, ho = IC
+    t = np.arange(NT) * dt
+    gr = (R_value - epsilon_value) / 2
+    w = np.sqrt(4 * F1_value * F2_value - (R_value + epsilon_value)**2) / 2
+
+    T_anal_out = np.zeros((NT, NE))
+    h_anal_out = np.zeros((NT, NE))
+
+    T_det = np.exp(gr * t) * (To * np.cos(w * t) +
+        To * (R_value + epsilon_value) / (2 * w) * np.sin(w * t) +
+        ho * (F1_value / w) * np.sin(w * t))
+
+    h_det = np.exp(gr * t) * (ho * np.cos(w * t) -
+        ho * (R_value + epsilon_value) / (2 * w) * np.sin(w * t) -
+        To * (F2_value / w) * np.sin(w * t))
+    ######################
+    ######################
+
+    ########################################################
+    # Noise generation (vectorized for all NE)
+    if noise_custom is None:
+        noise_array = np.random.randn(NT - 1, 4, NE)
+    elif np.isscalar(noise_custom):
+        np.random.seed(int(noise_custom))
+        noise_array = np.random.randn(NT - 1, 4, NE)
+    elif noise_custom.shape == (NT - 1, 4):
+        noise_array = np.repeat(noise_custom[..., None], NE, axis=2)
+    elif noise_custom.shape == (NT - 1, 4, NE):
+        noise_array = noise_custom
+    else:
+        raise ValueError(f"Invalid 'noise_custom' input: expected None, scalar, shape ({NT-1},4) or ({NT-1},4,{NE}).")
+
+    w_T = sigma_T_value * noise_array[:, 0, :] / np.sqrt(dt)  # (NT-1, NE)
+    w_h = sigma_h_value * noise_array[:, 1, :] / np.sqrt(dt)  # (NT-1, NE)
+
+    # Vectorized convolution-like integral for stochastic terms
+    T_sto = np.zeros((NT - 1, NE))
+    h_sto = np.zeros((NT - 1, NE))
+
+    ### analytical calculations
+    ########################################################
+    for i in range(1, NT - 1):
+        tau = t[:i+1]                     # (i+1,)
+        delta = t[i] - tau                 # (i+1,)
+        sin = np.sin(w * delta)            # (i+1,)
+        cos = np.cos(w * delta)            # (i+1,)
+        expgr = np.exp(gr * delta)         # (i+1,)
+
+        # (i+1, NE) via broadcasting
+        T_kernel = (
+            w_T[:i+1, :] * cos[:, None] +
+            w_T[:i+1, :] * (R_value + epsilon_value) / (2 * w) * sin[:, None] +
+            w_h[:i+1, :] * (F1_value / w) * sin[:, None]
+        )
+        h_kernel = (
+            w_h[:i+1, :] * cos[:, None] -
+            w_h[:i+1, :] * (R_value + epsilon_value) / (2 * w) * sin[:, None] -
+            w_T[:i+1, :] * (F2_value / w) * sin[:, None]
+        )
+
+        T_sto[i, :] = np.sum(dt * expgr[:, None] * T_kernel, axis=0)
+        h_sto[i, :] = np.sum(dt * expgr[:, None] * h_kernel, axis=0)
+
+    # Add deterministic and stochastic contributions
+    T_anal_out = np.vstack([T_det[:-1, None] + T_sto, np.full((1, NE), np.nan)])
+    h_anal_out = np.vstack([h_det[:-1, None] + h_sto, np.full((1, NE), np.nan)])
+
+    ### save data ###
+    #################
+    if savemethod == "sampling":
+        T_anal_out = T_anal_out[::step, :]
+        h_anal_out = h_anal_out[::step, :]
+    elif savemethod == "mean":
+        T_anal_out = T_anal_out[int(np.ceil(step/2)):NT-int(step/2),:]
+        h_anal_out = h_anal_out[int(np.ceil(step/2)):NT-int(step/2),:]
+        T_anal_out = T_anal_out.reshape(-1, step, NE).mean(axis=1)
+        h_anal_out = h_anal_out.reshape(-1, step, NE).mean(axis=1)
+        if NE > 1:
+            T_anal_out = np.vstack([T_anal_out[0:1, :], T_anal_out])
+            h_anal_out = np.vstack([h_anal_out[0:1, :], h_anal_out])
+        elif NE == 1:
+            T_anal_out = np.concatenate([[T_anal_out[0]], T_anal_out])
+            h_anal_out = np.concatenate([[h_anal_out[0]], h_anal_out])
+    
+    print("All steps are successfully completed!")
+    print("---------------------------------------------------------------------------------")
+    #################
+    #################
+    return T_anal_out, h_anal_out, noise_array

pyCRO/fit_LR.py

@@ -0,0 +1,341 @@
+import sys
+import numpy as np
+from .utils import wrap_to_pi, regress_std, heaviside
+from .fit_MAC import fit_MAC
+
+def fit_LR(T, h, par_option_T, par_option_h, par_option_noise, dt, 
+           tend_option, fitting_option_B, fitting_option_red):
+    """
+    Estimates Recharge Oscillator (RO) parameters with linear regression (LR)
+    given ENSO SST (T) and thermocline depth (h) anomaly time series.
+
+    Parameters
+    ----------
+    T : array_like
+        Time series of SST anomalies (1D).
+    h : array_like
+        Time series of thermocline depth anomalies (1D).
+    par_option_T : array_like
+        Seasonal fitting options for the T-equation terms.
+        Values must be one of:
+            - 0 : constant zero
+            - 1 : constant non-zero
+            - 3 : annual variation
+            - 5 : annual + semi-annual variation
+    par_option_h : array_like
+        Seasonal fitting options for the h-equation terms (same coding as above).
+    par_option_noise : array_like of length 3
+        Noise fitting options [n_T, n_h, n_g]:
+            - n_T : 0 = red, 1 = white
+            - n_h : 0 = red, 1 = white
+            - n_g : 0 = multiplicative (full),
+                    1 = multiplicative (Heaviside),
+                    2 = additive
+    dt : float
+        Time step (in months).
+    tend_option : {"F", "C"}
+        Option for numerical derivative:
+            - "F" : forward difference
+            - "C" : centered difference
+    fitting_option_B : {"MAC", "LR"}
+        Method for estimating multiplicative noise amplitude `B`:
+            - "MAC" : Moment Approximation Calibration
+            - "LR" : Linear regression
+    fitting_option_red : {"LR", "AR1", "ARn"}
+        Method for fitting red noise memory:
+            - "LR"  : regression of residual derivative
+            - "AR1" : lag-1 autocorrelation
+            - "ARn" : regression using autocorrelation decay
+
+    Returns
+    -------
+    final_par : ndarray, shape (N, 5)
+        Stacked parameter array including:
+            - Deterministic T-equation parameters (seasonalized)
+            - Deterministic h-equation parameters (seasonalized)
+            - Noise parameters:
+                [sigma_T, sigma_h, B, m_T, m_h, n_T, n_h, n_g]
+
+    Notes
+    -----
+    - Seasonal terms are expanded into constant, annual, and semi-annual
+      components depending on the fitting options.
+    - Noise fitting can handle white/red and additive/multiplicative cases.
+    - If multiplicative noise is selected and `fitting_option_B="MAC"`, 
+      the MAC method overwrites `sigma_T` and `B`.
+
+    """
+
+    T = np.ravel(T)
+    h = np.ravel(h)
+
+    par_option_T = np.array(par_option_T)
+    par_option_h = np.array(par_option_h)
+    par_option_noise = np.array(par_option_noise)
+
+    n_T, n_h, n_g = par_option_noise.astype(int)
+
+    par_option_T_onoff = (par_option_T >= 1).astype(int)
+    par_option_h_onoff = (par_option_h >= 1).astype(int)
+
+
+    # Derivative
+    if tend_option == "F":
+        Tdiff = np.diff(T) / dt
+        hdiff = np.diff(h) / dt
+        Ts = T[:-1]
+        hs = h[:-1]
+    elif tend_option == "C":
+        Tdiff = (T[2:] - T[:-2]) / (2 * dt)
+        hdiff = (h[2:] - h[:-2]) / (2 * dt)
+        Ts = T[1:-1]
+        hs = h[1:-1]
+
+    # Design matrices
+    Tv = np.column_stack([Ts, hs, Ts**2, -Ts**3, Ts * hs]) * par_option_T_onoff
+    hv = np.column_stack([-Ts, -hs, -Ts**2]) * par_option_h_onoff
+
+
+    # Seasonal features
+    w = 2 * np.pi / 12
+    t = np.arange(0, len(T)) * dt
+
+    if tend_option == "F":
+        t = t[:-1]
+        t_shift = t + 0.5 * dt
+    elif tend_option == "C": 
+        t = t[1:-1]
+        t_shift = t 
+
+
+
+    def add_seasonal(X=Tv, par_option=par_option_T, option="T"):
+        X_seasonal_extended = []
+        for i in range(0, len(par_option)):
+            if par_option[i] == 0:
+                on_and_off = np.array([0, 0, 0, 0, 0])
+            elif par_option[i] == 1:
+                on_and_off = np.array([1, 0, 0, 0, 0])
+            elif par_option[i] == 3:
+                on_and_off = np.array([1, 1, 1, 0, 0])
+            elif par_option[i] == 5:
+                on_and_off = np.array([1, 1, 1, 1, 1])
+            else:
+                raise ValueError(f"Invalid {option} fitting option {par_option}: "
+                                  "Value must be one of the following — 0 (constant zero), "
+                                  "1 (constant non-zero), 3 (annual variation), or 5 (annual + semi-annual variation).")
+
+
+       
+            X_seasonal = np.column_stack([
+                        X[:, i] * np.sin(w * t_shift),
+                        X[:, i] * np.cos(w * t_shift),
+                        X[:, i] * np.sin(2 * w * t_shift),
+                        X[:, i] * np.cos(2 * w * t_shift)])
+            X_seasonal = np.hstack([X[:, i].reshape(-1, 1), X_seasonal])
+            X_seasonal = X_seasonal * on_and_off
+            X_seasonal_extended.append(X_seasonal)
+
+        X_seasonal_extended = np.array(X_seasonal_extended)
+        X_seasonal_extended = np.transpose(X_seasonal_extended, (1, 0, 2))  # shape becomes (1199, 5, 5)
+        if option == "T":
+        	X_seasonal_extended = X_seasonal_extended.reshape(X.shape[0], 25) 
+        if option == "h":
+            X_seasonal_extended = X_seasonal_extended.reshape(X.shape[0], 15)
+        return X_seasonal_extended
+
+    Tv = add_seasonal(Tv, par_option_T, option="T")
+    hv = add_seasonal(hv, par_option_h, option="h")
+
+    par_T, _, res_T = regress_std(Tdiff, Tv)
+    par_h, _, res_h = regress_std(hdiff, hv)
+
+
+    def extract_seasonal(params, option):
+        my_params = []
+        if option == "T":
+            index = 5
+        elif option == "h":
+            index = 3
+        
+        params = params.reshape(index, 5)
+        for i in range(0, index):
+            A_a = np.sqrt(params[i,1]**2 + params[i,2]**2)
+            phi_a = np.mod(np.arctan2(params[i,2], params[i,1]), 2*np.pi)
+ 
+            A_sa = np.sqrt(params[i,3]**2 + params[i,4]**2)
+            phi_sa = np.mod(np.arctan2(params[i,4], params[i,3]), 2*np.pi)
+
+            if np.pi < phi_a:
+                phi_a = phi_a - 2*np.pi
+            if np.pi < phi_sa:
+                phi_sa = phi_sa - 2*np.pi
+            my_param = [params[i,0], A_a, phi_a, A_sa, phi_sa]
+            my_params.append(my_param)
+        my_params = np.array(my_params) 
+
+        return my_params
+
+
+    par_T = extract_seasonal(par_T, option="T")
+    par_h = extract_seasonal(par_h, option="h")
+
+
+    # Noise fitting
+    res_T_norm = res_T * np.sqrt(dt)
+    res_h_norm = res_h * np.sqrt(dt)
+
+    if n_T == 1 and n_h == 1 and n_g == 2:    # (white, white, additive)
+        sigma_T = np.std(res_T_norm)
+        sigma_h = np.std(res_h_norm)
+        B = m_T = m_h = 0.0 # np.nan 
+
+    elif n_T == 0 and n_h == 0 and n_g == 2:  # (red, red, additive)
+        if fitting_option_red == "LR":
+            res_Tdiff = np.diff(res_T) / dt
+            res_Ts = res_T[:-1]
+            m_T, _, _ = regress_std(res_Tdiff, res_Ts)
+
+            res_hdiff = np.diff(res_h) / dt
+            res_hs = res_h[:-1]
+            m_h, _, _ = regress_std(res_hdiff, res_hs)
+            m_T = np.abs(m_T.item())
+            m_h = np.abs(m_h.item())
+        elif fitting_option_red == "AR1":
+            m_T = -np.log(np.corrcoef(res_T[:-1], res_T[1:])[0, 1]) / dt
+            m_h = -np.log(np.corrcoef(res_h[:-1], res_h[1:])[0, 1]) / dt
+        elif fitting_option_red == "ARn":
+            # m_T part
+            r, lags = np.correlate(res_T - np.mean(res_T), res_T - np.mean(res_T), mode='full'), np.arange(-len(res_T)+1, len(res_T))
+            r = r / np.max(r)
+            lags_dt = lags * dt
+            positive_lags = lags_dt >= 0
+            r = r[positive_lags]
+            lags = lags_dt[positive_lags]
+
+            try:
+                cutoff = np.where(r < 0)[0][0]
+                r = r[:cutoff]
+                lags = lags[:cutoff]
+            except IndexError:
+                pass
+
+            y = -np.log(r)
+            X = lags.reshape(-1, 1)
+            m_T, _, _ = regress_std(y, X)
+
+            # m_h part
+            r, lags = np.correlate(res_h - np.mean(res_h), res_h - np.mean(res_h), mode='full'), np.arange(-len(res_h)+1, len(res_h))
+            r = r / np.max(r)
+            lags_dt = lags * dt
+            positive_lags = lags_dt >= 0
+            r = r[positive_lags]
+            lags = lags_dt[positive_lags]
+
+            try:
+                cutoff = np.where(r < 0)[0][0]
+                r = r[:cutoff]
+                lags = lags[:cutoff]
+            except IndexError:
+                pass
+
+            y = -np.log(r)
+            X = lags.reshape(-1, 1)
+            m_h, _, _ = regress_std(y, X)
+            m_T = m_T.item() 
+            m_h = m_h.item()
+        else:
+            print(f"Invalid fitting method for red noise: {fitting_option_red}")
+            sys.exit()
+
+
+
+        sigma_T = np.std(res_T)
+        sigma_h = np.std(res_h)
+        B = 0 #np.nan
+
+    elif n_T == 1 and n_h == 1 and n_g in (0, 1): # (white, white, multiplicative (either full or Heaviside))
+        n = 10 * 12
+        res_T_var = np.zeros(10**5)
+        T_var = np.zeros(10**5)
+
+        for i in range(10**5):
+            ind = np.random.randint(0, len(res_T_norm), size=n)
+            res_T_var[i] = np.var(res_T_norm[ind], ddof=1)  # ddof=1 for unbiased (same as MATLAB)
+    
+            if n_g == 0.0:
+                T_var[i] = np.var(Ts[ind], ddof=1)
+            elif n_g == 1.0:
+                T_var[i] = np.var(np.heaviside(Ts[ind], 0) * Ts[ind], ddof=1)
+
+
+
+        X = np.column_stack([T_var, np.ones(len(T_var))])
+        par_T_res, _, _ = regress_std(res_T_var, X)
+        if par_T_res[0] < 0:
+            B = 0.0
+        else:
+            B = np.sqrt(par_T_res[0] / par_T_res[1])
+        sigma_T = np.sqrt(par_T_res[1])
+        sigma_h = np.std(res_h_norm)
+        m_T = m_h = 0.0 
+
+    elif n_T == 0 and n_h == 0 and n_g in (0, 1): # (red, red, multiplicative (either full or Heaviside))
+        n = 10 * 12
+        res_T_var, T_var = [], []
+        for _ in range(10000):
+            ind = np.random.randint(0, len(res_T_norm), size=n)
+            res_T_var.append(np.var(res_T_norm[ind]))
+            if n_g == 0:
+                T_var.append(np.var(Ts[ind]))
+            else:
+                T_var.append(np.var(heaviside(Ts[ind]) * Ts[ind]))
+        X = np.column_stack([T_var, np.ones(len(T_var))])
+        res_T_var = np.array(res_T_var) 
+        par_T_res, _, _ = regress_std(res_T_var, X)
+        sigma_T = np.sqrt(par_T_res[1])
+
+        if par_T_res[0] < 0:
+            B = 0.0
+        else:
+            B = np.sqrt(par_T_res[0] / par_T_res[1])
+        xi = res_T / (sigma_T * (1 + B * Ts))
+        m_T, _, _ = regress_std(np.diff(xi), xi[:-1].reshape(-1, 1))
+        m_T = np.abs(m_T.item())
+        res_hdiff = np.diff(res_h) / dt
+        res_hs = res_h[:-1]
+        m_h, _, _ = regress_std(res_hdiff, res_hs.reshape(-1, 1))
+        m_h = np.abs(m_h.item())
+        sigma_h = np.std(res_h)
+
+
+    #if (not np.isfinite(m_T) or m_T < 0.) or (not np.isfinite(m_h) or m_h < 0.):
+    #    raise ValueError(f"Estimated m_T = {m_T} or m_h = {m_h} is not a valid number.\n"
+    #                     f"This error typically occurs when trying to fit red noise expressions into the time series "
+    #                     f"generated with white noise.\nTry using the fitting option for white noise instead.")
+
+    #if np.isnan(B):
+    #        raise ValueError(f"B = np.sqrt({par_T_res[0]} / {par_T_res[1]}) is NaN.\n"
+    #                         f"This error typically occurs when trying to fit too many annual and semi-annual components "
+    #                         f"at the same time with multiplicative noise.\n"
+    #                         f"It can also occur when attempting to fit red noise expressions to a time series generated with white noise.\n"
+    #                         f"Try reducing the seasonality components and/or using the fitting option for white noise instead.")
+        #sys.exit(
+
+
+    par_noise = np.array([[sigma_T, 0, 0, 0, 0], 
+                 [sigma_h, 0, 0, 0, 0], 
+                 [B, 0, 0, 0, 0], 
+                 [m_T, 0, 0, 0, 0], 
+                 [m_h, 0, 0, 0, 0],
+                 [n_T, 0, 0, 0, 0],
+                 [n_h, 0, 0, 0, 0],
+                 [n_g, 0, 0, 0, 0]])
+
+    final_par = np.vstack([par_T, par_h, par_noise])
+
+
+    if (n_g == 0 or n_g == 1) and fitting_option_B == "MAC":
+        final_par = fit_MAC(T, h, final_par)  # Calculate sigma_T and B using MAC, overwrite par
+
+    return final_par

pyCRO/fit_MAC.py

@@ -0,0 +1,90 @@
+import sys
+import numpy as np
+
+def fit_MAC(T, h, par):
+    # Extract annual mean values (first element from each parameter)
+    R = par[0,0]
+    F1 = par[1,0]
+    epsilon = par[5,0]
+    F2 = par[6,0]
+    b_T = par[2,0]
+    c_T = par[3,0]
+    d_T = par[4,0]
+    B_ref = par[10,0]
+    sigma_T_ref = par[8,0]
+    n_g = par[15,0]
+
+    T = T - np.mean(T)
+    h = h - np.mean(h)
+
+    if n_g == 0:  # linear
+        T2 = np.mean(T**2)
+        T3 = np.mean(T**3)
+        T4 = np.mean(T**4)
+        T5 = np.mean(T**5)
+        T6 = np.mean(T**6)
+        hT = np.mean(h * T)
+        hT2 = np.mean(h * T**2)
+        hT3 = np.mean(h * T**3)
+        hT4 = np.mean(h * T**4)
+
+        TA =  R*T2 + F1*hT  + b_T*T3 - c_T*T4 + d_T*hT2
+        T2A = R*T3 + F1*hT2 + b_T*T4 - b_T*T2**2 - c_T*T5 + c_T*T3*T2 + d_T*hT3 - d_T*hT*T2
+        T3A = R*T4 + F1*hT3 + b_T*T5 - b_T*T3*T2 - c_T*T6 + c_T*T3**2 + d_T*hT4 - d_T*hT*T3
+
+        k = -(2/3)*(T3A - 1.5*T2A*T3/T2 - 3*TA*T2) / (T4 - T3**2/T2 - T2**2)
+        if np.isnan(np.sqrt(-2 * TA - k * T2)):
+            sigma_T = sigma_T_ref
+            B = B_ref
+        else:
+            sigma_T = np.sqrt(-2 * TA - k * T2)
+            B = -(T2A + k * T3) / (2 * sigma_T**2 * T2)
+
+
+    elif n_g == 1:  # Heaviside-linear
+        T2 = np.mean(T**2)
+        T3 = np.mean(T**3)
+        T4 = np.mean(T**4)
+        T5 = np.mean(T**5)
+        T6 = np.mean(T**6)
+        hT = np.mean(h * T)
+        hT2 = np.mean(h * T**2)
+        hT3 = np.mean(h * T**3)
+        hT4 = np.mean(h * T**4)
+
+        TA =  R*T2 + F1*hT +  b_T*T3 - c_T*T4 + d_T*hT2
+        T2A = R*T3 + F1*hT2 + b_T*T4 - b_T*T2**2 - c_T*T5 + c_T*T3*T2 + d_T*hT3 - d_T*hT*T2
+        T3A = R*T4 + F1*hT3 + b_T*T5 - b_T*T3*T2 - c_T*T6 + c_T*T3**2 + d_T*hT4 - d_T*hT*T3
+
+        T1p = T
+        T2p = T**2
+        T3p = T**3
+        T4p = T**4
+
+        T1p = T1p[T1p >= 0]
+        T2p = T2p[T2p >= 0]
+        T3p = T3p[T3p >= 0]
+        T4p = T4p[T4p >= 0]
+
+        T1p = np.mean(T1p)
+        T2p = np.mean(T2p)
+        T3p = np.mean(T3p)
+        T4p = np.mean(T4p)
+
+        k1 = (4*TA*T2 + T2A*(2*T3p/T2p - T1p*T2/T2p) - (4/3)*T3A) / (
+            2*T4p - 2*T2p*T2 - 2*(T3p**2)/T2p + T2*T3p*T1p/T2p
+        )
+        k2 = -(T2A + T3p*k1) / (2 * T2p)
+
+        B = k1 / k2
+        sigma_T = np.sqrt(k2 / B)
+
+        if np.isnan(B) or np.isnan(sigma_T):
+            sigma_T = sigma_T_ref
+            B = B_ref
+
+    par_out = par[:]
+    par_out[8, 0] = sigma_T  # index 9 in MATLAB (0-based here)
+    par_out[10, 0] = B       # index 11 in MATLAB
+
+    return par_out