Source code for pywasp.lincom.fetch

"""Routines for calculating water fetch

LINCOM requires maps of fetch for calculating water roughness. Here we calculate this
using a fortran function, and include a utility function to create vector and raster
masks of water points based on a roughness raster map.
"""

import numpy as np
from windkit.xarray_structures.metadata import _update_var_attrs, _MAP_TYPE_ATTRS

from ..core import Rvea0330
from ..io.raster_map import raster_to_vector, _raster_to_wasp_format
from ..io.vector_map import _vector_to_wasp_format
from .._errors import _handle_fort_errors

FORT_FETCHMAP = Rvea0330.externals_mod.ffvec




def calculate_fetchmap(vector_mask, raster_mask, wind):
    """Calculates a fetchmap for a given domain and wind

    .. warning::
        This function is experimental and its signature may change.

    Parameters
    ----------
    vector_mask : geopandas.GeoDataFrame
        vector map of landmask (water (0) and land (1))
    raster_mask : xarray.DataArray or WaspRasterMap
        raster map of landmask (water (0) and land (1))
    wind : obj
        LINCOM wind object

    Returns
    -------
    xarray.DataArray
        raster map of the fetch
    """
    fetch_map = _raster_to_wasp_format(raster_mask)
    vector_mask = _vector_to_wasp_format(vector_mask)

    # Get bounding box
    xmin, ymin, xmax, ymax = fetch_map.bbox.bounds()

    # Replace values with those from
    fetch_map.values, ierr = FORT_FETCHMAP(
        vector_mask.gx,
        vector_mask.gy,
        fetch_map.values,
        fetch_map.nx,
        fetch_map.ny,
        xmin,
        ymin,
        xmax,
        ymax,
        wind.latitude,
        wind.type,
        wind.sector,
        wind.wind_speed,
        wind.wind_height,
        wind.z0,
        wind.west_east,
        wind.south_north,
    )
    if ierr != 0:
        _handle_fort_errors(95 if ierr == 1 else ierr)

    return fetch_map._to_raster_map()





def roughness_to_landmask(lc_map, lctable):
    """Convert a roughness map to a landmask map

    All classes with z0 = 0.0 will be converted to water (0) points, while all other
    roughness lengths will be considered land.

    .. warning::
        This function is experimental and its signature may change.

    Parameters
    ----------
    lc_map: xarray.DataArray
        Landcover raster map
    lctable: LandCoverTable
        Map from LandCoverID to roughnesses

    Returns
    -------
    raster_mask: pw.RasterMap
        RasterMap of water (0) and land (1) mask
    vector_mask: pw.VectorMap
        VectorMap of water (0) and land (1) mask
    """
    # Create copy of data filled with ones
    raster_mask = lc_map.copy()
    raster_mask.values = np.ones_like(lc_map)
    raster_mask.name = "landmask"
    raster_mask = _update_var_attrs(raster_mask, _MAP_TYPE_ATTRS)

    # Replace ids that have roughness equal to 0 (vectorized)
    water_ids = np.array([k for k, v in lctable.items() if v["z0"] == 0])
    raster_mask.values[np.isin(lc_map.values, water_ids)] = 0

    # Create vector mask from raster map mask
    vector_mask = raster_to_vector(raster_mask, polygons=False)

    return raster_mask, vector_mask