Repository: jameschapman19/cca_zoo (language: python)

Issue:
Changed iterative error to warning. This stops gridsearch breaking if…
… just one parameter is in the wrong range


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
index ad001a9e..8405adcf 100644
--- a/.github/workflows/python-package.yml
+++ b/.github/workflows/python-package.yml
@@ -5,9 +5,9 @@ name: Python package
 
 on:
   push:
-    branches: [ main ]
+    branches: [ main, dev ]
   pull_request:
-    branches: [ main ]
+    branches: [ main, dev ]
 
 jobs:
   build:
@@ -41,11 +41,8 @@ jobs:
         # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Run Test
-      run: pytest --cov-report=xml --cov=cca_zoo ./cca_zoo/test/
+      run: pytest --doctest-modules --cov-report=xml --cov=cca_zoo ./cca_zoo/test/
     - name: "Upload coverage to Codecov"
       uses: codecov/codecov-action@v1
       with:
         fail_ci_if_error: true
-    - name: Run doctests
-      run: |
-        pytest --doctest-modules --ignore=$MODULE_NAME/tests $MODULE_NAME
diff --git a/cca_zoo/data/simulated.py b/cca_zoo/data/simulated.py
index 70dfc80a..e195e3ac 100644
--- a/cca_zoo/data/simulated.py
+++ b/cca_zoo/data/simulated.py
@@ -10,16 +10,16 @@
 
 
 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: List[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: List[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":
@@ -88,8 +88,8 @@ def generate_covariance_data(
                     ).T
                     random_state.shuffle(mask)
                     while (
-                            np.sum(np.unique(mask, axis=1, return_counts=True)[1] > 1) > 0
-                            or np.sum(np.sum(mask, axis=0) == 0) > 0
+                        np.sum(np.unique(mask, axis=1, return_counts=True)[1] > 1) > 0
+                        or np.sum(np.sum(mask, axis=0) == 0) > 0
                     ):
                         random_state.shuffle(mask)
                     weights = weights * mask
@@ -113,12 +113,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)))
@@ -133,12 +133,12 @@ def generate_covariance_data(
 
 
 def generate_simple_data(
-        n: int,
-        view_features: List[int],
-        view_sparsity: List[int] = None,
-        eps: float = 0,
-        transform=True,
-        random_state=None,
+    n: int,
+    view_features: List[int],
+    view_sparsity: List[int] = None,
+    eps: float = 0,
+    transform=True,
+    random_state=None,
 ):
     """
     Simple latent variable model to generate data with one latent factor
@@ -200,9 +200,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 236f8125..cd4c70bf 100644
--- a/cca_zoo/data/toy.py
+++ b/cca_zoo/data/toy.py
@@ -16,7 +16,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
     ):
         """
 
@@ -72,7 +72,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/architectures.py b/cca_zoo/deepmodels/architectures.py
index b1406539..31bf2f1f 100644
--- a/cca_zoo/deepmodels/architectures.py
+++ b/cca_zoo/deepmodels/architectures.py
@@ -31,11 +31,11 @@ def forward(self, x):
 
 class Encoder(BaseEncoder):
     def __init__(
-            self,
-            latent_dims: int,
-            variational: bool = False,
-            feature_size: int = 784,
-            layer_sizes: Iterable = None,
+        self,
+        latent_dims: int,
+        variational: bool = False,
+        feature_size: int = 784,
+        layer_sizes: Iterable = None,
     ):
         super(Encoder, self).__init__(latent_dims, variational=variational)
         if layer_sizes is None:
@@ -78,11 +78,11 @@ def forward(self, x):
 
 class Decoder(BaseDecoder):
     def __init__(
-            self,
-            latent_dims: int,
-            feature_size: int = 784,
-            layer_sizes: list = None,
-            norm_output: bool = False,
+        self,
+        latent_dims: int,
+        feature_size: int = 784,
+        layer_sizes: list = None,
+        norm_output: bool = False,
     ):
         super(Decoder, self).__init__(latent_dims)
         if layer_sizes is None:
@@ -126,14 +126,14 @@ def forward(self, x):
 
 class CNNEncoder(BaseEncoder):
     def __init__(
-            self,
-            latent_dims: int,
-            variational: bool = False,
-            feature_size: Iterable = (28, 28),
-            channels: list = None,
-            kernel_sizes: list = None,
-            stride: list = None,
-            padding: list = None,
+        self,
+        latent_dims: int,
+        variational: bool = False,
+        feature_size: Iterable = (28, 28),
+        channels: list = None,
+        kernel_sizes: list = None,
+        stride: list = None,
+        padding: list = None,
     ):
         super(CNNEncoder, self).__init__(latent_dims, variational=variational)
         if channels is None:
@@ -199,14 +199,14 @@ def forward(self, x):
 
 class CNNDecoder(BaseDecoder):
     def __init__(
-            self,
-            latent_dims: int,
-            feature_size: Iterable = (28, 28),
-            channels: list = None,
-            kernel_sizes=None,
-            strides=None,
-            paddings=None,
-            norm_output: bool = False,
+        self,
+        latent_dims: int,
+        feature_size: Iterable = (28, 28),
+        channels: list = None,
+        kernel_sizes=None,
+        strides=None,
+        paddings=None,
+        norm_output: bool = False,
     ):
         super(CNNDecoder, self).__init__(latent_dims)
         if channels is None:
@@ -229,7 +229,7 @@ def __init__(
         # Loop backward through decoding layers in order to work out the dimensions at each layer - in particular the first
         # linear layer needs to know B*current_size*current_size*channels
         for l_id, (channel, kernel, stride, padding) in reversed(
-                list(enumerate(zip(channels, kernel_sizes, strides, paddings)))
+            list(enumerate(zip(channels, kernel_sizes, strides, paddings)))
         ):
             conv_layers.append(
                 torch.nn.Sequential(
@@ -302,10 +302,10 @@ def forward(self, x):
 # BrainNetCNN Network for fitting Gold-MSI on LSD dataset
 class BrainNetEncoder(BaseEncoder):
     def __init__(
-            self,
-            latent_dims: int,
-            variational: bool = False,
-            feature_size: Tuple[int] = (200, ...),
+        self,
+        latent_dims: int,
+        variational: bool = False,
+        feature_size: Tuple[int] = (200, ...),
     ):
         super(BrainNetEncoder, self).__init__(latent_dims, variational=variational)
         _check_feature_size(feature_size)
diff --git a/cca_zoo/deepmodels/dcca.py b/cca_zoo/deepmodels/dcca.py
index 9a58c863..05afb8ce 100644
--- a/cca_zoo/deepmodels/dcca.py
+++ b/cca_zoo/deepmodels/dcca.py
@@ -17,12 +17,12 @@ class DCCA(_DCCA_base):
     """
 
     def __init__(
-            self,
-            latent_dims: int,
-            objective=objectives.MCCA,
-            encoders=None,
-            r: float = 0,
-            eps: float = 1e-3,
+        self,
+        latent_dims: int,
+        objective=objectives.MCCA,
+        encoders=None,
+        r: float = 0,
+        eps: float = 1e-3,
     ):
         """
         Constructor class for DCCA
diff --git a/cca_zoo/deepmodels/dcca_barlow_twins.py b/cca_zoo/deepmodels/dcca_barlow_twins.py
index 27c1da94..d1e16036 100644
--- a/cca_zoo/deepmodels/dcca_barlow_twins.py
+++ b/cca_zoo/deepmodels/dcca_barlow_twins.py
@@ -17,10 +17,10 @@ class BarlowTwins(DCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int,
-            encoders: Iterable[BaseEncoder] = [Encoder, Encoder],
-            lam=1,
+        self,
+        latent_dims: int,
+        encoders: Iterable[BaseEncoder] = [Encoder, Encoder],
+        lam=1,
     ):
         """
         Constructor class for Barlow Twins
diff --git a/cca_zoo/deepmodels/dcca_noi.py b/cca_zoo/deepmodels/dcca_noi.py
index 6ab09929..9b8ac318 100644
--- a/cca_zoo/deepmodels/dcca_noi.py
+++ b/cca_zoo/deepmodels/dcca_noi.py
@@ -16,14 +16,14 @@ class DCCA_NOI(DCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int,
-            N: int,
-            encoders=None,
-            r: float = 0,
-            rho: float = 0.2,
-            eps: float = 1e-9,
-            shared_target: bool = False,
+        self,
+        latent_dims: int,
+        N: int,
+        encoders=None,
+        r: float = 0,
+        rho: float = 0.2,
+        eps: float = 1e-9,
+        shared_target: bool = False,
     ):
         """
         Constructor class for DCCA
@@ -58,7 +58,7 @@ def forward(self, *args):
         z = []
         # Users architecture + final linear layer
         for i, (encoder, linear_layer) in enumerate(
-                zip(self.encoders, self.linear_layers)
+            zip(self.encoders, self.linear_layers)
         ):
             z.append(linear_layer(encoder(args[i])))
         return z
diff --git a/cca_zoo/deepmodels/dcca_sdl.py b/cca_zoo/deepmodels/dcca_sdl.py
index dcbde8ac..de5676ee 100644
--- a/cca_zoo/deepmodels/dcca_sdl.py
+++ b/cca_zoo/deepmodels/dcca_sdl.py
@@ -15,15 +15,15 @@ class DCCA_SDL(DCCA_NOI):
     """
 
     def __init__(
-            self,
-            latent_dims: int,
-            N: int,
-            encoders=None,
-            r: float = 0,
-            rho: float = 0.2,
-            eps: float = 1e-3,
-            shared_target: bool = False,
-            lam=0.5,
+        self,
+        latent_dims: int,
+        N: int,
+        encoders=None,
+        r: float = 0,
+        rho: float = 0.2,
+        eps: float = 1e-3,
+        shared_target: bool = False,
+        lam=0.5,
     ):
         """
         Constructor class for DCCA
diff --git a/cca_zoo/deepmodels/dccae.py b/cca_zoo/deepmodels/dccae.py
index f9185480..0a167e87 100644
--- a/cca_zoo/deepmodels/dccae.py
+++ b/cca_zoo/deepmodels/dccae.py
@@ -17,14 +17,14 @@ class DCCAE(_DCCA_base):
     """
 
     def __init__(
-            self,
-            latent_dims: int,
-            objective=objectives.MCCA,
-            encoders=None,
-            decoders=None,
-            r: float = 0,
-            eps: float = 1e-3,
-            lam=0.5,
+        self,
+        latent_dims: int,
+        objective=objectives.MCCA,
+        encoders=None,
+        decoders=None,
+        r: float = 0,
+        eps: float = 1e-3,
+        lam=0.5,
     ):
         """
         :param latent_dims: # latent dimensions
diff --git a/cca_zoo/deepmodels/dtcca.py b/cca_zoo/deepmodels/dtcca.py
index 6d0b0535..85209061 100644
--- a/cca_zoo/deepmodels/dtcca.py
+++ b/cca_zoo/deepmodels/dtcca.py
@@ -20,7 +20,7 @@ class DTCCA(DCCA):
     """
 
     def __init__(
-            self, latent_dims: int, encoders=None, r: float = 0, eps: float = 1e-3
+        self, latent_dims: int, encoders=None, r: float = 0, eps: float = 1e-3
     ):
         """
 
diff --git a/cca_zoo/deepmodels/dvcca.py b/cca_zoo/deepmodels/dvcca.py
index 65e1c311..b6fac805 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/objectives.py b/cca_zoo/deepmodels/objectives.py
index 7a094f57..0fb08caf 100644
--- a/cca_zoo/deepmodels/objectives.py
+++ b/cca_zoo/deepmodels/objectives.py
@@ -199,10 +199,10 @@ def loss(self, H1, H2):
 
         SigmaHat12 = (1.0 / (n - 1)) * H1bar.T @ H2bar
         SigmaHat11 = (1 - self.r) * (
-                1.0 / (n - 1)
+            1.0 / (n - 1)
         ) * H1bar.T @ H1bar + self.r * torch.eye(o1, device=H1.device)
         SigmaHat22 = (1 - self.r) * (
-                1.0 / (n - 1)
+            1.0 / (n - 1)
         ) * H2bar.T @ H2bar + self.r * torch.eye(o2, device=H2.device)
 
         SigmaHat11RootInv = torch.linalg.inv(
diff --git a/cca_zoo/deepmodels/trainers.py b/cca_zoo/deepmodels/trainers.py
index 1462b267..6e36f71e 100644
--- a/cca_zoo/deepmodels/trainers.py
+++ b/cca_zoo/deepmodels/trainers.py
@@ -12,20 +12,20 @@
 
 class CCALightning(LightningModule):
     def __init__(
-            self,
-            model: _DCCA_base,
-            optimizer: Union[torch.optim.Optimizer, str] = "Adam",
-            learning_rate: float = 1e-3,
-            weight_decay: float = 0.1,
-            lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
-            StepLR_step_size: float = None,
-            StepLR_gamma: float = None,
-            lr_factor: float = None,
-            lr_patience: float = None,
-            OneCycleLR_max_lr: float = None,
-            OneCycleLR_epochs: float = None,
-            train_trajectories: float = None,
-            T: float = None,
+        self,
+        model: _DCCA_base,
+        optimizer: Union[torch.optim.Optimizer, str] = "Adam",
+        learning_rate: float = 1e-3,
+        weight_decay: float = 0.1,
+        lr_scheduler: torch.optim.lr_scheduler._LRScheduler = None,
+        StepLR_step_size: float = None,
+        StepLR_gamma: float = None,
+        lr_factor: float = None,
+        lr_patience: float = None,
+        OneCycleLR_max_lr: float = None,
+        OneCycleLR_epochs: float = None,
+        train_trajectories: float = None,
+        T: float = None,
     ):
         """
 
@@ -86,7 +86,7 @@ def configure_optimizers(self):
         if self.hparams.lr_scheduler is None:
             return optimizer
         elif isinstance(
-                self.hparams.lr_scheduler, torch.optim.lr_scheduler._LRScheduler
+            self.hparams.lr_scheduler, torch.optim.lr_scheduler._LRScheduler
         ):
             scheduler = self.hparams.lr_scheduler
         elif self.hparams.lr_scheduler == "StepLR":
@@ -143,9 +143,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: bool = False,
+        self,
+        loader: torch.utils.data.DataLoader,
+        train: bool = False,
     ):
         """
 
@@ -158,16 +158,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: bool = False,
+        self,
+        loader: torch.utils.data.DataLoader,
+        train: bool = False,
     ):
         """
 
@@ -190,9 +190,9 @@ def transform(
         return z_list
 
     def score(
-            self,
-            loader: torch.utils.data.DataLoader,
-            train: bool = False,
+        self,
+        loader: torch.utils.data.DataLoader,
+        train: bool = False,
     ):
         """
 
@@ -207,13 +207,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 predict_view(
-            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/model_selection/_search.py b/cca_zoo/model_selection/_search.py
index 57b20826..acdf7e47 100644
--- a/cca_zoo/model_selection/_search.py
+++ b/cca_zoo/model_selection/_search.py
@@ -114,11 +114,11 @@ def __init__(self, param_distributions, n_iter, *, random_state=None):
             for key in dist:
                 if isinstance(dist[key], Iterable):
                     if any(
-                            [
-                                not isinstance(view_param, Iterable)
-                                and not hasattr(view_param, "rvs")
-                                for view_param in dist[key]
-                            ]
+                        [
+                            not isinstance(view_param, Iterable)
+                            and not hasattr(view_param, "rvs")
+                            for view_param in dist[key]
+                        ]
                     ):
                         raise TypeError(
                             "Parameter value for at least one view is not iterable "
@@ -148,11 +148,11 @@ def __iter__(self):
                 if isinstance(v, Iterable):
                     # if each element is a distribution for each view (i.e. it is a non-string Iterable) then call return_param for each view
                     if any(
-                            [
-                                (isinstance(v_, Iterable) and not isinstance(v_, str))
-                                or hasattr(v_, "rvs")
-                                for v_ in v
-                            ]
+                        [
+                            (isinstance(v_, Iterable) and not isinstance(v_, str))
+                            or hasattr(v_, "rvs")
+                            for v_ in v
+                        ]
                     ):
                         params[k] = [self.return_param(v_) for v_ in v]
                     # if the parameter is shared across views then the list will just contain non-iterable values
@@ -180,17 +180,17 @@ def __len__(self):
 
 class BaseSearchCV(SKBaseSearchCV):
     def __init__(
-            self,
-            estimator,
-            *,
-            scoring=None,
-            n_jobs=None,
-            refit=True,
-            cv=None,
-            verbose=0,
-            pre_dispatch="2*n_jobs",
-            error_score=np.nan,
-            return_train_score=True,
+        self,
+        estimator,
+        *,
+        scoring=None,
+        n_jobs=None,
+        refit=True,
+        cv=None,
+        verbose=0,
+        pre_dispatch="2*n_jobs",
+        error_score=np.nan,
+        return_train_score=True,
     ):
         super().__init__(
             estimator=estimator,
@@ -625,18 +625,18 @@ class GridSearchCV(BaseSearchCV):
     _required_parameters = ["estimator", "param_grid"]
 
     def __init__(
-            self,
-            estimator,
-            param_grid,
-            *,
-            scoring=None,
-            n_jobs=None,
-            refit=True,
-            cv=None,
-            verbose=0,
-            pre_dispatch="2*n_jobs",
-            error_score=np.nan,
-            return_train_score=False,
+        self,
+        estimator,
+        param_grid,
+        *,
+        scoring=None,
+        n_jobs=None,
+        refit=True,
+        cv=None,
+        verbose=0,
+        pre_dispatch="2*n_jobs",
+        error_score=np.nan,
+        return_train_score=False,
     ):
         super().__init__(
             estimator=estimator,
@@ -889,20 +889,20 @@ class RandomizedSearchCV(BaseSearchCV):
     _required_parameters = ["estimator", "param_distributions"]
 
     def __init__(
-            self,
-            estimator,
-            param_distributions,
-            *,
-            n_iter=10,
-            scoring=None,
-            n_jobs=None,
-            refit=True,
-            cv=None,
-            verbose=0,
-            pre_dispatch="2*n_jobs",
-            random_state=None,
-            error_score=np.nan,
-            return_train_score=False,
+        self,
+        estimator,
+        param_distributions,
+        *,
+        n_iter=10,
+        scoring=None,
+        n_jobs=None,
+        refit=True,
+        cv=None,
+        verbose=0,
+        pre_dispatch="2*n_jobs",
+        random_state=None,
+        error_score=np.nan,
+        return_train_score=False,
     ):
         self.param_distributions = param_distributions
         self.n_iter = n_iter
diff --git a/cca_zoo/models/cca_base.py b/cca_zoo/models/cca_base.py
index f7b3da35..f0294084 100644
--- a/cca_zoo/models/cca_base.py
+++ b/cca_zoo/models/cca_base.py
@@ -23,13 +23,13 @@ class _CCA_Base(BaseEstimator, MultiOutputMixin, RegressorMixin):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale=True,
-            centre=True,
-            copy_data=True,
-            accept_sparse=False,
-            random_state: Union[int, np.random.RandomState] = None,
+        self,
+        latent_dims: int = 1,
+        scale=True,
+        centre=True,
+        copy_data=True,
+        accept_sparse=False,
+        random_state: Union[int, np.random.RandomState] = None,
     ):
         """
         Constructor for _CCA_Base
@@ -113,7 +113,7 @@ def correlations(self, views: Iterable[np.ndarray], y=None, **kwargs):
         all_corrs = []
         for x, y in itertools.product(transformed_views, repeat=2):
             all_corrs.append(
-                np.diag(np.corrcoef(x.T, y.T)[: self.latent_dims, self.latent_dims:])
+                np.diag(np.corrcoef(x.T, y.T)[: self.latent_dims, self.latent_dims :])
             )
         all_corrs = np.array(all_corrs).reshape(
             (len(views), len(views), self.latent_dims)
@@ -121,10 +121,10 @@ def correlations(self, views: Iterable[np.ndarray], y=None, **kwargs):
         return all_corrs
 
     def plot_latent(
-            self,
-            views: Iterable[np.ndarray],
-            test_views: Iterable[np.ndarray] = None,
-            title="",
+        self,
+        views: Iterable[np.ndarray],
+        test_views: Iterable[np.ndarray] = None,
+        title="",
     ):
         scores = self.transform(views)
         if test_views is not None:
@@ -140,8 +140,8 @@ def score(self, views: Iterable[np.ndarray], y=None, **kwargs):
         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 _centre_scale(self, views: Iterable[np.ndarray]):
diff --git a/cca_zoo/models/gcca.py b/cca_zoo/models/gcca.py
index dd5bc322..00fed385 100644
--- a/cca_zoo/models/gcca.py
+++ b/cca_zoo/models/gcca.py
@@ -41,14 +41,14 @@ class GCCA(rCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            c: Union[Iterable[float], float] = None,
-            view_weights: Iterable[float] = None,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        c: Union[Iterable[float], float] = None,
+        view_weights: Iterable[float] = None,
     ):
         """
         Constructor for GCCA
@@ -91,9 +91,9 @@ def _setup_evp(self, views: Iterable[np.ndarray], K=None):
             Q.append(view_weight * view @ np.linalg.inv(view_cov) @ view.T)
         Q = np.sum(Q, axis=0)
         Q = (
-                np.diag(np.sqrt(np.sum(K, axis=0)))
-                @ Q
-                @ np.diag(np.sqrt(np.sum(K, axis=0)))
+            np.diag(np.sqrt(np.sum(K, axis=0)))
+            @ Q
+            @ np.diag(np.sqrt(np.sum(K, axis=0)))
         )
         return views, Q, None
 
@@ -139,19 +139,19 @@ class KGCCA(GCCA):
     """
 
     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,
-            kernel: Iterable[Union[float, callable]] = None,
-            gamma: Iterable[float] = None,
-            degree: Iterable[float] = None,
-            coef0: Iterable[float] = None,
-            kernel_params: Iterable[dict] = None,
+        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,
+        kernel: Iterable[Union[float, callable]] = None,
+        gamma: Iterable[float] = None,
+        degree: Iterable[float] = None,
+        coef0: Iterable[float] = None,
+        kernel_params: Iterable[dict] = None,
     ):
         """
         Constructor for PLS
@@ -225,9 +225,9 @@ def _setup_evp(self, views: Iterable[np.ndarray], K=None):
             Q.append(view_weight * view @ np.linalg.inv(view_cov) @ view.T)
         Q = np.sum(Q, axis=0)
         Q = (
-                np.diag(np.sqrt(np.sum(K, axis=0)))
-                @ Q
-                @ np.diag(np.sqrt(np.sum(K, axis=0)))
+            np.diag(np.sqrt(np.sum(K, axis=0)))
+            @ Q
+            @ np.diag(np.sqrt(np.sum(K, axis=0)))
         )
         return kernels, Q, None
 
diff --git a/cca_zoo/models/innerloop.py b/cca_zoo/models/innerloop.py
index 1dc52afa..fe40eed2 100644
--- a/cca_zoo/models/innerloop.py
+++ b/cca_zoo/models/innerloop.py
@@ -6,9 +6,6 @@
 from sklearn.exceptions import ConvergenceWarning
 from sklearn.linear_model import (
     ElasticNet,
-    Lasso,
-    LinearRegression,
-    Ridge,
     SGDRegressor,
 )
 from sklearn.utils._testing import ignore_warnings
@@ -23,11 +20,11 @@
 
 class _InnerLoop:
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol: float = 1e-5,
-            initialization: str = "unregularized",
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol: float = 1e-5,
+        initialization: str = "unregularized",
+        random_state=None,
     ):
         """
         :param max_iter: maximum number of iterations to perform if tol is not reached
@@ -59,15 +56,15 @@ def _initialize(self):
                     random_state=self.random_state,
                     tol=self.tol,
                 )
-                    ._fit(*self.views)
-                    .scores
+                ._fit(*self.views)
+                .scores
             )
         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.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
@@ -85,12 +82,19 @@ def _fit(self, *views: np.ndarray):
         self._initialize()
 
         self.track = {}
+        self.track["converged"] = False
         # Iterate until convergence
         self.track["objective"] = []
         for _ in range(self.max_iter):
             self._inner_iteration()
+            if np.isnan(self.scores).sum() > 0:
+                warnings.warn(
+                    f"Some scores are nan. Usually regularisation is too high."
+                )
+                break
             self.track["objective"].append(self._objective())
-            if _ > 0 and self._early_stop():
+            if _ > 1 and self._early_stop():
+                self.track["converged"] = True
                 break
             self.old_scores = self.scores.copy()
         return self
@@ -118,11 +122,11 @@ def _objective(self) -> int:
 
 class PLSInnerLoop(_InnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-5,
-            initialization: str = "unregularized",
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-5,
+        initialization: str = "unregularized",
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
@@ -138,7 +142,9 @@ def _check_params(self):
     def _inner_iteration(self):
         # Update each view using loop update function
         for i, view in enumerate(self.views):
-            self._update_view(i)
+            # if no nans
+            if np.isnan(self.scores).sum() == 0:
+                self._update_view(i)
 
     @abstractmethod
     def _update_view(self, view_index: int):
@@ -153,7 +159,7 @@ 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]
@@ -161,8 +167,8 @@ def _update_view(self, view_index: int):
     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:
@@ -171,13 +177,13 @@ def _early_stop(self) -> bool:
 
 class PMDInnerLoop(PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-5,
-            initialization: str = "unregularized",
-            c=None,
-            positive=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-5,
+        initialization: str = "unregularized",
+        c=None,
+        positive=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
@@ -219,7 +225,7 @@ 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],
@@ -232,12 +238,12 @@ def _update_view(self, view_index: int):
 
 class ParkhomenkoInnerLoop(PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-5,
-            initialization: str = "unregularized",
-            c=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-5,
+        initialization: str = "unregularized",
+        c=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
@@ -261,26 +267,26 @@ def _update_view(self, view_index: int):
         targets = np.ma.array(self.scores, mask=False)
         targets.mask[view_index] = True
         w = self.views[view_index].T @ targets.sum(axis=0).filled()
-        _check_converged_weights(w, view_index)
         w /= np.linalg.norm(w)
-        w = _soft_threshold(w, self.c[view_index] / 2)
         _check_converged_weights(w, view_index)
+        w = _soft_threshold(w, self.c[view_index] / 2)
         self.weights[view_index] = w / np.linalg.norm(w)
+        _check_converged_weights(w, view_index)
         self.scores[view_index] = self.views[view_index] @ self.weights[view_index]
 
 
 class ElasticInnerLoop(PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-5,
-            initialization: str = "unregularized",
-            c=None,
-            l1_ratio=None,
-            maxvar=True,
-            stochastic=True,
-            positive=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-5,
+        initialization: str = "unregularized",
+        c=None,
+        l1_ratio=None,
+        maxvar=True,
+        stochastic=True,
+        positive=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
@@ -305,80 +311,28 @@ def _check_params(self):
         self.regressors = []
         for alpha, l1_ratio, positive in zip(self.c, self.l1_ratio, self.positive):
             if self.stochastic:
-                if l1_ratio == 0:
-                    self.regressors.append(
-                        SGDRegressor(
-                            penalty="l2",
-                            alpha=alpha / len(self.views),
-                            fit_intercept=False,
-                            tol=self.tol,
-                            warm_start=True,
-                            random_state=self.random_state,
-                        )
-                    )
-                elif l1_ratio == 1:
-                    self.regressors.append(
-                        SGDRegressor(
-                            penalty="l1",
-                            alpha=alpha / len(self.views),
-                            fit_intercept=False,
-                            tol=self.tol,
-                            warm_start=True,
-                            random_state=self.random_state,
-                        )
-                    )
-                else:
-                    self.regressors.append(
-                        SGDRegressor(
-                            penalty="elasticnet",
-                            alpha=alpha / len(self.views),
-                            l1_ratio=l1_ratio,
-                            fit_intercept=False,
-                            tol=self.tol,
-                            warm_start=True,
-                            random_state=self.random_state,
-                        )
+                self.regressors.append(
+                    SGDRegressor(
+                        penalty="elasticnet",
+                        alpha=alpha / len(self.views),
+                        l1_ratio=l1_ratio,
+                        fit_intercept=False,
+                        tol=self.tol,
+                        warm_start=True,
+                        random_state=self.random_state,
                     )
+                )
             else:
-                if alpha == 0:
-                    self.regressors.append(LinearRegression(fit_intercept=False))
-                elif l1_ratio == 0:
-                    if positive:
-                        self.regressors.append(
-                            ElasticNet(
-                                alpha=alpha / len(self.views),
-                                l1_ratio=0,
-                                fit_intercept=False,
-                                warm_start=True,
-                                positive=positive,
-                                random_state=self.random_state,
-                            )
-                        )
-                    else:
-                        self.regressors.append(
-                            Ridge(alpha=alpha / len(self.views), fit_intercept=False)
-                        )
-                elif l1_ratio == 1:
-                    self.regressors.append(
-                        Lasso(
-                            alpha=alpha / len(self.views),
-                            fit_intercept=False,
-                            warm_start=True,
-                            positive=positive,
-                            random_state=self.random_state,
-                        )
-                    )
-                else:
-                    self.regressors.append(
-                        ElasticNet(
-                            alpha=alpha / len(self.views),
-                            l1_ratio=l1_ratio,
-                            fit_intercept=False,
-                            warm_start=True,
-                            positive=positive,
-                            random_state=self.random_state,
-                        )
+                self.regressors.append(
+                    ElasticNet(
+                        alpha=alpha / len(self.views),
+                        l1_ratio=l1_ratio,
+                        fit_intercept=False,
+                        warm_start=True,
+                        positive=positive,
+                        random_state=self.random_state,
                     )
+                )
 
     def _update_view(self, view_index: int):
         """
@@ -386,22 +340,27 @@ def _update_view(self, view_index: int):
         :return: updated weights
         """
         if self.maxvar:
+            # For MAXVAR we rescale the targets
             target = self.scores.mean(axis=0)
-            target /= np.linalg.norm(target)
+            target /= np.linalg.norm(target) / np.sqrt(self.n)
         else:
             target = self.scores[view_index - 1]
-        self._elastic_solver(self.views[view_index], target, view_index)
-        _check_converged_weights(self.weights[view_index], view_index)
+        # Solve the elastic regression
+        self.weights[view_index] = self._elastic_solver(
+            self.views[view_index], target, view_index
+        )
+        # For SUMCOR we rescale the projections
+        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)
+            )
         self.scores[view_index] = self.views[view_index] @ self.weights[view_index]
 
     @ignore_warnings(category=ConvergenceWarning)
     def _elastic_solver(self, X, y, view_index):
-        self.weights[view_index] = np.expand_dims(
-            self.regressors[view_index].fit(X, y.ravel()).coef_, 1
-        )
-        self.weights[view_index] /= np.linalg.norm(
-            self.views[view_index] @ self.weights[view_index]
-        ) / np.sqrt(self.n)
+        return np.expand_dims(self.regressors[view_index].fit(X, y.ravel()).coef_, 1)
 
     def _objective(self):
         views = len(self.views)
@@ -410,13 +369,13 @@ def _objective(self):
         l1 = c * ratio
         l2 = c * (1 - ratio)
         total_objective = 0
+        target = self.scores.mean(axis=0)
         for i in range(views):
-            target = self.scores.mean(axis=0)
-            objective = (
-                    views
-                    * np.linalg.norm(self.views[i] @ self.weights[i] - target) ** 2
-                    / (2 * self.n)
-            )
+            if self.maxvar:
+                target /= np.linalg.norm(target) / np.sqrt(self.n)
+            objective = np.linalg.norm(
+                self.views[i] @ self.weights[i] - target
+            ) ** 2 / (2 * self.n)
             l1_pen = l1[i] * np.linalg.norm(self.weights[i], ord=1)
             l2_pen = l2[i] * np.linalg.norm(self.weights[i], ord=2)
             total_objective += objective + l1_pen + l2_pen
@@ -432,15 +391,15 @@ def _early_stop(self) -> bool:
 
 class ADMMInnerLoop(ElasticInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-5,
-            initialization: str = "unregularized",
-            mu=None,
-            lam=None,
-            c=None,
-            eta=None,
-            random_state=None,
+        self,
+        max_iter: int = 100,
+        tol=1e-5,
+        initialization: str = "unregularized",
+        mu=None,
+        lam=None,
+        c=None,
+        eta=None,
+        random_state=None,
     ):
         super().__init__(
             max_iter=max_iter,
@@ -499,9 +458,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,
@@ -517,9 +476,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(
@@ -551,15 +510,15 @@ def _prox_lam_g(self, x):
 
 class SpanCCAInnerLoop(_InnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-5,
-            initialization: str = "unregularized",
-            c=None,
-            regularisation="l0",
-            rank=1,
-            random_state=None,
-            positive=False,
+        self,
+        max_iter: int = 100,
+        tol=1e-5,
+        initialization: str = "unregularized",
+        c=None,
+        regularisation="l0",
+        rank=1,
+        random_state=None,
+        positive=False,
     ):
         super().__init__(
             max_iter=max_iter,
@@ -612,15 +571,15 @@ def _inner_iteration(self):
 
 class SWCCAInnerLoop(PLSInnerLoop):
     def __init__(
-            self,
-            max_iter: int = 100,
-            tol=1e-20,
-            initialization: str = "unregularized",
-            regularisation="l0",
-            c=None,
-            sample_support: int = None,
-            random_state=None,
-            positive=False,
+        self,
+        max_iter: int = 100,
+        tol=1e-20,
+        initialization: str = "unregularized",
+        regularisation="l0",
+        c=None,
+        sample_support: int = None,
+        random_state=None,
+        positive=False,
     ):
         super().__init__(
             max_iter=max_iter,
@@ -654,8 +613,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
+        ).T @ targets.sum(axis=0).filled()
         self.weights[view_index] = self.update(
             self.weights[view_index],
             self.c[view_index],
diff --git a/cca_zoo/models/iterative.py b/cca_zoo/models/iterative.py
index aab90457..bb89d060 100644
--- a/cca_zoo/models/iterative.py
+++ b/cca_zoo/models/iterative.py
@@ -25,16 +25,16 @@ class _Iterative(_CCA_Base):
     """
 
     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 = "unregularized",
-            tol: float = 1e-9,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        deflation="cca",
+        max_iter: int = 100,
+        initialization: str = "unregularized",
+        tol: float = 1e-9,
     ):
         """
         Constructor for _Iterative
@@ -99,6 +99,9 @@ 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")
+                break
         return self
 
     def _deflate(self, residual, score, loading):
@@ -111,8 +114,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)
@@ -156,15 +159,15 @@ class PLS_ALS(_Iterative):
     """
 
     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: str = "unregularized",
+        tol: float = 1e-9,
     ):
         """
         Constructor for PLS
@@ -222,27 +225,27 @@ class ElasticCCA(_Iterative):
     >>> rng=np.random.RandomState(0)
     >>> X1 = rng.random((10,5))
     >>> X2 = rng.random((10,5))
-    >>> model = ElasticCCA(c=[0.1,0.1],l1_ratio=[0.5,0.5], random_state=0)
+    >>> model = ElasticCCA(c=[1e-1,1e-1],l1_ratio=[0.5,0.5], random_state=0)
     >>> model.fit((X1,X2)).score((X1,X2))
-    array([0.95818397])
+    array([0.9316638])
     """
 
     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 = "unregularized",
-            tol: float = 1e-9,
-            c: Union[Iterable[float], float] = None,
-            l1_ratio: Union[Iterable[float], float] = None,
-            maxvar: bool = True,
-            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",
+        max_iter: int = 100,
+        initialization: str = "unregularized",
+        tol: float = 1e-9,
+        c: Union[Iterable[float], float] = None,
+        l1_ratio: Union[Iterable[float], float] = None,
+        maxvar: bool = True,
+        stochastic=False,
+        positive: Union[Iterable[bool], bool] = None,
     ):
         """
         Constructor for ElasticCCA
@@ -323,21 +326,21 @@ class CCA_ALS(ElasticCCA):
     >>> X2 = rng.random((10,3))
     >>> model = CCA_ALS(random_state=0)
     >>> model.fit((X1,X2)).score((X1,X2))
-    array([0.85890619])
+    array([0.858906])
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            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,
+        max_iter: int = 100,
+        initialization: str = "random",
+        tol: float = 1e-9,
+        stochastic=True,
+        positive: Union[Iterable[bool], bool] = None,
     ):
         """
         Constructor for CCA_ALS
@@ -365,7 +368,7 @@ def __init__(
             scale=scale,
             positive=positive,
             random_state=random_state,
-            c=1e-3,
+            c=1e-5,
             maxvar=False,
         )
 
@@ -399,19 +402,19 @@ class SCCA(ElasticCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            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,
+        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,
     ):
         """
         Constructor for SCCA
@@ -477,17 +480,17 @@ class PMD(_Iterative):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            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,
+        c: Union[Iterable[float], float] = None,
+        max_iter: int = 100,
+        initialization: str = "unregularized",
+        tol: float = 1e-9,
+        positive: Union[Iterable[bool], bool] = None,
     ):
         """
         Constructor for PMD
@@ -556,16 +559,16 @@ class ParkhomenkoCCA(_Iterative):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            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,
+        c: Union[Iterable[float], float] = None,
+        max_iter: int = 100,
+        initialization: str = "unregularized",
+        tol: float = 1e-9,
     ):
         """
         Constructor for ParkhomenkoCCA
@@ -631,19 +634,19 @@ class SCCA_ADMM(_Iterative):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            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,
+        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,
     ):
         """
         Constructor for SCCA_ADMM
@@ -719,19 +722,19 @@ 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,
+        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,
     ):
         """
 
@@ -797,19 +800,19 @@ class SWCCA(_Iterative):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            max_iter: int = 500,
-            initialization: str = "uniform",
-            tol: float = 1e-9,
-            regularisation="l0",
-            c: Union[Iterable[Union[float, int]], Union[float, int]] = None,
-            sample_support=None,
-            positive=False,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        max_iter: int = 500,
+        initialization: str = "uniform",
+        tol: float = 1e-9,
+        regularisation="l0",
+        c: Union[Iterable[Union[float, int]], Union[float, int]] = None,
+        sample_support=None,
+        positive=False,
     ):
         """
 
diff --git a/cca_zoo/models/mcca.py b/cca_zoo/models/mcca.py
index 4d0294c1..0762c21f 100644
--- a/cca_zoo/models/mcca.py
+++ b/cca_zoo/models/mcca.py
@@ -40,14 +40,14 @@ class MCCA(rCCA):
     """
 
     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,
+        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 MCCA
@@ -95,7 +95,7 @@ def _solve_evp(self, views: Iterable[np.ndarray], C, D=None, **kwargs):
         eigvecs = eigvecs * eigvals
         eigvecs = eigvecs[:, idx].real
         self.weights = [
-            eigvecs[split: self.splits[i + 1]]
+            eigvecs[split : self.splits[i + 1]]
             for i, split in enumerate(self.splits[:-1])
         ]
 
@@ -128,19 +128,19 @@ class KCCA(MCCA):
     """
 
     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,
-            kernel: Iterable[Union[float, callable]] = None,
-            gamma: Iterable[float] = None,
-            degree: Iterable[float] = None,
-            coef0: Iterable[float] = None,
-            kernel_params: Iterable[dict] = None,
+        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,
+        kernel: Iterable[Union[float, callable]] = None,
+        gamma: Iterable[float] = None,
+        degree: Iterable[float] = None,
+        coef0: Iterable[float] = None,
+        kernel_params: Iterable[dict] = None,
     ):
         """
         :param latent_dims: number of latent dimensions to fit
diff --git a/cca_zoo/models/ncca.py b/cca_zoo/models/ncca.py
index 2fe741c6..13473fc7 100644
--- a/cca_zoo/models/ncca.py
+++ b/cca_zoo/models/ncca.py
@@ -27,15 +27,15 @@ class NCCA(_CCA_Base):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale=True,
-            centre=True,
-            copy_data=True,
-            accept_sparse=False,
-            random_state: Union[int, np.random.RandomState] = None,
-            nearest_neighbors=None,
-            gamma: Iterable[float] = None,
+        self,
+        latent_dims: int = 1,
+        scale=True,
+        centre=True,
+        copy_data=True,
+        accept_sparse=False,
+        random_state: Union[int, np.random.RandomState] = None,
+        nearest_neighbors=None,
+        gamma: Iterable[float] = None,
     ):
         """
         Constructor for NCCA
