aggfly 0.1.0__py3-none-any.whl

aggfly/__init__.py

@@ -0,0 +1,32 @@
+from .aggregate import (
+    TemporalAggregator,
+    SpatialAggregator,
+    aggregate_dataset,
+    aggregate_time,
+    aggregate_space,
+    distributed_client,
+    is_distributed,
+    start_dask_client,
+    shutdown_dask_client
+)
+from .dataset import Dataset, Grid, dataset_from_path
+from .weights import (
+    CropWeights, 
+    PopWeights, 
+    GridWeights, 
+    SecondaryWeights,
+    weights_from_objects, 
+    pop_weights_from_path, 
+    crop_weights_from_path,
+    secondary_weights_from_path
+)
+from .regions import GeoRegions, georegions_from_path, georegions_from_name
+from .tests import (
+    georegion,
+    dataset_360,
+    secondary_weights,
+    weights,
+    test_weights,
+    test_aggregate_time,
+    test_aggregate
+)

aggfly/aggregate/__init__.py

@@ -0,0 +1,5 @@
+from .aggregate_utils import *
+from .aggregate import aggregate_dataset, aggregate_time, aggregate_space
+from .temporal import TemporalAggregator
+from .spatial import SpatialAggregator
+

aggfly/aggregate/aggregate.py

