disdrodb 0.0.21__py3-none-any.whl → 0.1.1__py3-none-any.whl

disdrodb/l0/l0b_processing.py

@@ -32,23 +32,25 @@ from disdrodb.l0.check_standards import (
 from disdrodb.l0.standards import (
     # get_valid_coordinates_names,
     get_bin_coords_dict,
-    get_coords_attrs_dict,
     get_data_range_dict,
     get_dims_size_dict,
     get_l0b_cf_attrs_dict,
     get_l0b_encodings_dict,
     get_raw_array_dims_order,
     get_raw_array_nvalues,
-    get_time_encoding,
+)
+from disdrodb.utils.attrs import (
+    set_coordinate_attributes,
     set_disdrodb_attrs,
 )
 from disdrodb.utils.directories import create_directory, remove_if_exists
+from disdrodb.utils.encoding import set_encodings
 from disdrodb.utils.logger import (
     # log_warning,
     # log_debug,
-    log_error,
     log_info,
 )
+from disdrodb.utils.time import ensure_sorted_by_time
 
 logger = logging.getLogger(__name__)
 
@@ -78,10 +80,7 @@ def infer_split_str(string: str) -> str:
         idx_delimiter = np.argmax(counts)
         # If don't find the delimiter, set to None
         # --> The array will not be split, and then raise an error later on
-        if counts[idx_delimiter] == 0:
-            split_str = None
-        else:
-            split_str = valid_delims[idx_delimiter]
+        split_str = None if counts[idx_delimiter] == 0 else valid_delims[idx_delimiter]
     else:
         split_str = None  # ''.split(None) output []
     return split_str
@@ -114,7 +113,6 @@ def _format_string_array(string: str, n_values: int) -> np.array:
     np.array
         array of float
     """
-
     split_str = infer_split_str(string)
     values = np.array(string.strip(split_str).split(split_str))
 
@@ -160,17 +158,20 @@ def _reshape_raw_spectrum(
     dims_order : list
         The order of dimension in the raw spectrum.
 
-        Examples:
-        - OTT Parsivel spectrum [v1d1 ... v1d32, v2d1, ..., v2d32]
+    Examples
+    --------
+        - OTT PARSIVEL spectrum [v1d1 ... v1d32, v2d1, ..., v2d32]
         --> dims_order = ["diameter_bin_center", "velocity_bin_center"]
         - Thies LPM spectrum [v1d1 ... v20d1, v1d2, ..., v20d2]
         --> dims_order = ["velocity_bin_center", "diameter_bin_center"]
     dims_size_dict : dict
         Dictionary with the number of bins for each dimension.
-        For OTT_Parsivel:
+        For PARSIVEL and PARSIVEL2:
         {"diameter_bin_center": 32, "velocity_bin_center": 32}
-        For This_LPM
+        For LPM
         {"diameter_bin_center": 22, "velocity_bin_center": 20}
+        For PWS100
+        {"diameter_bin_center": 34, "velocity_bin_center": 34}
     n_timesteps : int
         Number of timesteps.
 
@@ -185,14 +186,13 @@ def _reshape_raw_spectrum(
         Impossible to reshape the raw_spectrum matrix
     """
     # Define output dimensions
-    dims = ["time"] + dims_order
+    dims = ["time", *dims_order]
     # Retrieve reshaping dimensions as function of dimension order
     reshape_dims = [n_timesteps] + [dims_size_dict[dim] for dim in dims_order]
     try:
         arr = arr.reshape(reshape_dims)
     except Exception as e:
         msg = f"Impossible to reshape the raw_spectrum matrix. The error is: \n {e}"
-        log_error(logger=logger, msg=msg, verbose=False)
         raise ValueError(msg)
     return arr, dims
 
@@ -200,13 +200,14 @@ def _reshape_raw_spectrum(
 def retrieve_l0b_arrays(
     df: pd.DataFrame,
     sensor_name: str,
+    logger=None,
     verbose: bool = False,
 ) -> dict:
     """Retrieves the L0B data matrix.
 
     Parameters
     ----------
-    df : pd.DataFrame
+    df : pandas.DataFrame
         Input dataframe
     sensor_name : str
         Name of the sensor
@@ -217,12 +218,11 @@ def retrieve_l0b_arrays(
         Dictionary with data arrays.
 
     """
-
-    msg = " - Retrieval of L0B data arrays started."
+    msg = "Retrieval of L0B data arrays started."
     log_info(logger=logger, msg=msg, verbose=verbose)
     # ----------------------------------------------------------.
     # Check L0 raw field availability
-    _check_raw_fields_available(df=df, sensor_name=sensor_name)
+    _check_raw_fields_available(df=df, sensor_name=sensor_name, logger=logger, verbose=verbose)
 
     # Retrieve the number of values expected for each array
     n_values_dict = get_raw_array_nvalues(sensor_name=sensor_name)
@@ -258,8 +258,8 @@ def retrieve_l0b_arrays(
 
         # For key='raw_drop_number', if 2D spectrum, reshape to 2D matrix
         # Example:
-        # - This applies i.e for OTT_Parsivel* and Thies_LPM
-        # - This does not apply to RD_80
+        # - This applies i.e for PARSIVEL*, LPM, PWS100
+        # - This does not apply to RD80
         if key == "raw_drop_number" and len(dims_order) == 2:
             arr, dims = _reshape_raw_spectrum(
                 arr=arr,
@@ -269,14 +269,14 @@ def retrieve_l0b_arrays(
             )
         else:
             # Otherwise just define the dimensions of the array
-            dims = ["time"] + dims_order
+            dims = ["time", *dims_order]
 
         # Define dictionary to pass to xr.Dataset
         dict_data[key] = (dims, arr)
 
     # -------------------------------------------------------------------------.
     # Log
-    msg = " - Retrieval of L0B data arrays ended."
+    msg = "Retrieval of L0B data arrays ended."
     log_info(logger=logger, msg=msg, verbose=verbose)
     # Return
     return dict_data
@@ -287,16 +287,16 @@ def retrieve_l0b_arrays(
 
 
 def _convert_object_variables_to_string(ds: xr.Dataset) -> xr.Dataset:
-    """Convert variables with object dtype to string.
+    """Convert variables with ``object`` dtype to ``string``.
 
     Parameters
     ----------
-    ds : xr.Dataset
+    ds  : xarray.Dataset
         Input dataset.
 
     Returns
     -------
-    xr.Dataset
+    xarray.Dataset
         Output dataset.
     """
     for var in ds.data_vars:
@@ -306,11 +306,11 @@ def _convert_object_variables_to_string(ds: xr.Dataset) -> xr.Dataset:
 
 
 def _set_variable_attributes(ds: xr.Dataset, sensor_name: str) -> xr.Dataset:
-    """Set attributes to each xr.Dataset variable.
+    """Set attributes to each ``xr.Dataset`` variable.
 
     Parameters
     ----------
-    ds : xr.Dataset
+    ds  : xarray.Dataset
         Input dataset.
     sensor_name : str
         Name of the sensor.
@@ -333,35 +333,20 @@ def _set_variable_attributes(ds: xr.Dataset, sensor_name: str) -> xr.Dataset:
     return ds
 
 
-def _set_attrs_dict(ds, attrs_dict):
-    for var in attrs_dict.keys():
-        if var in ds:
-            ds[var].attrs.update(attrs_dict[var])
-    return ds
-
-
-def _set_coordinate_attributes(ds):
-    # Get attributes dictionary
-    attrs_dict = get_coords_attrs_dict(ds)
-    # Set attributes
-    ds = _set_attrs_dict(ds, attrs_dict)
-    return ds
-
-
 def _set_dataset_attrs(ds, sensor_name):
     """Set variable and coordinates attributes."""
     # - Add netCDF variable attributes
     # --> Attributes: long_name, units, descriptions, valid_min, valid_max
     ds = _set_variable_attributes(ds=ds, sensor_name=sensor_name)
     # - Add netCDF coordinate attributes
-    ds = _set_coordinate_attributes(ds=ds)
+    ds = set_coordinate_attributes(ds=ds)
     #  - Set DISDRODB global attributes
     ds = set_disdrodb_attrs(ds=ds, product="L0B")
     return ds
 
 
 def add_dataset_crs_coords(ds):
-    "Add the CRS coordinate to the xr.Dataset"
+    """Add the CRS coordinate to the xr.Dataset."""
     # TODO: define CF-compliant CRS !
     # - CF compliant
     # - wkt
@@ -377,75 +362,54 @@ def add_dataset_crs_coords(ds):
 #### L0B Raw DataFrame Preprocessing
 
 
-def _define_dataset_variables(df, sensor_name, verbose):
+def _define_dataset_variables(df, sensor_name, logger=None, verbose=False):
     """Define DISDRODB L0B netCDF variables."""
     # Preprocess raw_spectrum, diameter and velocity arrays if available
     raw_fields = ["raw_drop_concentration", "raw_drop_average_velocity", "raw_drop_number"]
     if np.any(np.isin(raw_fields, df.columns)):
         # Retrieve dictionary of raw data matrices for xarray Dataset
-        data_vars = retrieve_l0b_arrays(df, sensor_name, verbose=verbose)
+        data_vars = retrieve_l0b_arrays(df, sensor_name=sensor_name, logger=logger, verbose=verbose)
     else:
         raise ValueError("No raw fields available.")
 
     # Define other disdrometer 'auxiliary' variables varying over time dimension
+    # - Includes time
+    # - Includes longitude and latitude for moving sensors
     valid_core_fields = [
         "raw_drop_concentration",
         "raw_drop_average_velocity",
         "raw_drop_number",
-        "time",
-        # longitude and latitude too for moving sensors
     ]
     aux_columns = df.columns[np.isin(df.columns, valid_core_fields, invert=True)]
-    aux_data_vars = {column: (["time"], df[column].values) for column in aux_columns}
+    aux_data_vars = {column: (["time"], df[column].to_numpy()) for column in aux_columns}
     data_vars.update(aux_data_vars)
-
-    # Add key "time"
-    # - Is dropped in _define_coordinates !
-    data_vars["time"] = df["time"].values
-
     return data_vars
 
 
-def _define_coordinates(data_vars, attrs, sensor_name):
-    """Define DISDRODB L0B netCDF coordinates."""
-    # Note: attrs and data_vars are modified in place !
-
-    # - Diameter and velocity
-    coords = get_bin_coords_dict(sensor_name=sensor_name)
-
-    # - Geolocation + Time
-    geolocation_vars = ["time", "latitude", "longitude", "altitude"]
-    for var in geolocation_vars:
-        if var in data_vars:
-            coords[var] = data_vars[var]
-            _ = data_vars.pop(var)
-            _ = attrs.pop(var, None)
-        else:
-            coords[var] = attrs[var]
-            _ = attrs.pop(var)
-    return coords
-
-
 def create_l0b_from_l0a(
     df: pd.DataFrame,
-    attrs: dict,
+    metadata: dict,
+    logger=None,
     verbose: bool = False,
 ) -> xr.Dataset:
     """Transform the L0A dataframe to the L0B xr.Dataset.
 
     Parameters
     ----------
-    df : pd.DataFrame
+    df : pandas.DataFrame
         DISDRODB L0A dataframe.
-    attrs : dict
-        Station metadata.
+        The raw drop number spectrum is reshaped to a 2D(+time) array.
+        The raw drop concentration and velocity are reshaped to 1D(+time) arrays.
+    metadata : dict
+        DISDRODB station metadata.
+        To use this function outside the DISDRODB routines, the dictionary must
+        contain the fields: ``sensor_name``, ``latitude``, ``longitude``, ``altitude``, ``platform_type``.
     verbose : bool, optional
-        Whether to verbose the processing.
-        The default is False.
+        Whether to verbose the processing. The default value is ``False``.
 
     Returns
     -------
-    xr.Dataset
+    xarray.Dataset
         DISDRODB L0B dataset.
 
     Raises
@@ -454,27 +418,15 @@ def create_l0b_from_l0a(
         Error if the DISDRODB L0B xarray dataset can not be created.
     """
     # Retrieve sensor name
-    attrs = attrs.copy()
-    sensor_name = attrs["sensor_name"]
-    # -----------------------------------------------------------.
-    # Define Dataset variables and coordinates
-    data_vars = _define_dataset_variables(df, sensor_name=sensor_name, verbose=verbose)
+    metadata = metadata.copy()
+    sensor_name = metadata["sensor_name"]
 
-    # -----------------------------------------------------------.
-    # Define coordinates for xarray Dataset
-    # - attrs and data_vars are modified in place !
-    coords = _define_coordinates(data_vars, attrs=attrs, sensor_name=sensor_name)
+    # Define Dataset variables and coordinates
+    data_vars = _define_dataset_variables(df, sensor_name=sensor_name, logger=logger, verbose=verbose)
 
-    # -----------------------------------------------------------
     # Create xarray Dataset
-    ds = xr.Dataset(
-        data_vars=data_vars,
-        coords=coords,
-        attrs=attrs,
-    )
-    ds = finalize_dataset(ds, sensor_name=sensor_name)
-
-    # -----------------------------------------------------------
+    ds = xr.Dataset(data_vars=data_vars)
+    ds = finalize_dataset(ds, sensor_name=sensor_name, metadata=metadata)
     return ds
 
 
@@ -482,8 +434,43 @@ def create_l0b_from_l0a(
 #### L0B netCDF4 Writer
 
 
-def finalize_dataset(ds, sensor_name):
+def set_geolocation_coordinates(ds, metadata):
+    """Add geolocation coordinates to dataset."""
+    # Assumption
+    # - If coordinate is present in L0A, overrides the one specified in the attributes
+    # - If a station is fixed, discard the coordinates in the DISDRODB reader !
+
+    # Assign geolocation coordinates to dataset
+    coords = ["latitude", "longitude", "altitude"]
+    for coord in coords:
+        # If coordinate not present, add it from dictionary
+        if coord not in ds:
+            ds = ds.assign_coords({coord: metadata.pop(coord, np.nan)})
+        # Else if set coordinates the variable in the dataset (present in the raw data)
+        else:
+            ds = ds.set_coords(coord)
+            _ = metadata.pop(coord, None)
+
+    # Set -9999 flag value to np.nan
+    for coord in coords:
+        ds[coord] = xr.where(ds[coord] == -9999, np.nan, ds[coord])
+
+    # Set attributes without geolocation coordinates
+    ds.attrs = metadata
+    return ds
+
+
+def finalize_dataset(ds, sensor_name, metadata):
     """Finalize DISDRODB L0B Dataset."""
+    # Ensure sorted by time
+    ds = ensure_sorted_by_time(ds)
+
+    # Set diameter and velocity bin coordinates
+    ds = ds.assign_coords(get_bin_coords_dict(sensor_name=sensor_name))
+
+    # Set geolocation coordinates and attributes
+    ds = set_geolocation_coordinates(ds, metadata=metadata)
+
     # Add dataset CRS coordinate
     ds = add_dataset_crs_coords(ds)
 
@@ -501,88 +488,23 @@ def finalize_dataset(ds, sensor_name):
     return ds
 
 
-def sanitize_encodings_dict(encoding_dict: dict, ds: xr.Dataset) -> dict:
-    """Ensure chunk size to be smaller than the array shape.
+def set_l0b_encodings(ds: xr.Dataset, sensor_name: str):
+    """Apply the L0B encodings to the xarray Dataset.
 
     Parameters
     ----------
-    encoding_dict : dict
-        Dictionary containing the encoding to write DISDRODB L0B netCDFs.
-    ds : xr.Dataset
-        Input dataset.
-
-    Returns
-    -------
-    dict
-        Encoding dictionary.
-    """
-    for var in ds.data_vars:
-        shape = ds[var].shape
-        chunks = encoding_dict[var]["chunksizes"]
-        if chunks is not None:
-            chunks = [shape[i] if chunks[i] > shape[i] else chunks[i] for i in range(len(chunks))]
-            encoding_dict[var]["chunksizes"] = chunks
-    return encoding_dict
-
-
-def rechunk_dataset(ds: xr.Dataset, encoding_dict: dict) -> xr.Dataset:
-    """Coerce the dataset arrays to have the chunk size specified in the encoding dictionary.
-
-    Parameters
-    ----------
-    ds : xr.Dataset
-        Input xarray dataset
-    encoding_dict : dict
-        Dictionary containing the encoding to write the xarray dataset as a netCDF.
-
-    Returns
-    -------
-    xr.Dataset
-        Output xarray dataset
-    """
-
-    for var in ds.data_vars:
-        chunks = encoding_dict[var].pop("chunksizes")
-        dims = list(ds[var].dims)
-        chunks_dict = dict(zip(dims, chunks))
-        if chunks is not None:
-            ds[var] = ds[var].chunk(chunks_dict)
-    return ds
-
-
-def set_encodings(ds: xr.Dataset, sensor_name: str) -> xr.Dataset:
-    """Apply the encodings to the xarray Dataset.
-
-    Parameters
-    ----------
-    ds : xr.Dataset
+    ds  : xarray.Dataset
         Input xarray dataset.
     sensor_name : str
         Name of the sensor.
 
     Returns
     -------
-    xr.Dataset
+    xarray.Dataset
         Output xarray dataset.
     """
-    # Get encoding dictionary
     encoding_dict = get_l0b_encodings_dict(sensor_name)
-    encoding_dict = {k: encoding_dict[k] for k in ds.data_vars}
-
-    # Ensure chunksize smaller than the array shape
-    encoding_dict = sanitize_encodings_dict(encoding_dict, ds)
-
-    # Rechunk variables for fast writing !
-    # - This pop the chunksize argument from the encoding dict !
-    ds = rechunk_dataset(ds, encoding_dict)
-
-    # Set time encoding
-    ds["time"].encoding.update(get_time_encoding())
-
-    # Set the variable encodings
-    for var in ds.data_vars:
-        ds[var].encoding.update(encoding_dict[var])
-
+    ds = set_encodings(ds=ds, encoding_dict=encoding_dict)
     return ds
 
 
@@ -591,7 +513,7 @@ def write_l0b(ds: xr.Dataset, filepath: str, force=False) -> None:
 
     Parameters
     ----------
-    ds : xr.Dataset
+    ds  : xarray.Dataset
         Input xarray dataset.
     filepath : str
         Output file path.
@@ -599,8 +521,8 @@ def write_l0b(ds: xr.Dataset, filepath: str, force=False) -> None:
         Name of the sensor.
     force : bool, optional
         Whether to overwrite existing data.
-        If True, overwrite existing data into destination directories.
-        If False, raise an error if there are already data into destination directories. This is the default.
+        If ``True``, overwrite existing data into destination directories.
+        If ``False``, raise an error if there are already data into destination directories. This is the default.
     """
     # Create station directory if does not exist
     create_directory(os.path.dirname(filepath))
@@ -614,7 +536,7 @@ def write_l0b(ds: xr.Dataset, filepath: str, force=False) -> None:
     sensor_name = ds.attrs.get("sensor_name")
 
     # Set encodings
-    ds = set_encodings(ds=ds, sensor_name=sensor_name)
+    ds = set_l0b_encodings(ds=ds, sensor_name=sensor_name)
 
     # Write netcdf
     ds.to_netcdf(filepath, engine="netcdf4")