EIP4NNPotentials/mol.py at master · shuix007/EIP4NNPotentials · GitHub

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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import dgl
import torch
import numpy as np

LARGEST_ATOMIC_NUMBER = 119
SI_POTENTIALS = [
    'erhartalbe',
    'edip',
    'meam_spline',
    'sw',
    'sw_silicene1',
    'Tersoff',
    'Tersoff_T3',
    'CF_MOD'
]
AL_POTENTIALS = [
    'WKG',
    'Zha',
    'SL',
    'ZJW_NIST',
    "ZM",
    'EA',
    'JSN',
    'GW_High',
    'GW_Low',
    'GW_Med'
]

def remove_duplicate_edges(edge_index):
    edge_dict = {}
    edges = []
    for i, j in zip(edge_index[0], edge_index[1]):
        if (i, j) not in edge_dict and (j, i) not in edge_dict:
            edges.append((i, j))
            edge_dict[(i, j)] = 1

    edges = np.array(edges).T
    return (edges[0], edges[1])

def neighbor_list_to_molecular_graph(edge_index, pos, atomic_numbers, image, is_contributing):
    bidirectional_edge_index = (
        torch.cat([edge_index[0], edge_index[1]]),
        torch.cat([edge_index[1], edge_index[0]])
    )

    source_pos = pos[bidirectional_edge_index[0]]
    target_pos = pos[bidirectional_edge_index[1]]

    distance = torch.linalg.norm(
        source_pos - target_pos, dim=1).float()  # euclidean distance

    graph = dgl.graph(bidirectional_edge_index, num_nodes=is_contributing.numel())
    graph.ndata['node_feats'] = atomic_numbers.unsqueeze(1)
    graph.ndata['is_contributing'] = is_contributing.unsqueeze(1)
    graph.ndata['image'] = image.unsqueeze(1)
    graph.edata['edge_feats'] = distance.unsqueeze(1)

    return graph

def neighbor_list_to_directed_molecular_graph(edge_index, pos, atomic_numbers, image, is_contributing):
    bidirectional_edge_index = (
        edge_index[1],
        edge_index[0]
    )

    source_pos = pos[bidirectional_edge_index[0]]
    target_pos = pos[bidirectional_edge_index[1]]

    distance = torch.linalg.norm(
        source_pos - target_pos, dim=1).float()  # euclidean distance

    graph = dgl.graph(bidirectional_edge_index, num_nodes=is_contributing.numel())
    graph.ndata['node_feats'] = atomic_numbers.unsqueeze(1)
    graph.ndata['is_contributing'] = is_contributing.unsqueeze(1)
    graph.ndata['image'] = image.unsqueeze(1)
    graph.edata['edge_feats'] = distance.unsqueeze(1)

    return graph