exp_eigen.py

from tango import Step
from tango.common import DatasetDict

from beliefprobing.generate import GenOut

import torch
import torch.nn.functional as F
import numpy as np
import scipy
import pandas as pd

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.collections import LineCollection
import seaborn as sns


def plot_many_categories_seaborn(
    df, x_col="category", y_col="value",
    top_k_labels=8, tick_every=50, figsize=(5, 3),
    trend_window=25, use_log=False, title=None, outfile=None
):
    # order by y for ranking
    d = df[[x_col, y_col]].dropna().sort_values(y_col, ascending=True).reset_index(drop=True)
    d["rank"] = np.arange(len(d))

    # low-ink theme
    sns.set_theme(style="whitegrid", rc={
        "axes.spines.right": False, "axes.spines.top": False,
        "axes.grid.axis": "y", "grid.alpha": 0.25, "grid.linewidth": 0.6
    })

    fig, ax = plt.subplots(figsize=figsize, dpi=160)

    # thin stems (matplotlib primitive under seaborn theme)
    baseline = 0 if d[y_col].min() >= 0 else d[y_col].min()
    ax.vlines(d["rank"], baseline, d[y_col], linewidth=0.6, alpha=0.35)

    # small points
    sns.scatterplot(data=d, x="rank", y=y_col, s=12, alpha=0.7, edgecolor="none", ax=ax)

    # optional rolling-median trend to reveal shape
    if trend_window and trend_window > 1 and trend_window % 2 == 1 and len(d) >= trend_window:
        trend = d[y_col].rolling(trend_window, center=True).median()
        sns.lineplot(x=d["rank"], y=trend, ax=ax, linewidth=1.2)

    # sparse ticks & rotated labels
    if tick_every > 0:
        ticks = np.arange(0, len(d), tick_every)
        ax.set_xticks(ticks)
        ax.set_xticklabels(d[x_col].iloc[ticks], rotation=90, ha="center")
    else:
        ax.set_xticks([])

    # annotate extremes for context
    if top_k_labels > 0:
        extremes = pd.concat([d.head(top_k_labels), d.tail(top_k_labels)])
        for _, r in extremes.iterrows():
            ax.annotate(r[x_col], (r["rank"], r[y_col]),
                        xytext=(0, 6), textcoords="offset points", ha="center", va="bottom", fontsize=10)

    if use_log:
        ax.set_yscale("log",)

    ax.margins(x=0.05)
    ax.set_xlabel("")
    ax.set_ylabel(y_col)
    if title:
        ax.set_title(title, pad=10)

    ax.spines['bottom'].set_position(('data', 0))

    plt.close()
    return fig, ax


@Step.register('plot_basis_magn')
class PlotBasisMagnitudes(Step):
    VERSION = "013"
    # CACHEABLE = False

    def run(self, probe, n: int, prefix: str) -> None:
        if probe.direction_type == 'belief':
            eigvals = probe.INTERVENE_LAMBDA  # r
        elif probe.direction_type == 'normal':
            eigvals = probe.CLASSIFY_LAMBDA  # r
        else:
            raise

        values = eigvals[-n:].flip(0)
        n = min(len(values), n)

        df = pd.DataFrame({
            "category": [f"λ{i:02d}" for i in range(n)],
            "value": values.tolist()
        })

        fig, ax = plot_many_categories_seaborn(
            df, tick_every=0, top_k_labels=8, trend_window=25, use_log=False, title=""
        )
        fig.savefig(f"local_outputs/plot_basis_magn/{prefix}_basis_lengths.pdf", bbox_inches="tight")


