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

tfv_get_tools/fvc/_atmos.py

@@ -0,0 +1,230 @@
+from pathlib import Path
+from typing import Union, List, Optional, Dict
+
+import xarray as xr
+
+from tfv_get_tools.fvc._fvc import FVCWriter
+
+
+class AtmosFVCWriter(FVCWriter):
+    """Writer for ATMOS FVC files."""
+
+    # Default mappings of dataset variables to TUFLOW-FV variables
+    DEFAULT_VAR_MAPPINGS = {
+        "t2m": {
+            "tfv_var": "AIR_TEMP_GRID",
+            "bc_scale": 1.0,
+            "bc_offset": -273.15,  # K to C
+        },
+        "relhum": {"tfv_var": "REL_HUM_GRID", "bc_scale": 1.0, "bc_offset": 0.0},
+        "u10": {
+            "tfv_var": "W10_GRID",  # Special case - paired with vwnd10m
+            "bc_scale": 1.0,
+            "bc_offset": 0.0,
+        },
+        "v10": {
+            "tfv_var": "W10_GRID",  # Special case - paired with uwnd10m
+            "bc_scale": 1.0,
+            "bc_offset": 0.0,
+        },
+        "mslp": {
+            "tfv_var": "MSLP_GRID",  # Special case - paired with uwnd10m
+            "bc_scale": 0.01,
+            "bc_offset": 0.0,
+        },
+        "prate": {"tfv_var": "PRECIP_GRID", "bc_scale": 1.0, "bc_offset": 0.0},
+        "dlwrf": {"tfv_var": "LW_RAD_GRID", "bc_scale": 1.0, "bc_offset": 0.0},
+        "dswrf": {"tfv_var": "SW_RAD_GRID", "bc_scale": 1.0, "bc_offset": 0.0},
+    }
+
+    def __init__(self, var_mappings: Optional[Dict] = None, source=None, model=None, info_url=None):
+        """Initialize atmospheric FVC writer.
+
+        Args:
+            var_mappings: Optional dictionary overriding default variable mappings
+        """
+        super().__init__(requires_coordinates=True, source=source, model=model, info_url=info_url)
+        self.var_mappings = var_mappings or self.DEFAULT_VAR_MAPPINGS
+        self.available_vars = set()
+
+    def detect_variables(self, ds: xr.Dataset) -> None:
+        """Detect available variables in the dataset.
+
+        Args:
+            ds: Input xarray Dataset
+        """
+        self.available_vars = set(ds.data_vars) & set(self.var_mappings.keys())
+
+        # Verify wind components come in pairs
+        if ("uwnd10m" in self.available_vars) != ("vwnd10m" in self.available_vars):
+            raise ValueError(
+                "Both uwnd10m and vwnd10m must be present for wind configuration"
+            )
+
+    def write_grid_definition(self, lines: List[str], nc_path: str) -> None:
+        """Write grid definition block."""
+        xvar, yvar = self.coordinate_vars
+        lines.extend(
+            [
+                f"Grid Definition File == {nc_path}",
+                f"  Grid Definition Variables == {xvar}, {yvar}",
+                "  Grid Definition Label == atmos",
+                "End Grid",
+                "",
+            ]
+        )
+
+    def write_boundary_conditions(self, lines: List[str], nc_path: str) -> None:
+        """Write boundary conditions block."""
+        wind_written = False
+
+        for var_name in sorted(self.available_vars):
+            config = self.var_mappings[var_name]
+            tfv_var = config["tfv_var"]
+            bc_scale = config["bc_scale"]
+            bc_offset = config["bc_offset"]
+
+            # Handle wind components specially
+            if tfv_var == "W10_GRID":
+                if wind_written:
+                    continue
+                wind_written = True
+                lines.extend(
+                    [
+                        f"BC == {tfv_var}, atmos, {nc_path}",
+                        f"  BC Header == {self.time_var}, u10, v10",
+                    ]
+                )
+                if bc_scale != 1.0:
+                    lines.append(f"  BC Scale == {bc_scale}, {bc_scale}")
+                if bc_offset != 0.0:
+                    lines.append(f"  BC Offset == {bc_offset}, {bc_offset}")
+            else:
+                lines.extend(
+                    [
+                        f"BC == {tfv_var}, atmos, {nc_path}",
+                        f"  BC Header == {self.time_var}, {var_name}",
+                    ]
+                )
+                if bc_scale != 1.0:
+                    lines.append(f"  BC Scale == {bc_scale}")
+                if bc_offset != 0.0:
+                    lines.append(f"  BC Offset == {bc_offset}")
+
+            # Common BC settings
+            lines.extend(
+                [
+                    "  BC Update dt == 3600.",
+                    "  BC Time Units == hours",
+                    "  BC Reference Time == 01/01/1990 00:00",
+                    "  BC Default == NaN",
+                    "End BC",
+                    "",
+                ]
+            )
+
+    def generate(
+        self,
+        ds: xr.Dataset,
+        nc_path: str = "atmos_forcing.nc",
+    ) -> List[str]:
+        """Generate FVC configuration content for atmospheric forcing.
+
+        Args:
+            ds: Input xarray Dataset containing atmospheric variables
+            nc_path: Path or filename for the NetCDF file (referenced in FVC content)
+
+        Returns:
+            List[str]: FVC configuration content as a list of strings
+        """
+        # Process dataset to detect coordinates, time settings, and variables
+        self.process_dataset(ds)
+        self.detect_variables(ds)
+
+        if not self.available_vars:
+            raise ValueError("No supported atmospheric variables found in dataset")
+
+        lines = []
+
+        # Write header
+        self.write_header(lines, "Atmospheric Dataset", ds)
+
+        # Write grid definition
+        self.write_grid_definition(lines, nc_path)
+
+        # Write boundary conditions
+        self.write_boundary_conditions(lines, nc_path)
+
+        return lines
+
+    def write_file(
+        self,
+        lines: List[str],
+        output_path: Union[str, Path],
+        filename: Optional[str] = None,
+    ) -> None:
+        """Write FVC content to a file."""
+        output_path = Path(output_path)
+        filename = filename or "atmos_forcing.fvc"
+        if not filename.endswith(".fvc"):
+            filename = filename.replace(".nc", ".fvc")
+
+        with open(output_path / filename, "w") as f:
+            for line in lines:
+                f.write(line + "\n")
+
+
+def write_atmos_fvc(
+    ds: xr.Dataset,
+    nc_path: str = "atmos_forcing.nc",
+    output_path: Optional[Union[str, Path]] = None,
+    filename: Optional[str] = None,
+    var_mappings: Optional[Dict] = None,
+    source: Optional[str] = None,
+    model: Optional[str] = None,
+    info_url: Optional[str] = None,
+) -> List[str]:
+    """Generate (and optionally write) TUFLOW-FV atmospheric forcing configuration.
+
+    Args:
+        ds: Input xarray Dataset containing atmospheric variables
+        nc_path: Path or filename for the NetCDF file (referenced in FVC content)
+        output_path: Optional path to write FVC file (if None, no file is written)
+        filename: Optional filename for FVC file (if None, derives from nc_path)
+        var_mappings: Optional dictionary to override default variable mappings
+        source: Optional source string for the FVC header
+        model: Optional model string for the FVC header
+        info_url: Optional URL for source information to be printed in FVC
+
+    Returns:
+        List[str]: FVC configuration content as a list of strings
+
+    Examples:
+        # Just generate FVC content with auto-detected variables
+        >>> lines = write_atmos_fvc(dataset, nc_path="forcing.nc")
+
+        # Generate and write to file
+        >>> lines = write_atmos_fvc(
+        ...     dataset,
+        ...     nc_path="forcing.nc",
+        ...     output_path="path/to/output"
+        ... )
+
+        # Generate with custom variable mappings
+        >>> custom_mappings = {
+        ...     "temperature": {"tfv_var": "T2_GRID", "bc_scale": 1.0, "bc_offset": -273.15}
+        ... }
+        >>> lines = write_atmos_fvc(
+        ...     dataset,
+        ...     nc_path="forcing.nc",
+        ...     var_mappings=custom_mappings
+        ... )
+    """
+    writer = AtmosFVCWriter(var_mappings=var_mappings, source=source, model=model, info_url=info_url)
+
+    lines = writer.generate(ds, nc_path)
+
+    if output_path is not None:
+        writer.write_file(lines, output_path, filename)
+
+    return lines

