# SPDX-License-Identifier: MIT
# SPDX-FileCopyrightText: Copyright (c) 2026 TU Wien
# SPDX-FileContributor: For a full list of authors, see the AUTHORS file.
import datetime
import tempfile
from pathlib import Path
import numpy as np
import pandas as pd
import xarray as xr
from flox.xarray import xarray_reduce
from dask.array import unique as da_unique
from ascat.swath import SwathGridFiles
from ascat.product_info import get_swath_product_id
from ascat.utils import dtype_to_nan
[docs]
class TemporalSwathAggregator:
"""Class to aggregate ASCAT data its location ids over time."""
def __init__(
self,
filepath,
start_dt,
end_dt,
t_delta,
agg,
snow_cover_mask=80,
frozen_soil_mask=80,
subsurface_scattering_mask=5,
ssm_sensitivity_mask=1,
no_masking=False,
sat=None,
):
"""Initialize the class.
Parameters
----------
filepath : str
Path to the data.
start_dt : str
Start date and time (formatted e.g. 2020-01-01T00:00:00).
end_dt : str
End date and time (formatted e.g. 2020-02-01T00:00:00).
t_delta : str
Time period for aggregation (e.g. 1D, 1W, 1M, 1Y, 3M, 4Y, etc.).
agg : str
Aggregation method (e.g. mean, median, std).
snow_cover_mask : int, optional
Snow cover probability value above which to mask the source data.
frozen_soil_mask : int, optional
Frozen soil probability value above which to mask the source data.
subsurface_scattering_mask : int, optional
Subsurface scattering probability value above which to mask
the source data.
ssm_sensitivity_mask : float, optional
Soil moisture sensitivity value above which to mask
the source data.
no_masking : boolean, optional
Ignore all masks (default: False).
sat: str, optional
Regex pattern indicating which METOP satellite(s) to use, e.g. A, [BC], or '*'.
`None` will use the product's default satellite (usually all of them).
(default: None)
"""
self.filepath = filepath
fmt = "%Y-%m-%dT%H:%M:%S"
self.start_dt = datetime.datetime.strptime(start_dt, fmt)
self.end_dt = datetime.datetime.strptime(end_dt, fmt)
self.timedelta = pd.Timedelta(t_delta)
self.no_masking = no_masking
agg_methods = [
"mean", "median", "mode", "std", "min", "max", "argmin", "argmax",
"quantile", "first", "last"
]
if agg in agg_methods:
agg = "nan" + agg
self.agg = agg
# assumes ONLY swath files are in the folder
first_fname = str(next(Path(filepath).rglob("*.nc")).name)
product = get_swath_product_id(first_fname)
self.product = product
if sat is not None:
self.fmt_kwargs = {"sat": sat}
else:
self.fmt_kwargs = {}
self.collection = SwathGridFiles.from_product_id(
Path(filepath), product)
self.grid = self.collection.grid
self.agg_vars = {
"surface_soil_moisture": {
"dtype": np.dtype("int16"),
"scale_factor": 1e-2,
},
"backscatter40": {
"dtype": np.dtype("int32"),
"scale_factor": 1e-7,
},
}
self.mask_probs = {
"snow_cover_probability": snow_cover_mask,
"frozen_soil_probability": frozen_soil_mask,
"subsurface_scattering_probability": subsurface_scattering_mask,
"surface_soil_moisture_sensitivity": ssm_sensitivity_mask,
}
def _set_metadata(self, ds):
"""Add appropriate metadata to datasets."""
return ds
def _create_output_encoding(self):
"""Create NetCDF encoding."""
output_encoding = {
"latitude": {
"dtype": np.dtype("int32"),
"scale_factor": 1e-6,
"zlib": True,
"complevel": 4,
"_FillValue": dtype_to_nan[np.dtype("int32")],
"missing_value": dtype_to_nan[np.dtype("int32")],
},
"longitude": {
"dtype": np.dtype("int32"),
"scale_factor": 1e-6,
"zlib": True,
"complevel": 4,
"_FillValue": dtype_to_nan[np.dtype("int32")],
"missing_value": dtype_to_nan[np.dtype("int32")],
},
"time": {
"dtype": np.dtype("float64"),
"zlib": True,
"complevel": 4,
"_FillValue": 0,
"missing_value": 0,
},
}
for var in self.agg_vars:
if var in output_encoding:
continue
output_encoding[var] = {
"dtype": self.agg_vars[var]["dtype"],
"scale_factor": self.agg_vars[var]["scale_factor"],
"zlib": True,
"complevel": 4,
"_FillValue": dtype_to_nan[self.agg_vars[var]["dtype"]],
"missing_value": dtype_to_nan[self.agg_vars[var]["dtype"]],
}
return output_encoding
[docs]
def write_time_steps(self, outpath):
"""
Loop through time steps and write them to file.
Parameters
----------
outpath : str
Output path.
"""
product_id = self.product.lower().replace("_", "-")
grid_sampling = str(self.grid.res) + "km"
if self.agg is not None:
datasets = self.get_aggregated_time_steps()
agg_str = f"_{self.agg}"
else:
datasets = self.get_time_steps()
agg_str = "_data"
fmt = "%Y%m%d%H%M%S"
paths = []
output_encoding = self._create_output_encoding()
print("saving datasets...", end="\r")
for ds in datasets:
step_start_str = (
np.datetime64(ds.attrs["start_time"]).astype(
datetime.datetime).strftime(fmt))
step_end_str = (
np.datetime64(ds.attrs["end_time"]).astype(
datetime.datetime).strftime(fmt))
out_name = (f"ascat"
f"_{product_id}"
f"_{grid_sampling}"
f"{agg_str}"
f"_{step_start_str}"
f"_{step_end_str}.nc")
print("saving output")
ds.to_netcdf(
Path(outpath) / out_name,
encoding=output_encoding,
)
paths.append(Path(outpath) / out_name)
return paths
[docs]
def get_time_steps(self):
"""
Loop through time steps of the range, return the merged data
for each unmodified.
"""
time_steps = pd.date_range(
start=self.start_dt, end=self.end_dt, freq=self.timedelta)
for timestep in time_steps:
print("reading data for time step:", timestep, end="\r")
step_start = timestep
step_end = timestep + self.timedelta
ds_step = self.collection.read((step_start, step_end),
**self.fmt_kwargs)
step_end = step_end - pd.Timedelta("1s")
ds_step.attrs["start_time"] = np.datetime64(step_start).astype(
str)
ds_step.attrs["end_time"] = np.datetime64(step_end).astype(str)
yield ds_step
[docs]
def get_aggregated_time_steps(self):
"""Loop through data in time steps, aggregating it over time."""
for ds in self.get_time_steps():
present_agg_vars = [
var for var in self.agg_vars if var in ds.variables
]
print("masking data...", end="\r")
global_mask = (ds.surface_flag != 0)
ds = ds.where(~global_mask, drop=False)
if not self.no_masking:
variable_masks = {
"surface_soil_moisture": (
(ds["frozen_soil_probability"]
> self.mask_probs["frozen_soil_probability"])
| (ds["snow_cover_probability"]
> self.mask_probs["snow_cover_probability"])
| (ds["subsurface_scattering_probability"]
> self.mask_probs["subsurface_scattering_probability"])
| (ds["surface_soil_moisture_sensitivity"]
< self.mask_probs["surface_soil_moisture_sensitivity"])),
}
for var, var_mask in variable_masks.items():
ds[var] = ds[var].where(~var_mask, drop=False)
print("grouping data... ")
expected_location_ids = da_unique(ds["location_id"].data).compute()
# remove NaN from the expected location ids (this was introduced by the masking)
expected_location_ids = expected_location_ids[
~np.isnan(expected_location_ids)]
# grouped_ds the data by time_steps and location_id and aggregate it
grouped_ds = xarray_reduce(
ds[present_agg_vars],
ds["location_id"],
expected_groups=(expected_location_ids,),
func=self.agg)
# convert the location_id back to an integer
grouped_ds["location_id"] = grouped_ds["location_id"].astype(int)
step_start = ds.attrs["start_time"]
grouped_ds["time"] = np.datetime64(step_start, "ns")
lons, lats = self.grid.gpi2lonlat(grouped_ds.location_id.values)
grouped_ds["longitude"] = ("location_id", lons)
grouped_ds["latitude"] = ("location_id", lats)
grouped_ds = grouped_ds.set_coords(["longitude", "latitude"])
grouped_ds = self._set_metadata(grouped_ds)
yield grouped_ds