tfv-get-tools 0.2.0__py3-none-any.whl

tfv_get_tools/utilities/horizontal_padding.py

@@ -0,0 +1,89 @@
+"""
+Set of functions to managing horizontally padding an OCEAN dataset
+
+This will cast the nearest value horizontally (lat/lon) for each depth level, through all times
+
+"""
+import xarray as xr
+import numpy as np
+from scipy import spatial
+
+def horizontal_pad(ds):
+    """Pad an ocean dataset horizontally for each time.
+
+    Args:
+        ds (xr.Dataset): input merged ocean dataset
+
+    Returns:
+        xr.Dataset: resulting padded ocean dataset
+    """
+    for var in ds.data_vars:
+        if set(ds[var].dims) >= {'latitude', 'longitude', 'time'}:
+            # Compute nearest indices for each depth level
+            nearest_indices = xr.apply_ufunc(
+                _compute_nearest_indices,
+                ds[var].isel(time=0),
+                input_core_dims=[['latitude', 'longitude']],
+                output_core_dims=[['latitude', 'longitude']],
+                vectorize=True
+            )
+            
+            # Apply filling to all time steps for each depth
+            filled_var = xr.apply_ufunc(
+                _fill_nans_with_precomputed_indices,
+                ds[var],
+                nearest_indices,
+                input_core_dims=[['latitude', 'longitude'], ['latitude', 'longitude']],
+                output_core_dims=[['latitude', 'longitude']],
+                vectorize=True
+            )
+            
+            ds[var] = filled_var
+
+    return ds
+
+def _compute_nearest_indices(values: np.ndarray) -> np.ndarray:
+    """Find the indices of the nearest non nan values across longitude/latitude
+
+    Args:
+        values (np.ndarray): variable values
+
+    Returns:
+        np.ndarray: indices array
+    """
+    ny, nx = values.shape
+    lat = np.arange(ny)
+    lon = np.arange(nx)
+    lat, lon = np.meshgrid(lat, lon, indexing='ij')
+    
+    valid_mask = ~np.isnan(values)
+    valid_points = np.column_stack((lat[valid_mask], lon[valid_mask]))
+    
+    # Create KDTree for efficient nearest neighbor search
+    tree = spatial.cKDTree(valid_points)
+    
+    # Find nearest non-NaN point for each point
+    all_points = np.column_stack((lat.ravel(), lon.ravel()))
+    _, indices = tree.query(all_points)
+    
+    # Reshape indices to match the original shape
+    return indices.reshape(values.shape)
+
+def _fill_nans_with_precomputed_indices(values: np.ndarray, nearest_indices: np.ndarray) -> np.ndarray:
+    """Fill nan values in dataset wih the pre-calced indicies.
+
+    Args:
+        values (np.ndarray): variable values
+        nearest_indices (np.ndarray): indicies if nearest non-nan value across longitude/latitude
+
+    Returns:
+        np.ndarray: filled variable values
+    """
+    valid_values = values[~np.isnan(values)]
+    
+    # Apply the precomputed indices to fill NaN values
+    filled_values = values.copy()
+    nan_mask = np.isnan(filled_values)
+    filled_values[nan_mask] = valid_values[nearest_indices[nan_mask]]
+    
+    return filled_values

tfv_get_tools/utilities/land_masking.py