tfv_get_tools/fvc/_fvc.py

@@ -0,0 +1,218 @@
+from datetime import datetime
+from typing import Callable, List, Optional, Tuple, Union
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+class FVCWriter:
+    """Base class for writing FVC files."""
+
+    def __init__(
+        self, requires_coordinates: bool = True, source=None, model=None, info_url=None
+    ):
+        """Initialize FVC writer.
+
+        Args:
+            requires_coordinates: Whether this FVC type requires coordinate information
+        """
+        self.requires_coordinates = requires_coordinates
+        self.coordinate_vars: Optional[Tuple[str, str]] = ("longitude", "latitude")
+        self.time_var: str = "time"
+        self.tzlbl: Optional[str] = None
+        self.epsg: Optional[str] = None
+        self.crs_name: Optional[str] = None
+
+        # For metadata
+        self.source = source
+        self.model = model
+        self.info_url = info_url
+
+    def detect_coordinates(self, ds: xr.Dataset) -> Optional[Tuple[str, str]]:
+        """Detect coordinate system from dataset if present.
+
+        Args:
+            ds: Input xarray Dataset
+
+        Returns:
+            Optional[Tuple[str, str]]: x and y coordinate variable names if found
+        """
+        if not self.requires_coordinates:
+            return None
+
+        if "x" in ds and "y" in ds:
+            # Get CRS information from reprojected coordinates
+            try:
+                self.epsg = ds["x"].attrs.get("epsg")
+                self.crs_name = ds["x"].attrs.get("name")
+            except (KeyError, AttributeError):
+                pass
+            return "x", "y"
+        elif "longitude" in ds and "latitude" in ds:
+            # Get CRS information from geographic coordinates if available
+            try:
+                self.epsg = ds["longitude"].attrs.get("epsg")
+                self.crs_name = ds["longitude"].attrs.get("name")
+            except (KeyError, AttributeError):
+                pass
+            return "longitude", "latitude"
+        elif self.requires_coordinates:
+            raise ValueError(
+                "Dataset requires coordinate variables (x,y or longitude,latitude) but none were found"
+            )
+        return None
+
+    def detect_time_settings(self, ds: xr.Dataset) -> str:
+        """Detect time variable and settings from dataset.
+
+        Args:
+            ds: Input xarray Dataset
+
+        Returns:
+            str: Name of time variable to use
+        """
+        if "local_time" in ds:
+            try:
+                self.tzlbl = ds["local_time"].attrs.get("tz")
+            except (KeyError, AttributeError):
+                pass
+            return "local_time"
+        return "time"
+
+    def process_dataset(self, ds: xr.Dataset) -> None:
+        """Process dataset to detect coordinate system and time settings.
+
+        Args:
+            ds: Input xarray Dataset
+        """
+        self.coordinate_vars = self.detect_coordinates(ds)
+        self.time_var = self.detect_time_settings(ds)
+
+    def get_coordinate_info(self, ds: xr.Dataset) -> Optional[dict]:
+        """Extract coordinate information from dataset if available.
+
+        Args:
+            ds: Input xarray Dataset
+
+        Returns:
+            Optional[dict]: Dictionary containing coordinate variable names and limits if coordinates exist
+        """
+        if not self.coordinate_vars:
+            return None
+
+        xvar, yvar = self.coordinate_vars
+        return {
+            "xvar": xvar,
+            "yvar": yvar,
+            "xlims": format_limits(ds[xvar].values),
+            "ylims": format_limits(ds[yvar].values),
+        }
+
+    def write_header(
+        self, f: Union[List[str], any], title: str, ds: xr.Dataset
+    ) -> None:
+        """Write standard FVC file header.
+
+        Args:
+            f: File handle or list to write to
+            title: Title string for the header
+            ds: xarray Dataset for time/coordinate info
+        """
+        # Process dataset if not already done
+        self.process_dataset(ds)
+
+        lines = []
+        lines.append(f"! TUFLOW FV FVC File for {title}")
+        lines.append("! Written by TUFLOW FV `tfv-get-tools`")
+        lines.append("")
+        # These lines are too long, need to split over several lines
+        lines.append("! This control file has been prepared using the TUFLOW FV Get Tools (tfv-get-tools),")
+        lines.append("! a free set of Python tools designed to assist with the download and formatting of")
+        lines.append("! boundary condition data from global model sources such as ERA5 and CFSR for use in TUFLOW FV.")
+        lines.append("! These external model datasets are subject to change over time and are provided 'as is'.")
+        lines.append("! Users are responsible for reviewing and, where possible, verifying these inputs against")
+        lines.append("! observational data before use in any modelling application.")
+        lines.append("")
+
+        # Now bang in some standard source info and a sicko disclaimer
+        lines.append(f"! Source: {self.source if self.source else 'Unknown'}")
+        if self.model:
+            if self.model != "default":
+                lines.append(f"! Model: {self.model}")
+        if self.info_url:
+            lines.append(f"! Info: {self.info_url}")
+        lines.append("")
+
+        # Add timezone information if using local time
+        time_var = "local_time" if "local_time" in ds else "time"
+        if time_var == "local_time":
+            lines.append(f"! NetCDF time datum: {ds[time_var].attrs['tz']}")
+        else:
+            lines.append("! NetCDF time datum: UTC")
+
+        # Add time information
+        lines.append(f"! NetCDF start time: {ds_time_to_str(ds, 0)}")
+        lines.append(f"! NetCDF end time: {ds_time_to_str(ds, -1)}")
+        lines.append("")
+
+        # Add coordinate system information if available
+        if self.requires_coordinates and (self.epsg or self.crs_name):
+            if self.epsg:
+                lines.append(f"! Coordinate system EPSG: {self.epsg}")
+            if self.crs_name:
+                lines.append(f"! Coordinate system name: {self.crs_name}")
+
+        # Add coordinate limits if required and available
+        if self.requires_coordinates:
+            coords = self.get_coordinate_info(ds)
+            if coords:
+                lines.append(f"! NetCDF x-limits: {coords['xlims']}")
+                lines.append(f"! NetCDF y-limits: {coords['ylims']}")
+
+        lines.append("")
+
+        if isinstance(f, list):
+            f.extend(lines)
+        else:
+            for line in lines:
+                f.write(line + "\n")
+
+
+def format_limits(values: np.ndarray, funcs: List[Callable] = [np.min, np.max]) -> str:
+    """Format array limits into a string using specified functions."""
+    return ", ".join([f"{fn(values):0.4f}" for fn in funcs])
+
+
+def format_timestamp(
+    timestamp: Union[str, datetime, np.datetime64, pd.Timestamp],
+    fmt: str = "%Y-%m-%d %H:%M",
+) -> str:
+    """Format various timestamp types to a consistent string format.
+
+    Args:
+        timestamp: Input timestamp as string, datetime, numpy.datetime64, or pandas.Timestamp
+        fmt: Output format string (default: "%Y-%m-%d %H:%M")
+
+    Returns:
+        str: Formatted timestamp string
+
+    Examples:
+        >>> format_timestamp("2024-01-01")
+        '2024-01-01 00:00'
+        >>> format_timestamp(np.datetime64("2024-01-01"))
+        '2024-01-01 00:00'
+        >>> format_timestamp(pd.Timestamp("2024-01-01"))
+        '2024-01-01 00:00'
+    """
+    # Convert to pandas Timestamp
+    if isinstance(timestamp, (str, np.datetime64, datetime)):
+        timestamp = pd.Timestamp(timestamp)
+
+    return timestamp.strftime(fmt)
+
+
+def ds_time_to_str(ds: xr.Dataset, i: int, fmt: str = "%Y-%m-%d %H:%M") -> str:
+    """Convert a time from an xarray Dataset to a formatted string."""
+    time_var = "local_time" if "local_time" in ds else "time"
+    return format_timestamp(ds[time_var].values[i], fmt)

