Repository: yeatmanlab/AFQ-Insight (language: python)

Issue:
Add utilities for matching between participants
To create quasi-experiments based on certain match characteristics in large datasets. 

@36000 has written some code to get this started. 

Maybe this should go into a new module called `match`? It looks like the convention so far is for modules to be named as imperative verbs ("cross_validate", "transform", "plot").



Discussion / hints:
Agreed! I think `match` would be a good home for it. @36000, maybe you could get this started with a WIP PR by extracting the matching code you are using and modularizing it into various functions. I'm happy to help if you get a WIP PR started.

Fix patch:
diff --git a/afqinsight/match.py b/afqinsight/match.py
new file mode 100644
index 00000000..021423a7
--- /dev/null
+++ b/afqinsight/match.py
@@ -0,0 +1,142 @@
+"""Perform statistical matching comparing test and control units."""
+
+import numpy as np
+import pandas as pd
+from scipy.optimize import linear_sum_assignment
+
+
+__all__ = ["mahalonobis_dist_match"]
+
+
+def mahalonobis_dist_match(
+    data=None, test=None, ctrl=None, status_col=None, feature_cols=None, threshold=0.2
+):
+    """
+    Perform Mahalanobis distance matching (MDM) on a pandas dataframe or 2d array.
+
+    If data is None, then test and ctrl must be arrays.
+    If data is provided, then status_col and feature_cols are used.
+
+    Parameters
+    ----------
+    data : pandas Dataframe, optional
+        The feature samples for test subjects.
+        Default: None
+    test : array-like of shape (n_samples1, n_features), optional
+        The feature samples for test subjects.
+        Default: None
+    ctrl : array-like of shape (n_samples2, n_features), optional
+        Feature samples for control subjects from which
+        matches will be selected.
+        Default: None
+    status_col : str, optional
+        Column with subject status. Must have only two unique values,
+        evaluating to True and False, where True indicates a test subject
+        and False indicates a control.
+        Default: None
+    feature_cols : list of str, optional
+        List of columns to use as features for MDM.
+        If feature_columns is None then we use all of the columns
+        except status_col as the features.
+        Default: None
+    threshold : float, optional
+        Mahalnobis distance beyond which matches are not accepted.
+        Default: None
+
+    Returns
+    -------
+    selected_data : pandas Dataframe
+        If data is None, returns a Dataframe with columns named
+        'feature_0', 'feature_1', etc. and a 'status' column
+        where 1 corresponds to test and 0 corresponds to control.
+        If data is not None, returns data with only the selected rows.
+    """
+    if data is not None:
+        if status_col is None:
+            raise ValueError(
+                ("If data is provided," " status_col must also be provided")
+            )
+
+        if len(data[status_col].unique()) != 2:
+            raise ValueError(
+                ("status_col in data must " "contain only 2 unique values")
+            )
+
+        if feature_cols is None:
+            feature_cols = data.columns.drop(status_col)
+
+        df_test_rows = data[status_col].astype(bool)
+        df_ctrl_rows = ~data[status_col].astype(bool)
+        test = data.loc[df_test_rows, feature_cols].to_numpy()
+        ctrl = data.loc[df_ctrl_rows, feature_cols].to_numpy()
+    else:
+        test = np.asarray(test)
+        ctrl = np.asarray(ctrl)
+
+    if np.any(np.isnan(test)) or np.any(np.isnan(ctrl)):
+        raise ValueError(
+            (
+                "There are NaNs in test or ctrl data. "
+                "Please replace these NaNs using interpolation or by removing "
+                "the subjects with NaNs before calling mahalonobis_dist_match. "
+            )
+        )
+
+    # calculate Mahalonobis distance between test and control
+    nbrs = _mahalonobis_dist(test, ctrl)
+
+    # assign neighbors using Munkres algorithm
+    row_ind, col_ind = linear_sum_assignment(nbrs)
+
+    # remove matches that are too bad
+    filtered_test_ind = []
+    filtered_ctrl_ind = []
+    for row_idx, col_idx in zip(row_ind, col_ind):
+        if nbrs[row_idx, col_idx] <= threshold:
+            filtered_test_ind.append(row_idx)
+            filtered_ctrl_ind.append(col_idx)
+
+    # contstruct dataframe if no dataframe is provided
+    if data is None:
+        data = {}
+        filtered_test = test[filtered_test_ind]
+        filtered_ctrl = ctrl[filtered_ctrl_ind]
+        for i in range(test.shape[1]):
+            data[f"feature_{i}"] = np.concatenate(
+                (filtered_test[:, i], filtered_ctrl[:, i]), axis=0
+            )
+        data["status"] = np.concatenate(
+            (np.ones(len(filtered_test_ind)), np.zeros(len(filtered_ctrl_ind))), axis=0
+        )
+        return pd.DataFrame(data)
+    else:
+        # determine selected rows
+        df_test_idx = np.argwhere(df_test_rows.to_numpy()).flatten()
+        df_ctrl_idx = np.argwhere(df_ctrl_rows.to_numpy()).flatten()
+        all_idx = np.concatenate(
+            (df_test_idx[filtered_test_ind], df_ctrl_idx[filtered_ctrl_ind]), axis=0
+        )
+        return data.loc[all_idx]
+
+
+def _mahalonobis_dist(arr1, arr2):
+    """
+    Calculate the Mahalonobis distance between two 2d arrays along the first axis.
+
+    Parameters
+    ----------
+    arr1 : array-like of shape (n_samples1, n_features)
+        First input array.
+    arr2 : array-like of shape (n_samples2, n_features)
+        Second input array.
+
+    Returns
+    -------
+    nbrs : array-like of shape (n_samples1, n_samples2)
+        Mahalonobis distance between each sample in the
+        two input arrays.
+    """
+    v_inv = np.linalg.inv(np.cov(np.concatenate((arr1, arr2), axis=0).T, ddof=0))
+    diff = arr1[:, np.newaxis] - arr2[np.newaxis, :]
+    nbrs = np.sqrt(np.sum((diff @ v_inv) * diff, axis=2))
+    return nbrs


