Source code for pycif.plugins.obsoperators.standard.transforms.batch_computation

import copy
from logging import debug




def batch_computation(self, all_transforms, mapper, dask=False):
    """Adapt the transform pipeline for batch computation of Monte-Carlo samples.

    Traverses the forward period order in reverse and calls
    :meth:`Transform.mapper2batch` on each forward-direction transform to
    extend its input/output mapper so that a batch of *nsamples*
    perturbations can be computed simultaneously.

    Args:
        self (ObsOperator): the obs-operator plugin instance.  Must have
            ``self.batch_computation.nsamples``,
            ``self.batch_computation.dir_samples``, and
            ``self.batch_computation.file_samples`` set.
        all_transforms: the :class:`~pycif.utils.classes.transforms.Transform`
            object holding all transforms.
        mapper (dict): the pipeline mapper dictionary mapping transform IDs
            to their input/output/precursor/successor metadata; updated
            in-place.
        dask (bool): whether to use dask for parallel computation.
    """

    # Perturb transformations by propagating from the control/input vector
    nsamples = self.batch_computation.nsamples
    dir_samples = self.batch_computation.dir_samples
    file_samples = self.batch_computation.file_samples
    transform_pipe = self.transform_pipe
    perturbed_transforms = []
    for ddi, transform, direction in self.period_order_fwd[::-1]:
        if direction == "adjoint":
            continue

        debug(f"Perturbing {transform} for period {ddi}")
        transf_mapper = mapper[transform]
        transf = getattr(transform_pipe, transform)

        transf.dont_propagate = \
            getattr(self.batch_computation, "dont_propagate", [])

        transf.dont_propagate_obsvect = \
            getattr(self.batch_computation, "dont_propagate_obsvect", [])

        new_mapper = transf.mapper2batch(
            nsamples, transf_mapper, mapper, all_transforms,
            dir_samples, file_samples
        )
        mapper[transform] = new_mapper

        perturbed_transforms.append(transform)