diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 4249c05..b649ba6 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -79,7 +79,7 @@ class FancyPtycho(CDIModel):
             self.register_buffer('mask',
                                  t.tensor(mask, dtype=t.bool))
         
-        probe_guess = t.tensor(probe_guess, dtype=t.complex4)
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
         obj_guess = t.tensor(obj_guess, dtype=t.complex64)
 
         # We rescale the probe here so it learns at the same rate as the
@@ -533,18 +533,7 @@ class FancyPtycho(CDIModel):
         else:
             return translations
 
-
-    def get_rhos(self):
-        # If this is the general unified mode model
-        if self.weights.dim() >= 2:
-            Ws = self.weights.detach().cpu().numpy()
-            rhos_out = np.matmul(np.swapaxes(Ws, 1, 2), Ws.conj())
-            return rhos_out
-        # This is the purely incoherent case
-        else:
-            return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
-                            dtype=np.complex64)
-
+        
     def center_probes(self, iterations=4):
         """Centers the probes
         
@@ -557,95 +546,97 @@ class FancyPtycho(CDIModel):
         self.probe.data = centered_probe.to(device=self.probe.data.device)
 
 
-    def tidy_probes(self, normalize=False, tidy_each_frame=False):
+    def tidy_probes(self):
         """Tidies up the probes
         
         What we want to do here is use all the information on all the probes
         to calculate a natural basis for the experiment, and update all the
         density matrices to operate in that updated basis
+
+        As a first step, we calculate the state of the light field across the
+        full experiment, using the weight matrices and basis probes. Then, we
+        use an SVD to update the basis probes so they form an eigenbasis of
+        the implied density matrix for the full experiment.
+
+        Next, the weight matrices for each shot are recalculated so that the
+        probes generated by weights * basis_probes for each shot are themselves
+        an eigenbasis for that individual shot's density matrix.
         """
         
-        # First we treat the purely incoherent case
-
-        # I don't love this pattern of using an if statement with a return
-        # to catch this case, but because it's so much simpler than the
-        # unified mode case I think it's appropriate
+        # First we treat the incoherent but stable  case, where the weights are
+        # just one per-shot overall weight
         if self.weights.dim() == 1:
             probe = self.probe.detach().cpu().numpy()
-            ortho_probes = analysis.orthogonalize_probes(probe)
-            self.probe.data = t.as_tensor(
-                ortho_probes,
-                device=self.probe.device,
-                dtype=self.probe.dtype)
+            ortho_probes = analysis.orthogonalize_probes_t(self.probe.detach())
+            self.probe.data = ortho_probes
             return
 
-        # This is for the unified mode case
+        # What follows is for the unified OPRP and incoherent multi-mode model,
+        # where each shot has it's own matrix of weights such that the probe
+        # state for each shot is self.weights @ self.probe
         
-        # We concatenate all the weight matrices
-        all_weights = t.cat(t.unbind(self.weights.detach().cpu(), dim=0), dim=0)
-
-        # We use that to calculate the density matrix of the full experiment,
-        # normalized by number of exposures
-        overall_rho = (t.mm(all_weights.transpose(0,1), all_weights.conj())
-                       / self.weights.shape[0])
+        # We concatenate all the weight matrices, to come up with a state
+        # corresponding to the summed light field across all the exposures.
+        # This state will have a large number of modes, but all built from
+        # the same small number of basis modes
+        all_weights = t.cat(t.unbind(self.weights.detach(), dim=0), dim=0)
 
         # We generate the orthogonal probes based on this full-experiment
-        # density matrix. We also keep the transform matrix A
-        probe = self.probe.detach().cpu().numpy()
-        ortho_probes, A = analysis.orthogonalize_probes(
-            probe, density_matrix=overall_rho,
-            keep_transform=True, normalize=normalize)
+        # representation of the light field.
+        ortho_probes, reexpressed_weights = \
+            analysis.orthogonalize_probes_t(
+                self.probe.detach(),
+                weight_matrix=all_weights,
+                return_reexpressed_weights=True
+            )
 
-        # We apply A to the weight matrices to update them along with the
-        # probes
-        new_weights = t.matmul(
-            t.as_tensor(A).transpose(0,1),
-            self.weights.detach().cpu().transpose(-2,-1)).transpose(-2,-1)
-
-        self.probe.data = t.as_tensor(
-            ortho_probes, device=self.probe.device, dtype=self.probe.dtype)
+        # We just orthogonalized the incoherent sum of all the exposures
+        # across the full experiment, so the output probes are normalized so
+        # that their intensity matches the summed intensity across the full
+        # experiment. We divide their amplitudes by the square root of the
+        # number of shots so that we now have a set of probes corresponding
+        # to the mean shot
+        ortho_probes /= np.sqrt(self.weights.shape[0])
+        reexpressed_weights *= np.sqrt(self.weights.shape[0])
         
-        self.weights.data = new_weights.to(device=self.weights.device,
-                                           dtype=self.weights.dtype)
+        # We now replace the shot-to-shot weights with the versions that have
+        # been re-expressed in the new basis. 
+        new_weights = t.stack(t.split(reexpressed_weights,
+                                      self.weights.shape[1]), dim=0)
 
-        # At this point, we now have a new set of basis probes, which are the
-        # eigenbasis for the full-experiment density matrix, and we have
-        # re-expressed all the shot-to-shot weight matrices in that basis.
-        # But, the shot-to-shot probes (self.weights <matmul> self.probe)
-        # are still exactly the same as they were before.
-        #
-        # Oftentimes, we also want the shot-to-shot weights to be re-expressed
-        # so that the shot-to-shot probes are the eigenbasis for each
-        # individual shot's density matrix. That's what we do below.
-        #
-        if tidy_each_frame:
+        # And we save it back to the model
+        self.probe.data = ortho_probes.to(
+            device=self.probe.device, dtype=self.probe.dtype)        
+        self.weights.data = new_weights.to(
+            device=self.weights.device, dtype=self.weights.dtype)
 
-            # TODO: I need to check that this really works
-            
-            dm_rank = self.weights.shape[-2]
-            
-            for idx in range(self.weights.shape[0]):
-                weights = self.weights.data[idx].detach().cpu()
-                rho = t.mm(weights.transpose(0,1), weights.conj())
-                
-                ortho_probes, A = analysis.orthogonalize_probes(
-                    self.probe.detach().cpu(), density_matrix=rho,
-                    normalize=False, keep_transform=True)
+        # NOTE: I used to have this part as an option, with "tidy_each_frame",
+        # because it took such a long time. Now that I've rewritten it properly,
+        # it's quite fast and so I removed the kwarg because there's really
+        # no situation where you woudn't want to do this.
 
-                
-                new_weights = t.linalg.pinv(A)[:dm_rank, :].conj()
-                    
-                self.weights.data[idx] = new_weights.to(
-                    device=self.weights.device,
-                    dtype=self.weights.dtype)
+        # Now, we seek to edit the shot-to-shot weight matrices such that
+        # self.weights[i] @ self.probes will be properly orthogonalized for
+        # all i.
 
+        # All we need to know about the probes is that they are orthogonalized
+        # and the intensity within each probe mode
+        probe_sqrt_intensities = t.linalg.norm(self.probe.data, dim=(-2,-1))
 
+        # This does a super fast batched computation
+        U, S, Vh = t.linalg.svd(self.weights.data * probe_sqrt_intensities,
+                                full_matrices=False)
+
+        # We discard the U matrix and re-multiply S & Vh
+        self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities)
+
+    
     def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs):
         def get_probes(idx):
             basis_prs = self.probe * self.probe_support[..., :, :]
             prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs,
                         axis=-4)
-            ortho_probes = analysis.orthogonalize_probes(prs)
+            ortho_probes = analysis.orthogonalize_probes_t(prs)
 
             if mode.lower() == 'amplitude':
                 return np.abs(ortho_probes.detach().cpu().numpy())
@@ -737,15 +728,20 @@ class FancyPtycho(CDIModel):
              fig=fig,
              basis=self.probe_basis,
              units=self.units)),
-        ('Average Density Matrix Amplitudes',
+        ('Average Weight Matrix Amplitudes',
          lambda self, fig: p.plot_amplitude(
-             np.nanmean(np.abs(self.get_rhos()), axis=0),
+             np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
              fig=fig),
          lambda self: len(self.weights.shape) >= 2),
-        ('% Power in Top Mode (only accurate after tidy_probes)',
+        ('% of Power in Top Mode',
          lambda self, fig, dataset: p.plot_nanomap(
              self.corrected_translations(dataset),
-             analysis.calc_top_mode_fraction(self.get_rhos()),
+             100 * t.stack([
+                 analysis.calc_mode_power_fractions(
+                     self.probe.data,
+                     weight_matrix=self.weights.data[i])[0]
+                 for i in range(self.weights.shape[0])
+             ], dim=0),
              fig=fig,
              units=self.units),
          lambda self: len(self.weights.shape) >= 2),
diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py
index 741c040..b56b89b 100644
--- a/src/cdtools/tools/analysis/analysis.py
+++ b/src/cdtools/tools/analysis/analysis.py
@@ -13,69 +13,164 @@ from scipy import linalg as sla
 from scipy import special
 from scipy import optimize as opt
 
-__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
+__all__ = ['product_svd', 'orthogonalize_probes_t',
+           'orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
            'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
            'calc_frc', 'calc_vn_entropy', 'calc_top_mode_fraction',
+           'calc_mode_power_fractions',
            'calc_rms_error', 'calc_fidelity', 'calc_generalized_rms_error',
            'remove_phase_ramp', 'remove_amplitude_exponent',
            'standardize_reconstruction_set']
 
 
+def product_svd(A, B):
+    """ Computes the SVD of A @ B
+
+    This function uses a method which uses a QR decomposition of
+    A and B to calculate the final reduced SVD, without explicitly
+    calculating the full matrix. The output is defined such that
+    A B = U S Vh, and as a reduced SVD
+
+    Parameters
+    ----------
+    A : array
+        An nxr matrix
+    B : array
+        An rxm matrix
+
+    Returns
+    -------
+    U : array
+        An nxr matrix of left singular vectors
+    S : array
+        An length-r array, t.diag(S) is the diagonal matrix of singular values
+    Vh : array
+        And rxm matrix of the conjugate-transposed right singular vectors
+    """
+    # Handle the case of numpy input
+    return_np = False
+    if isinstance(A, np.ndarray):
+        A = t.as_tensor(A)
+        return_np = True
+    if isinstance(B, np.ndarray):
+        B = t.as_tensor(B)
+        return_np = True
+
+    # We take a QR decomposition of the two matrices
+    Qa, Ra = t.linalg.qr(A)
+    Qb, Rb = t.linalg.qr(B.conj().transpose(0,1))
+
+    # And now we take the SVD of the product of the two R matrices
+    U, S, Vh = t.linalg.svd(t.matmul(Ra, Rb.conj().transpose(0,1)),
+                            full_matrices=False)
+
+    # And build back the final SVD of the product matrix
+    U_final = t.matmul(Qa, U)
+    Vh_final = t.matmul(Vh, Qb.conj().transpose(0,1))
+
+    if return_np:
+        U_final = U_final.numpy()
+        S = S.numpy()
+        Vh_final = Vh_final.numpy()
+    
+    return U_final, S, Vh_final
+
+
 def orthogonalize_probes_t(
         probes,
-        density_matrix=None,
-        keep_transform=False,
-        normalize=False,
+        weight_matrix=None,
         n_probe_dims=2,
+        return_reexpressed_weights=False,
 ):
     """ Orthogonalizes a set of incoherently mixing probes
 
-    TODO: actually make this, and replace ortho_probes with a fully
-    pytorch-based function
+    ## TODO fully replace orthogonalize_probes with orthogonalize_probes_t
     
-    Any set of probe modes defines a density matrix (a.k.a mutual coherence
-    function) that is the ultimate description of the state of the light
-    field. This function takes any set of probe modes - not necessarily
-    orthogonalized - and returns an orthogonalized set of probe modes.
-    Formally, it returns the eigenbasis of the density matrix, ordered
-    from largest to smallest eigenvalue.
+    This function takes any set of probe modes for mixed-mode ptychography,
+    which are considered to define a mutual coherence function, and returns
+    an orthogonalized set probe modes which refer to the same mutual coherence
+    function. The orthogonalized modes are extracted via a singular value
+    decomposition and are unique up to a global per-mode phase factor.
 
-    If normalize is set to True, then it will return the normalized
-    eigenbasis. Otherwise, it will return a scaled version of the eigenbasis,
-    so that the returned probes can be used directly for multi-mode
-    ptychography.
+    If a weight matrix is explicitly given, the function will instead
+    orthogonalize the light field defined by weight_matrix @ probes. It
+    accomplishes this via a method which avoids explicitly constructing this
+    potentially large matrix. This can be useful for Orthogonal Probe
+    Relaxation ptychography. In this case, one may have a large stacked
+    matrix of shot-to-shot weights but a small basis set of probes.
 
-    If a density matrix is explicitly given, it will instead
-    consider the problem of extracting the eigenbasis of the matrix
-    probes * denstity_matrix * probes^dagger, where probes is the
-    column matrix of the given probe functions. This latter problem arises
-    in the generalization of the probe mixing model, and reduces to the
-    simpler case when the density matrix is equal to the identity matrix
+    In addition to returning the orthogonalized probe modes, this function
+    also returns a re-expression of the original weight matrix in the basis
+    of the orthogonalized probe modes, such that:
+
+    reexpressed_weight_matrix @ orthogonalized_probes = weight_matrix @ probes.
+
+    This re-expressed weight matrix is guaranteed to have orthonormalized rows,
+    such that:
+
+    reexpressed_weight_matrix^\\dagger @ reexpressed_weight_matrix = I.
+
+    There is usually no reason to use the re-expressed weight matrix.
+    However, it can be useful in situations where the individual rows in the 
+    weight matrix have a specific meaning, such as an exposure number, which
+    should be preserved.
     
-    If the parameter "keep_transform" is set, the function will additionally
-    return the matrix A such that A * ortho_probes^dagger = probes^dagger
-    TODO: is the above right, or are ortho_probes and probes flipped?
+    Warning! The shape of the output orthogonalized_probes may not be equal to
+    the shape of input probes, when a weight matrix is used. If the input
+    weight matrix has m < l rows, where l is the number of probe modes,
+    then the output orthogonalized probes will have length m, not length l.
 
     Parameters
     ----------
     probes : array
-        An l x (<n_probe_dims>) complex array representing  a stack of probes
-    density_matrix : array
-        An optional l x l density matrix further elaborating on the state
-    keep_transform : bool
-        Default False, whether to return the map from probes to ortho_probes
-    normalize : bool
-        Default False, whether to normalize the probe modes
+        An l x (<n_probe_pix>) array representing a stack of l probes
+    weight_matrix : array
+        Optional, an m x l weight matrix further elaborating on the state
     n_probe_dims : int
-        Default 2, the number of trailing dimensions defining each probe
+        Default is 2, the number of trailing dimensions for each probe state
 
     Returns
     -------
-    ortho_probes: array
-        An l x (<n_probe_dims>) complex array representing a stack of probes
+    orthogonalized_probes : array
+        A min(m,l) x (<n_probe_dims>) array representing a stack of probes
+    reexpressed_weight_matrix : array
+        A the original weight matrix, re-expressed to work with the new probes 
     """
-    pass
 
+    return_np = False
+    if isinstance(probes, np.ndarray):
+        probes = t.as_tensor(probes)
+        return_np = True
+    if weight_matrix is not None and isinstance(weight_matrix, np.ndarray):
+        weight_matrix = t.as_tensor(weight_matrix)
+        return_np = True
+    
+    n_probe_pix = np.prod(np.array(probes.shape[-n_probe_dims:]))
+    probes_mat = probes.reshape(probes.shape[:-n_probe_dims] +
+                                (n_probe_pix,))
+    
+    if weight_matrix is None:
+        # We just calculate a straight up SVD of the probes
+        U, S, Vh = t.linalg.svd(probes_mat, full_matrices=False)
+        
+    else:
+        U, S, Vh = product_svd(weight_matrix, probes_mat)
+
+    output_shape = (-1,) + tuple(n for n in probes.shape[1:])
+    orthogonalized_probes = (S[:,None] * Vh).reshape(output_shape)
+    reexpressed_weight_matrix = U
+
+    if return_np:
+        orthogonalized_probes = orthogonalized_probes.numpy()
+        reexpressed_weight_matrix = reexpressed_weight_matrix.numpy()
+
+    to_return = (orthogonalized_probes,)
+    if return_reexpressed_weights:
+        to_return += (reexpressed_weight_matrix,)
+    
+    return to_return
+
+        
 
 def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, normalize=False):
     """Orthogonalizes a set of incoherently mixing probes
@@ -668,7 +763,58 @@ def calc_vn_entropy(matrix):
         entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig)
         return entropy
 
+
+def calc_mode_power_fractions(
+        probes,
+        weight_matrix=None,
+        n_probe_dims=2,
+        assume_preorthogonalized=False,
+):
+    """Calculates the fraction of total power in each orthogonalized mode
+
+    This code first orthogonalizes the probe modes, so the result of this
+    function are independent of the particular way that the multi-mode
+    breakdown is expressed.
+
+
+    Parameters
+    ----------
+    probes : array
+        An l x (<n_probe_pix>) array representing a stack of l probes
+    weight_matrix : array
+        Optional, an m x l weight matrix further elaborating on the state
+    n_probe_dims : int
+        Default is 2, the number of trailing dimensions for each probe state
+    assume_preorthogonalized : bool
+        Default is False. If True, will not orthogonalize the probes
     
+    Returns
+    -------
+    power_fractions : array
+        The fraction of the total power in each mode
+    """
+
+    if not assume_preorthogonalized:
+        ortho_probes, reexpressed_weights = \
+            orthogonalize_probes_t(
+                probes,
+                weight_matrix=weight_matrix,
+                n_probe_dims=n_probe_dims,
+            )
+    else:
+        weight_slice = np.s_[...,] + np.s_[None,] * n_probe_dims
+        if weight_matrix is None:
+            ortho_probes = probes
+        else:
+            ortho_probes = t.sum(weight_matrix[weight_slice] * probes,
+                                 axis=-(n_probe_dims + 1))
+        
+    dims = [-d-1 for d in range(n_probe_dims)]
+    power = t.sum(t.abs(ortho_probes)**2, dim=dims)
+    power_fractions = power / t.sum(power)
+    return power_fractions
+
+
 def calc_top_mode_fraction(matrix):
     """Calculates the fraction of total power in the top mode of a density matrix
     
@@ -687,6 +833,7 @@ def calc_top_mode_fraction(matrix):
     entropy: float or np.array
         The fraction of power in the top mode of each matrix
     """
+    # TODO depricate
 
     if len(matrix.shape) == 3:
         # Get the eigenvalues
diff --git a/src/cdtools/tools/plotting/plotting.py b/src/cdtools/tools/plotting/plotting.py
index 5b08b68..8478c01 100644
--- a/src/cdtools/tools/plotting/plotting.py
+++ b/src/cdtools/tools/plotting/plotting.py
@@ -780,11 +780,13 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non
             im = im.reshape(-1,im.shape[-2],im.shape[-1])
             im_idx = axes[1].image_idx
 
-            if event.key == 'up' or event.button == 'up' \
-               or event.key == 'left':
+            if (event.key == 'up'
+                or (hasattr(event, 'button') and event.button == 'up') 
+                or event.key == 'left'):
                 im_idx = (im_idx - 1) % im.shape[0]
-            if event.key == 'down' or event.button == 'down' \
-               or event.key == 'right':
+            if (event.key == 'down'
+                or (hasattr(event, 'button') and event.button == 'down')
+                or event.key == 'right'):
                 im_idx = (im_idx + 1) % im.shape[0]
 
             axes[1].image_idx=im_idx
diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py
index 2c7d1f8..3c837a2 100644
--- a/tests/tools/test_analysis.py
+++ b/tests/tools/test_analysis.py
@@ -1,11 +1,157 @@
 import numpy as np
 from scipy import linalg as la
+from scipy.sparse import linalg as spla
 import torch as t
 from itertools import combinations
 
 from cdtools.tools import analysis, initializers
 
 
+def test_product_svd():
+
+    rank = 4
+    shape_A = (12, rank)
+    shape_B = (rank, 9)
+    A = np.random.rand(*shape_A) + 1j * np.random.rand(*shape_A)
+    B = np.random.rand(*shape_B) + 1j * np.random.rand(*shape_B)
+
+    AB = np.matmul(A,B)
+    U_1, S_1, Vh_1 = t.linalg.svd(t.as_tensor(AB), full_matrices=False)
+
+    U_2, S_2, Vh_2 = analysis.product_svd(t.as_tensor(A),t.as_tensor(B))
+    check_AB = U_2 @ t.diag_embed(S_2).to(dtype=Vh_2.dtype) @ Vh_2
+
+    # So, at a minimum, U S Vh = AB
+    assert np.allclose(AB, check_AB.numpy())
+
+    # SVD is only defined up to an arbitrary complex valued phase per
+    # singular vector, so all we can ask for in the comparison is that the
+    # magnitudes here are
+    assert np.allclose(S_1[:rank].numpy(), S_2.numpy())
+    prod_U = U_1[:,:rank].transpose(0,1).conj() @ U_2
+    prod_Vh = Vh_1[:rank,:] @ Vh_2.transpose(0,1).conj()
+    assert np.allclose(t.abs(prod_U).numpy(), np.eye(rank))
+    assert np.allclose(t.abs(prod_Vh).numpy(), np.eye(rank))
+    # Confirms that the phases are consistent between the two, I think
+    # it's redundant with the first check but I'm not sure
+    assert np.allclose(prod_Vh.numpy(), prod_U.numpy())
+    
+    # test with numpy
+    U_3, S_3, Vh_3 = analysis.product_svd(A,B)
+    assert isinstance(U_3, np.ndarray)
+    assert isinstance(S_3, np.ndarray)
+    assert isinstance(Vh_3, np.ndarray)
+    
+    assert np.allclose(S_1[:rank].numpy(), S_3)
+    prod_U = U_1[:,:rank].transpose(0,1).numpy().conj() @ U_3
+    prod_Vh = Vh_1[:rank,:].numpy() @ Vh_3.transpose().conj()
+    assert np.allclose(np.abs(prod_U), np.eye(rank))
+    assert np.allclose(np.abs(prod_Vh), np.eye(rank))
+    # Confirms that the phases are consistent between the two, I think
+    # it's redundant with the first check but I'm not sure
+    assert np.allclose(prod_Vh, prod_U)
+
+
+def test_orthogonalize_probes_t():
+
+    op = analysis.orthogonalize_probes_t
+    
+    probe_xs = np.arange(64) - 32
+    probe_ys = np.arange(76) - 38
+    probe_Ys, probe_Xs = np.meshgrid(probe_ys, probe_xs)
+    probe_Rs = np.sqrt(probe_Xs**2 + probe_Ys**2)
+
+    probes = np.array([10*np.exp(-probe_Rs**2 / (2 * 10**2 + 1j)),
+                       3*np.exp(-probe_Rs**2 / (2 * 12**2 - 3j)),
+                       1*np.exp(-probe_Rs**2 / (2 * 15**2))])
+
+    weight_matrix_none = None
+    weight_matrix_single = np.random.randn(1,3) + 1j * np.random.randn(1,3)
+    weight_matrix_small = np.random.randn(2,3) + 1j * np.random.randn(2,3)
+    weight_matrix_medium = np.random.randn(3,3) + 1j * np.random.randn(3,3)
+    weight_matrix_large = np.random.randn(7,3) + 1j * np.random.randn(7,3)
+
+    weight_matrices = [
+        weight_matrix_none,
+        weight_matrix_single,
+        weight_matrix_small,
+        weight_matrix_medium,
+        weight_matrix_large
+    ]
+    
+    for weight_matrix in weight_matrices:
+        ortho_probes_np, rwm_np = op(probes, weight_matrix=weight_matrix)
+        assert isinstance(ortho_probes_np, np.ndarray)
+        assert isinstance(rwm_np, np.ndarray)
+        
+        probes_t = t.as_tensor(probes)
+        wm_t = (t.as_tensor(weight_matrix) if weight_matrix is not None
+                else weight_matrix)
+
+        ortho_probes_t, rwm_t = op(probes_t, weight_matrix=wm_t)
+        assert t.is_tensor(ortho_probes_t)
+        assert t.is_tensor(rwm_t)
+
+        assert np.allclose(ortho_probes_np, ortho_probes_t.numpy())
+        assert np.allclose(rwm_np, rwm_t.numpy())
+
+        # Now we test a if wm @ ortho_probes is actually the original
+        # input
+        if weight_matrix is not None:
+            realized_probes = np.tensordot(weight_matrix, probes, axes=1)
+        else:
+            realized_probes = probes
+
+        calculated_probes = np.tensordot(rwm_np, ortho_probes_np, axes=1)
+            
+        assert np.allclose(realized_probes, calculated_probes)
+        
+        # Now we test if the orthogonalized probes are orthogonalized
+        reshaped_probes = ortho_probes_np.reshape(
+            (ortho_probes_np.shape[0],
+             ortho_probes_np.shape[1] * ortho_probes_np.shape[2]))
+        products = np.matmul(reshaped_probes,
+                             reshaped_probes.conj().transpose())
+
+        if weight_matrix is not None:
+            output_nmodes = min(weight_matrix.shape[0], probes.shape[0])
+        else:
+            output_nmodes = probes.shape[0]
+
+        for i in range(output_nmodes):
+            for j in range(i+1, output_nmodes):
+                assert np.isclose(products[i,j], 0)
+
+        # And now we test if they multiply to the same density matrix as
+        # the original probes + weight matrix
+        reshaped_realized_probes = realized_probes.reshape(
+            (realized_probes.shape[0],
+             realized_probes.shape[1] * realized_probes.shape[2]))
+        
+        dm_original = np.matmul(
+            reshaped_realized_probes.conj().transpose(),
+            reshaped_realized_probes
+        )
+        dm_output = np.matmul(
+            reshaped_probes.conj().transpose(),
+            reshaped_probes
+        )
+        assert np.allclose(dm_original, dm_output)
+
+        # And finally, we confirm that what we have are the eigenvectors/values
+        # of that density matrix
+        w, v = spla.eigsh(dm_original, k=output_nmodes)
+        assert np.allclose(w, np.diag(products))
+
+        cross_products = np.matmul(reshaped_probes, np.sqrt(w) * v)
+        # The abs accounts for the fact that the phase of the eigenvectors
+        # is undefined
+        assert np.allclose(np.abs(cross_products), np.abs(products))
+
+            
+        
+
+    
 def test_orthogonalize_probes():
 
     # The test strategy should be to define a few non-orthogonal probes