Reactivation¶

# load theta epochs
state_dict = loading.load_SleepState_states(basepath)
theta_epochs = nel.EpochArray(
    state_dict["THETA"],
)
nrem_epochs = nel.EpochArray(
    state_dict["NREMstate"],
)
theta_epochs, nrem_epochs

# load need data (spikes, ripples, epochs)
assembly_react.load_data()
assembly_react

assembly_react.epoch_df

assembly_react.get_weights(epoch=assembly_react.epochs[1] & theta_epochs)
assembly_react

assembly_react.plot()
plt.show()

assembly_act = assembly_react.get_assembly_act()
assembly_act

fig, ax = plt.subplots(2, 3, figsize=(15, 8), sharey=False, sharex=False)
ax = ax.flatten()

# share y axis of first row
ax[0] = plt.subplot(231, sharey=ax[1])
ax[2] = plt.subplot(233, sharey=ax[0])

# plot assembly ripple psth
psth_swr_pre.plot(ax=ax[0], legend=False)
psth_swr_post.plot(ax=ax[1], legend=False)
(psth_swr_post - psth_swr_pre).plot(ax=ax[2])

# plot mean assembly ripple psth
psth_swr_pre.mean(axis=1).plot(ax=ax[0], color="k", legend=False)
psth_swr_post.mean(axis=1).plot(ax=ax[1], color="k", legend=False)
(psth_swr_post - psth_swr_pre).mean(axis=1).plot(ax=ax[2], color="k")

# plot assembly ripple psth heatmap
sns.heatmap(psth_swr_pre.T, ax=ax[3], cbar=False, vmin=0, vmax=5)
sns.heatmap(psth_swr_post.T, ax=ax[4], cbar=False, vmin=0, vmax=5)
sns.heatmap(
    (psth_swr_post - psth_swr_pre).T,
    ax=ax[5],
    cbar=False,
    vmin=-5,
    vmax=5,
    cmap="coolwarm",
)

for ax_ in ax[:3]:
    # dashed line at zero
    ax_.axvline(0, linestyle="--", color="k", linewidth=1)
    # set x axis limits
    ax_.set_xlim(-0.5, 0.5)
    # add grid lines
    ax_.grid()

ax[0].set_title("Pre")
ax[1].set_title("Post")
ax[2].set_title("Post - Pre")

# move legend
ax[2].legend(
    bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0, frameon=False, title="assembly"
)

# add labels
ax[0].set_ylabel("assembly activity")
ax[3].set_ylabel("assembly")
ax[3].set_xlabel("time from SWR start (s)")

# clean axis using seaborn
sns.despine()

plt.show()

# Standard assembly detection (finds both within-region and cross-region assemblies)
assembly_react_standard = AssemblyReact(weight_dt=0.05, z_mat_dt=0.005)
assembly_react_standard.add_st(st)
assembly_react_standard.epochs = beh_epochs
assembly_react_standard.get_weights(
    epoch=assembly_react_standard.epochs[pre_task_post[0][1].item()] & theta_epochs
)

print(f"Standard detection found {assembly_react_standard.n_assemblies()} assemblies")

# Cross-structural assembly detection (only finds assemblies spanning multiple regions)
assembly_react_cross = AssemblyReact(
    weight_dt=0.05,
    z_mat_dt=0.005,
    cross_structural=brain_regions,  # This enables cross-structural detection
)
assembly_react_cross.add_st(st)
assembly_react_cross.epochs = beh_epochs
assembly_react_cross.get_weights(
    epoch=assembly_react_cross.epochs[pre_task_post[0][1].item()] & theta_epochs
)

print(
    f"Cross-structural detection found {assembly_react_cross.n_assemblies()} assemblies"
)

if assembly_react_cross.n_assemblies() > 0:
    # Plot cross-structural assembly weights
    fig, axes = assembly_react_cross.plot(figsize=(12, 6))

    # Add vertical line to separate brain regions
    ca1_neurons = np.sum(brain_regions == "CA1")
    for ax in axes.flat:
        ax.axhline(
            ca1_neurons - 0.5, color="red", linestyle="--", alpha=0.7, linewidth=2
        )
        ax.text(
            0.02,
            1.05,
            f"PFC (top)\nCA1 (bottom)",
            transform=ax.transAxes,
            va="top",
            ha="left",
            bbox=dict(boxstyle="round", facecolor="white", alpha=0.8),
        )

    plt.suptitle("Cross-Structural Assemblies (CA1-PFC)", fontsize=14, y=1.02)
    plt.tight_layout()
    plt.show()

    # Analyze which regions participate in each assembly
    assembly_react_cross.find_members()

    print("\n Cross-structural assembly analysis:")
    for i in range(assembly_react_cross.n_assemblies()):
        ca1_active = assembly_react_cross.assembly_members[i, :ca1_neurons].sum()
        pfc_active = assembly_react_cross.assembly_members[i, ca1_neurons:].sum()
        print(
            f"  Assembly {i + 1}: CA1 neurons: {ca1_active}, PFC neurons: {pfc_active}"
        )

else:
    print("No cross-structural assemblies detected in this dataset.")
    print("This could mean:")
    print("1. There are no assemblies spanning both regions")
    print("2. The assemblies are primarily within-region")
    print("3. More data or different parameters may be needed")