@@ -87,9 +87,9 @@ def fit(self, views: Iterable[np.ndarray], y=None, **kwargs):
         ]
         S = self.Ws[0] @ self.Ws[1]
         U, S, Vt = np.linalg.svd(S)
-        self.f = U[:, 1: self.latent_dims + 1] * np.sqrt(self.n)
-        self.g = Vt[1: self.latent_dims + 1, :].T * np.sqrt(self.n)
-        self.S = S[1: self.latent_dims + 1]
+        self.f = U[:, 1 : self.latent_dims + 1] * np.sqrt(self.n)
+        self.g = Vt[1 : self.latent_dims + 1, :].T * np.sqrt(self.n)
+        self.S = S[1 : self.latent_dims + 1]
         return self
 
     def transform(self, views: Iterable[np.ndarray], y=None, **kwargs):
diff --git a/cca_zoo/models/rcca.py b/cca_zoo/models/rcca.py
index 4247240b..7d8b45ea 100644
--- a/cca_zoo/models/rcca.py
+++ b/cca_zoo/models/rcca.py
@@ -45,15 +45,15 @@ class rCCA(_CCA_Base):
     """
 
     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,
-            accept_sparse=None,
+        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,
+        accept_sparse=None,
     ):
         """
         Constructor for rCCA