@@ -0,0 +1,228 @@
+"""
+This module contains functions for aggregating datasets over time and space.
+
+Functions:
+- transform_dataset(dataset: Dataset, key: str, **kwargs: Union[int, List[int]]) -> Dict[str, Dataset]:
+    Transform a given dataset by raising it to a power or interacting with another dataset.
+
+- aggregate_time(dataset: Dataset, aggregator_dict: Dict[str, Union[List[Tuple], TemporalAggregator]] = None, weights: GridWeights = None, **kwargs) -> Dict[str, Dataset]:
+    Aggregate a dataset over time using the specified temporal aggregators.
+
+- aggregate_space(dataset_dict: Dict[str, Dataset], weights: GridWeights) -> pd.DataFrame:
+    Aggregate a dictionary of datasets over space.
+
+- aggregate_dataset(dataset: Dataset, weights: GridWeights, aggregator_dict: Dict[str, Union[List[Tuple], TemporalAggregator]] = None, **kwargs) -> pd.DataFrame:
+    Aggregate a dataset over time and space.
+"""
+
+from typing import List, Dict, Union, Tuple
+import pandas as pd
+import dask
+from dask.distributed import LocalCluster, Client, progress
+
+from .temporal import TemporalAggregator
+from .spatial import SpatialAggregator
+from ..dataset import Dataset
+from ..weights import GridWeights
+from .aggregate_utils import distributed_client, is_distributed, start_dask_client, shutdown_dask_client
+
+
+from dask.diagnostics import ProgressBar
+ProgressBar().register()
+
+
+def transform_dataset(
+    dataset: Dataset, key: str, **kwargs: Union[int, List[int]]
+) -> Dict[str, Dataset]:
+    """
+    Transform a given dataset by raising it to a power or interacting with another dataset.
+
+    Args:
+        dataset (Dataset): The dataset to transform.
+        key (str): The name of the dataset.
+        **kwargs:   Keyword arguments for the transformation.
+                    If 'exp' is provided, raise the dataset to the power of the provided exponent(s).
+                    If 'inter' is provided, interact the dataset with another dataset.
+
+    Returns:
+        dict: A dictionary containing the transformed dataset(s).
+    """
+
+    if "exp" in kwargs:
+        if not isinstance(kwargs["exp"], list):
+            kwargs["exp"] = [kwargs["exp"]]
+        dataset = [dataset.power(exp) for exp in kwargs["exp"][0]]
+        new_keys = [f"{key}_{exp}" for exp in kwargs["exp"][0]]
+        output_dict = dict(zip(new_keys, dataset))
+    elif "inter" in kwargs:
+        dataset = dataset.interact(kwargs["inter"])
+        output_dict = {key: dataset}
+    else:
+        raise ValueError("No valid transform argument provided.")
+    return output_dict.values(), output_dict.keys()
+
+
+def aggregate_time(
+    dataset: Dataset,
+    weights: GridWeights = None,
+    aggregator_dict: Dict[str, Union[List[Tuple], TemporalAggregator]] = None,
+    **kwargs,
+) -> Dict[str, Dataset]:
+    """
+    Aggregate a dataset over time using the specified temporal aggregators.
+
+    Args:
+        dataset (Dataset): The dataset to aggregate.
+        weights (GridWeights, optional): The weights to use for aggregation. Defaults to None.
+        aggregator_dict (Dict[str, Union[List[Tuple], TemporalAggregator]], optional): A dictionary of temporal aggregators to apply to the dataset. Defaults to None.
+        **kwargs: Additional keyword arguments to use if `aggregator_dict` is not provided.
+
+    Returns:
+        Dict[str, Dataset]: A dictionary of aggregated datasets, with keys corresponding to the keys in `aggregator_dict`.
+    """
+
+
+def aggregate_time(
+    dataset: Dataset,
+    weights: GridWeights = None,
+    aggregator_dict: Dict[str, Union[List[Tuple], TemporalAggregator]] = None,
+    **kwargs,
+) -> Dict[str, Dataset]:
+    if aggregator_dict is None:
+        if kwargs is None:
+            raise ValueError("No arguments provided.")
+        else:
+            aggregator_dict = kwargs
+    out_dict = {}
+    for key, value in aggregator_dict.items():
+        keys = [key]
+        data = [dataset.deepcopy()]
+        for key2, value2 in value:
+            if key2 == "aggregate":
+                if not isinstance(value2, TemporalAggregator):
+                    value2 = TemporalAggregator(**value2)
+                data = [value2.execute(x, weights) for x in data]
+
+                if value2.multi_dd:
+                    if len(data) > 1:
+                        raise ValueError(
+                            "Cannot aggregate multiple datasets with multiple ddargs, e.g., multiple polynomials for multiple bins"
+                        )
+                    data, keys = multi_dd_to_dict(data[0], key, value2.ddargs)
+
+            elif key2 == "transform":
+                transformed_data, transformed_keys = [], []
+                for d, k in zip(data, keys):
+                    d2, k2 = transform_dataset(d, k, **value2)
+                    transformed_data.extend(d2)
+                    transformed_keys.extend(k2)
+                data, keys = transformed_data, transformed_keys
+
+        data_dict = dict(zip(keys, data))
+        out_dict = out_dict | data_dict
+    return out_dict
+
+
+def aggregate_space(
+    dataset_dict: Dict[str, Dataset], weights: GridWeights,
+    npartitions=None, **kwargs
+) -> pd.DataFrame:
+    """
+    Aggregate a dictionary of datasets over space.
+
+    Args:
+        dataset_dict (dict): A dictionary containing the datasets to aggregate, where
+                            the keys are the names of the datasets and the values are the datasets themselves.
+        weights (GridWeights): The weights to use for aggregation.
+
+    Returns:
+        df: A dataframe containing the aggregated data.
+    """
+    
+    
+    dataset_list = list(dataset_dict.values())
+    
+    # client = distributed_client()
+    # if client is None and npartitions is None:
+    #     npartitions=1
+    # else:
+        # npartitions=len(client.scheduler_info()["workers"])
+        # da_list = dask.persist([x.da for x in dataset_list])[0]
+        # progress(da_list)
+        # for i, dataset in enumerate(dataset_list):
+        #     dataset.da = da_list[i]
+    df = SpatialAggregator(
+        dataset_list, weights, names=list(dataset_dict.keys()),
+    ).compute(npartitions=npartitions)
+    return df
+
+
+def aggregate_dataset(
+    weights: GridWeights,
+    dataset: Dataset = None,
+    aggregator_dict: Dict[str, Union[List[Tuple], TemporalAggregator]] = None,
+    dataset_dict = None,
+    n_workers = 50,
+    threads_per_worker = 1,
+    processes = True,
+    memory_limit=None,
+    cluster_args = {},
+    **kwargs,
+) -> pd.DataFrame:
+    """
+    Aggregate a dataset over time and space.
+
+    Args:
+        dataset (Dataset): The dataset to aggregate.
+        weights (GridWeights): The weights to use for aggregation.
+        agg_dict (dict): A dictionary containing the arguments for creating TemporalAggregator objects.
+                        The keys of the dictionary are names, and the values are a list of either tuples or TemporalAggregator objects.
+                        If the list contains tuples, use them as arguments to instantiate a temporal aggregator.
+
+    Returns:
+        df: A dataframe containing the aggregated data.
+    """
+
+    if dataset is None:
+        raise ValueError("No dataset provided.")
+    
+    if aggregator_dict is None and kwargs is not None:
+        aggregator_dict = kwargs
+    
+    if aggregator_dict is not None:
+        dataset_dict = aggregate_time(dataset, weights, aggregator_dict)
+    elif dataset_dict is None:
+        dataset_dict = {"variable": dataset}
+        if dataset_dict is None and dataset is None:
+            raise ValueError("No aggregator dict or dataset dict provided.")
+
+    df = aggregate_space(dataset_dict, weights)
+    df = (
+        weights.georegions.shp[[weights.georegions.regionid]].merge(
+            df, left_index=True, right_on="region_id"
+        )
+    ).drop(columns="region_id")
+    
+    # client.shutdown()
+
+    # _ = shutdown_dask_client()
+    
+    return df
+
+
+def multi_dd_to_dict(data, key, ddargs):
+    """
+    Converts a multi-variable list of datasets to a dictionary with keys
+    generated from the given key and ddargs.
+
+    Args:
+        data (list): The list of Datasets to convert to a dictionary.
+        key (str): The base key to use for generating the dictionary keys.
+        ddargs (list): A list of tuples representing the dimensions of the array.
+
+    Returns:
+        dict: A dictionary with keys generated from the given key and ddargs, and values from the list of datasets.
+    """
+    keys = [f"{key}_{x[0]}_{x[1]}" for x in ddargs]
+    # data_dict = dict(zip(keys, data))
+    return data, keys