@@ -0,0 +1,93 @@
+"""
+Module to handle land masking of an atmospheric dataset.
+This has been straight air-lifted from GetAtmos with no modifications.
+"""
+
+import lzma
+import pickle
+import importlib.resources
+
+import xarray as xr
+import numpy as np
+import dask.array as da
+import pandas as pd
+import logging
+from scipy.ndimage import distance_transform_edt
+
+from tqdm.auto import trange
+
+
+def load_atmos_masks():
+   """Returns a dictionary containing landmasks for ERA5, CFSR, CFSv2 and BARRA_R.
+
+   Returns:
+       lm (dict): Dict containing keys 'cfsr', 'cfsv2', 'barra_r', 'era5'.
+           Each of these contains a sub-dict that has 'x', 'y' and 'lm' keys.
+   """
+   with importlib.resources.files(__name__).parent.joinpath("providers/atmos/data/atmos_masks.xz").open("rb") as stream:
+       with lzma.open(stream, "rb") as f:
+           lmdict = pickle.load(f)
+   
+   return lmdict
+
+
+def mask_land_data(ds, source):
+    lms = load_atmos_masks()
+
+    # Load all land masks (TODO: BETTER HANDLING OF CFSR and CFSv2 SWITCH)
+    if "cfs" in source:
+        if ds["time"][-1] > pd.Timestamp(2011, 1, 1):
+            k = "cfsv2"
+        else:
+            k = "cfsr"
+    else:
+        k = source.lower()
+    lm = lms[k]
+
+    dsm = xr.DataArray(
+        np.squeeze(lm["lm"]), coords=dict(latitude=lm["y"], longitude=lm["x"])
+    )
+
+    ds["lsm"] = dsm.reindex(
+        longitude=ds["longitude"], latitude=ds["latitude"], method="nearest"
+    )
+    ds["lsm"].attrs = {"units": "-", "long_name": f"{k} land mask index"}
+      
+    # If the Atmos dataset is in negative longitudes, we need to wrap the land mask around 180
+    # (i.e., -180 to 180), not (0 to 360)
+    if ds["longitude"].min() < 0.0:
+        lon = dsm['longitude'].values
+        lon[lon > 180] = (lon - 360)[lon > 180]
+        dsm['longitude'] = lon    
+        dsm = dsm.sortby('longitude')
+    
+    mask = (ds["lsm"] > 0.2).values
+    mask_3d = repeat_2d_array(mask, ds.sizes["time"])
+
+    for var in ds.data_vars.keys():
+        dims = ds[var].dims
+        if ('longitude' in dims) & ('latitude' in dims) & ('time' in dims):
+            logging.debug(f"{var} in ds, running land mask")
+            ds[var].data = fill(ds[var].values, mask_3d)
+
+    return ds
+
+
+def fill(data, mask):
+    """Replace "masked" data with the nearest valid data cell
+
+    Args:
+        data (da.ndarray): 2D input data array
+        invalid (da.ndarray): 2D boolean array, where True cells are replaced by the nearest data.
+
+    Returns:
+        filled_data: 2D output array with filled data
+    """
+
+    ind = distance_transform_edt(mask, return_distances=False, return_indices=True)
+    return data[tuple(ind)]
+
+
+def repeat_2d_array(array, n):
+    repeated_array = da.tile(array, (n, 1, 1))
+    return repeated_array

tfv_get_tools/utilities/parsers.py

@@ -0,0 +1,44 @@
+from datetime import datetime
+import pandas as pd
+from pathlib import Path
+from typing import Union
+
+
+def _parse_date(date: Union[str, datetime, pd.Timestamp]) -> datetime:
+    if isinstance(date, str):
+        if len(date) == 10:
+            fmt = "%Y-%m-%d"
+        elif len(date) == 13:
+            fmt = "%Y-%m-%d %H"
+        elif len(date) == 16:
+            fmt = "%Y-%m-%d %H:%M"
+        elif len(date) == 19:
+            fmt = "%Y-%m-%d %H:%M:%S"
+        elif len(date) == 8:
+            fmt = "%Y%m%d"
+        elif len(date) == 15:
+            fmt = "%Y%m%d.%H%M%S"
+
+        try:
+            date = datetime.strptime(date, fmt)
+        except:
+            raise ValueError(f'Failed to convert date `{date}` to a Timestamp, please check format (e.g., YYYY-mm-dd)')
+        return date
+    
+    elif isinstance(date, pd.Timestamp):
+        return date.to_pydatetime()
+    elif isinstance(date, datetime):
+        return date
+    else:
+        raise ValueError(
+            "Date must be a string, datetime, or pandas Timestamp object"
+        )
+
+def _parse_path(path: Union[str, Path]) -> Path:
+    path = Path(path)
+    if path.is_dir():
+        return path
+    else:
+        raise NotADirectoryError(
+            f"`{path.as_posix()}` does not exist, please check."
+        )

tfv_get_tools/utilities/warnings.py

