def _get_attribute_lookup(atom_features: str = "cgcnn"):
"""Build a lookup array indexed by atomic number."""
max_z = max(v["Z"] for v in chem_data.values())
# get feature shape (referencing Carbon)
template = get_node_attributes("C", atom_features)
features = np.zeros((1 + max_z, len(template)))
for element, v in chem_data.items():
z = v["Z"]
x = get_node_attributes(element, atom_features)
if x is not None:
features[z, :] = x
return features
def atom_dgl_multigraph(
atoms=None,
neighbor_strategy="k-nearest",
cutoff=8.0,
max_neighbors=12,
atom_features="cgcnn",
max_attempts=3,
id: Optional[str] = None,
compute_line_graph: bool = True,
use_canonize: bool = False,
):
"""Obtain a DGLGraph for Atoms object."""
if neighbor_strategy == "k-nearest":
edges = nearest_neighbor_edges(
atoms=atoms,
cutoff=cutoff,
max_neighbors=max_neighbors,
id=id,
use_canonize=use_canonize,
)
else:
raise ValueError("Not implemented yet", neighbor_strategy)
# elif neighbor_strategy == "voronoi":
# edges = voronoi_edges(structure)
u, v, r = build_undirected_edgedata(atoms, edges)
# build up atom attribute tensor
sps_features = []
for ii, s in enumerate(atoms.elements):
feat = list(get_node_attributes(s, atom_features=atom_features))
# if include_prdf_angles:
# feat=feat+list(prdf[ii])+list(adf[ii])
sps_features.append(feat)
sps_features = np.array(sps_features)
node_features = torch.tensor(sps_features).type(
torch.get_default_dtype())
g = dgl.graph((u, v))
g.ndata["atom_features"] = node_features
g.edata["r"] = r
if compute_line_graph:
# construct atomistic line graph
# (nodes are bonds, edges are bond pairs)
# and add bond angle cosines as edge features
lg = g.line_graph(shared=True)
lg.apply_edges(compute_bond_cosines)
return g, lg
else:
return g
def atom_dgl_multigraph(
atoms=None,
cutoff=8.0,
max_neighbors=12,
atom_features="cgcnn",
enforce_undirected=False,
max_attempts=3,
include_prdf_angles=False,
partial_rcut=4.0,
id=None,
):
"""Obtain a DGLGraph for Atoms object."""
dists = atoms.raw_distance_matrix
def cos_formula(a, b, c):
"""Get angle between three edges for oblique triangles."""
res = (a**2 + b**2 - c**2) / (2 * a * b)
res = -1.0 if res < -1.0 else res
res = 1.0 if res > 1.0 else res
return np.arccos(res)
def bond_to_bond_feats(nb):
tmp = 0
angles_tmp = []
for ii, i in enumerate(nb):
tmp = ii + 1
if tmp > len(nb) - 1:
tmp = 0
ang = 0
try:
ang = cos_formula(i[2], nb[tmp][2], dists[i[1],
nb[tmp][1]])
except Exception as exp:
# print("Setting angle zeros", id, exp)
pass
angles_tmp.append(ang)
return np.array(angles_tmp)
if include_prdf_angles:
(
all_neighbors,
prdf_arr,
pangle_arr,
pval,
aval,
nbor,
) = atoms.atomwise_angle_and_radial_distribution(r=cutoff)
pval = np.fliplr(np.sort(pval))[:, 0:max_neighbors]
aval = np.fliplr(np.sort(aval))[:, 0:max_neighbors]
else:
all_neighbors = atoms.get_all_neighbors(r=cutoff)
# if a site has too few neighbors, increase the cutoff radius
min_nbrs = min(len(neighborlist) for neighborlist in all_neighbors)
# print('min_nbrs,max_neighbors=',min_nbrs,max_neighbors)
attempt = 0
while min_nbrs < max_neighbors:
print("extending cutoff radius!", attempt, cutoff, id)
lat = atoms.lattice
r_cut = max(cutoff, lat.a, lat.b, lat.c)
attempt += 1
if attempt >= max_attempts:
atoms = atoms.make_supercell([2, 2, 2])
print(
"Making supercell, exceeded,attempts",
max_attempts,
"cutoff",
r_cut,
id,
)
cutoff = r_cut
all_neighbors = atoms.get_all_neighbors(r=cutoff)
min_nbrs = min(len(neighborlist) for neighborlist in all_neighbors)
# return Graph.atom_dgl_multigraph(
# atoms, r_cut, max_neighbors, atom_features
# )
# build up edge list
# Currently there's no guarantee that this creates undirected graphs
# An undirected solution would build the full edge list where nodes are
# keyed by (index,image), and ensure each edge has a complementary edge
# indeed,JVASP-59628 is an example of a calculation where this produces
# a graph where one site has no incident edges!
# build an edge dictionary u -> v
# so later we can run through the dictionary
# and remove all pairs of edges
# so what's left is the odd ones out
edges = defaultdict(list)
u, v, r, w, prdf, adf = [], [], [], [], [], []
for site_idx, neighborlist in enumerate(all_neighbors):
# sort on distance
neighborlist = sorted(neighborlist, key=lambda x: x[2])
ids = np.array([nbr[1] for nbr in neighborlist])
distances = np.array([nbr[2] for nbr in neighborlist])
c = np.array([nbr[3] for nbr in neighborlist])
# find the distance to the k-th nearest neighbor
max_dist = distances[max_neighbors - 1]
# keep all edges out to the neighbor shell of the k-th neighbor
ids = ids[distances <= max_dist]
new_angles = bond_to_bond_feats(neighborlist)
try:
new_angles = new_angles[ids - 1]
except Exception as exp:
new_angles = np.zeros(len(ids))
pass
c = c[distances <= max_dist]
distances = distances[distances <= max_dist]
u.append([site_idx] * len(ids))
v.append(ids)
r.append(distances)
w.append(new_angles)
if include_prdf_angles:
prdf.append(pval[site_idx])
adf.append(aval[site_idx])
# keep track of cell-resolved edges
# to enforce undirected graph construction
for dst, cell_id in zip(ids, c):
u_key = f"{site_idx}-(0.0, 0.0, 0.0)"
v_key = f"{dst}-{tuple(cell_id)}"
edge_key = tuple(sorted((u_key, v_key)))
edges[edge_key].append((site_idx, dst))
if enforce_undirected:
# add complementary edges to unpaired edges
for edge_pair in edges.values():
if len(edge_pair) == 1:
src, dst = edge_pair[0]
u.append(dst) # swap the order!
v.append(src)
r.append(atoms.raw_distance_matrix[src, dst])
u = np.hstack(u)
v = np.hstack(v)
r = np.hstack(r)
w = np.hstack(w)
u = torch.tensor(u)
v = torch.tensor(v)
w = torch.tensor(w)
if include_prdf_angles:
prdf = np.array(prdf)
adf = np.array(adf)
prdf = np.hstack(prdf)
adf = np.cos(np.hstack(adf))
if len(r) != len(prdf):
prdf = np.append(prdf, np.zeros(len(r) - len(prdf)))
if len(r) != len(adf):
adf = np.append(adf, np.zeros(len(r) - len(adf)))
prdf = torch.tensor(np.array(np.hstack(prdf))).type(
torch.get_default_dtype())
adf = torch.tensor(np.array(np.hstack(adf))).type(
torch.get_default_dtype())
r = torch.tensor(np.array(r)).type(torch.get_default_dtype())
w = torch.tensor(np.array(w)).type(torch.get_default_dtype())
# build up atom attribute tensor
species = atoms.elements
sps_features = []
for ii, s in enumerate(species):
feat = list(get_node_attributes(s, atom_features=atom_features))
# if include_prdf_angles:
# feat=feat+list(prdf[ii])+list(adf[ii])
sps_features.append(feat)
sps_features = np.array(sps_features)
node_features = torch.tensor(sps_features).type(
torch.get_default_dtype())
g = dgl.graph((u, v))
g.ndata["atom_features"] = node_features
g.edata["bondlength"] = r
g.edata["bondangle"] = w
if include_prdf_angles:
g.edata["partial_distance"] = prdf
g.edata["partial_angle"] = adf
return g
def atom_dgl_multigraph(
atoms=None,
cutoff=8.0,
max_neighbors=12,
atom_features="cgcnn",
enforce_undirected=False,
max_attempts=3,
id=None,
):
"""Obtain a DGLGraph for Atoms object."""
all_neighbors = atoms.get_all_neighbors(r=cutoff)
# if a site has too few neighbors, increase the cutoff radius
min_nbrs = min(len(neighborlist) for neighborlist in all_neighbors)
# print('min_nbrs,max_neighbors=',min_nbrs,max_neighbors)
attempt = 0
while min_nbrs < max_neighbors:
print("extending cutoff radius!", attempt, cutoff, id)
lat = atoms.lattice
r_cut = max(cutoff, lat.a, lat.b, lat.c)
attempt += 1
if attempt >= max_attempts:
atoms = atoms.make_supercell([2, 2, 2])
print(
"Making supercell, exceeded,attempts",
max_attempts,
"cutoff",
r_cut,
id,
)
cutoff = r_cut
all_neighbors = atoms.get_all_neighbors(r=cutoff)
min_nbrs = min(len(neighborlist) for neighborlist in all_neighbors)
# return Graph.atom_dgl_multigraph(
# atoms, r_cut, max_neighbors, atom_features
# )
# build up edge list
# Currently there's no guarantee that this creates undirected graphs
# An undirected solution would build the full edge list where nodes are
# keyed by (index,image), and ensure each edge has a complementary edge
# indeed,JVASP-59628 is an example of a calculation where this produces
# a graph where one site has no incident edges!
# build an edge dictionary u -> v
# so later we can run through the dictionary
# and remove all pairs of edges
# so what's left is the odd ones out
edges = defaultdict(list)
u, v, r = [], [], []
for site_idx, neighborlist in enumerate(all_neighbors):
# sort on distance
neighborlist = sorted(neighborlist, key=lambda x: x[2])
ids = np.array([nbr[1] for nbr in neighborlist])
distances = np.array([nbr[2] for nbr in neighborlist])
c = np.array([nbr[3] for nbr in neighborlist])
# find the distance to the k-th nearest neighbor
max_dist = distances[max_neighbors - 1]
# keep all edges out to the neighbor shell of the k-th neighbor
ids = ids[distances <= max_dist]
c = c[distances <= max_dist]
distances = distances[distances <= max_dist]
u.append([site_idx] * len(ids))
v.append(ids)
r.append(distances)
# keep track of cell-resolved edges
# to enforce undirected graph construction
for dst, cell_id in zip(ids, c):
u_key = f"{site_idx}-(0.0, 0.0, 0.0)"
v_key = f"{dst}-{tuple(cell_id)}"
edge_key = tuple(sorted((u_key, v_key)))
edges[edge_key].append((site_idx, dst))
if enforce_undirected:
# add complementary edges to unpaired edges
for edge_pair in edges.values():
if len(edge_pair) == 1:
src, dst = edge_pair[0]
u.append(dst) # swap the order!
v.append(src)
r.append(atoms.raw_distance_matrix[src, dst])
u = torch.tensor(np.hstack(u))
v = torch.tensor(np.hstack(v))
r = torch.tensor(np.hstack(r)).type(torch.get_default_dtype())
# build up atom attribute tensor
species = atoms.elements
node_features = torch.tensor([
get_node_attributes(s, atom_features=atom_features)
for s in species
]).type(torch.get_default_dtype())
g = dgl.graph((u, v))
g.ndata["atom_features"] = node_features
g.edata["bondlength"] = r
return g