aggfly/aggregate/aggregate_utils.py

@@ -0,0 +1,76 @@
+import dask.distributed
+from dask.distributed import Client
+import logging
+
+def is_distributed():
+    """
+    Returns True if the code is running in a distributed environment, False otherwise.
+    
+    This function checks if the code is running in a distributed environment by attempting to get the global Dask client.
+    If the client is not None, it means that the code is running in a distributed environment.
+    
+    Returns:
+        bool: True if the code is running in a distributed environment, False otherwise.
+    """
+    client = dask.distributed.client._get_global_client()
+    if client is not None:
+        return True
+    else:
+        return False
+    
+def distributed_client():
+    """
+    Returns the global Dask distributed client object.
+
+    Returns:
+    --------
+    client : dask.distributed.Client
+        The global Dask distributed client object.
+    """
+    client = dask.distributed.client._get_global_client()
+    return client 
+
+        
+def start_dask_client(n_workers: int = 2, threads_per_worker: int = 2, **kwargs):
+    """
+    Start a dask distributed cluster.
+
+    Args:
+        n_workers (int, optional): The number of workers to use. Defaults to 2.
+        threads_per_worker (int, optional): The number of threads per worker. Defaults to 2.
+        **kwargs: Additional keyword arguments to pass to the dask Client constructor.
+
+    Returns:
+        client: A dask distributed client.
+    """
+    client = Client(
+        n_workers=n_workers,
+        threads_per_worker=threads_per_worker,
+        **kwargs
+    )
+        
+    arg_dict = {
+        "n_workers": n_workers, 
+        "threads_per_worker": threads_per_worker,
+    }
+    all_dict = {**arg_dict, **kwargs}
+    for k in all_dict.keys():
+        client.set_metadata(['args', k], all_dict[k])
+    
+    return client
+
+
+def shutdown_dask_client():
+    """
+    Shutdown the global Dask distributed client object.
+    """
+    client = dask.distributed.client._get_global_client()
+    if client is not None:
+        try:
+            args = client.get_metadata('args')
+        except:
+            raise ValueError("Please start Dask client with af.start_dask_client() or run weight calculation without Dask distributed client.")
+        client.shutdown()
+        return args
+    else:
+        return None

