Removed GPU acceleration

eks/core.py

@@ -528,166 +528,6 @@ def single_timestep_nll(innovation, innovation_cov):
     nll_increment = 0.5 * jnp.abs(log_det_S + quadratic_term + c)
     return nll_increment
 
-
-# ----- Parallel Functions for GPU -----
-
-def first_filtering_element(C, A, Q, R, m0, P0, y):
-    # model.F = A, model.H = C,
-    S = C @ Q @ C.T + R
-    CF, low = jsc.linalg.cho_factor(S)  # note the jsc
-
-    m1 = A @ m0
-    P1 = A @ P0 @ A.T + Q
-    S1 = C @ P1 @ C.T + R
-    K1 = jsc.linalg.solve(S1, C @ P1, assume_a='pos').T  # note the jsc
-
-    A_updated = jnp.zeros_like(A)
-    b = m1 + K1 @ (y - C @ m1)
-    C_updated = P1 - K1 @ S1 @ K1.T
-
-    # note the jsc
-    eta = A.T @ C.T @ jsc.linalg.cho_solve((CF, low), y)
-    J = A.T @ C.T @ jsc.linalg.cho_solve((CF, low), C @ A)
-    return A_updated, b, C_updated, J, eta
-
-
-def generic_filtering_element(C, A, Q, R, y):
-    S = C @ Q @ C.T + R
-    CF, low = jsc.linalg.cho_factor(S)  # note the jsc
-    K = jsc.linalg.cho_solve((CF, low), C @ Q).T  # note the jsc
-    A_updated = A - K @ C @ A
-    b = K @ y
-    C_updated = Q - K @ C @ Q
-
-    # note the jsc
-    eta = A.T @ C.T @ jsc.linalg.cho_solve((CF, low), y)
-    J = A.T @ C.T @ jsc.linalg.cho_solve((CF, low), C @ A)
-    return A_updated, b, C_updated, J, eta
-
-
-def make_associative_filtering_elements(C, A, Q, R, m0, P0, observations):
-    first_elems = first_filtering_element(C, A, Q, R, m0, P0, observations[0])
-    generic_elems = vmap(lambda o: generic_filtering_element(C, A, Q, R, o))(observations[1:])
-    return tuple(jnp.concatenate([jnp.expand_dims(first_e, 0), gen_es])
-                 for first_e, gen_es in zip(first_elems, generic_elems))
-
-
-@partial(vmap)
-def filtering_operator(elem1, elem2):
-    # # note the jsc everywhere
-    A1, b1, C1, J1, eta1 = elem1
-    A2, b2, C2, J2, eta2 = elem2
-    dim = A1.shape[0]
-    I_var = jnp.eye(dim)  # note the jnp
-
-    I_C1J2 = I_var + C1 @ J2
-    temp = jsc.linalg.solve(I_C1J2.T, A2.T).T
-    A = temp @ A1
-    b = temp @ (b1 + C1 @ eta2) + b2
-    C = temp @ C1 @ A2.T + C2
-
-    I_J2C1 = I_var + J2 @ C1
-    temp = jsc.linalg.solve(I_J2C1.T, A1).T
-
-    eta = temp @ (eta2 - J2 @ b1) + eta1
-    J = temp @ J2 @ A1 + J1
-
-    return A, b, C, J, eta
-
-
-def pkf(y, m0, cov0, A, Q, C, R):
-    initial_elements = make_associative_filtering_elements(C, A, Q, R, m0, cov0, y)
-    final_elements = associative_scan(filtering_operator, initial_elements)
-    return final_elements
-
-
-pkf_func = jit(pkf)
-
-
-def get_kalman_means(A_scan, b_scan, m0):
-    """
-    Computes the Kalman mean at a single timepoint, the result is:
-    A_scan @ m0 + b_scan
-
-    Returned shape: (state_dimension, 1)
-    """
-    return A_scan @ jnp.expand_dims(m0, axis=1) + jnp.expand_dims(b_scan, axis=1)
-
-
-def get_kalman_variances(C):
-    return C
-
-
-def get_next_cov(A, C, Q, R, filter_cov, filter_mean):
-    """
-    Given the moments of p(x_t | y_1, ..., y_t) (normal filter distribution),
-    compute the moments of the distribution for:
-    p(y_{t+1} | y_1, ..., y_t)
-
-    Params:
-        A (np.ndarray): Shape (state_dimension, state_dimension) Process coeff matrix
-        C (np.ndarray): Shape (obs_dimension, state_dimension) Observation coeff matrix
-        Q (np.ndarray): Shape (state_dimension, state_dimension). Process noise covariance matrix.
-        R (np.ndarray): Shape (obs_dimension, obs_dimension). Observation noise covariance matrix.
-        filter_cov (np.ndarray). Shape (state_dimension, state_dimension). Filtered covariance
-        filter_mean (np.ndarray). Shape (state_dimension, 1). Filter mean
-
-    Returns:
-        mean (np.ndarray). Shape (obs_dimension, 1)
-        cov (np.ndarray). Shape (obs_dimension, obs_dimension).
-    """
-    mean = C @ A @ filter_mean
-    cov = C @ (A @ filter_cov @ A.T + Q) @ C.T + R
-    return mean, cov
-
-
-def compute_marginal_nll(value, mean, covariance):
-    return -1 * jax.scipy.stats.multivariate_normal.logpdf(value, mean, covariance)
-
-
-def parallel_loss_single(A_scan, b_scan, C_scan, A, C, Q, R, next_observation, m0):
-    curr_mean = get_kalman_means(A_scan, b_scan, m0)
-    curr_cov = get_kalman_variances(C_scan)  # Placeholder; just returns identity
-
-    next_mean, next_cov = get_next_cov(A, C, Q, R, curr_cov, curr_mean)
-    return jnp.squeeze(curr_mean), curr_cov, compute_marginal_nll(jnp.squeeze(next_observation),
-                                                                  jnp.squeeze(next_mean), next_cov)
-
-
-parallel_loss_func_vmap = jit(
-    vmap(parallel_loss_single, in_axes=(0, 0, 0, None, None, None, None, 0, None),
-         out_axes=(0, 0, 0)))
-
-
-@partial(jit)
-def y1_given_x0_nll(C, A, Q, R, m0, cov0, obs):
-    y1_predictive_mean = C @ A @ jnp.expand_dims(m0, axis=1)
-    y1_predictive_cov = C @ (A @ cov0 @ A.T + Q) @ C.T + R
-    addend = -1 * jax.scipy.stats.multivariate_normal.logpdf(obs, jnp.squeeze(y1_predictive_mean),
-                                                             y1_predictive_cov)
-    return addend
-
-
-def pkf_and_loss(y, m0, cov0, A, Q, C, R):
-    A_scan, b_scan, C_scan, _, _ = pkf_func(y, m0, cov0, A, Q, C, R)
-
-    # Gives us the NLL for p(y_i | y_1, ..., y_{i-1}) for i > 1.
-    # Need to use the parallel scan outputs for this. i = 1 handled below
-    filtered_states, filtered_covariances, losses = parallel_loss_func_vmap(A_scan[:-1],
-                                                                            b_scan[:-1],
-                                                                            C_scan[:-1], A, C, Q,
-                                                                            R, y[1:], m0)
-
-    # Gives us the NLL for p_y(y_1 | x_0)
-    addend = y1_given_x0_nll(C, A, Q, R, m0, cov0, y[0])
-
-    final_mean = get_kalman_means(A_scan[-1], b_scan[-1], m0).T
-    final_covariance = jnp.expand_dims(get_kalman_variances(C_scan[-1]), axis=0)
-    filtered_states = jnp.concatenate([filtered_states, final_mean], axis=0)
-    filtered_variances = jnp.concatenate([filtered_covariances, final_covariance], axis=0)
-    return filtered_states, filtered_variances, jnp.sum(losses) + addend
-
-
 # -------------------------------------------------------------------------------------
 # Misc: These miscellaneous functions generally have specific computations used by the
 # core functions or the smoothers

