Source code for pycif.plugins.transforms.basic.clipcrop.forward

import copy

import xarray as xr
import numpy as np
from logging import debug
import os
from .....utils.path import init_dir




def forward(
        transform,
        inout_datastore,
        controlvect,
        obsvect,
        mapper,
        di,
        df,
        mode,
        runsubdir,
        workdir,
        onlyinit=False,
        **kwargs
):
    """Zero out grid cells inside (clip) or outside (crop) a lon/lat rectangle.

    Builds a boolean mask from the domain coordinates and
    ``transform.{xmin, xmax, ymin, ymax}``, then sets all masked cells to
    zero in both ``'spec'`` and (in TL mode) ``'incr'``.

    * ``mode = 'crop'``: mask = cells **inside** the rectangle → zeroed.
    * ``mode = 'clip'``: mask = cells **outside** the rectangle → zeroed.

    Args:
        transform (Plugin): clipcrop instance (carries ``mode``, ``xmin``,
            ``xmax``, ``ymin``, ``ymax``).
        inout_datastore (dict): mutable datastore.
        controlvect: unused.
        obsvect: unused.
        mapper (dict): transform mapper (provides domain coordinates).
        di (datetime): sub-simulation start date.
        df (datetime): sub-simulation end date.
        mode (str): ``'fwd'`` or ``'tl'``.
        runsubdir (str): unused.
        workdir (str): unused.
        onlyinit (bool): unused.
        **kwargs: unused.
    """
    ddi = min(di, df)

    xmod_in = inout_datastore["inputs"]
    xmod_out = inout_datastore["outputs"]

    keys = ["spec"]
    if mode == "tl":
        keys.append("incr")

    for trid in xmod_in:
        xmod_out[trid][di] = {k: 0 for k in keys}

        # Find pixels to clip/crop
        domain = mapper["inputs"][trid]["domain"]
        mask = (domain.zlon >= transform.xmin) & (domain.zlon <= transform.xmax) \
            & (domain.zlat >= transform.ymin) & (domain.zlat <= transform.ymax)
        if transform.mode == "crop":
            mask = ~mask

        # Apply values to outputs
        for k in keys:
            if k in xmod_in[trid][di]:
                xmod_out[trid][di][k] = copy.deepcopy(xmod_in[trid][di][k])

                xmod_out[trid][di][k].values[..., mask] = 0.