diff --git a/src/example_sequence_matching.py b/src/example_sequence_matching.py index deb2d41..30d8413 100644 --- a/src/example_sequence_matching.py +++ b/src/example_sequence_matching.py @@ -10,204 +10,16 @@ """ import argparse -import os import sys -import numpy as np -import pandas as pd -import matplotlib.pyplot as plt -from tqdm import tqdm -from loading_utils import load_spike_data -from loading_utils import filter_dataframe from loading_utils import group_dataframes_by_time -from sequence_matching_utils import lccs_rabin_karp - - -# build complex file paths for test or full data set -def get_base_directory(individual): - """ - Return correct data root depending on individual. - - Parameters - ---------- - individual :: str - - Returns - ------- - str - Base directory path. - """ - if individual == "test": - return "data/test_data" - return "data/full_data" - - -def build_file_paths(individual, stimuli): - """ - Construct all required file paths depending on individual. - - For real animals (7742, 7744): - data/full_data//stimuli/ - - For test: - data/test_data/ - """ - base = get_base_directory(individual) - - # files for test data sets - if individual == "test": - root = base - return { - "spike_times": os.path.join(root, "test_spike_times.npy"), - "spike_clusters": os.path.join(root, "test_spike_clusters.npy"), - "kslabels": os.path.join(root, "test_cluster_KSLabel.tsv"), - "behavior_csv": os.path.join( - root, "test_events_with_indices.csv" - ), - "swr_csv": os.path.join( - root, "test_SWRs_ca2.csv" - ), - } - - # files from full_data: like 7742 or 7744 - stimuli = stimuli - extra_tag = "_sleepyvole" if individual == "7742" else "" - root = os.path.join(base, individual, stimuli) - - return { - "spike_times": os.path.join(root, "spike_times.npy"), - "spike_clusters": os.path.join(root, "spike_clusters.npy"), - "kslabels": os.path.join(root, "cluster_KSLabel.tsv"), - "behavior_csv": os.path.join( - root, - f"{individual}_{stimuli}{extra_tag}_events_with_indices.csv", - ), - "swr_csv": os.path.join( - root, - f"{individual}_{stimuli}{extra_tag}_SWRs_ca2.csv", - ), - } - - -# load data -def load_data(paths): - """ - Load spike, behavior, and SWR data. - - Parameters - ---------- - paths :: dict - - Returns - ------- - (DataFrame, DataFrame, DataFrame) - """ - try: - spike_df = load_spike_data( - time_dir=paths["spike_times"], - cluster_dir=paths["spike_clusters"], - label_dir=paths["kslabels"], - ) - except FileNotFoundError: - print("Error: Spike data files not found.") - sys.exit(1) - - spike_df = filter_dataframe(spike_df, {"KSLabel": ["good"]}) - spike_df = spike_df.reset_index(drop=True) - - try: - behavior_df = pd.read_csv(paths["behavior_csv"]) - except FileNotFoundError: - print("Error: Behavior CSV not found.") - sys.exit(1) - - try: - swr_df = pd.read_csv(paths["swr_csv"]) - except FileNotFoundError: - print("Error: SWR CSV not found.") - sys.exit(1) - - return spike_df, behavior_df, swr_df - - -# filter behavior dataframe for relevant events -def preprocess_behavior_df(df): - """ - Filter and rename time columns in behavior dataframe. - - Parameters - ---------- - df :: pandas.DataFrame - - Returns - ------- - pandas.DataFrame - """ - df = filter_dataframe( - df, - {"EventType": ["social interaction", "cup interaction"]}, - ) - df = df.reset_index(drop=True) - - df[["indexStart", "indexEnd"]] /= 2500.0 - df = df.rename(columns={"indexStart": "Start", "indexEnd": "Stop"}) - return df - - -def compute_overlap_matrix(behavior_lists, swr_lists): - """ - Compute normalized LCCS overlap matrix. - - Parameters - ---------- - behavior_lists :: list - swr_lists :: list - - Returns - ------- - numpy.ndarray - """ - overlap = np.zeros((len(swr_lists), len(behavior_lists))) - - for bx, bseq in enumerate(tqdm(behavior_lists)): - if not bseq: - continue - - bset = set(bseq) - - for sy, sseq in enumerate(swr_lists): - if not sseq: - continue - - if bset.isdisjoint(sseq): - continue - - overlap[sy, bx] = lccs_rabin_karp(sseq, bseq, normalize=True) - - return overlap - - -def plot_overlap_matrix(matrix): - """ - Visualize the SWR behavior replay of behavior activity. - - Parameters - ---------- - matrix :: numpy.ndarray - """ - plt.figure(figsize=(8, 8)) - plt.imshow( - matrix, - cmap="inferno", - aspect="auto", - interpolation="none", - origin="lower", - ) - plt.xlabel("Behavior Window Index") - plt.ylabel("SWR Window Index") - plt.colorbar(label="Normalized Overlap") - plt.tight_layout() - plt.show() +from sequence_matching_utils import ( + build_file_paths, + load_data, + preprocess_behavior_df, + compute_overlap_matrix, + plot_overlap_matrix, +) def parse_arguments(): @@ -226,9 +38,9 @@ def parse_arguments(): parser.add_argument( "--stimuli", type=str, - required=True, - help="Social stimuli folder name for full data " - "(default: PartnerIntro).", + required=False, + help="Social stimuli folder name for full data.", + default="PartnerIntro", ) args = parser.parse_args() diff --git a/src/sequence_matching_utils.py b/src/sequence_matching_utils.py index 404b6f9..6e2c4ef 100644 --- a/src/sequence_matching_utils.py +++ b/src/sequence_matching_utils.py @@ -1,7 +1,11 @@ import numpy as np import pandas as pd from tqdm import tqdm +import matplotlib.pyplot as plt import os +import sys + +from loading_utils import load_spike_data, filter_dataframe def _preprocess_lists( @@ -273,3 +277,190 @@ def lccs_rabin_karp( seq = _recover_sequence(L, l1, l2, H1, H2, n1, n2, pow_base, mod) length_out = L / n1 if normalize else L return length_out, seq + + +# build complex file paths for test or full data set +def get_base_directory(individual): + """ + Return correct data root depending on individual. + + Parameters + ---------- + individual :: str + + Returns + ------- + str + Base directory path. + """ + if individual == "test": + return "data/test_data" + return "data/full_data" + + +def build_file_paths(individual, stimuli): + """ + Construct all required file paths depending on individual. + + For real animals (7742, 7744): + data/full_data//stimuli/ + + For test: + data/test_data/ + """ + base = get_base_directory(individual) + + # files for test data sets + if individual == "test": + root = base + return { + "spike_times": os.path.join(root, "test_spike_times.npy"), + "spike_clusters": os.path.join(root, "test_spike_clusters.npy"), + "kslabels": os.path.join(root, "test_cluster_KSLabel.tsv"), + "behavior_csv": os.path.join( + root, "test_events_with_indices.csv" + ), + "swr_csv": os.path.join( + root, "test_SWRs_ca2.csv" + ), + } + + # files from full_data: like 7742 or 7744 + stimuli = stimuli + extra_tag = "_sleepyvole" if individual == "7742" else "" + root = os.path.join(base, individual, stimuli) + + return { + "spike_times": os.path.join(root, "spike_times.npy"), + "spike_clusters": os.path.join(root, "spike_clusters.npy"), + "kslabels": os.path.join(root, "cluster_KSLabel.tsv"), + "behavior_csv": os.path.join( + root, + f"{individual}_{stimuli}{extra_tag}_events_with_indices.csv", + ), + "swr_csv": os.path.join( + root, + f"{individual}_{stimuli}{extra_tag}_SWRs_ca2.csv", + ), + } + + +# load data +def load_data(paths): + """ + Load spike, behavior, and SWR data. + + Parameters + ---------- + paths :: dict + + Returns + ------- + (DataFrame, DataFrame, DataFrame) + """ + try: + spike_df = load_spike_data( + time_dir=paths["spike_times"], + cluster_dir=paths["spike_clusters"], + label_dir=paths["kslabels"], + ) + except FileNotFoundError: + print("Error: Spike data files not found.") + sys.exit(1) + + spike_df = filter_dataframe(spike_df, {"KSLabel": ["good"]}) + spike_df = spike_df.reset_index(drop=True) + + try: + behavior_df = pd.read_csv(paths["behavior_csv"]) + except FileNotFoundError: + print("Error: Behavior CSV not found.") + sys.exit(1) + + try: + swr_df = pd.read_csv(paths["swr_csv"]) + except FileNotFoundError: + print("Error: SWR CSV not found.") + sys.exit(1) + + return spike_df, behavior_df, swr_df + + +# filter behavior dataframe for relevant events +def preprocess_behavior_df(df): + """ + Filter and rename time columns in behavior dataframe. + + Parameters + ---------- + df :: pandas.DataFrame + + Returns + ------- + pandas.DataFrame + """ + df = filter_dataframe( + df, + {"EventType": ["social interaction", "cup interaction"]}, + ) + df = df.reset_index(drop=True) + + df[["indexStart", "indexEnd"]] /= 2500.0 + df = df.rename(columns={"indexStart": "Start", "indexEnd": "Stop"}) + return df + + +def compute_overlap_matrix(behavior_lists, swr_lists): + """ + Compute normalized LCCS overlap matrix. + + Parameters + ---------- + behavior_lists :: list + swr_lists :: list + + Returns + ------- + numpy.ndarray + """ + overlap = np.zeros((len(swr_lists), len(behavior_lists))) + + for bx, bseq in enumerate(tqdm(behavior_lists)): + if not bseq: + continue + + bset = set(bseq) + + for sy, sseq in enumerate(swr_lists): + if not sseq: + continue + + if bset.isdisjoint(sseq): + continue + + overlap[sy, bx] = lccs_rabin_karp(sseq, bseq, normalize=True) + + return overlap + + +def plot_overlap_matrix(matrix): + """ + Visualize the SWR behavior replay of behavior activity. + + Parameters + ---------- + matrix :: numpy.ndarray + """ + plt.figure(figsize=(8, 8)) + plt.imshow( + matrix, + cmap="inferno", + aspect="auto", + interpolation="none", + origin="lower", + ) + plt.xlabel("Behavior Window Index") + plt.ylabel("SWR Window Index") + plt.colorbar(label="Normalized Overlap") + plt.tight_layout() + plt.show() diff --git a/test/test_sequence_matching_utils.py b/test/test_sequence_matching_utils.py index 766b23e..96290cc 100644 --- a/test/test_sequence_matching_utils.py +++ b/test/test_sequence_matching_utils.py @@ -9,7 +9,6 @@ sys.path.append("src/") # noqa import sequence_matching_utils as utils # noqa -import example_sequence_matching as matrix # noqa # Default parameters used in 'lccs_rabin_karp' BASE = 257 @@ -133,19 +132,19 @@ def mock_group_dataframes_by_time(window_df, event_df, event_time_column, class TestPipelineFunctions(unittest.TestCase): def test_get_base_directory(self): - self.assertEqual(matrix.get_base_directory("test"), "data/test_data") - self.assertEqual(matrix.get_base_directory("7742"), "data/full_data") - self.assertEqual(matrix.get_base_directory("7744"), "data/full_data") + self.assertEqual(utils.get_base_directory("test"), "data/test_data") + self.assertEqual(utils.get_base_directory("7742"), "data/full_data") + self.assertEqual(utils.get_base_directory("7744"), "data/full_data") def test_build_file_paths_test(self): - paths = matrix.build_file_paths("test", "PartnerIntro") + paths = utils.build_file_paths("test", "PartnerIntro") self.assertIn("test_spike_times.npy", paths["spike_times"]) self.assertIn("test_events_with_indices.csv", paths["behavior_csv"]) self.assertIn("test_SWRs_ca2.csv", paths["swr_csv"]) self.assertTrue(paths["spike_times"].startswith("data/test_data")) def test_build_file_paths_fulldata_7742(self): - paths = matrix.build_file_paths("7742", "PartnerIntro") + paths = utils.build_file_paths("7742", "PartnerIntro") # Check directory structure for key in paths: self.assertIn("data/full_data/7742/PartnerIntro", paths[key]) @@ -153,7 +152,7 @@ def test_build_file_paths_fulldata_7742(self): self.assertIn("7742_PartnerIntro_sleepyvole", paths["behavior_csv"]) def test_build_file_paths_fulldata_7744(self): - paths = matrix.build_file_paths("7744", "SSIntro") + paths = utils.build_file_paths("7744", "SSIntro") for key in paths: self.assertIn("data/full_data/7744/SSIntro", paths[key]) self.assertIn("7744_SSIntro_events_with_indices.csv", @@ -170,7 +169,7 @@ def test_preprocess_behavior_df(self): original_filter = sys.modules["loading_utils"].filter_dataframe sys.modules["loading_utils"].filter_dataframe = mock_filter_dataframe - out = matrix.preprocess_behavior_df(df) + out = utils.preprocess_behavior_df(df) sys.modules["loading_utils"].filter_dataframe = original_filter self.assertEqual(len(out), 1) @@ -185,7 +184,7 @@ def test_compute_overlap_matrix(self): # First SWR row overlaps with first behavior seq → lccs = 2/3 # Third SWR row overlaps with first behavior seq → lccs = 3/3 = 1.0 - mat = matrix.compute_overlap_matrix(behavior_lists, swr_lists) + mat = utils.compute_overlap_matrix(behavior_lists, swr_lists) self.assertEqual(mat.shape, (3, 3)) self.assertGreater(mat[0, 0], 0)