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

disdrodb/utils/time.py

@@ -33,7 +33,7 @@ logger = logging.getLogger(__name__)
 #### Sampling Interval Acronyms
 
 
-def seconds_to_acronym(seconds):
+def seconds_to_temporal_resolution(seconds):
     """
     Convert a duration in seconds to a readable string format (e.g., "1H30", "1D2H").
 
@@ -57,27 +57,27 @@ def seconds_to_acronym(seconds):
         parts.append(f"{components.minutes}MIN")
     if components.seconds > 0:
         parts.append(f"{components.seconds}S")
-    acronym = "".join(parts)
-    return acronym
+    temporal_resolution = "".join(parts)
+    return temporal_resolution
 
 
-def get_resampling_information(sample_interval_acronym):
+def get_resampling_information(temporal_resolution):
     """
-    Extract resampling information from the sample interval acronym.
+    Extract resampling information from the temporal_resolution string.
 
     Parameters
     ----------
-    sample_interval_acronym: str
-      A string representing the sample interval: e.g., "1H30MIN", "ROLL1H30MIN".
+    temporal_resolution: str
+      A string representing the product temporal resolution: e.g., "1H30MIN", "ROLL1H30MIN".
 
     Returns
     -------
     sample_interval_seconds, rolling: tuple
         Sample_interval in seconds and whether rolling is enabled.
     """
-    rolling = sample_interval_acronym.startswith("ROLL")
+    rolling = temporal_resolution.startswith("ROLL")
     if rolling:
-        sample_interval_acronym = sample_interval_acronym[4:]  # Remove "ROLL"
+        temporal_resolution = temporal_resolution[4:]  # Remove "ROLL"
 
     # Allowed pattern: one or more occurrences of "<number><unit>"
     # where unit is exactly one of D, H, MIN, or S.
@@ -85,15 +85,15 @@ def get_resampling_information(sample_interval_acronym):
     pattern = r"^(\d+(?:D|H|MIN|S))+$"
 
     # Check if the entire string matches the pattern
-    if not re.match(pattern, sample_interval_acronym):
+    if not re.match(pattern, temporal_resolution):
         raise ValueError(
-            f"Invalid sample interval acronym '{sample_interval_acronym}'. "
+            f"Invalid temporal resolution '{temporal_resolution}'. "
             "Must be composed of one or more <number><unit> groups, where unit is D, H, MIN, or S.",
         )
 
     # Regular expression to match duration components and extract all (value, unit) pairs
     pattern = r"(\d+)(D|H|MIN|S)"
-    matches = re.findall(pattern, sample_interval_acronym)
+    matches = re.findall(pattern, temporal_resolution)
 
     # Conversion factors for each unit
     unit_to_seconds = {
@@ -112,21 +112,21 @@ def get_resampling_information(sample_interval_acronym):
     return sample_interval, rolling
 
 
-def acronym_to_seconds(acronym):
+def temporal_resolution_to_seconds(temporal_resolution):
     """
-    Extract the interval in seconds from the duration acronym.
+    Extract the measurement interval in seconds from the temporal resolution string.
 
     Parameters
     ----------
-    acronym: str
-      A string representing a duration: e.g., "1H30MIN", "ROLL1H30MIN".
+    temporal_resolution: str
+      A string representing the product measurement interval: e.g., "1H30MIN", "ROLL1H30MIN".
 
     Returns
     -------
     seconds
         Duration in seconds.
     """
-    seconds, _ = get_resampling_information(acronym)
+    seconds, _ = get_resampling_information(temporal_resolution)
     return seconds
 
 
@@ -262,6 +262,7 @@ def regularize_dataset(
         Regularized dataset.
 
     """
+    attrs = xr_obj.attrs.copy()
     xr_obj = _check_time_sorted(xr_obj, time_dim=time_dim)
     start_time, end_time = get_dataset_start_end_time(xr_obj, time_dim=time_dim)
 
@@ -289,11 +290,14 @@ def regularize_dataset(
         # tolerance=tolerance,  # mismatch in seconds
         fill_value=fill_value,
     )
+
+    # Ensure attributes are preserved
+    xr_obj.attrs = attrs
     return xr_obj
 
 
 ####------------------------------------------
-#### Sampling interval utilities
+#### Interval utilities
 
 
 def ensure_sample_interval_in_seconds(sample_interval):  # noqa: PLR0911
