Skip to content

cagpjax.policies

Modules:

Classes:

AbstractBatchLinearSolverPolicy

Bases: AbstractLinearSolverPolicy

Abstract base class for policies that produce action matrices.

Methods:

  • to_actions –

    Compute all actions used to solve the linear system \(Ax=b\).

to_actions abstractmethod 

to_actions(A: LinearOperator, *, key: PRNGKeyArray | None = None) -> LinearOperator

Compute all actions used to solve the linear system \(Ax=b\).

For a matrix \(A\) with shape (n, n), the action matrix has shape (n, n_actions).

Parameters:

  • A

    (LinearOperator) –

    Linear operator representing the linear system.

  • key

    (PRNGKeyArray | None, default: None ) –

    Optional random key used by stochastic policies.

Returns:

Source code in src/cagpjax/policies/base.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
@abc.abstractmethod
def to_actions(
    self, A: LinearOperator, *, key: PRNGKeyArray | None = None
) -> LinearOperator:
    r"""Compute all actions used to solve the linear system $Ax=b$.

    For a matrix $A$ with shape ``(n, n)``, the action matrix has shape
    ``(n, n_actions)``.

    Args:
        A: Linear operator representing the linear system.
        key: Optional random key used by stochastic policies.

    Returns:
        Linear operator representing the action matrix.
    """
    ...

AbstractLinearSolverPolicy

Bases: Module

Abstract base class for all linear solver policies.

Policies define actions used to solve a linear system \(A x = b\), where \(A\) is a square linear operator.

BlockSparsePolicy

Bases: AbstractBatchLinearSolverPolicy

Block-sparse linear solver policy.

This policy uses a fixed block-diagonal sparse structure to define independent learnable actions. The matrix has the following structure:

\[ S = \begin{bmatrix} s_1 & 0 & \cdots & 0 \\ 0 & s_2 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & s_{\text{n_actions}} \end{bmatrix} \]

These are stacked and stored as a single trainable parameter nz_values.

Attributes:

  • n_actions (int) –

    Number of actions to use.

  • nz_values (Float[Array, N] | AbstractUnwrappable[Float[Array, N]]) –

    Non-zero values of the block-diagonal sparse matrix.

Methods:

  • from_random –

    Initialize policy from block-normalized random samples.

  • to_actions –

    Convert to block diagonal sparse action operators.

from_random classmethod 

from_random(key: PRNGKeyArray, num_datapoints: int, n_actions: int, *, sampler: Callable[[PRNGKeyArray, tuple[int, ...], Any], Float[Array, ' N']] = jax.random.normal, dtype: Any = None) -> BlockSparsePolicy

Initialize policy from block-normalized random samples.

Parameters:

  • key

    (PRNGKeyArray) –

    Random key used to sample initial values.

  • num_datapoints

    (int) –

    Number of rows in the resulting operator.

  • n_actions

    (int) –

    Number of action columns in the resulting operator.

  • sampler

    (Callable[[PRNGKeyArray, tuple[int, ...], Any], Float[Array, ' N']], default: normal ) –

    Callable with signature (key, shape, dtype) -> values.

  • dtype

    (Any, default: None ) –

    Optional dtype forwarded to sampler.

Source code in src/cagpjax/policies/block_sparse.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@classmethod
def from_random(
    cls,
    key: PRNGKeyArray,
    num_datapoints: int,
    n_actions: int,
    *,
    sampler: Callable[
        [PRNGKeyArray, tuple[int, ...], Any], Float[Array, " N"]
    ] = jax.random.normal,
    dtype: Any = None,
) -> "BlockSparsePolicy":
    """Initialize policy from block-normalized random samples.

    Args:
        key: Random key used to sample initial values.
        num_datapoints: Number of rows in the resulting operator.
        n_actions: Number of action columns in the resulting operator.
        sampler: Callable with signature ``(key, shape, dtype) -> values``.
        dtype: Optional dtype forwarded to ``sampler``.
    """
    if num_datapoints < 1:
        raise ValueError("num_datapoints must be at least 1")
    nz_values = sampler(key, (num_datapoints,), dtype)
    nz_values = _normalize_by_blocks(nz_values, n_actions)
    return cls(n_actions=n_actions, nz_values=nz_values)

to_actions 

to_actions(A: LinearOperator, *, key: PRNGKeyArray | None = None) -> LinearOperator

Convert to block diagonal sparse action operators.

