initial commit

mne/_fiff/matrix.py

@@ -0,0 +1,137 @@
+# Authors: The MNE-Python contributors.
+# License: BSD-3-Clause
+# Copyright the MNE-Python contributors.
+
+from ..utils import logger
+from .constants import FIFF
+from .tag import find_tag, has_tag
+from .write import (
+    end_block,
+    start_block,
+    write_float_matrix,
+    write_int,
+    write_name_list,
+)
+
+
+def _transpose_named_matrix(mat):
+    """Transpose mat inplace (no copy)."""
+    mat["nrow"], mat["ncol"] = mat["ncol"], mat["nrow"]
+    mat["row_names"], mat["col_names"] = mat["col_names"], mat["row_names"]
+    mat["data"] = mat["data"].T
+
+
+def _read_named_matrix(fid, node, matkind, indent="    ", transpose=False):
+    """Read named matrix from the given node.
+
+    Parameters
+    ----------
+    fid : file
+        The opened file descriptor.
+    node : dict
+        The node in the tree.
+    matkind : int
+        The type of matrix.
+    transpose : bool
+        If True, transpose the matrix. Default is False.
+    %(verbose)s
+
+    Returns
+    -------
+    mat: dict
+        The matrix data
+    """
+    #   Descend one level if necessary
+    if node["block"] != FIFF.FIFFB_MNE_NAMED_MATRIX:
+        for k in range(node["nchild"]):
+            if node["children"][k]["block"] == FIFF.FIFFB_MNE_NAMED_MATRIX:
+                if has_tag(node["children"][k], matkind):
+                    node = node["children"][k]
+                    break
+        else:
+            logger.info(
+                f"{indent}Desired named matrix (kind = {matkind}) not available"
+            )
+            return None
+    else:
+        if not has_tag(node, matkind):
+            logger.info(
+                f"{indent}Desired named matrix (kind = {matkind}) not available"
+            )
+            return None
+
+    #   Read everything we need
+    tag = find_tag(fid, node, matkind)
+    if tag is None:
+        raise ValueError("Matrix data missing")
+    else:
+        data = tag.data
+
+    nrow, ncol = data.shape
+    tag = find_tag(fid, node, FIFF.FIFF_MNE_NROW)
+    if tag is not None and tag.data != nrow:
+        raise ValueError(
+            "Number of rows in matrix data and FIFF_MNE_NROW tag do not match"
+        )
+
+    tag = find_tag(fid, node, FIFF.FIFF_MNE_NCOL)
+    if tag is not None and tag.data != ncol:
+        raise ValueError(
+            "Number of columns in matrix data and FIFF_MNE_NCOL tag do not match"
+        )
+
+    tag = find_tag(fid, node, FIFF.FIFF_MNE_ROW_NAMES)
+    row_names = tag.data.split(":") if tag is not None else []
+
+    tag = find_tag(fid, node, FIFF.FIFF_MNE_COL_NAMES)
+    col_names = tag.data.split(":") if tag is not None else []
+
+    mat = dict(
+        nrow=nrow, ncol=ncol, row_names=row_names, col_names=col_names, data=data
+    )
+    if transpose:
+        _transpose_named_matrix(mat)
+    return mat
+
+
+def write_named_matrix(fid, kind, mat):
+    """Write named matrix from the given node.
+
+    Parameters
+    ----------
+    fid : file
+        The opened file descriptor.
+    kind : int
+        The kind of the matrix.
+    matkind : int
+        The type of matrix.
+    """
+    # let's save ourselves from disaster
+    n_tot = mat["nrow"] * mat["ncol"]
+    if mat["data"].size != n_tot:
+        ratio = n_tot / float(mat["data"].size)
+        if n_tot < mat["data"].size and ratio > 0:
+            ratio = 1 / ratio
+        raise ValueError(
+            f"Cannot write matrix: row ({mat['nrow']}) and column ({mat['ncol']}) "
+            f"total element ({n_tot}) mismatch with data size ({mat['data'].size}), "
+            f"appears to be off by a factor of {ratio:g}x"
+        )
+    start_block(fid, FIFF.FIFFB_MNE_NAMED_MATRIX)
+    write_int(fid, FIFF.FIFF_MNE_NROW, mat["nrow"])
+    write_int(fid, FIFF.FIFF_MNE_NCOL, mat["ncol"])
+
+    if len(mat["row_names"]) > 0:
+        # let's prevent unintentional stupidity
+        if len(mat["row_names"]) != mat["nrow"]:
+            raise ValueError('len(mat["row_names"]) != mat["nrow"]')
+        write_name_list(fid, FIFF.FIFF_MNE_ROW_NAMES, mat["row_names"])
+
+    if len(mat["col_names"]) > 0:
+        # let's prevent unintentional stupidity
+        if len(mat["col_names"]) != mat["ncol"]:
+            raise ValueError('len(mat["col_names"]) != mat["ncol"]')
+        write_name_list(fid, FIFF.FIFF_MNE_COL_NAMES, mat["col_names"])
+
+    write_float_matrix(fid, kind, mat["data"])
+    end_block(fid, FIFF.FIFFB_MNE_NAMED_MATRIX)