Source code for pywasp.lincom.grid

"""Stores gridded topographic-data as an xarray object"""

import pyproj
import xarray as xr
import numpy as np
import windkit as wk

LINCOM_MAP_TYPE = {"elevation": 0, "roughness": 1, "fetch": 2}




def open_raster(mapfilename, crs=None):
    """Reads grid from raster file

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

    Parameters
    ----------
    mapfilename: Path
        Path to the rasterfile to be opened
    crs : int, dict, str, or pyproj.CRS
        Value to create pyproj.CRS object or an existing pyproj.CRS object.
        (Default: use the CRS from the file.)

    Returns
    -------
    xr.DataArray
        DataArray containing the topographic data of interest

    Notes
    ----
    Open a rasterio file and rename dimensions to match PyWAsP Requirements.
    Also ensure that the raster is ordered from lower-left not upper-left.
    """

    grid = xr.open_rasterio(mapfilename).squeeze().drop_vars("band")
    grid = grid.rename({"x": "west_east", "y": "south_north"})

    # Flip y-axis if negative transform
    if grid.attrs["transform"][4] < 0:
        grid = grid.sel(south_north=slice(None, None, -1))
        trans = list(grid.attrs["transform"])
        trans[4] *= -1
        trans[-1] = float(grid.south_north[0]) - 0.5 * trans[4]
        grid.attrs["transform"] = tuple(trans)

    # Add CRS to object using CF conventions
    if crs is None:
        if "crs" not in grid.attrs:
            err = "File does not have a CRS, so you must provide an crs."
            raise ValueError(err)
        else:
            crs = grid.attrs["crs"]
    else:
        # Check that file crs matches provided
        if "crs" in grid.attrs:
            if pyproj.CRS.from_string(grid.crs) != pyproj.CRS.from_user_input(crs):
                err = "Provided crs, does not match file CRS."
                raise ValueError(err)
            del grid.attr["crs"]
    if crs.startswith("+init="):
        crs = crs[6:]
    wk.spatial.set_crs(grid, crs)
    return grid



def _grid_from_wasp_raster_map(raster_map):
    """Convert a WAsP RasterMap object to a Lincom Gridmap xarray object

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

    Parameters
    ----------
    raster_map: pw.io.rastermap._WaspRasterMap
        Rastermap to be converted

    Returns
    -------
    xarray.DataArray
        DataArray containing the topographic map
    """
    # Create coordinate arrays
    x_points = np.arange(raster_map.nx) * raster_map.res + raster_map.minx
    west_east = xr.DataArray(x_points, dims=("west_east"))
    y_points = np.arange(raster_map.ny) * raster_map.res + raster_map.miny
    south_north = xr.DataArray(y_points, dims=("south_north"))

    out_da = xr.DataArray(
        data=np.reshape(raster_map.values, (raster_map.ny, raster_map.nx)),
        dims=("south_north", "west_east"),
        attrs={"res": (raster_map.res, raster_map.res)},
        coords={"west_east": west_east, "south_north": south_north},
    )

    wk.spatial.set_crs(out_da, raster_map.crs)
    return out_da