graphld.precision

Precision matrix operations using LDGM graphical models.

The PrecisionOperator class is the core abstraction for working with LD matrices. It subclasses SciPy's LinearOperator and represents an LDGM precision matrix or its Schur complement.

Key Concepts

• To compute correlation_matrix[indices, indices] @ vector, use ldgm[indices].solve(vector)
• To compute inv(correlation_matrix[indices, indices]) @ vector, use ldgm[indices] @ vector

You cannot do this indexing manually - the submatrix of an inverse is not the inverse of a submatrix. See Section 5 of the supplementary material.

For conceptual usage and examples, see the Matrix Operations guide.

precision

Precision matrix operations for LDGM.

This module implements sparse precision matrix operations using scipy's LinearOperator interface for efficient matrix-vector operations.

PrecisionOperator dataclass

PrecisionOperator(_matrix: csc_matrix, variant_info: DataFrame, _which_indices: Optional[ndarray] = None, _solver: Optional[cholesky] = None, _cholesky_is_up_to_date: bool = False, _matrix_version: Optional[list[int]] = None, _factor_version: int = -1)

Bases: LinearOperator

LDGM precision matrix class implementing the LinearOperator interface.

This class provides an efficient implementation of precision matrix operations using the scipy.sparse.linalg.LinearOperator interface. It supports matrix-vector multiplication and other essential operations for working with LDGM precision matrices.

Attributes:

Name	Type	Description
_matrix	csc_matrix	The precision matrix in sparse format
variant_info	DataFrame	Polars DataFrame containing variant information
_which_indices	Optional[ndarray]	Array of indices for current selection
_solver	Optional[cholesky]	Previously computed Cholesky factorization
_cholesky_is_up_to_date	bool	Flag indicating whether the Cholesky factorization is up to date

shape property

shape

Get the shape of the matrix, accounting for partial indexing.

dtype property

dtype: dtype

Return the dtype of the precision matrix.

nbytes property

nbytes: int

Return the total memory usage in bytes.

matrix property

matrix

Get the precision matrix.

variant_indices property

variant_indices

Get the indices of the variants in the precision matrix.

diagonal_indices cached property

diagonal_indices: ndarray

Get indices of diagonal elements corresponding to _which_indices in _matrix.data.

Returns:

Type	Description
ndarray	Array of indices into self._matrix.data where diagonal elements are stored

inv property

inv: LinearOperator

A LinearOperator representing the LD correlation matrix.

times_scalar

times_scalar(multiplier: float) -> None

Multiply the precision matrix by a scalar in place.

Source code in src/graphld/precision.py

def times_scalar(self, multiplier: float) -> None:
    """Multiply the precision matrix by a scalar in place."""
    self._matrix *= multiplier
    self._bump_matrix_version()
    self.del_factor()

update_matrix

update_matrix(update: ndarray) -> None

Update the precision matrix by adding values to its diagonal.

If which_indices is set, only updates the corresponding diagonal elements.

Parameters:

Name	Type	Description	Default
update	ndarray	Vector of values to add to the diagonal elements	required

Raises:

Type	Description
ValueError	If update shape doesn't match or would make diagonal non-positive

Source code in src/graphld/precision.py

def update_matrix(self, update: np.ndarray) -> None:
    """Update the precision matrix by adding values to its diagonal.

    If which_indices is set, only updates the corresponding diagonal elements.

    Args:
        update: Vector of values to add to the diagonal elements

    Raises:
        ValueError: If update shape doesn't match or would make diagonal non-positive
    """
    update = np.asarray(update).reshape(-1)
    if len(update) != self.shape[0]:
        msg = f"Update vector length {len(update)} does not match matrix shape {self.shape}"
        raise ValueError(msg)

    diagonal_indices = self.diagonal_indices
    updated_diagonal = self._matrix.data[diagonal_indices] + update

    if np.any(updated_diagonal <= 0):
        raise ValueError("Update would make a diagonal element non-positive")

    self._matrix.data[diagonal_indices] = updated_diagonal
    self._bump_matrix_version()
    self._cholesky_is_up_to_date = False
    self._factor_version = -1

update_element

update_element(index: int, value: float) -> None

Update a single diagonal element of the precision matrix.

If which_indices is set, only updates the corresponding element.

Parameters:

Name	Type	Description	Default
index	int	The updated diagonal element is self._which_indices[index]	required
value	float	Value to add to the diagonal element	required