aggfly/aggregate/spatial.py

@@ -0,0 +1,154 @@
+"""
+SpatialAggregator class that aggregates climate data over regions using weights.
+
+Attributes:
+    clim (list): List of xarray DataArrays containing climate data.
+    weights (xarray Dataset): Dataset containing region weights.
+    names (list): List of names for the climate data variables.
+
+Methods:
+    compute(npartitions=30): Aggregates climate data over regions using weights.
+    weighted_average(ddf, names): Computes weighted average of climate data.
+"""
+import warnings
+
+import xarray as xr
+import dask
+from dask.distributed import progress
+
+from .aggregate_utils import distributed_client, is_distributed
+from ..dataset import Dataset
+from ..weights import GridWeights
+
+import logging
+from typing import List, Union
+import dask.dataframe
+import pandas as pd
+import xarray as xr
+
+class SpatialAggregator:
+    """
+    A class for spatially aggregating climate data using weights.
+
+    Parameters:
+    -----------
+    clim : list or Dataset
+        A list of Dataset objects containing climate data to be aggregated.
+    weights : GridWeights 
+        A GridWeights object containing the weights to be used for aggregation.
+    names : str or list of str, optional
+        The name(s) of the climate variable(s) to be aggregated. Default is "climate".
+
+    Methods:
+    --------
+    compute(npartitions: int = 30) -> pd.DataFrame:
+        Compute the spatial aggregation.
+
+    weighted_average(ddf: dask.dataframe.DataFrame, names: List[str]) -> pd.DataFrame:
+        Compute the weighted average of the climate data.
+
+    """
+    def __init__(self, dataset: Union[list, Dataset], weights: GridWeights, names: Union[str, List[str]] = "climate") -> None:
+        """
+        Initialize a SpatialAggregator object.
+
+        Parameters
+        ----------
+        dataset : list or Dataset
+            A list of Dataset objects, each containing the climate data
+            for a different temporal aggregation. Alternatively, a single Dataset
+            object can be passed.
+        weights : GridWeights
+            A GridWeights object containing the spatial weights used to
+            aggregate the climate data.
+        names : str or list of str, optional
+            The name(s) of the climate variable(s) being aggregated. If a single
+            variable is being aggregated, a string can be passed. If multiple
+            variables are being aggregated, a list of strings can be passed.
+
+        Returns
+        -------
+        None
+
+        """
+        if type(dataset) != list:
+            self.dataset = [dataset]
+        else:
+            self.dataset = dataset
+        _ = [x.rescale_longitude() for x in self.dataset if x.lon_is_360]
+        self.grid = weights.grid
+        self.weights = weights.weights
+        self.names = [names] if isinstance(names, str) else names
+
+    def compute(self, npartitions: int = 30) -> pd.DataFrame:
+        """
+        Compute the weighted average of the climate data over the regions defined by the weights.
+
+        Parameters:
+        -----------
+        npartitions : int, optional
+            The number of partitions to use for the Dask DataFrame. Default is 30.
+
+        Returns:
+        --------
+        aggregated : pandas.DataFrame
+            A DataFrame containing the weighted average of the climate data over the regions and time periods.
+        """
+        # with dask.config.set({"multiprocessing.context": "forkserver"}):
+        print("Computing...")
+        clim_ds = dask.compute([x.da for x in self.dataset])[0] #, scheduler='processes'
+        
+        print("Combining datasets...")
+        clim_ds = xr.combine_by_coords(
+            [x.to_dataset(name=self.names[i]) for i, x in enumerate(clim_ds)]
+        )
+
+        print("Stacking...")
+        clim_df = (
+            clim_ds.stack({"cell_id": ["latitude", "longitude"]})
+            .drop_vars(["cell_id", "latitude", "longitude"])
+            .assign_coords(coords={"cell_id": ("cell_id", self.dataset[0].grid.cell_id)})
+            .to_dataframe()
+            .reset_index("time")
+            .dropna(subset=self.names)
+        )
+
+        print("Merging...")
+        self.weights["region_id"] = self.weights.index_right
+        merged_df = clim_df.merge(self.weights, how="inner", on="cell_id")
+        merged_df = merged_df.dropna(subset=self.names)
+
+        print("Grouping...")
+        group_key = (
+            merged_df[["region_id", "time"]]
+            .drop_duplicates()
+            .reset_index(drop=True)
+            .reset_index()
+            .rename(columns={"index": "group_ID"})
+        )
+        
+        print("Merging again...")
+        merged_df = merged_df.merge(group_key, on=["region_id", "time"]).set_index(
+            "group_ID"
+        )[["weight", *self.names]]
+
+        print("Creating Dask DataFrame...")
+        ddf = dask.dataframe.from_pandas(merged_df, npartitions=50)
+        
+        print("Aggregating...")
+        out = self.weighted_average(ddf, self.names).compute()
+        aggregated = (
+            out.merge(group_key, how="right", left_index=True, right_on="group_ID")
+            .drop(columns="group_ID")[["region_id", "time"] + self.names]
+            .reset_index(drop=True)
+        )
+
+        return aggregated
+    
+    @staticmethod
+    def weighted_average(ddf: dask.dataframe.DataFrame, names: List[str]) -> pd.DataFrame:
+        out = ddf[names].mul(ddf["weight"], axis=0)
+        out["weight"] = ddf["weight"]
+        out = out.groupby(out.index).sum()
+        out = out[names].div(out["weight"], axis=0)
+        return out