@@ -376,7 +380,7 @@ def ensure_sample_interval_in_seconds(sample_interval):  # noqa: PLR0911
             raise TypeError("Float array sample_interval must contain only whole numbers.")
         return sample_interval.astype(int)
 
-    # Deal with xarray.DataArrayy of floats that are all integer-valued (with optionally some NaN)
+    # Deal with xarray.DataArray of floats that are all integer-valued (with optionally some NaN)
     if isinstance(sample_interval, xr.DataArray) and np.issubdtype(sample_interval.dtype, np.floating):
         arr = sample_interval.copy()
         data = arr.data
@@ -397,6 +401,17 @@ def ensure_sample_interval_in_seconds(sample_interval):  # noqa: PLR0911
     )
 
 
+def ensure_timedelta_seconds_interval(interval):
+    """Return interval as numpy.timedelta64 in seconds."""
+    if isinstance(interval, (xr.DataArray, np.ndarray)):
+        return ensure_sample_interval_in_seconds(interval).astype("m8[s]")
+    return np.array(ensure_sample_interval_in_seconds(interval), dtype="m8[s]")
+
+
+####------------------------------------------
+#### Sample Interval Utilities
+
+
 def infer_sample_interval(ds, robust=False, verbose=False, logger=None):
     """Infer the sample interval of a dataset.
 
@@ -655,3 +670,134 @@ def regularize_timesteps(ds, sample_interval, robust=False, add_quality_flag=Tru
         ds = ds.isel(time=idx_valid_timesteps)
     # Return dataset
     return ds
+
+
+####---------------------------------------------------------------------------------
+#### Time blocks
+
+
+def check_freq(freq: str) -> None:
+    """Check validity of freq argument."""
+    valid_freq = ["none", "year", "season", "quarter", "month", "day", "hour"]
+    if not isinstance(freq, str):
+        raise TypeError("'freq' must be a string.")
+    if freq not in valid_freq:
+        raise ValueError(
+            f"'freq' '{freq}' is not possible. Must be one of: {valid_freq}.",
+        )
+    return freq
+
+
+def generate_time_blocks(start_time: np.datetime64, end_time: np.datetime64, freq: str) -> np.ndarray:  # noqa: PLR0911
+    """Generate time blocks between `start_time` and `end_time` for a given frequency.
+
+    Parameters
+    ----------
+    start_time : numpy.datetime64
+        Inclusive start of the overall time range.
+    end_time : numpy.datetime64
+        Inclusive end of the overall time range.
+    freq : str
+        Frequency specifier. Accepted values are:
+        - 'none'    : return a single block [start_time, end_time]
+        - 'day'     : split into daily blocks
+        - 'month'   : split into calendar months
+        - 'quarter' : split into calendar quarters
+        - 'year'    : split into calendar years
+        - 'season'  : split into meteorological seasons (MAM, JJA, SON, DJF)
+
+    Returns
+    -------
+    numpy.ndarray
+        Array of shape (n, 2) with dtype datetime64[s], where each row is [block_start, block_end].
+
+    """
+    freq = check_freq(freq)
+    if freq == "none":
+        return np.array([[start_time, end_time]], dtype="datetime64[s]")
+
+    if freq == "hour":
+        periods = pd.period_range(start=start_time, end=end_time, freq="h")
+        blocks = np.array(
+            [
+                [
+                    period.start_time.to_datetime64().astype("datetime64[s]"),
+                    period.end_time.to_datetime64().astype("datetime64[s]"),
+                ]
+                for period in periods
+            ],
+            dtype="datetime64[s]",
+        )
+        return blocks
+
+    if freq == "day":
+        periods = pd.period_range(start=start_time, end=end_time, freq="d")
+        blocks = np.array(
+            [
+                [
+                    period.start_time.to_datetime64().astype("datetime64[s]"),
+                    period.end_time.to_datetime64().astype("datetime64[s]"),
+                ]
+                for period in periods
+            ],
+            dtype="datetime64[s]",
+        )
+        return blocks
+
+    if freq == "month":
+        periods = pd.period_range(start=start_time, end=end_time, freq="M")
+        blocks = np.array(
+            [
+                [
+                    period.start_time.to_datetime64().astype("datetime64[s]"),
+                    period.end_time.to_datetime64().astype("datetime64[s]"),
+                ]
+                for period in periods
+            ],
+            dtype="datetime64[s]",
+        )
+        return blocks
+
+    if freq == "year":
+        periods = pd.period_range(start=start_time, end=end_time, freq="Y")
+        blocks = np.array(
+            [
+                [
+                    period.start_time.to_datetime64().astype("datetime64[s]"),
+                    period.end_time.to_datetime64().astype("datetime64[s]"),
+                ]
+                for period in periods
+            ],
+            dtype="datetime64[s]",
+        )
+        return blocks
+
+    if freq == "quarter":
+        periods = pd.period_range(start=start_time, end=end_time, freq="Q")
+        blocks = np.array(
+            [
+                [
+                    period.start_time.to_datetime64().astype("datetime64[s]"),
+                    period.end_time.floor("s").to_datetime64().astype("datetime64[s]"),
+                ]
+                for period in periods
+            ],
+            dtype="datetime64[s]",
+        )
+        return blocks
+
+    if freq == "season":
+        # Fiscal quarter frequency ending in Feb → seasons DJF, MAM, JJA, SON
+        periods = pd.period_range(start=start_time, end=end_time, freq="Q-FEB")
+        blocks = np.array(
+            [
+                [
+                    period.start_time.to_datetime64().astype("datetime64[s]"),
+                    period.end_time.to_datetime64().astype("datetime64[s]"),
+                ]
+                for period in periods
+            ],
+            dtype="datetime64[s]",
+        )
+        return blocks
+    raise NotImplementedError(f"Frequency '{freq}' is not implemented.")

disdrodb/utils/writer.py

@@ -22,11 +22,29 @@ import os
 
 import xarray as xr
 
-from disdrodb.utils.attrs import set_disdrodb_attrs
+from disdrodb.utils.attrs import get_attrs_dict, set_attrs, set_disdrodb_attrs
 from disdrodb.utils.directories import create_directory, remove_if_exists
+from disdrodb.utils.encoding import get_encodings_dict, set_encodings
 
 
-def write_product(ds: xr.Dataset, filepath: str, product: str, force: bool = False) -> None:
+def finalize_product(ds, product=None) -> xr.Dataset:
+    """Finalize DISDRODB product."""
+    # Add variables attributes
+    attrs_dict = get_attrs_dict()
+    ds = set_attrs(ds, attrs_dict=attrs_dict)
+
+    # Add variables encoding
+    encodings_dict = get_encodings_dict()
+    ds = set_encodings(ds, encodings_dict=encodings_dict)
+
+    # Add DISDRODB global attributes
+    # - e.g. in generate_l2_radar it inherit from input dataset !
+    if product is not None:
+        ds = set_disdrodb_attrs(ds, product=product)
+    return ds
+
+
+def write_product(ds: xr.Dataset, filepath: str, force: bool = False) -> None:
     """Save the xarray dataset into a NetCDF file.
 
     Parameters