eks/ibl_paw_multiview_smoother.py

@@ -6,7 +6,6 @@
 from eks.core import backward_pass, eks_zscore, ensemble, forward_pass
 from eks.utils import make_dlc_pandas_index
 
-
 # TODO:
 # - allow conf_weighted_mean for ensemble variance computation
 

eks/ibl_pupil_smoother.py

@@ -6,8 +6,8 @@
 import pandas as pd
 from scipy.optimize import minimize
 
-from eks.core import backward_pass, compute_nll, eks_zscore, ensemble, forward_pass
-from eks.utils import crop_frames, make_dlc_pandas_index, format_data
+from eks.core import backward_pass, compute_nll, ensemble, forward_pass
+from eks.utils import crop_frames, format_data, make_dlc_pandas_index
 
 
 def get_pupil_location(dlc):
@@ -276,7 +276,8 @@ def ensemble_kalman_smoother_ibl_pupil(
         # compute zscore for EKS to see how it deviates from the ensemble
         # eks_predictions = \
         #     np.asarray([processed_arr_dict[key_pair[0]], processed_arr_dict[key_pair[1]]]).T
-        # ensemble_preds_curr = ensemble_preds[:, ensemble_indices[i][0]: ensemble_indices[i][1] + 1]
+        # ensemble_preds_curr = ensemble_preds
+        #                                   [:, ensemble_indices[i][0]: ensemble_indices[i][1] + 1]
         # ensemble_vars_curr = ensemble_vars[:, ensemble_indices[i][0]: ensemble_indices[i][1] + 1]
         # zscore, _ = eks_zscore(
         #     eks_predictions,