enerzyme.models.physnet.core.PhysNetCore

class enerzyme.models.physnet.core.PhysNetCore(dim_embedding: int, num_rbf: int, num_blocks: int = 3, num_residual_atomic: int = 2, num_residual_interaction: int = 2, num_residual_output: int = 1, activation_fn: Literal['shifted_softplus', 'swish'] = 'shifted_softplus', activation_params: Dict[Literal['dim_feature', 'initial_alpha', 'initial_beta', 'learnable'], int | float | bool] = {}, dropout_rate: float = 0.0, shallow_ensemble_size: int = 1)

Bases: BaseFFCore

__init__(dim_embedding: int, num_rbf: int, num_blocks: int = 3, num_residual_atomic: int = 2, num_residual_interaction: int = 2, num_residual_output: int = 1, activation_fn: Literal['shifted_softplus', 'swish'] = 'shifted_softplus', activation_params: Dict[Literal['dim_feature', 'initial_alpha', 'initial_beta', 'learnable'], int | float | bool] = {}, dropout_rate: float = 0.0, shallow_ensemble_size: int = 1) → None

build(built_layers: List[Module]) → None

get_output(rbf: Tensor, atom_embedding: Tensor, idx_i_sr: Tensor, idx_j_sr: Tensor) → Dict[str, Tensor]

Compute raw atomic properties