针对pulse-transit的工具

dist/client/mne/_fiff/ctf_comp.py

@@ -0,0 +1,189 @@
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+from copy import deepcopy
+
+import numpy as np
+
+from ..utils import _pl, logger, verbose
+from .constants import FIFF
+from .matrix import _read_named_matrix, write_named_matrix
+from .tag import read_tag
+from .tree import dir_tree_find
+from .write import end_block, start_block, write_int
+
+
+def _add_kind(one):
+    """Convert CTF kind to MNE kind."""
+    if one["ctfkind"] == int("47314252", 16):
+        one["kind"] = 1
+    elif one["ctfkind"] == int("47324252", 16):
+        one["kind"] = 2
+    elif one["ctfkind"] == int("47334252", 16):
+        one["kind"] = 3
+    else:
+        one["kind"] = int(one["ctfkind"])
+
+
+def _calibrate_comp(
+    comp, chs, row_names, col_names, mult_keys=("range", "cal"), flip=False
+):
+    """Get row and column cals."""
+    ch_names = [c["ch_name"] for c in chs]
+    row_cals = np.zeros(len(row_names))
+    col_cals = np.zeros(len(col_names))
+    for names, cals, inv in zip(
+        (row_names, col_names), (row_cals, col_cals), (False, True)
+    ):
+        for ii in range(len(cals)):
+            p = ch_names.count(names[ii])
+            if p != 1:
+                raise RuntimeError(
+                    f"Channel {names[ii]} does not appear exactly once "
+                    f"in data, found {p:d} instance{_pl(p)}"
+                )
+            idx = ch_names.index(names[ii])
+            val = chs[idx][mult_keys[0]] * chs[idx][mult_keys[1]]
+            val = float(1.0 / val) if inv else float(val)
+            val = 1.0 / val if flip else val
+            cals[ii] = val
+    comp["rowcals"] = row_cals
+    comp["colcals"] = col_cals
+    comp["data"]["data"] = row_cals[:, None] * comp["data"]["data"] * col_cals[None, :]
+
+
+@verbose
+def read_ctf_comp(fid, node, chs, verbose=None):
+    """Read the CTF software compensation data from the given node.
+
+    Parameters
+    ----------
+    fid : file
+        The file descriptor.
+    node : dict
+        The node in the FIF tree.
+    chs : list
+        The list of channels from info['chs'] to match with
+        compensators that are read.
+    %(verbose)s
+
+    Returns
+    -------
+    compdata : list
+        The compensation data
+    """
+    return _read_ctf_comp(fid, node, chs, None)
+
+
+def _read_ctf_comp(fid, node, chs, ch_names_mapping):
+    """Read the CTF software compensation data from the given node.
+
+    Parameters
+    ----------
+    fid : file
+        The file descriptor.
+    node : dict
+        The node in the FIF tree.
+    chs : list
+        The list of channels from info['chs'] to match with
+        compensators that are read.
+    ch_names_mapping : dict | None
+        The channel renaming to use.
+    %(verbose)s
+
+    Returns
+    -------
+    compdata : list
+        The compensation data
+    """
+    from .meas_info import _rename_comps
+
+    ch_names_mapping = dict() if ch_names_mapping is None else ch_names_mapping
+    compdata = []
+    comps = dir_tree_find(node, FIFF.FIFFB_MNE_CTF_COMP_DATA)
+
+    for node in comps:
+        #   Read the data we need
+        mat = _read_named_matrix(fid, node, FIFF.FIFF_MNE_CTF_COMP_DATA)
+        for p in range(node["nent"]):
+            kind = node["directory"][p].kind
+            pos = node["directory"][p].pos
+            if kind == FIFF.FIFF_MNE_CTF_COMP_KIND:
+                tag = read_tag(fid, pos)
+                break
+        else:
+            raise Exception("Compensation type not found")
+
+        #   Get the compensation kind and map it to a simple number
+        one = dict(ctfkind=tag.data.item())
+        del tag
+        _add_kind(one)
+        for p in range(node["nent"]):
+            kind = node["directory"][p].kind
+            pos = node["directory"][p].pos
+            if kind == FIFF.FIFF_MNE_CTF_COMP_CALIBRATED:
+                tag = read_tag(fid, pos)
+                calibrated = tag.data
+                break
+        else:
+            calibrated = False
+
+        one["save_calibrated"] = bool(calibrated)
+        one["data"] = mat
+        _rename_comps([one], ch_names_mapping)
+        if not calibrated:
+            #   Calibrate...
+            _calibrate_comp(one, chs, mat["row_names"], mat["col_names"])
+        else:
+            one["rowcals"] = np.ones(mat["data"].shape[0], dtype=np.float64)
+            one["colcals"] = np.ones(mat["data"].shape[1], dtype=np.float64)
+
+        compdata.append(one)
+
+    if len(compdata) > 0:
+        logger.info(f"    Read {len(compdata)} compensation matrices")
+
+    return compdata
+
+
+###############################################################################
+# Writing
+
+
+def write_ctf_comp(fid, comps):
+    """Write the CTF compensation data into a fif file.
+
+    Parameters
+    ----------
+    fid : file
+        The open FIF file descriptor
+
+    comps : list
+        The compensation data to write
+    """
+    if len(comps) <= 0:
+        return
+
+    #  This is very simple in fact
+    start_block(fid, FIFF.FIFFB_MNE_CTF_COMP)
+    for comp in comps:
+        start_block(fid, FIFF.FIFFB_MNE_CTF_COMP_DATA)
+        #    Write the compensation kind
+        write_int(fid, FIFF.FIFF_MNE_CTF_COMP_KIND, comp["ctfkind"])
+        if comp.get("save_calibrated", False):
+            write_int(fid, FIFF.FIFF_MNE_CTF_COMP_CALIBRATED, comp["save_calibrated"])
+
+        if not comp.get("save_calibrated", True):
+            # Undo calibration
+            comp = deepcopy(comp)
+            data = (
+                (1.0 / comp["rowcals"][:, None])
+                * comp["data"]["data"]
+                * (1.0 / comp["colcals"][None, :])
+            )
+            comp["data"]["data"] = data
+        write_named_matrix(fid, FIFF.FIFF_MNE_CTF_COMP_DATA, comp["data"])
+        end_block(fid, FIFF.FIFFB_MNE_CTF_COMP_DATA)
+
+    end_block(fid, FIFF.FIFFB_MNE_CTF_COMP)