@Step.register('plot_cc_data')
class PlotCCData(Step):
    VERSION = "010"
    # CACHEABLE = False

    def run(self, probe, data: DatasetDict[GenOut], max_nr: int, prefix: str, n_axes: int = 4) -> None:
        """
        Plots projections onto pairs of axes from the last n_axes of the basis.
        Each subplot shows a pair (axis_i, axis_{i-1}), starting from the last two axes.
        """
        import math

        hs, _, y = data['eval']
        idx = torch.randperm(hs.size(1))[:max_nr]
        neg_hs, pos_hs, y = hs[0, idx], hs[1, idx], y[idx]

        if probe.direction_type == 'belief':
            eigvals = probe.INTERVENE_LAMBDA
            basis = probe.INTERVENE             # D x r
        elif probe.direction_type == 'normal':
            eigvals = probe.CLASSIFY_LAMBDA
            basis = probe.CLASSIFY              # D x r
        else:
            raise

        basis = eigvals * basis

        r = basis.shape[1]
        n_axes = min(n_axes, r)
        n_pairs = n_axes // 2

        # Prepare axis pairs: (r-1, r-2), (r-3, r-4), ...
        axis_indices = []
        for i in range(n_pairs):
            idx1 = r - 1 - 2 * i
            idx2 = r - 2 - 2 * i
            if idx2 < 0:
                break
            axis_indices.append((idx1, idx2))

        # Set up subplots
        ncols = min(2, n_pairs)
        nrows = math.ceil(n_pairs / ncols)
        fig, axs = plt.subplots(nrows, ncols, figsize=(6 * ncols, 5 * nrows))
        if n_pairs == 1:
            axs = np.array([[axs]])
        elif nrows == 1:
            axs = np.array([axs])
        axs = axs.flatten()

        for i, (ax1, ax2) in enumerate(axis_indices):
            # Project activations onto the selected axes
            pos_coords = pos_hs @ basis[:, [ax1, ax2]]
            neg_coords = neg_hs @ basis[:, [ax1, ax2]]

            coords = torch.cat([pos_coords, neg_coords], dim=0)

            df = pd.DataFrame({
                "axis1": coords[:, 0].tolist(),
                "axis2": coords[:, 1].tolist(),
                "Label": torch.cat([y==1, y==0]).tolist(),
                "Polarity": ['Pos'] * len(y) + ['Neg'] * len(y),
            })

            # If labels are numeric, make them categorical for discrete coloring
            if np.issubdtype(df["Label"].dtype, np.number):
                df["Label"] = df["Label"].astype(int).astype(str)

            ax = axs[i]
            sns.scatterplot(
                data=df, x="axis1", y="axis2", hue="Label",
                style='Polarity', style_order=['Pos', 'Neg'], markers=[r'$+$', r'$-$'],
                s=20, alpha=0.8, linewidth=0, ax=ax,
            )
            ax.set_xlabel(f"Axis {r-ax1}")
            ax.set_ylabel(f"Axis {r-ax2}")
            ax.set_title(f"Projection onto Axes {r-ax1} & {r-ax2}")

            # ---- draw a line between each (pos, neg) pair ----
            pos_np = pos_coords.detach().cpu().numpy()
            neg_np = neg_coords.detach().cpu().numpy()
            segments = np.stack([pos_np, neg_np], axis=1)

            # Only for the first subplot (last two axes), show error via line color
            if i == 0:
                # Project error (pos_hs + neg_hs) onto the last axis (r-1)
                error_vec = pos_hs + neg_hs  # shape: (N, D)
                last_axis_vec = basis[:, r-1]  # shape: (D,)
                error_proj = (error_vec @ last_axis_vec).abs().detach().cpu().numpy()  # shape: (N,)

                norm = mcolors.Normalize(vmin=np.nanmin(error_proj), vmax=np.nanmax(error_proj))
                cmap = plt.get_cmap('viridis')

                lc = LineCollection(segments, linewidths=0.2, alpha=0.4, zorder=0, norm=norm, cmap=cmap)
                lc.set_array(error_proj)
                ax.add_collection(lc)
                ax.autoscale_view()

                # add a colorbar for the line colors
                sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
                sm.set_array([])
                plt.colorbar(sm, ax=ax, pad=0.02, shrink=0.5, anchor=(0.9, 0.0))
            else:
                # For other subplots, draw lines in a neutral color (no error display)
                lc = LineCollection(segments, linewidths=0.2, alpha=0.2, zorder=0, color='gray')
                ax.add_collection(lc)
                ax.autoscale_view()

            ax.legend(bbox_to_anchor=(1.02, 1), loc="upper left")

        # Hide any unused subplots
        for j in range(i+1, len(axs)):
            axs[j].set_visible(False)

        plt.tight_layout()
        plt.savefig(f"local_outputs/cc_act_plot/{prefix}_cc_data_axes_{n_axes}.pdf", bbox_inches="tight")
        plt.close()