Note

Updates the Cholesky factorization if it exists using a rank-1 update/downdate.

Source code in src/graphld/precision.py

def update_element(self, index: int, value: float) -> None:
    """Update a single diagonal element of the precision matrix.

    If which_indices is set, only updates the corresponding element.

    Args:
        index: The updated diagonal element is self._which_indices[index]
        value: Value to add to the diagonal element

    Note:
        Updates the Cholesky factorization if it exists using a rank-1 update/downdate.
    """
    # Find diagonal element in sparse matrix
    diag_pos = self.diagonal_indices[index]

    # Convert to global index for Cholesky update
    global_index = index if self._which_indices is None else self._which_indices[index]

    # Update matrix value
    new_val = self._matrix.data[diag_pos] + value
    if new_val <= 0:
        msg = f"Update would make diagonal element at index {global_index} non-positive"
        raise ValueError(msg)

    factor_was_current = self._factor_is_current()
    self._matrix.data[diag_pos] = new_val
    self._bump_matrix_version()

    if not factor_was_current:
        self._cholesky_is_up_to_date = False
        self._factor_version = -1
        return

    # Update Cholesky factorization
    shape = (self._matrix.shape[0], 1)
    e_sparse = csc_matrix(([np.sqrt(np.abs(value))], ([global_index], [0])), shape=shape)
    self._solver.update_inplace(e_sparse, value < 0)
    self._cholesky_is_up_to_date = True
    self._factor_version = self._current_matrix_version()

factor

factor() -> None

Update the Cholesky factorization of the precision matrix.

Source code in src/graphld/precision.py

def factor(self) -> None:
    """Update the Cholesky factorization of the precision matrix."""
    if self._solver is None:
        self._solver = cholesky(self._matrix)
    else:
        self._solver.cholesky_inplace(self._matrix)
    self._cholesky_is_up_to_date = True
    self._factor_version = self._current_matrix_version()

del_factor

del_factor() -> None

Free the memory used by the Cholesky factorization.

Source code in src/graphld/precision.py

def del_factor(self) -> None:
    """Free the memory used by the Cholesky factorization."""
    self._solver = None
    self._cholesky_is_up_to_date = False
    self._factor_version = -1

logdet

logdet() -> float

Compute log determinant of the Schur complement.

Returns:

Type	Description
float	Log determinant of the Schur complement

Source code in src/graphld/precision.py

def logdet(self) -> float:
    """Compute log determinant of the Schur complement.

    Returns:
        Log determinant of the Schur complement
    """
    if not self._factor_is_current():
        self.factor()

    # Get boolean mask for current selection
    mask = self._get_mask
    if np.all(mask):
        return self._solver.logdet()

    # Get P11 submatrix for missing indices
    P11 = self._matrix[~mask][:, ~mask]
    logdet_P11 = cholesky(P11).logdet()

    # Return logdet(P) - logdet(P11)
    return self._solver.logdet() - logdet_P11

copy

copy()

Copy the current LDGM instance.

Source code in src/graphld/precision.py

def copy(self):
    """Copy the current LDGM instance."""
    which_indices = self._which_indices.copy() if self._which_indices is not None else None
    matrix = self._matrix.copy()
    matrix._graphld_matrix_version = [0]
    return PrecisionOperator(matrix, self.variant_info.clone(),
                            which_indices, None, False, matrix._graphld_matrix_version)

__getitem__

__getitem__(key: Union[list, slice, ndarray]) -> PrecisionOperator

Sets the _which_indices class attribute.

Parameters:

Name	Type	Description	Default
key	Union[list, slice, ndarray]	Index, slice, or tuple of indices/slices to access. Can be: - List of indices - Array of indices - Boolean mask array - Slice object	required

Returns:

Type	Description
PrecisionOperator	New PrecisionOperator instance that shares the underlying matrix

Source code in src/graphld/precision.py

def __getitem__(self, key: Union[list, slice, np.ndarray]) -> 'PrecisionOperator':
    """Sets the _which_indices class attribute.

    Args:
        key: Index, slice, or tuple of indices/slices to access. Can be:
            - List of indices
            - Array of indices
            - Boolean mask array
            - Slice object

    Returns:
        New PrecisionOperator instance that shares the underlying matrix
    """
    # Create new instance sharing the same matrix and variant_info
    result = PrecisionOperator(
        self._matrix,
        self.variant_info,
        self._which_indices,
        None,
        False,
        self._matrix_state()
    )
    result.set_which_indices(key)
    return result

