jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def precision_weighted_inner_prod(self, b, reduce_dim=True):
        assert b.shape[-1] == 1
        prod = (torch.linalg.solve(self.scale_tril, b)**2).sum(-2)
        return prod.squeeze(-1) if reduce_dim else prod
//  [python]
def estimate_hscic(X, Y, Z, ridge_lambda=1e-2, use_median=False, normalize_kernel=False, sigma2=None):
    '''X ind. Y | Z '''
    # (1) action regularization version : X = imitator action
    # (2) regularized representation version : X = varphi(Obs)    

    if sigma2 is None:
        if use_median:
            sigma2_ = None
        else:
            sigma2_ = 1.
    else:
        sigma2_ = sigma2

    Kx = gaussian_kernel(X, sigma2=sigma2_, normalized=normalize_kernel)
    Ky = gaussian_kernel(Y, sigma2=sigma2_, normalized=normalize_kernel)
    Kz = gaussian_kernel(Z, sigma2=sigma2_, normalized=normalize_kernel)
    
    n = Kz.shape[0]    
    
    WtKzz = torch.solve(Kz, Kz + ridge_lambda * n * torch.eye(n).to(Kz.device))[0]     
    term_1 = (WtKzz * ((Kx * Ky) @ WtKzz)).sum()    # tr(WtKzz.T @ (Kx * Ky) @ WtKzz)
    WkKxWk = WtKzz * (Kx @ WtKzz)
    KyWk = Ky @ WtKzz
    term_2 = (WkKxWk * KyWk).sum()        
    term_3 = (WkKxWk.sum(dim=0) * (WtKzz * KyWk).sum(dim=0)).sum()    
    result = (term_1 - 2 * term_2 + term_3) / n    
    
    # W = (Kz + ridge_lambda * n * torch.eye(n)).inverse()
    # A1 = Kz.T @ W @ ( Kx * Ky ) @ W.T @ Kz
    # A2 = Kz.T @ W @ ( (Kx@W.T@Kz) * (Ky@W.T@Kz) )
    # A3 = (Kz.T @ W @ Kx @ W.T @ Kz) * (Kz.T @ W @ Ky @ W.T @ Kz)
    # result2 = (A1-2*A2+A3).trace()/n
    
    # W = (Kz + ridge_lambda  * n * torch.eye(n)).inverse()
    # total_term = 0.
    # for i in range(n):
    #     kz = Kz[i]  #(n,1)
    #     t1 = kz.T @ W @ (Kx * Ky) @ W.T @ kz
    #     t2 = kz.T @ W @ ( (Kx@W.T@kz) * (Ky@W.T@kz) )
    #     t3 = (kz.T @ W @ Kx @ W.T @ kz) * (kz.T @ W @ Ky @ W.T @ kz)
    #     total_term += t1 - 2 * t2 + t3
    # result3 = total_term / n
    
    return result
//  [python]
def estimate_hscic(X, Y, Z, ridge_lambda=1e-2, use_median=False, normalize_kernel=False, sigma2=None):
    '''X ind. Y | Z '''
    # (1) action regularization version : X = imitator action
    # (2) regularized representation version : X = varphi(Obs)    

    if sigma2 is None:
        if use_median:
            sigma2_ = None
        else:
            sigma2_ = 1.
    else:
        sigma2_ = sigma2

    Kx = gaussian_kernel(X, sigma2=sigma2_, normalized=normalize_kernel)
    Ky = gaussian_kernel(Y, sigma2=sigma2_, normalized=normalize_kernel)
    Kz = gaussian_kernel(Z, sigma2=sigma2_, normalized=normalize_kernel)
    
    n = Kz.shape[0]    
    
    WtKzz = torch.linalg.solve(Kz + ridge_lambda * n * torch.eye(n).to(Kz.device), Kz)     
    term_1 = (WtKzz * ((Kx * Ky) @ WtKzz)).sum()    # tr(WtKzz.T @ (Kx * Ky) @ WtKzz)
    WkKxWk = WtKzz * (Kx @ WtKzz)
    KyWk = Ky @ WtKzz
    term_2 = (WkKxWk * KyWk).sum()        
    term_3 = (WkKxWk.sum(dim=0) * (WtKzz * KyWk).sum(dim=0)).sum()    
    result = (term_1 - 2 * term_2 + term_3) / n    
    
    # W = (Kz + ridge_lambda * n * torch.eye(n)).inverse()
    # A1 = Kz.T @ W @ ( Kx * Ky ) @ W.T @ Kz
    # A2 = Kz.T @ W @ ( (Kx@W.T@Kz) * (Ky@W.T@Kz) )
    # A3 = (Kz.T @ W @ Kx @ W.T @ Kz) * (Kz.T @ W @ Ky @ W.T @ Kz)
    # result2 = (A1-2*A2+A3).trace()/n
    
    # W = (Kz + ridge_lambda  * n * torch.eye(n)).inverse()
    # total_term = 0.
    # for i in range(n):
    #     kz = Kz[i]  #(n,1)
    #     t1 = kz.T @ W @ (Kx * Ky) @ W.T @ kz
    #     t2 = kz.T @ W @ ( (Kx@W.T@kz) * (Ky@W.T@kz) )
    #     t3 = (kz.T @ W @ Kx @ W.T @ kz) * (kz.T @ W @ Ky @ W.T @ kz)
    #     total_term += t1 - 2 * t2 + t3
    # result3 = total_term / n
    
    return result
//  [python]
def geodesic(self, x, y, t):
        """Computes the Riemannian geodesic A exp(t log(A^{-1}B)).
        x: (..., D)
        y: (..., D)
        t: (...)
        """
        if self.Riem_geodesic:
            A, B = self.devectorize(x), self.devectorize(y)
            dtype = A.dtype
            A, B = A.double(), B.double()
            Ainv_B = torch.solve(B, A)[1]
            U = t[..., None, None] * matrix_logarithm(Ainv_B)
            G_t = torch.matmul(A, torch.matrix_exp(U))
            return self.vectorize(G_t).to(dtype)
        else:
            return x + t[..., None] * (y - x)
//  [python]
def geodesic(self, x, y, t):
        """Computes the Riemannian geodesic A exp(t log(A^{-1}B)).
        x: (..., D)
        y: (..., D)
        t: (...)
        """
        if self.Riem_geodesic:
            A, B = self.devectorize(x), self.devectorize(y)
            dtype = A.dtype
            A, B = A.double(), B.double()
            Ainv_B = torch.linalg.solve(A, B)
            U = t[..., None, None] * matrix_logarithm(Ainv_B)
            G_t = torch.matmul(A, torch.matrix_exp(U))
            return self.vectorize(G_t).to(dtype)
        else:
            return x + t[..., None] * (y - x)
