Analysis Tools

@trait_docs
class Delete(Operator):
    """Class to purge data from observations.

    This operator takes lists of shared, detdata, intervals and meta keys to delete from
    observations.

    """

    # Class traits

    API = Int(0, help="Internal interface version for this operator")

    meta = List([], help="List of Observation dictionary keys to delete")

    detdata = List([], help="List of Observation detdata keys to delete")

    shared = List([], help="List of Observation shared keys to delete")

    intervals = List(
        [],
        help="List of tuples of Observation intervals keys to delete",
    )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @function_timer
    def _exec(self, data, detectors=None, **kwargs):
        log = Logger.get()
        for ob in data.obs:
            for key in self.detdata:
                # This ignores non-existant keys
                del ob.detdata[key]
            for key in self.shared:
                # This ignores non-existant keys
                del ob.shared[key]
            for key in self.intervals:
                # This ignores non-existant keys
                del ob.intervals[key]
            for key in self.meta:
                try:
                    del ob[key]
                except KeyError:
                    pass
        return

    def _finalize(self, data, **kwargs):
        return None

    def _requires(self):
        # Although we could require nothing, since we are deleting keys only if they
        # exist, providing these as requirements allows us to catch dependency issues
        # in pipelines.
        req = dict()
        if self.meta is not None:
            req["meta"] = list(self.meta)
        if self.detdata is not None:
            req["detdata"] = list(self.detdata)
        if self.shared is not None:
            req["shared"] = list(self.shared)
        if self.intervals is not None:
            req["intervals"] = list(self.intervals)
        return req

    def _provides(self):
        return dict()

@trait_docs
class Reset(Operator):
    """Class to reset data from observations.

    This operator takes lists of shared, detdata, intervals, and meta keys to reset.
    Numerical data objects and arrays are set to zero.  String objects are set to an
    empty string.  Any object that defines a `clear()` method will have that called.
    Any object not matching those criteria will be set to None.  Since an IntervalList
    is not mutable, any specified intervals will simply be deleted.

    """

    # Class traits

    API = Int(0, help="Internal interface version for this operator")

    meta = List([], help="List of Observation dictionary keys to reset")

    detdata = List([], help="List of Observation detdata keys to reset")

    shared = List([], help="List of Observation shared keys to reset")

    intervals = List(
        [],
        help="List of tuples of Observation intervals keys to reset",
    )

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @function_timer
    def _exec(self, data, detectors=None, **kwargs):
        log = Logger.get()
        for ob in data.obs:
            if len(self.detdata) > 0:
                # Get the detectors we are using for this observation
                dets = ob.select_local_detectors(detectors)
                if len(dets) == 0:
                    # Nothing to do for this observation
                    continue
                for key in self.detdata:
                    for d in dets:
                        ob.detdata[key][d, :] = 0
            for key in self.shared:
                scomm = ob.shared[key].nodecomm
                if scomm is None:
                    # No MPI, just set to zero
                    ob.shared[key].data[:] = 0
                else:
                    # Only rank zero on each node resets
                    if scomm.rank == 0:
                        ob.shared[key]._flat[:] = 0
                    scomm.barrier()
            for key in self.intervals:
                # This ignores non-existant keys
                del ob.intervals[key]
            for key in self.meta:
                if isinstance(ob[key], np.ndarray):
                    # This is an array, set to zero
                    ob[key][:] = 0
                elif hasattr(ob[key], "clear"):
                    # This is some kind of container (list, dict, etc).  Clear it.
                    ob[key].clear()
                elif isinstance(ob[key], bool):
                    # Boolean scalar, set to False
                    ob[key] = False
                elif isinstance(ob[key], numbers.Number):
                    # This is a scalar numeric value
                    ob[key] = 0
                elif isinstance(ob[key], (str, bytes)):
                    # This is string-like
                    ob[key] = ""
                else:
                    # This is something else.  Set to None
                    ob[key] = None
        return

    def _finalize(self, data, **kwargs):
        return None

    def _requires(self):
        req = dict()
        if self.meta is not None:
            req["meta"] = list(self.meta)
        if self.detdata is not None:
            req["detdata"] = list(self.detdata)
        if self.shared is not None:
            req["shared"] = list(self.shared)
        if self.intervals is not None:
            req["intervals"] = list(self.intervals)
        return req

    def _provides(self):
        return dict()

