enerzyme.models.spookynet.core.SpookyNetCore

class enerzyme.models.spookynet.core.SpookyNetCore(dim_embedding: int, num_rbf: int, num_modules: int, num_residual_pre: int, num_residual_local_x: int, num_residual_local_s: int, num_residual_local_p: int, num_residual_local_d: int, num_residual_local: int, num_residual_nonlocal_q: int, num_residual_nonlocal_k: int, num_residual_nonlocal_v: int, num_residual_post: int, num_residual_output: int, activation_fn: Literal['shifted_softplus', 'swish'], use_irreps: bool, dropout_rate: float = 0.0, shallow_ensemble_size: int = 1)

Bases: BaseFFCore

__init__(dim_embedding: int, num_rbf: int, num_modules: int, num_residual_pre: int, num_residual_local_x: int, num_residual_local_s: int, num_residual_local_p: int, num_residual_local_d: int, num_residual_local: int, num_residual_nonlocal_q: int, num_residual_nonlocal_k: int, num_residual_nonlocal_v: int, num_residual_post: int, num_residual_output: int, activation_fn: Literal['shifted_softplus', 'swish'], use_irreps: bool, dropout_rate: float = 0.0, shallow_ensemble_size: int = 1) → None

build(built_layers: List[Module]) → None

get_output(Dij_sr: Tensor, vij_sr: Tensor, idx_i_sr: Tensor, idx_j_sr: Tensor, rbf: Tensor, atom_embedding: Tensor, batch_seg: Tensor | None = None) → Dict[Literal['Ea', 'Qa'], Tensor]