| 211 | } |
| 212 | |
| 213 | void MBTR::getK3(py::array_t<double> &descriptor, py::array_t<double> &derivatives, bool return_descriptor, bool return_derivatives, const vector<int> &Z, const vector<vector<double>> &positions, const vector<vector<double>> &distances, const vector<vector<int>> &neighbours, const string &geomFunc, const string &weightFunc, const map<string, double> ¶meters, double min, double max, double sigma, int n) |
| 214 | { |
| 215 | // Create mutable and unchecked versions |
| 216 | auto descriptor_mu = descriptor.mutable_unchecked<1>(); |
| 217 | auto derivatives_mu = derivatives.mutable_unchecked<3>(); |
| 218 | |
| 219 | int nAtoms = Z.size(); |
| 220 | int nElem = this->atomicNumberToIndexMap.size(); |
| 221 | double dx = (max-min)/(n-1); |
| 222 | double sigmasqrt2 = sigma*sqrt(2.0); |
| 223 | double start = min-dx/2; |
| 224 | |
| 225 | for (int i = 0; i < nAtoms; ++i) { |
| 226 | |
| 227 | // For each atom we loop only over the atoms triplets that are |
| 228 | // within the neighbourhood |
| 229 | const vector<int> &i_neighbours = neighbours[i]; |
| 230 | for (const int &j : i_neighbours) { |
| 231 | const vector<int> &j_neighbours = neighbours[j]; |
| 232 | for (const int &k : j_neighbours) { |
| 233 | // Only consider triplets that have one atom in the original |
| 234 | // cell |
| 235 | if (i >= this->interactionLimit && j >= this->interactionLimit && k >= this->interactionLimit) { |
| 236 | continue; |
| 237 | } |
| 238 | // Calculate angle for all index permutations from choosing |
| 239 | // three out of nAtoms. The same atom cannot be present twice |
| 240 | // in the permutation. |
| 241 | if (j == i || k == j || k == i) { |
| 242 | continue; |
| 243 | } |
| 244 | // The angles are symmetric: ijk = kji. The value is |
| 245 | // calculated only for the triplet where k > i. |
| 246 | if (k <= i){ |
| 247 | continue; |
| 248 | } |
| 249 | |
| 250 | // Find distance vectors |
| 251 | vector<double> r_ji{ positions[j][0] - positions[i][0], |
| 252 | positions[j][1] - positions[i][1], |
| 253 | positions[j][2] - positions[i][2]}; |
| 254 | vector<double> r_ik{ positions[i][0] - positions[k][0], |
| 255 | positions[i][1] - positions[k][1], |
| 256 | positions[i][2] - positions[k][2]}; |
| 257 | vector<double> r_jk{ positions[j][0] - positions[k][0], |
| 258 | positions[j][1] - positions[k][1], |
| 259 | positions[j][2] - positions[k][2]}; |
| 260 | |
| 261 | // Distances |
| 262 | double d_ji = distances[j][i]; |
| 263 | double d_ik = distances[i][k]; |
| 264 | double d_jk = distances[j][k]; |
| 265 | |
| 266 | // Calculate geometry value and its derivatives. |
| 267 | // "angle" is not supported because it is not differentiable. |
| 268 | double geom; |
| 269 | vector<vector<double>> geom_d(3); |
| 270 | if (geomFunc == "cosine") { |