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
):
    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.