continuous wavelet transform and wavelet transform coherence implementation

hypyp/plots.py

@@ -0,0 +1,252 @@
+import math
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+import itertools as itertools
+from matplotlib.ticker import FuncFormatter
+from mne_connectivity.viz import plot_connectivity_circle
+from mne.viz import circular_layout
+
+# Define a custom locator and formatter for periods
+def custom_locator_freqs(ymin, ymax):
+    ticks = []
+    ticks.extend(range(math.ceil(ymin), 11))
+    ticks.extend(range(12, 21, 2))
+    ticks.extend(range(25, 40 + 1, 5))
+    ticks.extend(range(50, int(ymax) + 1, 10))
+    return ticks
+
+def plot_wavelet_transform_weights(
+    W,
+    times,
+    freqs,
+    coif,
+    sfreq,
+    bin_seconds=None,
+    frequency_cuts=None,
+    title=None,
+    ax=None,
+    show_colorbar=True,
+    show_cone_of_influence=True,
+    show_nyquist=True,
+    show_bins=True,
+):
+    # create the figure if needed
+    if ax is None:
+        fig, ax = plt.subplots()
+    else:
+        fig = ax.get_figure()
+
+    xx, yy = np.meshgrid(times, freqs)
+
+    #im = ax.pcolor(xx, yy, W, vmin=0, vmax=1)
+    im = ax.pcolor(xx, yy, np.abs(W))
+    ax.set_yscale('log')
+    ax.set_xlabel('Time (s)')
+    ax.set_ylabel('Frequency (Hz)')
+
+    color_invalid = 'C0'
+    # Cone of influence
+    if show_cone_of_influence:
+        ax.plot(times, coif, color=color_invalid)
+        ax.fill_between(times, coif, y2=np.min(freqs), step="mid", color=color_invalid, alpha=0.4)
+
+    if show_nyquist:
+        nyquist = np.ones((len(times),)) * (sfreq / 2)
+        ax.plot(times, nyquist, color=color_invalid)
+        ax.fill_between(times, nyquist, y2=np.max(freqs), step="mid", color=color_invalid, alpha=0.4)
+
+    if show_bins:
+        if bin_seconds is not None:
+            for time_cut in np.arange(0, max(times), bin_seconds):
+                plt.axvline(x=time_cut, color='red', lw=0.5)
+
+        if frequency_cuts is not None:
+            for frequency_cut in frequency_cuts:
+                plt.axhline(y=frequency_cut, color='red', lw=0.5)
+
+    # Dynamically set ticks based on the current range
+    ymin, ymax = ax.get_ylim()  # Get the y-axis limits
+    ax.set_yticks(custom_locator_freqs(ymin, ymax))
+    ax.yaxis.set_major_formatter(FuncFormatter(lambda y,_: f"{int(y)}" if y >= 1 else f"{y:.1f}"))
+    #ax.yaxis.get_major_formatter().set_scientific(False)  # Disable scientific notation
+
+    ax.set_xlim(times.min(), times.max())
+    ax.set_ylim(freqs.min(), freqs.max())
+
+    #ax.invert_yaxis()
+
+    if show_colorbar:
+        fig.colorbar(im)
+
+    if title is not None:
+        ax.set_title(title)
+    else:
+        ax.set_title('CWT Weights')
+
+
+    return fig
+
+def subplot_heatmap_from_pivot(pivot, ordered_fields, ax):
+    index_order = [i for i in ordered_fields if i in pivot.index]
+    # Append indices not in ordered_fields
+    for i in pivot.index:
+        if i not in index_order:
+            index_order.append(i)
+    column_order = [c for c in ordered_fields if c in pivot.columns]
+    for c in pivot.columns:
+        if c not in column_order:
+            column_order.append(c)
+
+    pivot_reordered = pivot.reindex(index=index_order, columns=column_order)
+    heatmap = sns.heatmap(pivot_reordered, cmap='viridis', vmin=0, vmax=1, cbar=False, ax=ax)
+
+    ax.set_xticks(ticks=range(len(pivot_reordered.columns)))
+    ax.set_xticklabels(pivot_reordered.columns, rotation=90, ha='left', fontsize=6 if len(column_order)>15 else 10)
+    ax.set_yticks(ticks=range(len(pivot_reordered.index)))
+    ax.set_yticklabels(pivot_reordered.index, rotation=0, va='top', fontsize=6 if len(index_order)>15 else 10)
+    ax.tick_params(axis='both', which='both', length=0)
+
+    return heatmap
+
+def plot_coherence_matrix(
+    df,
+    s1_label,
+    s2_label,
+    field1, # roi1 or channel1
+    field2, # roi2 or channel2
+    ordered_fields,
+):
+    # We don't sharex and sharey because the list of channels might be different in the 2 subjects
+
+    dyad_selector = (df['is_intra']==False)
+    s1_selector = (df['is_intra']==True) & (df['is_intra_of']==1) & (df['channel1']!=df['channel2'])
+    s2_selector = (df['is_intra']==True) & (df['is_intra_of']==2) & (df['channel1']!=df['channel2'])
+    df_dyad = df[dyad_selector]
+    df_s1 = df[s1_selector]
+    df_s2 = df[s2_selector]
+
+    pivot_s1 = df_s1.pivot_table(index=field1, columns=field2, values='coherence', aggfunc='mean', observed=False)
+    pivot_s2 = df_s2.pivot_table(index=field2, columns=field1, values='coherence', aggfunc='mean', observed=False)
+    pivot_dyad = df_dyad.pivot_table(index=field1, columns=field2, values='coherence', aggfunc='mean', observed=False)
+
+    if np.all(df['is_intra']):
+        fig, ax = plt.subplots(1, 1, figsize=(8, 8), sharex=False, sharey=False)
+        subplot_heatmap_from_pivot(pivot_s1.rename_axis(index=s1_label, columns=s1_label), ordered_fields=ordered_fields, ax=ax)
+    else:
+        fig, axes = plt.subplots(2, 2, figsize=(8, 8), sharex=False, sharey=False)
+        subplot_heatmap_from_pivot(pivot_s1.rename_axis(index=s1_label, columns=s1_label), ordered_fields=ordered_fields, ax=axes[0,0]) # top left
+        subplot_heatmap_from_pivot(pivot_dyad.rename_axis(index=s1_label, columns=s2_label), ordered_fields=ordered_fields, ax=axes[0,1]) # top right
+        subplot_heatmap_from_pivot(pivot_dyad.T.rename_axis(index=s2_label, columns=s1_label), ordered_fields=ordered_fields, ax=axes[1,0]) # bottom left
+        subplot_heatmap_from_pivot(pivot_s2.rename_axis(index=s2_label, columns=s2_label), ordered_fields=ordered_fields, ax=axes[1,1]) # bottom right
+
+    #fig.subplots_adjust(wspace=0.1, hspace=0.1)
+
+    plt.tight_layout()
+    return fig
+
+
+def plot_coherence_connectogram(df_pivot, title='', ax=None):
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, subplot_kw={'projection': 'polar'})
+    else:
+        fig = ax.get_figure()
+
+    #node_angles = circular_layout(
+    #    df.columns, list(df.columns), start_pos=90, group_boundaries=[0, len(df.columns) // 2]
+    #)
+    plot_connectivity_circle(df_pivot.values,
+        df_pivot.columns,
+        #node_angles=node_angles,
+        title=title,
+        #vmin=0,
+        #vmax=1,
+        colormap='Greys',
+        ax=ax,
+        facecolor='white',
+        textcolor='black',
+        node_edgecolor='black',
+        node_colors=['white'],
+    )
+    return fig
+
+def plot_coherence_connectogram_split(df_pivot, title='', ax=None):
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, subplot_kw={'projection': 'polar'})
+    else:
+        fig = ax.get_figure()
+
+    values = df_pivot.values.flatten()
+    s1_roi_list = df_pivot.index.to_list()
+    s2_roi_list = df_pivot.columns.to_list()
+    node_names = s1_roi_list + s2_roi_list
+
+    s1_roi_idx = []
+    s2_roi_idx = []
+
+    for i in range(len(s1_roi_list)):
+        for j in range(len(s2_roi_list)):
+            s1_roi_idx.append(i)
+            s2_roi_idx.append(j+len(s1_roi_list))
+
+    indices = (
+        np.array(s1_roi_idx),
+        np.array(s2_roi_idx),
+    )
+
+    node_names_ordered = []
+    for i in range(len(s1_roi_list)):
+        node_names_ordered.append(s1_roi_list[i])
+    for i in range(len(s2_roi_list)):
+        node_names_ordered.append(s2_roi_list[len(s2_roi_list)-i-1])
+
+    node_angles = circular_layout(
+        node_names, node_names_ordered, start_pos=90, group_boundaries=[0, len(s1_roi_list)]
+    )
+
+    plot_connectivity_circle(
+        values,
+        node_names,
+        indices = indices,
+        node_angles=node_angles,
+        title=title,
+        #vmin=0,
+        #vmax=1,
+        colormap='Greys',
+        ax=ax,
+        facecolor='white',
+        textcolor='black',
+        node_edgecolor='black',
+        node_colors=['white'],
+    )
+    return fig
+
+
+def plot_coherence_bars_per_task(df):
+    selector = (df['roi1'] == df['roi2'])
+
+    filtered_df = df[selector]
+
+    # remove intra-subject "same channel" coherence to avoid counting them in means
+    intra_same_ch_selector = (filtered_df['is_intra']==True) & (filtered_df['channel1']==filtered_df['channel2'])
+    filtered_df[intra_same_ch_selector] = np.nan
+
+    p = sns.catplot(
+        data=filtered_df, kind="bar",
+        x="roi1", y="coherence", hue="task",
+        col="is_intra",
+        #palette="dark", alpha=.6, height=6
+    )
+    p.despine(left=True)
+    p.set_axis_labels("", "Coherence")
+    p.set_xticklabels(rotation=45)
+    p.set(ylim=(0, 1))
+
+    p.legend.set_title("Task")
+
+    p.set_titles('Is intra: {col_name}')
+
+    plt.subplots_adjust(bottom=0.5)
+    return p
+
+