@trait_docs
class Copy(Operator):
    """Class to copy data.

    This operator takes lists of shared, detdata, and meta keys to copy to a new
    location in each observation.

    Each list contains tuples specifying the input and output key names.

    """

    # Class traits

    API = Int(0, help="Internal interface version for this operator")

    meta = List([], help="List of tuples of Observation meta keys to copy")

    detdata = List([], help="List of tuples of Observation detdata keys to copy")

    shared = List([], help="List of tuples of Observation shared keys to copy")

    intervals = List(
        [],
        help="List of tuples of Observation intervals keys to copy",
    )

    @traitlets.validate("meta")
    def _check_meta(self, proposal):
        val = proposal["value"]
        for v in val:
            if not isinstance(v, (tuple, list)):
                raise traitlets.TraitError("trait should be a list of tuples")
            if len(v) != 2:
                raise traitlets.TraitError("key tuples should have 2 values")
            if not isinstance(v[0], str) or not isinstance(v[1], str):
                raise traitlets.TraitError("key tuples should have string values")
        return val

    @traitlets.validate("detdata")
    def _check_detdata(self, proposal):
        val = proposal["value"]
        for v in val:
            if not isinstance(v, (tuple, list)):
                raise traitlets.TraitError("trait should be a list of tuples")
            if len(v) != 2:
                raise traitlets.TraitError("key tuples should have 2 values")
            if not isinstance(v[0], str) or not isinstance(v[1], str):
                raise traitlets.TraitError("key tuples should have string values")
        return val

    @traitlets.validate("shared")
    def _check_shared(self, proposal):
        val = proposal["value"]
        for v in val:
            if not isinstance(v, (tuple, list)):
                raise traitlets.TraitError("trait should be a list of tuples")
            if len(v) != 2:
                raise traitlets.TraitError("key tuples should have 2 values")
            if not isinstance(v[0], str) or not isinstance(v[1], str):
                raise traitlets.TraitError("key tuples should have string values")
        return val

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @function_timer
    def _exec(self, data, detectors=None, **kwargs):
        log = Logger.get()
        for ob in data.obs:
            for in_key, out_key in self.meta:
                if out_key in ob:
                    # The key exists- issue a warning before overwriting.
                    msg = "Observation key {} already exists- overwriting".format(
                        out_key
                    )
                    log.warning(msg)
                ob[out_key] = ob[in_key]

            for in_key, out_key in self.shared:
                # Although this is an internal function, the input arguments come
                # from existing shared objects and so should already be valid.
                ob.shared.assign_mpishared(
                    out_key, ob.shared[in_key], ob.shared.comm_type(in_key)
                )

            if len(self.detdata) > 0:
                # Get the detectors we are using for this observation.
                # We copy the full set of detectors, even if they are flagged.
                dets = ob.select_local_detectors(
                    detectors,
                    flagmask=0,
                )
                if len(dets) == 0:
                    # Nothing to do for this observation
                    continue
                for in_key, out_key in self.detdata:
                    if in_key not in ob.detdata:
                        continue
                    if out_key in ob.detdata:
                        # The key exists- verify that dimensions / dtype match
                        in_dtype = ob.detdata[in_key].dtype
                        out_dtype = ob.detdata[out_key].dtype
                        if out_dtype != in_dtype:
                            msg = f"Cannot copy to existing detdata key {out_key}"
                            msg += f" with different dtype ({out_dtype}) != {in_dtype}"
                            log.error(msg)
                            raise RuntimeError(msg)
                        in_shape = ob.detdata[in_key].detector_shape
                        out_shape = ob.detdata[out_key].detector_shape
                        if out_shape != in_shape:
                            msg = f"Cannot copy to existing detdata key {out_key}"
                            msg += f" with different detector shape ({out_shape})"
                            msg += f" != {in_shape}"
                            log.error(msg)
                            raise RuntimeError(msg)
                        if ob.detdata[out_key].detectors != dets:
                            # The output has a different set of detectors.  Reallocate.
                            ob.detdata[out_key].change_detectors(dets)
                        # Copy units
                        ob.detdata[out_key].update_units(ob.detdata[in_key].units)
                    else:
                        sample_shape = None
                        shp = ob.detdata[in_key].detector_shape
                        if len(shp) > 1:
                            sample_shape = shp[1:]
                        ob.detdata.create(
                            out_key,
                            sample_shape=sample_shape,
                            dtype=ob.detdata[in_key].dtype,
                            detectors=dets,
                            units=ob.detdata[in_key].units,
                        )
                    # Copy detector data
                    for d in dets:
                        ob.detdata[out_key][d, :] = ob.detdata[in_key][d, :]
        return

    def _finalize(self, data, **kwargs):
        return None

    def _requires(self):
        req = dict()
        if self.meta is not None:
            req["meta"] = [x[0] for x in self.meta]
        if self.detdata is not None:
            req["detdata"] = [x[0] for x in self.detdata]
        if self.shared is not None:
            req["shared"] = [x[0] for x in self.shared]
        if self.intervals is not None:
            req["intervals"] = [x[0] for x in self.intervals]
        return req

    def _provides(self):
        prov = dict()
        if self.meta is not None:
            prov["meta"] = [x[1] for x in self.meta]
        if self.detdata is not None:
            prov["detdata"] = [x[1] for x in self.detdata]
        if self.shared is not None:
            prov["shared"] = [x[1] for x in self.shared]
        if self.intervals is not None:
            prov["intervals"] = [x[1] for x in self.intervals]
        return prov

@trait_docs
class Combine(Operator):
    """Arithmetic with detector data.

    Two detdata objects are combined element-wise using addition, subtraction,
    multiplication, or division.  The desired operation is specified by the "op"
    trait as a string.  The result is stored in the specified detdata object:

    result = first (op) second

    If the result name is the same as the first or second input, then this
    input will be overwritten.

    """

    # Class traits

    API = Int(0, help="Internal interface version for this operator")

    op = Unicode(
        None,
        allow_none=True,
        help="Operation on the timestreams: 'subtract', 'add', 'multiply', or 'divide'",
    )

    first = Unicode(None, allow_none=True, help="The first detdata object")

    second = Unicode(None, allow_none=True, help="The second detdata object")

    result = Unicode(None, allow_none=True, help="The resulting detdata object")

    @traitlets.validate("op")
    def _check_op(self, proposal):
        val = proposal["value"]
        if val is not None:
            if val not in ["add", "subtract", "multiply", "divide"]:
                raise traitlets.TraitError("op must be one of the 4 allowed strings")
        return val

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @function_timer
    def _exec(self, data, detectors=None, **kwargs):
        log = Logger.get()

        for check_name, check_val in [
            ("first", self.first),
            ("second", self.second),
            ("result", self.result),
            ("op", self.op),
        ]:
            if check_val is None:
                msg = f"The {check_name} trait must be set before calling exec"
                log.error(msg)
                raise RuntimeError(msg)

        for ob in data.obs:
            # Get the detectors we are using for this observation
            local_dets = ob.select_local_detectors(detectors)
            if len(local_dets) == 0:
                # Nothing to do for this observation
                continue
            if self.first not in ob.detdata:
                msg = f"The first detdata key '{self.first}' does not exist in"
                msg += f" observation {ob.name}, skipping"
                log.verbose(msg)
                continue
            if self.second not in ob.detdata:
                msg = f"The second detdata key '{self.first}' does not exist in"
                msg += f" observation {ob.name}, skipping"
                log.verbose(msg)
                continue

            first_units = ob.detdata[self.first].units
            second_units = ob.detdata[self.second].units

            # Operate on the intersection of detectors
            dets = list(
                sorted(
                    set.intersection(
                        set(ob.detdata[self.first].detectors),
                        set(ob.detdata[self.second].detectors),
                    )
                )
            )

            if self.result == self.first:
                result_units = first_units
                scale_first = 1.0
                scale_second = unit_conversion(second_units, result_units)
            elif self.result == self.second:
                result_units = second_units
                scale_first = unit_conversion(first_units, result_units)
                scale_second = 1.0
            else:
                # We are creating a new field for the output.  Use units of first field.
                result_units = first_units
                scale_first = 1.0
                scale_second = unit_conversion(second_units, result_units)
                exists = ob.detdata.ensure(
                    self.result,
                    sample_shape=ob.detdata[self.first].detector_shape[1:],
                    dtype=ob.detdata[self.first].dtype,
                    detectors=ob.detdata[self.first].detectors,
                    create_units=result_units,
                )
            if self.op == "add":
                for d in dets:
                    ob.detdata[self.result][d, :] = (
                        scale_first * ob.detdata[self.first][d, :]
                    ) + (scale_second * ob.detdata[self.second][d, :])
            elif self.op == "subtract":
                for d in dets:
                    ob.detdata[self.result][d, :] = (
                        scale_first * ob.detdata[self.first][d, :]
                    ) - (scale_second * ob.detdata[self.second][d, :])
            elif self.op == "multiply":
                for d in dets:
                    ob.detdata[self.result][d, :] = (
                        scale_first * ob.detdata[self.first][d, :]
                    ) * (scale_second * ob.detdata[self.second][d, :])
            elif self.op == "divide":
                for d in dets:
                    ob.detdata[self.result][d, :] = (
                        scale_first * ob.detdata[self.first][d, :]
                    ) / (scale_second * ob.detdata[self.second][d, :])

    def _finalize(self, data, **kwargs):
        return None

    def _requires(self):
        req = {"detdata": [self.first, self.second]}
        if self.result is not None:
            if (self.result != self.first) and (self.result != self.second):
                req["detdata"].append(self.result)
        return req

    def _provides(self):
        prov = {"detdata": list()}
        if (self.result != self.first) and (self.result != self.second):
            prov["detdata"].append(self.result)
        return prov

