File size: 1,397 Bytes
fecdc11 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | """Graph helpers shared by the OneScience eSEN backbones."""
from __future__ import annotations
from types import SimpleNamespace
import torch
from onescience.modules.func_utils.uma_graph.compute import generate_graph
class GraphModelMixin:
"""Provide FairChem-v1 compatible graph output using OneScience graph code."""
def generate_graph(self, data, cutoff=None, max_neighbors=None, **kwargs):
cutoff = cutoff or self.cutoff
max_neighbors = max_neighbors or self.max_neighbors
pbc = kwargs.pop("pbc", None)
if pbc is None:
pbc = getattr(data, "pbc", None)
if pbc is None:
pbc = torch.ones(
(data.natoms.numel(), 3), dtype=torch.bool, device=data.pos.device
)
elif pbc.ndim == 1:
pbc = pbc.view(1, 3).expand(data.natoms.numel(), -1)
graph = generate_graph(
data,
cutoff=cutoff,
max_neighbors=max_neighbors,
enforce_max_neighbors_strictly=getattr(
self, "enforce_max_neighbors_strictly", False
),
radius_pbc_version=getattr(self, "radius_pbc_version", 1),
pbc=pbc,
)
graph.update(
batch_full=data.batch,
atomic_numbers_full=data.atomic_numbers,
node_offset=0,
)
return SimpleNamespace(**graph)
|