tfv_get_tools/fvc/_ocean.py

@@ -0,0 +1,171 @@
+from typing import Union, List, Optional
+from pathlib import Path
+
+import xarray as xr
+
+from tfv_get_tools.fvc._fvc import FVCWriter
+
+
+class OceanFVCWriter(FVCWriter):
+    """Writer for OGCM FVC include files."""
+
+    DEFAULT_VAR_ORDER = ["surf_el", "water_u", "water_v", "salinity", "water_temp"]
+
+    def __init__(self, var_order: Optional[List[str]] = None, source=None, model=None, info_url=None):
+        """Initialize ocean FVC writer.
+
+        Args:
+            var_order: Optional list specifying order of variables in the boundary condition
+        """
+        super().__init__(requires_coordinates=True, source=source, model=model, info_url=info_url)
+        self.var_order = var_order or self.DEFAULT_VAR_ORDER
+
+    def write_grid_definition(self, lines: List[str], nc_path: str):
+        """Write grid definition block.
+
+        Args:
+            lines: List to append lines to
+            nc_path: Path to the NetCDF file
+        """
+        xvar, yvar = self.coordinate_vars
+        lines.extend(
+            [
+                f"Grid Definition File == {nc_path}",
+                f"  Grid Definition Variables == {xvar}, {yvar}, depth",
+                "  Grid Definition Label == ocean",
+                "  Boundary Gridmap == 1",
+                "End Grid",
+                "",
+            ]
+        )
+
+    def write_boundary_conditions(self, lines: List[str], nc_path: str):
+        """Write boundary conditions block.
+
+        Args:
+            lines: List to append lines to
+            nc_path: Path to the NetCDF file
+        """
+        vlist = ",".join(self.var_order)
+        lines.extend(
+            [
+                f"BC == OBC_GRID, ocean, {nc_path}",
+                f"  BC Nodestrings == #  ! Please supply open boundary ns list",
+                "  Sub-type == 6",
+                f"  BC Header == {self.time_var},{vlist}",
+                "  BC Update dt == 900.",
+                "  BC Time Units == hours",
+                "  BC Reference Time == 01/01/1990 00:00",
+                "  BC Offset == -0.0, 0.0, 0.0, 0.0, 0.0   ! Check Offset -0.0",
+                "  BC Default == NaN",
+                "  Vertical Coordinate Type == depth",
+                "End BC",
+                "",
+            ]
+        )
+
+    def generate(
+        self,
+        ds: xr.Dataset,
+        nc_path: str = "ocean_forcing.nc",
+    ) -> List[str]:
+        """Generate FVC configuration content for ocean forcing.
+
+        Args:
+            ds: Input xarray Dataset containing ocean variables
+            nc_path: Path or filename for the NetCDF file (referenced in FVC content)
+
+        Returns:
+            List[str]: FVC configuration content as a list of strings
+        """
+        # Process dataset to detect coordinates and time settings
+        self.process_dataset(ds)
+
+        lines = []
+
+        # Write header
+        self.write_header(lines, "Ocean Dataset", ds)
+
+        # Write grid definition
+        self.write_grid_definition(lines, nc_path)
+
+        # Write boundary conditions
+        self.write_boundary_conditions(lines, nc_path)
+
+        return lines
+
+    def write_file(
+        self,
+        lines: List[str],
+        output_path: Union[str, Path],
+        filename: Optional[str] = None,
+    ) -> None:
+        """Write FVC content to a file.
+
+        Args:
+            lines: FVC configuration content
+            output_path: Directory to write the file to
+            filename: Optional filename (if None, will use 'ocean_forcing.fvc')
+        """
+        output_path = Path(output_path)
+        filename = filename or "ocean_forcing.fvc"
+        if not filename.endswith(".fvc"):
+            filename = filename.replace(".nc", ".fvc")
+
+        output_path.mkdir(parents=True, exist_ok=True)
+        with open(output_path / filename, "w") as f:
+            for line in lines:
+                f.write(line + "\n")
+
+
+def write_ocean_fvc(
+    ds: xr.Dataset,
+    nc_path: str = "ocean_forcing.nc",
+    output_path: Optional[Union[str, Path]] = None,
+    filename: Optional[str] = None,
+    var_order: Optional[List[str]] = None,
+    source: Optional[str] = None,
+    model: Optional[str] = None,
+    info_url: Optional[str] = None,
+) -> List[str]:
+    """Generate (and optionally write) TUFLOW-FV ocean forcing configuration.
+
+    Args:
+        ds: Input xarray Dataset containing ocean variables
+        nc_path: Path or filename for the NetCDF file (referenced in FVC content)
+        output_path: Optional path to write FVC file (if None, no file is written)
+        filename: Optional filename for FVC file (if None, derives from nc_path)
+        var_order: Optional list specifying order of variables
+        source: Optional source string for the FVC header
+        model: Optional model string for the FVC header
+        info_url: Optional URL for source information to be printed in FVC
+
+    Returns:
+        List[str]: FVC configuration content as a list of strings
+
+    Examples:
+        # Just generate FVC content
+        >>> lines = write_ocean_fvc(dataset, nc_path="forcing.nc")
+
+        # Generate and write to file
+        >>> lines = write_ocean_fvc(
+        ...     dataset,
+        ...     nc_path="forcing.nc",
+        ...     output_path="path/to/output"
+        ... )
+
+        # Generate with custom variable order
+        >>> lines = write_ocean_fvc(
+        ...     dataset,
+        ...     nc_path="forcing.nc",
+        ...     var_order=["water_temp", "salinity", "water_u", "water_v", "surf_el"]
+        ... )
+    """
+    writer = OceanFVCWriter(var_order=var_order, source=source, model=model, info_url=info_url)
+
+    lines = writer.generate(ds, nc_path)
+
+    if output_path is not None:
+        writer.write_file(lines, output_path, filename)
+
+    return lines