@@ -124,12 +124,12 @@ def _solve_evp(self, views: Iterable[np.ndarray], C, D=None, **kwargs):
             eigvecs = eigvecs[:, idx].real
             w_y = views[1].T @ np.diag(1 / np.sqrt(self.Bs[1])) @ eigvecs
             w_x = (
-                    views[0].T
-                    @ np.diag(1 / self.Bs[0])
-                    @ self.R_12
-                    @ np.diag(1 / np.sqrt(self.Bs[1]))
-                    @ eigvecs
-                    / np.sqrt(eigvals[idx])
+                views[0].T
+                @ np.diag(1 / self.Bs[0])
+                @ self.R_12
+                @ np.diag(1 / np.sqrt(self.Bs[1]))
+                @ eigvecs
+                / np.sqrt(eigvals[idx])
             )
             self.weights = [w_x, w_y]
         else:
@@ -141,7 +141,7 @@ def _solve_evp(self, views: Iterable[np.ndarray], C, D=None, **kwargs):
             self.weights = [
                 Vt.T
                 @ np.diag(1 / np.sqrt(B))
-                @ eigvecs[split: self.splits[i + 1], : self.latent_dims]
+                @ eigvecs[split : self.splits[i + 1], : self.latent_dims]
                 for i, (split, Vt, B) in enumerate(
                     zip(self.splits[:-1], views, self.Bs)
                 )
@@ -151,11 +151,11 @@ def _two_view_evp(self, Us, Ss):
         Rs = [U @ np.diag(S) for U, S in zip(Us, Ss)]
         self.R_12 = Rs[0].T @ Rs[1]
         M = (
-                np.diag(1 / np.sqrt(self.Bs[1]))
-                @ self.R_12.T
-                @ np.diag(1 / self.Bs[0])
-                @ self.R_12
-                @ np.diag(1 / np.sqrt(self.Bs[1]))
+            np.diag(1 / np.sqrt(self.Bs[1]))
+            @ self.R_12.T
+            @ np.diag(1 / self.Bs[0])
+            @ self.R_12
+            @ np.diag(1 / np.sqrt(self.Bs[1]))
         )
         return M, None
 
@@ -207,12 +207,12 @@ class CCA(rCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
     ):
         """
         Constructor for CCA
@@ -264,12 +264,12 @@ class PLS(rCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
     ):
         """
         Constructor for PLS
diff --git a/cca_zoo/models/tcca.py b/cca_zoo/models/tcca.py
index 740ab737..2aba4047 100644
--- a/cca_zoo/models/tcca.py
+++ b/cca_zoo/models/tcca.py
@@ -43,13 +43,13 @@ class TCCA(_CCA_Base):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale=True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            c: Union[Iterable[float], float] = None,
+        self,
+        latent_dims: int = 1,
+        scale=True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        c: Union[Iterable[float], float] = None,
     ):
         """
         Constructor for TCCA
@@ -190,19 +190,19 @@ class KTCCA(TCCA):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            scale: bool = True,
-            centre=True,
-            copy_data=True,
-            random_state=None,
-            eps=1e-3,
-            c: Union[Iterable[float], float] = None,
-            kernel: Iterable[Union[float, callable]] = None,
-            gamma: Iterable[float] = None,
-            degree: Iterable[float] = None,
-            coef0: Iterable[float] = None,
-            kernel_params: Iterable[dict] = None,
+        self,
+        latent_dims: int = 1,
+        scale: bool = True,
+        centre=True,
+        copy_data=True,
+        random_state=None,
+        eps=1e-3,
+        c: Union[Iterable[float], float] = None,
+        kernel: Iterable[Union[float, callable]] = None,
+        gamma: Iterable[float] = None,
+        degree: Iterable[float] = None,
+        coef0: Iterable[float] = None,
+        kernel_params: Iterable[dict] = None,
     ):
         """
         Constructor for TCCA
