Skip to content

cagpjax.policies.lanczos

Lanczos-based policies.

Classes:

  • LanczosPolicy –

    Lanczos-based policy for eigenvalue decomposition approximation.

LanczosPolicy 

LanczosPolicy(n_actions: int | None, key: PRNGKeyArray | None = None, grad_rtol: float | None = 0.0)

Bases: AbstractBatchLinearSolverPolicy

Lanczos-based policy for eigenvalue decomposition approximation.

This policy uses the Lanczos algorithm to compute the top n_actions eigenvectors of the linear operator \(A\).

Attributes:

  • n_actions (int) –

    Number of Lanczos vectors/actions to compute.

  • key (PRNGKeyArray | None) –

    Random key for reproducible Lanczos iterations.

Initialize the Lanczos policy.

Parameters:

  • n_actions

    (int | None) –

    Number of Lanczos vectors to compute.

  • key

    (PRNGKeyArray | None, default: None ) –

    Random key for initialization.

  • grad_rtol

    (float | None, default: 0.0 ) –

    Specifies the cutoff for similar eigenvalues, used to improve gradient computation for (almost-)degenerate matrices. If not provided, the default is 0.0. If None or negative, all eigenvalues are treated as distinct. (see cagpjax.linalg.eigh for more details)

Methods:

Source code in src/cagpjax/policies/lanczos.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def __init__(
    self,
    n_actions: int | None,
    key: PRNGKeyArray | None = None,
    grad_rtol: float | None = 0.0,
):
    """Initialize the Lanczos policy.

    Args:
        n_actions: Number of Lanczos vectors to compute.
        key: Random key for initialization.
        grad_rtol: Specifies the cutoff for similar eigenvalues, used to improve
            gradient computation for (almost-)degenerate matrices.
            If not provided, the default is 0.0.
            If None or negative, all eigenvalues are treated as distinct.
            (see [`cagpjax.linalg.eigh`][] for more details)
    """
    self._n_actions: int = n_actions
    self.key = key
    self.grad_rtol = grad_rtol

to_actions 

to_actions(A: LinearOperator) -> LinearOperator

Compute action matrix.

Parameters:

  • A

    (LinearOperator) –

    Symmetric linear operator representing the linear system.

Returns:

  • LinearOperator –

    Linear operator containing the Lanczos vectors as columns.

Source code in src/cagpjax/policies/lanczos.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
@override
def to_actions(self, A: LinearOperator) -> LinearOperator:
    """Compute action matrix.

    Args:
        A: Symmetric linear operator representing the linear system.

    Returns:
        Linear operator containing the Lanczos vectors as columns.
    """
    vecs = eigh(
        A, alg=Lanczos(self.n_actions, key=self.key), grad_rtol=self.grad_rtol
    ).eigenvectors
    return vecs