Repository: hpc4cmb/toast (language: python)

Issue:
Extend the logger to support *args
I found that Python logger has additional arguments such that the string with not be formatted if the log level is higher: https://docs.python.org/3/library/logging.html#logging.debug

But then I found that the toast logger only supports the first argument of the Python logger.

MWE:

```python
from logging import getLogger

logger = getLogger('pantable')
logger.warning('abc: %s, %s', 'efg', 'hij')
# ok: abc: efg, hij
```

```python
from toast.utils import Logger

logger = Logger.get()
logger.warning('abc: %s, %s', 'efg', 'hij')
# not ok
```

The advantage of `*args` is that when the log level is not needed, the string is not constructed in the first place. So for example

```python
logger.debug('this dictionary is ok: %s', medium_sized_data)`
```

would not have constructed a long string if the log level is higher than DEBUG.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 4598d4da7..cff10a44d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,7 +11,10 @@
 # *.py[cod]
 
 # Compiled extensions
-# *.so
+*.so
+
+# Python egg
+toast_cmb.egg-info
 
 /examples/job_*
 /examples/data/
@@ -23,6 +26,7 @@
 
 # generated files
 src/libtoast/src/version.cpp
+src/libtoast/version.cpp
 
 # vim
 *.swp
@@ -36,3 +40,6 @@ src/libtoast/src/version.cpp
 # Jupyter notebooks
 .ipynb_checkpoints
 __pycache__
+
+# VSCode
+.vscode/
diff --git a/docs/op_processing.inc b/docs/op_processing.inc
index 1a53026c6..430423d24 100644
--- a/docs/op_processing.inc
+++ b/docs/op_processing.inc
@@ -26,6 +26,12 @@ This operator applies a set of gains to the timestreams:
 .. autoclass:: toast.tod.OpApplyGain
     :members:
 
+Crosstalk
+~~~~~~~~~
+
+.. automodule:: toast.tod.crosstalk
+    :members:
+
 Utilities
 ~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/docs/utils.rst b/docs/utils.rst
index f08d11c69..ab7f91e6b 100644
--- a/docs/utils.rst
+++ b/docs/utils.rst
@@ -12,3 +12,5 @@ TOAST contains a variety of utilities for controlling the runtime environment, l
 .. include:: utils_rng.inc
 
 .. include:: utils_vmath.inc
+
+.. include:: utils_fma.inc
diff --git a/docs/utils_fma.inc b/docs/utils_fma.inc
new file mode 100644
index 000000000..10dabb9dc
--- /dev/null
+++ b/docs/utils_fma.inc
@@ -0,0 +1,8 @@
+.. _utilsfma:
+
+Vector Operations
+-----------------
+
+The following function(s) is simple vector operations compiled both for speeding up and reduce memory use, comparing to their pure Python counterpart.
+
+.. autofunction:: toast.utils.inplace_weighted_sum
diff --git a/src/libtoast/CMakeLists.txt b/src/libtoast/CMakeLists.txt
index f708bb3fa..7ce12d14a 100644
--- a/src/libtoast/CMakeLists.txt
+++ b/src/libtoast/CMakeLists.txt
@@ -27,6 +27,7 @@ set(toast_SOURCES
     src/toast_math_rng.cpp
     src/toast_math_qarray.cpp
     src/toast_math_fft.cpp
+    src/toast_math_fma.cpp
     src/toast_math_healpix.cpp
     src/toast_map_cov.cpp
     src/toast_fod_psd.cpp
@@ -44,6 +45,7 @@ set(toast_SOURCES
     tests/toast_test_rng.cpp
     tests/toast_test_qarray.cpp
     tests/toast_test_fft.cpp
+    tests/toast_test_fma.cpp
     tests/toast_test_healpix.cpp
     tests/toast_test_cov.cpp
     tests/toast_test_polyfilter.cpp
diff --git a/src/libtoast/include/toast.hpp b/src/libtoast/include/toast.hpp
index 327839d10..3dad21019 100644
--- a/src/libtoast/include/toast.hpp
+++ b/src/libtoast/include/toast.hpp
@@ -13,6 +13,7 @@
 #include <toast/math_rng.hpp>
 #include <toast/math_qarray.hpp>
 #include <toast/math_fft.hpp>
+#include <toast/math_fma.hpp>
 #include <toast/math_healpix.hpp>
 #include <toast/fod_psd.hpp>
 #include <toast/map_pixels.hpp>
diff --git a/src/libtoast/include/toast/math_fma.hpp b/src/libtoast/include/toast/math_fma.hpp
new file mode 100644
index 000000000..0e2235851
--- /dev/null
+++ b/src/libtoast/include/toast/math_fma.hpp
@@ -0,0 +1,21 @@
+/*
+   Copyright (c) 2021 by the parties listed in the AUTHORS file.
+   All rights reserved.  Use of this source code is governed by
+   a BSD-style license that can be found in the LICENSE file.
+ */
+
+#ifndef TOAST_MATH_FMA_HPP
+#define TOAST_MATH_FMA_HPP
+
+
+namespace toast {
+void inplace_weighted_sum(
+    int const n_out,
+    int const n_weights,
+    double * const out,
+    double const * const weights,
+    double const * const * const arrays
+    );
+}
+
+#endif // ifndef TOAST_MATH_FMA_HPP
diff --git a/src/libtoast/src/toast_math_fma.cpp b/src/libtoast/src/toast_math_fma.cpp
new file mode 100644
index 000000000..f314bca4d
--- /dev/null
+++ b/src/libtoast/src/toast_math_fma.cpp
@@ -0,0 +1,60 @@
+
+// Copyright (c) 2021 by the parties listed in the AUTHORS file.
+// All rights reserved.  Use of this source code is governed by
+// a BSD-style license that can be found in the LICENSE file.
+
+#include <toast/sys_utils.hpp>
+#include <toast/math_fma.hpp>
+
+
+// be careful of incorrect result if the pragma is put outside
+void toast::inplace_weighted_sum(
+    int const n_out,
+    int const n_weights,
+    double * const out,
+    double const * const weights,
+    double const * const * const arrays
+    ) {
+    if (!toast::is_aligned(out)) {
+        for (int i = 0; i < n_weights; ++i) {
+            double const weight = weights[i];
+            double const * const array = arrays[i];
+            #pragma omp parallel for default(shared) schedule(static)
+            for (int j = 0; j < n_out; ++j) {
+                out[j] += weight * array[j];
+            }
+        }
+    } else {
+        for (int i = 0; i < n_weights; ++i) {
+            double const weight = weights[i];
+            double const * const array = arrays[i];
+            if (toast::is_aligned(array)) {
+                #pragma omp parallel for simd default(shared) schedule(static)
+                for (int j = 0; j < n_out; ++j) {
+                    out[j] += weight * array[j];
+                }
+            } else {
+                #pragma omp parallel for default(shared) schedule(static)
+                for (int j = 0; j < n_out; ++j) {
+                    out[j] += weight * array[j];
+                }
+            }
+        }
+    }
+}
+
+// the following produces correct results, but slower.
+// void toast::inplace_weighted_sum(
+//     int const n_out,
+//     int const n_weights,
+//     double * const out,
+//     double const * const weights,
+//     double const * const * const arrays
+//     ) {
+//     #pragma omp parallel for simd
+//     for (int j = 0; j < n_out; ++j) {
+//         for (int i = 0; i < n_weights; ++i) {
+//             out[j] += weights[i] * arrays[i][j];
+//         }
+//     }
+// }
diff --git a/src/toast/CMakeLists.txt b/src/toast/CMakeLists.txt
index 7c088d3dd..55ac12e2c 100644
--- a/src/toast/CMakeLists.txt
+++ b/src/toast/CMakeLists.txt
@@ -35,6 +35,7 @@ pybind11_add_module(_libtoast MODULE
     _libtoast_math_rng.cpp
     _libtoast_math_qarray.cpp
     _libtoast_math_fft.cpp
+    _libtoast_math_fma.cpp
     _libtoast_math_healpix.cpp
     _libtoast_fod_psd.cpp
     _libtoast_tod_filter.cpp
diff --git a/src/toast/_libtoast.cpp b/src/toast/_libtoast.cpp
index 4cb11c916..e449d072a 100644
--- a/src/toast/_libtoast.cpp
+++ b/src/toast/_libtoast.cpp
@@ -32,6 +32,7 @@ PYBIND11_MODULE(_libtoast, m) {
     init_math_qarray(m);
     init_math_healpix(m);
     init_math_fft(m);
+    init_math_fma(m);
     init_fod_psd(m);
     init_tod_filter(m);
     init_tod_pointing(m);
diff --git a/src/toast/_libtoast.hpp b/src/toast/_libtoast.hpp
index 40b465864..1833e026c 100644
--- a/src/toast/_libtoast.hpp
+++ b/src/toast/_libtoast.hpp
@@ -211,6 +211,7 @@ void init_math_rng(py::module & m);
 void init_math_qarray(py::module & m);
 void init_math_healpix(py::module & m);
 void init_math_fft(py::module & m);
+void init_math_fma(py::module & m);
 void init_fod_psd(py::module & m);
 void init_tod_filter(py::module & m);
 void init_tod_pointing(py::module & m);
diff --git a/src/toast/_libtoast_math_fma.cpp b/src/toast/_libtoast_math_fma.cpp
new file mode 100644
index 000000000..42d31a768
--- /dev/null
+++ b/src/toast/_libtoast_math_fma.cpp
@@ -0,0 +1,104 @@
+
+// Copyright (c) 2021 by the parties listed in the AUTHORS file.
+// All rights reserved.  Use of this source code is governed by
+// a BSD-style license that can be found in the LICENSE file.
+
+#include <_libtoast.hpp>
+
+
+void init_math_fma(py::module & m) {
+    m.def(
+        "inplace_weighted_sum",
+        [](py::buffer out, py::buffer weights, py::args a) {
+            py::buffer_info info_out = out.request();
+            py::buffer_info info_weights = weights.request();
+
+            if ((unsigned)info_weights.size != a.size()) {
+                auto log = toast::Logger::get();
+                std::ostringstream o;
+                o << "weights and args are of different sizes.";
+                log.error(o.str().c_str());
+                throw std::runtime_error(o.str().c_str());
+            }
+
+            if (info_out.format != py::format_descriptor <double>::format()) {
+                auto log = toast::Logger::get();
+                std::ostringstream o;
+                o << "Incompatible format: expecting out to be a double array!";
+                log.error(o.str().c_str());
+                throw std::runtime_error(o.str().c_str());
+            }
+            double * const out_raw = reinterpret_cast <double *> (info_out.ptr);
+
+            if (info_weights.format != py::format_descriptor <double>::format()) {
+                auto log = toast::Logger::get();
+                std::ostringstream o;
+                o << "Incompatible format: expecting weights to be a double array!";
+                log.error(o.str().c_str());
+                throw std::runtime_error(o.str().c_str());
+            }
+            double const * const weights_raw =
+                reinterpret_cast <double *> (info_weights.ptr);
+
+            double ** arrays = new double *[a.size()];
+            for (size_t i = 0; i < a.size(); ++i) {
+// Use raw Python API here to avoid an extra, intermediate incref on the tuple item:
+                py::handle array =
+                    PyTuple_GET_ITEM(a.ptr(), static_cast <py::ssize_t> (i));
+                py::buffer array_buffer = array.cast <py::buffer>();
+                py::buffer_info info_array = array_buffer.request();
+
+                if (info_array.size != info_out.size) {
+                    auto log = toast::Logger::get();
+                    std::ostringstream o;
+                    o << "An array from args does not have the same size with out.";
+                    log.error(o.str().c_str());
+                    throw std::runtime_error(o.str().c_str());
+                }
+
+                if (info_array.format != py::format_descriptor <double>::format()) {
+                    auto log = toast::Logger::get();
+                    std::ostringstream o;
+                    o << "Incompatible format: expecting all array from args to be a double array!";
+                    log.error(o.str().c_str());
+                    throw std::runtime_error(o.str().c_str());
+                }
+                arrays[i] = reinterpret_cast <double *> (info_array.ptr);
+            }
+
+            toast::inplace_weighted_sum(
+                info_out.size,
+                info_weights.size,
+                out_raw,
+                weights_raw,
+                arrays
+                );
+            delete[] arrays;
+        },
+        py::arg("out"),
+        py::arg(
+            "weights"
+            ),
+        R"(
+            Compute a weighted sum of float64 values.
+
+            The results are stored in the output buffer. To guarantee SIMD
+            vectorization, the input (args) and output arrays (out) should be aligned
+            (i.e. use an AlignedF64).
+
+            Args:
+                out (array_like):  1D array of float64 values.
+                weights (array_like):  1D array of float64 values.
+                *args (array_like):  1D array of float64 values.
+
+            Returns:
+                None
+
+            Effectively equivalent to:
+
+            >>> for weight, array in zip(weights, args):
+            >>>     out += weight * array
+        )");
+
+    return;
+}
diff --git a/src/toast/tod/CMakeLists.txt b/src/toast/tod/CMakeLists.txt
index 6eb83c2cf..58f1a2397 100644
--- a/src/toast/tod/CMakeLists.txt
+++ b/src/toast/tod/CMakeLists.txt
@@ -4,6 +4,7 @@
 install(FILES
     __init__.py
     applygain.py
+    crosstalk.py
     gainscrambler.py
     interval.py
     memorycounter.py
diff --git a/src/toast/tod/__init__.py b/src/toast/tod/__init__.py
index e74b964ba..35966bc94 100644
--- a/src/toast/tod/__init__.py
+++ b/src/toast/tod/__init__.py
@@ -39,6 +39,7 @@
 
 from .gainscrambler import OpGainScrambler
 from .applygain import OpApplyGain, write_calibration_file
+from .crosstalk import OpCrosstalk, SimpleCrosstalkMatrix
 
 from .memorycounter import OpMemoryCounter
 
diff --git a/src/toast/tod/crosstalk.py b/src/toast/tod/crosstalk.py
new file mode 100644
index 000000000..23125e8c4
--- /dev/null
+++ b/src/toast/tod/crosstalk.py
@@ -0,0 +1,622 @@
+# py37+
+# from __future__ import annotations
+
+"""Simulate crosstalk between detectors in ToD.
+
+A typical use of this is that in your pipeline script,
+
+1. Add `add_crosstalk_args(parser)` to your argparse parser to add cli args,
+2. Read & create `op_crosstalk` by adding
+
+
+    .. highlight:: python
+    .. code-block:: python
+
+        if args.crosstalk_matrix is not None:
+            op_crosstalk = OpCrosstalk.read(args)
+
+3. Apply the crosstalk matrix to your data after the ToD is prepared
+(e.g. signal & noise) by adding
+
+
+    .. highlight:: python
+    .. code-block:: python
+
+        if args.crosstalk_matrix is not None:
+            if comm.comm_world is not None:
+                comm.comm_world.barrier()
+            op_crosstalk.exec(data, "tot_signal")
+            if comm.comm_world is not None:
+                comm.comm_world.barrier()
+
+Note that you need to create your crosstalk matri(x|ces)
+in HDF5 container(s) beforehand. A simple class `SimpleCrosstalkMatrix`
+can assist you to write such file,
+with the simple conventions that
+it has the following datasets:
+
+1. names: ASCII names of your detectors, say of length ``n``
+
+2. data: ``n`` by ``n`` float64 array of your crosstalk matrix ``m``,
+where you expect
+
+    ToD_crosstalked = m @ ToD_original
+
+    (Assuming row-major as is standard in Python.)
+
+Note that each crosstalk matrix is dense,
+and when multiple matrices are supplied, they are effectively block-diagonal
+(where each given dense matrix is a block.)
+I.e. in order to provide a block-diagonal crosstalk matrix,
+each block (together with the names of the detectors in that block)
+should write to a seperate HDF5 file to avoid unnessary computation.
+"""
+
+import argparse
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import h5py
+import numpy as np
+
+from ..mpi import get_world
+from ..op import Operator
+from ..timing import GlobalTimers
+from ..utils import Logger, inplace_weighted_sum
+
+if TYPE_CHECKING:
+    from typing import List, Optional
+
+    from .dist import Data
+    from .mpi import Comm
+
+
+def add_crosstalk_args(parser: "argparse.ArgumentParser"):
+    """Add crosstalk args to argparse.ArgumentParser object.
+
+    Args:
+      parser: 'argparse.ArgumentParser': argparse' argument parser
+
+    Returns:
+        None
+    """
+    parser.add_argument(
+        "--crosstalk-matrix",
+        type=Path,
+        nargs="*",
+        required=False,
+        help="input path(s) to crosstalk matrix in HDF5 container.",
+    )
+
+
+class SimpleCrosstalkMatrix:
+    """A thin crosstalk matrix class.
+
+    Args:
+      names: detector names.
+      data: crosstalk matrix.
+    The length of `names` should match the dimension of `data`,
+    which should be a square array.
+
+    This is a simple container storing the crosstalk matrix,
+    but not generating it.
+    """
+
+    def __init__(
+        self,
+        names: "np.ndarray['S']",
+        data: "np.ndarray[np.float64]",
+    ):
+        self.names = names
+        self.data = data
+
+        # py37+: we keep the dataclass semantics here in anticipation to use it in the future
+        self.__post_init__()
+
+    def __post_init__(self):
+        logger = Logger.get()
+        try:
+            names = self.names = np.asarray(self.names, dtype="S")
+            data = self.data = np.asarray(self.data, dtype=np.float64)
+        except ValueError as e:
+            logger.critical("name and data has to be Numpy-array-like")
+            raise e
+        if data.ndim != 2:
+            raise TypeError("data should be 2-dimensional")
+        if names.ndim != 1:
+            raise TypeError("names should be 1-dimensional")
+        shape = data.shape
+        if not (names.size == shape[0] == shape[1]):
+            raise TypeError("The dimensions of names and/or data not matched.")
+
+    @property
+    def names_str(self) -> "List[str]":
+        """names in list of str."""
+        return [name.decode() for name in self.names]
+
+    @classmethod
+    def load(cls, path: "Path") -> "SimpleCrosstalkMatrix":
+        """Load from an HDF5 file.
+
+        Args:
+          path: 'Path': path to an HDF5 file containing the data and names.
+
+        Returns:
+            SimpleCrosstalkMatrix
+        """
+        with h5py.File(path, "r") as f:
+            names = f["names"][:]
+            data = f["data"][:]
+        return cls(names, data)
+
+    def dump(
+        self,
+        path: "Path",
+        *,
+        libver: "str" = "latest",
+        # h5py.File.create_dataset args
+        compression: "str" = "gzip",
+        shuffle: "bool" = True,
+        fletcher32: "bool" = True,
+        track_times: "bool" = False,
+        compression_opts: "int" = 9,
+        **kwargs,
+    ):
+        """Dump to an HDF5 file.
+
+        `libver` will be passed to h5py.File,
+        and the rest keyword arguments will be passed to `h5py.File.create_dataset`
+        with sensible defaults.
+
+        Args:
+          path: 'Path': output file path.
+
+        Returns:
+            None
+        """
+        with h5py.File(path, "w", libver=libver) as f:
+            f.create_dataset(
+                "names",
+                data=self.names,
+                compression=compression,
+                shuffle=shuffle,
+                fletcher32=fletcher32,
+                track_times=track_times,
+                compression_opts=compression_opts,
+                **kwargs,
+            )
+            f.create_dataset(
+                "data",
+                data=self.data,
+                compression=compression,
+                shuffle=shuffle,
+                fletcher32=fletcher32,
+                track_times=track_times,
+                compression_opts=compression_opts,
+                **kwargs,
+            )
+
+
+class OpCrosstalk(Operator):
+    """Operator that applies crosstalk matrix to detector ToDs.
+
+    Args:
+      n_crosstalk_matrices: total no. of crosstalk matrices
+      crosstalk_matrices: the crosstalk matrices.
+        In MPI case, this should holds only those matrices owned by a rank dictate
+        by the condition `i % world_procs == world_rank`
+      name: this name is used to save data in tod.cache, so better be unique from other cache
+
+    In a typical scenario, the classmethod `read` is used to create an object
+    instead of initiating directly. This classmethod guarantees the condition above holds.
+
+    Note that in the MPI case, the matrices are distributed in the world communicator, while the
+    `exec` method will uses the grid commuicator specified per ToD. This is because at the time
+    of initiating, we have no knowledge on how the data is going to be distributed yet,
+    nor how that is "partitioned" in relation with multiple crosstalk matrix files given.
+    These 2 processes (of finding the crosstalk matrix, and finding the ToD) are orthogonal though
+    so the operator is completely general.
+    """
+
+    def __init__(
+        self,
+        n_crosstalk_matrices: "int",
+        crosstalk_matrices: "List[SimpleCrosstalkMatrix]",
+        *,
+        name: "str" = "crosstalk",
+    ):
+        self.n_crosstalk_matrices = n_crosstalk_matrices
+        self.crosstalk_matrices = crosstalk_matrices
+        self.name = name
+
+        self.world_comm: "Optional[Comm]"
+        self.world_procs: "int"
+        self.world_rank: "int"
+        self.world_comm, self.world_procs, self.world_rank = get_world()
+
+    @property
+    def is_serial(self) -> "bool":
+        return self.world_procs == 1
+
+    def _get_crosstalk_matrix(self, i: "int") -> "SimpleCrosstalkMatrix":
+        """Get the i-th crosstalk matrix, used this with MPI only.
+
+        Args:
+          i: 'int': the i-th crosstalk matrix.
+
+        Returns:
+            SimpleCrosstalkMatrix
+        """
+        logger = Logger.get()
+
+        rank_owner = i % self.world_procs
+        # index of the i-th matrix in the local rank
+        idx = i // self.world_procs
+
+        if self.world_rank == rank_owner:
+            crosstalk_matrix = self.crosstalk_matrices[idx]
+
+            names = crosstalk_matrix.names
+            data = crosstalk_matrix.data
+
+            # cast to int for boardcasting
+            names_int = names.view(np.uint8)
+            # the data from HDF5 is already float64
+            # this is needed for comm.Bcast below
+            data = data.view(np.float64)
+
+            # prepare lengths for creating arrays
+            lengths = np.array([names.size, names.dtype.itemsize], dtype=np.int64)
+        else:
+            lengths = np.empty(2, dtype=np.int64)
+        self.world_comm.Bcast(lengths, root=rank_owner)
+        # logger.debug(f'crosstalk: Rank {self.world_rank} receives lengths {lengths}')
+
+        # broadcast arrays
+        if self.world_rank != rank_owner:
+            n = lengths[0]
+            name_len = lengths[1]
+            names_int = np.empty(n * name_len, dtype=np.uint8)
+            names = names_int.view(f"S{name_len}")
+            data = np.empty((n, n), dtype=np.float64)
+        self.world_comm.Bcast(names_int, root=rank_owner)
+        # logger.debug(f'crosstalk: Rank {self.world_rank} receives names {names}')
+        self.world_comm.Bcast(data, root=rank_owner)
+        # logger.debug(f'crosstalk: Rank {self.world_rank} receives data {data}')
+
+        logger.info(
+            f"Obtained the {i + 1}-th crosstalk matrix from world-rank {rank_owner}."
+        )
+        return (
+            crosstalk_matrix
+            if self.world_rank == rank_owner
+            else SimpleCrosstalkMatrix(names, data)
+        )
+
+    @classmethod
+    def read(
+        cls,
+        args: "argparse.Namespace",
+        *,
+        name: "str" = "crosstalk",
+    ) -> "OpCrosstalk":
+        """Read crosstalk matri(x|ces) from HDF5 file(s).
+
+        This holds only those matrices owned by a rank
+        dictate by the condition `i % world_procs == world_rank`.
+
+        Args:
+          args: 'argparse.Namespace': namespace object from argparse' parser
+          name: 'str':  (Default value = "crosstalk"): this name is used to
+            save data in tod.cache, so better be unique from other cache
+
+        Returns:
+
+        """
+        gt = GlobalTimers.get()
+        gt.start("OpCrosstalk_read")
+
+        _, world_procs, world_rank = get_world()
+
+        paths = args.crosstalk_matrix
+
+        N = len(paths)
+
+        crosstalk_matrices = [
+            SimpleCrosstalkMatrix.load(paths[i])
+            for i in range(world_rank, N, world_procs)
+        ]
+
+        gt.stop("OpCrosstalk_read")
+
+        return cls(N, crosstalk_matrices, name=name)
+
+    def _exec_serial(
+        self,
+        data: "Data",
+        signal_name: "str",
+    ):
+        """Apply crosstalk matrix on ToD in data serially."""
+        crosstalk_name = self.name
+        logger = Logger.get()
+
+        # loop over crosstalk matrices
+        for crosstalk_matrix in self.crosstalk_matrices:
+            names = crosstalk_matrix.names_str
+            names_set = set(names)
+            crosstalk_data = crosstalk_matrix.data
+            for obs in data.obs:
+                tod = obs["tod"]
+
+                detectors_set = set(tod.detectors)
+                if not (names_set & detectors_set):
+                    logger.info(
+                        f"Crosstalk: skipping tod {tod} as "
+                        "it does not include detectors from crosstalk matrix "
+                        f"with these detectors: {names}."
+                    )
+                    continue
+                elif not (names_set <= detectors_set):
+                    raise ValueError(
+                        f"Crosstalk: tod {tod} only include some detectors "
+                        "from the crosstalk matrix "
+                        f"with these detectors: {names}."
+                    )
+                del detectors_set
+
+                n_samples = tod.local_samples[1]
+
+                # mat-mul
+                # This follows the _exec_mpi mat-mul algorithm
+                # but not put them in a contiguous array and use real mat-mul @
+                # The advantage is to reduce memory use
+                # (if creating an intermediate contiguous array
+                # that would requires one more copy of tod than needed below)
+                # and perhaps served as a easier-to-understand version of _exec_mpi below
+                for name, row in zip(names, crosstalk_data):
+                    row_global_total = tod.cache.create(
+                        f"{crosstalk_name}_{name}",
+                        np.float64,
+                        (n_samples,),
+                    )
+                    tods_list = [
+                        tod.cache.reference(f"{signal_name}_{name_j}")
+                        for name_j in names
+                    ]
+                    inplace_weighted_sum(row_global_total, row, *tods_list)
+                for name in names:
+                    # overwrite it in-place
+                    # not using tod.cache.put as that will destroy and create
+                    tod.cache.reference(f"{signal_name}_{name}")[
+                        :
+                    ] = tod.cache.reference(f"{crosstalk_name}_{name}")
+                    tod.cache.destroy(f"{crosstalk_name}_{name}")
+
+    def _exec_mpi(
+        self,
+        data: "Data",
+        signal_name: "str",
+    ):
+        """Apply crosstalk matrix on ToD in data with MPI."""
+        crosstalk_name = self.name
+        logger = Logger.get()
+        gt = GlobalTimers.get()
+
+        # loop over crosstalk matrices
+        for idx_crosstalk_matrix in range(self.n_crosstalk_matrices):
+            crosstalk_matrix = self._get_crosstalk_matrix(idx_crosstalk_matrix)
+            names = crosstalk_matrix.names_str
+            names_set = set(names)
+            crosstalk_data = crosstalk_matrix.data
+            n = crosstalk_data.shape[0]
+            for obs_i, obs in enumerate(data.obs):
+                tod = obs["tod"]
+                comm = tod.grid_comm_col
+                procs = tod.grid_size[0]
+                rank = tod.grid_ranks[0]
+
+                gt.start(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}"
+                )
+                gt.start(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_1-check"
+                )
+
+                # all ranks need to check this as they need to perform the same action
+                detectors_set = set(tod.detectors)
+                if not (names_set & detectors_set):
+                    logger.info(
+                        f"Crosstalk: skipping tod {tod} as "
+                        "it does not include detectors from crosstalk matrix "
+                        f"with these detectors: {names}."
+                    )
+                    continue
+                elif not (names_set <= detectors_set):
+                    raise ValueError(
+                        f"Crosstalk: tod {tod} only include some detectors "
+                        f"from the crosstalk matrix with these detectors: {names}."
+                    )
+                del detectors_set
+
+                gt.stop(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_1-check"
+                )
+                gt.start(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_2-detector-lut"
+                )
+
+                n_samples = tod.local_samples[1]
+                local_crosstalk_dets_set = set(tod.local_dets) & names_set
+                n_local_dets = len(local_crosstalk_dets_set)
+
+                # this is easier to understand and shorter
+                # but uses allgather instead of the more efficient Allgather
+                # construct detector lookup table
+                # local_dets = tod.local_dets
+                # global_dets = comm.allgather(local_dets)
+                # det_lut = {}
+                # for i, dets in enumerate(global_dets):
+                #     for det in dets:
+                #         det_lut[det] = i
+                # log.debug(f'dets lookup table: {dets_lut}')
+
+                # construct det_lut, a lookup table to know which rank holds a detector
+                local_has_det = tod.cache.create(
+                    f"{crosstalk_name}_local_has_det_{rank}",
+                    np.uint8,
+                    (n,),
+                ).view(np.bool_)
+                for i, name in enumerate(names):
+                    if name in local_crosstalk_dets_set:
+                        local_has_det[i] = True
+
+                global_has_det = tod.cache.create(
+                    f"{crosstalk_name}_global_has_det_{rank}",
+                    np.uint8,
+                    (procs, n),
+                ).view(np.bool_)
+                comm.Allgather(local_has_det, global_has_det)
+
+                # np.testing.assert_array_equal(local_has_det, global_has_det[rank])
+                del local_has_det
+                tod.cache.destroy(f"{crosstalk_name}_local_has_det_{rank}")
+
+                det_lut = {}
+                for i in range(procs):
+                    for j in range(n):
+                        if global_has_det[i, j]:
+                            det_lut[names[j]] = i
+                del global_has_det, i, j
+                tod.cache.destroy(f"{crosstalk_name}_global_has_det_{rank}")
+
+                # logger.debug(f'Rank {rank} has detectors lookup table: {det_lut}')
+                # for name in local_crosstalk_dets_set:
+                #     if det_lut[name] != rank:
+                #         raise RuntimeError(
+                #             'Error in creating a lookup table '
+                #             f'from detector name to rank: {det_lut}'
+                #         )
+
+                gt.stop(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_2-detector-lut"
+                )
+                gt.start(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_3-mat-mul"
+                )
+
+                # mat-mul
+                row_local_total = tod.cache.create(
+                    f"{crosstalk_name}_row_local_total_{rank}",
+                    np.float64,
+                    (n_samples,),
+                )
+                if n_local_dets > 0:
+                    row_local_weights = tod.cache.create(
+                        f"{crosstalk_name}_row_local_weights_{rank}",
+                        np.float64,
+                        (n_local_dets,),
+                    )
+                    local_det_idxs = tod.cache.create(
+                        f"{crosstalk_name}_local_det_idxs_{rank}",
+                        np.int64,
+                        (n_local_dets,),
+                    )
+                for i, name in enumerate(local_crosstalk_dets_set):
+                    local_det_idxs[i] = names.index(name)
+                # row-loop
+                # * potentially the tod can have more detectors
+                # * than SimpleCrosstalkMatrix.names_str has
+                # * and they will be skipped
+                for name, row in zip(names, crosstalk_data):
+                    try:
+                        rank_owner = det_lut[name]
+                    # in the unlikely scenario this happened,
+                    # it is likely the detector LUT is incorrect.
+                    # check the code containing `det_lut` above
+                    except KeyError:
+                        raise RuntimeError(
+                            f"Detector look-up table cannot find rank-owner of {name}: det_lut = {det_lut}"
+                        )
+                    if n_local_dets > 0:
+                        row_local_total[:] = 0.0
+                        row_local_weights[:] = row[local_det_idxs]
+                        tods_list = [
+                            tod.cache.reference(
+                                f"{signal_name}_{names[local_det_idxs[i]]}"
+                            )
+                            for i in range(n_local_dets)
+                        ]
+                        inplace_weighted_sum(
+                            row_local_total,
+                            row_local_weights,
+                            *tods_list,
+                        )
+                    if rank == rank_owner:
+                        row_global_total = tod.cache.create(
+                            f"{crosstalk_name}_{name}",
+                            np.float64,
+                            (n_samples,),
+                        )
+                        comm.Reduce(
+                            row_local_total,
+                            row_global_total,
+                            root=rank_owner,
+                        )
+                        # it is reduced into tod.cache and
+                        # the python reference can be safely deleted
+                        del row_global_total
+                    else:
+                        comm.Reduce(row_local_total, None, root=rank_owner)
+                del det_lut, row_local_total, name, row, rank_owner
+                tod.cache.destroy(f"{crosstalk_name}_row_local_total_{rank}")
+                if n_local_dets > 0:
+                    del row_local_weights, local_det_idxs, tods_list
+                    tod.cache.destroy(f"{crosstalk_name}_row_local_weights_{rank}")
+                    tod.cache.destroy(f"{crosstalk_name}_local_det_idxs_{rank}")
+
+                gt.stop(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_3-mat-mul"
+                )
+                gt.start(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_4-copy-inplace"
+                )
+
+                for name in local_crosstalk_dets_set:
+                    # overwrite it in-place
+                    # not using tod.cache.put as that will destroy and create
+                    tod.cache.reference(f"{signal_name}_{name}")[
+                        :
+                    ] = tod.cache.reference(f"{crosstalk_name}_{name}")
+                    tod.cache.destroy(f"{crosstalk_name}_{name}")
+
+                gt.stop(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}_4-copy-inplace"
+                )
+                gt.stop(
+                    f"OpCrosstalk_matrix-{idx_crosstalk_matrix}_observation-{obs_i}"
+                )
+
+    def exec(
+        self,
+        data: "Data",
+        signal_name: "str",
+    ):
+        """Apply crosstalk matrix on ToD in data.
+
+        It is dispatched depending if MPI is used.
+
+        Args:
+          data: 'Data': the data to be exec on.
+          signal_name: 'str': the name of the signal.
+
+        Returns:
+
+        """
+        gt = GlobalTimers.get()
+        gt.start("OpCrosstalk_exec")
+        (
+            self._exec_serial(data, signal_name)
+            if self.is_serial
+            else self._exec_mpi(data, signal_name)
+        )
+        gt.stop("OpCrosstalk_exec")
diff --git a/src/toast/utils.py b/src/toast/utils.py
index 0d1b14041..85bf8e5e2 100644
--- a/src/toast/utils.py
+++ b/src/toast/utils.py
@@ -36,6 +36,8 @@
     vfast_erfinv,
 )
 
+from ._libtoast import inplace_weighted_sum
+
 from .mpi import MPI, use_mpi
 
 


Test patch:
diff --git a/src/libtoast/tests/toast_test.hpp b/src/libtoast/tests/toast_test.hpp
index 4638abfdf..967109a92 100644
--- a/src/libtoast/tests/toast_test.hpp
+++ b/src/libtoast/tests/toast_test.hpp
@@ -149,6 +149,26 @@ class TOASTfftTest : public ::testing::Test {
 };
 
 
+class TOASTfmaTest : public ::testing::Test {
+    public:
+
+        TOASTfmaTest() {}
+
+        ~TOASTfmaTest() {}
+
+        virtual void SetUp();
+
+        virtual void TearDown() {}
+
+        static const int n_out;
+        static const int n_weights;
+        toast::AlignedVector <double> weights;
+        toast::AlignedVector <double> answer;
+        toast::AlignedVector <double> arrays_flat;
+        toast::AlignedVector <double> out;
+};
+
+
 class TOASTcovTest : public ::testing::Test {
     public:
 
diff --git a/src/libtoast/tests/toast_test_fma.cpp b/src/libtoast/tests/toast_test_fma.cpp
new file mode 100644
index 000000000..2ed22badd
--- /dev/null
+++ b/src/libtoast/tests/toast_test_fma.cpp
@@ -0,0 +1,51 @@
+
+// Copyright (c) 2021 by the parties listed in the AUTHORS file.
+// All rights reserved.  Use of this source code is governed by
+// a BSD-style license that can be found in the LICENSE file.
+
+#include <toast_test.hpp>
+
+
+const int TOASTfmaTest::n_out = 128;
+const int TOASTfmaTest::n_weights = 64;
+
+
+void TOASTfmaTest::SetUp() {
+    weights.resize(n_weights);
+    answer.resize(n_out);
+    arrays_flat.resize(n_weights * n_out);
+
+    // create some inputs that are different for different i, j
+    // and also the final result would be integer
+    for (int i = 0; i < n_weights; ++i) {
+        weights[i] = (double)12 * (i + 1);
+    }
+    for (int i = 0; i < n_weights; ++i) {
+        for (int j = 0; j < n_out; ++j) {
+            arrays_flat[i * n_out + j] = (double)(i + 1) * (i + j + 1) * (j + 1);
+        }
+    }
+
+    // and calculate the analytically expected value and compare later
+    const int n = n_weights;
+    for (int j = 0; j < n_out; ++j) {
+        answer[j] = (double)(1 + j) * n * (1 + n) * (3 * n * (1 + n) + j * (2 + 4 * n));
+    }
+    return;
+}
+
+TEST_F(TOASTfmaTest, inplace_weighted_sum) {
+    // not in SetUp: this should be done just before passing to the function below that
+    // modifies it
+    out.assign(n_out, 0);
+    double ** arrays = new double *[n_weights];
+    for (int i = 0; i < n_weights; ++i) {
+        arrays[i] = &arrays_flat[i * n_out];
+    }
+    toast::inplace_weighted_sum(n_out, n_weights, out.data(), weights.data(), arrays);
+    delete[] arrays;
+
+    for (int j = 0; j < n_out; ++j) {
+        EXPECT_DOUBLE_EQ(out[j], answer[j]);
+    }
+}
diff --git a/src/toast/tests/CMakeLists.txt b/src/toast/tests/CMakeLists.txt
index 6ad9a6fed..ba6c92bb3 100644
--- a/src/toast/tests/CMakeLists.txt
+++ b/src/toast/tests/CMakeLists.txt
@@ -18,6 +18,7 @@ install(FILES
     intervals.py
     tod_satellite.py
     ops_applygain.py
+    ops_crosstalk.py
     ops_simnoise.py
     cov.py
     ops_pmat.py
diff --git a/src/toast/tests/ops_crosstalk.py b/src/toast/tests/ops_crosstalk.py
new file mode 100644
index 000000000..14c4e9df8
--- /dev/null
+++ b/src/toast/tests/ops_crosstalk.py
@@ -0,0 +1,232 @@
+# py37+
+# from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+from numpy.random import default_rng
+
+from ..dist import Data
+from ..mpi import get_world
+from ..tod import TODCache
+from ..tod.crosstalk import OpCrosstalk, SimpleCrosstalkMatrix
+from ..utils import Logger
+from ._helpers import create_outdir
+from .mpi import MPITestCase
+
+if TYPE_CHECKING:
+    from typing import List, Optional
+
+    from .mpi import Comm
+
+
+class FakeData(Data):
+    def __init__(
+        self,
+        obs: "List[dict]",
+    ):
+        self.obs = obs
+
+
+class OpCrosstalkTest(MPITestCase):
+    def setUp(self):
+        fixture_name: "str" = os.path.splitext(os.path.basename(__file__))[0]
+        self.outdir: "Path" = Path(create_outdir(self.comm, fixture_name))
+
+        self.world_comm: "Optional[Comm]"
+        self.world_procs: "int"
+        self.world_rank: "int"
+        self.world_comm, self.world_procs, self.world_rank = get_world()
+
+        logger = Logger.get()
+
+        # just to have more than 1 detectors per self.world_procs, and is irregular
+        n_detectors: "int" = 4 * self.world_procs + 3
+        self.n_samples: "int" = 100
+        self.signal_name: "str" = "signal"
+
+        self.detranks: "List[int]" = [1]
+        for i in range(2, self.world_procs + 1):
+            if self.world_procs % i == 0:
+                self.detranks.append(i)
+        logger.info(
+            f'Testing OpCrosstalk against detranks {", ".join(map(str, self.detranks))}.'
+        )
+
+        self.names_strs: "List[List[str]]" = []
+        self.names: 'List[np.ndarray["S"]]' = []
+        self.tod: "List[np.ndarray[np.float64]]" = []
+        self.data: "List[np.ndarray[np.float64]]" = []
+        self.tod_crosstalked: "List[np.ndarray[np.float64]]" = []
+
+        temp = [f"A{i}" for i in range(n_detectors)]
+        self.names_strs.append(temp)
+        self.names.append(np.array(temp, dtype="S"))
+
+        temp = [f"B{i}" for i in range(n_detectors)]
+        self.names_strs.append(temp)
+        self.names.append(np.array(temp, dtype="S"))
+
+        x = np.arange(n_detectors * self.n_samples, dtype=np.float64).reshape(
+            (n_detectors, self.n_samples)
+        )
+        self.tod.append(x)
+
+        A = np.arange(n_detectors * n_detectors, dtype=np.float64).reshape(
+            (n_detectors, n_detectors)
+        )
+        self.data.append(A)
+
+        y = (
+            np.arange(n_detectors * n_detectors).reshape((n_detectors, n_detectors))
+            @ np.arange(n_detectors * self.n_samples).reshape(
+                (n_detectors, self.n_samples)
+            )
+        ).astype(np.float64)
+        self.tod_crosstalked.append(y)
+
+        rng = default_rng()
+        x = rng.standard_normal((n_detectors, self.n_samples))
+        self.tod.append(x)
+
+        A = np.identity(n_detectors) + np.reciprocal(
+            np.arange(10, 10 + n_detectors * n_detectors)
+        ).reshape((n_detectors, n_detectors))
+        self.data.append(A)
+
+        y = A @ x
+        self.tod_crosstalked.append(y)
+
+    def _each_op_crosstalk(
+        self,
+        tod_array: "np.ndarray[np.float64]",
+        crosstalk_data: "np.ndarray[np.float64]",
+        tod_crosstalked: "np.ndarray[np.float64]",
+        detranks: "int",
+    ):
+        names = self.names[0]
+        names_str = self.names_strs[0]
+
+        if self.world_rank == 0:
+            crosstalk_matrices = [SimpleCrosstalkMatrix(names, crosstalk_data)]
+        else:
+            crosstalk_matrices = []
+        op_crosstalk = OpCrosstalk(1, crosstalk_matrices)
+
+        tod = TODCache(self.world_comm, names_str, self.n_samples, detranks=detranks)
+
+        # write local detectors
+        start, length = tod.local_samples
+        local_dets_set = set(tod.local_dets)
+        for i, name in enumerate(names_str):
+            if name in local_dets_set:
+                tod.write(detector=name, data=tod_array[i, start : start + length])
+        data = FakeData(
+            [
+                {
+                    "tod": tod,
+                }
+            ]
+        )
+        op_crosstalk.exec(data, signal_name=self.signal_name)
+
+        # check output local detectors
+        # half float64 precision (float64 has 53-bit precison)
+        # will not produce bit-identical output as mat-mul
+        # is sensitive to the sum-ordering
+        for i, name in enumerate(names_str):
+            if name in local_dets_set:
+                output = tod.cache.reference(f"{self.signal_name}_{name}")
+                answer = tod_crosstalked[i, start : start + length]
+                np.testing.assert_allclose(answer, output)
+
+    def test_op_crosstalk(self):
+        for i, (tod, data, tod_crosstalked) in enumerate(
+            zip(self.tod, self.data, self.tod_crosstalked)
+        ):
+            for detranks in self.detranks:
+                with self.subTest(
+                    msg=f"OpCrosstalk Test case {i + 1} with detranks {detranks}"
+                ):
+                    self._each_op_crosstalk(
+                        tod,
+                        data,
+                        tod_crosstalked,
+                        detranks,
+                    )
+
+    def _each_op_crosstalk_multiple_matrices(
+        self,
+        names: "List[np.ndarray['S']]",
+        names_strs: "List[List[str]]",
+        tods_array: "List[np.ndarray[np.float64]]",
+        crosstalk_datas: "List[np.ndarray[np.float64]]",
+        tods_crosstalked: "List[np.ndarray[np.float64]]",
+        detranks: "int",
+    ):
+        n_crosstalk_matrices = len(names)
+        idxs_per_rank = range(self.world_rank, n_crosstalk_matrices, self.world_procs)
+        crosstalk_matrices = [
+            SimpleCrosstalkMatrix(names[i], crosstalk_datas[i]) for i in idxs_per_rank
+        ]
+        op_crosstalk = OpCrosstalk(n_crosstalk_matrices, crosstalk_matrices)
+
+        tod = TODCache(
+            self.world_comm, sum(names_strs, []), self.n_samples, detranks=detranks
+        )
+
+        # write local detectors
+        start, length = tod.local_samples
+        local_dets_set = set(tod.local_dets)
+        for names_str, tod_array in zip(names_strs, tods_array):
+            for i, name in enumerate(names_str):
+                if name in local_dets_set:
+                    tod.write(detector=name, data=tod_array[i, start : start + length])
+        data = FakeData(
+            [
+                {
+                    "tod": tod,
+                }
+            ]
+        )
+        op_crosstalk.exec(data, signal_name=self.signal_name)
+
+        # check output local detectors
+        # half float64 precision (float64 has 53-bit precison)
+        # will not produce bit-identical output as mat-mul
+        # is sensitive to the sum-ordering
+        for names_str, tod_crosstalked in zip(names_strs, tods_crosstalked):
+            for i, name in enumerate(names_str):
+                if name in local_dets_set:
+                    output = tod.cache.reference(f"{self.signal_name}_{name}")
+                    answer = tod_crosstalked[i, start : start + length]
+                    np.testing.assert_allclose(answer, output)
+
+    def test_op_crosstalk_multiple_matrices(self):
+        for detranks in self.detranks:
+            with self.subTest(
+                msg=f"Test OpCrosstalk with multiple matrices and detranks {detranks}"
+            ):
+                self._each_op_crosstalk_multiple_matrices(
+                    self.names,
+                    self.names_strs,
+                    self.tod,
+                    self.data,
+                    self.tod_crosstalked,
+                    detranks,
+                )
+
+    def test_op_crosstalk_io(self):
+        names = self.names[0]
+        crosstalk_data = self.data[0]
+        if self.world_rank == 0:
+            path = self.outdir / "simple_crosstalk_matrix.hdf5"
+            crosstalk_matrix = SimpleCrosstalkMatrix(names, crosstalk_data)
+
+            crosstalk_matrix.dump(path)
+            crosstalk_matrix_read = SimpleCrosstalkMatrix.load(path)
+
+            np.testing.assert_array_equal(names, crosstalk_matrix_read.names)
+            np.testing.assert_array_equal(crosstalk_data, crosstalk_matrix_read.data)
diff --git a/src/toast/tests/runner.py b/src/toast/tests/runner.py
index 1bd48a238..5bc64aa85 100644
--- a/src/toast/tests/runner.py
+++ b/src/toast/tests/runner.py
@@ -47,6 +47,8 @@
 from . import ops_mapmaker as testopsmapmaker
 from . import ops_filterbin as testopsfilterbin
 
+from . import ops_crosstalk as testopscrosstalk
+
 from . import map_satellite as testmapsatellite
 
 from . import map_ground as testmapground
@@ -138,6 +140,7 @@ def test(name=None, verbosity=2):
         suite.addTest(loader.loadTestsFromModule(testopsmadam))
         suite.addTest(loader.loadTestsFromModule(testopsmapmaker))
         suite.addTest(loader.loadTestsFromModule(testopsfilterbin))
+        suite.addTest(loader.loadTestsFromModule(testopscrosstalk))
         suite.addTest(loader.loadTestsFromModule(testmapsatellite))
         suite.addTest(loader.loadTestsFromModule(testmapground))
         suite.addTest(loader.loadTestsFromModule(testbinned))