//  [python]
def inner(
        self, x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False
    ) -> torch.Tensor:
        if self.Riem_norm:
            P = self.devectorize(x)
            A = self.devectorize(u)
            B = self.devectorize(v)
            dtype = P.dtype
            P, A, B = P.double(), A.double(), B.double()
            Pinv_A = torch.solve(A, P)[0]
            Pinv_B = torch.solve(B, P)[0]
            return torch.diagonal(torch.matmul(Pinv_A, Pinv_B), dim1=-2, dim2=-1).to(dtype)
        else:
            return torch.sum(u * v, dim=-1, keepdim=keepdim)
//  [python]
def inner(
        self, x: torch.Tensor, u: torch.Tensor, v: torch.Tensor = None, *, keepdim=False
    ) -> torch.Tensor:
        if self.Riem_norm:
            P = self.devectorize(x)
            A = self.devectorize(u)
            B = self.devectorize(v)
            dtype = P.dtype
            P, A, B = P.double(), A.double(), B.double()
            Pinv_A = torch.linalg.solve(P, A)
            Pinv_B = torch.linalg.solve(P, B)
            return torch.diagonal(torch.matmul(Pinv_A, Pinv_B), dim1=-2, dim2=-1).to(dtype)
        else:
            return torch.sum(u * v, dim=-1, keepdim=keepdim)
//  [python]
def anderson(f, x0, m=5, lam=1e-4, max_iter=50, tol=1e-5, beta = 1.0):
    """ Anderson acceleration for fixed point iteration. """
    bsz, d, H, W = x0.shape
    X = torch.zeros(bsz, m, d*H*W, dtype=x0.dtype, device=x0.device)
    F = torch.zeros(bsz, m, d*H*W, dtype=x0.dtype, device=x0.device)
    X[:,0], F[:,0] = x0.view(bsz, -1), f(x0).view(bsz, -1)
    X[:,1], F[:,1] = F[:,0], f(F[:,0].view_as(x0)).view(bsz, -1)
    
    H = torch.zeros(bsz, m+1, m+1, dtype=x0.dtype, device=x0.device)
    H[:,0,1:] = H[:,1:,0] = 1
    y = torch.zeros(bsz, m+1, 1, dtype=x0.dtype, device=x0.device)
    y[:,0] = 1
    
    res = []
    for k in range(2, max_iter):
        n = min(k, m)
        G = F[:,:n]-X[:,:n]
        H[:,1:n+1,1:n+1] = torch.bmm(G,G.transpose(1,2)) + lam*torch.eye(n, dtype=x0.dtype,device=x0.device)[None]
  
        alpha = torch.solve(H[:, :n + 1, :n + 1], y[:, :n + 1])[0][:,1:n+1,0]
   
        X[:,k%m] = beta * (alpha[:,None] @ F[:,:n])[:,0] + (1-beta)*(alpha[:,None] @ X[:,:n])[:,0]
        F[:,k%m] = f(X[:,k%m].view_as(x0)).view(bsz, -1)
        res.append((F[:,k%m] - X[:,k%m]).norm().item()/(1e-5 + F[:,k%m].norm().item()))
        if (res[-1] < tol):
            break

    return X[:,k%m].view_as(x0), k
//  [python]
def anderson(f, x0, m=5, lam=1e-4, max_iter=50, tol=1e-5, beta = 1.0):
    """ Anderson acceleration for fixed point iteration. """
    bsz, d, H, W = x0.shape
    X = torch.zeros(bsz, m, d*H*W, dtype=x0.dtype, device=x0.device)
    F = torch.zeros(bsz, m, d*H*W, dtype=x0.dtype, device=x0.device)
    X[:,0], F[:,0] = x0.view(bsz, -1), f(x0).view(bsz, -1)
    X[:,1], F[:,1] = F[:,0], f(F[:,0].view_as(x0)).view(bsz, -1)
    
    H = torch.zeros(bsz, m+1, m+1, dtype=x0.dtype, device=x0.device)
    H[:,0,1:] = H[:,1:,0] = 1
    y = torch.zeros(bsz, m+1, 1, dtype=x0.dtype, device=x0.device)
    y[:,0] = 1
    
    res = []
    for k in range(2, max_iter):
        n = min(k, m)
        G = F[:,:n]-X[:,:n]
        H[:,1:n+1,1:n+1] = torch.bmm(G,G.transpose(1,2)) + lam*torch.eye(n, dtype=x0.dtype,device=x0.device)[None]
  
        alpha = torch.linalg.solve(H[:, :n + 1, :n + 1], y[:, :n + 1])[:,1:n+1,0]
   
        X[:,k%m] = beta * (alpha[:,None] @ F[:,:n])[:,0] + (1-beta)*(alpha[:,None] @ X[:,:n])[:,0]
        F[:,k%m] = f(X[:,k%m].view_as(x0)).view(bsz, -1)
        res.append((F[:,k%m] - X[:,k%m]).norm().item()/(1e-5 + F[:,k%m].norm().item()))
        if (res[-1] < tol):
            break

    return X[:,k%m].view_as(x0), k
