Calculators¶

def __init__(
    self,
    model: str | nn.Module = "aimnet2",
    nb_threshold: int = 120,
    needs_coulomb: bool | None = None,
    needs_dispersion: bool | None = None,
    device: str | None = None,
    compile_model: bool = False,
    compile_kwargs: dict | None = None,
    train: bool = False,
):
    # Device selection: use provided or auto-detect
    if device is not None:
        self.device = device
    else:
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
    self.external_coulomb: LRCoulomb | None = None
    self.external_dftd3: DFTD3 | None = None
    # Default cutoffs for LR modules
    self._default_dsf_cutoff = 15.0
    self._default_dftd3_cutoff = 15.0
    self._default_dftd3_smoothing = 0.2

    # Load model and get metadata
    metadata: dict | None = None
    if isinstance(model, str):
        p = get_model_path(model)
        self.model, metadata = load_model(p, device=self.device)
        self.cutoff = metadata["cutoff"]
    elif isinstance(model, nn.Module):
        self.model = model.to(self.device)
        self.cutoff = getattr(self.model, "cutoff", 5.0)
        metadata = getattr(self.model, "_metadata", None)
    else:
        raise TypeError("Invalid model type/name.")

    # Compile model if requested
    if compile_model:
        kwargs = compile_kwargs or {}
        self.model = torch.compile(self.model, **kwargs)

    # Resolve final flags (explicit overrides metadata)
    final_needs_coulomb = (
        needs_coulomb
        if needs_coulomb is not None
        else (metadata.get("needs_coulomb", False) if metadata is not None else False)
    )
    final_needs_dispersion = (
        needs_dispersion
        if needs_dispersion is not None
        else (metadata.get("needs_dispersion", False) if metadata is not None else False)
    )

    # Set up external Coulomb if needed
    if final_needs_coulomb:
        sr_embedded = metadata.get("coulomb_mode") == "sr_embedded" if metadata is not None else False
        # For PBC, user can switch to DSF/Ewald via set_lrcoulomb_method()
        # When sr_embedded=True: model has SRCoulomb which subtracts SR, so external
        # should compute FULL (subtract_sr=False) to give: (NN - SR) + FULL = NN + LR
        # When sr_embedded=False: model has no SR embedded, so external should compute
        # LR only (subtract_sr=True) to avoid double-counting
        self.external_coulomb = LRCoulomb(
            key_in="charges",
            key_out="energy",
            method="simple",
            rc=metadata.get("coulomb_sr_rc", 4.6) if metadata is not None else 4.6,
            envelope=metadata.get("coulomb_sr_envelope", "exp") if metadata is not None else "exp",
            subtract_sr=not sr_embedded,
        )
        self.external_coulomb = self.external_coulomb.to(self.device)

    # Set up external DFTD3 if needed
    if final_needs_dispersion:
        d3_params = metadata.get("d3_params") if metadata else None
        if d3_params is None:
            raise ValueError(
                "needs_dispersion=True but d3_params not found in metadata. "
                "Provide d3_params in model metadata or set needs_dispersion=False."
            )
        self.external_dftd3 = DFTD3(
            s8=d3_params["s8"],
            a1=d3_params["a1"],
            a2=d3_params["a2"],
            s6=d3_params.get("s6", 1.0),
        )
        self.external_dftd3 = self.external_dftd3.to(self.device)

    # Determine if model has long-range modules (embedded or external)
    has_embedded_lr = metadata.get("has_embedded_lr", False) if metadata is not None else False
    self.lr = (
        hasattr(self.model, "cutoff_lr")
        or self.external_coulomb is not None
        or self.external_dftd3 is not None
        or has_embedded_lr
    )
    # Set cutoff_lr based on model attribute or external modules
    if hasattr(self.model, "cutoff_lr"):
        self.cutoff_lr = getattr(self.model, "cutoff_lr", float("inf"))
    elif self.external_coulomb is not None:
        # For "simple" method, use inf (all pairs). For DSF, use dsf_cutoff.
        if self.external_coulomb.method == "simple":
            self.cutoff_lr = float("inf")
        else:
            self.cutoff_lr = self._default_dsf_cutoff
    elif self.external_dftd3 is not None:
        self.cutoff_lr = self._default_dftd3_cutoff
    elif has_embedded_lr:
        # Embedded LR modules (D3TS, SRCoulomb) need nbmat_lr
        self.cutoff_lr = self._default_dftd3_cutoff
    else:
        self.cutoff_lr = None
    self.nb_threshold = nb_threshold

    # Create adaptive neighbor list instances
    self._nblist = AdaptiveNeighborList(cutoff=self.cutoff)

    # Track separate cutoffs for LR modules
    self._coulomb_cutoff: float | None = None
    self._dftd3_cutoff: float = self._default_dftd3_cutoff
    if self.external_coulomb is not None:
        if self.external_coulomb.method == "simple":
            self._coulomb_cutoff = float("inf")
        elif self.external_coulomb.method == "ewald":
            self._coulomb_cutoff = None  # Ewald manages its own cutoff
        else:
            self._coulomb_cutoff = self.external_coulomb.dsf_rc
    if self.external_dftd3 is not None:
        self._dftd3_cutoff = self.external_dftd3.smoothing_off

    # Create long-range neighbor list(s) if LR modules present
    self._nblist_lr: AdaptiveNeighborList | None = None
    self._nblist_dftd3: AdaptiveNeighborList | None = None
    self._nblist_coulomb: AdaptiveNeighborList | None = None
    self._update_lr_nblists()

    # indicator if input was flattened
    self._batch = None
    self._max_mol_size: int = 0
    # placeholder for tensors that require grad
    self._saved_for_grad = {}
    # set flag of current Coulomb method
    self._coulomb_method: str | None = None
    if self.external_coulomb is not None:
        self._coulomb_method = self.external_coulomb.method
    elif self._has_embedded_coulomb():
        # Legacy models have embedded Coulomb with "simple" method
        self._coulomb_method = "simple"

    # Set training mode (default False for inference)
    self._train = train
    self.model.train(train)
    if not train:
        # Disable gradients on all parameters for inference mode
        for param in self.model.parameters():
            param.requires_grad_(False)
        if self.external_coulomb is not None:
            for param in self.external_coulomb.parameters():
                param.requires_grad_(False)
        if self.external_dftd3 is not None:
            for param in self.external_dftd3.parameters():
                param.requires_grad_(False)

def __init__(self, base_calc: AIMNet2Calculator | str = "aimnet2", charge=0, mult=1):
    super().__init__()
    if isinstance(base_calc, str):
        base_calc = AIMNet2Calculator(base_calc)
    self.base_calc = base_calc
    self.reset()
    self.charge = charge
    self.mult = mult
    self.update_tensors()
    # list of implemented species
    if hasattr(base_calc, "implemented_species"):
        self.implemented_species = base_calc.implemented_species.cpu().numpy()  # type: ignore
    else:
        self.implemented_species = None

def __init__(self, model: AIMNet2Calculator | str = "aimnet2", charge=0, mult=1, **kwargs):
    super().__init__(charge=charge, mult=mult, **kwargs)
    if isinstance(model, str):
        model = AIMNet2Calculator(model)
    self.model = model