diff --git a/bindings/rascal/representations/spherical_expansion.py b/bindings/rascal/representations/spherical_expansion.py index 1a8af2b51..22c1a0e0c 100644 --- a/bindings/rascal/representations/spherical_expansion.py +++ b/bindings/rascal/representations/spherical_expansion.py @@ -136,6 +136,12 @@ class SphericalExpansion(BaseIO): should contain all the species present in the structure for which invariants will be computed + central_species : list + list of central atom species that are only computed + + neighbour_species : list + list of neighbour atom species that are only computed + compute_gradients : bool control the computation of the representation's gradients w.r.t. atomic positions. @@ -178,6 +184,8 @@ def __init__( optimization_args=None, expansion_by_species_method="environment wise", global_species=None, + central_species=None, + neighbour_species=None, compute_gradients=False, cutoff_function_parameters=dict(), ): @@ -195,11 +203,23 @@ class documentation elif not isinstance(global_species, list): global_species = list(global_species) + if neighbour_species is None: + neighbour_species = [] + elif not isinstance(neighbour_species, list): + neighbour_species = list(neighbour_species) + + if central_species is None: + central_species = [] + elif not isinstance(central_species, list): + central_species = list(central_species) + self.update_hyperparameters( max_radial=max_radial, max_angular=max_angular, expansion_by_species_method=expansion_by_species_method, global_species=global_species, + central_species=central_species, + neighbour_species=neighbour_species, compute_gradients=compute_gradients, ) self.cutoff_function_parameters = deepcopy(cutoff_function_parameters) @@ -258,6 +278,8 @@ def update_hyperparameters(self, **hypers): "cutoff_function_parameters", "expansion_by_species_method", "global_species", + "central_species", + "neighbour_species", } hypers_clean = {key: hypers[key] for key in hypers if key in allowed_keys} self.hypers.update(hypers_clean) @@ -310,6 +332,8 @@ def _get_init_params(self): max_angular=self.hypers["max_angular"], expansion_by_species_method=self.hypers["expansion_by_species_method"], global_species=self.hypers["global_species"], + central_species=self.hypers["central_species"], + neighbour_species=self.hypers["neighbour_species"], compute_gradients=self.hypers["compute_gradients"], gaussian_sigma_type=gaussian_density["type"], gaussian_sigma_constant=gaussian_density["gaussian_sigma"]["value"], diff --git a/src/rascal/representations/calculator_spherical_expansion.hh b/src/rascal/representations/calculator_spherical_expansion.hh index c6c15acc6..0763d1e12 100644 --- a/src/rascal/representations/calculator_spherical_expansion.hh +++ b/src/rascal/representations/calculator_spherical_expansion.hh @@ -214,6 +214,7 @@ namespace rascal { RadialContributionBase & operator=(RadialContributionBase && other) = default; + using Hypers_t = CalculatorBase::Hypers_t; using Matrix_t = Eigen::Matrix; @@ -350,6 +351,19 @@ namespace rascal { void set_hyperparameters(const Hypers_t & hypers) override { this->hypers = hypers; + if (hypers.count("neighbour_species")) { + auto species = hypers.at("neighbour_species").get>(); + for (const auto & sp : species) { + this->neighbour_species.insert({sp}); + } + } + if (hypers.count("central_species")) { + auto species = hypers.at("central_species").get>(); + for (const auto & sp : species) { + this->central_species.insert({sp}); + } + } + this->max_radial = hypers.at("max_radial"); this->max_angular = hypers.at("max_angular"); @@ -610,9 +624,13 @@ namespace rascal { void finalize_coefficients_der(Coeffs & coefficients_gradient, Center & center) const { for (auto neigh : center.pairs_with_self_pair()) { - auto & coefficients_neigh_gradient = coefficients_gradient[neigh]; - coefficients_neigh_gradient.template lhs_dot_der( - this->ortho_norm_matrix); + if (!this->neighbour_species.size() || + (this->neighbour_species.size() && + internal::is_element_in(neigh.get_atom_type(), this->neighbour_species))) { + auto & coefficients_neigh_gradient = coefficients_gradient[neigh]; + coefficients_neigh_gradient.template lhs_dot_der( + this->ortho_norm_matrix); + } } // for (neigh : center) } @@ -724,6 +742,8 @@ namespace rascal { //! combination of radial_ortho_matrix and radial_norm_factors //! not symmetric Matrix_t ortho_norm_matrix{}; + std::set> neighbour_species{}; + std::set> central_species{}; }; /** @@ -766,6 +786,19 @@ namespace rascal { void set_hyperparameters(const Hypers_t & hypers) override { this->hypers = hypers; + if (hypers.count("neighbour_species")) { + auto species = hypers.at("neighbour_species").get>(); + for (const auto & sp : species) { + this->neighbour_species.insert({sp}); + } + } + if (hypers.count("central_species")) { + auto species = hypers.at("central_species").get>(); + for (const auto & sp : species) { + this->central_species.insert({sp}); + } + } + this->max_radial = hypers.at("max_radial"); this->max_angular = hypers.at("max_angular"); @@ -930,6 +963,8 @@ namespace rascal { size_t max_radial{}; size_t max_angular{}; bool compute_gradients{}; + std::set> neighbour_species{}; + std::set> central_species{}; Vector_t legendre_radial_factor{}; Vector_t legendre_points{}; @@ -1633,6 +1668,18 @@ namespace rascal { this->expansion_by_species = "environment wise"; } + if (hypers.count("neighbour_species")) { + auto species = hypers.at("neighbour_species").get(); + for (const auto & sp : species) { + this->neighbour_species.insert({sp}); + } + } + if (hypers.count("central_species")) { + auto species = hypers.at("central_species").get(); + for (const auto & sp : species) { + this->central_species.insert({sp}); + } + } if (hypers.count("global_species")) { auto species = hypers.at("global_species").get(); for (const auto & sp : species) { @@ -1905,6 +1952,9 @@ namespace rascal { //! user defined species appearing in the expansion indexing std::set global_species{}; + std::set> neighbour_species{}; + std::set> central_species{}; + internal::AtomicSmearingType atomic_smearing_type{}; std::shared_ptr radial_integral{}; @@ -2155,10 +2205,10 @@ namespace rascal { constexpr bool ExcludeGhosts{true}; const bool is_not_masked{manager->is_not_masked()}; const bool compute_gradients{this->compute_gradients}; - if (not is_not_masked and compute_gradients) { - throw std::logic_error("Can't compute spherical expansion gradients with " - "masked center atoms"); - } + //if (not is_not_masked and compute_gradients) { + // throw std::logic_error("Can't compute spherical expansion gradients with " + // "masked center atoms"); + //} if (not is_not_masked and IsHalfNL) { std::stringstream err_str{}; err_str << "Half neighbor list should only be used when all the " @@ -2233,6 +2283,36 @@ namespace rascal { // coeff C^{ij}_{nlm} auto c_ij_nlm = math::Matrix_t(n_row, n_col); + // if central_species is used then we need to keep track of all envs that + // need to be computed. These are all envs with an atom of type in + // central_species + std::vector central_species_in_env_mask{}; + + if (this->central_species.size()) { + // the mask of the central atoms that are have a atom type in central_species + // or a neighbour of type in central_species + int max_atom_tag = 0; + for (auto center : manager) { + max_atom_tag = std::max(max_atom_tag, center.get_atom_tag()); + // it could be that some centers are mask + // so we need to also go through all neighbours + // by including center pairs we are including the central atom + for (auto neigh : center.pairs()) { + max_atom_tag = std::max(max_atom_tag, neigh.get_atom_tag()); + } + } + central_species_in_env_mask.resize(max_atom_tag+1, 0); + for (auto center : manager) { + if (internal::is_element_in(center.get_atom_type(), this->central_species)) { + central_species_in_env_mask.at(center.get_atom_tag()) = true; + } + for (auto neigh : center.pairs()) { + if (internal::is_element_in(neigh.get_atom_type(), this->central_species)) { + central_species_in_env_mask.at(center.get_atom_tag()) = true; + } + } + } + } for (auto center : manager) { // c^{i} @@ -2240,179 +2320,216 @@ namespace rascal { // \grad_i c^{i} auto & coefficients_center_gradient = expansions_coefficients_gradient[center.get_atom_ii()]; - auto atom_i_tag = center.get_atom_tag(); - Key_t center_type{center.get_atom_type()}; - - // Start the accumulation with the central atom contribution - coefficients_center[center_type].col(0) += - radial_integral->template compute_center_contribution( - center, center.get_atom_type()) / - sqrt(4.0 * PI); - - for (auto neigh : center.pairs()) { - auto atom_j = neigh.get_atom_j(); - const int atom_j_tag = atom_j.get_atom_tag(); - const bool is_center_atom{manager->is_center_atom(neigh)}; - - const double & dist{manager->get_distance(neigh)}; - const auto direction{manager->get_direction_vector(neigh)}; - Key_t neigh_type{neigh.get_atom_type()}; - - // the typical definition of the expansion coefficients involves - // (Y^m_l)*, but we compute everything with actual _real_ harmonics, - // so we just compute the real-valued Y^m_l (conjugate=false) - this->spherical_harmonics.calc(direction, compute_gradients, false); - auto && harmonics{spherical_harmonics.get_harmonics()}; - auto && harmonics_gradients{ - spherical_harmonics.get_harmonics_derivatives()}; - auto && neighbour_contribution = - radial_integral->template compute_neighbour_contribution( - dist, neigh, neigh.get_atom_type()); - double f_c{cutoff_function->f_c(dist)}; - auto coefficients_center_by_type{coefficients_center[neigh_type]}; - - // compute the coefficients - size_t l_block_idx{0}; - for (size_t angular_l{0}; angular_l < this->max_angular + 1; - ++angular_l) { - size_t l_block_size{2 * angular_l + 1}; - c_ij_nlm.block(0, l_block_idx, this->max_radial, l_block_size) = - neighbour_contribution.col(angular_l) * - harmonics.segment(l_block_idx, l_block_size); - l_block_idx += l_block_size; + if (!this->central_species.size() || + central_species_in_env_mask.at(center.get_atom_tag())) { + auto atom_i_tag = center.get_atom_tag(); + Key_t center_type{center.get_atom_type()}; + + // Start the accumulation with the central atom contribution + if ((!this->neighbour_species.size() || + (this->neighbour_species.size() && + internal::is_element_in(center_type, this->neighbour_species)))) { + coefficients_center[center_type].col(0) += + radial_integral->template compute_center_contribution( + center, center.get_atom_type()) / + sqrt(4.0 * PI); } - c_ij_nlm *= f_c; - coefficients_center_by_type += c_ij_nlm; - - // half list branch for c^{ji} terms using - // c^{ij}_{nlm} = (-1)^l c^{ji}_{nlm}. - if (IsHalfNL) { - if (is_center_atom) { - auto & coefficients_neigh{expansions_coefficients[atom_j]}; - auto coefficients_neigh_by_type{coefficients_neigh[center_type]}; - l_block_idx = 0; - double parity{1.}; - for (size_t angular_l{0}; angular_l < this->max_angular + 1; - ++angular_l) { - size_t l_block_size{2 * angular_l + 1}; - coefficients_neigh_by_type.block(0, l_block_idx, this->max_radial, - l_block_size) += - parity * c_ij_nlm.block(0, l_block_idx, this->max_radial, - l_block_size); - l_block_idx += l_block_size; - parity *= -1.; + for (auto neigh : center.pairs()) { + bool central_species_condition = + !this->central_species.size() || + internal::is_element_in(center.get_atom_type(), this->central_species) || + internal::is_element_in(neigh.get_atom_type(), this->central_species); + if (central_species_condition) { + bool neighbour_species_condition = + !this->neighbour_species.size() || + (this->neighbour_species.size() && + internal::is_element_in(neigh.get_atom_type(), + this->neighbour_species)); + if (neighbour_species_condition) { + auto atom_j = neigh.get_atom_j(); + const int atom_j_tag = atom_j.get_atom_tag(); + const bool is_center_atom{manager->is_center_atom(neigh)}; + + const double & dist{manager->get_distance(neigh)}; + const auto direction{manager->get_direction_vector(neigh)}; + Key_t neigh_type{neigh.get_atom_type()}; + + // the typical definition of the expansion coefficients involves + // (Y^m_l)*, but we compute everything with actual _real_ harmonics, + // so we just compute the real-valued Y^m_l (conjugate=false) + this->spherical_harmonics.calc(direction, compute_gradients, false); + auto && harmonics{spherical_harmonics.get_harmonics()}; + auto && harmonics_gradients{ + spherical_harmonics.get_harmonics_derivatives()}; + auto && neighbour_contribution = + radial_integral->template compute_neighbour_contribution( + dist, neigh, neigh.get_atom_type()); + double f_c{cutoff_function->f_c(dist)}; + auto coefficients_center_by_type{coefficients_center[neigh_type]}; + + // compute the coefficients + size_t l_block_idx{0}; + for (size_t angular_l{0}; angular_l < this->max_angular + 1; + ++angular_l) { + size_t l_block_size{2 * angular_l + 1}; + c_ij_nlm.block(0, l_block_idx, this->max_radial, l_block_size) = + neighbour_contribution.col(angular_l) * + harmonics.segment(l_block_idx, l_block_size); + l_block_idx += l_block_size; + } + c_ij_nlm *= f_c; + coefficients_center_by_type += c_ij_nlm; + + // half list branch for c^{ji} terms using + // c^{ij}_{nlm} = (-1)^l c^{ji}_{nlm}. + if (IsHalfNL) { + if (is_center_atom) { + auto & coefficients_neigh{expansions_coefficients[atom_j]}; + auto coefficients_neigh_by_type{coefficients_neigh[center_type]}; + l_block_idx = 0; + double parity{1.}; + for (size_t angular_l{0}; angular_l < this->max_angular + 1; + ++angular_l) { + size_t l_block_size{2 * angular_l + 1}; + coefficients_neigh_by_type.block(0, l_block_idx, this->max_radial, + l_block_size) += + parity * c_ij_nlm.block(0, l_block_idx, this->max_radial, + l_block_size); + l_block_idx += l_block_size; + parity *= -1.; + } + } + } + + // compute the gradients of the coefficients with respect to + // atoms positions + // but only if the neighbour is _not_ an image of the center! + // (the periodic images move with the center, so their contribution to + // the center gradient is zero) + + if (compute_gradients) { // NOLINT + // \grad_j c^i + auto & coefficients_neigh_gradient = + expansions_coefficients_gradient[neigh]; + + auto && neighbour_derivative = + radial_integral->compute_neighbour_derivative( + dist, neigh, neigh.get_atom_type()); + double df_c{cutoff_function->df_c(dist)}; + // The type of the contribution c^{ij} to the coefficient c^{i} + // depends on the type of j (and it is the same for the gradients) + // In the following atom i is of type a and atom j is of type b + + // grad_i c^{ib} + auto && gradient_center_by_type{ + coefficients_center_gradient[neigh_type]}; + // grad_j c^{ib} + auto && gradient_neigh_by_type{ + coefficients_neigh_gradient[neigh_type]}; + + // clang-format off + // d/dr_{ij} (c_{ij} f_c{r_{ij}}) + Matrix_t pair_gradient_contribution_p1 = + ((neighbour_derivative * f_c) + + (neighbour_contribution * df_c)); + // grad_j c^{ij} + Matrix_t pair_gradient_contribution{this->max_radial, + this->max_angular + 1}; + for (int cartesian_idx{0}; cartesian_idx < ThreeD; + ++cartesian_idx) { + l_block_idx = 0; + for (size_t angular_l{0}; angular_l < this->max_angular + 1; + ++angular_l) { + size_t l_block_size{2 * angular_l + 1}; + pair_gradient_contribution.resize(this->max_radial, l_block_size); + pair_gradient_contribution = + pair_gradient_contribution_p1.col(angular_l) + * harmonics.segment(l_block_idx, l_block_size) + * direction(cartesian_idx); + pair_gradient_contribution += + neighbour_contribution.col(angular_l) + * harmonics_gradients.block(cartesian_idx, l_block_idx, + 1, l_block_size) + * f_c / dist; + + // Each Cartesian gradient component occupies a contiguous block + // (row-major storage) + // grad_i c^{ib} = - \sum_{j} grad_j c^{ijb} + if (atom_j_tag != atom_i_tag) { + gradient_center_by_type.block( + cartesian_idx * max_radial, l_block_idx, + max_radial, l_block_size) -= pair_gradient_contribution; + } + // grad_j c^{ib} = grad_j c^{ijb} + gradient_neigh_by_type.block( + cartesian_idx * max_radial, l_block_idx, + max_radial, l_block_size) = pair_gradient_contribution; + l_block_idx += l_block_size; + // clang-format on + } // for (angular_l) + } // for cartesian_idx + + // half list branch for accumulating parts of grad_j c^{j} using + // grad_j c^{ji a} = (-1)^l grad_j c^{ij b} + if (IsHalfNL) { + if (is_center_atom) { + // grad_j c^{j} + auto & coefficients_neigh_center_gradient = + expansions_coefficients_gradient[neigh.get_atom_jj()]; + // grad_j c^{j a} + auto gradient_neigh_center_by_type = + coefficients_neigh_center_gradient[center_type]; + + for (int cartesian_idx{0}; cartesian_idx < ThreeD; + ++cartesian_idx) { + l_block_idx = 0; + double parity{1}; + for (size_t angular_l{0}; angular_l < this->max_angular + 1; + ++angular_l) { + size_t l_block_size{2 * angular_l + 1}; + // clang-format off + gradient_neigh_center_by_type.block( + cartesian_idx * max_radial, l_block_idx, + max_radial, l_block_size) += parity * + gradient_neigh_by_type.block( + cartesian_idx * max_radial, l_block_idx, + max_radial, l_block_size); + + l_block_idx += l_block_size; + parity *= -1.; + // clang-format on + } // for (angular_l) + } // for cartesian_idx + } // if (is_center_atom) + } // if (IsHalfNL) + } // if (compute_gradients) } } - } - - // compute the gradients of the coefficients with respect to - // atoms positions - // but only if the neighbour is _not_ an image of the center! - // (the periodic images move with the center, so their contribution to - // the center gradient is zero) - if (compute_gradients) { // NOLINT - // \grad_j c^i - auto & coefficients_neigh_gradient = - expansions_coefficients_gradient[neigh]; - - auto && neighbour_derivative = - radial_integral->compute_neighbour_derivative( - dist, neigh, neigh.get_atom_type()); - double df_c{cutoff_function->df_c(dist)}; - // The type of the contribution c^{ij} to the coefficient c^{i} - // depends on the type of j (and it is the same for the gradients) - // In the following atom i is of type a and atom j is of type b - - // grad_i c^{ib} - auto && gradient_center_by_type{ - coefficients_center_gradient[neigh_type]}; - // grad_j c^{ib} - auto && gradient_neigh_by_type{ - coefficients_neigh_gradient[neigh_type]}; - - // clang-format off - // d/dr_{ij} (c_{ij} f_c{r_{ij}}) - Matrix_t pair_gradient_contribution_p1 = - ((neighbour_derivative * f_c) - + (neighbour_contribution * df_c)); - // grad_j c^{ij} - Matrix_t pair_gradient_contribution{this->max_radial, - this->max_angular + 1}; - for (int cartesian_idx{0}; cartesian_idx < ThreeD; - ++cartesian_idx) { - l_block_idx = 0; - for (size_t angular_l{0}; angular_l < this->max_angular + 1; - ++angular_l) { - size_t l_block_size{2 * angular_l + 1}; - pair_gradient_contribution.resize(this->max_radial, l_block_size); - pair_gradient_contribution = - pair_gradient_contribution_p1.col(angular_l) - * harmonics.segment(l_block_idx, l_block_size) - * direction(cartesian_idx); - pair_gradient_contribution += - neighbour_contribution.col(angular_l) - * harmonics_gradients.block(cartesian_idx, l_block_idx, - 1, l_block_size) - * f_c / dist; - - // Each Cartesian gradient component occupies a contiguous block - // (row-major storage) - // grad_i c^{ib} = - \sum_{j} grad_j c^{ijb} - if (atom_j_tag != atom_i_tag) { - gradient_center_by_type.block( - cartesian_idx * max_radial, l_block_idx, - max_radial, l_block_size) -= pair_gradient_contribution; - } - // grad_j c^{ib} = grad_j c^{ijb} - gradient_neigh_by_type.block( - cartesian_idx * max_radial, l_block_idx, - max_radial, l_block_size) = pair_gradient_contribution; - l_block_idx += l_block_size; - // clang-format on - } // for (angular_l) - } // for cartesian_idx - - // half list branch for accumulating parts of grad_j c^{j} using - // grad_j c^{ji a} = (-1)^l grad_j c^{ij b} - if (IsHalfNL) { - if (is_center_atom) { - // grad_j c^{j} - auto & coefficients_neigh_center_gradient = - expansions_coefficients_gradient[neigh.get_atom_jj()]; - // grad_j c^{j a} - auto gradient_neigh_center_by_type = - coefficients_neigh_center_gradient[center_type]; - - for (int cartesian_idx{0}; cartesian_idx < ThreeD; - ++cartesian_idx) { - l_block_idx = 0; - double parity{1}; - for (size_t angular_l{0}; angular_l < this->max_angular + 1; - ++angular_l) { - size_t l_block_size{2 * angular_l + 1}; - // clang-format off - gradient_neigh_center_by_type.block( - cartesian_idx * max_radial, l_block_idx, - max_radial, l_block_size) += parity * - gradient_neigh_by_type.block( - cartesian_idx * max_radial, l_block_idx, - max_radial, l_block_size); - - l_block_idx += l_block_size; - parity *= -1.; - // clang-format on - } // for (angular_l) - } // for cartesian_idx - } // if (is_center_atom) - } // if (IsHalfNL) - } // if (compute_gradients) - } // for (neigh : center) + } // for (neigh : center) + } // if (central_species_in_env_mask) // Normalize and orthogonalize the radial coefficients - radial_integral->finalize_coefficients(coefficients_center); - if (compute_gradients) { - radial_integral->template finalize_coefficients_der( - expansions_coefficients_gradient, center); + if (!this->central_species.size() || + internal::is_element_in(center.get_atom_type(), this->central_species)) { + radial_integral->finalize_coefficients(coefficients_center); + if (compute_gradients) { + radial_integral->template finalize_coefficients_der( + expansions_coefficients_gradient, center); + } + } else { + // If central atom not in central species we set features + // and features center gradient to zero here. + // We needed to compute them for the neighbour + // gradients (or at least they are too entangled with neighbour + // gradient computation that they can be separated) + for (const auto & key : coefficients_center.get_keys()) { + coefficients_center[key].setZero(); + } + if (compute_gradients) { + for (const auto & key : coefficients_center_gradient.get_keys()) { + coefficients_center_gradient[key].setZero(); + } + } } } // for (center : manager) } // compute() @@ -2436,8 +2553,8 @@ namespace rascal { keys_list.emplace_back(); if (compute_gradients) { for (auto neigh : center.pairs_with_self_pair()) { - (void)neigh; // to avoid compiler warning - keys_list_grad.emplace_back(); + (void)neigh; // to avoid compiler warning + keys_list_grad.emplace_back(); } } } @@ -2447,12 +2564,12 @@ namespace rascal { Key_t center_type{center.get_atom_type()}; for (auto neigh : center.pairs()) { - keys_list[i_center].insert({neigh.get_atom_type()}); - if (manager->is_center_atom(neigh) and IsHalfNL) { - auto atom_j = neigh.get_atom_j(); - auto j_center = center_tag2idx[atom_j.get_atom_tag()]; - keys_list[j_center].insert(center_type); - } + keys_list[i_center].insert({neigh.get_atom_type()}); + if (manager->is_center_atom(neigh) and IsHalfNL) { + auto atom_j = neigh.get_atom_j(); + auto j_center = center_tag2idx[atom_j.get_atom_tag()]; + keys_list[j_center].insert(center_type); + } } keys_list[i_center].insert({center_type}); if (compute_gradients) { @@ -2460,9 +2577,9 @@ namespace rascal { keys_list[i_center].end()); i_grad++; for (auto neigh : center.pairs()) { - Key_t neigh_type{neigh.get_atom_type()}; - keys_list_grad[i_grad].insert(neigh_type); - i_grad++; + Key_t neigh_type{neigh.get_atom_type()}; + keys_list_grad[i_grad].insert(neigh_type); + i_grad++; } } // if (compute_gradients) i_center++; @@ -2485,6 +2602,28 @@ namespace rascal { Property_t & expansions_coefficients, PropertyGradient_t & expansions_coefficients_gradient) { std::set keys{}; + + //bool neigh_has_masked_central_atom{false}; + //for (auto center : manager) { + // if (manager.center_atoms_mask(center.get_atom_tag())) { // this might need the root structure manager + // add center + // for (auto neigh : center.pairs()) { + // neigh.emplace_back ... + // } + // } else { + // neigh_has_masked_central_atom = false; + // for (auto neigh : center.pairs()) { + // if (manager.center_atoms_mask(neigh.get_atom_tag())) { + // neigh_has_masked_central_atom = true; + // neigh.emplace_back ... + // } + // if (neigh_has_masked_central_atom) { + // add center + // } + // } + // manager.center_atoms_mask(id) + //} + for (auto center : manager) { Key_t center_type{center.get_atom_type()}; keys.insert({center_type}); @@ -2531,6 +2670,7 @@ namespace rascal { // check that all species in the structure are present in global_species Key_t keys{}; + for (auto center : manager) { typename Key_t::value_type center_type{center.get_atom_type()}; keys.push_back(center_type); @@ -2565,6 +2705,7 @@ namespace rascal { } } + expansions_coefficients.resize(keys_list); expansions_coefficients.setZero(); diff --git a/tests/python/python_binding_tests.py b/tests/python/python_binding_tests.py index 620344344..db1449576 100755 --- a/tests/python/python_binding_tests.py +++ b/tests/python/python_binding_tests.py @@ -14,6 +14,10 @@ TestSphericalExpansionRepresentation, TestSphericalInvariantsRepresentation, ) +from python_representation_calculator_mask_test import ( + TestSphericalExpansionMask +) + from python_models_test import TestNumericalKernelGradient, TestCosineKernel from python_math_test import TestMath from test_filter import FPSTest, CURTest diff --git a/tests/python/python_representation_calculator_mask_test.py b/tests/python/python_representation_calculator_mask_test.py new file mode 100644 index 000000000..409b27f85 --- /dev/null +++ b/tests/python/python_representation_calculator_mask_test.py @@ -0,0 +1,248 @@ +from rascal.representations import ( + SortedCoulombMatrix, + SphericalExpansion, + SphericalInvariants, +) +from test_utils import load_json_frame, BoxList, Box, dot +import unittest +import numpy as np +from scipy.spatial.distance import cdist +import sys +import os +import json +from copy import copy, deepcopy +from scipy.stats import ortho_group +import pickle + +rascal_reference_path = "reference_data" +inputs_path = os.path.join(rascal_reference_path, "inputs") +dump_path = os.path.join(rascal_reference_path, "tests_only") + + + +class TestSphericalExpansionMask(unittest.TestCase): + def setUp(self): + """ + builds the test case. Test the order=1 structure manager implementation + against a triclinic crystal. + """ + self.seed = 18012021 + + fns = [ + os.path.join(inputs_path, "CaCrP2O7_mvc-11955_symmetrized.json"), + os.path.join(inputs_path, "SiC_moissanite_supercell.json"), + os.path.join(inputs_path, "methane.json"), + ] + self.species = [1, 6, 8, 14, 15, 20, 24] + self.frames = [load_json_frame(fn) for fn in fns] + + self.max_radial = 6 + self.max_angular = 4 + + self.hypers = { + "interaction_cutoff": 6.0, + "cutoff_smooth_width": 1.0, + "max_radial": self.max_radial, + "max_angular": self.max_angular, + "gaussian_sigma_type": "Constant", + "gaussian_sigma_constant": 0.5, + } + + def test_neighbour_mask_all_species_features(self): + rep = SphericalExpansion(**self.hypers) + ref_features = rep.transform(self.frames).get_features(rep) + + hypers_neighbour_mask = deepcopy(self.hypers) + hypers_neighbour_mask["neighbour_species"] = self.species + rep = SphericalExpansion(**hypers_neighbour_mask) + totest_features = rep.transform(self.frames).get_features(rep) + + self.assertTrue(np.allclose(ref_features, totest_features)) + + def test_neighbour_mask_oxygen_features(self): + rep = SphericalExpansion(**self.hypers) + ref_features = rep.transform(self.frames).get_features_by_species(rep) + + hypers_neighbour_mask = deepcopy(self.hypers) + hypers_neighbour_mask["neighbour_species"] = [8] + rep = SphericalExpansion(**hypers_neighbour_mask) + totest_features = rep.transform(self.frames).get_features_by_species(rep) + + self.assertTrue(np.allclose(ref_features[(8,)], totest_features[(8,)])) + for species_key in totest_features.keys(): + if species_key != (8,): + self.assertTrue(np.allclose(np.zeros(ref_features[species_key].shape), totest_features[species_key])) + + def test_neighbour_mask_all_species_features_gradient(self): + hypers = deepcopy(self.hypers) + hypers["compute_gradients"] = True + + rep = SphericalExpansion(**hypers) + ref_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + hypers_neighbour_mask = deepcopy(hypers) + hypers_neighbour_mask["neighbour_species"] = self.species + rep = SphericalExpansion(**hypers_neighbour_mask) + totest_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + self.assertTrue(np.allclose(ref_features_gradient, totest_features_gradient)) + + def test_neighbour_mask_oxygen_features_gradient(self): + hypers = deepcopy(self.hypers) + hypers["compute_gradients"] = True + + rep = SphericalExpansion(**hypers) + feature_size_per_species = int(hypers["max_radial"] * (hypers["max_angular"]+1)**2) + ref_features = rep.transform(self.frames).get_features_by_species(rep) + ref_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + hypers_neighbour_mask = deepcopy(hypers) + hypers_neighbour_mask["neighbour_species"] = [8] + rep = SphericalExpansion(**hypers_neighbour_mask) + totest_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + for i, species_key in enumerate(ref_features.keys()): + feature_species_slice = slice(i*feature_size_per_species, (i+1)*feature_size_per_species) + if species_key == (8,): + self.assertTrue(np.allclose(totest_features_gradient[:, feature_species_slice], totest_features_gradient[:, feature_species_slice])) + else: + self.assertTrue(np.allclose(np.zeros(totest_features_gradient[:, feature_species_slice].shape), totest_features_gradient[:, feature_species_slice])) + + def test_central_mask_all_species_features(self): + rep = SphericalExpansion(**self.hypers) + ref_features = rep.transform(self.frames).get_features(rep) + + hypers_central_mask = deepcopy(self.hypers) + hypers_central_mask["central_species"] = self.species + rep = SphericalExpansion(**hypers_central_mask) + totest_features = rep.transform(self.frames).get_features(rep) + + self.assertTrue(np.allclose(ref_features, totest_features)) + + def test_central_mask_oxygen_features(self): + rep = SphericalExpansion(**self.hypers) + ref_features = rep.transform(self.frames).get_features(rep) + + hypers_central_mask = deepcopy(self.hypers) + hypers_central_mask["central_species"] = [8] + rep = SphericalExpansion(**hypers_central_mask) + totest_features = rep.transform(self.frames).get_features(rep) + + central_species = np.concatenate([frame["atom_types"] for frame in self.frames]).flatten() + central_oxygen_mask = central_species == 8 + + self.assertTrue(np.allclose(ref_features[central_oxygen_mask], totest_features[central_oxygen_mask])) + self.assertTrue(np.allclose(np.zeros(ref_features[~central_oxygen_mask].shape), totest_features[~central_oxygen_mask])) + + + def test_central_mask_all_species_features_gradient(self): + hypers = deepcopy(self.hypers) + hypers["compute_gradients"] = True + + rep = SphericalExpansion(**hypers) + ref_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + hypers_central_mask = deepcopy(hypers) + hypers_central_mask["central_species"] = self.species + rep = SphericalExpansion(**hypers_central_mask) + totest_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + self.assertTrue(np.allclose(ref_features_gradient, totest_features_gradient)) + + def test_central_mask_oxygen_features_gradient(self): + hypers = deepcopy(self.hypers) + hypers["compute_gradients"] = True + + rep = SphericalExpansion(**hypers) + manager = rep.transform(self.frames) + ref_features_gradient = manager.get_features_gradient(rep) + gradients_info = manager.get_gradients_info() + # chooses all pair gradients that have oxygen as center or neighbour + computed_gradients = np.logical_or(gradients_info[:, 3] == 8, gradients_info[:, 4] == 8) + computed_gradients = np.concatenate([[pair_mask]*3 for pair_mask in computed_gradients]) + + hypers_central_mask = deepcopy(hypers) + hypers_central_mask["central_species"] = [8] + rep = SphericalExpansion(**hypers_central_mask) + totest_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + self.assertTrue(np.allclose(ref_features_gradient[computed_gradients], totest_features_gradient[computed_gradients])) + self.assertTrue(np.allclose(np.zeros(ref_features_gradient[~computed_gradients].shape), totest_features_gradient[~computed_gradients])) + + + def test_neighbour_and_central_mask_all_species_features(self): + rep = SphericalExpansion(**self.hypers) + ref_features = rep.transform(self.frames).get_features(rep) + + hypers_mask = deepcopy(self.hypers) + hypers_mask["neighbour_species"] = self.species + hypers_mask["central_species"] = self.species + rep = SphericalExpansion(**hypers_mask) + totest_features = rep.transform(self.frames).get_features(rep) + + self.assertTrue(np.allclose(ref_features, totest_features)) + + def test_neighbour_and_central_mask_oxygen_features(self): + rep = SphericalExpansion(**self.hypers) + ref_features = rep.transform(self.frames).get_features_by_species(rep) + + hypers_mask = deepcopy(self.hypers) + hypers_mask["neighbour_species"] = [8] + hypers_mask["central_species"] = [8] + rep = SphericalExpansion(**hypers_mask) + totest_features = rep.transform(self.frames).get_features_by_species(rep) + + central_species = np.concatenate([frame["atom_types"] for frame in self.frames]).flatten() + central_oxygen_mask = central_species == 8 + + self.assertTrue(np.allclose(ref_features[(8,)][central_oxygen_mask], totest_features[(8,)][central_oxygen_mask])) + for species_key in totest_features.keys(): + if species_key != (8,): + self.assertTrue(np.allclose(np.zeros(ref_features[species_key].shape), totest_features[species_key])) + + self.assertTrue(np.allclose(np.zeros(ref_features[species_key][~central_oxygen_mask].shape), totest_features[species_key][~central_oxygen_mask])) + + def test_neighbour_and_central_mask_all_species_features_gradient(self): + hypers = deepcopy(self.hypers) + hypers["compute_gradients"] = True + + rep = SphericalExpansion(**hypers) + ref_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + hypers_mask = deepcopy(hypers) + hypers_mask["neighbour_species"] = self.species + hypers_mask["central_species"] = self.species + rep = SphericalExpansion(**hypers_mask) + totest_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + self.assertTrue(np.allclose(ref_features_gradient, totest_features_gradient)) + + def test_neighbour_central_mask_oxygen_features_gradient(self): + hypers = deepcopy(self.hypers) + hypers["compute_gradients"] = True + + rep = SphericalExpansion(**hypers) + feature_size_per_species = int(hypers["max_radial"] * (hypers["max_angular"]+1)**2) + manager = rep.transform(self.frames) + ref_features = manager.get_features_by_species(rep) + ref_features_gradient = manager.get_features_gradient(rep) + + gradients_info = manager.get_gradients_info() + # chooses all pair gradients that have oxygen as center or neighbour + computed_gradients = np.logical_or(gradients_info[:, 3] == 8, gradients_info[:, 4] == 8) + computed_gradients = np.concatenate([[pair_mask]*3 for pair_mask in computed_gradients]) + + hypers_mask = deepcopy(hypers) + hypers_mask["neighbour_species"] = [8] + hypers_mask["central_species"] = [8] + rep = SphericalExpansion(**hypers_mask) + totest_features_gradient = rep.transform(self.frames).get_features_gradient(rep) + + self.assertTrue(np.allclose(np.zeros(ref_features_gradient[~computed_gradients].shape), totest_features_gradient[~computed_gradients])) + + for i, species_key in enumerate(ref_features.keys()): + feature_species_slice = slice(i*feature_size_per_species, (i+1)*feature_size_per_species) + if species_key == (8,): + self.assertTrue(np.allclose(totest_features_gradient[computed_gradients][:, feature_species_slice], totest_features_gradient[computed_gradients][:, feature_species_slice])) + else: + self.assertTrue(np.allclose(np.zeros(totest_features_gradient[computed_gradients][:, feature_species_slice].shape), totest_features_gradient[computed_gradients][:, feature_species_slice]))