set_which_indices

set_which_indices(key: Union[list, slice, ndarray, int]) -> None

Sets the _which_indices class attribute.

Parameters:

Name	Type	Description	Default
key	Union[list, slice, ndarray, int]	Index, slice, or tuple of indices/slices to access. Can be: - List of indices - Array of indices - Boolean mask array - Slice object - Integer index	required

Source code in src/graphld/precision.py

def set_which_indices(self, key: Union[list, slice, np.ndarray, int]) -> None:
    """Sets the _which_indices class attribute.

    Args:
        key: Index, slice, or tuple of indices/slices to access. Can be:
            - List of indices
            - Array of indices
            - Boolean mask array
            - Slice object
            - Integer index
    """
    # Convert key to indices array
    if isinstance(key, slice):
        # If we already have _which_indices, slice those instead of the full range
        if self._which_indices is not None:
            indices = self._which_indices[key]
        else:
            indices = np.arange(self._matrix.shape[0])[key]
    elif isinstance(key, (list, np.ndarray)):
        key = np.asarray(key)  # Convert list to array first
        if key.dtype == bool:
            # For boolean masks, convert to integer indices
            if self._which_indices is not None:
                # Apply mask to existing indices
                indices = self._which_indices[key]
            else:
                indices = np.where(key)[0]
        else:
            # For integer indices, index into existing _which_indices if it exists
            if self._which_indices is not None:
                indices = self._which_indices[key]
            else:
                indices = key
    elif isinstance(key, (int, np.integer)):
        # For scalar indices, convert to array
        if self._which_indices is not None:
            indices = np.array([self._which_indices[key]])
        else:
            indices = np.array([key])
    else:
        raise TypeError("Invalid key type. Use integer, list, slice, or numpy array.")

    self._which_indices = indices
    self._invalidate_selection_cache()

solve

solve(b: ndarray, method: str = 'direct', tol: float = 1e-05, callback: Optional[Callable] = None, initialization: Optional[ndarray] = None) -> np.ndarray

Solve the linear system Px = b.

Parameters:

Name	Type	Description	Default
b	ndarray	Right-hand side vector or matrix	required
method	str	Solver method ('direct' or 'pcg')	'direct'
tol	float	Tolerance for PCG solver	1e-05
callback	Optional[Callable]	Optional callback for PCG solver	None
initialization	Optional[ndarray]	Optional initial guess for pcg	None

Returns:

Type	Description
ndarray	Solution vector x

Source code in src/graphld/precision.py

def solve(self,
    b: np.ndarray,
    method: str = "direct",
    tol: float = 1e-5,
    callback: Optional[Callable] = None,
    initialization: Optional[np.ndarray] = None
) -> np.ndarray:
    """Solve the linear system Px = b.

    Args:
        b: Right-hand side vector or matrix
        method: Solver method ('direct' or 'pcg')
        tol: Tolerance for PCG solver
        callback: Optional callback for PCG solver
        initialization: Optional initial guess for pcg

    Returns:
        Solution vector x
    """
    original_shape = np.asarray(b).shape

    if method == "direct":
        y = self._expand_vector(b)
        if not self._factor_is_current():
            self.factor()
        solution = self._solver(y)
        if self._which_indices is not None:
            solution = solution[self._which_indices, :]
    elif method == "pcg":
        y = self._reshape_rhs(b)
        initialization = (
            self._reshape_rhs(initialization)
            if initialization is not None else None
        )
        if self._solver is None:
            self.factor()  # Some factorization is needed for use as a preconditioner
        solution = self._pcg(y, tol=tol, callback=callback, initialization=initialization)
    else:
        raise ValueError("Method must be either 'direct' or 'pcg'")

    return solution.reshape(original_shape)

solve_Lt

solve_Lt(b: ndarray) -> np.ndarray

Solve the triangular system L'x = b, where LL' = _matrix. If b ~ MVN(0, I), then x ~ MVN(0, R).

Parameters:

Name	Type	Description	Default
b	ndarray	Right-hand side vector	required

Returns:

Type	Description
ndarray	Solution vector x

Source code in src/graphld/precision.py

