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

disdrodb/l2/processing.py

@@ -19,12 +19,13 @@
 import numpy as np
 import xarray as xr
 
-from disdrodb.l1.encoding_attrs import get_attrs_dict, get_encoding_dict
+from disdrodb.constants import DIAMETER_DIMENSION
 from disdrodb.l1.fall_velocity import get_raindrop_fall_velocity
 from disdrodb.l1_env.routines import load_env_dataset
 from disdrodb.l2.empirical_dsd import (
+    BINS_METRICS,
+    add_bins_metrics,
     compute_integral_parameters,
-    compute_qc_bins_metrics,
     compute_spectrum_parameters,
     get_drop_average_velocity,
     get_drop_number_concentration,
@@ -35,10 +36,9 @@ from disdrodb.l2.empirical_dsd import (
 )
 from disdrodb.psd import create_psd, estimate_model_parameters
 from disdrodb.psd.fitting import compute_gof_stats
-from disdrodb.utils.attrs import set_attrs
 from disdrodb.utils.decorators import check_pytmatrix_availability
-from disdrodb.utils.encoding import set_encodings
 from disdrodb.utils.time import ensure_sample_interval_in_seconds
+from disdrodb.utils.writer import finalize_product
 
 
 def define_diameter_array(diameter_min=0, diameter_max=10, diameter_spacing=0.05):
@@ -109,26 +109,122 @@ def define_velocity_array(ds):
     return velocity
 
 
+####--------------------------------------------------------------------------
+#### Timesteps filtering functions
+
+
+def select_timesteps_with_drops(ds, minimum_ndrops=0):
+    """Select timesteps with at least the specified number of drops."""
+    valid_timesteps = ds["N"].to_numpy() >= minimum_ndrops
+    if not valid_timesteps.any().item():
+        raise ValueError(f"No timesteps with N >= {minimum_ndrops}.")
+    if "time" in ds.dims:
+        ds = ds.isel(time=valid_timesteps, drop=False)
+    return ds
+
+
+def select_timesteps_with_minimum_nbins(ds, minimum_nbins):
+    """Select timesteps with at least the specified number of diameter bins with drops."""
+    if minimum_nbins == 0:
+        return ds
+    valid_timesteps = ds["Nbins"].to_numpy() >= minimum_nbins
+    if not valid_timesteps.any().item():
+        raise ValueError(f"No timesteps with Nbins >= {minimum_nbins}.")
+    if "time" in ds.dims:
+        ds = ds.isel(time=valid_timesteps, drop=False)
+    return ds
+
+
+def select_timesteps_with_minimum_rain_rate(ds, minimum_rain_rate):
+    """Select timesteps with at least the specified rain rate."""
+    if minimum_rain_rate == 0:
+        return ds
+    # Ensure dimensionality of R is 1
+    # - Collapse velocity_method
+    dims_to_agg = set(ds["R"].dims) - {"time"}
+    da_r = ds["R"].max(dim=dims_to_agg)
+    # Determine valid timesteps
+    valid_timesteps = da_r.to_numpy() >= minimum_rain_rate
+    if not valid_timesteps.any().item():
+        raise ValueError(f"No timesteps with rain rate (R) >= {minimum_rain_rate} mm/hr.")
+    if "time" in ds.dims:
+        ds = ds.isel(time=valid_timesteps, drop=False)
+    return ds
+
+
 ####--------------------------------------------------------------------------
 #### L2 Empirical Parameters
 
 
-def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
+def _ensure_present(container, required, kind):
+    """Raise a ValueError if any of `required` are missing from the `container`."""
+    missing = [item for item in required if item not in container]
+    if missing:
+        raise ValueError(f"Dataset is missing required {kind}: {', '.join(missing)}")
+
+
+def check_l2e_input_dataset(ds):
+    """Check dataset validity for L2E production."""
+    from disdrodb.scattering import RADAR_OPTIONS
+
+    # Check minimum required variables, coordinates and dimensions are presents
+    required_variables = ["drop_number", "fall_velocity"]
+    required_coords = [
+        "diameter_bin_center",
+        "diameter_bin_width",
+        "sample_interval",
+    ]
+    required_attributes = ["sensor_name"]
+    required_dims = [DIAMETER_DIMENSION]
+    _ensure_present(list(ds.data_vars), required=required_variables, kind="variables")
+    _ensure_present(list(ds.coords), required=required_coords, kind="coords")
+    _ensure_present(list(ds.dims), required=required_dims, kind="dimensions")
+    _ensure_present(list(ds.attrs), required=required_attributes, kind="attributes")
+
+    # Remove dimensions and coordinates generated by L2E routine
+    # - This allow to recursively repass L2E product to the generate_l2e function
+    unallowed_dims = [dim for dim in ds.dims if dim in ["source", "velocity_method", *RADAR_OPTIONS]]
+    ds = ds.drop_dims(unallowed_dims)
+    unallowed_coords = [coord for coord in ds.coords if coord in ["source", "velocity_method", *RADAR_OPTIONS]]
+    ds = ds.drop_vars(unallowed_coords)
+    return ds
+
+
+def generate_l2e(
+    ds,
+    ds_env=None,
+    compute_spectra=False,
+    compute_percentage_contribution=False,
+    minimum_ndrops=1,
+    minimum_nbins=1,
+    minimum_rain_rate=0.01,
+):
     """Generate the DISDRODB L2E dataset from the DISDRODB L1 dataset.
 
     Parameters
     ----------
     ds : xarray.Dataset
         DISDRODB L1 dataset.
+        Alternatively, a xarray dataset with at least:
+
+            - variables: drop_number, fall_velocity
+            - dimension: DIAMETER_DIMENSION
+            - coordinates: diameter_bin_center, diameter_bin_width, sample_interval
+            - attributes: sensor_name
+
     ds_env : xarray.Dataset, optional
-    Environmental dataset used for fall velocity and water density estimates.
-    If None, a default environment dataset will be loaded.
+        Environmental dataset used for fall velocity and water density estimates.
+        If None, a default environment dataset will be loaded.
 
     Returns
     -------
     xarray.Dataset
-        DISRODB L2E dataset.
+        DISDRODB L2E dataset.
     """
+    # Check and prepapre input dataset
+    ds = check_l2e_input_dataset(ds)
+
+    # -------------------------------------------------------
     # Initialize L2E dataset
     ds_l2 = xr.Dataset()
 
@@ -137,24 +233,20 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
 
     # -------------------------------------------------------
     #### Preprocessing
-    # Discard all timesteps without measured drops
-    # - This allow to speed up processing
-    # - Regularization can be done at the end
-    ds = ds.isel(time=ds["N"] > 0)
+    # Select timesteps with at least the specified number of drops
+    ds = select_timesteps_with_drops(ds, minimum_ndrops=minimum_ndrops)
 
-    # Count number of diameter bins with data
-    if "Nbins" not in ds:
-        # Add bins statistics
-        ds.update(compute_qc_bins_metrics(ds))
+    # Add bins metrics to resampled data if missing
+    ds = add_bins_metrics(ds)
+
+    # Remove timesteps with not enough bins with drops
+    ds = select_timesteps_with_minimum_nbins(ds, minimum_nbins=minimum_nbins)
 
     # Retrieve ENV dataset or take defaults
     # --> Used for fall velocity and water density estimates
     if ds_env is None:
         ds_env = load_env_dataset(ds)
-
-    # TODO: Derive water density as function of ENV (temperature, ...)
-    # -->  (T == 10){density_water <- 999.7}else if(T == 20){density_water <- 998.2}else{density_water <- 995.7}
-    water_density = 1000  # kg / m3
+    water_density = ds_env.get("water_density", 1000)  # kg / m3
 
     # Determine if the velocity dimension is available
     has_velocity_dimension = "velocity_bin_center" in ds.dims
@@ -172,6 +264,7 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
 
     # Copy relevant L1 variables to L2 product
     variables = [
+        "raw_drop_number",  # 2D V x D
         "drop_number",  # 2D V x D
         "drop_counts",  # 1D D
         "sample_interval",
@@ -184,6 +277,7 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
 
     variables = [var for var in variables if var in ds]
     ds_l2.update(ds[variables])
+    ds_l2.update(ds[BINS_METRICS])
 
     # -------------------------------------------------------------------------------------------
     # Compute and add drop average velocity if an optical disdrometer (i.e OTT Parsivel or ThiesLPM)
@@ -206,7 +300,7 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
     ds_l2["drop_number_concentration"] = drop_number_concentration
 
     # -------------------------------------------------------
-    #### Compute L2 spectra
+    #### Compute R, LWC, KE and Z spectra
     if compute_spectra:
         ds_spectrum = compute_spectrum_parameters(
             drop_number_concentration,
@@ -217,6 +311,10 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
         )
         ds_l2.update(ds_spectrum)
 
+    if compute_percentage_contribution:
+        # TODO: Implement percentage contribution computation
+        pass
+
     # ----------------------------------------------------------------------------
     #### Compute L2 integral parameters from drop_number_concentration
     ds_parameters = compute_integral_parameters(
@@ -242,7 +340,7 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
 
     # -------------------------------------------------------
     #### Compute KE integral parameters directly from drop_number
-    # - The kinetic energy variables can be computed using the actual measured fall velocity by the sensor.
+    # The kinetic energy variables can be computed using the actual measured fall velocity by the sensor.
     if has_velocity_dimension:
         ds_ke = get_kinetic_energy_variables_from_drop_number(
             drop_number=drop_number,
@@ -261,24 +359,18 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
         for var in ke_vars:
             ds_parameters[var] = ds_ke[var]
 
-    # ----------------------------------------------------------------------------
-    #### Finalize L2 Dataset
     # Add DSD integral parameters
     ds_l2.update(ds_parameters)
 
-    # ----------------------------------------------------------------------------.
-    #### Add encodings and attributes
-    # Add variables attributes
-    attrs_dict = get_attrs_dict()
-    ds_l2 = set_attrs(ds_l2, attrs_dict=attrs_dict)
-
-    # Add variables encoding
-    encoding_dict = get_encoding_dict()
-    ds_l2 = set_encodings(ds_l2, encoding_dict=encoding_dict)
+    # ----------------------------------------------------------------------------
+    #### Finalize L2 Dataset
+    ds_l2 = select_timesteps_with_minimum_rain_rate(ds_l2, minimum_rain_rate=minimum_rain_rate)
 
     # Add global attributes
     ds_l2.attrs = attrs
 
+    # Add variables attributes and encodings
+    ds_l2 = finalize_product(ds_l2, product="L2E")
     return ds_l2
 
 
@@ -286,22 +378,74 @@ def generate_l2_empirical(ds, ds_env=None, compute_spectra=False):
 #### L2 Model Parameters
 
 
-def generate_l2_model(
+def _get_default_optimization(psd_model):
+    """PSD model defaults."""
+    defaults = {
+        "ExponentialPSD": "ML",
+        "GammaPSD": "ML",
+        "LognormalPSD": "ML",
+        "NormalizedGammaPSD": "GS",
+    }
+    optimization = defaults[psd_model]
+    return optimization
+
+
+def check_l2m_input_dataset(ds):
+    """Check dataset validity for L2M production."""
+    # Retrieve drop_number concentration (if not available) from drop_number
+    # --> This allow people to use directly L1 datasets to generate L2M datasets
+    if "drop_number_concentration" not in ds:
+        if "drop_number" in ds:
+            check_l2e_input_dataset(ds)
+            sample_interval = ensure_sample_interval_in_seconds(ds["sample_interval"])
+            sampling_area = get_effective_sampling_area(
+                sensor_name=ds.attrs["sensor_name"],
+                diameter=ds["diameter_bin_center"] / 1000,
+            )  # m2
+            # Compute drop number concentration (Nt) [#/m3/mm]
+            ds["drop_number_concentration"] = get_drop_number_concentration(
+                drop_number=ds["drop_number"],
+                velocity=define_velocity_array(ds),  # fall_velocity (and optionally also velocity_bin_center)
+                diameter_bin_width=ds["diameter_bin_width"],  # mm
+                sample_interval=sample_interval,
+                sampling_area=sampling_area,
+            )
+        else:
+            raise ValueError("Please provide DISDRODB L1 or L2E dataset !")
+
+    # Check minimum required variables, coordinates and dimensions are presents
+    required_variables = ["drop_number_concentration"]
+    required_coords = [
+        "diameter_bin_center",
+        "diameter_bin_width",
+        "diameter_bin_lower",
+        "diameter_bin_upper",
+        "sample_interval",
+    ]
+    required_dims = [DIAMETER_DIMENSION]
+    _ensure_present(list(ds.data_vars), required=required_variables, kind="variables")
+    _ensure_present(list(ds.coords), required=required_coords, kind="coords")
+    _ensure_present(list(ds.dims), required=required_dims, kind="dimensions")
+    return ds
+
+
+def generate_l2m(
     ds,
-    ds_env=None,
-    fall_velocity_method="Beard1976",
+    psd_model,
+    # Fitting options
+    optimization=None,
+    optimization_kwargs=None,
     # PSD discretization
     diameter_min=0,
     diameter_max=10,
     diameter_spacing=0.05,
-    # Fitting options
-    psd_model=None,
-    optimization=None,
-    optimization_kwargs=None,
+    # Processing options
+    ds_env=None,
+    fall_velocity_method="Beard1976",
     # Filtering options
-    min_nbins=4,
-    remove_timesteps_with_few_bins=False,
-    mask_timesteps_with_few_bins=False,
+    minimum_ndrops=1,
+    minimum_nbins=3,
+    minimum_rain_rate=0.01,
     # GOF metrics options
     gof_metrics=True,
 ):
@@ -316,6 +460,8 @@ def generate_l2_model(
     ----------
     ds : xarray.Dataset
         DISDRODB L2E dataset.
+    psd_model : str
+        The PSD model to fit. See ``disdrodb.psd.available_psd_models()``.
     ds_env : xarray.Dataset, optional
         Environmental dataset used for fall velocity and water density estimates.
         If None, a default environment dataset will be loaded.
@@ -325,13 +471,14 @@ def generate_l2_model(
         Maximum PSD diameter. The default value is 8 mm.
     diameter_spacing : float, optional
         PSD diameter spacing. The default value is 0.05 mm.
-    psd_model : str
-        The PSD model to fit. See ``available_psd_models()``.
     optimization : str, optional
         The fitting optimization procedure. Either "GS" (Grid Search), "ML (Maximum Likelihood)
         or "MOM" (Method of Moments).
     optimization_kwargs : dict, optional
         Dictionary with arguments to customize the fitting procedure.
+    minimum_nbins: int
+        Minimum number of bins with drops required to fit the PSD model.
+        The default value is 5.
     gof_metrics : bool, optional
         Whether to add goodness-of-fit metrics to the output dataset. The default is True.
 
@@ -340,51 +487,37 @@ def generate_l2_model(
     xarray.Dataset
         DISDRODB L2M dataset.
     """
-    # ----------------------------------------------------------------------------.
-    #### NOTES
-    # - Final processing: Optionally filter dataset only when PSD has fitted ?
-    # --> but good to have everything to compare across models
+    ####------------------------------------------------------.
+    #### Define default PSD model and optimization
+    psd_model = "NormalizedGammaPSD" if psd_model is None else psd_model
+    optimization = _get_default_optimization(psd_model) if optimization is None else optimization
 
     # ----------------------------------------------------------------------------.
+    #### Preprocessing
     # Retrieve attributes
     attrs = ds.attrs.copy()
 
-    # -------------------------------------------------------
-    # Derive water density as function of ENV (temperature, ...)
-    # TODO --> Add into ds_env !
-    # -->  (T == 10){density_water <- 999.7}else if(T == 20){density_water <- 998.2}else{density_water <- 995.7}
-    water_density = 1000  # kg / m3
+    # Check and prepare dataset
+    ds = check_l2m_input_dataset(ds)
 
-    ####------------------------------------------------------.
-    #### Preprocessing
-    # Count number of diameter bins with data
-    if "Nbins" not in ds:
-        # Add bins statistics
-        ds.update(compute_qc_bins_metrics(ds))
+    # Retrieve measurement interval
+    # - If dataset is opened with decode_timedelta=False, sample_interval is already in seconds !
+    sample_interval = ensure_sample_interval_in_seconds(ds["sample_interval"])
 
-    # Identify timesteps with enough diameter bins with counted trops
-    valid_timesteps = ds["Nbins"] >= min_nbins
+    # Select timesteps with at least the specified number of drops
+    ds = select_timesteps_with_drops(ds, minimum_ndrops=minimum_ndrops)
 
-    # Drop such timesteps if asked
-    if remove_timesteps_with_few_bins:
-        mask_timesteps_with_few_bins = False
-        ds = ds.isel(time=valid_timesteps, drop=False)
+    # Add bins metrics if missing
+    ds = add_bins_metrics(ds)
+
+    # Remove timesteps with not enough bins with drops
+    ds = select_timesteps_with_minimum_nbins(ds, minimum_nbins=minimum_nbins)
 
     # Retrieve ENV dataset or take defaults
     # --> Used for fall velocity and water density estimates
     if ds_env is None:
         ds_env = load_env_dataset(ds)
-
-    ####------------------------------------------------------.
-    #### Define default PSD optimization arguments
-    if psd_model is None and optimization is None:
-        psd_model = "NormalizedGammaPSD"
-        optimization = "GS"
-        optimization_kwargs = {
-            "target": "ND",
-            "transformation": "identity",
-            "error_order": 1,  # MAE
-        }
+    water_density = ds_env.get("water_density", 1000)  # kg / m3
 
     ####------------------------------------------------------.
     #### Retrieve PSD parameters
@@ -394,11 +527,7 @@ def generate_l2_model(
         optimization=optimization,
         optimization_kwargs=optimization_kwargs,
     )
-
-    ####------------------------------------------------------.
-    #### Mask timesteps with few bins if asked
-    if mask_timesteps_with_few_bins:
-        ds_psd_params = ds_psd_params.where(valid_timesteps)
+    psd_fitting_attrs = ds_psd_params.attrs
 
     ####-------------------------------------------------------
     #### Create PSD
@@ -415,10 +544,6 @@ def generate_l2_model(
     )
     diameter_bin_width = diameter["diameter_bin_width"]
 
-    # Retrieve time of integration
-    # - If dataset is opened with decode_timedelta=False, sample_interval is already in seconds !
-    sample_interval = ensure_sample_interval_in_seconds(ds["sample_interval"])
-
     # Retrieve drop number concentration
     drop_number_concentration = psd(diameter)
 
@@ -442,22 +567,29 @@ def generate_l2_model(
 
     # Add GOF statistics if asked
     if gof_metrics:
-        ds_gof = compute_gof_stats(drop_number_concentration=ds["drop_number_concentration"], psd=psd)
+        ds_gof = compute_gof_stats(
+            obs=ds["drop_number_concentration"],  # empirical N(D)
+            pred=psd(ds["diameter_bin_center"]),  # fitted N(D) on empirical diameter bins !
+        )
         ds_params.update(ds_gof)
 
-    #### ----------------------------------------------------------------------------.
-    #### Add encodings and attributes
-    # Add variables attributes
-    attrs_dict = get_attrs_dict()
-    ds_params = set_attrs(ds_params, attrs_dict=attrs_dict)
+    # Add empirical drop_number_concentration and fall velocity
+    # - To reuse output dataset to create another L2M dataset or to compute other GOF metrics
+    ds_params["drop_number_concentration"] = ds["drop_number_concentration"]
+    ds_params["fall_velocity"] = ds["fall_velocity"]
+    ds_params["N"] = ds["N"]
+    ds_params.update(ds[BINS_METRICS])
 
-    # Add variables encoding
-    encoding_dict = get_encoding_dict()
-    ds_params = set_encodings(ds_params, encoding_dict=encoding_dict)
+    #### ----------------------------------------------------------------------------.
+    #### Finalize dataset
+    ds_params = select_timesteps_with_minimum_rain_rate(ds_params, minimum_rain_rate=minimum_rain_rate)
 
     # Add global attributes
     ds_params.attrs = attrs
-    ds_params.attrs["disdrodb_psd_model"] = psd_name
+    ds_params.attrs.update(psd_fitting_attrs)
+
+    # Add variables attributes and encodings
+    ds_params = finalize_product(ds_params, product="L2M")
 
     # Return dataset
     return ds_params
@@ -470,10 +602,14 @@ def generate_l2_model(
 @check_pytmatrix_availability
 def generate_l2_radar(
     ds,
-    radar_band=None,
-    canting_angle_std=7,
+    frequency=None,
+    num_points=1024,
     diameter_max=10,
-    axis_ratio="Thurai2007",
+    canting_angle_std=7,
+    axis_ratio_model="Thurai2007",
+    permittivity_model="Turner2016",
+    water_temperature=10,
+    elevation_angle=0,
     parallel=True,
 ):
     """Simulate polarimetric radar variables from empirical drop number concentration or the estimated PSD.
@@ -482,15 +618,31 @@ def generate_l2_radar(
     ----------
     ds : xarray.Dataset
         Dataset containing the drop number concentration variable or the PSD parameters.
-    radar_band : str or list of str, optional
-        Radar band(s) to be used.
-        If ``None`` (the default), all available radar bands are used.
+    frequency : str, float, or list of str and float, optional
+        Frequencies in GHz for which to compute the radar parameters.
+        Alternatively, also strings can be used to specify common radar frequencies.
+        If ``None``, the common radar frequencies will be used.
+        See ``disdrodb.scattering.available_radar_bands()``.
+    num_points: int or list of integer, optional
+        Number of bins into which discretize the PSD.
+    diameter_max : float or list of float, optional
+        Maximum diameter. The default value is 10 mm.
     canting_angle_std : float or list of float, optional
         Standard deviation of the canting angle.  The default value is 7.
-    diameter_max : float or list of float, optional
-        Maximum diameter. The default value is 8 mm.
-    axis_ratio : str or list of str, optional
-        Method to compute the axis ratio. The default method is ``Thurai2007``.
+    axis_ratio_model : str or list of str, optional
+        Models to compute the axis ratio. The default model is ``Thurai2007``.
+        See available models with ``disdrodb.scattering.available_axis_ratio_models()``.
+    permittivity_model : str str or list of str, optional
+        Permittivity model to use to compute the refractive index and the
+        rayleigh_dielectric_factor. The default is ``Turner2016``.
+        See available models with ``disdrodb.scattering.available_permittivity_models()``.
+    water_temperature : float or list of float, optional
+        Water temperature in degree Celsius to be used in the permittivity model.
+        The default is 10 degC.
+    elevation_angle : float or list of float, optional
+        Radar elevation angles in degrees.
+        Specify 90 degrees for vertically pointing radars.
+        The default is 0 degrees.
     parallel : bool, optional
         Whether to compute radar variables in parallel.
         The default value is ``True``.
@@ -507,22 +659,21 @@ def generate_l2_radar(
     # Retrieve radar variables from L2E drop number concentration or from estimated L2M PSD model
     ds_radar = get_radar_parameters(
         ds=ds,
-        radar_band=radar_band,
-        canting_angle_std=canting_angle_std,
+        frequency=frequency,
+        num_points=num_points,
         diameter_max=diameter_max,
-        axis_ratio=axis_ratio,
+        canting_angle_std=canting_angle_std,
+        axis_ratio_model=axis_ratio_model,
+        permittivity_model=permittivity_model,
+        water_temperature=water_temperature,
+        elevation_angle=elevation_angle,
         parallel=parallel,
     )
 
     #### ----------------------------------------------------------------------------.
-    #### Add encodings and attributes
-    # Add variables attributes
-    attrs_dict = get_attrs_dict()
-    ds_radar = set_attrs(ds_radar, attrs_dict=attrs_dict)
-
-    # Add variables encoding
-    encoding_dict = get_encoding_dict()
-    ds_radar = set_encodings(ds_radar, encoding_dict=encoding_dict)
+    #### Finalize dataset
+    # Add variables attributes and encodings
+    ds_radar = finalize_product(ds_radar)
 
     # Return dataset
     return ds_radar