@Step.register('char_displ_err_rel')
class CharacterizeDisplacementErrorRelation(Step):
    VERSION = "001"
    CACHEABLE = False

    def run(self, data: GenOut, prefix: str) -> None:
        hs, _, _ = data['eval']
        N, P = hs[0].numpy(), hs[1].numpy()

        D = N - P                                   # (nr_samples, dims)     displacement
        E = N + P                                   # (nr_samples, dims)     error (contrast inconsistency)

        _, s_d, Vt_d = np.linalg.svd(D, full_matrices=False)        # _, (r_d,), (r_d, dims)
        eig_vals_d = s_d**2

        rE = E @ Vt_d.T                                             # (nr_samples, r_d)
        _, s_e, Vt_e = np.linalg.svd(rE, full_matrices=True)        # _, (r_d,), (r_d, r_d)
        eig_vals_e = s_e**2

        print(np.linalg.norm(scipy.linalg.logm(Vt_e)))
        print(np.linalg.norm(np.identity(len(s_e)) - Vt_e))

        min_abscos_per_pc = np.min(np.abs(Vt_e), axis=1)

        plt.hist(min_abscos_per_pc, log=True, bins='auto')  # arguments are passed to np.histogram
        plt.savefig(f'local_outputs/{prefix}_displ_err_min_abs_cos.pdf', bbox_inches="tight")
        plt.close()

        _, new_order = scipy.optimize.linear_sum_assignment(np.abs(Vt_e), maximize=True)
        order = np.argsort(new_order)
        min_rot, eig_vals_e = (Vt_e @ Vt_d)[order], eig_vals_e[order]

        angle_to_axes = np.arccos(np.abs(np.diag(Vt_e[order])))
        print(angle_to_axes.shape)

        plt.hist(angle_to_axes, range=(0, 1.58), bins='auto')  # arguments are passed to np.histogram
        plt.savefig(f'local_outputs/char_displ_err_rel/{prefix}_displ_err_min_angle_hist.pdf', bbox_inches="tight")
        plt.close()


@Step.register('intervene_classify_angle')
class InterveneClassifyAngle(Step):
    VERSION = "010"

    def run(self, probe, nr_dirs: int, prefix: str) -> None:
        basis1: torch.Tensor = probe.INTERVENE
        basis2: torch.Tensor = probe.CLASSIFY

        # Slice first N columns
        b1 = basis1[:, -nr_dirs:].flip(1)    # [D, N]
        b2 = basis2[:, -nr_dirs:].flip(1)    # [D, N]

        # Normalize columns and compute cosine similarities via dot product
        b1 = F.normalize(b1, dim=0)  # column-wise
        b2 = F.normalize(b2, dim=0)
        rads: torch.Tensor = (b1.T @ b2).clamp(-1, 1).abs().acos()
        degs = torch.rad2deg(rads)

        # Save numeric matrix
        degs_np = degs.detach().cpu().numpy()
        csv_path = f"local_outputs/angles/{prefix}_intervene_classify_angle.csv"
        np.savetxt(csv_path, degs_np, delimiter=",", fmt="%.6f")

        # Plot heatmap
        fig, ax = plt.subplots(figsize=(max(4, nr_dirs * 0.45), max(3.5, nr_dirs * 0.4)))
        im = ax.imshow(degs_np, vmin=0, vmax=90, cmap="coolwarm", origin="upper")
        ax.set_title(f"Cosine Similarity • INTERVENE vs CLASSIFY (first {nr_dirs} columns)")
        ax.set_xlabel("CLASSIFY column index")
        ax.set_ylabel("INTERVENE column index")
        ax.set_xticks(range(nr_dirs))
        ax.set_yticks(range(nr_dirs))
        ax.set_xticklabels([str(i) for i in range(nr_dirs)], rotation=90)
        ax.set_yticklabels([str(i) for i in range(nr_dirs)])
        cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        cbar.set_label("cosine")

        fig.tight_layout()
        out_path = f"local_outputs/angles/{prefix}_intervene_classify_cosine_heatmap.png"
        fig.savefig(out_path, dpi=200)
        plt.close(fig)