Repository: jameschapman19/cca_zoo (language: python)

Issue:
Fixed test for probabilistic
Fix as recommended in issue


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 9e166d43..db39fb7a 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,9 @@
-[![DOI](https://zenodo.org/badge/303801602.svg)](https://zenodo.org/badge/latestdoi/303801602)
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.5748062.svg)](https://doi.org/10.5281/zenodo.4382739)
 [![codecov](https://codecov.io/gh/jameschapman19/cca_zoo/branch/master/graph/badge.svg?token=JHG9VUB0L8)](https://codecov.io/gh/jameschapman19/cca_zoo)
 ![Build Status](https://github.com/jameschapman19/cca_zoo/actions/workflows/python-package.yml/badge.svg)
 [![Documentation Status](https://readthedocs.org/projects/cca-zoo/badge/?version=latest)](https://cca-zoo.readthedocs.io/en/latest/?badge=latest)
 [![version](https://img.shields.io/pypi/v/cca-zoo)](https://pypi.org/project/cca-zoo/)
 [![downloads](https://img.shields.io/pypi/dm/cca-zoo)](https://pypi.org/project/cca-zoo/)
-[![Scrutinizer Code Quality](https://scrutinizer-ci.com/g/jameschapman19/cca_zoo/badges/quality-score.png?b=master)](https://scrutinizer-ci.com/g/jameschapman19/cca_zoo/?branch=master)
 
 # Installation
 Note: for standard install use: 
@@ -34,9 +33,8 @@ In case this work is used as part of research I attach a DOI bibtex entry:
   month        = jun,
   year         = 2021,
   publisher    = {Zenodo},
-  version      = {v1.6.1},
-  doi          = {10.5281/zenodo.4925892},
-  url          = {https://doi.org/10.5281/zenodo.4925892}
+  doi          = {10.5281/zenodo.4382739},
+  url          = {https://doi.org/10.5281/zenodo.4382739}
 }
 ```
 
diff --git a/cca_zoo/data/simulated.py b/cca_zoo/data/simulated.py
index eb954d91..a32b92ab 100644
--- a/cca_zoo/data/simulated.py
+++ b/cca_zoo/data/simulated.py
@@ -6,20 +6,20 @@
 from scipy.linalg import block_diag
 from sklearn.utils.validation import check_random_state
 
-from ..utils.check_values import _process_parameter
+from ..utils import _process_parameter
 
 
 def generate_covariance_data(
-        n: int,
-        view_features: List[int],
-        latent_dims: int = 1,
-        view_sparsity: List[Union[int, float]] = None,
-        correlation: Union[List[float], float] = 1,
-        structure: Union[str, List[str]] = None,
-        sigma: Union[List[float], float] = None,
-        decay: float = 0.5,
-        positive=None,
-        random_state: Union[int, np.random.RandomState] = None,
+    n: int,
+    view_features: List[int],
+    latent_dims: int = 1,
+    view_sparsity: List[Union[int, float]] = None,
+    correlation: Union[List[float], float] = 1,
+    structure: Union[str, List[str]] = None,
+    sigma: Union[List[float], float] = None,
+    decay: float = 0.5,
+    positive=None,
+    random_state: Union[int, np.random.RandomState] = None,
 ):
     """
     Function to generate CCA dataset with defined population correlations
@@ -58,7 +58,7 @@ def generate_covariance_data(
             covs = []
             true_features = []
             for view_p, sparsity, view_structure, view_positive, view_sigma in zip(
-                    view_features, view_sparsity, structure, positive, sigma
+                view_features, view_sparsity, structure, positive, sigma
             ):
                 # Covariance Bit
                 if view_structure == "identity":
@@ -110,12 +110,12 @@ def generate_covariance_data(
                     # Cross Bit
                     cross += covs[i] @ A @ covs[j]
                 cov[
-                splits[i]: splits[i] + view_features[i],
-                splits[j]: splits[j] + view_features[j],
+                    splits[i] : splits[i] + view_features[i],
+                    splits[j] : splits[j] + view_features[j],
                 ] = cross
                 cov[
-                splits[j]: splits[j] + view_features[j],
-                splits[i]: splits[i] + view_features[i],
+                    splits[j] : splits[j] + view_features[j],
+                    splits[i] : splits[i] + view_features[i],
                 ] = cross.T
 
             X = np.zeros((n, sum(view_features)))
@@ -130,12 +130,12 @@ def generate_covariance_data(
 
 
 def generate_simple_data(
-        n: int,
-        view_features: List[int],
-        view_sparsity: List[Union[int, float]] = None,
-        eps: float = 0,
-        transform=True,
-        random_state=None,
+    n: int,
+    view_features: List[int],
+    view_sparsity: List[Union[int, float]] = None,
+    eps: float = 0,
+    transform=True,
+    random_state=None,
 ):
     """
     Simple latent variable model to generate data with one latent factor
@@ -196,9 +196,9 @@ def _gaussian(x, mu, sig, dn):
     :param dn:
     """
     return (
-            np.exp(-np.power(x - mu, 2.0) / (2 * np.power(sig, 2.0)))
-            * dn
-            / (np.sqrt(2 * np.pi) * sig)
+        np.exp(-np.power(x - mu, 2.0) / (2 * np.power(sig, 2.0)))
+        * dn
+        / (np.sqrt(2 * np.pi) * sig)
     )
 
 
diff --git a/cca_zoo/data/toy.py b/cca_zoo/data/toy.py
index 68b904e6..7c5aebc8 100644
--- a/cca_zoo/data/toy.py
+++ b/cca_zoo/data/toy.py
@@ -15,7 +15,7 @@ class Split_MNIST_Dataset(Dataset):
     """
 
     def __init__(
-            self, mnist_type: str = "MNIST", train: bool = True, flatten: bool = True
+        self, mnist_type: str = "MNIST", train: bool = True, flatten: bool = True
     ):
         """
 
@@ -71,7 +71,7 @@ class Noisy_MNIST_Dataset(Dataset):
     """
 
     def __init__(
-            self, mnist_type: str = "MNIST", train: bool = True, flatten: bool = True
+        self, mnist_type: str = "MNIST", train: bool = True, flatten: bool = True
     ):
         """
 
diff --git a/cca_zoo/deepmodels/dvcca.py b/cca_zoo/deepmodels/dvcca.py
index e50728de..b8b12696 100644
--- a/cca_zoo/deepmodels/dvcca.py
+++ b/cca_zoo/deepmodels/dvcca.py
@@ -23,11 +23,11 @@ class DVCCA(_DCCA_base):
     """
 
     def __init__(
-            self,
-            latent_dims: int,
-            encoders=None,
-            decoders=None,
-            private_encoders: Iterable[BaseEncoder] = None,
+        self,
+        latent_dims: int,
+        encoders=None,
+        decoders=None,
+        private_encoders: Iterable[BaseEncoder] = None,
     ):
         """
         :param latent_dims: # latent dimensions
diff --git a/cca_zoo/deepmodels/trainers.py b/cca_zoo/deepmodels/trainers.py
index 2192a113..328e1cfb 100644
--- a/cca_zoo/deepmodels/trainers.py
+++ b/cca_zoo/deepmodels/trainers.py
@@ -11,10 +11,10 @@
 
 class CCALightning(LightningModule):
     def __init__(
-            self,
-            model: _DCCA_base,
-            optimizer: torch.optim.Optimizer = None,
-            lr_scheduler: torch.optim.lr_scheduler = None,
+        self,
+        model: _DCCA_base,
+        optimizer: torch.optim.Optimizer = None,
+        lr_scheduler: torch.optim.lr_scheduler = None,
     ):
         """
 
@@ -78,9 +78,9 @@ def on_validation_epoch_end(self, unused: Optional = None) -> None:
             self.log("val corr", score)
 
     def correlations(
-            self,
-            loader: torch.utils.data.DataLoader,
-            train=False,
+        self,
+        loader: torch.utils.data.DataLoader,
+        train=False,
     ):
         """
 
@@ -92,16 +92,16 @@ def correlations(
             return None
         all_corrs = []
         for x, y in itertools.product(transformed_views, repeat=2):
-            all_corrs.append(np.diag(np.corrcoef(x.T, y.T)[: x.shape[1], y.shape[1]:]))
+            all_corrs.append(np.diag(np.corrcoef(x.T, y.T)[: x.shape[1], y.shape[1] :]))
         all_corrs = np.array(all_corrs).reshape(
             (len(transformed_views), len(transformed_views), -1)
         )
         return all_corrs
 
     def transform(
-            self,
-            loader: torch.utils.data.DataLoader,
-            train=False,
+        self,
+        loader: torch.utils.data.DataLoader,
+        train=False,
     ):
         """
 
@@ -124,9 +124,9 @@ def transform(
         return z_list
 
     def score(
-            self,
-            loader: torch.utils.data.DataLoader,
-            train=False,
+        self,
+        loader: torch.utils.data.DataLoader,
+        train=False,
     ):
         """
 
@@ -141,13 +141,13 @@ def score(
         n_views = pair_corrs.shape[0]
         # sum all the pairwise correlations for each dimension. Subtract the self correlations. Divide by the number of views. Gives average correlation
         dim_corrs = (
-                            pair_corrs.sum(axis=tuple(range(pair_corrs.ndim - 1))) - n_views
-                    ) / (n_views ** 2 - n_views)
+            pair_corrs.sum(axis=tuple(range(pair_corrs.ndim - 1))) - n_views
+        ) / (n_views ** 2 - n_views)
         return dim_corrs
 
     def recon(
-            self,
-            loader: torch.utils.data.DataLoader,
+        self,
+        loader: torch.utils.data.DataLoader,
     ):
         with torch.no_grad():
             for batch_idx, (data, label) in enumerate(loader):
diff --git a/cca_zoo/deepmodels/utils.py b/cca_zoo/deepmodels/utils.py
index 8c04edc3..8e23e120 100644
--- a/cca_zoo/deepmodels/utils.py
+++ b/cca_zoo/deepmodels/utils.py
@@ -8,11 +8,11 @@
 
 
 def process_data(
-        dataset: Union[torch.utils.data.Dataset, Iterable[np.ndarray]],
-        val_dataset: Union[torch.utils.data.Dataset, Iterable[np.ndarray]] = None,
-        labels=None,
-        val_labels=None,
-        val_split: float = 0,
+    dataset: Union[torch.utils.data.Dataset, Iterable[np.ndarray]],
+    val_dataset: Union[torch.utils.data.Dataset, Iterable[np.ndarray]] = None,
+    labels=None,
+    val_labels=None,
+    val_split: float = 0,
 ):
     # Ensure datasets are in the right form (e.g. if numpy arrays are passed turn them into
     if isinstance(dataset, tuple):
@@ -29,7 +29,7 @@ def process_data(
 
 
 def get_dataloaders(
-        dataset, val_dataset=None, batch_size=None, val_batch_size=None, num_workers=0
+    dataset, val_dataset=None, batch_size=None, val_batch_size=None, num_workers=0
 ):
     if batch_size is None:
         batch_size = len(dataset)
diff --git a/cca_zoo/models/__init__.py b/cca_zoo/models/__init__.py
index b8fade1e..0aafdd8e 100644
--- a/cca_zoo/models/__init__.py
+++ b/cca_zoo/models/__init__.py
@@ -1,5 +1,5 @@
-from .cca_base import _CCA_Base
 from .rcca import rCCA, CCA, PLS
+from .mcca import MCCA, KCCA
 from .gcca import GCCA, KGCCA
 from .iterative import (
     _Iterative,
@@ -13,6 +13,6 @@
     SpanCCA,
     SWCCA,
 )
-from .mcca import MCCA, KCCA
 from .ncca import NCCA
+from .partialcca import PartialCCA
 from .tcca import TCCA, KTCCA
diff --git a/cca_zoo/models/cca_base.py b/cca_zoo/models/_cca_base.py
similarity index 98%
rename from cca_zoo/models/cca_base.py
rename to cca_zoo/models/_cca_base.py
index f0294084..4c3437d6 100644
--- a/cca_zoo/models/cca_base.py
+++ b/cca_zoo/models/_cca_base.py
@@ -8,7 +8,8 @@
 from sklearn.utils.sparsefuncs import mean_variance_axis
 from sklearn.utils.validation import check_random_state, check_is_fitted
 
-from cca_zoo.utils import check_views, plot_latent_train_test
+from cca_zoo.utils.check_values import _check_views
+from cca_zoo.utils.plotting import plot_latent_train_test
 
 
 class _CCA_Base(BaseEstimator, MultiOutputMixin, RegressorMixin):
@@ -66,7 +67,7 @@ def transform(self, views: Iterable[np.ndarray], y=None, **kwargs):
         :param kwargs: any additional keyword arguments required by the given model
         """
         check_is_fitted(self, attributes=["weights"])
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale_transform(views)
diff --git a/cca_zoo/models/innerloop.py b/cca_zoo/models/_innerloop.py
similarity index 77%
rename from cca_zoo/models/innerloop.py
rename to cca_zoo/models/_innerloop.py
index 39f31f7b..1ceec0f0 100644
--- a/cca_zoo/models/innerloop.py
+++ b/cca_zoo/models/_innerloop.py
@@ -11,7 +11,7 @@
 from sklearn.utils._testing import ignore_warnings
 from sklearn.utils.validation import check_random_state
 
-from ..utils.check_values import (
+from cca_zoo.utils.check_values import (
     _check_converged_weights,
     _check_Parikh2014,
     _process_parameter,
@@ -20,18 +20,16 @@
 
 class _InnerLoop:
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol: float = 1e-9,
-            initialization: str = "unregularized",
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol: float = 1e-9,
+        random_state=None,
     ):
         """
         :param max_iter: maximum number of iterations to perform if tol is not reached
         :param tol: tolerance value used for stopping criteria
-        :param initialization: initialise the optimisation with either the 'unregularized' (CCA/PLS) solution, or a 'random' initialisation
         """
-        self.initialization = initialization
+        self.track = {"converged": False, "objective": []}
         self.max_iter = max_iter
         self.tol = tol
         self.random_state = check_random_state(random_state)
@@ -42,51 +40,15 @@ def _check_params(self):
         """
         pass
 
-    def _initialize(self):
-        if self.initialization == "random":
-            self.scores = np.array(
-                [
-                    self.random_state.normal(0, 1, size=(view.shape[0], 1))
-                    for view in self.views
-                ]
-            )
-        elif self.initialization == "uniform":
-            self.scores = np.array([np.ones((view.shape[0], 1)) for view in self.views])
-        elif self.initialization == "unregularized":
-            self.scores = (
-                PLSInnerLoop(
-                    initialization="random",
-                    random_state=self.random_state,
-                    tol=self.tol,
-                )
-                    ._fit(*self.views)
-                    .scores
-            )
-        elif callable(self.initialization):
-            self.scores = next(self.initialization())
-        else:
-            raise ValueError("initialize must be random, uniform or unregularized")
-        self.scores = (
-                self.scores
-                * np.sqrt(self.n - 1)
-                / np.linalg.norm(self.scores, axis=1)[:, np.newaxis]
-        )
-        self.weights = [
-            self.random_state.randn(view.shape[1], 1) for view in self.views
-        ]
-
-    def _fit(self, *views: np.ndarray):
+    def fit(self, *views: np.ndarray, initial_scores):
+        self.scores = initial_scores
         self.views = views
         self.n = views[0].shape[0]
 
         # Check that the parameters that have been passed are valid for these views given #views and #features
         self._check_params()
-        self._initialize()
-
-        self.track = {}
-        self.track["converged"] = False
+        self.weights = [self.random_state.randn(view.shape[1]) for view in self.views]
         # Iterate until convergence
-        self.track["objective"] = []
         for _ in range(self.max_iter):
             self._inner_iteration()
             if np.isnan(self.scores).sum() > 0:
@@ -122,18 +84,16 @@ def _objective(self) -> int:
         return obj.item()
 
 
-class PLSInnerLoop(_InnerLoop):
+class _PLSInnerLoop(_InnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-9,
-            initialization: str = "unregularized",
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-9,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
             tol=tol,
-            initialization=initialization,
             random_state=random_state,
         )
 
@@ -161,36 +121,36 @@ def _update_view(self, view_index: int):
         targets = np.ma.array(self.scores, mask=False)
         targets.mask[view_index] = True
         self.weights[view_index] = (
-                self.views[view_index].T @ targets.sum(axis=0).filled()
+            self.views[view_index].T @ targets.sum(axis=0).filled()
         )
         self.weights[view_index] /= np.linalg.norm(self.weights[view_index])
-        self.scores[view_index] = self.views[view_index] @ self.weights[view_index]
+        self.scores[view_index] = self.views[view_index] @ np.squeeze(
+            np.array(self.weights[view_index])
+        )
 
     def _early_stop(self) -> bool:
         # Some kind of early stopping
         if all(
-                _cosine_similarity(self.scores[n], self.old_scores[n]) > (1 - self.tol)
-                for n, view in enumerate(self.scores)
+            _cosine_similarity(self.scores[n], self.old_scores[n]) > (1 - self.tol)
+            for n, view in enumerate(self.scores)
         ):
             return True
         else:
             return False
 
 
-class PMDInnerLoop(PLSInnerLoop):
+class _PMDInnerLoop(_PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-9,
-            initialization: str = "unregularized",
-            c=None,
-            positive=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-9,
+        c=None,
+        positive=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
             tol=tol,
-            initialization=initialization,
             random_state=random_state,
         )
         self.c = c
@@ -202,18 +162,13 @@ def _check_params(self):
                 "c parameter not set. Setting to c=1 i.e. maximum regularisation of l1 norm"
             )
         self.c = _process_parameter("c", self.c, 1, len(self.views))
-        if any(c < 1 for c in self.c):
+        if any(c < 0 or c > 1 for c in self.c):
             raise ValueError(
-                "All regulariation parameters should be at least " f"1. c=[{self.c}]"
+                "All regularisation parameters should be between 0 and 1 "
+                f"1. c=[{self.c}]"
             )
         shape_sqrts = [np.sqrt(view.shape[1]) for view in self.views]
-        if any(c > shape_sqrt for c, shape_sqrt in zip(self.c, shape_sqrts)):
-            raise ValueError(
-                "All regulariation parameters should be less than"
-                " the square root of number of the respective"
-                f" view. c=[{self.c}], limit of each view: "
-                f"{shape_sqrts}"
-            )
+        self.t = [max(1, x * y) for x, y in zip(self.c, shape_sqrts)]
         self.positive = _process_parameter(
             "positive", self.positive, False, len(self.views)
         )
@@ -227,11 +182,11 @@ def _update_view(self, view_index: int):
         targets = np.ma.array(self.scores, mask=False)
         targets.mask[view_index] = True
         self.weights[view_index] = (
-                self.views[view_index].T @ targets.sum(axis=0).filled()
+            self.views[view_index].T @ targets.sum(axis=0).filled()
         )
         self.weights[view_index] = _delta_search(
             self.weights[view_index],
-            self.c[view_index],
+            self.t[view_index],
             positive=self.positive[view_index],
             tol=self.tol,
         )
@@ -239,19 +194,17 @@ def _update_view(self, view_index: int):
         self.scores[view_index] = self.views[view_index] @ self.weights[view_index]
 
 
-class ParkhomenkoInnerLoop(PLSInnerLoop):
+class _ParkhomenkoInnerLoop(_PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-9,
-            initialization: str = "unregularized",
-            c=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-9,
+        c=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
             tol=tol,
-            initialization=initialization,
             random_state=random_state,
         )
         self.c = c
@@ -278,23 +231,21 @@ def _update_view(self, view_index: int):
         self.scores[view_index] = self.views[view_index] @ self.weights[view_index]
 
 
-class ElasticInnerLoop(PLSInnerLoop):
+class _ElasticInnerLoop(_PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-9,
-            initialization: str = "unregularized",
-            c=None,
-            l1_ratio=None,
-            maxvar=True,
-            stochastic=True,
-            positive=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-9,
+        c=None,
+        l1_ratio=None,
+        maxvar=True,
+        stochastic=True,
+        positive=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
             tol=tol,
-            initialization=initialization,
             random_state=random_state,
         )
         self.stochastic = stochastic
@@ -356,14 +307,14 @@ def _update_view(self, view_index: int):
         if not self.maxvar:
             _check_converged_weights(self.weights[view_index], view_index)
             self.weights[view_index] = self.weights[view_index] / (
-                    np.linalg.norm(self.views[view_index] @ self.weights[view_index])
-                    / np.sqrt(self.n)
+                np.linalg.norm(self.views[view_index] @ self.weights[view_index])
+                / np.sqrt(self.n)
             )
         self.scores[view_index] = self.views[view_index] @ self.weights[view_index]
 
     @ignore_warnings(category=ConvergenceWarning)
     def _elastic_solver(self, X, y, view_index):
-        return np.expand_dims(self.regressors[view_index].fit(X, y.ravel()).coef_, 1)
+        return self.regressors[view_index].fit(X, y.ravel()).coef_
 
     def _objective(self):
         views = len(self.views)
@@ -392,22 +343,20 @@ def _early_stop(self) -> bool:
             return False
 
 
-class ADMMInnerLoop(ElasticInnerLoop):
+class _ADMMInnerLoop(_ElasticInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-9,
-            initialization: str = "unregularized",
-            mu=None,
-            lam=None,
-            c=None,
-            eta=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-9,
+        mu=None,
+        lam=None,
+        c=None,
+        eta=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
             tol=tol,
-            initialization=initialization,
             random_state=random_state,
         )
         self.c = c
@@ -433,9 +382,9 @@ def _check_params(self):
         _check_Parikh2014(self.mu, self.lam, self.views)
 
         self.eta = [
-            np.ones((view.shape[0], 1)) * eta for view, eta in zip(self.views, self.eta)
+            np.ones(view.shape[0]) * eta for view, eta in zip(self.views, self.eta)
         ]
-        self.z = [np.zeros((view.shape[0], 1)) for view in self.views]
+        self.z = [np.zeros(view.shape[0]) for view in self.views]
         self.l1_ratio = [1] * len(self.views)
 
     def _update_view(self, view_index: int):
@@ -461,9 +410,9 @@ def _update_view(self, view_index: int):
                 / lam
                 * self.views[view_index].T
                 @ (
-                        self.views[view_index] @ self.weights[view_index]
-                        - self.z[view_index]
-                        + self.eta[view_index]
+                    self.views[view_index] @ self.weights[view_index]
+                    - self.z[view_index]
+                    + self.eta[view_index]
                 ),
                 mu,
                 gradient,
@@ -479,9 +428,9 @@ def _update_view(self, view_index: int):
                 self.views[view_index] @ self.weights[view_index] + self.eta[view_index]
             )
             self.eta[view_index] = (
-                    self.eta[view_index]
-                    + self.views[view_index] @ self.weights[view_index]
-                    - self.z[view_index]
+                self.eta[view_index]
+                + self.views[view_index] @ self.weights[view_index]
+                - self.z[view_index]
             )
             norm_eta.append(np.linalg.norm(self.eta[view_index]))
             norm_proj.append(
@@ -511,22 +460,20 @@ def _prox_lam_g(self, x):
             return x / max(1, norm)
 
 
-class SpanCCAInnerLoop(_InnerLoop):
+class _SpanCCAInnerLoop(_InnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-9,
-            initialization: str = "unregularized",
-            c=None,
-            regularisation="l0",
-            rank=1,
-            random_state=None,
-            positive=False,
+        self,
+        max_iter: int = 100,
+        tol=1e-9,
+        c=None,
+        regularisation="l0",
+        rank=1,
+        random_state=None,
+        positive=False,
     ):
         super().__init__(
             max_iter=max_iter,
             tol=tol,
-            initialization=initialization,
             random_state=random_state,
         )
         self.c = c
@@ -556,7 +503,7 @@ def _check_params(self):
         )
 
     def _inner_iteration(self):
-        c = self.random_state.randn(self.rank, 1)
+        c = self.random_state.randn(self.rank)
         c /= np.linalg.norm(c)
         a = self.P @ np.diag(self.D) @ c
         u = self.update(a, self.c[0])
@@ -572,22 +519,20 @@ def _inner_iteration(self):
             self.weights[1] = v
 
 
-class SWCCAInnerLoop(PLSInnerLoop):
+class _SWCCAInnerLoop(_PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-9,
-            initialization: str = "unregularized",
-            regularisation="l0",
-            c=None,
-            sample_support: int = None,
-            random_state=None,
-            positive=False,
+        self,
+        max_iter: int = 100,
+        tol=1e-9,
+        regularisation="l0",
+        c=None,
+        sample_support: int = None,
+        random_state=None,
+        positive=False,
     ):
         super().__init__(
             max_iter=max_iter,
             tol=tol,
-            initialization=initialization,
             random_state=random_state,
         )
         self.c = c
@@ -601,7 +546,7 @@ def __init__(
     def _check_params(self):
         if self.sample_support is None:
             self.sample_support = self.views[0].shape[0]
-        self.sample_weights = np.ones((self.views[0].shape[0], 1))
+        self.sample_weights = np.ones(self.views[0].shape[0])
         self.sample_weights /= np.linalg.norm(self.sample_weights)
         self.c = _process_parameter("c", self.c, 2, len(self.views))
         self.positive = _process_parameter(
@@ -616,8 +561,8 @@ def _update_view(self, view_index: int):
         targets = np.ma.array(self.scores, mask=False)
         targets.mask[view_index] = True
         self.weights[view_index] = (
-                                           self.views[view_index] * self.sample_weights
-                                   ).T @ targets.sum(axis=0).filled()
+            self.views[view_index] * self.sample_weights[:, np.newaxis]
+        ).T @ targets.sum(axis=0).filled()
         self.weights[view_index] = self.update(
             self.weights[view_index],
             self.c[view_index],
@@ -706,7 +651,7 @@ def _delta_search(w, c, positive=False, init=0, tol=1e-9):
             current, previous, current_val, previous_val, min_, max_
         )
         previous_val = current_val
-        if np.abs(current_val) < tol or np.abs(max_ - min_) < tol or i == 50:
+        if np.abs(current_val) < tol or np.abs(max_ - min_) < tol or i == 150:
             converged = True
     return coef
 
diff --git a/cca_zoo/models/gcca.py b/cca_zoo/models/gcca.py
index 00fed385..ee248475 100644
--- a/cca_zoo/models/gcca.py
+++ b/cca_zoo/models/gcca.py
@@ -6,7 +6,7 @@
 from sklearn.utils.validation import check_is_fitted
 
 from cca_zoo.models import rCCA
-from cca_zoo.utils.check_values import _process_parameter, check_views
+from cca_zoo.utils.check_values import _process_parameter, _check_views
 
 
 class GCCA(rCCA):
@@ -239,7 +239,7 @@ def transform(self, views: np.ndarray, y=None, **kwargs):
         :param kwargs: any additional keyword arguments required by the given model
         """
         check_is_fitted(self, attributes=["weights"])
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale_transform(views)
diff --git a/cca_zoo/models/iterative.py b/cca_zoo/models/iterative.py
index 083359ef..321885ed 100644
--- a/cca_zoo/models/iterative.py
+++ b/cca_zoo/models/iterative.py
@@ -5,17 +5,18 @@
 
 import numpy as np
 
-from cca_zoo.models.cca_base import _CCA_Base
-from cca_zoo.models.innerloop import (
-    PLSInnerLoop,
-    PMDInnerLoop,
-    ParkhomenkoInnerLoop,
-    ElasticInnerLoop,
-    ADMMInnerLoop,
-    SpanCCAInnerLoop,
-    SWCCAInnerLoop,
+from cca_zoo.models import MCCA
+from cca_zoo.models._cca_base import _CCA_Base
+from cca_zoo.models._innerloop import (
+    _PLSInnerLoop,
+    _PMDInnerLoop,
+    _ParkhomenkoInnerLoop,
+    _ElasticInnerLoop,
+    _ADMMInnerLoop,
+    _SpanCCAInnerLoop,
+    _SWCCAInnerLoop,
 )
-from cca_zoo.utils import check_views
+from cca_zoo.utils.check_values import _check_views
 
 
 class _Iterative(_CCA_Base):
@@ -33,7 +34,7 @@ def __init__(
         random_state=None,
         deflation="cca",
         max_iter: int = 100,
-        initialization: str = "unregularized",
+        initialization: Union[str, callable] = "random",
         tol: float = 1e-9,
     ):
         """
@@ -46,8 +47,7 @@ def __init__(
         :param random_state: Pass for reproducible output across multiple function calls
         :param deflation: the type of deflation.
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         """
         super().__init__(
@@ -56,12 +56,12 @@ def __init__(
             centre=centre,
             copy_data=copy_data,
             accept_sparse=["csc", "csr"],
+            random_state=random_state,
         )
         self.max_iter = max_iter
         self.initialization = initialization
         self.tol = tol
         self.deflation = deflation
-        self.random_state = random_state
 
     def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
         """
@@ -69,13 +69,16 @@ def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
 
         :param views: list/tuple of numpy arrays or array likes with the same number of rows (samples)
         """
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale(views)
         self.n_views = len(views)
         self.n = views[0].shape[0]
-        self._set_loop_params()
+        if isinstance(self.initialization, str):
+            self.initializer = self._initializer(views, self.initialization)
+        else:
+            self.initializer = self.initialization()
         n = views[0].shape[0]
         p = [view.shape[1] for view in views]
         # List of d: p x k
@@ -86,11 +89,11 @@ def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
         self.scores = [np.zeros((n, self.latent_dims)) for _ in views]
 
         residuals = copy.deepcopy(list(views))
-
         self.track = []
         # For each of the dimensions
         for k in range(self.latent_dims):
-            self.loop = self.loop._fit(*residuals)
+            self._set_loop_params()
+            self.loop = self.loop.fit(*residuals, initial_scores=next(self.initializer))
             for i, residual in enumerate(residuals):
                 self.weights[i][:, k] = self.loop.weights[i].ravel()
                 self.scores[i][:, k] = self.loop.scores[i].ravel()
@@ -99,8 +102,10 @@ def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
                     residuals[i], self.scores[i][:, k], self.weights[i][:, k]
                 )
             self.track.append(self.loop.track)
-            if self.track[-1]["converged"] == False:
-                warnings.warn(f"Inner loop {k} did not converge or converged to nans")
+            if not self.track[-1]["converged"]:
+                warnings.warn(
+                    f"Inner loop {k} not converged. Increase number of iterations."
+                )
         return self
 
     def _deflate(self, residual, score, loading):
@@ -113,8 +118,8 @@ def _deflate(self, residual, score, loading):
         """
         if self.deflation == "cca":
             return (
-                    residual
-                    - np.outer(score, score) @ residual / np.dot(score, score).item()
+                residual
+                - np.outer(score, score) @ residual / np.dot(score, score).item()
             )
         elif self.deflation == "pls":
             return residual - np.outer(score, loading)
@@ -126,12 +131,41 @@ def _set_loop_params(self):
         """
         Sets up the inner optimization loop for the method. By default uses the PLS inner loop.
         """
-        self.loop = PLSInnerLoop(
+        self.loop = _PLSInnerLoop(
             max_iter=self.max_iter,
-            initialization=self.initialization,
             random_state=self.random_state,
         )
 
+    def _initializer(self, views, initialization):
+        if initialization == "random":
+            while True:
+                yield np.array(
+                    [
+                        self.random_state.normal(0, 1, size=(view.shape[0]))
+                        for view in views
+                    ]
+                )
+        elif initialization == "uniform":
+            while True:
+                yield np.array([np.ones(view.shape[0]) for view in views])
+        elif initialization == "pls":
+            latent_dim = 0
+            # use rCCA to use multiple views
+            pls_scores = MCCA(self.latent_dims, c=1).fit_transform(views)
+            while True:
+                yield np.stack(pls_scores)[:, :, latent_dim]
+                latent_dim += 1
+        elif initialization == "cca":
+            latent_dim = 0
+            cca_scores = MCCA(self.latent_dims).fit_transform(views)
+            while True:
+                yield np.stack(cca_scores)[:, :, latent_dim]
+                latent_dim += 1
+        else:
+            raise ValueError(
+                "Initialization {type} not supported. Pass a generator implementing this method"
+            )
+
 
 class PLS_ALS(_Iterative):
     r"""
@@ -156,19 +190,19 @@ class PLS_ALS(_Iterative):
     >>> X2 = rng.random((10,5))
     >>> model = PLS_ALS(random_state=0)
     >>> model.fit((X1,X2)).score((X1,X2))
-    array([0.81796873])
+    array([0.81796854])
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            max_iter: int = 100,
-            initialization: str = "unregularized",
-            tol: float = 1e-9,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        max_iter: int = 100,
+        initialization: Union[str, callable] = "random",
+        tol: float = 1e-9,
     ):
         """
         Constructor for PLS
@@ -179,7 +213,7 @@ def __init__(
         :param copy_data: If True, X will be copied; else, it may be overwritten
         :param random_state: Pass for reproducible output across multiple function calls
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         """
         super().__init__(
@@ -195,9 +229,8 @@ def __init__(
         )
 
     def _set_loop_params(self):
-        self.loop = PLSInnerLoop(
+        self.loop = _PLSInnerLoop(
             max_iter=self.max_iter,
-            initialization=self.initialization,
             tol=self.tol,
             random_state=self.random_state,
         )
@@ -240,7 +273,7 @@ def __init__(
         random_state=None,
         deflation="cca",
         max_iter: int = 100,
-        initialization: str = "unregularized",
+        initialization: Union[str, callable] = "pls",
         tol: float = 1e-9,
         c: Union[Iterable[float], float] = None,
         l1_ratio: Union[Iterable[float], float] = None,
@@ -258,7 +291,7 @@ def __init__(
         :param random_state: Pass for reproducible output across multiple function calls
         :param deflation: the type of deflation.
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         :param c: lasso alpha
         :param l1_ratio: l1 ratio in lasso subproblems
@@ -289,12 +322,11 @@ def __init__(
         )
 
     def _set_loop_params(self):
-        self.loop = ElasticInnerLoop(
+        self.loop = _ElasticInnerLoop(
             max_iter=self.max_iter,
             c=self.c,
             l1_ratio=self.l1_ratio,
             maxvar=self.maxvar,
-            initialization=self.initialization,
             tol=self.tol,
             stochastic=self.stochastic,
             positive=self.positive,
@@ -331,18 +363,18 @@ class CCA_ALS(ElasticCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            deflation="cca",
-            max_iter: int = 100,
-            initialization: str = "random",
-            tol: float = 1e-9,
-            stochastic=True,
-            positive: Union[Iterable[bool], bool] = None,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        deflation="cca",
+        max_iter: int = 100,
+        initialization: str = "random",
+        tol: float = 1e-9,
+        stochastic=True,
+        positive: Union[Iterable[bool], bool] = None,
     ):
         """
         Constructor for CCA_ALS
@@ -353,7 +385,7 @@ def __init__(
         :param copy_data: If True, X will be copied; else, it may be overwritten
         :param random_state: Pass for reproducible output across multiple function calls
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: initialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         :param stochastic: use stochastic regression optimisers for subproblems
         :param positive: constrain model weights to be positive
@@ -401,24 +433,24 @@ class SCCA(ElasticCCA):
     >>> X2 = rng.random((10,5))
     >>> model = SCCA(c=[0.001,0.001], random_state=0)
     >>> model.fit((X1,X2)).score((X1,X2))
-    array([0.99998919])
+    array([0.99998761])
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            deflation="cca",
-            c: Union[Iterable[float], float] = None,
-            max_iter: int = 100,
-            maxvar: bool = False,
-            initialization: str = "unregularized",
-            tol: float = 1e-9,
-            stochastic=False,
-            positive: Union[Iterable[bool], bool] = None,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        deflation="cca",
+        c: Union[Iterable[float], float] = None,
+        max_iter: int = 100,
+        maxvar: bool = False,
+        initialization: Union[str, callable] = "pls",
+        tol: float = 1e-9,
+        stochastic=False,
+        positive: Union[Iterable[bool], bool] = None,
     ):
         """
         Constructor for SCCA
@@ -430,7 +462,7 @@ def __init__(
         :param random_state: Pass for reproducible output across multiple function calls
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
         :param maxvar: use auxiliary variable "maxvar" form
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         :param c: lasso alpha
         :param stochastic: use stochastic regression optimisers for subproblems
@@ -481,22 +513,22 @@ class PMD(_Iterative):
     >>> X2 = rng.random((10,5))
     >>> model = PMD(c=[1,1],random_state=0)
     >>> model.fit((X1,X2)).score((X1,X2))
-    array([0.69792082])
+    array([0.81796873])
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            deflation="cca",
-            c: Union[Iterable[float], float] = None,
-            max_iter: int = 100,
-            initialization: str = "unregularized",
-            tol: float = 1e-9,
-            positive: Union[Iterable[bool], bool] = None,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        deflation="cca",
+        c: Union[Iterable[float], float] = None,
+        max_iter: int = 100,
+        initialization: Union[str, callable] = "pls",
+        tol: float = 1e-9,
+        positive: Union[Iterable[bool], bool] = None,
     ):
         """
         Constructor for PMD
@@ -508,7 +540,7 @@ def __init__(
         :param random_state: Pass for reproducible output across multiple function calls
         :param c: l1 regularisation parameter between 1 and sqrt(number of features) for each view
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         :param positive: constrain model weights to be positive
         """
@@ -527,10 +559,9 @@ def __init__(
         )
 
     def _set_loop_params(self):
-        self.loop = PMDInnerLoop(
+        self.loop = _PMDInnerLoop(
             max_iter=self.max_iter,
             c=self.c,
-            initialization=self.initialization,
             tol=self.tol,
             positive=self.positive,
             random_state=self.random_state,
@@ -562,21 +593,21 @@ class ParkhomenkoCCA(_Iterative):
     >>> X2 = rng.random((10,5))
     >>> model = ParkhomenkoCCA(c=[0.001,0.001],random_state=0)
     >>> model.fit((X1,X2)).score((X1,X2))
-    array([0.81803543])
+    array([0.81803527])
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            deflation="cca",
-            c: Union[Iterable[float], float] = None,
-            max_iter: int = 100,
-            initialization: str = "unregularized",
-            tol: float = 1e-9,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        deflation="cca",
+        c: Union[Iterable[float], float] = None,
+        max_iter: int = 100,
+        initialization: Union[str, callable] = "pls",
+        tol: float = 1e-9,
     ):
         """
         Constructor for ParkhomenkoCCA
@@ -588,7 +619,7 @@ def __init__(
         :param random_state: Pass for reproducible output across multiple function calls
         :param c: l1 regularisation parameter
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         """
         self.c = c
@@ -605,10 +636,9 @@ def __init__(
         )
 
     def _set_loop_params(self):
-        self.loop = ParkhomenkoInnerLoop(
+        self.loop = _ParkhomenkoInnerLoop(
             max_iter=self.max_iter,
             c=self.c,
-            initialization=self.initialization,
             tol=self.tol,
             random_state=self.random_state,
         )
@@ -637,26 +667,26 @@ class SCCA_ADMM(_Iterative):
     >>> rng=np.random.RandomState(0)
     >>> X1 = rng.random((10,5))
     >>> X2 = rng.random((10,5))
-    >>> model = SCCA_ADMM(random_state=0)
+    >>> model = SCCA_ADMM(random_state=0,c=[1e-1,1e-1])
     >>> model.fit((X1,X2)).score((X1,X2))
-    array([0.99999997])
+    array([0.84348183])
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            deflation="cca",
-            c: Union[Iterable[float], float] = None,
-            mu: Union[Iterable[float], float] = None,
-            lam: Union[Iterable[float], float] = None,
-            eta: Union[Iterable[float], float] = None,
-            max_iter: int = 100,
-            initialization: str = "unregularized",
-            tol: float = 1e-9,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        deflation="cca",
+        c: Union[Iterable[float], float] = None,
+        mu: Union[Iterable[float], float] = None,
+        lam: Union[Iterable[float], float] = None,
+        eta: Union[Iterable[float], float] = None,
+        max_iter: int = 100,
+        initialization: Union[str, callable] = "pls",
+        tol: float = 1e-9,
     ):
         """
         Constructor for SCCA_ADMM
@@ -668,7 +698,7 @@ def __init__(
         :param random_state: Pass for reproducible output across multiple function calls
         :param c: l1 regularisation parameter
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         :param mu:
         :param lam:
@@ -691,13 +721,12 @@ def __init__(
         )
 
     def _set_loop_params(self):
-        self.loop = ADMMInnerLoop(
+        self.loop = _ADMMInnerLoop(
             max_iter=self.max_iter,
             c=self.c,
             mu=self.mu,
             lam=self.lam,
             eta=self.eta,
-            initialization=self.initialization,
             tol=self.tol,
             random_state=self.random_state,
         )
@@ -733,20 +762,20 @@ class SpanCCA(_Iterative):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            max_iter: int = 100,
-            initialization: str = "uniform",
-            tol: float = 1e-9,
-            regularisation="l0",
-            c: Union[Iterable[Union[float, int]], Union[float, int]] = None,
-            rank=1,
-            positive: Union[Iterable[bool], bool] = None,
-            random_state=None,
-            deflation="cca",
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        max_iter: int = 100,
+        initialization: str = "uniform",
+        tol: float = 1e-9,
+        regularisation="l0",
+        c: Union[Iterable[Union[float, int]], Union[float, int]] = None,
+        rank=1,
+        positive: Union[Iterable[bool], bool] = None,
+        random_state=None,
+        deflation="cca",
     ):
         """
 
@@ -756,7 +785,7 @@ def __init__(
         :param copy_data: If True, X will be copied; else, it may be overwritten
         :param random_state: Pass for reproducible output across multiple function calls
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         :param regularisation:
         :param c: regularisation parameter
@@ -780,10 +809,9 @@ def __init__(
         self.positive = positive
 
     def _set_loop_params(self):
-        self.loop = SpanCCAInnerLoop(
+        self.loop = _SpanCCAInnerLoop(
             max_iter=self.max_iter,
             c=self.c,
-            initialization=self.initialization,
             tol=self.tol,
             regularisation=self.regularisation,
             rank=self.rank,
@@ -835,7 +863,7 @@ def __init__(
         :param copy_data: If True, X will be copied; else, it may be overwritten
         :param random_state: Pass for reproducible output across multiple function calls
         :param max_iter: the maximum number of iterations to perform in the inner optimization loop
-        :param initialization: intialization for optimisation. 'unregularized' uses CCA or PLS solution,'random' uses random initialization,'uniform' uses uniform initialization of weights and scores
+        :param initialization: either string from "pls", "cca", "random", "uniform" or callable to initialize the score variables for iterative methods
         :param tol: tolerance value used for early stopping
         :param regularisation: the type of regularisation on the weights either 'l0' or 'l1'
         :param c: regularisation parameter
@@ -859,9 +887,8 @@ def __init__(
         )
 
     def _set_loop_params(self):
-        self.loop = SWCCAInnerLoop(
+        self.loop = _SWCCAInnerLoop(
             max_iter=self.max_iter,
-            initialization=self.initialization,
             tol=self.tol,
             regularisation=self.regularisation,
             c=self.c,
diff --git a/cca_zoo/models/mcca.py b/cca_zoo/models/mcca.py
index 0762c21f..76ea6e22 100644
--- a/cca_zoo/models/mcca.py
+++ b/cca_zoo/models/mcca.py
@@ -6,7 +6,7 @@
 from sklearn.utils.validation import check_is_fitted
 
 from cca_zoo.models import rCCA
-from cca_zoo.utils.check_values import _process_parameter, check_views
+from cca_zoo.utils.check_values import _process_parameter, _check_views
 
 
 class MCCA(rCCA):
@@ -222,7 +222,7 @@ def transform(self, views: np.ndarray, y=None, **kwargs):
         :param kwargs: any additional keyword arguments required by the given model
         """
         check_is_fitted(self, attributes=["weights"])
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale_transform(views)
diff --git a/cca_zoo/models/ncca.py b/cca_zoo/models/ncca.py
index 13473fc7..c8c3c91e 100644
--- a/cca_zoo/models/ncca.py
+++ b/cca_zoo/models/ncca.py
@@ -5,8 +5,8 @@
 from sklearn.neighbors import NearestNeighbors
 from sklearn.utils.validation import check_is_fitted
 
-from cca_zoo.models.cca_base import _CCA_Base
-from cca_zoo.utils.check_values import _process_parameter, check_views
+from cca_zoo.models._cca_base import _CCA_Base
+from cca_zoo.utils.check_values import _process_parameter, _check_views
 
 
 class NCCA(_CCA_Base):
@@ -46,7 +46,7 @@ def __init__(
         :param copy_data: If True, X will be copied; else, it may be overwritten
         :param accept_sparse: Whether model can take sparse data as input
         :param random_state: Pass for reproducible output across multiple function calls
-        :param nearest_neighbors: Number of neaest neighbors (l2 distance) to consider when constructing affinity
+        :param nearest_neighbors: Number of nearest neighbors (l2 distance) to consider when constructing affinity
         :param gamma: Bandwidth parameter for rbf kernel
         """
         super().__init__(
@@ -63,7 +63,7 @@ def _check_params(self):
         self.kernel = _process_parameter("kernel", None, "rbf", self.n_views)
 
     def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale(views)
@@ -71,16 +71,16 @@ def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
         self.n = views[0].shape[0]
         self._check_params()
         self.train_views = views
-        self.KNs = [
+        self.knns = [
             NearestNeighbors(n_neighbors=self.nearest_neighbors[i]).fit(view)
             for i, view in enumerate(views)
         ]
         NNs = [
-            self.KNs[i].kneighbors(view, self.nearest_neighbors[i])
+            self.knns[i].kneighbors(view, self.nearest_neighbors[i])
             for i, view in enumerate(views)
         ]
         kernels = [self._get_kernel(i, view) for i, view in enumerate(self.train_views)]
-        self.Ws = [fill_W(kernel, inds) for kernel, (dists, inds) in zip(kernels, NNs)]
+        self.Ws = [fill_w(kernel, inds) for kernel, (dists, inds) in zip(kernels, NNs)]
         self.Ws = [
             self.Ws[0] / self.Ws[0].sum(axis=1, keepdims=True),
             self.Ws[1] / self.Ws[1].sum(axis=0, keepdims=True),
@@ -100,19 +100,19 @@ def transform(self, views: Iterable[np.ndarray], y=None, **kwargs):
         :param kwargs: any additional keyword arguments required by the given model
         """
         check_is_fitted(self, attributes=["f", "g"])
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale_transform(views)
-        NNs = [
-            self.KNs[i].kneighbors(view, self.nearest_neighbors[i])
+        nns = [
+            self.knns[i].kneighbors(view, self.nearest_neighbors[i])
             for i, view in enumerate(views)
         ]
         kernels = [
             self._get_kernel(i, self.train_views[i], Y=view)
             for i, view in enumerate(views)
         ]
-        Wst = [fill_W(kernel, inds) for kernel, (dists, inds) in zip(kernels, NNs)]
+        Wst = [fill_w(kernel, inds) for kernel, (dists, inds) in zip(kernels, nns)]
         Wst = [
             Wst[0] / Wst[0].sum(axis=1, keepdims=True),
             Wst[1] / Wst[1].sum(axis=1, keepdims=True),
@@ -129,8 +129,8 @@ def _get_kernel(self, view, X, Y=None):
         )
 
 
-def fill_W(kernels, inds):
-    W = np.zeros_like(kernels)
+def fill_w(kernels, inds):
+    w = np.zeros_like(kernels)
     for i, ind in enumerate(inds):
-        W[ind, i] = kernels[ind, i]
-    return W.T
+        w[ind, i] = kernels[ind, i]
+    return w.T
diff --git a/cca_zoo/models/partialcca.py b/cca_zoo/models/partialcca.py
new file mode 100644
index 00000000..851164be
--- /dev/null
+++ b/cca_zoo/models/partialcca.py
@@ -0,0 +1,130 @@
+from typing import Iterable, Union
+
+import numpy as np
+from scipy.linalg import block_diag
+from sklearn.utils.validation import check_is_fitted
+
+from cca_zoo.models import MCCA
+from cca_zoo.utils import _check_views
+
+
+class PartialCCA(MCCA):
+    r""" """
+
+    def __init__(
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        c: Union[Iterable[float], float] = None,
+        eps=1e-3,
+    ):
+        """
+        Constructor for Partial CCA
+
+        :param latent_dims: number of latent dimensions to fit
+        :param scale: normalize variance in each column before fitting
+        :param centre: demean data by column before fitting (and before transforming out of sample
+        :param copy_data: If True, X will be copied; else, it may be overwritten
+        :param random_state: Pass for reproducible output across multiple function calls
+        :param c: Iterable of regularisation parameters for each view (between 0:CCA and 1:PLS)
+        :param eps: epsilon for stability
+        """
+        super().__init__(
+            latent_dims=latent_dims,
+            scale=scale,
+            centre=centre,
+            copy_data=copy_data,
+            random_state=random_state,
+        )
+        self.c = c
+        self.eps = eps
+
+    def _setup_evp(self, views: Iterable[np.ndarray], confounds=None):
+        if confounds is None:
+            raise ValueError(
+                f"confounds is {confounds}. Require matching confounds to transform with"
+                f"partial CCA."
+            )
+        self.confound_betas = [np.linalg.pinv(confounds) @ view for view in views]
+        views = [
+            view - confounds @ np.linalg.pinv(confounds) @ view
+            for view, confound_beta in zip(views, self.confound_betas)
+        ]
+        all_views = np.concatenate(views, axis=1)
+        C = all_views.T @ all_views / self.n
+        # Can regularise by adding to diagonal
+        D = block_diag(
+            *[
+                (1 - self.c[i]) * m.T @ m / self.n + self.c[i] * np.eye(m.shape[1])
+                for i, m in enumerate(views)
+            ]
+        )
+        C -= block_diag(*[view.T @ view / self.n for view in views]) - D
+        D_smallest_eig = min(0, np.linalg.eigvalsh(D).min()) - self.eps
+        D = D - D_smallest_eig * np.eye(D.shape[0])
+        self.splits = np.cumsum([0] + [view.shape[1] for view in views])
+        return views, C, D
+
+    # TODO TRANSFORM
+    def transform(self, views: Iterable[np.ndarray], y=None, confounds=None):
+        """
+        Transforms data given a fit model
+
+        :param views: numpy arrays with the same number of rows (samples) separated by commas
+        """
+        if confounds is None:
+            raise ValueError(
+                f"confounds is {confounds}. Require matching confounds to transform with"
+                f"partial CCA."
+            )
+        check_is_fitted(self, attributes=["weights"])
+        views = _check_views(
+            *views, copy=self.copy_data, accept_sparse=self.accept_sparse
+        )
+        views = self._centre_scale_transform(views)
+        transformed_views = []
+        for i, (view) in enumerate(views):
+            transformed_view = (
+                view - confounds @ self.confound_betas[i]
+            ) @ self.weights[i]
+            transformed_views.append(transformed_view)
+        return transformed_views
+
+
+def main():
+    import numpy as np
+
+    n = 100
+    p = 10
+    q = 10
+    Z = np.random.rand(100, 1)
+    Z -= Z.mean(axis=0)
+    C = np.random.rand(100, 1)
+    C -= C.mean(axis=0)
+    Wx = np.random.rand(1, p)
+    Wxc = np.random.rand(1, p)
+    Wy = np.random.rand(1, q)
+    Wyc = np.random.rand(1, q)
+    X = Z @ Wx + C @ Wxc + 2 * np.random.rand(n, p)
+    Y = Z @ Wy + C @ Wyc + 2 * np.random.rand(n, q)
+    pcca = PartialCCA(centre=False, scale=False).fit((X, Y), confounds=C)
+    a = pcca.score((X, Y), confounds=C)
+    from cca_zoo.models import MCCA
+
+    cca = MCCA(centre=False, scale=False).fit((X, Y), confounds=C)
+    b = cca.score((X, Y))
+    bb = np.corrcoef(cca.transform((X, Y))[0], C, rowvar=False)
+    X_deconf = (np.eye(n) - C @ np.linalg.inv(C.T @ C) @ C.T) @ X
+    Y_deconf = (np.eye(n) - C @ np.linalg.inv(C.T @ C) @ C.T) @ Y
+    m = MCCA(centre=False, scale=False).fit((X_deconf, Y_deconf))
+    c = m.score((X_deconf, Y_deconf))
+    cc = np.corrcoef(m.transform((X_deconf, Y_deconf))[0], C, rowvar=False)
+    aa = np.corrcoef(pcca.transform((X, Y), confounds=C)[0], C, rowvar=False)
+    print()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/cca_zoo/models/rcca.py b/cca_zoo/models/rcca.py
index 7d8b45ea..da7063c3 100644
--- a/cca_zoo/models/rcca.py
+++ b/cca_zoo/models/rcca.py
@@ -4,8 +4,8 @@
 import numpy as np
 from scipy.linalg import block_diag, eigh
 
-from cca_zoo.models.cca_base import _CCA_Base
-from cca_zoo.utils.check_values import _process_parameter, check_views
+from cca_zoo.models._cca_base import _CCA_Base
+from cca_zoo.utils.check_values import _process_parameter, _check_views
 
 
 # from hyperopt import fmin, tpe, Trials
@@ -89,7 +89,7 @@ def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
 
         :param views: list/tuple of numpy arrays or array likes with the same number of rows (samples)
         """
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale(views)
@@ -285,7 +285,7 @@ def __init__(
             scale=scale,
             centre=centre,
             copy_data=copy_data,
-            c=[1.0, 1.0],
+            c=1,
             random_state=random_state,
         )
 
diff --git a/cca_zoo/models/tcca.py b/cca_zoo/models/tcca.py
index 2aba4047..cc518e21 100644
--- a/cca_zoo/models/tcca.py
+++ b/cca_zoo/models/tcca.py
@@ -7,8 +7,8 @@
 from sklearn.utils.validation import check_is_fitted
 from tensorly.decomposition import parafac
 
-from cca_zoo.models.cca_base import _CCA_Base
-from cca_zoo.utils.check_values import _process_parameter, check_views
+from cca_zoo.models._cca_base import _CCA_Base
+from cca_zoo.utils.check_values import _process_parameter, _check_views
 
 
 class TCCA(_CCA_Base):
@@ -79,7 +79,7 @@ def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
 
         :param views: list/tuple of numpy arrays or array likes with the same number of rows (samples)
         """
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale(views)
@@ -284,7 +284,7 @@ def transform(self, views: np.ndarray, y=None, **kwargs):
         :param kwargs: any additional keyword arguments required by the given model
         """
         check_is_fitted(self, attributes=["weights"])
-        views = check_views(
+        views = _check_views(
             *views, copy=self.copy_data, accept_sparse=self.accept_sparse
         )
         views = self._centre_scale_transform(views)
diff --git a/cca_zoo/probabilisticmodels/vcca.py b/cca_zoo/probabilisticmodels/vcca.py
index 5fef9b52..4ec08ed0 100644
--- a/cca_zoo/probabilisticmodels/vcca.py
+++ b/cca_zoo/probabilisticmodels/vcca.py
@@ -8,7 +8,7 @@
 from numpyro.infer import MCMC, NUTS, Predictive
 from sklearn.utils.validation import check_is_fitted
 
-from cca_zoo.models import _CCA_Base
+from cca_zoo.models._cca_base import _CCA_Base
 
 
 class VariationalCCA(_CCA_Base):
diff --git a/cca_zoo/utils/__init__.py b/cca_zoo/utils/__init__.py
index 72fe9ca4..f20b5b9c 100644
--- a/cca_zoo/utils/__init__.py
+++ b/cca_zoo/utils/__init__.py
@@ -1,2 +1,9 @@
-from .check_values import *
-from .plotting import plot_latent_train_test, plot_latent_label
+from .check_values import (
+    _check_views,
+    _check_batch_size,
+    _check_Parikh2014,
+    _process_parameter,
+    _check_parameter_number,
+    _check_converged_weights,
+)
+from .plotting import plot_latent_train_test, plot_latent_label, cv_plot
diff --git a/cca_zoo/utils/check_values.py b/cca_zoo/utils/check_values.py
index 451145cb..3531d507 100644
--- a/cca_zoo/utils/check_values.py
+++ b/cca_zoo/utils/check_values.py
@@ -5,7 +5,7 @@
 from sklearn.utils import check_array
 
 
-def check_views(*views: Iterable[np.ndarray], copy=False, accept_sparse=False):
+def _check_views(*views: Iterable[np.ndarray], copy=False, accept_sparse=False):
     """
 
     :param views:
diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst
index 78240b66..350803c2 100644
--- a/docs/source/api/models.rst
+++ b/docs/source/api/models.rst
@@ -140,7 +140,7 @@ Sparse Weighted CCA
 Base Class
 --------------------------------
 
-.. automodule:: cca_zoo.models.cca_base
+.. automodule:: cca_zoo.models._cca_base
     :members:
     :private-members: _CCA_Base
     :exclude-members: get_params, set_params
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 5321577f..6be517d9 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -19,7 +19,7 @@
     "ignore",
     category=UserWarning,
     message="Matplotlib is currently using agg, which is a"
-            " non-GUI backend, so cannot show the figure.",
+    " non-GUI backend, so cannot show the figure.",
 )
 # -- Project information -----------------------------------------------------
 
diff --git a/docs/source/documentation/maths.rst b/docs/source/documentation/maths.rst
index 1277cdbb..1ff252aa 100644
--- a/docs/source/documentation/maths.rst
+++ b/docs/source/documentation/maths.rst
@@ -12,13 +12,13 @@ The first principle component can be written as the solution to the convex optim
 
 .. math::
 
-    \w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_1w_1  \}
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_1w_1  \}
 
     \text{subject to:}
 
     w_1^Tw_1=1
 
-That is the singular vectors of the covariance matrix :math:`X^TX`
+WHich is optimized for the singular vectors of the covariance matrix :math:`X^TX`.
 
 PLS
 ----
@@ -28,7 +28,7 @@ It is tempting to write a slightly different optimisation problem:
 
 .. math::
 
-    \w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_2w_2  \}
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_2w_2  \}
 
     \text{subject to:}
 
@@ -42,11 +42,12 @@ Which is optimised for the left and right singular vectors of the cross covarian
 CCA
 ----
 
-To arrive at Canonical Correlation
+To arrive at Canonical Correlation we change the constraints slightly so that they now depend on the variance
+in each view as well as the weights. This makes them data dependent.
 
 .. math::
 
-    \w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_2w_2  \}
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_2w_2  \}
 
     \text{subject to:}
 
@@ -54,15 +55,104 @@ To arrive at Canonical Correlation
 
     w_2^TX_2^TX_2w_2=1
 
+Despite these more complex constraints, CCA is the solution to a generalized eigenvalue problem.
+
+Regularized CCA
+-----------------
+
+Notice that the constraints for PLS and CCA are identical in the case that :math:`X_i^TX_i=I`.
+This leads naturally to mixing the constraints and in particular the PLS constraint acts in a similar
+way to the 'ridge' regularisation in Ridge Regression.
+
+.. math::
+
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_2w_2  \}
+
+    \text{subject to:}
+
+    (1-c_1)w_1^TX_1^TX_1w_1+c_1w_1^Tw_1=1
+
+    (1-c_2)w_2^TX_2^TX_2w_2+c_2w_2^Tw_2=1
+
+Other regularized CCA and PLS
+--------------------------------
+
+There has been lots of research into more general forms of regularisation for CCA, in particular forms that induce
+sparsity on the weights. In general these can be written in a form somewhat similar to:
+
+.. math::
+
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{ w_1^TX_1^TX_2w_2 + \lambda_1f_1(w_1) + \lambda_2f_2(w_2)  \}
+
+    \text{subject to:}
+
+    (1-c_1)w_1^TX_1^TX_1w_1+c_1w_1^Tw_1=1
+
+    (1-c_2)w_2^TX_2^TX_2w_2+c_2w_2^Tw_2=1
+
+These problems usually have no closed form solution and are typically solved by alternately fixing :math:`w_1`
+and :math:`w_2`.
+
+Kernel CCA and PLS
+---------------------
+
+By expressing the weights as :math:`w_i=\alpha_iX_i` we can transform the CCA problem into its kernel form:
+
+.. math::
+
+    \alpha_{opt}=\underset{\alpha}{\mathrm{argmax}}\{ \alpha_1^TK_1^TK_2\alpha_2  \}
+
+    \text{subject to:}
+
+    \alpha_1^TK_1^TK_1\alpha_1=1
+
+    \alpha_2^TK_2^TK_2\alpha_2=1
+
+Finally we can also consider more general kernel matrices without loss of generalisation. A similar reparameterization
+exists for PLS and therefore for ridge regularized CCA.
+
+Multiset CCA
+----------------
+
+Ridge Regularized CCA can be generalized to find correlations between more than one view with the formulation:
+
+.. math::
+
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{\sum_i\sum_{j\neq i} w_i^TX_i^TX_jw_j  \}\\
+
+    \text{subject to:}
+
+    (1-c_i)w_i^TX_i^TX_iw_i+c_iw_i^Tw_i=1
+
+This form is often referred to as SUMCOR CCA because it optimizes for the pairwise sum of . This can be formulated as a
+generalized eigenvalue problem and therefore solved efficiently.
+
+Generalized CCA
+-----------------
+
+Ridge Regularized CCA can be generalized to find correlations between more than one view with the formulation:
+
+.. math::
+
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{ \sum_iw_i^TX_i^TT  \}\\
+
+    \text{subject to:}
+
+    T^TT=1
+
+This form is often referred to as MAXVAR CCA since it finds an auxiliary vector :math:`T` with fixed unit norm that has
+maximum sum of variance with each view. This can also be formulated as a generalized eigenvalue problem and
+therefore solved efficiently.
 
 Deep CCA
 ----------
 
-To arrive
+The ideas behind CCA can be extended to a general form where instead of linear weights, we consider functions
+:math:`f(X_i)`. Where these functions are parameterized by neural networks, we can consider a 'Deep' CCA:
 
 .. math::
 
-    \w_{opt}=\underset{w}{\mathrm{argmax}}\{ f(X_1)^Tf(X_2)  \}
+    w_{opt}=\underset{w}{\mathrm{argmax}}\{ f(X_1)^Tf(X_2)  \}
 
     \text{subject to:}
 
diff --git a/examples/plot_dcca_custom.py b/examples/plot_dcca_custom.py
index 8827cd81..754ae457 100644
--- a/examples/plot_dcca_custom.py
+++ b/examples/plot_dcca_custom.py
@@ -6,6 +6,7 @@
 """
 
 import numpy as np
+
 # %%
 import pytorch_lightning as pl
 from torch import optim
diff --git a/examples/plot_kernel_cca.py b/examples/plot_kernel_cca.py
index 581c4476..937116fa 100644
--- a/examples/plot_kernel_cca.py
+++ b/examples/plot_kernel_cca.py
@@ -51,8 +51,8 @@ def my_kernel(X, Y, param=0):
     GridSearchCV(
         KCCA(latent_dims=latent_dims), param_grid=param_grid, cv=cv, verbose=True
     )
-        .fit([X, Y])
-        .best_estimator_
+    .fit([X, Y])
+    .best_estimator_
 )
 
 # %%
@@ -64,8 +64,8 @@ def my_kernel(X, Y, param=0):
     GridSearchCV(
         KCCA(latent_dims=latent_dims), param_grid=param_grid, cv=cv, verbose=True
     )
-        .fit([X, Y])
-        .best_estimator_
+    .fit([X, Y])
+    .best_estimator_
 )
 
 # %%
@@ -77,8 +77,8 @@ def my_kernel(X, Y, param=0):
     GridSearchCV(
         KCCA(latent_dims=latent_dims), param_grid=param_grid, cv=cv, verbose=True
     )
-        .fit([X, Y])
-        .best_estimator_
+    .fit([X, Y])
+    .best_estimator_
 )
 
 
diff --git a/examples/plot_sparse_cca.py b/examples/plot_sparse_cca.py
index d73208a9..a900b1a0 100644
--- a/examples/plot_sparse_cca.py
+++ b/examples/plot_sparse_cca.py
@@ -62,7 +62,7 @@ def plot_model_weights(wx, wy, tx, ty):
 plot_model_weights(pls.weights[0], pls.weights[1], tx, ty)
 
 # %%
-pmd = PMD(c=[2, 2]).fit([X, Y])
+pmd = PMD(c=[0.5, 0.5]).fit([X, Y])
 plot_model_weights(pmd.weights[0], pmd.weights[1], tx, ty)
 
 # %%
@@ -73,8 +73,8 @@ def plot_model_weights(wx, wy, tx, ty):
 plt.xlabel("#iterations")
 
 # %%
-c1 = [1, 3, 7, 9]
-c2 = [1, 3, 7, 9]
+c1 = [0.1, 0.3, 0.7, 0.9]
+c2 = [0.1, 0.3, 0.7, 0.9]
 param_grid = {"c": [c1, c2]}
 pmd = GridSearchCV(PMD(), param_grid=param_grid, cv=3, verbose=True).fit([X, Y])
 


Test patch:
diff --git a/cca_zoo/test/test_deepmodels.py b/cca_zoo/test/test_deepmodels.py
index 61c6c37b..e7ea258a 100644
--- a/cca_zoo/test/test_deepmodels.py
+++ b/cca_zoo/test/test_deepmodels.py
@@ -53,10 +53,10 @@ def test_DCCA_methods():
     )
     trainer.fit(dcca_noi, train_loader)
     assert (
-            np.testing.assert_array_less(
-                cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
-            )
-            is None
+        np.testing.assert_array_less(
+            cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
+        )
+        is None
     )
     # Soft Decorrelation (stochastic Decorrelation Loss)
     encoder_1 = architectures.Encoder(latent_dims=latent_dims, feature_size=10)
@@ -67,10 +67,10 @@ def test_DCCA_methods():
     trainer = pl.Trainer(max_epochs=epochs, log_every_n_steps=10)
     trainer.fit(sdl, train_loader)
     assert (
-            np.testing.assert_array_less(
-                cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
-            )
-            is None
+        np.testing.assert_array_less(
+            cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
+        )
+        is None
     )
     # DCCA
     encoder_1 = architectures.Encoder(latent_dims=latent_dims, feature_size=10)
@@ -87,10 +87,10 @@ def test_DCCA_methods():
     )
     trainer.fit(dcca, train_loader, val_dataloaders=val_loader)
     assert (
-            np.testing.assert_array_less(
-                cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
-            )
-            is None
+        np.testing.assert_array_less(
+            cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
+        )
+        is None
     )
     # DGCCA
     encoder_1 = architectures.Encoder(latent_dims=latent_dims, feature_size=10)
@@ -107,10 +107,10 @@ def test_DCCA_methods():
     )
     trainer.fit(dgcca, train_loader)
     assert (
-            np.testing.assert_array_less(
-                cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
-            )
-            is None
+        np.testing.assert_array_less(
+            cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
+        )
+        is None
     )
     # DMCCA
     encoder_1 = architectures.Encoder(latent_dims=latent_dims, feature_size=10)
@@ -127,10 +127,10 @@ def test_DCCA_methods():
     )
     trainer.fit(dmcca, train_loader)
     assert (
-            np.testing.assert_array_less(
-                cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
-            )
-            is None
+        np.testing.assert_array_less(
+            cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
+        )
+        is None
     )
     # Barlow Twins
     encoder_1 = architectures.Encoder(latent_dims=latent_dims, feature_size=10)
@@ -146,10 +146,10 @@ def test_DCCA_methods():
     )
     trainer.fit(barlowtwins, train_loader)
     assert (
-            np.testing.assert_array_less(
-                cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
-            )
-            is None
+        np.testing.assert_array_less(
+            cca.score((X, Y)).sum(), trainer.model.score(train_loader).sum()
+        )
+        is None
     )
 
 
@@ -257,8 +257,8 @@ def test_linear():
     cca = CCA().fit((X, Y))
     # check linear encoder with SGD matches vanilla linear CCA
     assert (
-            np.testing.assert_array_almost_equal(
-                cca.score((X, Y)), trainer.model.score(loader), decimal=2
-            )
-            is None
+        np.testing.assert_array_almost_equal(
+            cca.score((X, Y)), trainer.model.score(loader), decimal=2
+        )
+        is None
     )
diff --git a/cca_zoo/test/test_models.py b/cca_zoo/test/test_models.py
index e9e17d0f..36bc65bc 100644
--- a/cca_zoo/test/test_models.py
+++ b/cca_zoo/test/test_models.py
@@ -2,7 +2,7 @@
 import scipy.sparse as sp
 from sklearn.utils.fixes import loguniform
 from sklearn.utils.validation import check_random_state
-
+import pytest
 from cca_zoo.model_selection import GridSearchCV, RandomizedSearchCV
 from cca_zoo.models import (
     rCCA,
@@ -24,6 +24,7 @@
     NCCA,
     ParkhomenkoCCA,
     SCCA_ADMM,
+    PartialCCA,
 )
 from cca_zoo.utils.plotting import cv_plot
 
@@ -40,12 +41,8 @@ def test_unregularized_methods():
     # Tests unregularized CCA methods. The idea is that all of these should give the same result.
     latent_dims = 2
     cca = CCA(latent_dims=latent_dims).fit([X, Y])
-    iter = CCA_ALS(
-        latent_dims=latent_dims, tol=1e-9, random_state=rng, stochastic=False
-    ).fit([X, Y])
-    iter_pls = PLS_ALS(
-        latent_dims=latent_dims, tol=1e-9, initialization="unregularized", centre=False
-    ).fit([X, Y])
+    iter = CCA_ALS(latent_dims=latent_dims, tol=1e-9, stochastic=False).fit([X, Y])
+    iter_pls = PLS_ALS(latent_dims=latent_dims, tol=1e-9, centre=False).fit([X, Y])
     gcca = GCCA(latent_dims=latent_dims).fit([X, Y])
     mcca = MCCA(latent_dims=latent_dims, eps=1e-9).fit([X, Y])
     kcca = KCCA(latent_dims=latent_dims).fit([X, Y])
@@ -65,56 +62,56 @@ def test_unregularized_methods():
     assert np.testing.assert_array_almost_equal(corr_cca, corr_kcca, decimal=1) is None
     assert np.testing.assert_array_almost_equal(corr_cca, corr_tcca, decimal=1) is None
     assert (
-            np.testing.assert_array_almost_equal(corr_kgcca, corr_gcca, decimal=1) is None
+        np.testing.assert_array_almost_equal(corr_kgcca, corr_gcca, decimal=1) is None
     )
     # Check standardized models have standard outputs
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(iter.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(iter.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(cca.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(cca.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(mcca.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(mcca.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(kcca.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(kcca.transform((X, Y))[0], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(iter.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(iter.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(cca.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(cca.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(mcca.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(mcca.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
     assert (
-            np.testing.assert_allclose(
-                np.linalg.norm(kcca.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
-            )
-            is None
+        np.testing.assert_allclose(
+            np.linalg.norm(kcca.transform((X, Y))[1], axis=0) ** 2, n, rtol=0.2
+        )
+        is None
     )
 
 
@@ -125,9 +122,9 @@ def test_sparse_input():
     iter = CCA_ALS(
         latent_dims=latent_dims, tol=1e-9, stochastic=False, centre=False
     ).fit((X_sp, Y_sp))
-    iter_pls = PLS_ALS(
-        latent_dims=latent_dims, tol=1e-9, initialization="unregularized", centre=False
-    ).fit((X_sp, Y_sp))
+    iter_pls = PLS_ALS(latent_dims=latent_dims, tol=1e-9, centre=False).fit(
+        (X_sp, Y_sp)
+    )
     gcca = GCCA(latent_dims=latent_dims, centre=False).fit((X_sp, Y_sp))
     mcca = MCCA(latent_dims=latent_dims, centre=False).fit((X_sp, Y_sp))
     kcca = KCCA(latent_dims=latent_dims, centre=False).fit((X_sp, Y_sp))
@@ -182,7 +179,7 @@ def test_regularized_methods():
     # Check the correlations from each unregularized method are the same
     assert np.testing.assert_array_almost_equal(corr_pls, corr_mcca, decimal=1) is None
     assert (
-            np.testing.assert_array_almost_equal(corr_pls, corr_kernel, decimal=1) is None
+        np.testing.assert_array_almost_equal(corr_pls, corr_kernel, decimal=1) is None
     )
     assert np.testing.assert_array_almost_equal(corr_pls, corr_rcca, decimal=1) is None
 
@@ -190,15 +187,21 @@ def test_regularized_methods():
 def test_non_negative_methods():
     latent_dims = 2
     nnscca = SCCA(
-        latent_dims=latent_dims, tol=1e-9, positive=True, c=[1e-1, 1e-1], random_state=0
+        latent_dims=latent_dims,
+        tol=1e-9,
+        positive=True,
+        c=[1e-6, 1e-6],
+        random_state=0,
+        initialization="random",
     ).fit((X, Y))
     nnelastic = ElasticCCA(
         latent_dims=latent_dims,
         tol=1e-9,
         positive=True,
         l1_ratio=[0.5, 0.5],
-        c=[1e-4, 1e-5],
+        c=[1e-5, 1e-5],
         random_state=0,
+        initialization="random",
     ).fit([X, Y])
     nnals = CCA_ALS(
         latent_dims=latent_dims, tol=1e-9, positive=True, random_state=0
@@ -212,8 +215,8 @@ def test_non_negative_methods():
 
 
 def test_sparse_methods():
-    c1 = [1, 3]
-    c2 = [1, 3]
+    c1 = [0.1, 0.2]
+    c2 = [0.1, 0.2]
     param_grid = {"c": [c1, c2]}
     pmd_cv = GridSearchCV(PMD(random_state=rng), param_grid=param_grid).fit([X, Y])
     cv_plot(pmd_cv.cv_results_)
@@ -266,10 +269,10 @@ def test_weighted_GCCA_methods():
     K[0, 200:] = 0
     unobserved_gcca = GCCA(latent_dims=latent_dims, c=[c, c]).fit((X, Y), K=K)
     assert (
-            np.testing.assert_array_almost_equal(
-                corr_unweighted_gcca, corr_deweighted_gcca, decimal=1
-            )
-            is None
+        np.testing.assert_array_almost_equal(
+            corr_unweighted_gcca, corr_deweighted_gcca, decimal=1
+        )
+        is None
     )
 
 
@@ -300,6 +303,18 @@ def test_l0():
     assert (np.abs(swcca.loop.sample_weights) > 1e-5).sum() == 5
 
 
+def test_partialcca():
+    # Tests that partial CCA scores are not correlated with confounds.
+    pcca = PartialCCA(latent_dims=3)
+    pcca.fit((X, Y), confounds=Z)
+    assert np.allclose(
+        np.corrcoef(pcca.transform((X, Y), confounds=Z)[0], Z, rowvar=False)[:3, :3]
+        - np.eye(3),
+        0,
+        atol=0.001,
+    )
+
+
 def test_pls():
     pls_als = PLS_ALS(latent_dims=3)
     pls = PLS(latent_dims=3)
@@ -309,25 +324,23 @@ def test_pls():
 
 
 def test_VCCA():
-    try:
-        from cca_zoo.probabilisticmodels import VariationalCCA
-        from cca_zoo.data import generate_simple_data
+    # some might not have access to jax/numpyro so leave this as an optional test locally.
+    numpyro = pytest.importorskip("numpyro")
+    from cca_zoo.probabilisticmodels import VariationalCCA
+    from cca_zoo.data import generate_simple_data
 
-        # Tests tensor CCA methods
-        (X, Y), (_) = generate_simple_data(20, [9, 9], random_state=rng, eps=0.1)
-        latent_dims = 1
-        cca = CCA(latent_dims=latent_dims).fit([X, Y])
-        vcca = VariationalCCA(
-            latent_dims=latent_dims, num_warmup=500, num_samples=500
-        ).fit([X, Y])
-        # Test that vanilla CCA and VCCA produce roughly similar latent space
-        assert (
-                np.corrcoef(
-                    cca.transform([X, Y])[1].T,
-                    vcca.posterior_samples["z"].mean(axis=0)[:, 0],
-                )[0, 1]
-                > 0.9
-        )
-    except:
-        # some might not have access to jax/numpyro so leave this as an optional test locally.
-        pass
+    # Tests tensor CCA methods
+    (X, Y), (_) = generate_simple_data(20, [9, 9], random_state=rng, eps=0.1)
+    latent_dims = 1
+    cca = CCA(latent_dims=latent_dims).fit([X, Y])
+    vcca = VariationalCCA(latent_dims=latent_dims, num_warmup=500, num_samples=500).fit(
+        [X, Y]
+    )
+    # Test that vanilla CCA and VCCA produce roughly similar latent space
+    assert (
+        np.corrcoef(
+            cca.transform([X, Y])[1].T,
+            vcca.posterior_samples["z"].mean(axis=0)[:, 0],
+        )[0, 1]
+        > 0.9
+    )