@trait_docs
class CalibrateDetectors(Operator):
    """Multiply detector data by calibration factors.

    The calibration factor can be specified as a fixed value for all detectors (for
    example when converting from raw ADC values) or per-detector values can be
    supplied as a dictionary within the observation metadata or a column of the
    focalplane data table.  Detectors which are missing in the calibration dictionary
    are flagged.

    This operator is frequently the first one applied to "raw" data loaded from disk.
    If the dtype of the input DetectorData is an integer type, it will be promoted
    to float64 before applying the calibration.

    The units of the DetectorData can be updated with the `cal_units` trait.

    """

    # Class traits

    API = Int(0, help="Internal interface version for this operator")

    cal_name = Unicode(
        "calibration", help="The observation or focalplane key containing the gains"
    )

    cal_value = Float(
        None,
        allow_none=True,
        help="Apply this constant value to all detectors.  Overrides `cal_name`",
    )

    cal_units = Unit(
        None,
        allow_none=True,
        help="Update the detector data units",
    )

    det_data = Unicode(
        defaults.det_data,
        help="Observation detdata key for data to calibrate",
    )

    det_mask = Int(
        defaults.det_mask_invalid,
        help="Bit mask value for per-detector flagging",
    )

    cal_mask = Int(
        defaults.det_mask_invalid,
        help="Bit mask to apply to detectors with no calibration information",
    )

    @traitlets.validate("det_mask")
    def _check_det_mask(self, proposal):
        check = proposal["value"]
        if check < 0:
            raise traitlets.TraitError("Det mask should be a positive integer")
        return check

    @traitlets.validate("cal_mask")
    def _check_cal_mask(self, proposal):
        check = proposal["value"]
        if check < 0:
            raise traitlets.TraitError("Calibration mask should be a positive integer")
        return check

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @function_timer
    def _exec(self, data, detectors=None, use_accel=None, **kwargs):
        log = Logger.get()

        for ob in data.obs:
            if self.det_data not in ob.detdata:
                continue

            dets = ob.select_local_detectors(detectors, flagmask=self.det_mask)
            if len(dets) == 0:
                continue

            fp = ob.telescope.focalplane
            if self.cal_value is not None:
                cal = {x: self.cal_value for x in dets}
            else:
                if self.cal_name in ob:
                    # Observation has a separate calibration table
                    cal = ob[self.cal_name]
                elif self.cal_name in fp.properties:
                    # Gains are in the focalplane database
                    cal = {}
                    for det in dets:
                        cal[det] = fp[det][self.cal_name]
                else:
                    msg = f"{ob.name}: Gains '{self.cal_name}' do not exist "
                    msg += "as a dictionary nor in the focalplane database"
                    raise RuntimeError(msg)

            # Check dtype of detector data
            input_dtype = ob.detdata[self.det_data].dtype
            if input_dtype == np.dtype(np.int32) or input_dtype == np.dtype(np.int64):
                # Create a new DetectorData object and copy the original.  Then delete
                # the original and move the new one into place.
                temp_name = f"{self.name}_{self.det_data}_TEMPORARY"
                if self.cal_units is None:
                    temp_units = ob.detdata[self.det_data].units
                else:
                    temp_units = self.cal_units
                ob.detdata.create(
                    temp_name,
                    sample_shape=ob.detdata[self.det_data].sample_shape,
                    dtype=np.float64,
                    detectors=ob.detdata[self.det_data].detectors,
                    units=temp_units,
                )
                for det in ob.detdata[self.det_data].detectors:
                    ob.detdata[temp_name][det] = ob.detdata[self.det_data][det].astype(
                        np.float64
                    )
                del ob.detdata[self.det_data]
                ob.detdata.rename(temp_name, self.det_data)
            else:
                # Just update units if needed
                if self.cal_units is not None:
                    ob.detdata[self.det_data].update_units(self.cal_units)

            # Process all detectors
            det_flags = dict(ob.local_detector_flags)
            for det in dets:
                if det not in cal:
                    # Flag this detector
                    det_flags[det] |= self.cal_mask
                    continue
                ob.detdata[self.det_data][det] *= cal[det]

            # Update flags
            ob.update_local_detector_flags(det_flags)

    def _finalize(self, data, **kwargs):
        return

    def _requires(self):
        req = {
            "detdata": [self.det_data],
        }
        return req

    def _provides(self):
        return {"detdata": [self.det_data]}

@trait_docs
class MemoryCounter(Operator):
    """Compute total memory used by Observations in a Data object.

    Every process group iterates over their observations and sums the total memory used
    by detector and shared data.  Metadata and interval lists are assumed to be
    negligible and are not counted.

    """

    # Class traits

    API = Int(0, help="Internal interface version for this operator")

    silent = Bool(
        False,
        help="If True, return the memory used but do not log the result",
    )

    prefix = Unicode("", help="Prefix for log messages")

    def __init__(self, **kwargs):
        self.total_bytes = 0
        self.sys_mem_str = ""
        super().__init__(**kwargs)

    @function_timer
    def _exec(self, data, detectors=None, **kwargs):
        for ob in data.obs:
            self.total_bytes += ob.memory_use()
        self.sys_mem_str = memreport(
            msg="(whole node)", comm=data.comm.comm_world, silent=True
        )
        return

    def _finalize(self, data, **kwargs):
        log = Logger.get()
        if not self.silent:
            total_gb = self.total_bytes / 2**30
            if data.comm.comm_group_rank is not None:
                total_gb = data.comm.comm_group_rank.allreduce(total_gb)
            if data.comm.world_rank == 0:
                msg = f"Total timestream memory use = {total_gb:.3f} GB"
                log.info(f"{self.prefix}:  {msg}")
                log.info(f"{self.prefix}:  {self.sys_mem_str}")
        total = self.total_bytes
        self.total_bytes = 0
        return total

    def _requires(self):
        return dict()

    def _provides(self):
        return dict()