hypyp/profiling.py

@@ -0,0 +1,80 @@
+import tracemalloc
+import time
+
+class MemoryMonitor(object):
+    def __init__(self):
+        """
+        Track memory usage during a code block execution. Usage:
+
+        ```
+        with MemoryMonitor():
+            // my code
+        ```
+
+        """
+        pass
+
+    def format_memory_size(self, size_bytes):
+        if size_bytes >= 1 << 30:  # Greater than or equal to 1 GiB
+            return f"{size_bytes / (1 << 30):.2f} G"
+        elif size_bytes >= 1 << 20:  # Greater than or equal to 1 MiB
+            return f"{size_bytes / (1 << 20):.2f} M"
+        elif size_bytes >= 1 << 10:  # Greater than or equal to 1 KiB
+            return f"{size_bytes / (1 << 10):.2f} k"
+        else:  # Less than 1 KiB
+            return f"{size_bytes} B"
+
+    def __enter__(self):
+        tracemalloc.start()
+
+    def __exit__(self, *args):
+        res = tracemalloc.get_traced_memory()
+        print(f"[MemoryMonitor] allocated: {self.format_memory_size(res[0])}, peak: {self.format_memory_size(res[1])}", )
+        tracemalloc.stop()
+
+class TimeTracker(object):
+    def __init__(self, label:str='time_tracker'):
+        """
+        Track code execution time. Usage:
+
+        ```
+        with TimeTracker('foo'):
+            // my code
+        ```
+
+        Args:
+            label (str): identifier for the code block (for display only)
+        """
+        self.start_time = None
+        self.stop_time = None
+        self.duration = None
+        self.label = label
+
+    @staticmethod
+    def human_readable_duration(seconds):
+        if seconds < 5:
+            return f"{seconds:.2f} seconds"
+        if seconds < 60:
+            return f"{seconds:.0f} seconds"
+        if seconds < 3600:
+            minutes = seconds / 60
+            return f"{minutes:.1f} minutes"
+        if seconds < 86400:
+            hours = seconds / 3600
+            return f"{hours:.1f} hours"
+        days = seconds / 86400
+        return f"{days:.1f} days"
+
+    def start(self):
+        self.start_time = time.time()
+
+    def stop(self):
+        self.stop_time = time.time()
+        self.duration = self.stop_time - self.start_time
+
+    def __enter__(self):
+        self.start()
+
+    def __exit__(self, *args):
+        self.stop()
+        print(f"--- [{self.label}] {self.duration} seconds ---")