//  [python]
def expmap(self, x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
        A, U = self.devectorize(x), self.devectorize(u)
        dtype = A.dtype
        A, U = A.double(), U.double()

        A_sqrt = sqrtmh(A) + 1e-6 * torch.eye(A.shape[-1]).to(A)

        B = torch.solve(U, A_sqrt)[1]
        B = torch.solve(B, A_sqrt)[1]
        B = torch.linalg.matrix_exp(B)
        B = A_sqrt @ B @ A_sqrt

        b = self.vectorize(B).to(dtype)
        return b
//  [python]
def expmap(self, x: torch.Tensor, u: torch.Tensor) -> torch.Tensor:
        A, U = self.devectorize(x), self.devectorize(u)
        dtype = A.dtype
        A, U = A.double(), U.double()

        A_sqrt = sqrtmh(A) + 1e-6 * torch.eye(A.shape[-1]).to(A)

        B = torch.linalg.solve(A_sqrt, torch.linalg.solve(A_sqrt, U), left=False)
        B = torch.linalg.matrix_exp(B)
        B = A_sqrt @ B @ A_sqrt

        b = self.vectorize(B).to(dtype)
        return b
//  [python]
def reconstruct_atoms(atomsobject: AtomsBatch,
                      mol_idx: list[np.array],
                      centre:int =None):
    '''
    Function to shift atoms when we create non-periodic system from periodic.  
    inputs:
    atomsobject: Atomsbatch object from NFF
    mol_idx: list of array of atom indices in molecules or atoms you want to keep together when changing to non-periodic
    system
    centre: by default the atoms in a molecule or set of close atoms are shifted so as to get them close to the centre which 
    is by default the first atom index in the array. For reconstructing molecules this is fine. However, for attribution,
    we may have to shift a whole molecule to come closer to the atoms with high attribution. In that case, we manually assign 
    the atom index. 
    '''
    
    sys_xyz = torch.Tensor(atomsobject.get_positions(wrap=True))
    box_len = torch.Tensor(atomsobject.get_cell_lengths_and_angles()[:3])
    
    for idx in mol_idx:
        mol_xyz = sys_xyz[idx]
        if any(atomsobject.pbc):
            center = mol_xyz.shape[0]//2
            if centre!=None:
                center=centre # changes the central atom to atom in focus
            intra_dmat = (mol_xyz[None, :,...] - mol_xyz[:, None, ...])[center]
            if np.count_nonzero(atomsobject.cell.T-np.diag(np.diagonal(atomsobject.cell.T)))!=0:
                M,N=intra_dmat.shape[0],intra_dmat.shape[1]
                f=torch.solve(intra_dmat.view(-1,3).T,torch.Tensor(atomsobject.cell.T))[1]
                g=f-torch.floor(f+0.5)
                intra_dmat=torch.matmul(g,torch.Tensor(atomsobject.cell))
                intra_dmat=intra_dmat.view(M,3)
                offsets=-torch.floor(f+0.5).view(M,3)
                traj_unwrap = mol_xyz+torch.matmul(offsets,torch.Tensor(atomsobject.cell))
            else:
                sub = (intra_dmat > 0.5 * box_len).to(torch.float) * box_len
                add = (intra_dmat <= -0.5 * box_len).to(torch.float) * box_len
                shift=torch.round(torch.divide(intra_dmat,box_len))
                offsets=-shift
                traj_unwrap = mol_xyz+offsets*box_len
        else:
            traj_unwrap = mol_xyz
        #traj_unwrap=mol_xyz+add-sub
        sys_xyz[idx] = traj_unwrap

    new_pos = sys_xyz.numpy()

    return new_pos
//  [python]
def reconstruct_atoms(atomsobject: AtomsBatch,
                      mol_idx: list[np.array],
                      centre:int =None):
    '''
    Function to shift atoms when we create non-periodic system from periodic.  
    inputs:
    atomsobject: Atomsbatch object from NFF
    mol_idx: list of array of atom indices in molecules or atoms you want to keep together when changing to non-periodic
    system
    centre: by default the atoms in a molecule or set of close atoms are shifted so as to get them close to the centre which 
    is by default the first atom index in the array. For reconstructing molecules this is fine. However, for attribution,
    we may have to shift a whole molecule to come closer to the atoms with high attribution. In that case, we manually assign 
    the atom index. 
    '''
    
    sys_xyz = torch.Tensor(atomsobject.get_positions(wrap=True))
    box_len = torch.Tensor(atomsobject.get_cell_lengths_and_angles()[:3])
    
    for idx in mol_idx:
        mol_xyz = sys_xyz[idx]
        if any(atomsobject.pbc):
            center = mol_xyz.shape[0]//2
            if centre!=None:
                center=centre # changes the central atom to atom in focus
            intra_dmat = (mol_xyz[None, :,...] - mol_xyz[:, None, ...])[center]
            if np.count_nonzero(atomsobject.cell.T-np.diag(np.diagonal(atomsobject.cell.T)))!=0:
                M,N=intra_dmat.shape[0],intra_dmat.shape[1]
                f=torch.linalg.solve(torch.Tensor(atomsobject.cell.T),(intra_dmat.view(-1,3).T)).T
                g=f-torch.floor(f+0.5)
                intra_dmat=torch.matmul(g,torch.Tensor(atomsobject.cell))
                intra_dmat=intra_dmat.view(M,3)
                offsets=-torch.floor(f+0.5).view(M,3)
                traj_unwrap = mol_xyz+torch.matmul(offsets,torch.Tensor(atomsobject.cell))
            else:
                sub = (intra_dmat > 0.5 * box_len).to(torch.float) * box_len
                add = (intra_dmat <= -0.5 * box_len).to(torch.float) * box_len
                shift=torch.round(torch.divide(intra_dmat,box_len))
                offsets=-shift
                traj_unwrap = mol_xyz+offsets*box_len
        else:
            traj_unwrap = mol_xyz
        #traj_unwrap=mol_xyz+add-sub
        sys_xyz[idx] = traj_unwrap

    new_pos = sys_xyz.numpy()

    return new_pos
//  [python]
def test_simple_vmap():
    def torch_fn(A, x):
        # return torch.linalg.solve_triangular(torch.linalg.cholesky(A, upper=True), x, upper=True)
        return torch.solve(x, A)[0]

    def expected_fn(A, x):
        return cho_solve(cho_factor(A), x)

    device_list = ["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"]
    dtype_list = [jnp.float32, jnp.float64]
    for use_torch_vmap in [True, False]:
        for device in device_list:
            for dtype in dtype_list:
                A = jax_randn((10, 5, 5), dtype=dtype, device=device)
                A = jnp.swapaxes(A, -1, -2) @ A + 1e-2 * jnp.tile(
                    jnp.diag(1.0 + 0.0 * jax_randn((A.shape[-1],), dtype=dtype, device=device)),
                    (A.shape[0], 1, 1),
                )
                x = jax_randn((5, 7), dtype=dtype, device=device)
                jax_fn = torch2jax_with_vjp(
                    torch_fn, A, x, output_shapes=x, depth=2, use_torch_vmap=use_torch_vmap
                )
                sol = jax.vmap(jax_fn, in_axes=(0, None))(A, x)
                sol_expected = jax.vmap(expected_fn, in_axes=(0, None))(A, x)
                err = jnp.linalg.norm(sol - sol_expected) / jnp.linalg.norm(sol_expected)
                assert err < 1e-3
//  [python]
def test_simple_vmap():
    def torch_fn(A, x):
        # return torch.linalg.solve_triangular(torch.linalg.cholesky(A, upper=True), x, upper=True)
        return torch.linalg.solve(A, x)

    def expected_fn(A, x):
        return cho_solve(cho_factor(A), x)

    device_list = ["cuda", "cpu"] if torch.cuda.is_available() else ["cpu"]
    dtype_list = [jnp.float32, jnp.float64]
    for use_torch_vmap in [True, False]:
        for device in device_list:
            for dtype in dtype_list:
                A = jax_randn((10, 5, 5), dtype=dtype, device=device)
                A = jnp.swapaxes(A, -1, -2) @ A + 1e-2 * jnp.tile(
                    jnp.diag(1.0 + 0.0 * jax_randn((A.shape[-1],), dtype=dtype, device=device)),
                    (A.shape[0], 1, 1),
                )
                x = jax_randn((5, 7), dtype=dtype, device=device)
                jax_fn = torch2jax_with_vjp(
                    torch_fn, A, x, output_shapes=x, depth=2, use_torch_vmap=use_torch_vmap
                )
                sol = jax.vmap(jax_fn, in_axes=(0, None))(A, x)
                sol_expected = jax.vmap(expected_fn, in_axes=(0, None))(A, x)
                err = jnp.linalg.norm(sol - sol_expected) / jnp.linalg.norm(sol_expected)
                assert err < 1e-3
//  [python]
def solve(A,B):
    return torch.solve(B,A)[0]
//  [python]
def solve(A,B):
    return torch.linalg.solve(A,B)
//  [python]
def transp(self, x: torch.Tensor, y: torch.Tensor, v: torch.Tensor):
        A, B, V = self.devectorize(x), self.devectorize(y), self.devectorize(v)
        dtype = A.dtype
        A, B, V = A.double(), B.double(), V.double()
        A_sqrt = sqrtmh(A)
        B_sqrt = sqrtmh(B)
        E = torch.solve(A_sqrt, B_sqrt)[0]
        U = E @ V @ E.mT
        return self.vectorize(U).to(dtype)
//  [python]
def transp(self, x: torch.Tensor, y: torch.Tensor, v: torch.Tensor):
        A, B, V = self.devectorize(x), self.devectorize(y), self.devectorize(v)
        dtype = A.dtype
        A, B, V = A.double(), B.double(), V.double()
        A_sqrt = sqrtmh(A)
        B_sqrt = sqrtmh(B)
        E = torch.linalg.solve(B_sqrt, A_sqrt, left=False)
        U = E @ V @ E.mT
        return self.vectorize(U).to(dtype)
//  [python]
def solve(A, B, C):
    C = torch.solve(B, A)[0]
    return C
//  [python]
def solve(A, B, C):
    C = torch.linalg.solve(A, B)
    return C
//  [python]
def logmap(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        A, B = self.devectorize(x), self.devectorize(y)
        dtype = A.dtype
        A, B = A.double(), B.double()

        A_sqrt = sqrtmh(A) + 1e-6 * torch.eye(A.shape[-1]).to(A)

        U = torch.solve(B, A_sqrt)[1]
        U = torch.solve(U, A_sqrt)[1]
        U = matrix_logarithm(U)
        U = A_sqrt @ U @ A_sqrt

        return self.vectorize(U).to(dtype)
//  [python]
def logmap(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        A, B = self.devectorize(x), self.devectorize(y)
        dtype = A.dtype
        A, B = A.double(), B.double()

        A_sqrt = sqrtmh(A) + 1e-6 * torch.eye(A.shape[-1]).to(A)

        U = torch.linalg.solve(A_sqrt, torch.linalg.solve(A_sqrt, B), left=False)
        U = matrix_logarithm(U)
        U = A_sqrt @ U @ A_sqrt

        return self.vectorize(U).to(dtype)
//  [python]
def score_parameters(self) -> List[Tensor]:
        """
        :return: List of Tensors the same shapes as the given Parameters where
            each Parameter's elements are scored based on the blockwise OBS
        """
        scores = [None] * len(self._params)
        block_finv_w = [None] * len(self._params)

        if self._is_main_proc:
            for i, finv in enumerate(self._finvs):
                if self._mask_type == "unstructured":
                    scores[i] = (
                        (self._params[i].data.reshape(-1) ** 2).to(self._devices[i])
                        / (2.0 * finv.diag() + self._eps)
                    ).reshape(self._params[i].shape)
                else:  # self._mask_type == "block4":
                    block_w = (
                        self._params[i].data.reshape(-1, 4).to(finv.dev)
                    )  # (d/Q, Q)
                    block_finv = (
                        torch.cat(
                            [
                                finv.f_inv[:, i : i + 4, i : i + 4]
                                for i in range(0, finv.B, 4)
                            ],
                            dim=1,
                        )
                        .reshape((finv.d // finv.B, finv.B // 4, 4, 4))
                        .reshape((finv.d // 4, 4, 4))
                    )  # (Q, d/Q, Q) -> (d/Q, Q, Q)
                    block_finv_w[i] = torch.solve(
                        block_w,
                        block_finv,
                    )  # (d/Q, Q)
                    score = 0.5 * torch.einsum(
                        "bi,bi->b", block_w, block_finv_w[i]
                    )  # d/Q
                    scores[i] = (
                        score.unsqueeze(1)
                        .expand(-1, 4)
                        .reshape(self._params[i].data.shape)
                    )

            # make sure pruned ones will stay pruned
            for i, score in enumerate(scores):
                score[self._masks[i] == 0] = float("-inf")

        self._broadcast_list_from_main(scores)
        if self._mask_type == "block4":
            self._broadcast_list_from_main(block_finv_w)
            self._block_finv_w = block_finv_w  # cache for OBS weight update

        return scores
//  [python]
def score_parameters(self) -> List[Tensor]:
        """
        :return: List of Tensors the same shapes as the given Parameters where
            each Parameter's elements are scored based on the blockwise OBS
        """
        scores = [None] * len(self._params)
        block_finv_w = [None] * len(self._params)

        if self._is_main_proc:
            for i, finv in enumerate(self._finvs):
                if self._mask_type == "unstructured":
                    scores[i] = (
                        (self._params[i].data.reshape(-1) ** 2).to(self._devices[i])
                        / (2.0 * finv.diag() + self._eps)
                    ).reshape(self._params[i].shape)
                else:  # self._mask_type == "block4":
                    block_w = (
                        self._params[i].data.reshape(-1, 4).to(finv.dev)
                    )  # (d/Q, Q)
                    block_finv = (
                        torch.cat(
                            [
                                finv.f_inv[:, i : i + 4, i : i + 4]
                                for i in range(0, finv.B, 4)
                            ],
                            dim=1,
                        )
                        .reshape((finv.d // finv.B, finv.B // 4, 4, 4))
                        .reshape((finv.d // 4, 4, 4))
                    )  # (Q, d/Q, Q) -> (d/Q, Q, Q)
                    block_finv_w[i] = torch.linalg.solve(
                        block_finv,
                        block_w,
                    )  # (d/Q, Q)
                    score = 0.5 * torch.einsum(
                        "bi,bi->b", block_w, block_finv_w[i]
                    )  # d/Q
                    scores[i] = (
                        score.unsqueeze(1)
                        .expand(-1, 4)
                        .reshape(self._params[i].data.shape)
                    )

            # make sure pruned ones will stay pruned
            for i, score in enumerate(scores):
                score[self._masks[i] == 0] = float("-inf")

        self._broadcast_list_from_main(scores)
        if self._mask_type == "block4":
            self._broadcast_list_from_main(block_finv_w)
            self._block_finv_w = block_finv_w  # cache for OBS weight update

        return scores
//  [python]
def forward_krr(self, x, y, t):
        k = t[0, 0]
        assert torch.all(t == k)

        # similarity between sleep and wake data
        G = self.kernels[k](self.xs[k], x)
        if len(self.Ks) > k:
            K = self.Ks[k]
        else:
            K = self.compute_kernel(self.xs[k], k)

        # this computes G.t() @ K^{-1}
        GKinv = torch.solve(G, K).solution.t()

        return GKinv @ self.outs[k]
//  [python]
def forward_krr(self, x, y, t):
        k = t[0, 0]
        assert torch.all(t == k)

        # similarity between sleep and wake data
        G = self.kernels[k](self.xs[k], x)
        if len(self.Ks) > k:
            K = self.Ks[k]
        else:
            K = self.compute_kernel(self.xs[k], k)

        # this computes G.t() @ K^{-1}
        GKinv = torch.linalg.solve(K, G).t()

        return GKinv @ self.outs[k]
//  [python]
def train(self, X, Q, E=None, F=None, cells=None, inv_cells=None, print_info=True, cpu_solve=False, ntiles=1, use_specialized_matmul=False):
        
        ZTZ, ZtrainY = self.build_Z_components(X, Q, E, F, cells, inv_cells, print_info, cpu_solve, ntiles, use_specialized_matmul=use_specialized_matmul)
        
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        
        for i in range(self.nfeatures()):
            ZTZ[i, i] += self.llambda
        
        start.record()
        
        if (cpu_solve):
            self.alpha = torch.solve(ZTZ.cpu(), ZtrainY.cpu()).solution[:, 0].cuda()
        else:
            self.alpha = torch.solve(ZTZ, ZtrainY).solution[:, 0]
            
        end.record()
        torch.cuda.synchronize()
        
        if (print_info):
            print("coefficients time: ", start.elapsed_time(end), "ms")
        
        del ZtrainY
        del ZTZ
        
        self.is_trained = True
        
        torch.cuda.empty_cache()
//  [python]
def train(self, X, Q, E=None, F=None, cells=None, inv_cells=None, print_info=True, cpu_solve=False, ntiles=1, use_specialized_matmul=False):
        
        ZTZ, ZtrainY = self.build_Z_components(X, Q, E, F, cells, inv_cells, print_info, cpu_solve, ntiles, use_specialized_matmul=use_specialized_matmul)
        
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        
        for i in range(self.nfeatures()):
            ZTZ[i, i] += self.llambda
        
        start.record()
        
        if (cpu_solve):
            self.alpha = torch.linalg.solve(ZTZ, ZtrainY.cpu())[:, 0].cuda()
        else:
            self.alpha = torch.linalg.solve(ZTZ, ZtrainY)[:, 0]
            
        end.record()
        torch.cuda.synchronize()
        
        if (print_info):
            print("coefficients time: ", start.elapsed_time(end), "ms")
        
        del ZtrainY
        del ZTZ
        
        self.is_trained = True
        
        torch.cuda.empty_cache()
//  [python]
def rbf_interpolation(
    rotations: Tensor,
    grid: Tensor,
    signal: Tensor,
    *,
    dist_fn: Tensor = geodesic_distance,
    rbf: Callable[[Tensor], Tensor] = _rbf_gauss2,
    width=2,
) -> Tensor:
    """
    Performs radial basis function interpolation of rotations.

    Arguments:
        Arguments:
        - rotations: Tensor of shape `(N, K, 4)` containing quaternions.
        - grid: Tensor of shape `(N, L, 4)` containing quaternions
                defining a grid on SO3.
        - signal: Tensor of shape `(N, L, S)` containing the signal
                  corresponding to `grid`.
        - dist_fn: Distance function used to calculate distance between
                   rotations.
        - rbf: Radial basis function, defaults to gaussian rbf.
        - rbf_args: Tuple containing parameter arguments for `rbf`.

    Returns:
        - Tensor of shape '(N, K, S)'.
    """
    m = rbf(dist_fn(grid[..., None, :], grid[..., None, :, :]), width)

    coeffs = torch.solve(signal, m)[0].transpose(-1, -2)
    p = rbf(dist_fn(rotations[..., None, :], grid[..., None, :, :]), width)

    return p @ coeffs.transpose(-1, -2)
//  [python]
def rbf_interpolation(
    rotations: Tensor,
    grid: Tensor,
    signal: Tensor,
    *,
    dist_fn: Tensor = geodesic_distance,
    rbf: Callable[[Tensor], Tensor] = _rbf_gauss2,
    width=2,
) -> Tensor:
    """
    Performs radial basis function interpolation of rotations.

    Arguments:
        Arguments:
        - rotations: Tensor of shape `(N, K, 4)` containing quaternions.
        - grid: Tensor of shape `(N, L, 4)` containing quaternions
                defining a grid on SO3.
        - signal: Tensor of shape `(N, L, S)` containing the signal
                  corresponding to `grid`.
        - dist_fn: Distance function used to calculate distance between
                   rotations.
        - rbf: Radial basis function, defaults to gaussian rbf.
        - rbf_args: Tuple containing parameter arguments for `rbf`.

    Returns:
        - Tensor of shape '(N, K, S)'.
    """
    m = rbf(dist_fn(grid[..., None, :], grid[..., None, :, :]), width)

    coeffs = torch.linalg.solve(m, signal).transpose(-1, -2)
    p = rbf(dist_fn(rotations[..., None, :], grid[..., None, :, :]), width)

    return p @ coeffs.transpose(-1, -2)
//  [python]
def resolvent(A, B, C, z):
    """
    A: (... N N)
    B: (... N)
    C: (... N)
    z: (... L)
    returns: (... L) represents C (z - A)^{-1} B
    """
    if A.is_complex() or B.is_complex() or C.is_complex() or z.is_complex():
        dtype = torch.cfloat
    else:
        dtype = torch.float
    N = A.shape[-1]
    I = torch.eye(N).to(device=A.device, dtype=dtype)
    A_ = I * z[...,None,None] - A.unsqueeze(-3)
    B = B.to(dtype)[..., None, :, None] # (... L N 1)
    r = torch.solve(B, A_)[0].squeeze(-1) # (... L N)
    r = torch.sum(r*C.to(dtype).unsqueeze(-2), dim=-1)
    return r
//  [python]
def resolvent(A, B, C, z):
    """
    A: (... N N)
    B: (... N)
    C: (... N)
    z: (... L)
    returns: (... L) represents C (z - A)^{-1} B
    """
    if A.is_complex() or B.is_complex() or C.is_complex() or z.is_complex():
        dtype = torch.cfloat
    else:
        dtype = torch.float
    N = A.shape[-1]
    I = torch.eye(N).to(device=A.device, dtype=dtype)
    A_ = I * z[...,None,None] - A.unsqueeze(-3)
    B = B.to(dtype)[..., None, :, None] # (... L N 1)
    r = torch.linalg.solve(A_, B).squeeze(-1) # (... L N)
    r = torch.sum(r*C.to(dtype).unsqueeze(-2), dim=-1)
    return r
//  [python]
def _depr_rbf_interpolation(
    rotations: Tensor,
    grid: Tensor,
    signal: Tensor,
    *,
    dist_fn: Tensor = geodesic_distance,
    rbf: Callable[[Tensor], Tensor] = _rbf_gauss2,
    width=2,
) -> Tensor:
    """
    Performs radial basis function interpolation of rotations.

    Arguments:
        Arguments:
        - rotations: Tensor of shape `(N, K, 4)` containing quaternions.
        - grid: Tensor of shape `(N, L, 4)` containing quaternions
                defining a grid on SO3.
        - signal: Tensor of shape `(N, L, S)` containing the signal
                  corresponding to `grid`.
        - dist_fn: Distance function used to calculate distance between
                   rotations.
        - rbf: Radial basis function, defaults to gaussian rbf.
        - rbf_args: Tuple containing parameter arguments for `rbf`.

    Returns:
        - Tensor of shape '(N, K, S)'.
    """
    m = rbf(dist_fn(grid[..., None, :], grid[..., None, :, :]), width)

    coeffs = torch.solve(signal, m)[0].transpose(-1, -2)
    p = rbf(dist_fn(rotations[..., None, :], grid[..., None, :, :]), width)

    # coeffs (N, P, G)    p (N, S, G) --> (N, P, S, G)
    return torch.sum(coeffs[:, None] * p[..., None, :], axis=-1, keepdim=True).squeeze(
        -1
    )
//  [python]
def _depr_rbf_interpolation(
    rotations: Tensor,
    grid: Tensor,
    signal: Tensor,
    *,
    dist_fn: Tensor = geodesic_distance,
    rbf: Callable[[Tensor], Tensor] = _rbf_gauss2,
    width=2,
) -> Tensor:
    """
    Performs radial basis function interpolation of rotations.

    Arguments:
        Arguments:
        - rotations: Tensor of shape `(N, K, 4)` containing quaternions.
        - grid: Tensor of shape `(N, L, 4)` containing quaternions
                defining a grid on SO3.
        - signal: Tensor of shape `(N, L, S)` containing the signal
                  corresponding to `grid`.
        - dist_fn: Distance function used to calculate distance between
                   rotations.
        - rbf: Radial basis function, defaults to gaussian rbf.
        - rbf_args: Tuple containing parameter arguments for `rbf`.

    Returns:
        - Tensor of shape '(N, K, S)'.
    """
    m = rbf(dist_fn(grid[..., None, :], grid[..., None, :, :]), width)

    coeffs = torch.linalg.solve(m, signal).transpose(-1, -2)
    p = rbf(dist_fn(rotations[..., None, :], grid[..., None, :, :]), width)

    # coeffs (N, P, G)    p (N, S, G) --> (N, P, S, G)
    return torch.sum(coeffs[:, None] * p[..., None, :], axis=-1, keepdim=True).squeeze(
        -1
    )
//  [python]
def ABC_DNPLR(self):
        assert self.rank == 1, "Not implemented for rank > 1"
        w = torch.view_as_complex(self.w)
        BpV = torch.view_as_complex(self.BV)
        CqV = torch.view_as_complex(self.CV)
        w = torch.cat([w, w.conj()], dim=-1)
        BpV = torch.cat([BpV, BpV.conj()], dim=-1)
        CqV = torch.cat([CqV, CqV.conj()], dim=-1)

        BV = BpV[..., 0, :]
        pV = BpV[..., 1, :]
        CV = CqV[..., 0, :]
        qV = CqV[..., 1, :]
        if self.use_length:
            dA = self.bilinear_DNPLR()
            dA_L = power(self.L, dA)
            I = torch.eye(dA.size(-1)).to(dA)
            CV = torch.solve(CV.unsqueeze(-1), I - dA_L.transpose(-1,-2))[0].squeeze(-1)
        AV = torch.diag_embed(w) - contract('... n, ... m -> ... n m', pV.conj(), qV)
        return AV, BV, CV
//  [python]
def ABC_DNPLR(self):
        assert self.rank == 1, "Not implemented for rank > 1"
        w = torch.view_as_complex(self.w)
        BpV = torch.view_as_complex(self.BV)
        CqV = torch.view_as_complex(self.CV)
        w = torch.cat([w, w.conj()], dim=-1)
        BpV = torch.cat([BpV, BpV.conj()], dim=-1)
        CqV = torch.cat([CqV, CqV.conj()], dim=-1)

        BV = BpV[..., 0, :]
        pV = BpV[..., 1, :]
        CV = CqV[..., 0, :]
        qV = CqV[..., 1, :]
        if self.use_length:
            dA = self.bilinear_DNPLR()
            dA_L = power(self.L, dA)
            I = torch.eye(dA.size(-1)).to(dA)
            CV = torch.linalg.solve(I - dA_L.transpose(-1,-2), CV.unsqueeze(-1)).squeeze(-1)
        AV = torch.diag_embed(w) - contract('... n, ... m -> ... n m', pV.conj(), qV)
        return AV, BV, CV
//  [python]
def test_solve(self, device, dtype):
        def run_test(n, batch, rhs):
            A_dims = (*batch, n, n)
            b_dims = (*batch, n, *rhs)
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)

            # Correctness test
            x = torch.solve(b, A)[0]
            if rhs == ():
                Ax = np.matmul(A.cpu(), x.unsqueeze(-1).cpu())
                Ax.squeeze_(-1)
            else:
                Ax = np.matmul(A.cpu(), x.cpu())
            self.assertEqual(b.expand_as(Ax), Ax)

            # Check against NumPy
            expected = np.linalg.solve(A.cpu().numpy(), b.expand_as(x).cpu().numpy())
            self.assertEqual(x, expected)

            # Check out= variant
            out = torch.empty_like(x)
            ans = torch.solve(b, A, out=out)[0]
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

            # Check out= variant with complex128 out tensor
            out = torch.empty_like(x).to(torch.complex128)
            ans = torch.solve(b, A, out=out)[0]
            self.assertEqual(ans, out)
            self.assertEqual(x.to(torch.complex128), out)

            # Check empty out
            out = torch.empty(0, dtype=dtype, device=device)
            ans = torch.solve(b, A, out=out)[0]
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

        batches = [(), (0, ), (3, ), (2, 3)]
        ns = [0, 5, 32]
        nrhs = [(), (1, ), (5, )]
        for n, batch, rhs in itertools.product(ns, batches, nrhs):
            run_test(n, batch, rhs)
//  [python]
def test_solve(self, device, dtype):
        def run_test(n, batch, rhs):
            A_dims = (*batch, n, n)
            b_dims = (*batch, n, *rhs)
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)

            # Correctness test
            x = torch.linalg.solve(A, b)
            if rhs == ():
                Ax = np.matmul(A.cpu(), x.unsqueeze(-1).cpu())
                Ax.squeeze_(-1)
            else:
                Ax = np.matmul(A.cpu(), x.cpu())
            self.assertEqual(b.expand_as(Ax), Ax)

            # Check against NumPy
            expected = np.linalg.solve(A.cpu().numpy(), b.expand_as(x).cpu().numpy())
            self.assertEqual(x, expected)

            # Check out= variant
            out = torch.empty_like(x)
            ans = torch.linalg.solve(A, b, out=out)
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

            # Check out= variant with complex128 out tensor
            out = torch.empty_like(x).to(torch.complex128)
            ans = torch.linalg.solve(A, b, out=out)
            self.assertEqual(ans, out)
            self.assertEqual(x.to(torch.complex128), out)

            # Check empty out
            out = torch.empty(0, dtype=dtype, device=device)
            ans = torch.linalg.solve(A, b, out=out)
            self.assertEqual(ans, out)
            self.assertEqual(x, out)

        batches = [(), (0, ), (3, ), (2, 3)]
        ns = [0, 5, 32]
        nrhs = [(), (1, ), (5, )]
        for n, batch, rhs in itertools.product(ns, batches, nrhs):
            run_test(n, batch, rhs)
//  [python]
def forward(self, X: torch.Tensor) -> torch.Tensor:
        n, k = X.size(-2), X.size(-1)
        transposed = n < k
        if transposed:
            X = X.mT
            n, k = k, n
        # Here n > k and X is a tall matrix
        if self.orthogonal_map == _OrthMaps.matrix_exp or self.orthogonal_map == _OrthMaps.cayley:
            # We just need n x k - k(k-1)/2 parameters
            X = X.tril()
            if n != k:
                # Embed into a square matrix
                X = torch.cat([X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1)
            A = X - X.mH
            # A is skew-symmetric (or skew-hermitian)
            if self.orthogonal_map == _OrthMaps.matrix_exp:
                Q = torch.matrix_exp(A)
            elif self.orthogonal_map == _OrthMaps.cayley:
                # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1}
                Id = torch.eye(n, dtype=A.dtype, device=A.device)
                Q = torch.solve(torch.add(Id, A, alpha=0.5), torch.add(Id, A, alpha=-0.5))[1]
            # Q is now orthogonal (or unitary) of size (..., n, n)
            if n != k:
                Q = Q[..., :k]
            # Q is now the size of the X (albeit perhaps transposed)
        else:
            # X is real here, as we do not support householder with complex numbers
            A = X.tril(diagonal=-1)
            tau = 2. / (1. + (A * A).sum(dim=-2))
            Q = torch.linalg.householder_product(A, tau)
            # The diagonal of X is 1's and -1's
            # We do not want to differentiate through this or update the diagonal of X hence the casting
            Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2)

        if hasattr(self, "base"):
            Q = self.base @ Q
        if transposed:
            Q = Q.mT
        return Q
//  [python]
def forward(self, X: torch.Tensor) -> torch.Tensor:
        n, k = X.size(-2), X.size(-1)
        transposed = n < k
        if transposed:
            X = X.mT
            n, k = k, n
        # Here n > k and X is a tall matrix
        if self.orthogonal_map == _OrthMaps.matrix_exp or self.orthogonal_map == _OrthMaps.cayley:
            # We just need n x k - k(k-1)/2 parameters
            X = X.tril()
            if n != k:
                # Embed into a square matrix
                X = torch.cat([X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1)
            A = X - X.mH
            # A is skew-symmetric (or skew-hermitian)
            if self.orthogonal_map == _OrthMaps.matrix_exp:
                Q = torch.matrix_exp(A)
            elif self.orthogonal_map == _OrthMaps.cayley:
                # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1}
                Id = torch.eye(n, dtype=A.dtype, device=A.device)
                Q = torch.linalg.solve(torch.add(Id, A, alpha=-0.5), torch.add(Id, A, alpha=0.5))
            # Q is now orthogonal (or unitary) of size (..., n, n)
            if n != k:
                Q = Q[..., :k]
            # Q is now the size of the X (albeit perhaps transposed)
        else:
            # X is real here, as we do not support householder with complex numbers
            A = X.tril(diagonal=-1)
            tau = 2. / (1. + (A * A).sum(dim=-2))
            Q = torch.linalg.householder_product(A, tau)
            # The diagonal of X is 1's and -1's
            # We do not want to differentiate through this or update the diagonal of X hence the casting
            Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2)

        if hasattr(self, "base"):
            Q = self.base @ Q
        if transposed:
            Q = Q.mT
        return Q
//  [python]
def dABC(self): # TODO fix for use_length case
        # return torch.zeros(self.CV.size(0), self.w.size(-2), self.w.size(-2)), torch.zeros(self.CV.size(0), self.w.size(-2)),torch.zeros(self.CV.size(0), self.w.size(-2))
        dA, dB = self.dAB()
        dA_L = power(self.L, dA)
        I = torch.eye(dA.size(-1)).to(dA)

        CV = torch.view_as_complex(self.CV)
        V = torch.view_as_complex(self.V)
        C = contract('... r m, ... m n -> ... r n', CV, V.transpose(-1, -2).conj())
        C = 2 * C[..., 0, :].real

        dC = torch.solve(C.unsqueeze(-1), I - dA_L.transpose(-1,-2))[0].squeeze(-1)
        return dA, dB, dC
//  [python]
def dABC(self): # TODO fix for use_length case
        # return torch.zeros(self.CV.size(0), self.w.size(-2), self.w.size(-2)), torch.zeros(self.CV.size(0), self.w.size(-2)),torch.zeros(self.CV.size(0), self.w.size(-2))
        dA, dB = self.dAB()
        dA_L = power(self.L, dA)
        I = torch.eye(dA.size(-1)).to(dA)

        CV = torch.view_as_complex(self.CV)
        V = torch.view_as_complex(self.V)
        C = contract('... r m, ... m n -> ... r n', CV, V.transpose(-1, -2).conj())
        C = 2 * C[..., 0, :].real

        dC = torch.linalg.solve(I - dA_L.transpose(-1,-2), C.unsqueeze(-1)).squeeze(-1)
        return dA, dB, dC
//  [python]
def bilinear(cls, dt, A, B=None, separate=False):
        """
        dt: (...) timescales
        A: (... N N)
        B: (... N)
        """
        N = A.shape[-1]
        I = torch.eye(N).to(A)
        A_backwards = I - dt[:, None, None]/2 * A
        A_forwards = I + dt[:, None, None]/2 * A

        if B is None:
            dB = None
        else:
            dB = dt[..., None] * torch.solve(B.unsqueeze(-1), A_backwards)[0].squeeze(-1) # (... N)

        if separate:
            A_b = torch.solve(I, A_backwards)[0] # (... N N)
            return A_forwards, A_b, dB
        else:
            dA = torch.solve(A_forwards, A_backwards)[0] # (... N N)
            return dA, dB
//  [python]
def bilinear(cls, dt, A, B=None, separate=False):
        """
        dt: (...) timescales
        A: (... N N)
        B: (... N)
        """
        N = A.shape[-1]
        I = torch.eye(N).to(A)
        A_backwards = I - dt[:, None, None]/2 * A
        A_forwards = I + dt[:, None, None]/2 * A

        if B is None:
            dB = None
        else:
            dB = dt[..., None] * torch.linalg.solve(A_backwards, B.unsqueeze(-1)).squeeze(-1) # (... N)

        if separate:
            A_b = torch.linalg.solve(A_backwards, I) # (... N N)
            return A_forwards, A_b, dB
        else:
            dA = torch.linalg.solve(A_backwards, A_forwards) # (... N N)
            return dA, dB
//  [python]
def solve_system(self):
        """
        Solve the global system of equations to get the nodal displacements.
        """

        known_dof_indices, unknown_dof_indices = self._get_dof_indices()

        # Extract submatrix and subvector
        K_sub = self.global_stiffness[unknown_dof_indices, :][:, unknown_dof_indices]
        R_sub = self.global_load[unknown_dof_indices]

        # Solve the system
        u_sub = torch.solve(R_sub.unsqueeze(1), K_sub)[0]


        # Create a global displacement vector
        self.global_displacements = torch.zeros(self.num_dofs, device=self.device).to(dtype=torch.float64)
        self.global_displacements[unknown_dof_indices] = u_sub.squeeze()

        # Fill known displacements, if any
        if self.node_dof_disp.numel() > 0:
            self.global_displacements[known_dof_indices] = self.node_dof_disp[:, 2].to(dtype=torch.float64)
//  [python]
def solve_system(self):
        """
        Solve the global system of equations to get the nodal displacements.
        """

        # # Identify known DOFs from self.node_dof_disp
        # if self.node_dof_disp.numel() > 0:
        #     known_dof_indices = (self.node_dof_disp[:, 0] * self.parameters['num_dimensions'] + self.node_dof_disp[:, 1]).long()
        # else:
        #     known_dof_indices = torch.tensor([], dtype=torch.long, device=self.device)
        # # Identify unknown DOFs
        # all_dofs = torch.arange(self.num_dofs, device=self.device)
        # unknown_dof_indices = torch.tensor([idx for idx in all_dofs if idx not in known_dof_indices], device=self.device)
        known_dof_indices, unknown_dof_indices = self._get_dof_indices()

        # Extract submatrix and subvector
        K_sub = self.global_stiffness[unknown_dof_indices, :][:, unknown_dof_indices]
        R_sub = self.global_load[unknown_dof_indices]

        # Solve the system
        u_sub = torch.linalg.solve(K_sub, R_sub.unsqueeze(1))


        # residual = R_sub - K_sub @ u_sub.squeeze()
        # relative_residual_error = torch.norm(residual) / torch.norm(R_sub)
        # print(f"Relative Residual Error: {relative_residual_error.item():.8f}")



        # Create a global displacement vector
        self.global_displacements = torch.zeros(self.num_dofs, device=self.device).to(dtype=torch.float64)
        self.global_displacements[unknown_dof_indices] = u_sub.squeeze()

        # Fill known displacements, if any
        if self.node_dof_disp.numel() > 0:
            self.global_displacements[known_dof_indices] = self.node_dof_disp[:, 2].to(dtype=torch.float64)
//  [python]
def solve_system_nonlinear_finiteStrain(self):
        """
        Solve the global system of equations to get the nodal displacements.
        """

        # Identify known DOFs from self.node_dof_disp
        known_dof_indices, unknown_dof_indices = self._get_dof_indices()

        # Extract submatrix and subvector
        K_sub = self.global_stiffness[unknown_dof_indices, :][:, unknown_dof_indices]
        R_sub = self.global_load[unknown_dof_indices]
        F_sub = self.global_residual_stress[unknown_dof_indices]
        B_sub = R_sub - F_sub
        

        # Solve the system
        self.u_incr_sub = torch.solve(B_sub.unsqueeze(1), K_sub)[0]

        self.global_displacements_temp[unknown_dof_indices] += self.u_incr_sub.squeeze()

        if known_dof_indices.numel() > 0:
            self.global_displacements_temp[known_dof_indices] = self.current_known_disp
//  [python]
def solve_system_nonlinear_finiteStrain(self):
        """
        Solve the global system of equations to get the nodal displacements.
        """

        # Identify known DOFs from self.node_dof_disp
        known_dof_indices, unknown_dof_indices = self._get_dof_indices()

        # Extract submatrix and subvector
        K_sub = self.global_stiffness[unknown_dof_indices, :][:, unknown_dof_indices]
        R_sub = self.global_load[unknown_dof_indices]
        F_sub = self.global_residual_stress[unknown_dof_indices]
        B_sub = R_sub - F_sub
        

        # Solve the system
        self.u_incr_sub = torch.linalg.solve(K_sub, B_sub.unsqueeze(1))

        self.global_displacements_temp[unknown_dof_indices] += self.u_incr_sub.squeeze()

        if known_dof_indices.numel() > 0:
            self.global_displacements_temp[known_dof_indices] = self.current_known_disp
//  [python]
def _direct_solve_I_p_lambda_AT_A_x_eq_b(self, lambda_, b, **kwargs):
        if self.__ATA_matrix is None:
            A = self._matrix.clone()
            self.__ATA_matrix = A.T @ A

        if self.__lambda_cache != lambda_:
            lATA = lambda_ * (self.__ATA_matrix)
            torch.diagonal(lATA)[:] += 1
            self.__IplATA_matrix = lATA
        return torch.solve(self.__IplATA_matrix, b)[0]
//  [python]
def _direct_solve_I_p_lambda_AT_A_x_eq_b(self, lambda_, b, **kwargs):
        if self.__ATA_matrix is None:
            A = self._matrix.clone()
            self.__ATA_matrix = A.T @ A

        if self.__lambda_cache != lambda_:
            lATA = lambda_ * (self.__ATA_matrix)
            torch.diagonal(lATA)[:] += 1
            self.__IplATA_matrix = lATA
        return torch.linalg.solve(self.__IplATA_matrix, b)