def k1_geoms_and_weights(self, system):
"""Calculate the atom count for each element.
Args:
system (System): The atomic system.
Returns:
1D ndarray: The counts for each element in a list where the index
of atomic number x is self.atomic_number_to_index[x]
"""
if self._k1_geoms is None or self._k1_weights is None:
cmbtr = MBTRWrapper(system.get_positions(),
system.get_atomic_numbers(),
self.atomic_number_to_index,
interaction_limit=self._interaction_limit,
is_local=self._is_local)
# For k=1, the geometry function is given by the atomic number, and
# the weighting function is unity by default.
parameters = {}
self._k1_geoms, self._k1_weights = cmbtr.get_k1_geoms_and_weights(
geom_func=b"atomic_number",
weight_func=b"unity",
parameters=parameters)
return self._k1_geoms, self._k1_weights
def _get_k1(self, system):
"""Calculates the second order terms where the scalar mapping is the
inverse distance between atoms.
Returns:
1D ndarray: flattened K2 values.
"""
grid = self.k1["grid"]
start = grid["min"]
stop = grid["max"]
n = grid["n"]
sigma = grid["sigma"]
# Determine the geometry function
geom_func_name = self.k1["geometry"]["function"]
cmbtr = MBTRWrapper(self.atomic_number_to_index,
interaction_limit=self._interaction_limit,
indices=np.zeros((len(system), 3), dtype=int))
k1_map = cmbtr.get_k1(
system.get_atomic_numbers(),
geom_func=geom_func_name.encode(),
weight_func=b"unity",
parameters={},
start=start,
stop=stop,
sigma=sigma,
n=n,
)
# Depending on flattening, use either a sparse matrix or a dense one.
n_elem = self.n_elements
if self.flatten:
k1 = lil_matrix((1, n_elem * n), dtype=np.float32)
else:
k1 = np.zeros((n_elem, n), dtype=np.float32)
for key, gaussian_sum in k1_map.items():
i = key[0]
# Denormalize if requested
if not self.normalize_gaussians:
max_val = 1 / (sigma * math.sqrt(2 * math.pi))
gaussian_sum /= max_val
if self.flatten:
start = i * n
end = (i + 1) * n
k1[0, start:end] = gaussian_sum
else:
k1[i, :] = gaussian_sum
return k1
def k3_geoms_and_weights(self, system, cell_indices):
"""Calculates the value of the geometry function and corresponding
weights for unique three-body combinations.
Args:
system (System): The atomic system.
Returns:
tuple: (geoms, weights) Cosines of the angles (values between -1
and 1) in the form {(i,j,k): [list of angles] }. The weights
corresponding to the angles are stored in a similar dictionary.
"""
if self._k3_geoms is None or self._k2_weights is None:
# Calculate the angles with the C++ implementation
cmbtr = MBTRWrapper(system.get_positions(),
system.get_atomic_numbers(),
self.atomic_number_to_index,
interaction_limit=self._interaction_limit,
indices=cell_indices,
is_local=self._is_local)
# Determine the weighting function
weight_info = self.k3.get("weighting")
parameters = {}
if weight_info is not None:
weighting_function = weight_info["function"]
if weighting_function == "exponential" or weighting_function == "exp":
parameters = {
b"scale": weight_info["scale"],
b"cutoff": weight_info["cutoff"]
}
else:
weighting_function = "unity"
# Determine the geometry function
geom_func_name = self.k3["geometry"]["function"]
self._k3_geoms, self._k3_weights = cmbtr.get_k3_geoms_and_weights(
geom_func=geom_func_name.encode(),
weight_func=weighting_function.encode(),
parameters=parameters)
return self._k3_geoms, self._k3_weights
def k2_geoms_and_weights(self, system, cell_indices):
"""Calculates the value of the geometry function and corresponding
weights for unique two-body combinations.
Args:
system (System): The atomic system.
Returns:
dict: Inverse distances in the form: {(i, j): [list of angles] }.
The dictionaries are filled so that the entry for pair i and j is
in the entry where j>=i.
"""
if self._k2_geoms is None or self._k2_weights is None:
cmbtr = MBTRWrapper(system.get_positions(),
system.get_atomic_numbers(),
self.atomic_number_to_index,
interaction_limit=self._interaction_limit,
indices=cell_indices,
is_local=self._is_local)
# Determine the weighting function
weighting = self.k2.get("weighting")
parameters = {}
if weighting is not None:
weighting_function = weighting["function"]
if weighting_function == "exponential" or weighting_function == "exp":
parameters = {
b"scale": weighting["scale"],
b"cutoff": weighting["cutoff"]
}
else:
weighting_function = "unity"
# Determine the geometry function
geom_func_name = self.k2["geometry"]["function"]
self._k2_geoms, self._k2_weights = cmbtr.get_k2_geoms_and_weights(
geom_func=geom_func_name.encode(),
weight_func=weighting_function.encode(),
parameters=parameters)
return self._k2_geoms, self._k2_weights
def _get_k3(self, system):
"""Calculates the third order terms.
Returns:
1D ndarray: flattened K3 values.
"""
grid = self.k3["grid"]
start = grid["min"]
stop = grid["max"]
n = grid["n"]
sigma = grid["sigma"]
# Determine the weighting function and possible radial cutoff
radial_cutoff = None
weighting = self.k3.get("weighting")
parameters = {}
if weighting is not None:
weighting_function = weighting["function"]
if weighting_function == "exponential" or weighting_function == "exp":
scale = weighting["scale"]
cutoff = weighting["cutoff"]
if scale != 0:
radial_cutoff = -0.5 * math.log(cutoff) / scale
parameters = {b"scale": scale, b"cutoff": cutoff}
else:
weighting_function = "unity"
# Determine the geometry function
geom_func_name = self.k3["geometry"]["function"]
# If needed, create the extended system
if self.periodic:
centers = system.get_positions()
ext_system, cell_indices = self.create_extended_system(
system, centers, radial_cutoff)
else:
ext_system = system
cell_indices = np.zeros((len(system), 3), dtype=int)
cmbtr = MBTRWrapper(self.atomic_number_to_index,
interaction_limit=self._interaction_limit,
indices=cell_indices)
# If radial cutoff is finite, use it to calculate the sparse
# distance matrix to reduce computational complexity from O(n^2) to
# O(n log(n))
n_atoms = len(ext_system)
if radial_cutoff is not None:
dmat = ext_system.get_distance_matrix_within_radius(radial_cutoff)
adj_list = dscribe.utils.geometry.get_adjacency_list(dmat)
dmat_dense = np.full(
(n_atoms, n_atoms), sys.float_info.max
) # The non-neighbor values are treated as "infinitely far".
dmat_dense[dmat.col, dmat.row] = dmat.data
# If no weighting is used, the full distance matrix is calculated
else:
dmat_dense = ext_system.get_distance_matrix()
adj_list = np.tile(np.arange(n_atoms), (n_atoms, 1))
k3_map = cmbtr.get_k3(
ext_system.get_atomic_numbers(),
distances=dmat_dense,
neighbours=adj_list,
geom_func=geom_func_name.encode(),
weight_func=weighting_function.encode(),
parameters=parameters,
start=start,
stop=stop,
sigma=sigma,
n=n,
)
# Depending of flattening, use either a sparse matrix or a dense one.
n_elem = self.n_elements
if self.flatten:
k3 = lil_matrix((1, int(n_elem * n_elem * (n_elem + 1) / 2 * n)),
dtype=np.float32)
else:
k3 = np.zeros((n_elem, n_elem, n_elem, n), dtype=np.float32)
for key, gaussian_sum in k3_map.items():
i = key[0]
j = key[1]
k = key[2]
# This is the index of the spectrum. It is given by enumerating the
# elements of a three-dimensional array where for valid elements
# k>=i. The enumeration begins from [0, 0, 0], and ends at [n_elem,
# n_elem, n_elem], looping the elements in the order j, i, k.
m = int(j * n_elem * (n_elem + 1) / 2 + k + i * n_elem - i *
(i + 1) / 2)
# Denormalize if requested
if not self.normalize_gaussians:
max_val = 1 / (sigma * math.sqrt(2 * math.pi))
gaussian_sum /= max_val
if self.flatten:
start = m * n
end = (m + 1) * n
k3[0, start:end] = gaussian_sum
else:
k3[i, j, k, :] = gaussian_sum
return k3
def _get_k3(self, system, new_system, indices):
"""Calculates the second order terms where the scalar mapping is the
inverse distance between atoms.
Returns:
1D ndarray: flattened K2 values.
"""
grid = self.k3["grid"]
start = grid["min"]
stop = grid["max"]
n = grid["n"]
sigma = grid["sigma"]
# Determine the weighting function and possible radial cutoff
radial_cutoff = None
weighting = self.k3.get("weighting")
parameters = {}
if weighting is not None:
weighting_function = weighting["function"]
if weighting_function == "exponential" or weighting_function == "exp":
scale = weighting["scale"]
cutoff = weighting["cutoff"]
if scale != 0:
radial_cutoff = -0.5 * math.log(cutoff) / scale
parameters = {b"scale": scale, b"cutoff": cutoff}
else:
weighting_function = "unity"
# Determine the geometry function
geom_func_name = self.k3["geometry"]["function"]
# Calculate extended system
if self.periodic:
centers_new = new_system.get_positions()
centers_existing = system.get_positions()[indices]
centers = np.concatenate((centers_new, centers_existing), axis=0)
ext_system, cell_indices = dscribe.utils.geometry.get_extended_system(
system,
radial_cutoff,
centers,
return_cell_indices=True,
)
ext_system = System.from_atoms(ext_system)
else:
ext_system = system
cell_indices = np.zeros((len(system), 3), dtype=int)
cmbtr = MBTRWrapper(self.atomic_number_to_index,
interaction_limit=self._interaction_limit,
indices=cell_indices)
# If radial cutoff is finite, use it to calculate the sparse
# distance matrix to reduce computational complexity from O(n^2) to
# O(n log(n))
fin_system = ext_system + new_system
n_atoms_ext = len(ext_system)
n_atoms_fin = len(fin_system)
n_atoms_new = len(new_system)
ext_pos = ext_system.get_positions()
new_pos = new_system.get_positions()
if radial_cutoff is not None:
# Calculate distance within the extended system
dmat_ext_to_ext = ext_system.get_distance_matrix_within_radius(
radial_cutoff, pos=ext_pos)
col = dmat_ext_to_ext.col
row = dmat_ext_to_ext.row
data = dmat_ext_to_ext.data
dmat = scipy.sparse.coo_matrix((data, (row, col)),
shape=(n_atoms_fin, n_atoms_fin))
# Calculate the distances from the new positions to atoms in the
# extended system using the cutoff
if len(new_pos) != 0:
dmat_ext_to_new = ext_system.get_distance_matrix_within_radius(
radial_cutoff, pos=new_pos)
col = dmat_ext_to_new.col
row = dmat_ext_to_new.row
data = dmat_ext_to_new.data
dmat.col = np.append(dmat.col, col + n_atoms_ext)
dmat.row = np.append(dmat.row, row)
dmat.data = np.append(dmat.data, data)
dmat.col = np.append(dmat.col, row)
dmat.row = np.append(dmat.row, col + n_atoms_ext)
dmat.data = np.append(dmat.data, data)
# Calculate adjacencies and transform to the dense matrix for
# sending information to C++
adj_list = dscribe.utils.geometry.get_adjacency_list(dmat)
dmat_dense = np.full(
(n_atoms_fin, n_atoms_fin), sys.float_info.max
) # The non-neighbor values are treated as "infinitely far".
dmat_dense[dmat.row, dmat.col] = dmat.data
# If no weighting is used, the full distance matrix is calculated
else:
dmat = scipy.sparse.lil_matrix((n_atoms_fin, n_atoms_fin))
# Fill in block for extended system
dmat_ext_to_ext = ext_system.get_distance_matrix()
dmat[0:n_atoms_ext, 0:n_atoms_ext] = dmat_ext_to_ext
# Fill in block for extended system to new system
dmat_ext_to_new = scipy.spatial.distance.cdist(ext_pos, new_pos)
dmat[0:n_atoms_ext,
n_atoms_ext:n_atoms_ext + n_atoms_new] = dmat_ext_to_new
dmat[n_atoms_ext:n_atoms_ext + n_atoms_new,
0:n_atoms_ext] = dmat_ext_to_new.T
# Calculate adjacencies and the dense version
dmat = dmat.tocoo()
adj_list = dscribe.utils.geometry.get_adjacency_list(dmat)
dmat_dense = np.full(
(n_atoms_fin, n_atoms_fin), sys.float_info.max
) # The non-neighbor values are treated as "infinitely far".
dmat_dense[dmat.row, dmat.col] = dmat.data
# Form new indices that include the existing atoms and the newly added
# ones
indices = np.append(indices,
[n_atoms_ext + i for i in range(n_atoms_new)])
k3_list = cmbtr.get_k3_local(
Z=fin_system.get_atomic_numbers(),
distances=dmat_dense,
neighbours=adj_list,
indices=indices,
geom_func=geom_func_name.encode(),
weight_func=weighting_function.encode(),
parameters=parameters,
start=start,
stop=stop,
sigma=sigma,
n=n,
)
# Depending on flattening, use either a sparse matrix or a dense one.
n_elem = self.n_elements
n_loc = len(indices)
if self.flatten:
k3 = lil_matrix((n_loc, int((n_elem * (3 * n_elem - 1) * n / 2))),
dtype=np.float32)
for i_loc, k3_map in enumerate(k3_list):
for key, gaussian_sum in k3_map.items():
i = key[0]
j = key[1]
k = key[2]
# This is the index of the spectrum. It is given by enumerating the
# elements of a three-dimensional array and only considering
# elements for which k>=i and i || j == 0. The enumeration begins
# from [0, 0, 0], and ends at [n_elem, n_elem, n_elem], looping the
# elements in the order k, i, j.
if j == 0:
m = k + i * n_elem - i * (i + 1) / 2
else:
m = n_elem * (n_elem + 1) / 2 + (j - 1) * n_elem + k
start = int(m * n)
end = int((m + 1) * n)
# Denormalize if requested
if not self.normalize_gaussians:
max_val = 1 / (sigma * math.sqrt(2 * math.pi))
gaussian_sum /= max_val
k3[i_loc, start:end] = gaussian_sum
else:
k3 = np.zeros((n_loc, n_elem, n_elem, n_elem, n), dtype=np.float32)
for i_loc, k3_map in enumerate(k3_list):
for key, gaussian_sum in k3_map.items():
i = key[0]
j = key[1]
k = key[2]
# Denormalize if requested
if not self.normalize_gaussians:
max_val = 1 / (sigma * math.sqrt(2 * math.pi))
gaussian_sum /= max_val
k3[i_loc, i, j, k, :] = gaussian_sum
return k3
def _get_k2(self, system, new_system, indices):
"""Calculates the second order terms where the scalar mapping is the
inverse distance between atoms.
Returns:
1D ndarray: flattened K2 values.
"""
grid = self.k2["grid"]
start = grid["min"]
stop = grid["max"]
n = grid["n"]
sigma = grid["sigma"]
# Determine the weighting function and possible radial cutoff
radial_cutoff = None
weighting = self.k2.get("weighting")
parameters = {}
if weighting is not None:
weighting_function = weighting["function"]
if weighting_function == "exponential" or weighting_function == "exp":
scale = weighting["scale"]
cutoff = weighting["cutoff"]
if scale != 0:
radial_cutoff = -math.log(cutoff) / scale
parameters = {
b"scale": weighting["scale"],
b"cutoff": weighting["cutoff"]
}
else:
weighting_function = "unity"
# Determine the geometry function
geom_func_name = self.k2["geometry"]["function"]
# Calculate extended system
if self.periodic:
centers = new_system.get_positions()
ext_system, cell_indices = dscribe.utils.geometry.get_extended_system(
system,
radial_cutoff,
centers,
return_cell_indices=True,
)
ext_system = System.from_atoms(ext_system)
else:
ext_system = system
cell_indices = np.zeros((len(system), 3), dtype=int)
cmbtr = MBTRWrapper(self.atomic_number_to_index,
interaction_limit=self._interaction_limit,
indices=cell_indices)
# If radial cutoff is finite, use it to calculate the sparse distance
# matrix to reduce computational complexity from O(n^2) to O(n log(n)).
# If radial cutoff is not available, calculate full matrix.
n_atoms_ext = len(ext_system)
n_atoms_new = len(new_system)
ext_pos = ext_system.get_positions()
new_pos = new_system.get_positions()
if radial_cutoff is not None:
dmat = new_system.get_distance_matrix_within_radius(radial_cutoff,
pos=ext_pos)
adj_list = dscribe.utils.geometry.get_adjacency_list(dmat)
dmat_dense = np.full(
(n_atoms_new, n_atoms_ext), sys.float_info.max
) # The non-neighbor values are treated as "infinitely far".
dmat_dense[dmat.row, dmat.col] = dmat.data
else:
dmat_dense = scipy.spatial.distance.cdist(new_pos, ext_pos)
adj_list = np.tile(np.arange(n_atoms_ext), (n_atoms_new, 1))
# Form new indices that include the existing atoms and the newly added
# ones
indices = np.append(
indices,
[n_atoms_ext + i for i in range(n_atoms_new - len(indices))])
k2_list = cmbtr.get_k2_local(
indices=indices,
Z=ext_system.get_atomic_numbers(),
distances=dmat_dense,
neighbours=adj_list,
geom_func=geom_func_name.encode(),
weight_func=weighting_function.encode(),
parameters=parameters,
start=start,
stop=stop,
sigma=sigma,
n=n,
)
# Depending on flattening, use either a sparse matrix or a dense one.
n_elem = self.n_elements
n_loc = len(indices)
if self.flatten:
k2 = lil_matrix((n_loc, n_elem * n), dtype=np.float32)
for i_loc, k2_map in enumerate(k2_list):
for key, gaussian_sum in k2_map.items():
i = key[1]
m = i
start = int(m * n)
end = int((m + 1) * n)
# Denormalize if requested
if not self.normalize_gaussians:
max_val = 1 / (sigma * math.sqrt(2 * math.pi))
gaussian_sum /= max_val
k2[i_loc, start:end] = gaussian_sum
else:
k2 = np.zeros((n_loc, n_elem, n), dtype=np.float32)
for i_loc, k2_map in enumerate(k2_list):
for key, gaussian_sum in k2_map.items():
i = key[1]
# Denormalize if requested
if not self.normalize_gaussians:
max_val = 1 / (sigma * math.sqrt(2 * math.pi))
gaussian_sum /= max_val
k2[i_loc, i, :] = gaussian_sum
return k2
