Source code for pycif.plugins.obsvects.standard.rinvprod

from __future__ import annotations

import numpy as np
from numpy.typing import NDArray




def rinvprod(
    obsvect,
    dy: NDArray[np.floating],
    inverse: bool = True,
    mask: NDArray[np.bool_] | None = None,
) -> NDArray[np.floating]:
    # At the moment, this is valid for diagonal observation matrices only
    # TODO: Include these lines into rinvprod.py, with option for
    # non-diagonal matrices, eventually
    # err = np.where(obsvect.yobs_err == 0, np.nan, obsvect.yobs_err)

    dy_out = dy.copy()
    yobs = obsvect.yobs
    yobs_err = obsvect.yobs_err

    # Raise exception if some errors are 0 or NaN
    if np.any((yobs_err == 0) | np.isnan(yobs_err)):
        raise ValueError(
            "Your observation errors have some zero or NaN values. "
            "If will make the departures infinite or NaNs in the cost "
            "function. Please check your observation errors in your monitor files."
        )

    # Apply mask if any
    if mask is not None:
        if mask.shape != dy.shape:
            raise ValueError(
                f"Mask shape {mask.shape} does not match dy shape {dy.shape} in rinvprod"
            )

        dy_out = dy_out[mask]
        yobs_err = yobs_err[mask]
        yobs = yobs[mask]

    if inverse:
        dy_out = dy_out / yobs_err**2
    else:
        dy_out = dy_out * yobs_err + yobs

    # Filling masked values with zeros
    if mask is not None:
        dy_tmp = np.zeros_like(dy)
        dy_tmp[mask] = dy_out
        dy_out = dy_tmp

    return dy_out