diff --git a/cca_zoo/probabilisticmodels/vcca.py b/cca_zoo/probabilisticmodels/vcca.py
index eaf5ed3a..5fef9b52 100644
--- a/cca_zoo/probabilisticmodels/vcca.py
+++ b/cca_zoo/probabilisticmodels/vcca.py
@@ -22,12 +22,12 @@ class VariationalCCA(_CCA_Base):
     """
 
     def __init__(
-            self,
-            latent_dims: int = 1,
-            copy_data=True,
-            random_state: int = 0,
-            num_samples=100,
-            num_warmup=100,
+        self,
+        latent_dims: int = 1,
+        copy_data=True,
+        random_state: int = 0,
+        num_samples=100,
+        num_warmup=100,
     ):
         super().__init__(
             latent_dims=latent_dims,
@@ -101,6 +101,6 @@ def _model(self, views: Iterable[np.ndarray]):
                     obs=X_,
                 )
                 for i, (X_, psi_, mu_, W_) in enumerate(
-                zip(views, psi, mu, self.weights_list)
-            )
+                    zip(views, psi, mu, self.weights_list)
+                )
             ]
diff --git a/cca_zoo/utils/plotting.py b/cca_zoo/utils/plotting.py
index 2857fb7c..42a691e0 100644
--- a/cca_zoo/utils/plotting.py
+++ b/cca_zoo/utils/plotting.py
@@ -47,7 +47,7 @@ def cv_plot(cv_results_):
         param_pairs = list(itertools.product(unique_x, unique_y))
         for pair in param_pairs:
             mask = (cv_results_[param_cols[-2]] == pair[0]) & (
-                    cv_results_[param_cols[-1]] == pair[1]
+                cv_results_[param_cols[-1]] == pair[1]
             )
             sub_dfs.append(cv_results_.loc[mask].iloc[:, :-2])
             sub_scores.append(cv_results_[mask].mean_test_score)
@@ -85,9 +85,9 @@ def cv_plot(cv_results_):
 
 
 def plot_latent_train_test(
-        train_scores: Union[Tuple[np.ndarray], List[np.ndarray]],
-        test_scores: Union[Tuple[np.ndarray], List[np.ndarray]] = None,
-        title="",
+    train_scores: Union[Tuple[np.ndarray], List[np.ndarray]],
+    test_scores: Union[Tuple[np.ndarray], List[np.ndarray]] = None,
+    title="",
 ):
     """
     Makes a pair plot showing the projections of each view against each other for each dimensions. Coloured by train and test
