disdrodb 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

disdrodb/psd/fitting.py

@@ -20,98 +20,150 @@ import scipy.stats as ss
 import xarray as xr
 from scipy.integrate import quad
 from scipy.optimize import minimize
-from scipy.special import gamma, gammainc, gammaln  # Regularized lower incomplete gamma function
-
+from scipy.special import gamma, gammaln  # Regularized lower incomplete gamma function
+
+from disdrodb.constants import DIAMETER_DIMENSION
+from disdrodb.l1.fall_velocity import get_dataset_fall_velocity
+from disdrodb.l2.empirical_dsd import (
+    get_median_volume_drop_diameter,
+    get_moment,
+    get_normalized_intercept_parameter_from_moments,
+    get_total_number_concentration,
+)
 from disdrodb.psd.models import ExponentialPSD, GammaPSD, LognormalPSD, NormalizedGammaPSD
+from disdrodb.utils.manipulations import get_diameter_bin_edges
 from disdrodb.utils.warnings import suppress_warnings
 
 # gamma(>171) return inf !
 
+####--------------------------------------------------------------------------------------.
+#### Notes
+## Variable requirements for fitting PSD Models
+# - drop_number_concentration and diameter coordinates
+# - Always recompute other parameters to ensure not use model parameters of L2M
+
+# ML: None
+
+# MOM: moments
+# --> get_moment(drop_number_concentration, diameter, diameter_bin_width, moment)
+
+# GS: fall_velocity if target optimization is R (rain)
+# - NormalizedGamma: "Nw", "D50"
+# --> get_normalized_intercept_parameter_from_moments(moment_3, moment_4)
+# --> get_median_volume_drop_diameter(drop_number_concentration, diameter, diameter_bin_width):
+# --> get_mean_volume_drop_diameter(moment_3, moment_4)  (Dm)
+
+# - LogNormal,Exponential, Gamma: Nt
+# --> get_total_number_concentration(drop_number_concentration, diameter_bin_width)
+
 
 ####--------------------------------------------------------------------------------------.
 #### Goodness of fit (GOF)
-def compute_gof_stats(drop_number_concentration, psd):
+def compute_gof_stats(obs, pred, dim=DIAMETER_DIMENSION):
     """
-    Compute various goodness-of-fit (GoF) statistics between observed and predicted values.
+    Compute various goodness-of-fit (GoF) statistics between obs and predicted values.
 
     Parameters
     ----------
-    - drop_number_concentration: xarray.DataArray with dimensions ('time', 'diameter_bin_center')
-    - psd: instance of PSD class
+    obs: xarray.DataArray
+        Observations DataArray with at least dimension ``dim``.
+    pred: xarray.DataArray
+        Predictions DataArray with at least dimension ``dim``.
+    dim: str
+        DataArray dimension over which to compute GOF statistics.
+        The default is DIAMETER_DIMENSION.
 
     Returns
     -------
-    - ds: xarray.Dataset containing the computed GoF statistics
+    ds: xarray.Dataset
+        Dataset containing the computed GoF statistics.
     """
     from disdrodb.l2.empirical_dsd import get_mode_diameter
 
-    # Retrieve diameter bin width
-    diameter = drop_number_concentration["diameter_bin_center"]
-    diameter_bin_width = drop_number_concentration["diameter_bin_width"]
+    # Retrieve diameter and diameter bin width
+    diameter = obs["diameter_bin_center"]
+    diameter_bin_width = obs["diameter_bin_width"]
+
+    # Compute errors
+    error = obs - pred
 
-    # Define observed and predicted values and compute errors
-    observed_values = drop_number_concentration
-    fitted_values = psd(diameter)  # .transpose(*observed_values.dims)
-    error = observed_values - fitted_values
+    # Compute max obs and pred
+    obs_max = obs.max(dim=dim, skipna=False)
+    pred_max = pred.max(dim=dim, skipna=False)
+
+    # Compute NaN mask
+    mask_nan = np.logical_or(np.isnan(obs_max), np.isnan(pred_max))
 
     # Compute GOF statistics
     with suppress_warnings():
-        # Compute Pearson correlation
-        pearson_r = xr.corr(observed_values, fitted_values, dim="diameter_bin_center")
-
-        # Compute MSE
-        mse = (error**2).mean(dim="diameter_bin_center")
+        # Compute Pearson Correlation
+        pearson_r = xr.corr(obs, pred, dim=dim)
+
+        # Compute Mean Absolute Error (MAE)
+        mae = np.abs(error).mean(dim=dim, skipna=False)
+
+        # Compute maximum absolute error
+        max_error = np.abs(error).max(dim=dim, skipna=False)
+        relative_max_error = xr.where(max_error == 0, 0, xr.where(obs_max == 0, np.nan, max_error / obs_max))
+
+        # Compute deviation of N(D) at distribution mode
+        mode_deviation = obs_max - pred_max
+        mode_relative_deviation = xr.where(
+            mode_deviation == 0,
+            0,
+            xr.where(obs_max == 0, np.nan, mode_deviation / obs_max),
+        )
 
-        # Compute maximum error
-        max_error = error.max(dim="diameter_bin_center")
-        relative_max_error = error.max(dim="diameter_bin_center") / observed_values.max(dim="diameter_bin_center")
+        # Compute diameter difference of the distribution mode
+        diameter_mode_pred = get_mode_diameter(pred, diameter)
+        diameter_mode_obs = get_mode_diameter(obs, diameter)
+        diameter_mode_deviation = diameter_mode_obs - diameter_mode_pred
 
         # Compute difference in total number concentration
-        total_number_concentration_obs = (observed_values * diameter_bin_width).sum(dim="diameter_bin_center")
-        total_number_concentration_pred = (fitted_values * diameter_bin_width).sum(dim="diameter_bin_center")
+        total_number_concentration_obs = (obs * diameter_bin_width).sum(dim=dim, skipna=False)
+        total_number_concentration_pred = (pred * diameter_bin_width).sum(dim=dim, skipna=False)
         total_number_concentration_difference = total_number_concentration_pred - total_number_concentration_obs
 
         # Compute Kullback-Leibler divergence
         # - Compute pdf per bin
-        pk_pdf = observed_values / total_number_concentration_obs
-        qk_pdf = fitted_values / total_number_concentration_pred
+        pk_pdf = obs / total_number_concentration_obs
+        qk_pdf = pred / total_number_concentration_pred
 
         # - Compute probabilities per bin
         pk = pk_pdf * diameter_bin_width
-        pk = pk / pk.sum(dim="diameter_bin_center")  # this might not be necessary
+        pk = pk / pk.sum(dim=dim, skipna=False)  # this might not be necessary
         qk = qk_pdf * diameter_bin_width
-        qk = qk / qk.sum(dim="diameter_bin_center")  # this might not be necessary
+        qk = qk / qk.sum(dim=dim, skipna=False)  # this might not be necessary
 
-        # - Compute divergence
+        # - Compute log probability ratio
+        epsilon = 1e-10
+        pk = xr.where(pk == 0, epsilon, pk)
+        qk = xr.where(qk == 0, epsilon, qk)
         log_prob_ratio = np.log(pk / qk)
         log_prob_ratio = log_prob_ratio.where(np.isfinite(log_prob_ratio))
-        kl_divergence = (pk * log_prob_ratio).sum(dim="diameter_bin_center")
 