def solve_Lt(self, b: np.ndarray) -> np.ndarray:
    """Solve the triangular system L'x = b, where LL' = _matrix.
    If b ~ MVN(0, I), then x ~ MVN(0, R).

    Args:
        b: Right-hand side vector

    Returns:
        Solution vector x
    """
    if not self._factor_is_current():
        self.factor()

    y = self._expand_vector(b)
    solution = self._solver.solve_Lt(y, use_LDLt_decomposition=False)

    if self._which_indices is not None:
        solution = solution[self._which_indices, :]
    return solution.reshape(b.shape)

variant_solve

variant_solve(b: ndarray) -> np.ndarray

Computes correlation_matrix @ b where the dimension is number of variants, not number of indices. If two variants i,j have the same index (and are in perfect LD), b[i] and b[j] are summed, then solve() is called, then the solution vector is assigned the same values in entries i,j.

Parameters:

Name	Type	Description	Default
b	ndarray	Right-hand side vector	required

Returns:

Type	Description
ndarray	Solution vector z

Source code in src/graphld/precision.py

def variant_solve(self, b: np.ndarray) -> np.ndarray:
    """Computes correlation_matrix @ b where the dimension is number of variants, not number of indices.
    If two variants i,j have the same index (and are in perfect LD), b[i] and b[j] are summed, then
    solve() is called, then the solution vector is assigned the same values in entries i,j.

    Args:
        b: Right-hand side vector

    Returns:
        Solution vector z
    """
    if len(b) != len(self.variant_info):
        raise ValueError("b must have the same length as the number of variants")

    y = np.zeros(self.shape[0])
    np.add.at(y, self.variant_indices, b)

    y = self.solve(y)

    return y[self.variant_indices]

inverse_diagonal

inverse_diagonal(method: str = 'xdiag', initialization: Optional[Tuple[ndarray, ndarray]] = None, n_samples: int = 100, seed: Optional[int] = None) -> np.ndarray

Compute the diagonal elements of the inverse of the precision matrix.

Parameters:

Name	Type	Description	Default
initialization	Optional[Tuple[ndarray, ndarray]]	Optional tuple of matrices containing: - Initial guess for the diagonal elements - Initial guess for the off-diagonal elements	None
method	str	Method to use for computing diagonal elements ('exact', 'hutchinson', or 'xdiag')	'xdiag'
n_samples	int	Number of probe vectors for Hutchinson's method or xdiag	100
seed	Optional[int]	Random seed for generating probe vectors	None

Returns:

Type	Description
ndarray	Array of diagonal elements of the inverse

Source code in src/graphld/precision.py

def inverse_diagonal(
    self,
    method: str = "xdiag",
    initialization: Optional[Tuple[np.ndarray, np.ndarray]] = None,
    n_samples: int = 100,
    seed: Optional[int] = None,
) -> np.ndarray:
    """Compute the diagonal elements of the inverse of the precision matrix.

    Args:
        initialization: Optional tuple of matrices containing:
            - Initial guess for the diagonal elements
            - Initial guess for the off-diagonal elements
        method: Method to use for computing diagonal elements
            ('exact', 'hutchinson', or 'xdiag')
        n_samples: Number of probe vectors for Hutchinson's method or xdiag
        seed: Random seed for generating probe vectors

    Returns:
        Array of diagonal elements of the inverse
    """
    if method not in ["exact", "hutchinson", "xdiag"]:
        raise ValueError(f"Unknown method: {method}")

    # Slow, exact inversion
    if method.lower() == "exact":
        if initialization is not None:
            raise ValueError("Initialization not supported for exact method")

        dense_matrix = self._matrix.toarray()
        inv_matrix = np.linalg.inv(dense_matrix)
        diag = np.diag(inv_matrix)

        return diag if self._which_indices is None else diag[self._which_indices]

    # Use a stochastic estimator
    if initialization is not None:
        v, pv = initialization
    else:
        rng = np.random.RandomState(seed)

        # Rademacher random vectors
        v = rng.choice([-1.0, 1.0], size=(self.shape[0], min(self.shape[0], n_samples)))
        pv = v.copy()

    if method.lower() == "hutchinson":
        # Solve M * y = v for each probe vector
        y = self.solve(v, method="pcg", initialization=pv)

        # Estimate diagonal elements as average element-wise product of v_i * y_i
        diag_estimate = np.mean(v * y, axis=1)

    elif method.lower() == "xdiag":
        diag_estimate, y = self._xdiag_estimator(v, initialization=pv)

    return (diag_estimate, y) if initialization is not None else diag_estimate