Skip to content

cagpjax.policies.lanczos

Lanczos-based policies.

LanczosPolicy

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:

Name Type Description
n_actions int

Number of Lanczos vectors/actions to compute.
key PRNGKeyArray | None

Random key for reproducible Lanczos iterations.
Source code in src/cagpjax/policies/lanczos.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class LanczosPolicy(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: Number of Lanczos vectors/actions to compute.
        key: Random key for reproducible Lanczos iterations.
    """

    def __init__(self, n_actions: int, key: PRNGKeyArray | None = None):
        """Initialize the Lanczos policy.

        Args:
            n_actions: Number of Lanczos vectors to compute.
            key: Random key for initialization.
        """
        self._n_actions: int = n_actions
        self.key: PRNGKeyArray | None = key

    @property
    @override
    def n_actions(self) -> int:
        return self._n_actions

    @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 = cola.linalg.eig(
            cola.SelfAdjoint(A),
            self.n_actions,
            which="LM",
            alg=cola.linalg.Lanczos(key=self.key),
        )[1]
        if not isinstance(vecs, LinearOperator):
            vecs = Dense(vecs)
        return vecs

__init__(n_actions, key=None)

Initialize the Lanczos policy.

Parameters:

Name Type Description Default
n_actions int

Number of Lanczos vectors to compute.

 required
key PRNGKeyArray | None

Random key for initialization.

 None
Source code in src/cagpjax/policies/lanczos.py 
23
24
25
26
27
28
29
30
31
def __init__(self, n_actions: int, key: PRNGKeyArray | None = None):
    """Initialize the Lanczos policy.

    Args:
        n_actions: Number of Lanczos vectors to compute.
        key: Random key for initialization.
    """
    self._n_actions: int = n_actions
    self.key: PRNGKeyArray | None = key

to_actions(A)

Compute action matrix.

Parameters:

Name Type Description Default
A LinearOperator

Symmetric linear operator representing the linear system.

 required

Returns:

Type Description
LinearOperator

Linear operator containing the Lanczos vectors as columns.
Source code in src/cagpjax/policies/lanczos.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
@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 = cola.linalg.eig(
        cola.SelfAdjoint(A),
        self.n_actions,
        which="LM",
        alg=cola.linalg.Lanczos(key=self.key),
    )[1]
    if not isinstance(vecs, LinearOperator):
        vecs = Dense(vecs)
    return vecs