aenet.torch_featurize.ChebyshevDescriptor

class aenet.torch_featurize.ChebyshevDescriptor(*args: Any, **kwargs: Any)[source]

Complete featurization pipeline using Chebyshev descriptors.

Implements the typespin architecture from aenet’s Fortran code: - Single radial and angular basis functions - Two feature sets: unweighted + typespin-weighted - Typespin coefficients centered around zero

This exactly matches the behavior of aenet’s Fortran implementation.

Notes

  • Under PBC, neighbor displacements are reconstructed using wrapped fractional positions and integer image offsets provided by the ghost PBC backend of TorchNeighborList. Specifically: positions_frac = remainder(positions @ inv(cell), 1.0) r_ij = ((positions_frac[j] + offsets) @ cell)-(positions_frac[i] @ cell) This aligns featurization with the neighbor list semantics and ensures PBC vs explicit supercell consistency.

  • The neighbor cutoff used internally for neighbor list construction is max(rad_cutoff, ang_cutoff) to ensure complete coverage for both basis sets.

__init__(species: List[str], rad_order: int, rad_cutoff: float, ang_order: int, ang_cutoff: float, min_cutoff: float = 0.55, device: str = 'cpu', dtype: torch.dtype = torch.float64, profile_timing: bool = False)[source]

Initialize Chebyshev descriptor.

Parameters:

  • species – List of atomic species (e.g., [‘O’, ‘H’])

  • rad_order – Maximum radial Chebyshev order

  • rad_cutoff – Radial cutoff radius (Angstroms)

  • ang_order – Maximum angular Chebyshev order

  • ang_cutoff – Angular cutoff radius (Angstroms)

  • min_cutoff – Minimum distance cutoff (Angstroms)

  • device – ‘cpu’ or ‘cuda’

  • dtype – torch.float64 for double precision

  • profile_timing – If True, track separate timings for neighbor list construction vs feature computation

Methods

__init__(species, rad_order, rad_cutoff, ...)

Initialize Chebyshev descriptor.

compute_angular_features(positions, ...)

Compute angular features from pre-computed neighbor information.

compute_feature_gradients(positions, species)

Compute features and their gradients w.r.t.

compute_feature_gradients_from_neighbor_info(...)

Compute features and their gradients using precomputed neighbor info.

compute_feature_gradients_with_graph(...[, ...])

Vectorized features and gradients using CSR neighbors and triplets.

compute_features_and_local_derivatives_with_graph(...)

Compute features plus sparse local derivative blocks from graph data.

compute_forces_from_energy(positions, ...[, ...])

Compute atomic forces from an energy model.

compute_radial_features(positions, ...)

Compute radial features from pre-computed neighbor information.

featurize_with_neighbor_info(positions, species)

Featurize structure and extract neighbor information for force training.

forward(positions, species[, ...])

Featurize atomic structure.

forward_from_positions(positions, species[, ...])

Convenience method that computes neighbors then calls forward().

forward_with_graph(positions, ...[, ...])

Vectorized forward using CSR neighbor graph and optional triplets.

get_n_features()

Get number of features (same for all species).

get_timing_stats()

Get timing statistics for neighbor list vs feature computation.

reset_timing_stats()

Reset timing statistics to zero.