fmaussion / salem

Add geolocalised subsetting, masking, and plotting operations to xarray
http://salem.readthedocs.io
Other
186 stars 44 forks source link

Can geogrid simulator add 'd01','d02' labels to it? #214

Closed singledoggy closed 6 days ago

singledoggy commented 3 years ago

If salem has plan to add it in default or there is a simple way to do it?

fmaussion commented 2 years ago

yes, this sounds like a good idea. Since users have access to the maps and do the plot themselves, they could add it to the axes themselves, but arguably this could be made a bit easier in salem itself.

The way to go would be to add a call to m.set_text() after having added the geometry here:

https://github.com/fmaussion/salem/blob/0ae5885e4f466fc333a14544df851501f11102fd/salem/wrftools.py#L791-L794

singledoggy commented 2 years ago

Thanks for your work! And this solution is great, I can adjust the postion myself for now.

singledoggy commented 6 days ago
import matplotlib.pyplot as plt

from salem import lazy_property, wgs84, gis

def geogrid_simulator(fpath, do_maps=True, map_kwargs=None):
    """Emulates geogrid.exe, which is useful when defining new WRF domains.

    Parameters
    ----------
    fpath: str
       path to a namelist.wps file
    do_maps: bool
       if you want the simulator to return you maps of the grids as well
    map_kwargs: dict
       kwargs to pass to salem.Map()

    Returns
    -------
    (grids, maps) with:
        - grids: a list of Grids corresponding to the domains
          defined in the namelist
        - maps: a list of maps corresponding to the grids (if do_maps==True)
    """

    with open(fpath) as f:
        lines = f.readlines()

    pargs = dict()
    for l in lines:
        s = l.split("=")
        if len(s) < 2:
            continue
        s0 = s[0].strip().upper()
        s1 = list(filter(None, s[1].strip().replace("\n", "").split(",")))

        if s0 == "PARENT_ID":
            parent_id = [int(s) for s in s1]
        if s0 == "PARENT_GRID_RATIO":
            parent_ratio = [int(s) for s in s1]
        if s0 == "I_PARENT_START":
            i_parent_start = [int(s) for s in s1]
        if s0 == "J_PARENT_START":
            j_parent_start = [int(s) for s in s1]
        if s0 == "E_WE":
            e_we = [int(s) for s in s1]
        if s0 == "E_SN":
            e_sn = [int(s) for s in s1]
        if s0 == "DX":
            dx = float(s1[0])
        if s0 == "DY":
            dy = float(s1[0])
        if s0 == "MAP_PROJ":
            map_proj = s1[0].replace("'", "").strip().upper()
        if s0 == "REF_LAT":
            pargs["lat_0"] = float(s1[0])
        if s0 == "REF_LON":
            pargs["ref_lon"] = float(s1[0])
        if s0 == "TRUELAT1":
            pargs["lat_1"] = float(s1[0])
        if s0 == "TRUELAT2":
            pargs["lat_2"] = float(s1[0])
        if s0 == "STAND_LON":
            pargs["lon_0"] = float(s1[0])

    # Sometimes files are not complete
    pargs.setdefault("lon_0", pargs["ref_lon"])

    # define projection
    if map_proj == "LAMBERT":
        pwrf = (
            "+proj=lcc +lat_1={lat_1} +lat_2={lat_2} "
            "+lat_0={lat_0} +lon_0={lon_0} "
            "+x_0=0 +y_0=0 +a=6370000 +b=6370000"
        )
        pwrf = pwrf.format(**pargs)
    elif map_proj == "MERCATOR":
        pwrf = (
            "+proj=merc +lat_ts={lat_1} +lon_0={lon_0} "
            "+x_0=0 +y_0=0 +a=6370000 +b=6370000"
        )
        pwrf = pwrf.format(**pargs)
    elif map_proj == "POLAR":
        pwrf = (
            "+proj=stere +lat_ts={lat_1} +lat_0=90.0 +lon_0={lon_0} "
            "+x_0=0 +y_0=0 +a=6370000 +b=6370000"
        )
        pwrf = pwrf.format(**pargs)
    else:
        raise NotImplementedError(
            "WRF proj not implemented yet: " "{}".format(map_proj)
        )
    pwrf = gis.check_crs(pwrf)

    # get easting and northings from dom center (probably unnecessary here)
    e, n = gis.transform_proj(wgs84, pwrf, pargs["ref_lon"], pargs["lat_0"])

    # LL corner
    nx, ny = e_we[0] - 1, e_sn[0] - 1
    x0 = -(nx - 1) / 2.0 * dx + e  # -2 because of staggered grid
    y0 = -(ny - 1) / 2.0 * dy + n

    # parent grid
    grid = gis.Grid(nxny=(nx, ny), x0y0=(x0, y0), dxdy=(dx, dy), proj=pwrf)

    # child grids
    out = [grid]
    for ips, jps, pid, ratio, we, sn in zip(
        i_parent_start, j_parent_start, parent_id, parent_ratio, e_we, e_sn
    ):
        if ips == 1:
            continue
        ips -= 1
        jps -= 1
        we -= 1
        sn -= 1
        nx = we / ratio
        ny = sn / ratio
        if nx != (we / ratio):
            raise RuntimeError(
                "e_we and ratios are incompatible: "
                "(e_we - 1) / ratio must be integer!"
            )
        if ny != (sn / ratio):
            raise RuntimeError(
                "e_sn and ratios are incompatible: "
                "(e_sn - 1) / ratio must be integer!"
            )

        prevgrid = out[pid - 1]
        xx, yy = prevgrid.corner_grid.x_coord, prevgrid.corner_grid.y_coord
        dx = prevgrid.dx / ratio
        dy = prevgrid.dy / ratio
        grid = gis.Grid(
            nxny=(we, sn),
            x0y0=(xx[ips], yy[jps]),
            dxdy=(dx, dy),
            pixel_ref="corner",
            proj=pwrf,
        )
        out.append(grid.center_grid)

    maps = None
    if do_maps:
        from salem import Map
        import shapely.geometry as shpg

        if map_kwargs is None:
            map_kwargs = {}

        maps = []
        for i, g in enumerate(out):
            m = Map(g, **map_kwargs)
            left, right, bottom, top = g.extent
            x_offset = (right - left) * 0.05
            y_offset = -(top - bottom) * 0.05

            m.set_text(
                left + x_offset,
                top + y_offset,
                "D01",
                zorder=6,
                fontsize=11,
                fontweight="bold",
                crs=g.proj,
            )

            for j in range(i + 1, len(out)):
                cg = out[j]
                left, right, bottom, top = cg.extent

                s = np.array(
                    [(left, bottom), (right, bottom), (right, top), (left, top)]
                )
                l1 = shpg.LinearRing(s)
                m.set_geometry(
                    l1,
                    crs=cg.proj,
                    linewidth=(len(out) - j),
                    zorder=5,
                )
                x_offset = (right - left) * 0.05 * j
                y_offset = -(top - bottom) * 0.15 * j
                m.set_text(
                    left + x_offset,
                    top + y_offset,
                    f"D0{j+1}",
                    fontsize=11,
                    fontweight="bold",
                    zorder=6,
                    crs=cg.proj,
                )

            maps.append(m)

    return out, maps

For those still interested in this thread, the previous approach was rather rudimentary, but I recently had some time to make a simple update to the script. The offset ratio might need to be passed as a parameter.