plot_all_units_delta_distributions.py

#!/usr/bin/env python3
from __future__ import annotations

from pathlib import Path

import matplotlib.pyplot as plt
import pandas as pd


ROOT = Path(__file__).resolve().parent
IN_CSV = (
    ROOT
    / "NWB"
    / "analysis_NP-SC12-jh"
    / "trimmed_transition_control"
    / "per_unit_rates_trimmed_all85.csv"
)
OUT_DIR = ROOT / "NWB" / "analysis_NP-SC12-jh" / "trimmed_transition_control"


def main():
    df = pd.read_csv(IN_CSV)

    cols = [
        "delta_movie_minus_b1_hz",
        "delta_movie_minus_b2_hz",
        "delta_movie_minus_bcombined_hz",
    ]
    labels = [
        "movie - baseline1 (trimmed)",
        "movie - baseline2 (trimmed)",
        "movie - baseline_combined (trimmed)",
    ]

    fig, axes = plt.subplots(1, 3, figsize=(14, 4.5), sharey=True)
    for ax, col, label in zip(axes, cols, labels):
        ax.hist(df[col], bins=20, alpha=0.85)
        ax.axvline(0, color="k", linestyle="--", linewidth=1)
        ax.set_title(label)
        ax.set_xlabel("Delta rate (Hz)")
        ax.text(
            0.02,
            0.95,
            f"mean={df[col].mean():.3f}\nmedian={df[col].median():.3f}",
            transform=ax.transAxes,
            va="top",
            ha="left",
            fontsize=9,
        )
    axes[0].set_ylabel("# units")
    fig.tight_layout()
    fig.savefig(OUT_DIR / "all85_delta_distributions_triplet.png", dpi=180)
    plt.close(fig)

    long_df = pd.DataFrame(
        {
            "unit_id": list(df["unit_id"]) * 3,
            "comparison": (
                ["movie_minus_b1"] * len(df)
                + ["movie_minus_b2"] * len(df)
                + ["movie_minus_bcombined"] * len(df)
            ),
            "delta_hz": list(df["delta_movie_minus_b1_hz"])
            + list(df["delta_movie_minus_b2_hz"])
            + list(df["delta_movie_minus_bcombined_hz"]),
        }
    )
    long_df.to_csv(OUT_DIR / "all85_delta_distributions_long.csv", index=False)
    print(f"Saved: {OUT_DIR / 'all85_delta_distributions_triplet.png'}")
    print(f"Saved: {OUT_DIR / 'all85_delta_distributions_long.csv'}")


if __name__ == "__main__":
    main()