Test patch:
diff --git a/afqinsight/tests/test_match.py b/afqinsight/tests/test_match.py
new file mode 100644
index 00000000..0ccf0217
--- /dev/null
+++ b/afqinsight/tests/test_match.py
@@ -0,0 +1,112 @@
+"""Test functions in the match module"""
+
+import afqinsight.match as aim
+import numpy as np
+import pandas as pd
+from sklearn.utils._testing import assert_array_almost_equal, assert_array_equal
+from scipy.spatial.distance import mahalanobis
+import pytest
+
+
+test_features = np.asarray([[1, 2, 3, 4], [4, 2, 1, 3], [1, 4, 5, 6]])
+
+
+ctrl_features = np.asarray(
+    [
+        [3, 2, 3, 4],
+        [4, 5, 1, 3],
+        [1, 10, 5, 6],
+        [1, 2, 2, 4],
+        [4, 2, 0, 3],
+        [1, 4, 9, 6],
+    ]
+)
+
+
+data = pd.DataFrame(
+    {
+        "eid": [34, 35, 36, 37, 38, 39, 40, 41, 42],
+        "age": [12, 345, 63, 4, 5, 6, 89, 88, 10],
+        "status": [1, 1, 1, 0, 0, 0, 0, 0, 0],
+    }
+)
+for ii in range(4):
+    data[f"feature_{ii}"] = [*test_features[:, ii], *ctrl_features[:, ii]]
+
+
+def test_mahalonobis_dist():
+    nbrs = aim._mahalonobis_dist(test_features, ctrl_features)
+
+    v_inv = np.linalg.inv(
+        np.cov(np.concatenate((test_features, ctrl_features), axis=0).T, ddof=0)
+    )
+    nbrs_scipy = np.zeros((test_features.shape[0], ctrl_features.shape[0]))
+    for ii, test_sub in enumerate(test_features):
+        for jj, ctrl_sub in enumerate(ctrl_features):
+            nbrs_scipy[ii, jj] = mahalanobis(test_sub, ctrl_sub, v_inv)
+
+    assert_array_almost_equal(nbrs, nbrs_scipy)
+
+
+def test_mahalonobis_dist_match_df():
+    matched_df = aim.mahalonobis_dist_match(
+        data=data,
+        status_col="status",
+        feature_cols=["feature_0", "feature_1", "feature_2", "feature_3"],
+        threshold=1,
+    )
+
+    assert_array_equal(
+        matched_df.loc[matched_df.status.astype(bool)].eid.to_numpy(), [34, 35]
+    )
+    assert_array_equal(
+        matched_df.loc[~matched_df.status.astype(bool)].eid.to_numpy(), [40, 41]
+    )
+
+
+def test_mahalonobis_dist_match_df_err():
+    with pytest.raises(ValueError):  # no status column
+        aim.mahalonobis_dist_match(data=data)
+    with pytest.raises(ValueError):  # status has more than two unique values
+        aim.mahalonobis_dist_match(data=data, status_col="feature_0")
+
+
+def test_mahalonobis_dist_match_df_feauture_none():
+    data_wo_extras = data.drop(columns=["eid", "age"])
+    matched_df = aim.mahalonobis_dist_match(
+        data=data_wo_extras,
+        status_col="status",
+        threshold=1,
+    )
+
+    assert_array_equal(
+        matched_df.loc[matched_df.status.astype(bool)].feature_2.to_numpy(), [3, 1]
+    )
+    assert_array_equal(
+        matched_df.loc[~matched_df.status.astype(bool)].feature_2.to_numpy(), [2, 0]
+    )
+
+
+def test_mahalonobis_dist_match():
+    matched_df = aim.mahalonobis_dist_match(
+        test=test_features, ctrl=ctrl_features, threshold=1
+    )
+
+    assert_array_equal(
+        matched_df.loc[matched_df.status.astype(bool)].feature_2.to_numpy(), [3, 1]
+    )
+    assert_array_equal(
+        matched_df.loc[~matched_df.status.astype(bool)].feature_2.to_numpy(), [2, 0]
+    )
+    assert np.sum(matched_df.status.to_numpy() == 1) == np.sum(
+        matched_df.status.to_numpy() == 0
+    )
+
+
+def test_mahalonobis_dist_match_err():
+    with pytest.raises(ValueError):  # test if nan error is raised
+        test_features_with_nans = test_features.copy()
+        test_features_with_nans[0, 2] = np.nan
+        aim.mahalonobis_dist_match(
+            test=test_features_with_nans, ctrl=ctrl_features, threshold=1
+        )