@@ -116,10 +116,10 @@ def plot_latent_train_test(
 
 
 def plot_latent_label(
-        scores: Union[Tuple[np.ndarray], List[np.ndarray]],
-        labels=None,
-        label_name=None,
-        title="",
+    scores: Union[Tuple[np.ndarray], List[np.ndarray]],
+    labels=None,
+    label_name=None,
+    title="",
 ):
     """
     Makes a pair plot showing the projections of each view against each other for each dimensions. Coloured by categorical label
diff --git a/setup.py b/setup.py
index 6406732c..8b6d1f29 100644
--- a/setup.py
+++ b/setup.py
@@ -11,6 +11,7 @@
     EXTRA_PACKAGES["deep"] = f.read()
 with open("./requirements/probabilistic.txt", "r") as f:
     EXTRA_PACKAGES["probabilistic"] = f.read()
+EXTRA_PACKAGES["all"] = EXTRA_PACKAGES["deep"]+'\n'+EXTRA_PACKAGES["probabilistic"]
 
 setup(
     name="cca_zoo",
@@ -22,7 +23,7 @@
     license="MIT",
     author="jameschapman",
     description=(
-        "Canonical Correlation Analysis Zoo: CCA, GCCA, MCCA, DCCA, DGCCA, DVCCA, DCCAE, KCCA and regularised variants including sparse CCA , ridge CCA and elastic CCA"
+        "Canonical Correlation Analysis Zoo: A collection of Regularized, Deep Learning based, Kernel, and Probabilistic methods in a scikit-learn style framework"
     ),
     long_description=long_description,
     long_description_content_type="text/markdown",


Test patch:
diff --git a/cca_zoo/test/test_deepmodels.py b/cca_zoo/test/test_deepmodels.py
index dcad9bb9..01cd4124 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)
     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
     )
 
 
@@ -256,8 +256,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 22f00b08..51e750be 100644
--- a/cca_zoo/test/test_models.py
+++ b/cca_zoo/test/test_models.py
@@ -65,56 +65,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
     )
 
 
@@ -182,24 +182,27 @@ 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
 
 
 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
+    ).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],
+        random_state=0,
+    ).fit([X, Y])
+    nnals = CCA_ALS(
+        latent_dims=latent_dims, tol=1e-9, positive=True, random_state=0
     ).fit([X, Y])
-    nnals = CCA_ALS(latent_dims=latent_dims, tol=1e-9, positive=True).fit([X, Y])
-    nnscca = SCCA(latent_dims=latent_dims, tol=1e-9, positive=True, c=[1e-4, 1e-5]).fit(
-        (X, Y)
-    )
     assert np.testing.assert_array_less(-1e-9, nnelastic.weights[0]) is None
     assert np.testing.assert_array_less(-1e-9, nnelastic.weights[1]) is None
     assert np.testing.assert_array_less(-1e-9, nnals.weights[0]) is None
@@ -221,11 +224,15 @@ def test_sparse_methods():
     c1 = [1e-1]
     c2 = [1e-1]
     param_grid = {"c": [c1, c2]}
-    parkhomenko_cv = GridSearchCV(ParkhomenkoCCA(random_state=rng), param_grid=param_grid).fit([X, Y])
+    parkhomenko_cv = GridSearchCV(
+        ParkhomenkoCCA(random_state=rng), param_grid=param_grid
+    ).fit([X, Y])
     c1 = [2e-2]
     c2 = [1e-2]
     param_grid = {"c": [c1, c2]}
-    admm_cv = GridSearchCV(SCCA_ADMM(random_state=rng), param_grid=param_grid).fit([X, Y])
+    admm_cv = GridSearchCV(SCCA_ADMM(random_state=rng), param_grid=param_grid).fit(
+        [X, Y]
+    )
     c1 = loguniform(1e-1, 2e-1)
     c2 = loguniform(1e-1, 2e-1)
     param_grid = {"c": [c1, c2], "l1_ratio": [[0.9], [0.9]]}
@@ -259,10 +266,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
     )
 
 
@@ -307,11 +314,11 @@ def test_VCCA():
         ).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
+            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.