lgatr.primitives.config.PrimitivesConfig

class lgatr.primitives.config.PrimitivesConfig(subgroup=True, bivector=True, geometric_product=True, sparse_gp=True, sparse_linear=False)[source]

Bases: object

Symmetry-group and bilinear-layer toggles for an L-GATr model.

A PrimitivesConfig is passed to LGATr (and to the layers and primitive functions it contains) at construction time. Multiple models with different configs can coexist in the same process.

Parameters:
  • subgroup (bool) – If True, the model is only equivariant with respect to the connected subgroup of the Lorentz group, the proper orthochronous Lorentz group \(SO^+(1,3)\), which excludes parity and time reversal. This setting affects how the EquiLinear maps work: for \(SO^+(1,3)\) they additionally mix scalars with pseudoscalars, vectors with axialvectors, and among bivectors, effectively treating the pseudoscalar and axialvector representations like another scalar and vector. Defaults to True, because parity-odd representations are usually not important in high-energy physics simulations.

  • bivector (bool) – If False, the bivector components are set to zero after they are created in the GeometricBilinear layer. This is a toy switch to explore the effect of higher-order representations.

  • geometric_product (bool) – If False, the GeometricBilinear layer is replaced by a ScalarGatedNonlinearity followed by an EquiLinear layer. This is a toy switch to explore the effect of the geometric product.

  • sparse_gp (bool) – If True, route geometric_product() through the gather-and-reduce kernel that exploits the basis sparsity (the dense path otherwise spends most of its FLOPs on zero entries). Under torch.compile this is both faster and far lighter than the dense product.

  • sparse_linear (bool) – If True, route equi_linear() through the per-grade kernel that exploits the basis sparsity. This has fewer FLOPs than the dense path but uses less optimized kernels (no single fused BLAS GEMM), so on FLOP-rich GPUs (e.g. H100) it is typically slower and heavier than dense; it mainly helps on FLOP-bound hardware. Sparse outputs match the dense path within standard test tolerances but are not bit-identical.

bivector: bool = True
classmethod cast(config)[source]

Cast a PrimitivesConfig or mapping to a PrimitivesConfig.

Return type:

PrimitivesConfig

geometric_product: bool = True
property num_pin_linear_basis_elements: int

Number of equivariant linear basis elements (10 for the subgroup, 5 for full Lorentz).

sparse_gp: bool = True
sparse_linear: bool = False
subgroup: bool = True