Python torch.cholesky方法代码示例

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def _get_sps(self, mean: torch.Tensor, cov: torch.Tensor):
        """
        Constructs the Sigma points used for propagation.
        :return: Sigma points
        """

        self._mean[..., self._sslc] = mean
        self._cov[..., self._sslc, self._sslc] = cov

        cholcov = sqrt(self._lam + self._ndim) * torch.cholesky(self._cov)

        spx = self._mean.unsqueeze(-2)
        sph = self._mean[..., None, :] + cholcov
        spy = self._mean[..., None, :] - cholcov

        return torch.cat((spx, sph, spy), -2) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def _kernel_2d(self, y, loc, h_var_inv, o_var_inv, c):
        tc = c if self._model.obs_ndim > 0 else c.unsqueeze(-2)

        # ===== Define covariance ===== #
        ttc = tc.transpose(-2, -1)
        diag_o_var_inv = construct_diag(o_var_inv if self._model.observable.ndim > 0 else o_var_inv.unsqueeze(-1))
        t2 = torch.matmul(ttc, torch.matmul(diag_o_var_inv, tc))

        cov = (construct_diag(h_var_inv) + t2).inverse()

        # ===== Get mean ===== #
        t1 = h_var_inv * loc

        t2 = torch.matmul(diag_o_var_inv, y if y.dim() > 0 else y.unsqueeze(-1))
        t3 = torch.matmul(ttc, t2.unsqueeze(-1))[..., 0]

        m = torch.matmul(cov, (t1 + t3).unsqueeze(-1))[..., 0]

        return MultivariateNormal(m, scale_tril=torch.cholesky(cov)) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_discrete_mvn_log_prob(init_shape, trans_shape, obs_shape, state_dim):
    event_size = 4
    init_logits = torch.randn(init_shape + (state_dim,))
    trans_logits = torch.randn(trans_shape + (state_dim, state_dim))
    loc = torch.randn(obs_shape + (state_dim, event_size))
    cov = torch.randn(obs_shape + (state_dim, event_size, 2 * event_size))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.cholesky(cov)
    obs_dist = dist.MultivariateNormal(loc, scale_tril=scale_tril)

    actual_dist = DiscreteHMM(init_logits, trans_logits, obs_dist)
    expected_dist = dist.DiscreteHMM(init_logits, trans_logits, obs_dist)
    assert actual_dist.event_shape == expected_dist.event_shape
    assert actual_dist.batch_shape == expected_dist.batch_shape

    batch_shape = broadcast_shape(init_shape + (1,), trans_shape, obs_shape)
    data = obs_dist.expand(batch_shape + (state_dim,)).sample()
    data = data[(slice(None),) * len(batch_shape) + (0,)]
    actual_log_prob = actual_dist.log_prob(data)
    expected_log_prob = expected_dist.log_prob(data)
    assert_close(actual_log_prob, expected_log_prob)
    check_expand(actual_dist, data) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_update_cholesky():
    """Test that the update cholesky function returns correct values."""
    n = 6
    new_A = torch.rand(n, n, dtype=torch.float64)
    new_A = new_A @ new_A.t()
    new_A += torch.eye(len(new_A), dtype=torch.float64)

    A = new_A[:n - 1, :n - 1]

    old_chol = torch.cholesky(A, upper=False)
    new_row = new_A[-1]

    # Test updateing overall
    new_chol = update_cholesky(old_chol, new_row)
    error = new_chol - torch.cholesky(new_A, upper=False)
    assert torch.all(torch.abs(error) <= 1e-15)

    # Test updating inplace
    new_chol = torch.zeros(n, n, dtype=torch.float64)
    new_chol[:n - 1, :n - 1] = old_chol

    update_cholesky(old_chol, new_row, chol_row_out=new_chol[-1])
    error = new_chol - torch.cholesky(new_A, upper=False)
    assert torch.all(torch.abs(error) <= 1e-15) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_cg_with_tridiag(self):
        size = 10
        matrix = torch.randn(size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.transpose(-1, -2))
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(size, 50, dtype=torch.float64)
        solves, t_mats = linear_cg(
            matrix.matmul, rhs=rhs, n_tridiag=5, max_tridiag_iter=10, max_iter=size, tolerance=0, eps=1e-15
        )

        # Check cg
        matrix_chol = matrix.cholesky()
        actual = torch.cholesky_solve(rhs, matrix_chol)
        self.assertTrue(torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        # Check tridiag
        eigs = matrix.symeig()[0]
        for i in range(5):
            approx_eigs = t_mats[i].symeig()[0]
            self.assertTrue(torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4)) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_solve_qr(self, dtype=torch.float64, tol=1e-8):
        size = 50
        X = torch.rand((size, 2)).to(dtype=dtype)
        y = torch.sin(torch.sum(X, 1)).unsqueeze(-1).to(dtype=dtype)
        with settings.min_preconditioning_size(0):
            noise = torch.DoubleTensor(size).uniform_(math.log(1e-3), math.log(1e-1)).exp_().to(dtype=dtype)
            lazy_tsr = RBFKernel().to(dtype=dtype)(X).evaluate_kernel().add_diag(noise)
            precondition_qr, _, logdet_qr = lazy_tsr._preconditioner()

            F = lazy_tsr._piv_chol_self
            M = noise.diag() + F.matmul(F.t())

        x_exact = torch.solve(y, M)[0]
        x_qr = precondition_qr(y)

        self.assertTrue(approx_equal(x_exact, x_qr, tol))

        logdet = 2 * torch.cholesky(M).diag().log().sum(-1)
        self.assertTrue(approx_equal(logdet, logdet_qr, tol)) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_solve_qr_constant_noise(self, dtype=torch.float64, tol=1e-8):
        size = 50
        X = torch.rand((size, 2)).to(dtype=dtype)
        y = torch.sin(torch.sum(X, 1)).unsqueeze(-1).to(dtype=dtype)

        with settings.min_preconditioning_size(0):
            noise = 1e-2 * torch.ones(size, dtype=dtype)
            lazy_tsr = RBFKernel().to(dtype=dtype)(X).evaluate_kernel().add_diag(noise)
            precondition_qr, _, logdet_qr = lazy_tsr._preconditioner()

            F = lazy_tsr._piv_chol_self
        M = noise.diag() + F.matmul(F.t())

        x_exact = torch.solve(y, M)[0]
        x_qr = precondition_qr(y)

        self.assertTrue(approx_equal(x_exact, x_qr, tol))

        logdet = 2 * torch.cholesky(M).diag().log().sum(-1)
        self.assertTrue(approx_equal(logdet, logdet_qr, tol)) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def deprecate_task_noise_corr(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
    if prefix + "task_noise_corr_factor" in state_dict:
        # Remove after 1.0
        warnings.warn(
            "Loading a deprecated parameterization of _MultitaskGaussianLikelihoodBase. Consider re-saving your model.",
            OldVersionWarning,
        )
        # construct the task correlation matrix from the factors using the old parameterization
        corr_factor = state_dict.pop(prefix + "task_noise_corr_factor").squeeze(0)
        corr_diag = state_dict.pop(prefix + "task_noise_corr_diag").squeeze(0)
        num_tasks, rank = corr_factor.shape[-2:]
        M = corr_factor.matmul(corr_factor.transpose(-1, -2))
        idx = torch.arange(M.shape[-1], dtype=torch.long, device=M.device)
        M[..., idx, idx] += corr_diag
        sem_inv = 1 / torch.diagonal(M, dim1=-2, dim2=-1).sqrt().unsqueeze(-1)
        C = M * sem_inv.matmul(sem_inv.transpose(-1, -2))
        # perform a Cholesky decomposition and extract the required entries
        L = torch.cholesky(C)
        tidcs = torch.tril_indices(num_tasks, rank)[:, 1:]
        task_noise_corr = L[..., tidcs[0], tidcs[1]]
        state_dict[prefix + "task_noise_corr"] = task_noise_corr 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def _add_jitter(self, X: Tensor) -> Tensor:
        jitter_prev = 0
        Eye = torch.eye(X.size(-1)).expand(X.shape)
        for i in range(3):
            jitter_new = self._jitter * (10 ** i)
            X = X + (jitter_new - jitter_prev) * Eye
            jitter_prev = jitter_new
            # This may be VERY slow given upstream pytorch issue:
            # https://github.com/pytorch/pytorch/issues/34272
            try:
                _ = torch.cholesky(X)
                warnings.warn(
                    "X is not a p.d. matrix; "
                    f"Added jitter of {jitter_new:.2e} to the diagonal",
                    RuntimeWarning,
                )
                return X
            except RuntimeError:
                continue
        warnings.warn(
            f"Failed to render X p.d. after adding {jitter_new:.2e} jitter",
            RuntimeWarning,
        )
        return X 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def _construct_mvn(x: torch.Tensor, w: torch.Tensor):
    """
    Constructs a multivariate normal distribution of weighted samples.
    :param x: The samples
    :param w: The weights
    """

    mean = (x * w.unsqueeze(-1)).sum(0)
    centralized = x - mean
    cov = torch.matmul(w * centralized.t(), centralized)

    return MultivariateNormal(mean, scale_tril=torch.cholesky(cov)) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def _cholesky(x):
    """
    Like :func:`torch.cholesky` but uses sqrt for scalar matrices.
    Works around https://github.com/pytorch/pytorch/issues/24403 often.
    """
    if x.size(-1) == 1:
        return x.sqrt()
    return x.cholesky() 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_dist_to_funsor_mvn(batch_shape, event_size):
    loc = torch.randn(batch_shape + (event_size,))
    cov = torch.randn(batch_shape + (event_size, 2 * event_size))
    cov = cov.matmul(cov.transpose(-1, -2))
    scale_tril = torch.cholesky(cov)
    d = dist.MultivariateNormal(loc, scale_tril=scale_tril)
    f = dist_to_funsor(d)
    assert isinstance(f, Funsor)

    value = d.sample()
    actual_log_prob = f(value=tensor_to_funsor(value, event_output=1))
    expected_log_prob = tensor_to_funsor(d.log_prob(value))
    assert_close(actual_log_prob, expected_log_prob) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def to_log_cholesky(self) -> Tuple[torch.Tensor, torch.Tensor]:
        batch_dim = self.shape[:-2]
        rank = self.shape[-1]
        L = torch.cholesky(self)

        n_off = int(rank * (rank - 1) / 2)
        off_diag = torch.empty(batch_dim + (n_off,))
        idx = 0
        for i in range(rank):
            for j in range(i):
                off_diag[..., idx] = L[..., i, j]
                idx += 1
        log_diag = torch.log(torch.diagonal(L, dim1=-2, dim2=-1))
        return log_diag, off_diag 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_lkj_cholesky_factor_prior_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJCholeskyFactorPrior(2, torch.tensor(0.5, device=device))
        S = torch.eye(2, device=device)
        S_chol = torch.cholesky(S)
        self.assertAlmostEqual(prior.log_prob(S_chol).item(), -1.86942, places=4)
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S_chol.device)])
        S_chol = torch.stack([torch.cholesky(Si) for Si in S])
        self.assertTrue(approx_equal(prior.log_prob(S_chol), torch.tensor([-1.86942, -1.72558], device=S_chol.device)))
        with self.assertRaises(ValueError):
            prior.log_prob(torch.eye(3, device=device))

        # For eta=1.0 log_prob is flat over all covariance matrices
        prior = LKJCholeskyFactorPrior(2, torch.tensor(1.0, device=device))
        self.assertTrue(torch.all(prior.log_prob(S_chol) == prior.C)) 

# 需要导入模块: import torch [as 别名]
# 或者: from torch import cholesky [as 别名]
def test_lkj_cholesky_factor_prior_batch_log_prob(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        prior = LKJCholeskyFactorPrior(2, torch.tensor([0.5, 1.5], device=device))

        S = torch.eye(2, device=device)
        S_chol = torch.cholesky(S)
        self.assertTrue(approx_equal(prior.log_prob(S_chol), torch.tensor([-1.86942, -0.483129], device=S_chol.device)))
        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
        S_chol = torch.stack([torch.cholesky(Si) for Si in S])
        self.assertTrue(approx_equal(prior.log_prob(S_chol), torch.tensor([-1.86942, -0.62697], device=S_chol.device)))
        with self.assertRaises(ValueError):
            prior.log_prob(torch.eye(3, device=device))