@@ -0,0 +1,38 @@
+import warnings
+import functools
+
+
+def deprecated(new_func):
+    """
+    Decorator to mark functions as deprecated.
+    
+    This decorator will result in a warning being emitted
+    when the decorated function is used.
+    
+    Args:
+        new_func (callable): The new function to be used instead of the deprecated function.
+    
+    Returns:
+        callable: A decorator function.
+    
+    Example:
+        >>> def new_function(x, y):
+        ...     return x + y
+        ...
+        >>> @deprecated(new_function)
+        ... def old_function(x, y):
+        ...     return new_function(x, y)
+        ...
+        >>> old_function(1, 2)
+        3
+        # Warning: Call to deprecated function old_function. Use new_function instead.
+    """
+    def decorator(old_func):
+        @functools.wraps(old_func)
+        def wrapper(*args, **kwargs):
+            warnings.warn(f"Call to deprecated function {old_func.__name__}. "
+                          f"Use {new_func.__name__} instead.",
+                          category=DeprecationWarning, stacklevel=2)
+            return new_func(*args, **kwargs)
+        return wrapper
+    return decorator

tfv_get_tools/wave.py

@@ -0,0 +1,179 @@
+from datetime import datetime
+from pathlib import Path
+from typing import Union, Tuple, Optional, List
+import pandas as pd
+
+from tfv_get_tools.utilities._tfv_bc import write_tuflowfv_fvc
+from tfv_get_tools.providers._downloader import BatchDownloadResult
+
+def DownloadWave(
+    start_date: Union[str, datetime, pd.Timestamp],
+    end_date: Union[str, datetime, pd.Timestamp],
+    xlims: Tuple[float, float],
+    ylims: Tuple[float, float],
+    out_path: Union[str, Path] = Path("./raw"),
+    source: str = "CAWCR",
+    model: str = "default",
+    prefix: Optional[str] = None,
+    verbose: bool = False,
+    variables: Optional[List[str]] = None,
+    skip_check: bool = False,
+    **kwargs,
+) -> BatchDownloadResult:
+    """Download Wave Data
+    
+    Users should call this function, not the individual downloader classes directly.
+    
+    This module will download wave data from several possible sources to facilitate 
+    TUFLOW FV and SWAN modelling.
+
+    The following sources have been implemented:
+        - `CAWCR` - CSIRO Ocean wave hindcast forced using NCEP's CFSR Atmospheric Model
+            - `glob_24m` - Global domain at 24' resolution
+            - `aus_10m` - Australia ribbon model domain at 10' resolution
+            - `aus_4m` - Australia ribbon model domain at 4' resolution
+            - `pac_10m` - Pacific island domain at 10' resolution
+            - `pac_4m` - Pacific island domain at 4' resolution
+        - `Copernicus` - Various models from the Copernicus Marine Service
+            - `GLO` - Global wave model
+        - `ERA5` - Global 0.5deg Wave model (WAM) from the European Centre for 
+            Medium-Range Weather Forecasts (ECMWF)
+
+    Args:
+        start_date: Start date. The string format is `%Y-%m-%d` (e.g., '2011-01-01')
+        end_date: End date. The string format is `%Y-%m-%d` (e.g., '2011-02-01')
+        xlims: Minimum and maximum longitude, as floats. e.g., (115, 120)
+        ylims: Minimum and maximum latitude, as floats. e.g., (-40, -35)
+        out_path: Output directory for data files
+        source: Data source to download. One of {'CAWCR', 'Copernicus', 'ERA5'}
+        model: Choice of model, depending on "source"
+        prefix: Extra file name prefix
+        verbose: Print extra program information
+        variables: List of variables to download
+        skip_check: Skip user confirmation
+        **kwargs: Additional arguments
+        
+    Returns:
+        BatchDownloadResult: Results of the download operation
+    """
+    
+    if source.lower() == "cawcr":
+        from tfv_get_tools.providers.wave.cawcr import DownloadCAWCR
+        
+        downloader = DownloadCAWCR(
+            start_date=start_date,
+            end_date=end_date,
+            xlims=xlims,
+            ylims=ylims,
+            out_path=out_path,
+            model=model,
+            prefix=prefix,
+            verbose=verbose,
+            variables=variables,
+            skip_check=skip_check,
+            **kwargs
+        )
+        return downloader.execute_download()
+    
+    elif source.lower() == "copernicus":
+        from tfv_get_tools.providers.wave.copernicus_wave import DownloadCopernicusWave
+        
+        downloader = DownloadCopernicusWave(
+            start_date=start_date,
+            end_date=end_date,
+            xlims=xlims,
+            ylims=ylims,
+            out_path=out_path,
+            model=model,
+            prefix=prefix,
+            verbose=verbose,
+            variables=variables,
+            skip_check=skip_check,
+            **kwargs
+        )
+        return downloader.execute_download()
+
+    elif source.lower() == "era5":
+        from tfv_get_tools.providers.wave.era5 import DownloadERA5Wave
+        
+        downloader = DownloadERA5Wave(
+            start_date=start_date,
+            end_date=end_date,
+            xlims=xlims,
+            ylims=ylims,
+            out_path=out_path,
+            model=model,
+            prefix=prefix,
+            verbose=verbose,
+            variables=variables,
+            skip_check=skip_check,
+            **kwargs
+        )
+        return downloader.execute_download()
+        
+    else:
+        raise ValueError(f'Unrecognised source {source}. Must be one of: CAWCR, Copernicus, ERA5')
+
+
+
+def MergeWave(
+    in_path: Path = Path("./raw"),
+    out_path: Path = Path("."),
+    fname: str = None,
+    source: str = 'CAWCR',
+    model: str = 'default',
+    time_start: str = None,
+    time_end: str = None,
+    reproject: int = None,
+    local_tz: Tuple[float, str] = None,
+    wrapto360=False,
+):
+    """
+    Merge raw downloaded wave datafiles into a single netcdf file. 
+    
+    **Use the same `source` and `model` that was supplied to the Downloader function**
+
+    Args:
+        in_path (Path, optional): Directory of the raw ocean data-files. Defaults to Path(".").
+        out_path (Path, optional): Output directory for the merged ocean netcdf and (opt) the fvc. Defaults to Path(".").
+        fname (str, optional): Merged ocean netcdf filename. Defaults to None.
+        source (str, optional): Source to be merged, defaults to "CAWCR". 
+        model (str, optional): Model for source to be merged. Defaults to 'default' (which is 'glob_24m' for CAWCR).
+        time_start (str, optional): Start time limit of the merged dataset (str: "YYYY-mm-dd HH:MM"). Defaults to None.
+        time_end (str, optional): End time limit of the merged dataset (str: "YYYY-mm-dd HH:MM"). Defaults to None.
+        reproject (int, optional): Optionally reproject based, based on EPSG code. Defaults to None.
+        local_tz: (Tuple(float, str): optional): Add local timezone format is a tuple with Offset[float] and Label[str]
+        source: (str, optional): Ocean data source. Defaults to Hycom.
+        wrapto360: (bool, optional): Optionally wrap longitude to (0, 360) rather than (-180, 180). Defaults to False.
+    """
+    
+    args = tuple()
+
+    kwargs = dict(
+        in_path=in_path,
+        out_path=out_path,
+        fname=fname,
+        source=source,
+        model=model,
+        time_start=time_start,
+        time_end=time_end,
+        reproject=reproject,
+        local_tz=local_tz,
+        pad_dry=False,
+        wrapto360=wrapto360,
+    )
+    
+    if 'write_fvc' in kwargs:
+        print('Writing an FVC include file is not implemented for wave data. Skipping this flag')
+
+    if source.lower() == "cawcr":
+        from tfv_get_tools.providers.wave.cawcr import MergeCAWCR
+        MergeCAWCR(*args, **kwargs)
+    
+    elif source.lower() == "copernicus":
+        from tfv_get_tools.providers.wave.copernicus_wave import MergeCopernicusWave
+        MergeCopernicusWave(*args, **kwargs)
+
+    elif source.lower() == "era5":
+        from tfv_get_tools.providers.wave.era5 import MergeERA5Wave
+        MergeERA5Wave(*args, **kwargs)