@trait_docs
class Statistics(Operator):
    """Operator to measure and write out data statistics"""

    API = Int(0, help="Internal interface version for this operator")

    det_data = Unicode(defaults.det_data, help="Observation detdata key to analyze")

    det_mask = Int(
        defaults.det_mask_nonscience,
        help="Bit mask value for per-detector flagging",
    )

    det_flags = Unicode(
        defaults.det_flags,
        allow_none=True,
        help="Observation detdata key for flags to use",
    )

    det_flag_mask = Int(
        defaults.det_mask_nonscience,
        help="Bit mask value for detector sample flagging",
    )

    shared_flags = Unicode(
        defaults.shared_flags,
        allow_none=True,
        help="Observation shared key for telescope flags to use",
    )

    shared_flag_mask = Int(
        defaults.shared_mask_nonscience,
        help="Bit mask value for optional shared flagging",
    )

    view = Unicode(
        None, allow_none=True, help="Use this view of the data in all observations"
    )

    output_dir = Unicode(
        None,
        allow_none=True,
        help="If specified, write output data products to this directory",
    )

    @traitlets.validate("det_mask")
    def _check_det_mask(self, proposal):
        check = proposal["value"]
        if check < 0:
            raise traitlets.TraitError("Det mask should be a positive integer")
        return check

    @traitlets.validate("shared_flag_mask")
    def _check_shared_flag_mask(self, proposal):
        check = proposal["value"]
        if check < 0:
            raise traitlets.TraitError("Shared flag mask should be a positive integer")
        return check

    @traitlets.validate("det_flag_mask")
    def _check_det_flag_mask(self, proposal):
        check = proposal["value"]
        if check < 0:
            raise traitlets.TraitError("Det flag mask should be a positive integer")
        return check

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        return

    @function_timer
    def _exec(self, data, detectors=None, **kwargs):
        """Measure the statistics

        Args:
            data (toast.Data): The distributed data.

        """
        log = Logger.get()
        nstat = 3  # Variance, Skewness, Kurtosis

        if self.output_dir is not None:
            if not os.path.isdir(self.output_dir):
                os.makedirs(self.output_dir, exist_ok=True)

        for obs in data.obs:
            # NOTE:  We could use the session name / uid in the filename
            # too for easy sorting.
            if obs.name is None:
                fname_out = f"{self.name}_{obs.uid}.h5"
            else:
                fname_out = f"{self.name}_{obs.name}.h5"
            if self.output_dir is not None:
                fname_out = os.path.join(self.output_dir, fname_out)

            # Get the list of all detectors that are not cut
            obs_dets = obs.select_local_detectors(detectors, flagmask=self.det_mask)
            if obs.comm.group_size == 1:
                all_dets = obs_dets
            else:
                proc_dets = obs.comm.comm_group.gather(obs_dets, root=0)
                all_dets = None
                if obs.comm.group_rank == 0:
                    all_set = set()
                    for pdets in proc_dets:
                        for d in pdets:
                            all_set.add(d)
                    all_dets = list(sorted(all_set))
                all_dets = obs.comm.comm_group.bcast(all_dets, root=0)

            ndet = len(all_dets)
            hits = np.zeros([ndet], dtype=int)
            means = np.zeros([ndet], dtype=float)
            stats = np.zeros([nstat, ndet], dtype=float)

            views = obs.view[self.view]

            # Measure the mean separately to simplify the math
            for iview, view in enumerate(views):
                if view.start is None:
                    # This is a view of the whole obs
                    nsample = obs.n_local_samples
                else:
                    nsample = view.stop - view.start
                if self.shared_flags is not None:
                    shared_flags = views.shared[self.shared_flags][iview]
                    shared_mask = (shared_flags & self.shared_flag_mask) == 0
                else:
                    shared_mask = np.ones(nsample, dtype=bool)

                for det in obs_dets:
                    if self.det_flags is not None:
                        det_flags = views.detdata[self.det_flags][iview][det]
                        det_mask = (det_flags & self.det_flag_mask) == 0
                        mask = np.logical_and(shared_mask, det_mask)
                    else:
                        mask = shared_mask
                    ngood = np.sum(mask)
                    if ngood == 0:
                        continue
                    signal = views.detdata[self.det_data][iview][det]
                    good_signal = signal[mask].copy()
                    idet = all_dets.index(det)
                    # Valid samples
                    hits[idet] += ngood
                    # Mean
                    means[idet] += np.sum(good_signal)

            if obs.comm.comm_group is not None:
                hits = obs.comm.comm_group.allreduce(hits, op=MPI.SUM)
                means = obs.comm.comm_group.allreduce(means, op=MPI.SUM)

            good = hits != 0
            means[good] /= hits[good]

            # Now evaluate the moments

            for iview, view in enumerate(views):
                if view.start is None:
                    # This is a view of the whole obs
                    nsample = obs.n_local_samples
                else:
                    nsample = view.stop - view.start
                if self.shared_flags is not None:
                    shared_flags = views.shared[self.shared_flags][iview]
                    shared_mask = (shared_flags & self.shared_flag_mask) == 0
                else:
                    shared_mask = np.ones(nsample, dtype=bool)

                for det in obs_dets:
                    if self.det_flags is not None:
                        det_flags = views.detdata[self.det_flags][iview][det]
                        det_mask = (det_flags & self.det_flag_mask) == 0
                        mask = np.logical_and(shared_mask, det_mask)
                    else:
                        mask = shared_mask
                    ngood = np.sum(mask)
                    if ngood == 0:
                        continue
                    idet = all_dets.index(det)
                    signal = views.detdata[self.det_data][iview][det]
                    good_signal = signal[mask].copy() - means[idet]
                    # Variance
                    stats[0, idet] += np.sum(good_signal**2)
                    # Skewness
                    stats[1, idet] += np.sum(good_signal**3)
                    # Kurtosis
                    stats[2, idet] += np.sum(good_signal**4)

            if obs.comm.comm_group is not None:
                stats = obs.comm.comm_group.reduce(stats, op=MPI.SUM)

            if obs.comm.group_rank == 0:
                # Central moments
                m2 = stats[0]
                m3 = stats[1]
                m4 = stats[2]
                for m in m2, m3, m4:
                    m[good] /= hits[good]
                # Variance
                var = m2.copy()
                # Skewness
                skew = m3.copy()
                skew[good] /= m2[good] ** 1.5
                # Kurtosis
                kurt = m4.copy()
                kurt[good] /= m2[good] ** 2
                # Write the results
                with h5py.File(fname_out, "w") as fout:
                    fout.attrs["UID"] = obs.uid
                    if obs.name is not None:
                        fout.attrs["name"] = obs.name
                    fout.attrs["nsample"] = obs.n_all_samples
                    fout.create_dataset(
                        "detectors", data=all_dets, dtype=h5py.string_dtype()
                    )
                    fout["ngood"] = hits
                    fout["mean"] = means
                    fout["variance"] = var
                    fout["skewness"] = skew
                    fout["kurtosis"] = kurt
                log.debug(f"Wrote data statistics to {fname_out}")

        return

    def _finalize(self, data, **kwargs):
        return

    def _requires(self):
        req = {
            "meta": list(),
            "shared": list(),
            "detdata": [self.det_data],
            "intervals": list(),
        }
        if self.shared_flags is not None:
            req["shared"].append(self.shared_flags)
        if self.det_flags is not None:
            req["detdata"].append(self.det_flags)
        if self.view is not None:
            req["intervals"].append(self.view)
        return req

    def _provides(self):
        prov = {
            "meta": list(),
            "shared": list(),
            "detdata": list(),
        }
        return prov