assembly_act_standard = assembly_react_standard.get_assembly_act(
    epoch=assembly_react_standard.epochs[pre_task_post[0][2].item()]
)

# Compute assembly activity for cross-structural assemblies
assembly_act_cross = assembly_react_cross.get_assembly_act(
    epoch=assembly_react_cross.epochs[pre_task_post[0][2].item()]
)

nrem_ripples = ripples & nrem_epochs

psth_standard_post = npy.process.event_triggered_average(
    timestamps=assembly_act_standard.abscissa_vals,
    signal=assembly_act_standard.data.T,
    events=nrem_ripples[
        assembly_react_standard.epochs[pre_task_post[0][2].item()]
    ].starts,
    sampling_rate=assembly_act_standard.fs,
    window=[-0.5, 0.5],
    return_average=True,
    return_pandas=True,
)
psth_cross_post = npy.process.event_triggered_average(
    timestamps=assembly_act_cross.abscissa_vals,
    signal=assembly_act_cross.data.T,
    events=nrem_ripples[assembly_react_cross.epochs[pre_task_post[0][2].item()]].starts,
    sampling_rate=assembly_act_cross.fs,
    window=[-0.5, 0.5],
    return_average=True,
    return_pandas=True,
)

# Plot comparison
fig, axes = plt.subplots(1, 2, figsize=(10, 4), sharey=True)

# Plot standard assemblies
psth_standard_post.plot(ax=axes[0], legend=False)

axes[0].legend().set_visible(False)
# axes[0].set_title(f"Standard Assemblies (first {n_compare})")
axes[0].set_ylabel("Assembly Activity")
axes[0].axvline(0, linestyle="--", color="k", alpha=0.7)
axes[0].grid(True, alpha=0.3)
axes[0].set_title(f"Standard Assemblies (n={assembly_react_standard.n_assemblies()})")

# Plot cross-structural assemblies
psth_cross_post.plot(ax=axes[1], legend=False)
axes[1].set_title(
    f"Cross-Structural Assemblies (n={assembly_react_cross.n_assemblies()})"
)
axes[1].set_ylabel("Assembly Activity")
axes[1].axvline(0, linestyle="--", color="k", alpha=0.7)
axes[1].grid(True, alpha=0.3)

plt.suptitle("Assembly Activity During Post-Task Ripples", fontsize=14)
plt.tight_layout()
sns.despine()
plt.show()

# Compare peak activation
standard_peak = psth_standard_post.max().mean()
cross_peak = psth_cross_post.max().mean()

print(f"Peak activation comparison:")
print(f"Standard assemblies (mean): {standard_peak:.2f}")
print(f"Cross-structural assemblies (mean): {cross_peak:.2f}")

# Visualize cross-structural correlation matrix with within-structure filtered out
# Recompute the z-scored binned spike matrix using the method from AssemblyReact
zmat, _ = assembly_react_cross.get_z_mat(
    assembly_react_cross.st[
        assembly_react_cross.epochs[pre_task_post[0][1].item()] & theta_epochs
    ]
)

corr = np.corrcoef(zmat)

# Mask within-structure correlations
region_labels = brain_regions
mask = np.equal.outer(region_labels, region_labels)
corr_masked = corr.copy()
corr_masked[mask] = np.nan  # set within-structure to NaN

fig, ax = plt.subplots(figsize=(7, 6))
im = ax.imshow(corr_masked, cmap="coolwarm", vmin=-0.015, vmax=0.015)

# Add lines to separate regions
ca1_neurons = np.sum(brain_regions == "CA1")
ax.axhline(ca1_neurons - 0.5, color="red", linestyle="--", linewidth=2, alpha=0.7)
ax.axvline(ca1_neurons - 0.5, color="red", linestyle="--", linewidth=2, alpha=0.7)

# Axis labels and ticks
ax.set_xlabel("Neuron (sorted by region)")
ax.set_ylabel("Neuron (sorted by region)")
ax.set_xticks([0, ca1_neurons, len(brain_regions)])
ax.set_yticks([0, ca1_neurons, len(brain_regions)])
ax.set_xticklabels(["0", f"{ca1_neurons}", f"{len(brain_regions)}"])
ax.set_yticklabels(["0", f"{ca1_neurons}", f"{len(brain_regions)}"])

# Add region labels using text outside the plot area for visibility
ax.text(
    ca1_neurons / 2,
    -10,
    "CA1",
    ha="center",
    va="bottom",
    fontsize=14,
    fontweight="bold",
    color="black",
    clip_on=False,
)
ax.text(
    ca1_neurons + (len(brain_regions) - ca1_neurons) / 2,
    -10,
    "PFC",
    ha="center",
    va="bottom",
    fontsize=14,
    fontweight="bold",
    color="black",
    clip_on=False,
)

# Add the title above the region annotations by using fig.suptitle
fig.suptitle(
    "Cross-Structural Correlation Matrix (within-structure filtered)",
    fontsize=12,
    y=0.94,
)

# Colorbar
cbar = plt.colorbar(im, ax=ax, label="Correlation", pad=0.02)
plt.tight_layout(rect=[0, 0, 1, 0.93])
plt.show()