-        # Other statistics that can be computed also from different diameter discretization
-        # - Compute max deviation at distribution mode
-        max_deviation = observed_values.max(dim="diameter_bin_center") - fitted_values.max(dim="diameter_bin_center")
-        max_relative_deviation = max_deviation / fitted_values.max(dim="diameter_bin_center")
-
-        # - Compute diameter difference of the distribution mode
-        diameter_mode_deviation = get_mode_diameter(observed_values, diameter) - get_mode_diameter(
-            fitted_values,
-            diameter,
-        )
+        # - Compute divergence
+        kl_divergence = (pk * log_prob_ratio).sum(dim=dim, skipna=False)
+        kl_divergence = xr.where((error == 0).all(dim=dim), 0, kl_divergence)
 
     # Create an xarray.Dataset to hold the computed statistics
     ds = xr.Dataset(
         {
-            "r2": pearson_r**2,  # Squared Pearson correlation coefficient
-            "mse": mse,  # Mean Squared Error
-            "max_error": max_error,  # Maximum Absolute Error
-            "relative_max_error": relative_max_error,  # Relative Maximum Error
-            "total_number_concentration_difference": total_number_concentration_difference,
-            "kl_divergence": kl_divergence,  # Kullback-Leibler divergence
-            "max_deviation": max_deviation,  # Deviation at distribution mode
-            "max_relative_deviation": max_relative_deviation,  # Relative deviation at mode
-            "diameter_mode_deviation": diameter_mode_deviation,  # Difference in mode diameters
+            "R2": pearson_r**2,  # Squared Pearson correlation coefficient
+            "MAE": mae,  # Mean Absolute Error
+            "MaxAE": max_error,  # Maximum Absolute Error
+            "RelMaxAE": relative_max_error,  # Relative Maximum Absolute Error
+            "PeakDiff": mode_deviation,  # Difference at distribution peak
+            "RelPeakDiff": mode_relative_deviation,  # Relative difference at peak
+            "DmodeDiff": diameter_mode_deviation,  # Difference in mode diameters
+            "NtDiff": total_number_concentration_difference,
+            "KLDiv": kl_divergence,  # Kullback-Leibler divergence
         },
     )
+    # Round
+    ds = ds.round(2)
+    # Mask where input obs or pred is NaN
+    ds = ds.where(~mask_nan)
     return ds
 
 
