diff --git a/mllam_data_prep/config.py b/mllam_data_prep/config.py index b7a0ab4..72c774f 100644 --- a/mllam_data_prep/config.py +++ b/mllam_data_prep/config.py @@ -12,7 +12,7 @@ class InvalidConfigException(Exception): pass -def validate_config(config_inputs): +def validate_config(config_inputs: Dict[str, "InputDataset"]) -> None: """ Validate that, in the config: - either `variables` or `derived_variables` are present in the config diff --git a/mllam_data_prep/ops/chunking.py b/mllam_data_prep/ops/chunking.py index 9df27e9..cc61312 100644 --- a/mllam_data_prep/ops/chunking.py +++ b/mllam_data_prep/ops/chunking.py @@ -1,11 +1,14 @@ +from typing import Dict + import numpy as np +import xarray as xr from loguru import logger # Max chunk size warning CHUNK_MAX_SIZE_WARNING = 1 * 1024**3 # 1GB -def check_chunk_size(ds, chunks): +def check_chunk_size(ds: xr.Dataset, chunks: Dict[str, int]) -> None: """ Check the chunk size and warn if it exceeds CHUNK_MAX_SIZE_WARNING. @@ -44,7 +47,7 @@ def check_chunk_size(ds, chunks): ) -def chunk_dataset(ds, chunks): +def chunk_dataset(ds: xr.Dataset, chunks: Dict[str, int]) -> xr.Dataset: """ Check the chunk size and chunk the dataset. diff --git a/mllam_data_prep/ops/loading.py b/mllam_data_prep/ops/loading.py index f6bfc34..b104c5d 100644 --- a/mllam_data_prep/ops/loading.py +++ b/mllam_data_prep/ops/loading.py @@ -1,7 +1,7 @@ import xarray as xr -def load_input_dataset(fp): +def load_input_dataset(fp: str) -> xr.Dataset: """ Load the dataset diff --git a/mllam_data_prep/ops/mapping.py b/mllam_data_prep/ops/mapping.py index 9482ff8..79fe57c 100644 --- a/mllam_data_prep/ops/mapping.py +++ b/mllam_data_prep/ops/mapping.py @@ -1,14 +1,23 @@ +from typing import Any, Dict, List + +import xarray as xr + +from ..config import DimMapping from .stacking import stack_variables_as_coord_values, stack_variables_by_coord_values -def _check_for_malformed_list_arg(s): +def _check_for_malformed_list_arg(s: Any) -> None: if isinstance(s, str) and "," in s: raise Exception( "Rather than writing `{s}` to define a list you would `[{s}]` in the config file." ) -def map_dims_and_variables(ds, dim_mapping, expected_input_var_dims): +def map_dims_and_variables( + ds: xr.Dataset, + dim_mapping: Dict[str, DimMapping], + expected_input_var_dims: List[str], +) -> xr.DataArray: """ Map the input dimensions to the architecture dimensions using the `dim_mapping` dictionary. Each key in the `dim_mapping` diff --git a/mllam_data_prep/ops/selection.py b/mllam_data_prep/ops/selection.py index 37b91c1..e73530f 100644 --- a/mllam_data_prep/ops/selection.py +++ b/mllam_data_prep/ops/selection.py @@ -1,11 +1,15 @@ import datetime +from typing import List, Optional, Union import pandas as pd +import xarray as xr from ..config import Range -def _normalize_slice_startstop(s): +def _normalize_slice_startstop( + s: Union[pd.Timestamp, str, int, float] +) -> Union[pd.Timestamp, str, int, float]: if isinstance(s, pd.Timestamp): return s elif isinstance(s, str): @@ -17,7 +21,9 @@ def _normalize_slice_startstop(s): return s -def _normalize_slice_step(s): +def _normalize_slice_step( + s: Optional[Union[pd.Timedelta, str, int, float]] +) -> Optional[Union[pd.Timedelta, str, int, float]]: if isinstance(s, pd.Timedelta): return s elif isinstance(s, str): @@ -29,7 +35,9 @@ def _normalize_slice_step(s): return s -def select_by_kwargs(ds, **coord_ranges): +def select_by_kwargs( + ds: xr.Dataset, **coord_ranges: Union[Range, List[Any]] +) -> xr.Dataset: """ Do `xr.Dataset.sel` on `ds` using the `coord_ranges` to select the coordinates, for each entry in the dictionary, the key is the coordinate name and the value is the selection @@ -92,7 +100,7 @@ def select_by_kwargs(ds, **coord_ranges): return ds -def check_point_in_dataset(coord, point, ds): +def check_point_in_dataset(coord: str, point: Any, ds: xr.Dataset) -> None: """ check that the requested point is in the data. """ @@ -102,7 +110,11 @@ def check_point_in_dataset(coord, point, ds): ) -def check_step(sel_step, coord, ds): +def check_step( + sel_step: Union[pd.Timedelta, datetime.timedelta], + coord: str, + ds: xr.Dataset, +) -> None: """ check that the step requested is exactly what the data has """ diff --git a/mllam_data_prep/ops/stacking.py b/mllam_data_prep/ops/stacking.py index a56e0fd..4a436ef 100644 --- a/mllam_data_prep/ops/stacking.py +++ b/mllam_data_prep/ops/stacking.py @@ -1,7 +1,9 @@ import xarray as xr -def stack_variables_as_coord_values(ds, name_format, combined_dim_name): +def stack_variables_as_coord_values( + ds: xr.Dataset, name_format: str, combined_dim_name: str +) -> xr.DataArray: """ combine all variables in an xr.Dataset into a single xr.DataArray by stacking the variables along a new coordinate with the name given @@ -50,7 +52,12 @@ def stack_variables_as_coord_values(ds, name_format, combined_dim_name): return da_combined -def stack_variables_by_coord_values(ds, coord, name_format, combined_dim_name): +def stack_variables_by_coord_values( + ds: xr.Dataset, + coord: str, + name_format: str, + combined_dim_name: str, +) -> xr.DataArray: """ combine all variables in an xr.Dataset on all coordinate values of `coord` into a single xr.DataArray diff --git a/mllam_data_prep/ops/subsetting.py b/mllam_data_prep/ops/subsetting.py index 80f2ce1..7c7cdcd 100644 --- a/mllam_data_prep/ops/subsetting.py +++ b/mllam_data_prep/ops/subsetting.py @@ -1,4 +1,15 @@ -def extract_variable(ds, var_name, coords_to_sample=dict()): +from typing import Dict, Optional + +import xarray as xr + +from ..config import ValueSelection + + +def extract_variable( + ds: xr.Dataset, + var_name: str, + coords_to_sample: Optional[Dict[str, ValueSelection]] = None, +) -> xr.DataArray: """ Extract specified variable from the provided input dataset. If coordinates for subsetting are defined, then subset the variable along @@ -8,19 +19,20 @@ def extract_variable(ds, var_name, coords_to_sample=dict()): ---------- ds : xr.Dataset Input dataset - var_name : Union[Dict, List] - Either a list or dictionary with variables to extract. - If a dictionary the keys are the variable name and the values are - entries for each coordinate and coordinate values to extract - coords_to_sample: Dict + var_name : str + Name of the variable to extract from the dataset + coords_to_sample: Dict[str, ValueSelection], optional Optional argument for subsetting/sampling along the specified - coordinates + coordinates. Keys are coordinate names, values are ValueSelection + objects defining the values to select and optionally the units. Returns ---------- da: xr.DataArray Extracted variable (subsetted along the specified coordinates) """ + if coords_to_sample is None: + coords_to_sample = {} try: da = ds[var_name]