docs/changes/latest.inc

@@ -22,9 +22,12 @@ Enhancements
 - Implemented SPM Auditory testing datagrabber X (:gh:`52` by `Fede Raimondo`_).
 
 - Created a the repository based on the mockup by by `Fede Raimondo`_.
+
 - Added comments to datalad grabber and changed to use datalad-clone instead of
   datalad-install (:gh: `55` by `Benjamin Poldrack`_).
 
+- Implement matrix storage in SQliteFeatureStorage (:gh:`42` by `Fede Raimondo`_).
+
 Bugs
 ~~~~
 

junifer/storage/base.py

@@ -156,6 +156,8 @@ class BaseFeatureStorage(ABC):
         meta: Dict,
         col_names: Optional[Iterable[str]] = None,
         row_names: Optional[Iterable[str]] = None,
+        kind: Optional[str] = "full",
+        diagonal: bool = True,
     ) -> None:
         """Store 2D matrix.
 
@@ -168,6 +170,15 @@ class BaseFeatureStorage(ABC):
             The column names (default None).
         row_names : list of tuple of str, optional
             The row names (default None).
+        kind : str, optional
+            The kind of matrix:
+            - 'triu': store upper triangular only.
+            - 'tril': store lower triangular.
+            - 'full': full matrix (default 'full').
+        diagonal : bool, optional
+            Whether to store the diagonal (default True).
+            If kind == 'full', setting this to false will raise
+            an error
 
         """
         raise_error(

junifer/storage/sqlite.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Union
 
+import numpy as np
 import pandas as pd
 from pandas.core.base import NoNewAttributesMixin
 from pandas.io.sql import pandasSQL_builder
@@ -77,7 +78,7 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             uri.parent.mkdir(parents=True, exist_ok=True)
         super().__init__(uri=uri, single_output=single_output, **kwargs)
         self._upsert = upsert
-        self._valid_inputs = ["table", "timeseries"]
+        self._valid_inputs = ["table", "timeseries", "matrix"]
 
     def get_engine(self, meta: Optional[Dict] = None) -> "Engine":
         """Get engine.
@@ -403,8 +404,10 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         self,
         data,
         meta: Dict,
-        col_names: Optional[Iterable[str]] = None,
-        rows_col_name: Optional[str] = None,
+        col_names: Optional[List[str]] = None,
+        row_names: Optional[List[str]] = None,
+        kind: Optional[str] = "full",
+        diagonal: bool = True,
     ) -> None:
         """Implement 2D matrix storing.
 
@@ -415,16 +418,77 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             The metadata as a dictionary.
         col_names : list or tuple of str, optional
             The column names (default None).
-        rows_col_name : str, optional
+        row_names : str, optional
             The column name to use in case number of rows greater than 1.
             If None and number of rows greater than 1, then the name will be
             "index" (default None).
-
+        kind : str, optional
+            The kind of matrix:
+            - 'triu: store upper triangular only.
+            - 'tril': store lower triangular.
+            - 'full': full matrix (default 'full').
+        diagonal : bool, optional
+            Whether to store the diagonal (default True).
+            If kind == 'full', setting this to false will raise
+            an error
         """
-        # Same as store_2d, but order is important
-        raise_error(
-            msg="store_matrix2d() not implemented", klass=NotImplementedError
-        )
+        if diagonal is False and kind not in ["triu", "tril"]:
+            raise_error(
+                msg="Diagonal cannot be False if kind is not full",
+                klass=ValueError,
+            )
+
+        if kind in ["triu", "tril"]:
+            if data.shape[0] != data.shape[1]:
+                raise_error(
+                    "Cannot store a non-square matrix as a triangular matrix",
+                    klass=ValueError,
+                )
+
+        n_rows = 1
+        # Convert element metadata to index
+        idx = element_to_index(meta=meta, n_rows=n_rows, rows_col_name=None)
+
+        if kind == "triu":
+            k = 0 if diagonal is True else 1
+            data_idx = np.triu_indices(data.shape[0], k=k)
+        elif kind == "tril":
+            k = 0 if diagonal is True else -1
+            data_idx = np.tril_indices(data.shape[0], k=k)
+        elif kind == "full":
+            data_idx = (
+                np.repeat(np.arange(data.shape[0]), data.shape[1]),
+                np.tile(np.arange(data.shape[1]), data.shape[0]),
+            )
+        else:
+            raise_error(msg=f"Invalid kind {kind}", klass=ValueError)
+        if row_names is None:
+            row_names = [f"r{i}" for i in range(data.shape[0])]
+        elif len(row_names) != data.shape[0]:
+            raise_error(
+                msg="Number of row names does not match number of rows",
+                klass=ValueError,
+            )
+
+        if col_names is None:
+            col_names = [f"c{i}" for i in range(data.shape[1])]
+        elif len(col_names) != data.shape[1]:
+            raise_error(
+                msg="Number of column names does not match number of columns",
+                klass=ValueError,
+            )
+
+        flat_data = data[data_idx]
+        columns = [
+            f"{row_names[i]}~{col_names[j]}"
+            for i, j in zip(data_idx[0], data_idx[1])
+        ]
+        # Prepare new dataframe
+        data_df = pd.DataFrame(
+            flat_data[None, :], columns=columns, index=idx
+        )  # type: ignore
+        # Store dataframe
+        self.store_df(df=data_df, meta=meta)
 
     # TODO: complete type annotations
     def store_table(

junifer/storage/tests/test_sqlite.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from typing import List, Union
 
 import numpy as np
+from numpy.testing import assert_array_equal
 import pandas as pd
 import pytest
 from pandas.testing import assert_frame_equal
@@ -405,6 +406,177 @@ def test_store_table(tmp_path: Path) -> None:
     assert_frame_equal(df_new, c_df_new)
 
 
+def test_store_matrix2d(tmp_path: Path) -> None:
+    """Test 2D Matrix store.
+
+    Parameters
+    ----------
+    tmp_path : pathlib.Path
+        The path to the test directory.
+
+    """
+    uri = tmp_path / "test_store_table.db"
+    storage = SQLiteFeatureStorage(uri=uri, single_output=True)
+    # Metadata to store
+    meta = {"element": "test", "version": "0.0.1", "marker": {"name": "fc"}}
+
+    # Store 4 x 3 full matrix
+    data = np.array(
+        [[1, 2, 3], [11, 22, 33], [111, 222, 333], [1111, 2222, 3333]]
+    )
+    row_names = ["row1", "row2", "row3", "row4"]
+    col_names = ["col1", "col2", "col3"]
+
+    # Store table
+    storage.store_matrix2d(
+        data, meta, row_names=row_names, col_names=col_names
+    )
+
+    stored_names = [f"{i}~{j}" for i in row_names for j in col_names]
+
+    features = storage.list_features()
+    feature_md5 = list(features.keys())[0]
+    assert "fc" == features[feature_md5]["name"]
+
+    read_df = storage.read_df(feature_md5=feature_md5)
+    assert read_df.shape == (1, 12)
+    assert_array_equal(read_df.values[0], data.flatten())
+    assert list(read_df.columns) == stored_names
+    # Store without row and column names
+    uri = tmp_path / "test_store_table_nonames.db"
+    storage = SQLiteFeatureStorage(uri=uri, single_output=True)
+    storage.store_matrix2d(data, meta)
+    stored_names = [
+        f"r{i}~c{j}"
+        for i in range(data.shape[0])
+        for j in range(data.shape[1])
+    ]
+    features = storage.list_features()
+    feature_md5 = list(features.keys())[0]
+    assert "fc" == features[feature_md5]["name"]
+    read_df = storage.read_df(feature_md5=feature_md5)
+    assert list(read_df.columns) == stored_names
+
+    with pytest.raises(ValueError, match="Invalid kind"):
+        storage.store_matrix2d(data, meta, kind="wrong")
+
+    with pytest.raises(ValueError, match="non-square"):
+        storage.store_matrix2d(data, meta, kind="triu")
+
+    with pytest.raises(ValueError, match="cannot be False"):
+        storage.store_matrix2d(data, meta, kind="full", diagonal=False)
+
+    # Store upper triangular matrix
+    data = np.array([[1, 2, 3], [11, 22, 33], [111, 222, 333]])
+    row_names = ["row1", "row2", "row3"]
+    col_names = ["col1", "col2", "col3"]
+    uri = tmp_path / "test_store_table_triu.db"
+    storage = SQLiteFeatureStorage(uri=uri, single_output=True)
+    storage.store_matrix2d(
+        data, meta, kind="triu", row_names=row_names, col_names=col_names
+    )
+
+    stored_names = [
+        "row1~col1",
+        "row1~col2",
+        "row1~col3",
+        "row2~col2",
+        "row2~col3",
+        "row3~col3",
+    ]
+
+    features = storage.list_features()
+    feature_md5 = list(features.keys())[0]
+    assert "fc" == features[feature_md5]["name"]
+    read_df = storage.read_df(feature_md5=feature_md5)
+    assert list(read_df.columns) == stored_names
+    assert_array_equal(
+        read_df.values, data[np.triu_indices(n=data.shape[0])][None, :]
+    )
+
+    # Store upper triangular matrix without diagonal
+    uri = tmp_path / "test_store_table_triu_nodiagonal.db"
+    storage = SQLiteFeatureStorage(uri=uri, single_output=True)
+    storage.store_matrix2d(
+        data,
+        meta,
+        kind="triu",
+        row_names=row_names,
+        col_names=col_names,
+        diagonal=False,
+    )
+
+    stored_names = [
+        "row1~col2",
+        "row1~col3",
+        "row2~col3",
+    ]
+
+    features = storage.list_features()
+    feature_md5 = list(features.keys())[0]
+    assert "fc" == features[feature_md5]["name"]
+    read_df = storage.read_df(feature_md5=feature_md5)
+    assert list(read_df.columns) == stored_names
+    assert_array_equal(
+        read_df.values, data[np.triu_indices(n=data.shape[0], k=1)][None, :]
+    )
+
+    # Store lower triangular matrix
+    data = np.array([[1, 2, 3], [11, 22, 33], [111, 222, 333]])
+    row_names = ["row1", "row2", "row3"]
+    col_names = ["col1", "col2", "col3"]
+    uri = tmp_path / "test_store_table_tril.db"
+    storage = SQLiteFeatureStorage(uri=uri, single_output=True)
+    storage.store_matrix2d(
+        data, meta, kind="tril", row_names=row_names, col_names=col_names
+    )
+
+    stored_names = [
+        "row1~col1",
+        "row2~col1",
+        "row2~col2",
+        "row3~col1",
+        "row3~col2",
+        "row3~col3",
+    ]
+
+    features = storage.list_features()
+    feature_md5 = list(features.keys())[0]
+    assert "fc" == features[feature_md5]["name"]
+    read_df = storage.read_df(feature_md5=feature_md5)
+    assert list(read_df.columns) == stored_names
+    assert_array_equal(
+        read_df.values, data[np.tril_indices(n=data.shape[0])][None, :]
+    )
+
+    # Store lower triangular matrix without diagonal
+    uri = tmp_path / "test_store_table_tril_nodiagonal.db"
+    storage = SQLiteFeatureStorage(uri=uri, single_output=True)
+    storage.store_matrix2d(
+        data,
+        meta,
+        kind="tril",
+        row_names=row_names,
+        col_names=col_names,
+        diagonal=False,
+    )
+
+    stored_names = [
+        "row2~col1",
+        "row3~col1",
+        "row3~col2",
+    ]
+
+    features = storage.list_features()
+    feature_md5 = list(features.keys())[0]
+    assert "fc" == features[feature_md5]["name"]
+    read_df = storage.read_df(feature_md5=feature_md5)
+    assert list(read_df.columns) == stored_names
+    assert_array_equal(
+        read_df.values, data[np.tril_indices(n=data.shape[0], k=-1)][None, :]
+    )
+
+
 # TODO: can the test be parametrized?
 def test_store_multiple_output(tmp_path: Path):
     """Test storing using single_output=False.