TensorSketch

class franken.rf.heads.TensorSketch(input_dim, num_random_features=2**10, num_species=None, chemically_informed_ratio=None, degree=2, bias=0.0, length_scale=1.0, rng_seed=None)

Bases: RandomFeaturesHead

TensorSketch random-features approximation of the polynomial kernel

\[\left(\frac{\langle x, y \rangle}{\ell^{2}} + b\right)^{d}\]

from “Fast and scalable polynomial kernels via explicit feature maps” by Pham and Pagh.

Parameters:

  • input_dim (int) – The dimensionality of the input features.

  • num_random_features (int) – The number of random features to use in the sketching process. Defaults to \(2^{10} = 1024\).

  • num_species (int | None) – The number of chemical species for which the kernel is computed. This parameter is relevant for systems with multiple chemical species. Defaults to None.

  • chemically_informed_ratio (float | None) – The relative weight of chemically-informed kernels with respect to the all-species kernel. Ignored if num_species is None. Defaults to None.

  • degree (int) – Degree \(d\) of the polynomial kernel to sketch. Defaults to 2.

  • bias (float) – The bias term \(b\) added to the kernel function. This allows for shifting the kernel values, which can be useful for certain applications. Defaults to 0.0.

  • length_scale (float) – The length scale parameter \(\ell\) that controls the smoothness of the kernel function. It affects how quickly the kernel values decay with distance. Defaults to 1.0.

  • rng_seed (int | None) – A seed for the random number generator used in generating random features. Setting this ensures reproducibility of results. Defaults to None.

feature_map(h, atomic_numbers=None, batch_ids=None)

Computes the random-feature map for a given configuration h

Parameters:

  • h (torch.Tensor) – descriptors for a single configuration ~[natoms, descriptors]

  • atomic_numbers (torch.Tensor) – atomic numbers for a single configuration ~[natoms]