fig2 notebook cleanup

Author: JesseLivezey

dca/analysis.py

@@ -1,11 +1,15 @@
+import scipy
 import numpy as np
 from sklearn.linear_model import LinearRegression as LR
 from sklearn.decomposition import PCA
 from scipy.stats import special_ortho_group as sog
 
-from .data_util import CrossValidate, form_lag_matrix
+from .base import init_coef
+from .cov_util import calc_pi_from_cross_cov_mats, form_lag_matrix
+from .data_util import CrossValidate
 from .methods_comparison import SlowFeatureAnalysis as SFA
-from .dca import DynamicalComponentsAnalysis
+from .dca import DynamicalComponentsAnalysis as DCA, DynamicalComponentsAnalysisFFT as DCAFFT
+
 
 import warnings
 warnings.simplefilter(action='ignore', category=FutureWarning)
@@ -113,7 +117,7 @@ def run_analysis(X, Y, T_pi_vals, dim_vals, offset_vals, num_cv_folds, decoding_
         Y_test_ctd = Y_test - Y_mean
 
         # compute cross-cov mats for DCA
-        dca_model = DynamicalComponentsAnalysis(T=np.max(T_pi_vals))
+        dca_model = DCA(T=np.max(T_pi_vals))
         dca_model.estimate_data_statistics(X_train_ctd)
 
         # do PCA/SFA
@@ -220,7 +224,7 @@ def run_dim_analysis_dca(X, Y, T_pi, dim_vals, offset, num_cv_folds, decoding_wi
 
         # make DCA object
         # compute cross-cov mats for DCA
-        dca_model = DynamicalComponentsAnalysis(T=T_pi)
+        dca_model = DCA(T=T_pi)
         dca_model.estimate_data_statistics(X_train_ctd)
 
         # loop over dimensionalities
@@ -250,3 +254,145 @@ def run_dim_analysis_dca(X, Y, T_pi, dim_vals, offset, num_cv_folds, decoding_wi
                     null_results[fold_idx, dim_idx, ii] = r2_dca
 
     return results, null_results
+
+
+def gen_pi_heatmap(calc_pi_fn, N=100):
+    theta_vals = np.linspace(0, np.pi, N)
+    phi_vals = np.linspace(0, np.pi, N)
+    heatmap = np.zeros((N, N))
+    for theta_idx in range(N):
+        if theta_idx % 10 == 0:
+            print("theta_idx =", theta_idx)
+        for phi_idx in range(N):
+            theta, phi = theta_vals[theta_idx], phi_vals[phi_idx]
+            x = np.cos(phi) * np.sin(theta)
+            y = np.sin(phi) * np.sin(theta)
+            z = np.cos(theta)
+            V = np.array([x, y, z]).reshape((3, 1))
+            heatmap[theta_idx, phi_idx] = calc_pi_fn(V)
+    return heatmap
+
+
+def make_pi_fn_gp(cross_cov_mats):
+    def calc_pi_fn_gp(V):
+        pi = calc_pi_from_cross_cov_mats(cross_cov_mats, proj=V)
+        return pi
+    return calc_pi_fn_gp
+
+
+def make_pi_fn_knn(X, T_pi, n_jobs=-1):
+    from info_measures.continuous import kraskov_stoegbauer_grassberger as ksg
+
+    def calc_pi_fn_knn(V):
+        X_proj = np.dot(X, V)
+        X_proj_lags = form_lag_matrix(X_proj, 2 * T_pi)
+        mi = ksg.MutualInformation(X_proj_lags[:, :T_pi], X_proj_lags[:, T_pi:], add_noise=True)
+        pi = mi.mutual_information(n_jobs=n_jobs)
+        return pi
+    return calc_pi_fn_knn
+
+
+def random_proj_pi_comparison(calc_pi_fn_1, cal_pi_fn_2, N, d=1,
+                              n_samples=10000, seed=20210412):
+    rng = np.random.RandomState(seed)
+    pi_1, pi_2 = np.zeros(n_samples), np.zeros(n_samples)
+    for i in range(n_samples):
+        if i % 100 == 0:
+            print("sample {} of {}".format(i, n_samples))
+        V = init_coef(N, d, rng=rng, init='random_ortho')
+        pi_1[i] = calc_pi_fn_1(V)
+        pi_2[i] = cal_pi_fn_2(V)
+    pi_12 = np.vstack((pi_1, pi_2)).T  # (n_samples, 2)
+    return pi_12
+
+
+def gp_knn_trajectories(num_traj, cross_cov_mats, X, T_pi, d):
+    f_gp = make_pi_fn_gp(cross_cov_mats)
+    f_knn = make_pi_fn_knn(X, T_pi=T_pi)
+    trajectories = []
+    for traj_idx in range(num_traj):
+        print("traj_idx =", traj_idx)
+        opt = DCA(d=d, T=T_pi)
+        opt.cross_covs = cross_cov_mats
+        opt.fit_projection(record_V=True)
+        V_seq = opt.V_seq
+        num_dca_iter = len(V_seq)
+        pi_gp_knn_traj = np.zeros((num_dca_iter, 2))
+        for i in range(num_dca_iter):
+            if i % 50 == 0:
+                print("{} of {}".format(i, num_dca_iter))
+            V = V_seq[i]
+            pi_gp_knn_traj[i, 0] = f_gp(V)
+            pi_gp_knn_traj[i, 1] = f_knn(V)
+        trajectories.append(pi_gp_knn_traj)
+    return trajectories
+
+
+def dca_deflation(cross_cov_mats, n_proj, n_init=1):
+    N = cross_cov_mats.shape[1]
+    T = cross_cov_mats.shape[0] // 2
+    F = np.eye(N)
+    cov_proj = np.copy(cross_cov_mats)
+    basis = np.zeros((N, n_proj))
+    opt = DCA(T=T)
+    for i in range(n_proj):
+        if i % 10 == 0:
+            print(i)
+        # run DCA
+        opt.cross_covs = cov_proj
+        opt.fit_projection(d=1, n_init=n_init)
+        v = opt.coef_.flatten()
+        # get full-dim v
+        v_full = np.dot(F, v)
+        basis[:, i] = v_full
+        # update U, F, cov_proj
+        U = scipy.linalg.orth(np.eye(N - i) - np.outer(v, v))
+        F = np.dot(F, U)
+        cov_proj = np.array([U.T.dot(C).dot(U) for C in cov_proj])
+    return basis
+
+
+def dca_fft_deflation(X, T, n_proj, n_init=1):
+    N = X.shape[1]
+    F = np.eye(N)
+    X_proj = np.copy(X)
+    basis = np.zeros((N, n_proj))
+    opt = DCAFFT(T=T, d=1)
+    for i in range(n_proj):
+        if i % 10 == 0:
+            print(i)
+        # run DCA
+        opt.fit(X_proj, n_init=n_init)
+        v = opt.coef_.flatten()
+        # get full-dim v
+        v_full = np.dot(F, v)
+        basis[:, i] = v_full
+        # update U, F, X
+        U = scipy.linalg.orth(np.eye(N - i) - np.outer(v, v))
+        F = np.dot(F, U)
+        X_proj = np.dot(X_proj, U)
+    return basis
+
+
+def dca_full(cross_cov_mats, n_proj, n_init=1):
+    T = cross_cov_mats.shape[0] // 2
+    opt = DCA(T=T)
+    opt.cross_covs = cross_cov_mats
+    V_seq = []
+    for i in range(n_proj):
+        if i % 10 == 0:
+            print(i)
+        opt.fit_projection(d=i + 1, n_init=n_init)
+        V = opt.coef_
+        V_seq.append(V)
+    return V_seq
+
+
+def calc_pi_vs_dim(cross_cov_mats, V=None, V_seq=None):
+    if V_seq is None:
+        V_seq = [V[:, :i + 1] for i in range(V.shape[1])]
+    pi_vals = np.zeros(len(V_seq))
+    for i in range(len(V_seq)):
+        V = V_seq[i]
+        pi_vals[i] = calc_pi_from_cross_cov_mats(cross_cov_mats, proj=V)
+    return pi_vals

dca/dca.py

@@ -5,6 +5,7 @@
 from scipy.signal.windows import hann
 
 import torch
+import torch.fft
 import torch.nn.functional as F
 
 from .base import SingleProjectionComponentsAnalysis, ortho_reg_fn, init_coef, ObjectiveWrapper
@@ -292,7 +293,7 @@ def make_cepts2(X, T_pi):
 
     # Compute the power spectral density
     window = torch.Tensor(hann(Y.shape[-1])[np.newaxis, np.newaxis]).type(Y.dtype)
-    Yf = torch.fft.rfft(Y * window, dim=1)
+    Yf = torch.fft.rfft(Y * window, dim=-1)
     spect = abs(Yf)**2
     spect = spect.mean(dim=1)
     spect = torch.cat([torch.flip(spect[:, 1:], dims=(1,)), spect], dim=1)

dca/plotting/__init__.py

@@ -0,0 +1,160 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
+
+from .. import style
+
+
+def make_axes(fig_width, wpad_edge=0, wpad_mid=0.05, hpad_top=0.05, hpad_bottom=0.05,
+              small_sq_width=0.07):
+    sq_width = (1. - 2 * wpad_edge - small_sq_width - 4 * wpad_mid) / 4.
+    sq_height = 1. - hpad_top - hpad_bottom
+    fig_height = sq_width * fig_width / sq_height
+    small_sq_height = small_sq_width * fig_width / fig_height
+    fig = plt.figure(figsize=(fig_width, fig_height))
+    # 2 small squares
+    ax2 = fig.add_axes((wpad_edge, hpad_bottom, small_sq_width, small_sq_height))
+    ax1 = fig.add_axes((wpad_edge, 1. - hpad_top - small_sq_height,
+                        small_sq_width, small_sq_height))
+    # 3 big squares
+    ax3 = fig.add_axes((wpad_edge + small_sq_width + wpad_mid, hpad_bottom, sq_width, sq_height))
+    ax4 = fig.add_axes((wpad_edge + small_sq_width + 2 * wpad_mid + sq_width, hpad_bottom,
+                        sq_width, sq_height))
+    ax5 = fig.add_axes((wpad_edge + small_sq_width + 3 * wpad_mid + 2 * sq_width, hpad_bottom,
+                        sq_width, sq_height))
+    ax6 = fig.add_axes((wpad_edge + small_sq_width + 4 * wpad_mid + 3 * sq_width, hpad_bottom,
+                        sq_width, sq_height))
+    axes = [ax1, ax2, ax3, ax4, ax5, ax6]
+    for ax in axes:
+        ax.set_xticks([])
+        ax.set_yticks([])
+
+    label_dx = -0.02
+    label_dy = 0.05
+    label_y = hpad_bottom + sq_height + label_dy
+    fig.text(wpad_edge + label_dx, label_y,
+             "A", va="bottom", ha="right", color="black",
+             **style.panel_letter_fontstyle)
+    fig.text(wpad_edge + small_sq_width + wpad_mid, label_y,
+             "B", va="bottom", ha="center", color="black",
+             **style.panel_letter_fontstyle)
+    fig.text(wpad_edge + small_sq_width + 2 * wpad_mid + sq_width, label_y,
+             "C", va="bottom", ha="center", color="black",
+             **style.panel_letter_fontstyle)
+    fig.text(wpad_edge + small_sq_width + 3 * wpad_mid + 2 * sq_width, label_y,
+             "D", va="bottom", ha="center", color="black",
+             **style.panel_letter_fontstyle)
+    fig.text(wpad_edge + small_sq_width + 4 * wpad_mid + 3 * sq_width, label_y,
+             "E", va="bottom", ha="center", color="black",
+             **style.panel_letter_fontstyle)
+    return axes
+
+
+def disp_heatmap(ax, heatmap, show_xlabels=True, show_ylabels=True, title=None):
+    N_theta, N_phi = heatmap.shape
+    ax.imshow(heatmap, origin="lower left", cmap="gray", aspect="equal")
+    if show_xlabels:
+        ax.set_xlabel("$\phi$", fontsize=style.axis_label_fontsize, labelpad=-8.5)
+        ax.set_xticks([0, N_phi - 1])
+        ax.set_xticklabels(["0", "$\pi$"], fontsize=style.ticklabel_fontsize)
+    else:
+        ax.set_xticks([])
+    if show_ylabels:
+        ax.set_ylabel("$\\theta$", fontsize=style.axis_label_fontsize, labelpad=-8.5)
+        ax.set_yticks([0, N_theta - 1])
+        ax.set_yticklabels(["0", "$\pi$"], fontsize=style.ticklabel_fontsize)
+    else:
+        ax.set_yticks([])
+    ax.set_xlim([0, heatmap.shape[1] - 1])
+    ax.set_ylim([0, heatmap.shape[0] - 1])
+    if title is not None:
+        ax.set_title(title, fontsize=style.axis_label_fontsize * 0.8, pad=1)
+
+
+def disp_scatter(ax, pi_gp, pi_knn, trajectories=None, diag_text=False,
+                 arrow=True, xlabel="full PI", ylabel="Gaussian PI"):
+    # Note that gp=y and knn=x, but 0 index is gp and 1 is knn in data arrays!
+    traj_color = "#C63F3A"
+    all_gp_vals = [pi_gp]
+    all_knn_vals = [pi_knn]
+    if trajectories is not None:
+        all_gp_vals += [traj[:, 0] for traj in trajectories]
+        all_knn_vals += [traj[:, 1] for traj in trajectories]
+    all_gp_vals = np.concatenate(all_gp_vals)
+    all_knn_vals = np.concatenate(all_knn_vals)
+    min_gp, max_gp = all_gp_vals.min(), all_gp_vals.max()
+    range_gp = max_gp - min_gp
+    min_knn, max_knn = all_knn_vals.min(), all_knn_vals.max()
+    range_knn = max_knn - min_knn
+    pi_gp_norm = (pi_gp - min_gp) / range_gp
+    pi_knn_norm = (pi_knn - min_knn) / range_knn
+    ax.hexbin(pi_knn_norm, pi_gp_norm, gridsize=50, extent=(0, 1, 0, 1),
+              cmap="gray_r", bins="log", linewidth=0.05)
+    if trajectories is not None:
+        for traj_idx in range(len(trajectories)):
+            traj = np.copy(trajectories[traj_idx])
+            traj[:, 0] = (traj[:, 0] - min_gp) / range_gp
+            traj[:, 1] = (traj[:, 1] - min_knn) / range_knn
+            ax.plot(traj[:, 1], traj[:, 0], linewidth=0.5, color=traj_color)
+    ax.set_xlim([0, 1.025])
+    ax.set_ylim([0, 1.025])
+    ax.set_xticks([0, 1])
+    ax.set_xticklabels([0, 1], fontsize=style.ticklabel_fontsize)
+    ax.set_yticks([0, 1])
+    ax.set_yticklabels([0, 1], fontsize=style.ticklabel_fontsize)
+    ax.spines['left'].set_bounds(0, 1)
+    ax.spines['bottom'].set_bounds(0, 1)
+    ax.spines['right'].set_visible(False)
+    ax.spines['top'].set_visible(False)
+    ax.set_xlabel(xlabel, fontsize=style.axis_label_fontsize, labelpad=-9.5)
+    ax.set_ylabel(ylabel, fontsize=style.axis_label_fontsize, labelpad=-8)
+
+    theta_deg = 47
+    if diag_text:
+        ax.text(0.5, 0.65, "DCA trajectories", fontsize=style.ticklabel_fontsize * 0.8,
+                rotation=theta_deg, rotation_mode="anchor", ha="center", va="center",
+                color=traj_color)
+    if arrow:
+        len_x = np.cos(np.deg2rad(theta_deg))
+        len_y = np.sin(np.deg2rad(theta_deg))
+        mag = 0.425
+        ax.quiver(0.475, 0.40, mag * len_x, mag * len_y,
+                  angles='xy', scale_units='xy', scale=1, width=0.015,
+                  color=traj_color)
+
+
+def plot_deflation_results(ax, pi_regular, pi_def, pi_fft):
+    dim_vals = np.arange(len(pi_def) + 1)
+    pi_vals = (pi_regular, pi_def, pi_fft)
+    labels = ["DCA", "deflation", "FFT deflation"]
+    markersize = 1.5
+    colors = ["#C63F3A", "gray", "black"]
+    for i in range(len(pi_vals)):
+        if i < 2:
+            ax.plot(dim_vals, [0] + list(pi_vals[i]), label=labels[i],
+                    linewidth=0.85, color=colors[i], linestyle="-")
+        else:
+            ax.plot(dim_vals, [0] + list(pi_vals[i]), label=labels[i],
+                    linewidth=0, marker=".", markersize=markersize,
+                    color=colors[i])
+
+    ax.legend(fontsize=style.ticklabel_fontsize * 0.8, frameon=False,
+              labelspacing=0.1, bbox_to_anchor=(0.2, 0, 1, 1))
+    ax.set_xlabel("dimension", fontsize=style.axis_label_fontsize, labelpad=-9.5)
+    ax.set_ylabel("PI (nats)", fontsize=style.axis_label_fontsize, labelpad=-13)
+    ax.spines["right"].set_visible(False)
+    ax.spines["top"].set_visible(False)
+
+    max_dim = len(pi_regular)
+    max_dim_padded = max_dim * 1.025
+    ax.set_xticks([0, max_dim])
+    ax.set_xticklabels([0, max_dim], fontsize=style.ticklabel_fontsize)
+    ax.set_xlim([0, max_dim_padded])
+    ax.spines["bottom"].set_bounds(0, max_dim)
+
+    max_pi = np.max(np.concatenate(pi_vals))
+    max_pi_padded = max_pi * 1.025
+    ax.set_yticks([0, max_pi])
+    ax.set_yticklabels([0, np.round(max_pi, 1)], fontsize=style.ticklabel_fontsize)
+    ax.set_ylim([0, max_pi_padded])
+    ax.spines["left"].set_bounds(0, max_pi)

dca/plotting.py dca/plotting/fig1.pydca/plotting.py renamed to dca/plotting/fig1.py

@@ -13,7 +13,7 @@
     "from dca import DynamicalComponentsAnalysis as DCA, style\n",
     "from dca.cov_util import calc_cross_cov_mats_from_data\n",
     "from dca.synth_data import embedded_lorenz_cross_cov_mats, gen_lorenz_data, random_basis, median_subspace\n",
-    "from dca.plotting import lorenz_fig_axes, plot_3d, plot_lorenz_3d, plot_traces, plot_dca_demo, plot_r2, plot_cov\n",
+    "from dca.plotting.fig1 import lorenz_fig_axes, plot_3d, plot_lorenz_3d, plot_traces, plot_dca_demo, plot_r2, plot_cov\n",
     "\n",
     "RESULTS_FILENAME = \"lorenz_results.hdf5\""
    ]
@@ -283,5 +283,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }