jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def pose_loss(self, pairwise_corr,  pairwise_t_matrix, pairwise_t_matrix_gt, record_len):
        """
        Args:
            pairwise_corr: [B, L, L, 3]
            pairwise_t_matrix/pairwise_t_matrix_gt: [B,L,L,4,4]
            record_len: list, shape [B]
        """

        pairwise_t_matrix_gt = pairwise_t_matrix_gt.float()
        B, L = pairwise_t_matrix.shape[:2]
        mask = torch.zeros((B, L, L), device = pairwise_t_matrix.device)

        for b in range(B):
            N = record_len[b]
            mask[b,:N,:N] = 1
            mask[b,range(N), range(N)] = 0
        
        pair_corr_gt = torch.solve(pairwise_t_matrix_gt.transpose(-2,-1), pairwise_t_matrix.transpose(-2,-1))[0].transpose(-2,-1)
 
        yaw = pairwise_corr[..., 2] # [B,L,L]
        yaw_gt = torch.rad2deg(torch.atan2(pair_corr_gt[...,1,0], pair_corr_gt[...,0,0])) # [B,L,L]

        x = pairwise_corr[..., 0] # [B,L,L]
        x_gt = pair_corr_gt[..., 0,3]

        y = pairwise_corr[..., 1] # [B,L,L]
        y_gt = pair_corr_gt[..., 1,3]

        mask = mask.bool()
        mask = mask.view(B,L,L) # [B, L, L, ]

        input_x = torch.masked_select(x, mask)
        target_x = torch.masked_select(x_gt, mask)

        input_y = torch.masked_select(y, mask)
        target_y = torch.masked_select(y_gt, mask)

        input_yaw = torch.masked_select(yaw, mask)
        target_yaw = torch.masked_select(yaw_gt, mask)

        loss_x = self.pose_loss_func(input_x, target_x)
        loss_y = self.pose_loss_func(input_y, target_y)
        loss_yaw = self.pose_loss_func(input_yaw, target_yaw)

        lambda_trans = 2/3
        lambda_rot = 1/3

        return lambda_trans * (loss_x + loss_y) + lambda_rot * loss_yaw
//  [python]
def pose_loss(self, pairwise_corr,  pairwise_t_matrix, pairwise_t_matrix_gt, record_len):
        """
        Args:
            pairwise_corr: [B, L, L, 3]
            pairwise_t_matrix/pairwise_t_matrix_gt: [B,L,L,4,4]
            record_len: list, shape [B]
        """

        pairwise_t_matrix_gt = pairwise_t_matrix_gt.float()
        B, L = pairwise_t_matrix.shape[:2]
        mask = torch.zeros((B, L, L), device = pairwise_t_matrix.device)

        for b in range(B):
            N = record_len[b]
            mask[b,:N,:N] = 1
            mask[b,range(N), range(N)] = 0
        
        pair_corr_gt = torch.linalg.solve(pairwise_t_matrix.transpose(-2,-1), pairwise_t_matrix_gt.transpose(-2,-1)).transpose(-2,-1)
 
        yaw = pairwise_corr[..., 2] # [B,L,L]
        yaw_gt = torch.rad2deg(torch.atan2(pair_corr_gt[...,1,0], pair_corr_gt[...,0,0])) # [B,L,L]

        x = pairwise_corr[..., 0] # [B,L,L]
        x_gt = pair_corr_gt[..., 0,3]

        y = pairwise_corr[..., 1] # [B,L,L]
        y_gt = pair_corr_gt[..., 1,3]

        mask = mask.bool()
        mask = mask.view(B,L,L) # [B, L, L, ]

        input_x = torch.masked_select(x, mask)
        target_x = torch.masked_select(x_gt, mask)

        input_y = torch.masked_select(y, mask)
        target_y = torch.masked_select(y_gt, mask)

        input_yaw = torch.masked_select(yaw, mask)
        target_yaw = torch.masked_select(yaw_gt, mask)

        loss_x = self.pose_loss_func(input_x, target_x)
        loss_y = self.pose_loss_func(input_y, target_y)
        loss_yaw = self.pose_loss_func(input_yaw, target_yaw)

        lambda_trans = 2/3
        lambda_rot = 1/3

        return lambda_trans * (loss_x + loss_y) + lambda_rot * loss_yaw
//  [python]
def test_solve_errors(self, device, dtype):
        # solve expects batches of square matrices as input
        with self.assertRaisesRegex(RuntimeError, "must be batches of square matrices"):
            a = torch.randn(2, 3, 4, 3, dtype=dtype, device=device)
            b = torch.randn(2, 3, 4, 1, dtype=dtype, device=device)
            torch.solve(a, b)

        # solve expects compatible shapes for A x = b
        with self.assertRaisesRegex(RuntimeError, "Incompatible matrix sizes"):
            a = torch.randn(2, 3, 3, 3, dtype=dtype, device=device)
            b = torch.randn(2, 3, 2, 1, dtype=dtype, device=device)
            torch.solve(a, b)

        # if input is not solvable, RuntimeError is raised mentioning the first non-solvable batch
        def run_test_singular_input(batch_dim, n):
            a = torch.eye(3, 3, dtype=dtype, device=device).reshape((1, 3, 3)).repeat(batch_dim, 1, 1)
            a[n, -1, -1] = 0
            b = torch.randn(batch_dim, 3, 1, dtype=dtype, device=device)
            with self.assertRaisesRegex(RuntimeError, rf'For batch {n}: U\(3,3\) is zero'):
                torch.solve(a, b)

        for params in [(1, 0), (2, 0), (2, 1), (4, 0), (4, 2), (10, 2)]:
            run_test_singular_input(*params)

        # if out is non-empty then it should have correct sizes
        with self.assertRaisesRegex(RuntimeError, r'does not match broadcasted other shape'):
            out = torch.empty(1, dtype=dtype, device=device)
            A = torch.eye(3, dtype=dtype, device=device)
            b = torch.randn(3, 1, dtype=dtype, device=device)
            torch.solve(A, b, out=out)

        # if out is non-empty then it should also be Fortran contiguous
        with self.assertRaisesRegex(RuntimeError, r'tensor must be in batched column major'):
            out = torch.zeros(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
            self.assertFalse(out.is_contiguous())
            A = torch.eye(2, dtype=dtype, device=device).reshape((1, 2, 2)).repeat(2, 1, 1)
            b = torch.randn(2, 2, 2, dtype=dtype, device=device)
            torch.solve(A, b, out=out)
//  [python]
def test_solve_errors(self, device, dtype):
        # solve expects batches of square matrices as input
        with self.assertRaisesRegex(RuntimeError, "must be batches of square matrices"):
            a = torch.randn(2, 3, 4, 3, dtype=dtype, device=device)
            b = torch.randn(2, 3, 4, 1, dtype=dtype, device=device)
            torch.linalg.solve(a, b)

        # solve expects compatible shapes for A x = b
        with self.assertRaisesRegex(RuntimeError, "Incompatible matrix sizes"):
            a = torch.randn(2, 3, 3, 3, dtype=dtype, device=device)
            b = torch.randn(2, 3, 2, 1, dtype=dtype, device=device)
            torch.linalg.solve(a, b)

        # if input is not solvable, RuntimeError is raised mentioning the first non-solvable batch
        def run_test_singular_input(batch_dim, n):
            a = torch.eye(3, 3, dtype=dtype, device=device).reshape((1, 3, 3)).repeat(batch_dim, 1, 1)
            a[n, -1, -1] = 0
            b = torch.randn(batch_dim, 3, 1, dtype=dtype, device=device)
            with self.assertRaisesRegex(RuntimeError, rf'For batch {n}: U\(3,3\) is zero'):
                torch.linalg.solve(a, b)

        for params in [(1, 0), (2, 0), (2, 1), (4, 0), (4, 2), (10, 2)]:
            run_test_singular_input(*params)

        # if out is non-empty then it should have correct sizes
        with self.assertRaisesRegex(RuntimeError, r'does not match broadcasted other shape'):
            out = torch.empty(1, dtype=dtype, device=device)
            A = torch.eye(3, dtype=dtype, device=device)
            b = torch.randn(3, 1, dtype=dtype, device=device)
            torch.linalg.solve(A, b, out=out)

        # if out is non-empty then it should also be Fortran contiguous
        with self.assertRaisesRegex(RuntimeError, r'tensor must be in batched column major'):
            out = torch.zeros(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
            self.assertFalse(out.is_contiguous())
            A = torch.eye(2, dtype=dtype, device=device).reshape((1, 2, 2)).repeat(2, 1, 1)
            b = torch.randn(2, 2, 2, dtype=dtype, device=device)
            torch.linalg.solve(A, b, out=out)
//  [python]
def _setup_linear(self):
        """ Create parameters that allow fast linear stepping of state """
        w = self._w()
        B = _r2c(self.B) # (H N)
        P = _r2c(self.P)
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, '1 t n -> 1 (v t) n', v=self.repeat)
        P = repeat(P, 'r t n -> r (v t) n', v=self.repeat)
        Q = repeat(Q, 'r t n -> r (v t) n', v=self.repeat)
        w = repeat(w, 't n -> (v t) n', v=self.repeat)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=w.dtype, device=w.device) + 2*contract('r h n, h n, s h n -> h r s', Q, D, P).real) # (H R R)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        try:
            R = torch.solve(Q_D, R)[0] # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().cpu())).to(Q_D)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D, # (H N)
            "R": R, # (R H N)
            "P": P, # (R H N)
            "Q": Q, # (R H N)
            "B": B, # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w, # (H N)
        }
//  [python]
def _setup_linear(self):
        """ Create parameters that allow fast linear stepping of state """
        w = self._w()
        B = _r2c(self.B) # (H N)
        P = _r2c(self.P)
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, '1 t n -> 1 (v t) n', v=self.repeat)
        P = repeat(P, 'r t n -> r (v t) n', v=self.repeat)
        Q = repeat(Q, 'r t n -> r (v t) n', v=self.repeat)
        w = repeat(w, 't n -> (v t) n', v=self.repeat)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (torch.eye(self.rank, dtype=w.dtype, device=w.device) + 2*contract('r h n, h n, s h n -> h r s', Q, D, P).real) # (H R R)
        Q_D = rearrange(Q*D, 'r h n -> h r n')
        try:
            R = torch.linalg.solve(R, Q_D) # (H R N)
        except:
            R = torch.tensor(np.linalg.solve(R.to(Q_D).contiguous().detach().cpu(), Q_D.contiguous().detach().cpu())).to(Q_D)
        R = rearrange(R, 'h r n -> r h n')

        self.step_params = {
            "D": D, # (H N)
            "R": R, # (R H N)
            "P": P, # (R H N)
            "Q": Q, # (R H N)
            "B": B, # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w, # (H N)
        }
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")
            self.xp = np
            self.xsp = nsp

        self.RTOL = 1e-3

        self.A_shape = (4, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A = self.A + self.A.t()
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (4, 2)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.solve(self.B, self.A)[0]
//  [python]
def setUp(self) -> None:
        # The device can be specialised by a daughter class
        if not hasattr(self, "device"):
            self.device = torch.device("cpu")
            self.xp = np
            self.xsp = nsp

        self.RTOL = 1e-3

        self.A_shape = (4, 4)
        self.A = torch.randn(self.A_shape, dtype=torch.float64, device=self.device)
        self.A = self.A + self.A.t()
        self.A_csr = self.A.to_sparse_csr()
        self.B_shape = (4, 2)
        self.B = torch.randn(self.B_shape, dtype=torch.float64, device=self.device)

        self.x_ref = torch.linalg.solve(self.A, self.B)
//  [python]
def log_prob(self, X):
        logdetp = torch.logdet(X)
        pinvK = torch.solve(self.K, X)[1]
        trpinvK = torch.diagonal(pinvK, dim1=-2, dim2=-1).sum(-1)
        return self.C - 0.5 * ((self.nu + 2 * self.n) * logdetp + trpinvK)
//  [python]
def log_prob(self, X):
        logdetp = torch.logdet(X)
        pinvK = torch.linalg.solve(X, self.K)
        trpinvK = torch.diagonal(pinvK, dim1=-2, dim2=-1).sum(-1)
        return self.C - 0.5 * ((self.nu + 2 * self.n) * logdetp + trpinvK)
//  [python]
def test_solver_gradient_c4t(self):
        # Sparse solver:
        As1 = self.A_csr.detach().to(torch.float32).clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().to(torch.float32).clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = tsgucupy.sparse_solve_c4t(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().to(torch.float32).clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().to(torch.float32).clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.solve(Bd2, Ad2)[0]
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def test_solver_gradient_c4t(self):
        # Sparse solver:
        As1 = self.A_csr.detach().to(torch.float32).clone()
        As1.requires_grad = True
        Bd1 = self.B.detach().to(torch.float32).clone()
        Bd1.requires_grad = True
        As1.retain_grad()
        Bd1.retain_grad()
        x = tsgucupy.sparse_solve_c4t(As1, Bd1)
        loss = x.sum()
        loss.backward()

        # torch dense solver:
        Ad2 = self.A.detach().to(torch.float32).clone()
        Ad2.requires_grad = True
        Bd2 = self.B.detach().to(torch.float32).clone()
        Bd2.requires_grad = True
        Ad2.retain_grad()
        Bd2.retain_grad()
        x2 = torch.linalg.solve(Ad2, Bd2)
        loss_torch = x2.sum()
        loss_torch.backward()

        self.assertTrue(torch.isclose(x, x2, rtol=self.RTOL).all())

        nz_mask = As1.grad.to_dense() != 0.0
        self.assertTrue(torch.isclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], rtol=self.RTOL).all())
        self.assertTrue(torch.isclose(Bd1.grad, Bd2.grad, rtol=self.RTOL).all())
//  [python]
def get_c(self, *args, **kwargs):
        assert self.bdir == False
        device = self.ab.device
        N = self.N
        A = torch.roll(torch.eye(self.N, device=device),1,0)
        A = torch.clone(repeat(A, "N M -> C N M",C=self.num_a))
        A[:,0] = -self.ab[:self.a_channels] # construct A matrix
        I_AL = repeat(torch.eye(N, device=device) - torch.matrix_power(A, self.L), "C N M -> (C R) N M", R = self.D//self.num_a) # (I-A^L)
        return torch.solve(self.ab[self.a_channels:], I_AL)[0]
//  [python]
def get_c(self, *args, **kwargs):
        assert self.bdir == False
        device = self.ab.device
        N = self.N
        A = torch.roll(torch.eye(self.N, device=device),1,0)
        A = torch.clone(repeat(A, "N M -> C N M",C=self.num_a))
        A[:,0] = -self.ab[:self.a_channels] # construct A matrix
        I_AL = repeat(torch.eye(N, device=device) - torch.matrix_power(A, self.L), "C N M -> (C R) N M", R = self.D//self.num_a) # (I-A^L)
        return torch.linalg.solve(I_AL, self.ab[self.a_channels:], left=True)
//  [python]
def _setup_linear(self):
        """Create parameters that allow fast linear stepping of state"""
        w = self._w()
        B = _r2c(self.B)  # (H N)
        P = _r2c(self.P)
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, "1 t n -> 1 (v t) n", v=self.broadcast)
        P = repeat(P, "r t n -> r (v t) n", v=self.broadcast)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.broadcast)
        w = repeat(w, "t n -> (v t) n", v=self.broadcast)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (
            torch.eye(self.rank, dtype=w.dtype, device=w.device)
            + 2 * contract("r h n, h n, s h n -> h r s", Q, D, P).real
        )  # (H R R)
        Q_D = rearrange(Q * D, "r h n -> h r n")
        try:
            R = torch.solve(Q_D, R)[0]  # (H R N)
        except:
            R = torch.tensor(
                np.linalg.solve(
                    R.to(Q_D).contiguous().detach().cpu(),
                    Q_D.contiguous().detach().cpu(),
                )
            ).to(Q_D)
        R = rearrange(R, "h r n -> r h n")

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (R H N)
            "P": P,  # (R H N)
            "Q": Q,  # (R H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def _setup_linear(self):
        """Create parameters that allow fast linear stepping of state"""
        w = self._w()
        B = _r2c(self.B)  # (H N)
        P = _r2c(self.P)
        Q = P.conj()

        # Repeat w shape properly
        B = repeat(B, "1 t n -> 1 (v t) n", v=self.broadcast)
        P = repeat(P, "r t n -> r (v t) n", v=self.broadcast)
        Q = repeat(Q, "r t n -> r (v t) n", v=self.broadcast)
        w = repeat(w, "t n -> (v t) n", v=self.broadcast)

        # Prepare Linear stepping
        dt = torch.exp(self.log_dt)
        D = (2.0 / dt.unsqueeze(-1) - w).reciprocal()  # (H, N)
        R = (
            torch.eye(self.rank, dtype=w.dtype, device=w.device)
            + 2 * contract("r h n, h n, s h n -> h r s", Q, D, P).real
        )  # (H R R)
        Q_D = rearrange(Q * D, "r h n -> h r n")
        try:
            R = torch.linalg.solve(R, Q_D)  # (H R N)
        except:
            R = torch.tensor(
                np.linalg.solve(
                    R.to(Q_D).contiguous().detach().cpu(),
                    Q_D.contiguous().detach().cpu(),
                )
            ).to(Q_D)
        R = rearrange(R, "h r n -> r h n")

        self.step_params = {
            "D": D,  # (H N)
            "R": R,  # (R H N)
            "P": P,  # (R H N)
            "Q": Q,  # (R H N)
            "B": B,  # (1 H N)
            "E": 2.0 / dt.unsqueeze(-1) + w,  # (H N)
        }
//  [python]
def LM_opt(self, R_init):
        #Set current rotation based on initialisation conditions
        R_opt = R_init

        #Calculate intial total loss (sum of square errors, f(R)^t @ f(R)) and cost function values 
        loss, fR = self.cost_func(R_opt)

        #Calculate jacobian of d(f(R)) / d(delta(phi))
        J = self.jacobian_func(R_opt)
     
        #Get hessian matrix, J^t @ J   
        hess = J.t() @ J
        #Get gradient vector, J^t @ cost
        grad = J.t() @ fR

        #Set initial damping scaling
        u = self.tau * torch.max(hess.diag())

        #Optimisation loop
        for iter in range(self.max_iters-1):
            #Set damping based on current Hessian max value
            D = torch.eye(J.shape[1], device=J.device)
            D *= torch.max(torch.max(hess.diagonal(), D.diagonal()))

            #Scale damping
            damping = u*D
    
            #Solve Ax=b, (J^t@J + damping*I) = J^t@f(R)
            delta = -torch.solve((hess + damping), grad)[0]
 
            #Calculate SO(3) delta rotation to apply to current estimate
            R_delta = self.so3exp(delta)
            R_step = R_delta @ R_opt

            #Get new cost values and loss for the new rotation 
            loss_step, f_step = self.cost_func(R_step)
 
            #Setting previous to current for the next iteration
            fR_prev = fR.clone()
            
            #Work out LM rho value to determine if we accept or reject this new rotation
            rho_denom = delta.t() @ (u*delta-grad)
            rho_nom = loss - loss_step
            rho = rho_nom / rho_denom if rho_denom > 0 else torch.inf if rho_nom > 0 else -torch.inf

            #If happy then set the current values to the new setpoint
            if rho > 0:
                R_opt = R_step
                J = self.jacobian_func(R_opt)
                grad = J.T @ f_step
                hess = J.T @ J
                loss = loss_step
                fR = f_step

            #Change damping scaling based on the rho value above
            u = u*self.beta if rho < self.rho1 else u/self.gamma if rho > self.rho2 else u
            torch.clamp(u, 1e-8, 1e8)

            # stop conditions
            gcon = max(abs(grad)) < self.gtol #Very small gradients (stationary point)
            pcon = (delta**2).sum()**0.5 < self.ptol #Very small changes to rotation (rotation is negligible)
            fcon = ((fR_prev-fR)**2).sum() < ((self.ftol*fR)**2).sum() if rho > 0 else False #+ve rho but small change in cost (cost stationary point)

            if gcon or pcon or fcon:
                break

        #recover covariance from optimal hessian
        cov = torch.linalg.inv(hess)
        
        return R_opt, cov
//  [python]
def LM_opt(self, R_init):
        #Set current rotation based on initialisation conditions
        R_opt = R_init

        #Calculate intial total loss (sum of square errors, f(R)^t @ f(R)) and cost function values 
        loss, fR = self.cost_func(R_opt)

        #Calculate jacobian of d(f(R)) / d(delta(phi))
        J = self.jacobian_func(R_opt)
     
        #Get hessian matrix, J^t @ J   
        hess = J.t() @ J
        #Get gradient vector, J^t @ cost
        grad = J.t() @ fR

        #Set initial damping scaling
        u = self.tau * torch.max(hess.diag())

        #Optimisation loop
        for iter in range(self.max_iters-1):
            #Set damping based on current Hessian max value
            D = torch.eye(J.shape[1], device=J.device)
            D *= torch.max(torch.max(hess.diagonal(), D.diagonal()))

            #Scale damping
            damping = u*D
    
            #Solve Ax=b, (J^t@J + damping*I) = J^t@f(R)
            delta = -torch.linalg.solve((hess + damping), grad)
 
            #Calculate SO(3) delta rotation to apply to current estimate
            R_delta = self.so3exp(delta)
            R_step = R_delta @ R_opt

            #Get new cost values and loss for the new rotation 
            loss_step, f_step = self.cost_func(R_step)
 
            #Setting previous to current for the next iteration
            fR_prev = fR.clone()
            
            #Work out LM rho value to determine if we accept or reject this new rotation
            rho_denom = delta.t() @ (u*delta-grad)
            rho_nom = loss - loss_step
            rho = rho_nom / rho_denom if rho_denom > 0 else torch.inf if rho_nom > 0 else -torch.inf

            #If happy then set the current values to the new setpoint
            if rho > 0:
                R_opt = R_step
                J = self.jacobian_func(R_opt)
                grad = J.T @ f_step
                hess = J.T @ J
                loss = loss_step
                fR = f_step

            #Change damping scaling based on the rho value above
            u = u*self.beta if rho < self.rho1 else u/self.gamma if rho > self.rho2 else u
            torch.clamp(u, 1e-8, 1e8)

            # stop conditions
            gcon = max(abs(grad)) < self.gtol #Very small gradients (stationary point)
            pcon = (delta**2).sum()**0.5 < self.ptol #Very small changes to rotation (rotation is negligible)
            fcon = ((fR_prev-fR)**2).sum() < ((self.ftol*fR)**2).sum() if rho > 0 else False #+ve rho but small change in cost (cost stationary point)

            if gcon or pcon or fcon:
                break

        #recover covariance from optimal hessian
        cov = torch.linalg.inv(hess)
        
        return R_opt, cov
//  [python]
def _solve(self, A, b):
        if np.ndim(b) == 1: b = np.expand_dims(b, axis=1)
        if self.compute_device == "cpu":
            try:
                x = np.linalg.solve(A, b)
            except:
                x, res, rank, s = np.linalg.lstsq(A, b, rcond=None)
            return x
        elif self.compute_device == "gpu" or A.ndim < 3:
            try:
                import torch
                A = torch.from_numpy(A).cuda()
                b = torch.from_numpy(b).cuda()
                x = torch.solve(b, A)[0]
                return x.cpu().numpy()
            except Exception as e:
                warnings.warn(
                    "I encountered a problem using the GPU via pytorch. Falling back to Numpy and the CPU.")
                try:
                    x = np.linalg.solve(A, b)
                except:
                    x, res, rank, s = np.linalg.lstsq(A, b, rcond=None)
                return x
        elif self.compute_device == "multi-gpu":
            try:
                import torch
                n = min(len(A), torch.cuda.device_count())
                split_A = np.array_split(A, n)
                split_b = np.array_split(b, n)
                results = []
                for i, (tmp_A, tmp_b) in enumerate(zip(split_A, split_b)):
                    cur_device = torch.device("cuda:" + str(i))
                    tmp_A = torch.from_numpy(tmp_A).cuda(cur_device)
                    tmp_b = torch.from_numpy(tmp_b).cuda(cur_device)
                    results.append(torch.solve(tmp_b, tmp_A)[0])
                total = results[0].cpu().numpy()
                for i in range(1, len(results)):
                    total = np.append(total, results[i].cpu().numpy(), 0)
                return total
            except Exception as e:
                warnings.warn(
                    "I encountered a problem using the GPU via pytorch. Falling back to Numpy and the CPU.")
                try:
                    x = np.linalg.solve(A, b)
                except:
                    x, res, rank, s = np.linalg.lstsq(A, b, rcond=None)
                return x
        else:
            raise Exception("No valid solve method specified")
//  [python]
def _solve(self, A, b):
        if np.ndim(b) == 1: b = np.expand_dims(b, axis=1)
        if self.compute_device == "cpu":
            try:
                x = np.linalg.solve(A, b)
            except:
                x, res, rank, s = np.linalg.lstsq(A, b, rcond=None)
            return x
        elif self.compute_device == "gpu" or A.ndim < 3:
            try:
                import torch
                A = torch.from_numpy(A).cuda()
                b = torch.from_numpy(b).cuda()
                x = torch.linalg.solve(A, b)
                return x.cpu().numpy()
            except Exception as e:
                warnings.warn(
                    "I encountered a problem using the GPU via pytorch. Falling back to Numpy and the CPU.")
                try:
                    x = np.linalg.solve(A, b)
                except:
                    x, res, rank, s = np.linalg.lstsq(A, b, rcond=None)
                return x
        elif self.compute_device == "multi-gpu":
            try:
                import torch
                n = min(len(A), torch.cuda.device_count())
                split_A = np.array_split(A, n)
                split_b = np.array_split(b, n)
                results = []
                for i, (tmp_A, tmp_b) in enumerate(zip(split_A, split_b)):
                    cur_device = torch.device("cuda:" + str(i))
                    tmp_A = torch.from_numpy(tmp_A).cuda(cur_device)
                    tmp_b = torch.from_numpy(tmp_b).cuda(cur_device)
                    results.append(torch.linalg.solve(tmp_A, tmp_b)[0])
                total = results[0].cpu().numpy()
                for i in range(1, len(results)):
                    total = np.append(total, results[i].cpu().numpy(), 0)
                return total
            except Exception as e:
                warnings.warn(
                    "I encountered a problem using the GPU via pytorch. Falling back to Numpy and the CPU.")
                try:
                    x = np.linalg.solve(A, b)
                except:
                    x, res, rank, s = np.linalg.lstsq(A, b, rcond=None)
                return x
        else:
            raise Exception("No valid solve method specified")
//  [python]
def precompute_backward(self, delta):
        return torch.solve(self.I - delta*self.A, self.I)[1]
//  [python]
def precompute_backward(self, delta):
        return torch.linalg.solve(self.I - delta*self.A, self.I)[0]
//  [python]
def solve(self, x: torch.tensor, FO: Linear) -> torch.tensor:
        r"""
        Agrs:
            - :math:`x`: measurement vector
            - :math:`FO`: Linear

        Shape:
            - Input1: :math:`(bc, M)`
            - Ouput: :math:`(bc, M)`

        Example:
            >>> from spyrit.core.forwop import Linear
            >>> Hsub = np.array(np.random.random([400,32*32]))
            >>> FO = Linear(Hsub)
            >>> x = torch.tensor(np.random.random([10,400]), dtype=torch.float)
            >>> y = TIK_solve.solve(x, FO)
            >>> print(y.shape)
            torch.Size([10, 400])
        """
        A = FO.Mat() @ torch.transpose(FO.Mat(), 0, 1) + self.mu * torch.eye(FO.M)
        # Can precompute H@H.T to save time!
        A = A.view(1, FO.M, FO.M)
        # Instead of reshaping A, reshape x in the batch-final dimension
        # A = A.repeat(x.shape[0],1, 1); # Not optimal in terms of memory
        A = A.expand(x.shape[0], -1, -1)
        # Not optimal in terms of memory
        x = torch.solve(x, A)[0]
        return x
//  [python]
def solve(self, x: torch.tensor, FO: Linear) -> torch.tensor:
        r"""
        Agrs:
            - :math:`x`: measurement vector
            - :math:`FO`: Linear

        Shape:
            - Input1: :math:`(bc, M)`
            - Ouput: :math:`(bc, M)`

        Example:
            >>> from spyrit.core.forwop import Linear
            >>> Hsub = np.array(np.random.random([400,32*32]))
            >>> FO = Linear(Hsub)
            >>> x = torch.tensor(np.random.random([10,400]), dtype=torch.float)
            >>> y = TIK_solve.solve(x, FO)
            >>> print(y.shape)
            torch.Size([10, 400])
        """
        A = FO.Mat() @ torch.transpose(FO.Mat(), 0, 1) + self.mu * torch.eye(FO.M)
        # Can precompute H@H.T to save time!
        A = A.view(1, FO.M, FO.M)
        # Instead of reshaping A, reshape x in the batch-final dimension
        # A = A.repeat(x.shape[0],1, 1); # Not optimal in terms of memory
        A = A.expand(x.shape[0], -1, -1)
        # Not optimal in terms of memory
        x = torch.linalg.solve(A, x)
        return x
//  [python]
def solve(self, x: torch.tensor, var: torch.tensor, FO: Linear) -> torch.tensor:
        r"""
        Args:
            - :math:`x`: measurement vector
            - :math:`var`: measurement variance
            - :math:`FO`: Forwrad_operator

        Shape:
            - Input1: :math:`(bc, M)`
            - Input2: :math:`(bc, M)`
            - Output: :math:`(bc, M)`

        Example:
            >>> from spyrit.core.forwop import Linear
            >>> Hsub = np.array(np.random.random([400,32*32]))
            >>> FO = Linear(Hsub)
            >>> x = torch.tensor(np.random.random([10,400]), dtype=torch.float)
            >>> var = torch.tensor(np.random.random([10,400]), dtype=torch.float)
            >>> y = General_TIK_solve.solve(x, var, FO)
            >>> print(y.shape)
            torch.Size([10, 400])
        """
        A = FO.Mat() @ self.Sigma_prior @ torch.transpose(FO.Mat(), 0, 1)
        A = A.view(1, FO.M, FO.M)
        # A = A.repeat(x.shape[0],1,1);# this could be precomputed maybe
        # A += torch.diag_embed(var);
        A = A.expand(x.shape[0], -1, -1) + torch.diag_embed(var)
        x = torch.solve(x, A)[0]
        return x
//  [python]
def solve(self, x: torch.tensor, var: torch.tensor, FO: Linear) -> torch.tensor:
        r"""
        Args:
            - :math:`x`: measurement vector
            - :math:`var`: measurement variance
            - :math:`FO`: Forwrad_operator

        Shape:
            - Input1: :math:`(bc, M)`
            - Input2: :math:`(bc, M)`
            - Output: :math:`(bc, M)`

        Example:
            >>> from spyrit.core.forwop import Linear
            >>> Hsub = np.array(np.random.random([400,32*32]))
            >>> FO = Linear(Hsub)
            >>> x = torch.tensor(np.random.random([10,400]), dtype=torch.float)
            >>> var = torch.tensor(np.random.random([10,400]), dtype=torch.float)
            >>> y = General_TIK_solve.solve(x, var, FO)
            >>> print(y.shape)
            torch.Size([10, 400])
        """
        A = FO.Mat() @ self.Sigma_prior @ torch.transpose(FO.Mat(), 0, 1)
        A = A.view(1, FO.M, FO.M)
        # A = A.repeat(x.shape[0],1,1);# this could be precomputed maybe
        # A += torch.diag_embed(var);
        A = A.expand(x.shape[0], -1, -1) + torch.diag_embed(var)
        x = torch.linalg.solve(A, x)
        return x
//  [python]
def fit_plane(points, weights=None):
        """Expects points to be of shape (..., 3).
        Returns [a, b, c] such that the plane is defined as
            ax + by + c = z
        """
        D = torch.cat([points[..., :2], torch.ones_like(points[..., 2:3])], dim=-1)
        z = points[..., 2]
        if weights is not None:
            Dtrans = D.transpose(-1, -2)
        else:
            DW = D * weights
            Dtrans = DW.transpose(-1, -2)
        w = torch.solve(
            torch.matmul(Dtrans, D), torch.matmul(Dtrans, z.unsqueeze(-1))
        ).solution.squeeze(-1)
        return w
//  [python]
def fit_plane(points, weights=None):
        """Expects points to be of shape (..., 3).
        Returns [a, b, c] such that the plane is defined as
            ax + by + c = z
        """
        D = torch.cat([points[..., :2], torch.ones_like(points[..., 2:3])], dim=-1)
        z = points[..., 2]
        if weights is not None:
            Dtrans = D.transpose(-1, -2)
        else:
            DW = D * weights
            Dtrans = DW.transpose(-1, -2)
        w = torch.linalg.solve(
            torch.matmul(Dtrans, D), torch.matmul(Dtrans, z.unsqueeze(-1))
        ).squeeze(-1)
        return w
//  [python]
def inverse_mult(self, u, delta, transpose=False):
        """ Computes (I - d A)^-1 u """

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)
        _A = self.I - delta * self.A
        if transpose: _A = _A.transpose(-1, -2)

        # x = torch.linalg.solve(_A, u.unsqueeze(-1)).squeeze(-1)

        # TODO pass in a flag to toggle the two codepaths depending on how big the problem is
        xs = []
        for _A_, u_ in zip(*torch.broadcast_tensors(_A, u.unsqueeze(-1))):
            x_ = torch.solve(_A_, u_[...,:1])[0].squeeze(-1)
            xs.append(x_)
        x = torch.stack(xs, dim=0)

        return x
//  [python]
def inverse_mult(self, u, delta, transpose=False):
        """ Computes (I - d A)^-1 u """

        if isinstance(delta, torch.Tensor):
            delta = delta.unsqueeze(-1).unsqueeze(-1)
        _A = self.I - delta * self.A
        if transpose: _A = _A.transpose(-1, -2)

        # x = torch.linalg.solve(_A, u.unsqueeze(-1)).squeeze(-1)

        # TODO pass in a flag to toggle the two codepaths depending on how big the problem is
        xs = []
        for _A_, u_ in zip(*torch.broadcast_tensors(_A, u.unsqueeze(-1))):
            x_ = torch.linalg.solve(_A_, u_[...,:1]).squeeze(-1)
            xs.append(x_)
        x = torch.stack(xs, dim=0)

        return x
//  [python]
def volume_gradient_tester(self, P: torch.Tensor, q: torch.Tensor,
                               r: torch.Tensor):
        n = P.shape[0]
        V = n * torch.log(r + q.dot(torch.solve(q.unsqueeze(1), P)[0].squeeze())) - torch.log(
            torch.det(P))
        V.backward()
        P_grad_expected = P.grad.clone()
        q_grad_expected = q.grad.clone()
        r_grad_expected = r.grad.clone()
        # We call the matrix [r qᵀ] as X
        #                    [q -P]
        X = torch.empty((n + 1, n + 1), dtype=P.dtype)
        X[0, 0] = r
        X[0, 1:] = q.T
        X[1:, 0] = q
        X[1:, 1:] = -P
        X_inv = torch.linalg.inv(X)
        P_inv = torch.linalg.inv(P)
        r_grad = n * X_inv[0, 0]
        q_grad = n * X_inv[0, 1:] * 2
        P_grad = n * -X_inv[1:, 1:] - (n + 1) * P_inv
        self.assertAlmostEqual(r_grad.item(), r_grad_expected.item())
        np.testing.assert_allclose(q_grad.detach().numpy(),
                                   q_grad_expected.detach().numpy())
        np.testing.assert_allclose(P_grad.detach().numpy(),
                                   P_grad_expected.detach().numpy())
//  [python]
def volume_gradient_tester(self, P: torch.Tensor, q: torch.Tensor,
                               r: torch.Tensor):
        n = P.shape[0]
        V = n * torch.log(r + q.dot(torch.linalg.solve(P, q))) - torch.log(
            torch.linalg.det(P))
        V.backward()
        P_grad_expected = P.grad.clone()
        q_grad_expected = q.grad.clone()
        r_grad_expected = r.grad.clone()
        # We call the matrix [r qᵀ] as X
        #                    [q -P]
        X = torch.empty((n + 1, n + 1), dtype=P.dtype)
        X[0, 0] = r
        X[0, 1:] = q.T
        X[1:, 0] = q
        X[1:, 1:] = -P
        X_inv = torch.linalg.inv(X)
        P_inv = torch.linalg.inv(P)
        r_grad = n * X_inv[0, 0]
        q_grad = n * X_inv[0, 1:] * 2
        P_grad = n * -X_inv[1:, 1:] - (n + 1) * P_inv
        self.assertAlmostEqual(r_grad.item(), r_grad_expected.item())
        np.testing.assert_allclose(q_grad.detach().numpy(),
                                   q_grad_expected.detach().numpy())
        np.testing.assert_allclose(P_grad.detach().numpy(),
                                   P_grad_expected.detach().numpy())
//  [python]
def bell_core(A, B, J, j, H, P, p, q, R=None):
    """
    Core linear algebra manipulations (akin to compute a Schur's complement)
    See bell_step for more details.
    """
    if R is None:
        R = torch.zeros(p.shape[0], q.shape[0])
    K = - torch.solve(R.t() + B.t().mm(J.mm(A)),H)[0]
    k = - torch.solve( (q + B.t().mv(j)),H)[0]
    next_J = P + A.t().mm(J.mm(A)) + (R + A.t().mm(J.mm(B))).mm(K)
    next_j = p + A.t().mv(j) + A.t().mv(J.mv(B.mv(k))) + R.mv(k)
    rest = 0.5 * (q + B.t().mv(j)).dot(k)
    return next_J, next_j, rest, K, k
//  [python]
def bell_core(A, B, J, j, H, P, p, q, R=None):
    """
    Core linear algebra manipulations (akin to compute a Schur's complement)
    See bell_step for more details.
    """
    if R is None:
        R = torch.zeros(p.shape[0], q.shape[0])
    K = - torch.linalg.solve(H, R.t() + B.t().mm(J.mm(A)))
    k = - torch.linalg.solve(H, (q + B.t().mv(j)))
    next_J = P + A.t().mm(J.mm(A)) + (R + A.t().mm(J.mm(B))).mm(K)
    next_j = p + A.t().mv(j) + A.t().mv(J.mv(B.mv(k))) + R.mv(k)
    rest = 0.5 * (q + B.t().mv(j)).dot(k)
    return next_J, next_j, rest, K, k
//  [python]
def dipole_x_beta(self, input_ptp, x, log_y, z):
        # B x C x 3
        duv = torch.stack(
            [
                x[:, None] - self.local_geom[None, :, 0],
                torch.broadcast_to(
                    torch.exp(2 * log_y)[:, None], input_ptp.shape
                ),
                z[:, None] - self.local_geom[None, :, 1],
            ],
            axis=2,
        )
        # B x C x 3
        X = duv / (torch.square(duv).sum(axis=2, keepdim=True) + 0.01)
        beta = torch.solve(
            torch.einsum("bci,bcj->bij", X, X),
            torch.einsum("bck,bc->bk", X, input_ptp),
        )[1]
        return X, beta
//  [python]
def dipole_x_beta(self, input_ptp, x, log_y, z):
        # B x C x 3
        duv = torch.stack(
            [
                x[:, None] - self.local_geom[None, :, 0],
                torch.broadcast_to(
                    torch.exp(2 * log_y)[:, None], input_ptp.shape
                ),
                z[:, None] - self.local_geom[None, :, 1],
            ],
            axis=2,
        )
        # B x C x 3
        X = duv / (torch.square(duv).sum(axis=2, keepdim=True) + 0.01)
        beta = torch.linalg.solve(
            torch.einsum("bci,bcj->bij", X, X),
            torch.einsum("bck,bc->bk", X, input_ptp),
        )
        return X, beta
//  [python]
def newton(f, z_init, jac_f=None, max_iter=40, tol=1e-4):
    """Basic Newton method for finding roots in k variables.

    When `jac_f` is None, the backward hook defined in the forward pass of
    `rootfind.RootFind` needs `create_graph=True` to compute numerical gradients.
    This can be avoided if an analytical expression for the Jacobian is available,
    which can be passed along to the solver via the `solver_fwd_jac_f` argument.
    """

    def jacobian(f, z):
        return jac_f(z) if jac_f is not None else batch_jacobian(f, z)

    def g(z):
        return z - torch.solve(_check_singular_jacobian(jacobian(f, z)), f(z))[0]

    z_prev, z, n_steps, trace = z_init, g(z_init), 0, []
    trace.append(torch.linalg.norm(f(z_init)).detach())
    trace.append(torch.linalg.norm(f(z)).detach())

    while torch.linalg.norm(z_prev - z) > tol and n_steps < max_iter:
        z_prev, z = z, g(z)
        trace.append(torch.linalg.norm(f(z)).detach())
        n_steps += 1

    return {
        'result': z,
        'n_steps': n_steps,
        'trace': trace,
        'max_iter': max_iter,
        'tol': tol,
    }
//  [python]
def newton(f, z_init, jac_f=None, max_iter=40, tol=1e-4):
    """Basic Newton method for finding roots in k variables.

    When `jac_f` is None, the backward hook defined in the forward pass of
    `rootfind.RootFind` needs `create_graph=True` to compute numerical gradients.
    This can be avoided if an analytical expression for the Jacobian is available,
    which can be passed along to the solver via the `solver_fwd_jac_f` argument.
    """

    def jacobian(f, z):
        return jac_f(z) if jac_f is not None else batch_jacobian(f, z)

    def g(z):
        return z - torch.linalg.solve(_check_singular_jacobian(jacobian(f, z)), f(z))

    z_prev, z, n_steps, trace = z_init, g(z_init), 0, []
    trace.append(torch.linalg.norm(f(z_init)).detach())
    trace.append(torch.linalg.norm(f(z)).detach())

    while torch.linalg.norm(z_prev - z) > tol and n_steps < max_iter:
        z_prev, z = z, g(z)
        trace.append(torch.linalg.norm(f(z)).detach())
        n_steps += 1

    return {
        'result': z,
        'n_steps': n_steps,
        'trace': trace,
        'max_iter': max_iter,
        'tol': tol,
    }
//  [python]
def _contains(self, points, params=Points.empty()):
        translate_values = self.rotate_around(points.join(params)).reshape(-1, self.space.dim)
        rotation_matrix = self.rotation_fn(points.join(params)).reshape(-1, self.space.dim, 
                                                                        self.space.dim)
        shifted_points = points[:, list(self.space.keys())].as_tensor \
                        - translate_values
        # here apply inverse rotation -> solve: Matrix * x = shifted_points
        rotated_points = torch.solve(shifted_points.unsqueeze(-1), rotation_matrix)[0]
        shifted_points = rotated_points.squeeze(-1) + translate_values
        return self.domain._contains(Points(shifted_points, self.space), params)
//  [python]
def _contains(self, points, params=Points.empty()):
        translate_values = self.rotate_around(points.join(params)).reshape(-1, self.space.dim)
        rotation_matrix = self.rotation_fn(points.join(params)).reshape(-1, self.space.dim, 
                                                                        self.space.dim)
        shifted_points = points[:, list(self.space.keys())].as_tensor \
                        - translate_values
        # here apply inverse rotation -> solve: Matrix * x = shifted_points
        rotated_points = torch.linalg.solve(rotation_matrix, shifted_points.unsqueeze(-1))
        shifted_points = rotated_points.squeeze(-1) + translate_values
        return self.domain._contains(Points(shifted_points, self.space), params)
//  [python]
def analytic_linear(self, x, dx, G):
        Gt = torch.transpose(G.to(x), -2, -1).unsqueeze(1)
        x_masked = Gt * x
        A_est = []
        for p in range(self.num_linear):
            w_est = torch.solve(
                torch.transpose(x_masked[:, :, p, :], -2, -1) @ dx[:, :, p],
                (torch.transpose(x_masked[:, :, p, :], -2, -1) @ x_masked[:, :, p, :])
                + self.beta * torch.eye(self.num_linear).unsqueeze(0).type_as(x_masked),
            )[0]
            A_est.append(w_est)
        A_est = torch.cat(A_est, dim=2)
        # A_est = A_est * G # trying something
        return A_est
//  [python]
def analytic_linear(self, x, dx, G):
        Gt = torch.transpose(G.to(x), -2, -1).unsqueeze(1)
        x_masked = Gt * x
        A_est = []
        for p in range(self.num_linear):
            w_est = torch.linalg.solve(
                (torch.transpose(x_masked[:, :, p, :], -2, -1) @ x_masked[:, :, p, :])
                + self.beta * torch.eye(self.num_linear).unsqueeze(0).type_as(x_masked),
                torch.transpose(x_masked[:, :, p, :], -2, -1) @ dx[:, :, p],
            )
            A_est.append(w_est)
        A_est = torch.cat(A_est, dim=2)
        # A_est = A_est * G # trying something
        return A_est
//  [python]
def _apply_g(g, var_g, cov, mu, mass_rhs):  # noqa: WPS210
    _, _, nx, _ = mu.shape
    _, _, ng, _ = g.shape
    eye = torch.eye(nx, device=g.device)
    eye = rearrange(eye, "nx1 nx2 -> 1 1 nx1 nx2")
    eye_g = torch.eye(ng, device=g.device)
    eye_g = rearrange(eye_g, "ng1 ng2 -> 1 1 ng1 ng2")
    gtr = g.transpose(3, 2)
    small_a = eye_g * var_g + (g.matmul(cov * gtr))
    rinv1 = torch.solve(g.matmul(cov * eye), small_a)[1]
    new_cov = cov * (eye - gtr.matmul(rinv1))

    b = mass_rhs.unsqueeze(-1).unsqueeze(-1)
    rinv2 = torch.solve(g.matmul(mu) - b, small_a)[1]
    new_mu = mu - cov * gtr.matmul(rinv2)
    return new_mu, new_cov
//  [python]
def _apply_g(g, var_g, cov, mu, mass_rhs):  # noqa: WPS210
    _, _, nx, _ = mu.shape
    _, _, ng, _ = g.shape
    eye = torch.eye(nx, device=g.device)
    eye = rearrange(eye, "nx1 nx2 -> 1 1 nx1 nx2")
    eye_g = torch.eye(ng, device=g.device)
    eye_g = rearrange(eye_g, "ng1 ng2 -> 1 1 ng1 ng2")
    gtr = g.transpose(3, 2)
    small_a = eye_g * var_g + (g.matmul(cov * gtr))
    rinv1 = torch.linalg.solve(small_a, g.matmul(cov * eye))
    new_cov = cov * (eye - gtr.matmul(rinv1))

    b = mass_rhs.unsqueeze(-1).unsqueeze(-1)
    rinv2 = torch.linalg.solve(small_a, g.matmul(mu) - b)
    new_mu = mu - cov * gtr.matmul(rinv2)
    return new_mu, new_cov
//  [python]
def get_solution(self, area):
        '''
        Solve a truss problem

        Args:
            area : torch.Tensor, size = (N, n_bar), dtype = float64
                area of bars

        Returns:
            u : torch.Tensor, size=(n_node, n_dim), dtype=float64
                displacement of nodes
            tension : torch.Tensor, size=(n_bar, ), dtype=float64
                tension of bars
            volume : torch.Tensor, size=(n_bar, ), dtype=float64
                volume of bars
        '''

        N = area.shape[0]
        sys = torch.zeros((N, self.n_dim * self.n_node, self.n_dim * self.n_node), dtype=torch.float, device=self.device)  # system matrix
        volume = self.length * area  # volume of bars
        G = self.modulus * area / self.length  # a temporary coefficient to be used later

        local_matrix = G[:, :, None, None] * self.matrix  # local stiffness matrix
        for i in range(self.n_bar):
            sys[:, self.idx[i].unsqueeze(0), self.idx[i].unsqueeze(-1)] += local_matrix[:, i]  # global stiffness matrix

        u = torch.zeros((N, self.n_node, self.n_dim), dtype=torch.float, device=self.device)  # displacement

        u_free = torch.solve(self.load[self.free_mask].unsqueeze(-1), sys[:, self.free_list][:, :, self.free_list])[0]
        u[:, self.free_mask] = u_free.squeeze()

        TG = torch.einsum('ij,jk->ijk', G, self.T)  # projection coefficient
        tension = torch.sum(TG * (u[:, self.con[:, 1], :] - u[:, self.con[:, 0], :]), dim=-1)
        return u, tension, volume