@@ -176,13 +228,13 @@ def get_expected_probabilities(params, cdf_func, pdf_func, bin_edges, probabilit
 def get_adjusted_nt(cdf, params, Nt, bin_edges):
     """Adjust Nt for the proportion of missing drops. See Johnson's et al., 2013 Eqs. 3 and 4."""
     # Estimate proportion of missing drops (Johnson's 2011 Eqs. 3)
-    # --> Alternative: p = 1 - np.sum(pdf(diameter, params)* diameter_bin_width)  # [-]
+    # --> Alternative:
+    # - p = 1 - np.sum(pdf(diameter, params)* diameter_bin_width)  # [-]
+    # - p = 1 - np.sum((Lambda ** (mu + 1)) / gamma(mu + 1) * D**mu * np.exp(-Lambda * D) * diameter_bin_width)  # [-]
     p = 1 - np.diff(cdf([bin_edges[0], bin_edges[-1]], params)).item()  # [-]
-    # Adjusts Nt for the proportion of drops not observed
-    #   p = np.clip(p, 0, 1 - 1e-12)
-    if np.isclose(p, 1, atol=1e-12):
-        return np.nan
-    return Nt / (1 - p)  # [m-3]
+    # Adjusts Nt for the proportion of missing drops
+    nt_adj = np.nan if np.isclose(p, 1, atol=1e-12) else Nt / (1 - p)  # [m-3]
+    return nt_adj
 
 
 def compute_negative_log_likelihood(
@@ -206,7 +258,7 @@ def compute_negative_log_likelihood(
     bin_edges : array-like
         Edges of the bins (length N+1).
     counts : array-like
-        Observed counts in each bin (length N).
+        obs counts in each bin (length N).
     cdf_func : callable
         Cumulative distribution function of the distribution.
     pdf_func : callable
@@ -259,6 +311,8 @@ def compute_negative_log_likelihood(
 
 def estimate_lognormal_parameters(
     counts,
+    mu,
+    sigma,
     bin_edges,
     probability_method="cdf",
     likelihood="multinomial",
@@ -273,6 +327,12 @@ def estimate_lognormal_parameters(
     ----------
     counts : array-like
         The counts for each bin in the histogram.
+    mu: float
+        The initial guess of the mean of the log of the distribution.
+        A good default value is 0.
+    sigma: float
+        The initial guess of the standard deviation of the log distribution.
+        A good default value is 1.
     bin_edges : array-like
         The edges of the bins.
     probability_method : str, optional
@@ -306,9 +366,9 @@ def estimate_lognormal_parameters(
     ----------
     .. [1] https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.lognorm.html#scipy.stats.lognorm
     """
-    # LogNormal
-    # - mu = log(scale)
-    # - loc = 0
+    # Definite initial guess for the parameters
+    scale = np.exp(mu)  # mu = np.log(scale)
+    initial_params = [sigma, scale]
 
     # Initialize bad results
     null_output = (
@@ -329,9 +389,6 @@ def estimate_lognormal_parameters(
         sigma, scale = params
         return sigma > 0 and scale > 0
 
-    # Definite initial guess for the parameters
-    initial_params = [1.0, 1.0]  # sigma, scale
-
     # Define bounds for sigma and scale
     bounds = [(1e-6, None), (1e-6, None)]
 
@@ -375,6 +432,7 @@ def estimate_lognormal_parameters(
 
 def estimate_exponential_parameters(
     counts,
+    Lambda,
     bin_edges,
     probability_method="cdf",
     likelihood="multinomial",
@@ -389,6 +447,10 @@ def estimate_exponential_parameters(
     ----------
     counts : array-like
         The counts for each bin in the histogram.
+    Lambda : float
+        The initial guess of the scale parameter.
+        scale = 1 / lambda correspond to the scale parameter of the scipy.stats.expon distribution.
+        A good default value is 1.
     bin_edges : array-like
         The edges of the bins.
     probability_method : str, optional
@@ -421,6 +483,10 @@ def estimate_exponential_parameters(
     ----------
     .. [1] https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.expon.html
     """
+    # Definite initial guess for parameters
+    scale = 1 / Lambda
+    initial_params = [scale]
+
     # Initialize bad results
     null_output = {"N0": np.nan, "Lambda": np.nan} if output_dictionary else np.array([np.nan, np.nan])
 
@@ -438,9 +504,6 @@ def estimate_exponential_parameters(
         scale = params[0]
         return scale > 0
 
-    # Definite initial guess for the scale parameter
-    initial_params = [1.0]  # scale
-
     # Define bounds for scale
     bounds = [(1e-6, None)]
 
@@ -485,8 +548,8 @@ def estimate_exponential_parameters(
 
 def estimate_gamma_parameters(
     counts,
-    a,
-    scale,
+    mu,
+    Lambda,
     bin_edges,
     probability_method="cdf",
     likelihood="multinomial",
@@ -501,11 +564,13 @@ def estimate_gamma_parameters(
     ----------
     counts : array-like
         The counts for each bin in the histogram.
-    a: float
-        The shape parameter of the scipy.stats.gamma distribution.
-        A good default value is 1.
-    scale: float
-        The scale parameter of the scipy.stats.gamma distribution.
+    mu: float
+        The initial guess of the shape parameter.
+        a = mu + 1 correspond to the shape parameter of the scipy.stats.gamma distribution.
+        A good default value is 0.
+    lambda: float
+        The initial guess of the scale parameter.
+        scale = 1 / lambda correspond to the scale parameter of the scipy.stats.gamma distribution.
         A good default value is 1.
     bin_edges : array-like
         The edges of the bins.
@@ -541,6 +606,11 @@ def estimate_gamma_parameters(
     .. [1] https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.gamma.html
 
     """
+    # Define initial guess for parameters
+    a = mu + 1  # (mu = a-1, a = mu+1)
+    scale = 1 / Lambda
+    initial_params = [a, scale]
+
     # Initialize bad results
     null_output = (
         {"N0": np.nan, "mu": np.nan, "lambda": np.nan} if output_dictionary else np.array([np.nan, np.nan, np.nan])
@@ -561,9 +631,6 @@ def estimate_gamma_parameters(
         a, scale = params
         return a > 0.1 and scale > 0  # using a > 0 cause some troubles
 
-    # Definite initial guess for the parameters
-    initial_params = [a, scale]  # (mu=a-1, a=mu+1)
-
     # Define bounds for a and scale
     bounds = [(1e-6, None), (1e-6, None)]
 
@@ -603,7 +670,7 @@ def estimate_gamma_parameters(
 
     # Compute N0
     # - Use logarithmic computations to prevent overflow
-    # - N0 = Nt * Lambda ** (mu + 1) / gamma(mu + 1)
+    # - N0 = Nt * Lambda ** (mu + 1) / gamma(mu + 1)  # [m-3 * mm^(-mu-1)]
     with suppress_warnings():
         log_N0 = np.log(Nt) + (mu + 1) * np.log(Lambda) - gammaln(mu + 1)
         N0 = np.exp(log_N0)
@@ -617,12 +684,59 @@ def estimate_gamma_parameters(
     return output
 
 
+def _get_initial_lognormal_parameters(ds, mom_method=None):
+    default_mu = 0  # mu = np.log(scale)
+    default_sigma = 1
+    if mom_method is None or mom_method == "None":
+        ds_init = xr.Dataset(
+            {
+                "mu": default_mu,
+                "sigma": default_sigma,
+            },
+        )
+    else:
+        ds_init = get_mom_parameters(
+            ds=ds,
+            psd_model="LognormalPSD",
+            mom_methods=mom_method,
+        )
+        # If initialization results in some not finite number, set default value
+        ds_init["mu"] = xr.where(
+            np.logical_and(np.isfinite(ds_init["mu"]), ds_init["mu"] > 0),
+            ds_init["mu"],
+            default_mu,
+        )
+        ds_init["sigma"] = xr.where(np.isfinite(ds_init["sigma"]), ds_init["sigma"], default_sigma)
+    return ds_init
+
+
+def _get_initial_exponential_parameters(ds, mom_method=None):
+    default_lambda = 1  # lambda = 1 /scale
+    if mom_method is None or mom_method == "None":
+        ds_init = xr.Dataset(
+            {
+                "Lambda": default_lambda,
+            },
+        )
+    else:
+        ds_init = get_mom_parameters(
+            ds=ds,
+            psd_model="ExponentialPSD",
+            mom_methods=mom_method,
+        )
+        # If initialization results in some not finite number, set default value
+        ds_init["Lambda"] = xr.where(np.isfinite(ds_init["Lambda"]), ds_init["Lambda"], default_lambda)
+    return ds_init
+
+
 def _get_initial_gamma_parameters(ds, mom_method=None):
-    if mom_method is None:
+    default_mu = 0  #  a = mu + 1  |   mu = a - 1
+    default_lambda = 1  #  scale = 1 / Lambda
+    if mom_method is None or mom_method == "None":
         ds_init = xr.Dataset(
             {
-                "a": xr.ones_like(ds["M1"]),
-                "scale": xr.ones_like(ds["M1"]),
+                "mu": default_mu,
+                "Lambda": default_lambda,
             },
         )
     else:
@@ -631,11 +745,13 @@ def _get_initial_gamma_parameters(ds, mom_method=None):
             psd_model="GammaPSD",
             mom_methods=mom_method,
         )
-        ds_init["a"] = ds_init["mu"] + 1
-        ds_init["scale"] = 1 / ds_init["Lambda"]
         # If initialization results in some not finite number, set default value
-        ds_init["a"] = xr.where(np.isfinite(ds_init["a"]), ds_init["a"], ds["M1"])
-        ds_init["scale"] = xr.where(np.isfinite(ds_init["scale"]), ds_init["scale"], ds["M1"])
+        ds_init["mu"] = xr.where(
+            np.logical_and(np.isfinite(ds_init["mu"]), ds_init["mu"] > -1),
+            ds_init["mu"],
+            default_mu,
+        )
+        ds_init["Lambda"] = xr.where(np.isfinite(ds_init["Lambda"]), ds_init["Lambda"], default_lambda)
     return ds_init
 
 
@@ -663,13 +779,14 @@ def get_gamma_parameters(
           with the specified mom_method.
     init_method: str or list
         The method(s) of moments used to initialize the gamma parameters.
-        If None, the scale parameter is set to 1 and mu to 0 (a=1).
+        If None (or 'None'), the scale parameter is set to 1 and mu to 0 (a=1).
     probability_method : str, optional
         Method to compute probabilities. The default value is ``cdf``.
     likelihood : str, optional
         Likelihood function to use for fitting. The default value is ``multinomial``.
     truncated_likelihood : bool, optional
         Whether to use truncated likelihood. The default value is ``True``.
+        See Johnson et al., 2011 and 2011 for more information.
     optimizer : str, optional
         Optimization method to use. The default value is ``Nelder-Mead``.
 
@@ -687,12 +804,21 @@ def get_gamma_parameters(
     The function uses `xr.apply_ufunc` to fit the lognormal distribution parameters
     in parallel, leveraging Dask for parallel computation.
 
+    References
+    ----------
+    Johnson, R. W., D. V. Kliche, and P. L. Smith, 2011: Comparison of Estimators for Parameters of Gamma Distributions
+    with Left-Truncated Samples. J. Appl. Meteor. Climatol., 50, 296-310, https://doi.org/10.1175/2010JAMC2478.1
+
+    Johnson, R.W., Kliche, D., & Smith, P.L. (2010).
+    Maximum likelihood estimation of gamma parameters for coarsely binned and truncated raindrop size data.
+    Quarterly Journal of the Royal Meteorological Society, 140. DOI:10.1002/qj.2209
+
     """
     # Define inputs
     counts = ds["drop_number_concentration"] * ds["diameter_bin_width"]
-    diameter_breaks = np.append(ds["diameter_bin_lower"].data, ds["diameter_bin_upper"].data[-1])
+    diameter_breaks = get_diameter_bin_edges(ds)
 
-    # Define initial parameters (a, scale)
+    # Define initial parameters (mu, Lambda)
     ds_init = _get_initial_gamma_parameters(ds, mom_method=init_method)
 
     # Define kwargs
@@ -709,10 +835,10 @@ def get_gamma_parameters(
     da_params = xr.apply_ufunc(
         estimate_gamma_parameters,
         counts,
-        ds_init["a"],
-        ds_init["scale"],
+        ds_init["mu"],
+        ds_init["Lambda"],
         kwargs=kwargs,
-        input_core_dims=[["diameter_bin_center"], [], []],
+        input_core_dims=[[DIAMETER_DIMENSION], [], []],
         output_core_dims=[["parameters"]],
         vectorize=True,
         dask="parallelized",
@@ -720,8 +846,6 @@ def get_gamma_parameters(
         output_dtypes=["float64"],
     )
 
-    ds_init.isel(velocity_method=0, time=-3)
-
     # Add parameters coordinates
     da_params = da_params.assign_coords({"parameters": ["N0", "mu", "Lambda"]})
 
@@ -735,7 +859,7 @@ def get_gamma_parameters(
 
 def get_lognormal_parameters(
     ds,
-    init_method=None,  # noqa: ARG001
+    init_method=None,
     probability_method="cdf",
     likelihood="multinomial",
     truncated_likelihood=True,
@@ -779,7 +903,10 @@ def get_lognormal_parameters(
     """
     # Define inputs
     counts = ds["drop_number_concentration"] * ds["diameter_bin_width"]
-    diameter_breaks = np.append(ds["diameter_bin_lower"].data, ds["diameter_bin_upper"].data[-1])
+    diameter_breaks = get_diameter_bin_edges(ds)
+
+    # Define initial parameters (mu, sigma)
+    ds_init = _get_initial_lognormal_parameters(ds, mom_method=init_method)
 
     # Define kwargs
     kwargs = {
@@ -795,8 +922,10 @@ def get_lognormal_parameters(
     da_params = xr.apply_ufunc(
         estimate_lognormal_parameters,
         counts,
+        ds_init["mu"],
+        ds_init["sigma"],
         kwargs=kwargs,
-        input_core_dims=[["diameter_bin_center"]],
+        input_core_dims=[[DIAMETER_DIMENSION], [], []],
         output_core_dims=[["parameters"]],
         vectorize=True,
         dask="parallelized",
@@ -818,7 +947,7 @@ def get_lognormal_parameters(
 
 def get_exponential_parameters(
     ds,
-    init_method=None,  # noqa: ARG001
+    init_method=None,
     probability_method="cdf",
     likelihood="multinomial",
     truncated_likelihood=True,
@@ -863,8 +992,11 @@ def get_exponential_parameters(
 
     """
     # Define inputs
-    counts = ds["drop_number_concentration"] * ds["diameter_bin_width"]
-    diameter_breaks = np.append(ds["diameter_bin_lower"].data, ds["diameter_bin_upper"].data[-1])
+    counts = ds["drop_number_concentration"] * ds["diameter_bin_width"]  # mm-1 m-3 --> m-3
+    diameter_breaks = get_diameter_bin_edges(ds)
+
+    # Define initial parameters (Lambda)
+    ds_init = _get_initial_exponential_parameters(ds, mom_method=init_method)
 
     # Define kwargs
     kwargs = {
@@ -880,8 +1012,9 @@ def get_exponential_parameters(
     da_params = xr.apply_ufunc(
         estimate_exponential_parameters,
         counts,
+        ds_init["Lambda"],
         kwargs=kwargs,
-        input_core_dims=[["diameter_bin_center"]],
+        input_core_dims=[[DIAMETER_DIMENSION], []],
         output_core_dims=[["parameters"]],
         vectorize=True,
         dask="parallelized",
@@ -900,162 +1033,6 @@ def get_exponential_parameters(
     return ds_params
 
 
-####-------------------------------------------------------------------------------------------------------------------.
-
-
-def _estimate_gamma_parameters_johnson(
-    drop_number_concentration,
-    diameter,
-    diameter_breaks,
-    output_dictionary=True,
-    method="Nelder-Mead",
-    mu=0.5,
-    Lambda=3,
-    **kwargs,
-):
-    """Deprecated Maximum likelihood estimation of Gamma model.
-
-    N(D) = N_t * lambda**(mu+1) / gamma(mu+1) D**mu exp(-lambda*D)
-
-    Args:
-        spectra: The DSD for which to find parameters [mm-1 m-3].
-        widths: Class widths for each DSD bin [mm].
-        diams: Class-centre diameters for each DSD bin [mm].
-        mu: Initial value for shape parameter mu [-].
-        lambda_param: Initial value for slope parameter lambda [mm^-1].
-        kwargs: Extra arguments for the optimization process.
-
-    Returns
-    -------
-        Dictionary with estimated mu, lambda, and N0.
-        mu (shape) N0 (scale) lambda(slope)
-
-    Notes
-    -----
-    The last bin counts are not accounted in the fitting procedure !
-
-    References
-    ----------
-    Johnson, R. W., D. V. Kliche, and P. L. Smith, 2011: Comparison of Estimators for Parameters of Gamma Distributions
-    with Left-Truncated Samples. J. Appl. Meteor. Climatol., 50, 296-310, https://doi.org/10.1175/2010JAMC2478.1
-
-    Johnson, R.W., Kliche, D., & Smith, P.L. (2010).
-    Maximum likelihood estimation of gamma parameters for coarsely binned and truncated raindrop size data.
-    Quarterly Journal of the Royal Meteorological Society, 140. DOI:10.1002/qj.2209
-
-    """
-    # Initialize bad results
-    if output_dictionary:
-        null_output = {"mu": np.nan, "lambda": np.nan, "N0": np.nan}
-    else:
-        null_output = np.array([np.nan, np.nan, np.nan])
-
-    # Initialize parameters
-    # --> Ideally with method of moments estimate
-    # --> See equation 8 of Johnson's 2013
-    x0 = [mu, Lambda]
-
-    # Compute diameter_bin_width
-    diameter_bin_width = np.diff(diameter_breaks)
-
-    # Convert drop_number_concentration from mm-1 m-3 to m-3.
-    spectra = np.asarray(drop_number_concentration) * diameter_bin_width
-
-    # Define cost function
-    # - Parameter to be optimized on first positions
-    def _cost_function(parameters, spectra, diameter_breaks):
-        # Assume spectra to be in unit [m-3] (drop_number_concentration*diameter_bin_width) !
-        mu, Lambda = parameters
-        # Precompute gamma integrals between various diameter bins
-        # - gamminc(mu+1) already divides the integral by gamma(mu+1) !
-        pgamma_d = gammainc(mu + 1, Lambda * diameter_breaks)
-        # Compute probability with interval
-        delta_pgamma_bins = pgamma_d[1:] - pgamma_d[:-1]
-        # Compute normalization over interval
-        denominator = pgamma_d[-1] - pgamma_d[0]
-        # Compute cost function
-        # a = mu - 1, x = lambda
-        if mu > -1 and Lambda > 0:
-            cost = np.sum(-spectra * np.log(delta_pgamma_bins / denominator))
-            return cost
-        return np.inf
-
-    # Minimize the cost function
-    with suppress_warnings():
-        bounds = [(0, None), (0, None)]  # Force mu and lambda to be non-negative
-        res = minimize(
-            _cost_function,
-            x0=x0,
-            args=(spectra, diameter_breaks),
-            method=method,
-            bounds=bounds,
-            **kwargs,
-        )
-
-    # Check if the fit had success
-    if not res.success:
-        return null_output
-
-    # Extract parameters
-    mu = res.x[0]  # [-]
-    Lambda = res.x[1]  # [mm-1]
-
-    # Estimate tilde_N_T using the total drop concentration
-    tilde_N_T = np.sum(drop_number_concentration * diameter_bin_width)  # [m-3]
-
-    # Estimate proportion of missing drops (Johnson's 2011 Eqs. 3)
-    with suppress_warnings():
-        D = diameter
-        p = 1 - np.sum((Lambda ** (mu + 1)) / gamma(mu + 1) * D**mu * np.exp(-Lambda * D) * diameter_bin_width)  # [-]
-
-        # Convert tilde_N_T to N_T using Johnson's 2013 Eqs. 3 and 4.
-        # - Adjusts for the proportion of drops not observed
-        N_T = tilde_N_T / (1 - p)  # [m-3]
-
-        # Compute N0
-        N0 = N_T * (Lambda ** (mu + 1)) / gamma(mu + 1)  # [m-3 * mm^(-mu-1)]
-
-    # Compute Dm
-    # Dm = (mu + 4)/ Lambda
-
-    # Compute Nw
-    # Nw = N0* D^mu / f(mu) , with f(mu of the Normalized PSD)
-
-    # Define output
-    output = {"mu": mu, "Lambda": Lambda, "N0": N0} if output_dictionary else np.array([mu, Lambda, N0])
-    return output
-
-
-def get_gamma_parameters_johnson2014(ds, method="Nelder-Mead"):
-    """Deprecated model. See Gamma Model with truncated_likelihood and 'pdf'."""
-    drop_number_concentration = ds["drop_number_concentration"]
-    diameter = ds["diameter_bin_center"]
-    diameter_breaks = np.append(ds["diameter_bin_lower"].data, ds["diameter_bin_upper"].data[-1])
-    # Define kwargs
-    kwargs = {
-        "output_dictionary": False,
-        "diameter_breaks": diameter_breaks,
-        "method": method,
-    }
-    da_params = xr.apply_ufunc(
-        _estimate_gamma_parameters_johnson,
-        drop_number_concentration,
-        diameter,
-        # diameter_bin_width,
-        kwargs=kwargs,
-        input_core_dims=[["diameter_bin_center"], ["diameter_bin_center"]],  # ["diameter_bin_center"],
-        output_core_dims=[["parameters"]],
-        vectorize=True,
-    )
-
-    # Add parameters coordinates
-    da_params = da_params.assign_coords({"parameters": ["mu", "Lambda", "N0"]})
-
-    # Convert to skill Dataset
-    ds_params = da_params.to_dataset(dim="parameters")
-    return ds_params
-
-
 ####-----------------------------------------------------------------------------------------.
 #### Grid Search (GS)
 
@@ -1079,24 +1056,60 @@ def _compute_z(ND, D, dD):
     return Z
 
 
+def _compute_target_variable_error(target, ND_obs, ND_preds, D, dD, V):
+    if target == "Z":
+        errors = np.abs(_compute_z(ND_obs, D, dD) - _compute_z(ND_preds, D, dD))
+    elif target == "R":
+        errors = np.abs(_compute_rain_rate(ND_obs, D, dD, V) - _compute_rain_rate(ND_preds, D, dD, V))
+    else:  # if target == "LWC":
+        errors = np.abs(_compute_lwc(ND_obs, D, dD) - _compute_lwc(ND_preds, D, dD))
+    return errors
+
+
 def _compute_cost_function(ND_obs, ND_preds, D, dD, V, target, transformation, error_order):
     # Assume ND_obs of shape (D bins) and ND_preds of shape (# params, D bins)
     if target == "ND":
         if transformation == "identity":
             errors = np.mean(np.abs(ND_obs[None, :] - ND_preds) ** error_order, axis=1)
+            return errors
         if transformation == "log":
             errors = np.mean(np.abs(np.log(ND_obs[None, :] + 1) - np.log(ND_preds + 1)) ** error_order, axis=1)
-        if transformation == "np.sqrt":
+            return errors
+        if transformation == "sqrt":
             errors = np.mean(np.abs(np.sqrt(ND_obs[None, :]) - np.sqrt(ND_preds)) ** error_order, axis=1)
-    elif target == "Z":
-        errors = np.abs(_compute_z(ND_obs, D, dD) - _compute_z(ND_preds, D, dD))
-    elif target == "R":
-        errors = np.abs(_compute_rain_rate(ND_obs, D, dD, V) - _compute_rain_rate(ND_preds, D, dD, V))
-    elif target == "LWC":
-        errors = np.abs(_compute_lwc(ND_obs, D, dD) - _compute_lwc(ND_preds, D, dD))
-    else:
-        raise ValueError("Invalid target")
-    return errors
+            return errors
+    # if target in ["Z", "R", "LWC"]:
+    return _compute_target_variable_error(target, ND_obs, ND_preds, D, dD, V)
+
+
+def define_param_range(center, step, bounds, factor=2, refinement=20):
+    """
+    Create a refined parameter search range around a center value, constrained to bounds.
+
+    Parameters
+    ----------
+    center : float
+        Center of the range (e.g., current best estimate).
+    step : float
+        Coarse step size used in the first search.
+    bounds : tuple of (float, float)
+        Lower and upper bounds (can include -np.inf, np.inf).
+    factor : float, optional
+        How wide the refined range extends from the center (in multiples of step).
+        Default = 2.
+    refinement : int, optional
+        Factor to refine the step size (smaller step = finer grid).
+        Default = 20.
+
+    Returns
+    -------
+    np.ndarray
+        Array of values constrained to bounds.
+    """
+    lower = max(center - factor * step, bounds[0])
+    upper = min(center + factor * step, bounds[1])
+    new_step = step / refinement
+    return np.arange(lower, upper, new_step)
 
 
 def apply_exponential_gs(
@@ -1132,9 +1145,15 @@ def apply_exponential_gs(
             transformation=transformation,
             error_order=error_order,
         )
+    # Replace inf with NaN
+    errors[~np.isfinite(errors)] = np.nan
+
+    # If all invalid, return NaN parameters
+    if np.all(np.isnan(errors)):
+        return np.array([np.nan, np.nan])
 
-    # Identify best parameter set
-    best_index = np.argmin(errors)
+    # Otherwise, choose the best index
+    best_index = np.nanargmin(errors)
     return np.array([N0_arr[best_index].item(), lambda_arr[best_index].item()])
 
 
@@ -1163,8 +1182,15 @@ def _apply_gamma_gs(mu_values, lambda_values, Nt, ND_obs, D, dD, V, target, tran
             error_order=error_order,
         )
 
-    # Best parameter
-    best_index = np.argmin(errors)
+    # Replace inf with NaN
+    errors[~np.isfinite(errors)] = np.nan
+
+    # If all invalid, return NaN parameters
+    if np.all(np.isnan(errors)):
+        return np.array([np.nan, np.nan, np.nan])
+
+    # Otherwise, choose the best index
+    best_index = np.nanargmin(errors)
     return N0[best_index].item(), mu_arr[best_index].item(), lambda_arr[best_index].item()
 
 
@@ -1181,6 +1207,10 @@ def apply_gamma_gs(
     error_order,
 ):
     """Estimate GammaPSD model parameters using Grid Search."""
+    # Define parameters bounds
+    mu_bounds = (0.01, 20)
+    lambda_bounds = (0.01, 60)
+
     # Define initial set of parameters
     mu_step = 0.5
     lambda_step = 0.5
@@ -1200,10 +1230,13 @@ def apply_gamma_gs(
         transformation=transformation,
         error_order=error_order,
     )
+    if np.isnan(N0):  # if np.nan, return immediately
+        return np.array([N0, mu, Lambda])
 
     # Second round of GS
-    mu_values = np.arange(mu - mu_step * 2, mu + mu_step * 2, step=mu_step / 20)
-    lambda_values = np.arange(Lambda - lambda_step * 2, Lambda + lambda_step * 2, step=lambda_step / 20)
+    mu_values = define_param_range(mu, mu_step, bounds=mu_bounds)
+    lambda_values = define_param_range(Lambda, lambda_step, bounds=lambda_bounds)
+
     N0, mu, Lambda = _apply_gamma_gs(
         mu_values=mu_values,
         lambda_values=lambda_values,
@@ -1244,8 +1277,15 @@ def _apply_lognormal_gs(mu_values, sigma_values, Nt, ND_obs, D, dD, V, target, t
             error_order=error_order,
         )
 
-    # Best parameter
-    best_index = np.argmin(errors)
+    # Replace inf with NaN
+    errors[~np.isfinite(errors)] = np.nan
+
+    # If all invalid, return NaN parameters
+    if np.all(np.isnan(errors)):
+        return np.array([np.nan, np.nan, np.nan])
+
+    # Otherwise, choose the best index
+    best_index = np.nanargmin(errors)
     return Nt, mu_arr[best_index].item(), sigma_arr[best_index].item()
 
 
@@ -1262,10 +1302,16 @@ def apply_lognormal_gs(
     error_order,
 ):
     """Estimate LognormalPSD model parameters using Grid Search."""
+    # Define parameters bounds
+    sigma_bounds = (0, np.inf)  # > 0
+    scale_bounds = (0.1, np.inf)  # > 0
+    # mu_bounds = (- np.inf, np.inf) # mu = np.log(scale)
+
     # Define initial set of parameters
-    mu_step = 0.5
-    sigma_step = 0.5
-    mu_values = np.arange(0.01, 20, step=mu_step)  # TODO: define realistic values
+    scale_step = 0.2
+    sigma_step = 0.2
+    scale_values = np.arange(0.1, 20, step=scale_step)
+    mu_values = np.log(scale_values)  # TODO: define realistic values
     sigma_values = np.arange(0, 20, step=sigma_step)  # TODO: define realistic values
 
     # First round of GS
@@ -1281,10 +1327,13 @@ def apply_lognormal_gs(
         transformation=transformation,
         error_order=error_order,
     )
+    if np.isnan(mu):  # if np.nan, return immediately
+        return np.array([Nt, mu, sigma])
 
     # Second round of GS
-    mu_values = np.arange(mu - mu_step * 2, mu + mu_step * 2, step=mu_step / 20)
-    sigma_values = np.arange(sigma - sigma_step * 2, sigma + sigma_step * 2, step=sigma_step / 20)
+    sigma_values = define_param_range(sigma, sigma_step, bounds=sigma_bounds)
+    scale_values = define_param_range(np.exp(mu), scale_step, bounds=scale_bounds)
+    mu_values = np.log(scale_values)
     Nt, mu, sigma = _apply_lognormal_gs(
         mu_values=mu_values,
         sigma_values=sigma_values,
@@ -1335,8 +1384,16 @@ def apply_normalized_gamma_gs(
             error_order=error_order,
         )
 
-    # Identify best parameter set
-    mu = mu_arr[np.argmin(errors)]
+    # Replace inf with NaN
+    errors[~np.isfinite(errors)] = np.nan
+
+    # If all invalid, return NaN parameters
+    if np.all(np.isnan(errors)):
+        return np.array([np.nan, np.nan, np.nan])
+
+    # Otherwise, choose the best index
+    best_index = np.nanargmin(errors)
+    mu = mu_arr[best_index]
     return np.array([Nw, mu, D50])
 
 
@@ -1346,6 +1403,12 @@ def get_exponential_parameters_gs(ds, target="ND", transformation="log", error_o
     # "transformation": "log", "identity", "sqrt",  # only for drop_number_concentration
     # "error_order": 1,     # MAE/MSE ... only for drop_number_concentration
 
+    # Compute required variables
+    ds["Nt"] = get_total_number_concentration(
+        drop_number_concentration=ds["drop_number_concentration"],
+        diameter_bin_width=ds["diameter_bin_width"],
+    )
+
     # Define kwargs
     kwargs = {
         "D": ds["diameter_bin_center"].data,
@@ -1365,7 +1428,7 @@ def get_exponential_parameters_gs(ds, target="ND", transformation="log", error_o
         # Other options
         kwargs=kwargs,
         # Settings
-        input_core_dims=[[], ["diameter_bin_center"], ["diameter_bin_center"]],
+        input_core_dims=[[], [DIAMETER_DIMENSION], [DIAMETER_DIMENSION]],
         output_core_dims=[["parameters"]],
         vectorize=True,
         dask="parallelized",
@@ -1390,6 +1453,12 @@ def get_gamma_parameters_gs(ds, target="ND", transformation="log", error_order=1
     # "transformation": "log", "identity", "sqrt",  # only for drop_number_concentration
     # "error_order": 1,     # MAE/MSE ... only for drop_number_concentration
 
+    # Compute required variables
+    ds["Nt"] = get_total_number_concentration(
+        drop_number_concentration=ds["drop_number_concentration"],
+        diameter_bin_width=ds["diameter_bin_width"],
+    )
+
     # Define kwargs
     kwargs = {
         "D": ds["diameter_bin_center"].data,
@@ -1409,7 +1478,7 @@ def get_gamma_parameters_gs(ds, target="ND", transformation="log", error_order=1
         # Other options
         kwargs=kwargs,
         # Settings
-        input_core_dims=[[], ["diameter_bin_center"], ["diameter_bin_center"]],
+        input_core_dims=[[], [DIAMETER_DIMENSION], [DIAMETER_DIMENSION]],
         output_core_dims=[["parameters"]],
         vectorize=True,
         dask="parallelized",
@@ -1434,6 +1503,12 @@ def get_lognormal_parameters_gs(ds, target="ND", transformation="log", error_ord
     # "transformation": "log", "identity", "sqrt",  # only for drop_number_concentration
     # "error_order": 1,     # MAE/MSE ... only for drop_number_concentration
 
+    # Compute required variables
+    ds["Nt"] = get_total_number_concentration(
+        drop_number_concentration=ds["drop_number_concentration"],
+        diameter_bin_width=ds["diameter_bin_width"],
+    )
+
     # Define kwargs
     kwargs = {
         "D": ds["diameter_bin_center"].data,
@@ -1453,7 +1528,7 @@ def get_lognormal_parameters_gs(ds, target="ND", transformation="log", error_ord
         # Other options
         kwargs=kwargs,
         # Settings
-        input_core_dims=[[], ["diameter_bin_center"], ["diameter_bin_center"]],
+        input_core_dims=[[], [DIAMETER_DIMENSION], [DIAMETER_DIMENSION]],
         output_core_dims=[["parameters"]],
         vectorize=True,
         dask="parallelized",
@@ -1475,8 +1550,8 @@ def get_lognormal_parameters_gs(ds, target="ND", transformation="log", error_ord
 def get_normalized_gamma_parameters_gs(ds, target="ND", transformation="log", error_order=1):
     r"""Estimate $\mu$ of a Normalized Gamma distribution using Grid Search.
 
-    The D50 and Nw parameters of the Normalized Gamma distribution are derived empirically from the observed DSD.
-    $\mu$ is derived by minimizing the errors between the observed DSD and modelled Normalized Gamma distribution.
+    The D50 and Nw parameters of the Normalized Gamma distribution are derived empirically from the obs DSD.
+    $\mu$ is derived by minimizing the errors between the obs DSD and modelled Normalized Gamma distribution.
 
     Parameters
     ----------
@@ -1501,6 +1576,29 @@ def get_normalized_gamma_parameters_gs(ds, target="ND", transformation="log", er
     # "transformation": "log", "identity", "sqrt",  # only for drop_number_concentration
     # "error_order": 1,     # MAE/MSE ... only for drop_number_concentration
 
+    # Compute required variables
+    drop_number_concentration = ds["drop_number_concentration"]
+    diameter_bin_width = ds["diameter_bin_width"]
+    diameter = ds["diameter_bin_center"] / 1000  # conversion from mm to m
+    m3 = get_moment(
+        drop_number_concentration=drop_number_concentration,
+        diameter=diameter,  # m
+        diameter_bin_width=diameter_bin_width,  # mm
+        moment=3,
+    )
+    m4 = get_moment(
+        drop_number_concentration=drop_number_concentration,
+        diameter=diameter,  # m
+        diameter_bin_width=diameter_bin_width,  # mm
+        moment=4,
+    )
+    ds["Nw"] = get_normalized_intercept_parameter_from_moments(moment_3=m3, moment_4=m4)
+    ds["D50"] = get_median_volume_drop_diameter(
+        drop_number_concentration=drop_number_concentration,
+        diameter=diameter,  # m
+        diameter_bin_width=diameter_bin_width,  # mm
+    )
+
     # Define kwargs
     kwargs = {
         "D": ds["diameter_bin_center"].data,
@@ -1521,7 +1619,7 @@ def get_normalized_gamma_parameters_gs(ds, target="ND", transformation="log", er
         # Other options
         kwargs=kwargs,
         # Settings
-        input_core_dims=[[], [], ["diameter_bin_center"], ["diameter_bin_center"]],
+        input_core_dims=[[], [], [DIAMETER_DIMENSION], [DIAMETER_DIMENSION]],
         output_core_dims=[["parameters"]],
         vectorize=True,
         dask="parallelized",
@@ -1569,6 +1667,8 @@ def get_exponential_parameters_Zhang2008(moment_l, moment_m, l, m):  # noqa: E74
         Meteor. Climatol.,
         https://doi.org/10.1175/2008JAMC1876.1
     """
+    if l == m:
+        raise ValueError("Equal l and m moment orders are not allowed.")
     num = moment_l * gamma(m + 1)
     den = moment_m * gamma(l + 1)
     Lambda = np.power(num / den, (1 / (m - l)))
@@ -1593,21 +1693,21 @@ def get_exponential_parameters_M34(moment_3, moment_4):
     return N0, Lambda
 
 
-def get_gamma_parameters_M012(M0, M1, M2):
-    """Compute gamma distribution parameters following Cao et al., 2009.
+# def get_gamma_parameters_M012(M0, M1, M2):
+#     """Compute gamma distribution parameters following Cao et al., 2009.
 
-    References
-    ----------
-    Cao, Q., and G. Zhang, 2009:
-    Errors in Estimating Raindrop Size Distribution Parameters Employing Disdrometer  and Simulated Raindrop Spectra.
-    J. Appl. Meteor. Climatol., 48, 406-425, https://doi.org/10.1175/2008JAMC2026.1.
-    """
-    # TODO: really bad results. check formula !
-    G = M1**3 / M0 / M2
-    mu = 1 / (1 - G) - 2
-    Lambda = M0 / M1 * (mu + 1)
-    N0 = Lambda ** (mu + 1) * M0 / gamma(mu + 1)
-    return N0, mu, Lambda
+#     References
+#     ----------
+#     Cao, Q., and G. Zhang, 2009:
+#     Errors in Estimating Raindrop Size Distribution Parameters Employing Disdrometer  and Simulated Raindrop Spectra.
+#     J. Appl. Meteor. Climatol., 48, 406-425, https://doi.org/10.1175/2008JAMC2026.1.
+#     """
+#     # TODO: really bad results. check formula !
+#     G = M1**3 / M0 / M2
+#     mu = 1 / (1 - G) - 2
+#     Lambda = M0 / M1 * (mu + 1)
+#     N0 = Lambda ** (mu + 1) * M0 / gamma(mu + 1)
+#     return N0, mu, Lambda
 
 
 def get_gamma_parameters_M234(M2, M3, M4):
@@ -1735,12 +1835,25 @@ def get_lognormal_parameters_M346(M3, M4, M6):
     return Nt, mu, sigma
 
 
+def _compute_moments(ds, moments):
+    list_moments = [
+        get_moment(
+            drop_number_concentration=ds["drop_number_concentration"],
+            diameter=ds["diameter_bin_center"] / 1000,  # m
+            diameter_bin_width=ds["diameter_bin_width"],  # mm
+            moment=int(moment.replace("M", "")),
+        )
+        for moment in moments
+    ]
+    return list_moments
+
+
 def _get_gamma_parameters_mom(ds: xr.Dataset, mom_method: str) -> xr.Dataset:
     # Get the correct function and list of variables for the requested method
     func, needed_moments = MOM_METHODS_DICT["GammaPSD"][mom_method]
 
-    # Extract the required arrays from the dataset
-    arrs = [ds[var_name] for var_name in needed_moments]
+    # Compute required moments
+    arrs = _compute_moments(ds, moments=needed_moments)
 
     # Apply the function. This will produce (mu, Lambda, N0) with the same coords/shapes as input data
     N0, mu, Lambda = func(*arrs)
@@ -1761,8 +1874,8 @@ def _get_lognormal_parameters_mom(ds: xr.Dataset, mom_method: str) -> xr.Dataset
     # Get the correct function and list of variables for the requested method
     func, needed_moments = MOM_METHODS_DICT["LognormalPSD"][mom_method]
 
-    # Extract the required arrays from the dataset
-    arrs = [ds[var_name] for var_name in needed_moments]
+    # Compute required moments
+    arrs = _compute_moments(ds, moments=needed_moments)
 
     # Apply the function. This will produce (mu, Lambda, N0) with the same coords/shapes as input data
     Nt, mu, sigma = func(*arrs)
@@ -1783,8 +1896,8 @@ def _get_exponential_parameters_mom(ds: xr.Dataset, mom_method: str) -> xr.Datas
     # Get the correct function and list of variables for the requested method
     func, needed_moments = MOM_METHODS_DICT["ExponentialPSD"][mom_method]
 
-    # Extract the required arrays from the dataset
-    arrs = [ds[var_name] for var_name in needed_moments]
+    # Compute required moments
+    arrs = _compute_moments(ds, moments=needed_moments)
 
     # Apply the function. This will produce (mu, Lambda, N0) with the same coords/shapes as input data
     N0, Lambda = func(*arrs)
@@ -1803,6 +1916,79 @@ def _get_exponential_parameters_mom(ds: xr.Dataset, mom_method: str) -> xr.Datas
 ####--------------------------------------------------------------------------------------.
 #### Routines dictionary
 
+####--------------------------------------------------------------------------------------.
+ATTRS_PARAMS_DICT = {
+    "GammaPSD": {
+        "N0": {
+            "description": "Intercept parameter of the Gamma PSD",
+            "standard_name": "particle_size_distribution_intercept",
+            "units": "mm**(-1-mu) m-3",
+            "long_name": "GammaPSD intercept parameter",
+        },
+        "mu": {
+            "description": "Shape parameter of the Gamma PSD",
+            "standard_name": "particle_size_distribution_shape",
+            "units": "",
+            "long_name": "GammaPSD shape parameter",
+        },
+        "Lambda": {
+            "description": "Slope (rate) parameter of the Gamma PSD",
+            "standard_name": "particle_size_distribution_slope",
+            "units": "mm-1",
+            "long_name": "GammaPSD slope parameter",
+        },
+    },
+    "NormalizedGammaPSD": {
+        "Nw": {
+            "standard_name": "normalized_intercept_parameter",
+            "units": "mm-1 m-3",
+            "long_name": "NormalizedGammaPSD Normalized Intercept Parameter",
+        },
+        "mu": {
+            "description": "Dimensionless shape parameter controlling the curvature of the Normalized Gamma PSD",
+            "standard_name": "particle_size_distribution_shape",
+            "units": "",
+            "long_name": "NormalizedGammaPSD Shape Parameter ",
+        },
+        "D50": {
+            "standard_name": "median_volume_diameter",
+            "units": "mm",
+            "long_name": "NormalizedGammaPSD Median Volume Drop Diameter",
+        },
+    },
+    "LognormalPSD": {
+        "Nt": {
+            "standard_name": "number_concentration_of_rain_drops_in_air",
+            "units": "m-3",
+            "long_name": "Total Number Concentration",
+        },
+        "mu": {
+            "description": "Mean of the Lognormal PSD",
+            "units": "log(mm)",
+            "long_name": "Mean of the Lognormal PSD",
+        },
+        "sigma": {
+            "standard_name": "Standard Deviation of the Lognormal PSD",
+            "units": "",
+            "long_name": "Standard Deviation of the Lognormal PSD",
+        },
+    },
+    "ExponentialPSD": {
+        "N0": {
+            "description": "Intercept parameter of the Exponential PSD",
+            "standard_name": "particle_size_distribution_intercept",
+            "units": "mm-1 m-3",
+            "long_name": "ExponentialPSD intercept parameter",
+        },
+        "Lambda": {
+            "description": "Slope (rate) parameter of the Exponential PSD",
+            "standard_name": "particle_size_distribution_slope",
+            "units": "mm-1",
+            "long_name": "ExponentialPSD slope parameter",
+        },
+    },
+}
+
 
 MOM_METHODS_DICT = {
     "GammaPSD": {
@@ -1843,6 +2029,8 @@ OPTIMIZATION_ROUTINES_DICT = {
 
 def available_mom_methods(psd_model):
     """Implemented MOM methods for a given PSD model."""
+    if psd_model not in MOM_METHODS_DICT:
+        raise NotImplementedError(f"No MOM methods available for {psd_model}")
     return list(MOM_METHODS_DICT[psd_model])
 
 
@@ -1863,7 +2051,7 @@ def check_psd_model(psd_model, optimization):
             f"{optimization} optimization is not available for 'psd_model' {psd_model}. "
             f"Accepted PSD models are {valid_psd_models}."
         )
-        raise ValueError(msg)
+        raise NotImplementedError(msg)
 
 
 def check_target(target):
@@ -1921,11 +2109,14 @@ def check_optimizer(optimizer):
     return optimizer
 
 
-def check_mom_methods(mom_methods, psd_model):
+def check_mom_methods(mom_methods, psd_model, allow_none=False):
     """Check valid mom_methods arguments."""
-    if isinstance(mom_methods, str):
+    if isinstance(mom_methods, (str, type(None))):
         mom_methods = [mom_methods]
+    mom_methods = [str(v) for v in mom_methods]  # None --> 'None'
     valid_mom_methods = available_mom_methods(psd_model)
+    if allow_none:
+        valid_mom_methods = [*valid_mom_methods, "None"]
     invalid_mom_methods = np.array(mom_methods)[np.isin(mom_methods, valid_mom_methods, invert=True)]
     if len(invalid_mom_methods) > 0:
         raise ValueError(
@@ -1970,36 +2161,50 @@ def check_optimization_kwargs(optimization_kwargs, optimization, psd_model):
     expected_arguments = dict_arguments.get(optimization, {})
 
     # Check for missing arguments in optimization_kwargs
-    missing_args = [arg for arg in expected_arguments if arg not in optimization_kwargs]
-    if missing_args:
-        raise ValueError(f"Missing required arguments for {optimization} optimization: {missing_args}")
+    # missing_args = [arg for arg in expected_arguments if arg not in optimization_kwargs]
+    # if missing_args:
+    #     raise ValueError(f"Missing required arguments for {optimization} optimization: {missing_args}")
 
-    # Validate argument values
-    _ = [check(optimization_kwargs[arg]) for arg, check in expected_arguments.items() if callable(check)]
+    # Validate arguments values
+    _ = [
+        check(optimization_kwargs[arg])
+        for arg, check in expected_arguments.items()
+        if callable(check) and arg in optimization_kwargs
+    ]
 
     # Further special checks
-    if optimization == "MOM":
+    if optimization == "MOM" and "mom_methods" in optimization_kwargs:
         _ = check_mom_methods(mom_methods=optimization_kwargs["mom_methods"], psd_model=psd_model)
-    if optimization == "ML" and optimization_kwargs["init_method"] is not None:
-        _ = check_mom_methods(mom_methods=optimization_kwargs["init_method"], psd_model=psd_model)
+    if optimization == "ML" and optimization_kwargs.get("init_method", None) is not None:
+        _ = check_mom_methods(mom_methods=optimization_kwargs["init_method"], psd_model=psd_model, allow_none=True)
 
 
 ####--------------------------------------------------------------------------------------.
 #### Wrappers for fitting
 
 
-def get_mom_parameters(ds: xr.Dataset, psd_model: str, mom_methods: str) -> xr.Dataset:
+def _finalize_attributes(ds_params, psd_model, optimization, optimization_kwargs):
+    ds_params.attrs["disdrodb_psd_model"] = psd_model
+    ds_params.attrs["disdrodb_psd_optimization"] = optimization
+    ds_params.attrs["disdrodb_psd_optimization_kwargs"] = ", ".join(
+        [f"{k}: {v}" for k, v in optimization_kwargs.items()],
+    )
+    return ds_params
+
+
+def get_mom_parameters(ds: xr.Dataset, psd_model: str, mom_methods=None) -> xr.Dataset:
     """
     Compute PSD model parameters using various method-of-moments (MOM) approaches.
 
-    The method is specified by the `mom_methods` acronym, e.g. 'M012', 'M234', 'M246'.
+    The method is specified by the `mom_methods` abbreviations, e.g. 'M012', 'M234', 'M246'.
 
     Parameters
     ----------
     ds : xarray.Dataset
         An xarray Dataset with the required moments M0...M6 as data variables.
-    mom_methods: str or list
-        Valid MOM methods are {'M012', 'M234', 'M246', 'M456', 'M346'}.
+    mom_methods: str or list (optional)
+        See valid values with disdrodb.psd.available_mom_methods(psd_model)
+        If None (the default), compute model parameters with all available MOM methods.
 
     Returns
     -------
@@ -2010,6 +2215,8 @@ def get_mom_parameters(ds: xr.Dataset, psd_model: str, mom_methods: str) -> xr.D
     """
     # Check inputs
     check_psd_model(psd_model=psd_model, optimization="MOM")
+    if mom_methods is None:
+        mom_methods = available_mom_methods(psd_model)
     mom_methods = check_mom_methods(mom_methods, psd_model=psd_model)
 
     # Retrieve function
@@ -2017,13 +2224,21 @@ def get_mom_parameters(ds: xr.Dataset, psd_model: str, mom_methods: str) -> xr.D
 
     # Compute parameters
     if len(mom_methods) == 1:
-        ds = func(ds=ds, mom_method=mom_methods[0])
-        ds.attrs["mom_method"] = mom_methods[0]
-        return ds
-    list_ds = [func(ds=ds, mom_method=mom_method) for mom_method in mom_methods]
-    ds = xr.concat(list_ds, dim="mom_method")
-    ds = ds.assign_coords({"mom_method": mom_methods})
-    return ds
+        ds_params = func(ds=ds, mom_method=mom_methods[0])
+    else:
+        list_ds = [func(ds=ds, mom_method=mom_method) for mom_method in mom_methods]
+        ds_params = xr.concat(list_ds, dim="mom_method")
+        ds_params = ds_params.assign_coords({"mom_method": mom_methods})
+
+    # Add model attributes
+    optimization_kwargs = {"mom_methods": mom_methods}
+    ds_params = _finalize_attributes(
+        ds_params=ds_params,
+        psd_model=psd_model,
+        optimization="MOM",
+        optimization_kwargs=optimization_kwargs,
+    )
+    return ds_params
 
 
 def get_ml_parameters(
@@ -2052,7 +2267,7 @@ def get_ml_parameters(
         The PSD model to fit. See ``available_psd_models()``.
     init_method: str or list
         The method(s) of moments used to initialize the PSD model parameters.
-        See ``available_mom_methods(psd_model)``.
+        Multiple methods can be specified. See ``available_mom_methods(psd_model)``.
     probability_method : str, optional
         Method to compute probabilities. The default value is ``cdf``.
     likelihood : str, optional
@@ -2076,21 +2291,51 @@ def get_ml_parameters(
     optimizer = check_optimizer(optimizer)
 
     # Check valid init_method
-    if init_method is not None:
-        init_method = check_mom_methods(mom_methods=init_method, psd_model=psd_model)
+    init_method = check_mom_methods(mom_methods=init_method, psd_model=psd_model, allow_none=True)
 
     # Retrieve estimation function
     func = OPTIMIZATION_ROUTINES_DICT["ML"][psd_model]
 
-    # Retrieve parameters
-    ds_params = func(
-        ds=ds,
-        init_method=init_method,
-        probability_method=probability_method,
-        likelihood=likelihood,
-        truncated_likelihood=truncated_likelihood,
-        optimizer=optimizer,
+    # Compute parameters
+    if init_method is None or len(init_method) == 1:
+        ds_params = func(
+            ds=ds,
+            init_method=init_method[0],
+            probability_method=probability_method,
+            likelihood=likelihood,
+            truncated_likelihood=truncated_likelihood,
+            optimizer=optimizer,
+        )
+    else:
+        list_ds = [
+            func(
+                ds=ds,
+                init_method=method,
+                probability_method=probability_method,
+                likelihood=likelihood,
+                truncated_likelihood=truncated_likelihood,
+                optimizer=optimizer,
+            )
+            for method in init_method
+        ]
+        ds_params = xr.concat(list_ds, dim="init_method")
+        ds_params = ds_params.assign_coords({"init_method": init_method})
+
+    # Add model attributes
+    optimization_kwargs = {
+        "init_method": init_method,
+        "probability_method": "probability_method",
+        "likelihood": likelihood,
+        "truncated_likelihood": truncated_likelihood,
+        "optimizer": optimizer,
+    }
+    ds_params = _finalize_attributes(
+        ds_params=ds_params,
+        psd_model=psd_model,
+        optimization="ML",
+        optimization_kwargs=optimization_kwargs,
     )
+
     # Return dataset with parameters
     return ds_params
 
@@ -2106,26 +2351,70 @@ def get_gs_parameters(ds, psd_model, target="ND", transformation="log", error_or
     # Check valid transformation
     transformation = check_transformation(transformation)
 
+    # Check fall velocity is available if target R
+    if "fall_velocity" not in ds:
+        ds["fall_velocity"] = get_dataset_fall_velocity(ds)
+
     # Retrieve estimation function
     func = OPTIMIZATION_ROUTINES_DICT["GS"][psd_model]
 
     # Estimate parameters
     ds_params = func(ds, target=target, transformation=transformation, error_order=error_order)
 
+    # Add model attributes
+    optimization_kwargs = {
+        "target": target,
+        "transformation": transformation,
+        "error_order": error_order,
+    }
+    ds_params = _finalize_attributes(
+        ds_params=ds_params,
+        psd_model=psd_model,
+        optimization="GS",
+        optimization_kwargs=optimization_kwargs,
+    )
     # Return dataset with parameters
     return ds_params
 
 
+def sanitize_drop_number_concentration(drop_number_concentration):
+    """Sanitize drop number concentration array.
+
+    If N(D) is all zero or contain not finite values, set everything to np.nan
+    """
+    # Condition 1: all zeros along diameter_bin_center
+    all_zero = (drop_number_concentration == 0).all(dim="diameter_bin_center")
+
+    # Condition 2: any non-finite along diameter_bin_center
+    any_nonfinite = (~np.isfinite(drop_number_concentration)).any(dim="diameter_bin_center")
+
+    # Combine conditions
+    invalid = all_zero | any_nonfinite
+
+    # Replace entire profile with NaN where invalid
+    drop_number_concentration = drop_number_concentration.where(~invalid, np.nan)
+    return drop_number_concentration
+
+
 def estimate_model_parameters(
     ds,
     psd_model,
     optimization,
-    optimization_kwargs,
+    optimization_kwargs=None,
 ):
     """Routine to estimate PSD model parameters."""
+    # Check inputs arguments
+    optimization_kwargs = {} if optimization_kwargs is None else optimization_kwargs
     optimization = check_optimization(optimization)
     check_optimization_kwargs(optimization_kwargs=optimization_kwargs, optimization=optimization, psd_model=psd_model)
 
+    # Check N(D)
+    # --> If all 0, set to np.nan
+    # --> If any is not finite --> set to np.nan
+    if "drop_number_concentration" not in ds:
+        raise ValueError("'drop_number_concentration' variable not present in input xarray.Dataset.")
+    ds["drop_number_concentration"] = sanitize_drop_number_concentration(ds["drop_number_concentration"])
+
     # Define function
     dict_func = {
         "ML": get_ml_parameters,
@@ -2137,10 +2426,7 @@ def estimate_model_parameters(
     # Retrieve parameters
     ds_params = func(ds, psd_model=psd_model, **optimization_kwargs)
 
-    # Finalize attributes
-    ds_params.attrs["disdrodb_psd_model"] = psd_model
-    ds_params.attrs["disdrodb_psd_optimization"] = optimization
-    if optimization == "GS":
-        ds_params.attrs["disdrodb_psd_optimization_target"] = optimization_kwargs["target"]
-
+    # Add parameters attributes (and units)
+    for var, attrs in ATTRS_PARAMS_DICT[psd_model].items():
+        ds_params[var].attrs = attrs
     return ds_params