@trait_docs
class Pipeline(Operator):
    """Class representing a sequence of Operators.

    This runs a list of other operators over sets of detectors (default is all
    detectors in one shot).  By default all observations are passed to each operator,
    but the `observation_key` and `observation_value` traits can be used to run the
    operators on only observations which have a matching key / value pair.

    """

    # Class traits

    API = Int(0, help="Internal interface version for this operator")

    operators = List([], help="List of Operator instances to run.")

    detector_sets = List(
        ["ALL"],
        help="List of detector sets.  ['ALL'] and ['SINGLE'] are also valid values.",
    )

    use_hybrid = Bool(
        True,
        help="Should the pipeline be allowed to use the GPU when it has some cpu-only operators.",
    )

    @traitlets.validate("detector_sets")
    def _check_detsets(self, proposal):
        detsets = proposal["value"]
        if len(detsets) == 0:
            msg = "detector_sets must be a list with at least one entry "
            msg += "('ALL' and 'SINGLE' are valid entries)"
            raise traitlets.TraitError(msg)
        for dset in detsets:
            if (dset != "ALL") and (dset != "SINGLE"):
                # Not a built-in name, must be an actual list of detectors
                if isinstance(dset, str) or len(dset) == 0:
                    raise traitlets.TraitError(
                        "A detector set must be a list of detectors or 'ALL' / 'SINGLE'"
                    )
                for d in dset:
                    if not isinstance(d, str):
                        raise traitlets.TraitError(
                            "Each element of a det set should be a detector name"
                        )
        return detsets

    @traitlets.validate("operators")
    def _check_operators(self, proposal):
        ops = proposal["value"]
        for op in ops:
            if not isinstance(op, Operator):
                raise traitlets.TraitError(
                    "operators must be a list of Operator instances or None"
                )
        return ops

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # keeps track of the data that is on device
        self._staged_data = None
        # keep track of the data that had to move back to host due to a cpu-only operator
        # (for display / debugging purposes)
        self._unstaged_data = None

    @function_timer
    def _exec(self, data, detectors=None, use_accel=None, **kwargs):
        log = Logger.get()
        pstr = f"Proc ({data.comm.world_rank}, {data.comm.group_rank})"

        if len(self.operators) == 0:
            log.debug_rank(
                "Pipeline has no operators, nothing to do", comm=data.comm.comm_world
            )
            return

        # If the calling code passed use_accel=True, we assume that it will move
        # the data for us.  Otherwise, if possible / allowed, use the accelerator and
        # deal with data movement ourselves.
        self._staged_data = None
        self._unstaged_data = None
        pipe_accel = self._pipe_accel(use_accel)

        if pipe_accel:
            # some of our operators support using the accelerator
            msg = f"{self} supports accelerators."
            log.verbose_rank(msg, comm=data.comm.comm_world)
            use_accel = True
            # keeps track of the data that is on device
            self._staged_data = SetDict(
                {
                    key: set()
                    for key in ["global", "meta", "detdata", "shared", "intervals"]
                }
            )
            # keep track of the data that had to move back from device
            # (for display / debugging purposes)
            self._unstaged_data = SetDict(
                {
                    key: set()
                    for key in ["global", "meta", "detdata", "shared", "intervals"]
                }
            )

        if len(data.obs) == 0:
            # No observations for this group
            msg = f"{self} data, group {data.comm.group} has no observations."
            log.verbose_rank(msg, comm=data.comm.comm_group)

        # Ensure that all operators with a detector mask are using the same
        # mask.
        det_mask = None
        for op in self.operators:
            if hasattr(op, "det_mask"):
                if det_mask is None:
                    det_mask = op.det_mask
                else:
                    if op.det_mask != det_mask:
                        msg = "All operators in a Pipeline which use a det_mask"
                        msg += f" must have the same mask value ({op.det_mask} != "
                        msg += f"{det_mask}) in ops {[x.name for x in self.operators]}"
                        log.error(msg)
                        raise RuntimeError(msg)
        if det_mask is None:
            det_mask = 0

        if len(self.detector_sets) == 1 and self.detector_sets[0] == "ALL":
            # Run the operators with all detectors at once
            for op in self.operators:
                self._exec_operator(
                    op,
                    data,
                    detectors=None,
                    pipe_accel=pipe_accel,
                )
        elif len(self.detector_sets) == 1 and self.detector_sets[0] == "SINGLE":
            # Get superset of detectors across all observations
            all_local_dets = data.all_local_detectors(
                selection=detectors, flagmask=det_mask
            )
            if len(all_local_dets) == 0:
                all_local_dets = [None]
            # Run operators one detector at a time
            for det in all_local_dets:
                msg = f"{pstr} {self} SINGLE detector {det}"
                log.verbose(msg)
                if det is None:
                    dets = []
                else:
                    dets = [det]
                for op in self.operators:
                    self._exec_operator(
                        op,
                        data,
                        detectors=dets,
                        pipe_accel=pipe_accel,
                    )
        else:
            # We have explicit detector sets
            det_check = set(detectors)
            for det_set in self.detector_sets:
                selected_set = det_set
                if detectors is not None:
                    selected_set = list()
                    for det in det_set:
                        if det in det_check:
                            selected_set.append(det)
                if len(selected_set) == 0:
                    # Nothing in this detector set is being used, skip it
                    continue
                msg = f"{pstr} {self} detector set {selected_set}"
                log.verbose(msg)
                for op in self.operators:
                    self._exec_operator(
                        op,
                        data,
                        detectors=selected_set,
                        pipe_accel=pipe_accel,
                    )

        # notify user of device->host data movements introduced by CPU operators
        if (self._unstaged_data is not None) and (not self._unstaged_data.is_empty()):
            cpu_ops = {
                op.__class__.__qualname__
                for op in self.operators
                if not op.supports_accel()
            }
            log.debug(
                f"{pstr} {self}, had to move {self._unstaged_data} back to host as {cpu_ops} do not support accel."
            )

    @function_timer
    def _exec_operator(self, op, data, detectors, pipe_accel):
        """Runs an operator, dealing with data movement to/from device if needed."""
        # For this operator, we run on the accelerator if the pipeline has some
        # operators enabled and if this operator supports it.
        run_accel = pipe_accel and op.supports_accel()

        log = Logger.get()
        msg = f"Proc ({data.comm.world_rank}, {data.comm.group_rank}) {self} "
        msg += f"calling operator '{op.name}' exec(accelerator={run_accel})"
        if detectors is None:
            msg += " with ALL dets"
        log.verbose(msg)

        # Ensures data is where it should be for this operator
        if self._staged_data is not None:
            requires = SetDict(op.requires())
            if run_accel:
                # This operator will use the accelerator, stage data
                msg = f"Proc ({data.comm.world_rank}, {data.comm.group_rank}) {self} "
                msg += f"BEFORE staged = {self._staged_data}, unstaged = {self._unstaged_data}"
                log.verbose(msg)
                requires -= self._staged_data
                msg = f"Proc ({data.comm.world_rank}, {data.comm.group_rank}) {self} "
                msg += f"Staging objects {requires}"
                log.verbose(msg)
                data.accel_create(requires)
                data.accel_update_device(requires)
                # Update our record of data on device
                self._unstaged_data -= requires
                self._staged_data |= requires
                self._staged_data |= op.provides()
                self._unstaged_data -= op.provides()
                msg = f"Proc ({data.comm.world_rank}, {data.comm.group_rank}) {self} "
                msg += f"AFTER staged = {self._staged_data}, unstaged = {self._unstaged_data}"
                log.verbose(msg)
            else:
                # This operator is running on the host, unstage data
                msg = f"Proc ({data.comm.world_rank}, {data.comm.group_rank}) {self} "
                msg += f"BEFORE staged = {self._staged_data}, unstaged = {self._unstaged_data}"
                log.verbose(msg)
                requires &= self._staged_data  # intersection
                msg = f"Proc ({data.comm.world_rank}, {data.comm.group_rank}) {self} "
                msg += f"Un-staging objects {requires}"
                log.verbose(msg)
                data.accel_update_host(requires)
                # Update our record of data on the device
                self._staged_data -= requires
                self._unstaged_data |= requires  # union
                self._unstaged_data |= op.provides()
                # lets operator decide if it wants to move data and operate on device by itself
                run_accel = None
                msg = f"Proc ({data.comm.world_rank}, {data.comm.group_rank}) {self} "
                msg += f"AFTER staged = {self._staged_data}, unstaged = {self._unstaged_data}"
                log.verbose(msg)
        # runs operator
        op.exec(data, detectors=detectors, use_accel=run_accel)

    @function_timer
    def _finalize(self, data, use_accel=None, **kwargs):
        # FIXME:  We need to clarify in documentation that if using the
        # accelerator in _finalize() to produce output products, these
        # outputs should remain on the device so that they can be copied
        # out at the end automatically.

        log = Logger.get()
        pstr = f"Proc ({data.comm.world_rank}, {data.comm.group_rank})"
        msg = f"{pstr} {self} finalize"
        log.verbose(msg)

        # Are we running on the accelerator?
        pipe_accel = self._pipe_accel(use_accel)

        # run finalize on all the operators in the pipeline
        # NOTE: this might produce some output products
        result = list()
        if self.operators is not None:
            for op in self.operators:
                # Did we set use_accel to true when running with this operator
                use_accel_op = pipe_accel and op.supports_accel()
                result.append(op.finalize(data, use_accel=use_accel_op, **kwargs))

        # get outputs back and clean up data
        # if we are in charge of the data movement
        if self._staged_data is not None:
            # get outputs back from device
            provides = SetDict(self.provides())
            provides &= self._staged_data  # intersection
            log.verbose(f"{pstr} {self} copying out accel data outputs: {provides}")
            data.accel_update_host(provides)
            # deleting all data on device
            log.verbose(f"{pstr} {self} deleting accel data: {self._staged_data}")
            data.accel_delete(self._staged_data)
            self._staged_data = None
            self._unstaged_data = None

        return result

    def _pipe_accel(self, use_accel):
        if (use_accel is None) and accel_enabled():
            # Only allows hybrid pipelines if the environement variable and pipeline agree to it
            # (they both default to True)
            use_hybrid = self.use_hybrid and use_hybrid_pipelines
            # can we run this pipelines on accelerator
            supports_accel = (
                self._supports_accel_partial() if use_hybrid else self._supports_accel()
            )
            return supports_accel
        else:
            return use_accel

    def _requires(self):
        """
        Work through the operator list in reverse order and prune intermediate products
        (that will be provided by a previous operator).
        """
        # constructs the union of the requires minus the provides (in reverse order)
        req = SetDict(
            {key: set() for key in ["global", "meta", "detdata", "shared", "intervals"]}
        )
        for op in reversed(self.operators):
            # remove provides first as there can be an overlap between provides and requires
            req -= op.provides()
            req |= op.requires()
        # converts into a dictionary of lists
        req = {k: list(v) for (k, v) in req.items()}
        return req

    def _provides(self):
        """
        Work through the operator list and prune intermediate products
        (that are be provided to an intermediate operator).
        FIXME could a final result also be used by an intermediate operator?
        """
        # constructs the union of the provides minus the requires
        prov = SetDict(
            {key: set() for key in ["global", "meta", "detdata", "shared", "intervals"]}
        )
        for op in self.operators:
            # remove requires first as there can be an overlap between provides and requires
            prov -= op.requires()
            prov |= op.provides()
        # converts into a dictionary of lists
        prov = {k: list(v) for (k, v) in prov.items()}
        return prov

    def _implementations(self):
        """
        Find implementations supported by all the operators
        """
        implementations = {
            ImplementationType.DEFAULT,
            ImplementationType.COMPILED,
            ImplementationType.NUMPY,
            ImplementationType.JAX,
        }
        for op in self.operators:
            implementations.intersection_update(op.implementations())
        return list(implementations)

    def _supports_accel(self):
        """
        Returns True if *all* the operators are accelerator compatible.
        """
        for op in self.operators:
            if not op.supports_accel():
                return False
        return True

    def _supports_accel_partial(self):
        """
        Returns True if *at least one* of the operators is accelerator compatible.
        """
        for op in self.operators:
            if op.supports_accel():
                return True
        return False

    def __str__(self):
        """
        Converts the pipeline into a human-readable string.
        """
        return f"Pipeline{[op.__class__.__qualname__ for op in self.operators]}"