@@ -35,8 +53,6 @@ def write_product(ds: xr.Dataset, filepath: str, product: str, force: bool = Fal
         Input xarray dataset.
     filepath : str
         Output file path.
-    product: str
-        DISDRODB product name.
     force : bool, optional
         Whether to overwrite existing data.
         If ``True``, overwrite existing data into destination directories.
@@ -50,8 +66,5 @@ def write_product(ds: xr.Dataset, filepath: str, product: str, force: bool = Fal
     # - If force=False --> Raise error
     remove_if_exists(filepath, force=force)
 
-    # Update attributes
-    ds = set_disdrodb_attrs(ds, product=product)
-
     # Write netcdf
     ds.to_netcdf(filepath, engine="netcdf4")

disdrodb/utils/xarray.py

@@ -21,6 +21,8 @@ import numpy as np
 import xarray as xr
 from xarray.core import dtypes
 
+from disdrodb.constants import DIAMETER_COORDS, VELOCITY_COORDS
+
 
 def xr_get_last_valid_idx(da_condition, dim, fill_value=None):
     """
@@ -97,6 +99,89 @@ def xr_get_last_valid_idx(da_condition, dim, fill_value=None):
     return last_idx
 
 
+####-------------------------------------------------------------------
+#### Unstacking dimension
+
+
+def _check_coord_handling(coord_handling):
+    if coord_handling not in {"keep", "drop", "unstack"}:
+        raise ValueError("coord_handling must be one of 'keep', 'drop', or 'unstack'.")
+
+
+def _unstack_coordinates(xr_obj, dim, prefix, suffix):
+    # Identify coordinates that share the target dimension
+    coords_with_dim = _get_non_dimensional_coordinates(xr_obj, dim=dim)
+    ds = xr.Dataset()
+    for coord_name in coords_with_dim:
+        coord_da = xr_obj[coord_name]
+        # Split the coordinate DataArray along the target dimension, drop coordinate and merge
+        split_ds = unstack_datarray_dimension(coord_da, coord_handling="drop", dim=dim, prefix=prefix, suffix=suffix)
+        ds.update(split_ds)
+    return ds
+
+
+def _handle_unstack_non_dim_coords(ds, source_xr_obj, coord_handling, dim, prefix, suffix):
+    # Deal with coordinates sharing the target dimension
+    if coord_handling == "keep":
+        return ds
+    if coord_handling == "unstack":
+        ds_coords = _unstack_coordinates(source_xr_obj, dim=dim, prefix=prefix, suffix=suffix)
+        ds.update(ds_coords)
+    # Remove non dimensional coordinates (unstack and drop coord_handling)
+    ds = ds.drop_vars(_get_non_dimensional_coordinates(ds, dim=dim))
+    return ds
+
+
+def _get_non_dimensional_coordinates(xr_obj, dim):
+    return [coord_name for coord_name, coord_da in xr_obj.coords.items() if dim in coord_da.dims and coord_name != dim]
+
+
+def unstack_datarray_dimension(da, dim, coord_handling="keep", prefix="", suffix=""):
+    """
+    Split a DataArray along a specified dimension into a Dataset with separate prefixed and suffixed variables.
+
+    Parameters
+    ----------
+    da : xarray.DataArray
+        The DataArray to split.
+    dim : str
+        The dimension along which to split the DataArray.
+    coord_handling : str, optional
+        Option to handle coordinates sharing the target dimension.
+        Choices are 'keep', 'drop', or 'unstack'. Defaults to 'keep'.
+    prefix : str, optional
+        String to prepend to each new variable name.
+    suffix : str, optional
+        String to append to each new variable name.
+
+    Returns
+    -------
+    xarray.Dataset
+        A Dataset with each variable split along the specified dimension.
+        The Dataset variables are named  "{prefix}{name}{suffix}{dim_value}".
+        Coordinates sharing the target dimension are handled based on `coord_handling`.
+    """
+    # Retrieve DataArray name
+    name = da.name
+    # Unstack variables
+    ds = da.to_dataset(dim=dim)
+    rename_dict = {dim_value: f"{prefix}{name}{suffix}{dim_value}" for dim_value in list(ds.data_vars)}
+    ds = ds.rename_vars(rename_dict)
+    # Deal with coordinates sharing the target dimension
+    return _handle_unstack_non_dim_coords(
+        ds=ds,
+        source_xr_obj=da,
+        coord_handling=coord_handling,
+        dim=dim,
+        prefix=prefix,
+        suffix=suffix,
+    )
+
+
+####--------------------------------------------------------------------------
+#### Fill Values Utilities
+
+
 def define_dataarray_fill_value(da):
     """Define the fill value for a numerical xarray.DataArray."""
     if np.issubdtype(da.dtype, np.floating):
@@ -163,13 +248,9 @@ def define_fill_value_dictionary(xr_obj):
 
 def remove_diameter_coordinates(xr_obj):
     """Drop diameter coordinates from xarray object."""
-    from disdrodb import DIAMETER_COORDS
-
     return xr_obj.drop_vars(DIAMETER_COORDS, errors="ignore")
 
 
 def remove_velocity_coordinates(xr_obj):
     """Drop velocity coordinates from xarray object."""
-    from disdrodb import VELOCITY_COORDS
-
     return xr_obj.drop_vars(VELOCITY_COORDS, errors="ignore")

disdrodb/viz/__init__.py

@@ -15,3 +15,16 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 # -----------------------------------------------------------------------------.
 """DISDRODB Visualization Module."""
+from disdrodb.viz.plots import (
+    compute_dense_lines,
+    max_blend_images,
+    plot_nd,
+    to_rgba,
+)
+
+__all__ = [
+    "compute_dense_lines",
+    "max_blend_images",
+    "plot_nd",
+    "to_rgba",
+]