Parameters:

  • A

    (LinearOperator) –

    Linear operator (unused).

  • key

    (PRNGKeyArray | None, default: None ) –

    Optional random key (unused).

Returns:

  • BlockDiagonalSparse ( LinearOperator ) –

    Sparse action structure representing the blocks.

Source code in src/cagpjax/policies/block_sparse.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
@override
def to_actions(
    self, A: LinearOperator, *, key: PRNGKeyArray | None = None
) -> LinearOperator:
    """Convert to block diagonal sparse action operators.

    Args:
        A: Linear operator (unused).
        key: Optional random key (unused).

    Returns:
        BlockDiagonalSparse: Sparse action structure representing the blocks.
    """
    return BlockDiagonalSparse(paramax.unwrap(self.nz_values), self.n_actions)

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:

  • n_actions (int) –

    Number of Lanczos vectors/actions to compute.

  • grad_rtol (float | None) –

    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:

to_actions 

to_actions(A: LinearOperator, *, key: PRNGKeyArray | None = None) -> LinearOperator

Compute action matrix.

Parameters:

  • A

    (LinearOperator) –

    Symmetric linear operator representing the linear system.

  • key

    (PRNGKeyArray | None, default: None ) –

    Random key used to initialize the Lanczos run.

Returns:

  • LinearOperator –

    Linear operator containing the Lanczos vectors as columns.

Source code in src/cagpjax/policies/lanczos.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@override
def to_actions(
    self, A: LinearOperator, *, key: PRNGKeyArray | None = None
) -> LinearOperator:
    """Compute action matrix.

    Args:
        A: Symmetric linear operator representing the linear system.
        key: Random key used to initialize the Lanczos run.

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

OrthogonalizationPolicy

Bases: AbstractBatchLinearSolverPolicy

Orthogonalization policy.

This policy orthogonalizes (if necessary) the action operator produced by the base policy.

Parameters:

  • base_policy

    –

    The base policy that produces the action operator to be orthogonalized.

  • method

    –

    The method to use for orthogonalization.

  • n_reortho

    –

    The number of times to re-orthogonalize each column. Reorthogonalizing once is generally sufficient to improve orthogonality for Gram-Schmidt variants (see e.g. 10.1007/s00211-005-0615-4).

PseudoInputPolicy 

PseudoInputPolicy(pseudo_inputs: Float[Array, 'M D'] | AbstractUnwrappable[Float[Array, 'M D']], train_inputs_or_dataset: Float[Array, 'N D'] | Dataset, kernel: AbstractKernel)

Bases: AbstractBatchLinearSolverPolicy

Pseudo-input linear solver policy.

This policy constructs actions from the cross-covariance between the training inputs and pseudo-inputs in the same input space. These pseudo-inputs are conceptually similar to inducing points and can be marked as trainable.

Parameters:

  • pseudo_inputs

    (Float[Array, 'M D'] | AbstractUnwrappable[Float[Array, 'M D']]) –

    Pseudo-inputs for the kernel.

  • train_inputs

    –

    Training inputs or a dataset containing training inputs. These must be the same inputs in the same order as the training data used to condition the CaGP model.

  • kernel

    (AbstractKernel) –

    Kernel for the GP prior. It must be able to take train_inputs and pseudo_inputs as arguments to its cross_covariance method.

Note

When training with many pseudo-inputs, it is common for the cross-covariance matrix to become poorly conditioned. Performance can be significantly improved by orthogonalizing the actions using an OrthogonalizationPolicy.

Source code in src/cagpjax/policies/pseudoinput.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def __init__(
    self,
    pseudo_inputs: Float[Array, "M D"]
    | paramax.AbstractUnwrappable[Float[Array, "M D"]],
    train_inputs_or_dataset: Float[Array, "N D"] | gpjax.dataset.Dataset,
    kernel: gpjax.kernels.AbstractKernel,
):
    self.pseudo_inputs = pseudo_inputs
    if isinstance(train_inputs_or_dataset, gpjax.dataset.Dataset):
        train_data = train_inputs_or_dataset
        if train_data.X is None:
            raise ValueError("Dataset must contain training inputs.")
        train_inputs = train_data.X
    else:
        train_inputs = train_inputs_or_dataset
    self.train_inputs = paramax.non_trainable(jnp.atleast_2d(train_inputs))
    self.kernel = kernel
    self.n_actions = paramax.unwrap(self.pseudo_inputs).shape[0]