aggfly/aggregate/temporal.py

@@ -0,0 +1,198 @@
+from copy import deepcopy
+import os
+os.environ['USE_PYGEOS'] = '0'
+
+import numpy as np
+import dask.array as da
+from .aggregate_utils import *
+from ..dataset import Dataset, array_lon_to_360
+from ..weights import GridWeights
+
+import numpy as np
+from typing import List, Union
+from copy import deepcopy
+
+class TemporalAggregator:
+    """
+    A class for aggregating temporal data.
+
+    Parameters:
+    -----------
+    calc: str
+        The type of calculation to perform. Can be one of "mean", "sum", "dd", "bins", "min", or "max".
+    groupby: str
+        The time frequency to group the data by.
+    ddargs: List[Union[int, float]], optional
+        A list of values to use for the "dd" calculation. Only used if calc is "dd" or "bins".
+
+    Attributes:
+    -----------
+    calc: str
+        The type of calculation to perform.
+    groupby: str
+        The time frequency to group the data by.
+    kwargs: dict
+        Additional keyword arguments to pass to the reduce function.
+    ddargs: List[Union[int, float]], optional
+        A list of values to use for the "dd" calculation.
+    multi_dd: bool
+        Whether or not multiple "dd" values were provided.
+    func: function
+        The function to use for the calculation.
+
+    Methods:
+    --------
+    assign_func(self) -> function:
+        Assigns the appropriate function based on the value of self.calc.
+    get_ddargs(self, ddargs: List[Union[int, float]]) -> List[Union[int, float]]:
+        Returns the ddargs list if it is not None, and sets self.multi_dd to True if there are multiple values.
+    execute(self, dataset: Dataset, weights: Dataset, update: bool, **kwargs) -> Dataset:
+        Executes the aggregation and returns the result.
+    """
+
+    def __init__(self, calc: str, groupby: str, ddargs: List[Union[int, float]] = None):
+        self.calc = calc
+        self.groupby = translate_groupby(groupby)
+        self.kwargs = {}
+        self.ddargs = self.get_ddargs(ddargs)
+        self.func = self.assign_func()
+
+    def assign_func(self):
+        """
+        Assigns the appropriate function based on the value of self.calc.
+
+        Returns:
+        --------
+        function:
+            The function to use for the calculation.
+        """
+        if self.calc == "mean":
+            f = np.mean
+        if self.calc == "sum":
+            f = np.sum
+        elif self.calc == "dd":
+            if self.multi_dd:
+                f = _multi_dd
+            else:
+                f = _dd
+            self.kwargs = {"ddargs": self.ddargs}
+        elif self.calc == "bins":
+            if self.multi_dd:
+                f = _multi_bins
+            else:
+                f = _bins
+            self.kwargs = {"ddargs": self.ddargs}
+        if self.calc == "min":
+            f = np.min
+        if self.calc == "max":
+            f = np.max
+        return f
+
+    def get_ddargs(self, ddargs: List[Union[int, float]]) -> List[Union[int, float]]:
+        """
+        Returns the ddargs list if it is not None, and sets self.multi_dd to True if there are multiple values.
+
+        Parameters:
+        -----------
+        ddargs: List[Union[int, float]]
+            A list of values to use for the "dd" calculation.
+
+        Returns:
+        --------
+        List[Union[int, float]]:
+            The ddargs list.
+        """
+        if ddargs is None:
+            self.multi_dd = False
+            return None
+        else:
+            ddarr = np.array(ddargs)
+            if len(ddarr.shape) > 1:
+                self.multi_dd = True
+            else:
+                self.multi_dd = False
+            return ddargs
+
+    def execute(self, dataset: Dataset, weights: GridWeights = None, update: bool = False, **kwargs) -> Dataset:
+        """
+        Executes the aggregation and returns the result.
+
+        Parameters:
+        -----------
+        dataset: Dataset
+            The data to aggregate.
+        weights: GridWeights, optional
+            The weights to use for the aggregation.
+        update: bool, optional
+            Whether or not to update the input data with the result.
+        **kwargs:
+            Additional keyword arguments to pass to the reduce function.
+
+        Returns:
+        --------
+        Dataset:
+            The aggregated data.
+        """
+        ds = deepcopy(dataset.da)
+
+        # if weights is not None:
+        #     if dataset.grid.lon_is_360:
+        #         weights.nonzero_weight_mask = array_lon_to_360(weights.nonzero_weight_mask)
+        #     ds = ds.where(weights.nonzero_weight_mask)
+
+        if self.multi_dd:
+            ds = ds.expand_dims("dd", axis=-1)
+            if not update:
+                dataset_list = [deepcopy(dataset) for x in np.arange(len(self.ddargs))]
+        else:
+            if not update:
+                dataset = deepcopy(dataset)
+        with dask.config.set(**{'array.slicing.split_large_chunks': False}):
+            out = ds.resample(time=self.groupby).reduce(self.func, **self.kwargs)
+
+        if self.multi_dd:
+            out = out.to_dataset(dim="dd")
+            out = [out[var_name] for var_name in out.variables]
+
+            # Update object and return result
+            if type(dataset) == Dataset:
+                [x.update(y) for x, y in zip(dataset_list, out)]
+                [x.history.append(self.groupby) for x in dataset_list]
+                if len(dataset_list) == 1:
+                    return dataset_list[0]
+                else:
+                    return dataset_list
+            else:
+                return out
+        else:
+            # Update object and return result
+            if type(dataset) == Dataset:
+                dataset.update(out)
+                dataset.history.append(self.groupby)
+                return dataset
+            else:
+                return out
+
+
+def _dd(frame, axis, ddargs):
+    return (
+        (frame > ddargs[0])
+        * (frame < ddargs[1])
+        * np.absolute(frame - ddargs[ddargs[2]])
+    ).sum(axis=axis)
+
+
+def _multi_dd(frame, axis, ddargs):
+    return da.concatenate([_dd(frame, axis, ddarg) for ddarg in ddargs], axis=-1)
+
+
+def _bins(frame, axis, ddargs):
+    return ((frame > ddargs[0]) * (frame < ddargs[1])).sum(axis=axis)
+
+
+def _multi_bins(frame, axis, ddargs):
+    return da.concatenate([_bins(frame, axis, ddarg) for ddarg in ddargs], axis=-1)
+
+
+def translate_groupby(groupby):
+    return {"date": "1D", "month": "ME", "year": "YE"}[groupby]