Atomic density formula for multiple elements works now

mala-project · Dec 16, 2024 · 9dd9a34 · 9dd9a34
1 parent f84dd75
commit 9dd9a34
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 25 deletions.
diff --git a/external_modules/total_energy_module/total_energy.f90 b/external_modules/total_energy_module/total_energy.f90
@@ -664,12 +664,9 @@ SUBROUTINE set_positions_gauss(verbose, gaussian_descriptors,reference_gaussian_
   CALL start_clock( 'structure_factors' )
   ALLOCATE(rgd_of_g(ngm,1), rhon(ngm))
   DO isp = 1, nsp
-    PRINT *, SHAPE(gaussian_descriptors(:, isp:isp)), SHAPE(strf(:,isp:isp))
      CALL rho_r2g(dfftp, gaussian_descriptors(:, isp:isp), strf(:,isp:isp))
   ENDDO
 
-!  CALL rho_r2g(dfftp, gaussian_descriptors, strf)
-!
   CALL rho_r2g(dfftp, reference_gaussian_descriptors, rgd_of_g)
 
   DO isp = 1, nsp
@@ -868,10 +865,6 @@ SUBROUTINE set_rho_of_r(rho_of_r,nnr_in,nspin_in)
   !
   CALL mp_sum( deband, intra_bgrp_comm )
   !
-  DO isp = 1, nsp
-     print *, strf(1:10,isp)
-  ENDDO
-
 
   RETURN
 

diff --git a/mala/descriptors/atomic_density.py b/mala/descriptors/atomic_density.py
@@ -8,7 +8,7 @@
 import numpy as np
 from scipy.spatial import distance
 
-from mala.common.parallelizer import printout
+from mala.common.parallelizer import printout, parallel_warn
 from mala.descriptors.lammps_utils import extract_compute_np
 from mala.descriptors.descriptor import Descriptor
 
@@ -215,18 +215,24 @@ def __calculate_lammps(self, outdir, **kwargs):
         )
 
         self._clean_calculation(lmp, keep_logs)
-        gaussian_descriptors_np_test = gaussian_descriptors_np.copy()
-
-        if (
-            isinstance(self.parameters.minterpy_lp_norm, list)
-            and np.shape(gaussian_descriptors_np_test)[1] > 7
-        ):
-            for index in range(len(self.parameters.minterpy_lp_norm)):
-                gaussian_descriptors_np_test[:, 6 + index] = (
-                    gaussian_descriptors_np[
-                        :, self.parameters.minterpy_lp_norm[index] + 6
-                    ]
-                )
+
+        if len(set(self._atoms.numbers)) > 1:
+            parallel_warn(
+                "Atomic density formula and multielement system detected: "
+                "Quantum ESPRESSO has a different internal order "
+                "for structure factors and atomic positions. "
+                "MALA recovers the correct ordering for multielement "
+                "systems, but the algorithm to do so is still "
+                "experimental. Please test on a small system before "
+                "using the atomic density formula at scale."
+            )
+            symbols, indices = np.unique(
+                [atom.symbol for atom in self._atoms], return_index=True
+            )
+            permutation = np.concatenate(
+                ([0, 1, 2, 3, 4, 5], np.argsort(indices) + 6)
+            )
+            gaussian_descriptors_np = gaussian_descriptors_np[:, permutation]
 
         # In comparison to bispectrum, the atomic density always returns
         # in the "local mode". Thus we have to make some slight adjustments
@@ -244,7 +250,7 @@ def __calculate_lammps(self, outdir, **kwargs):
             # the descriptors, even in serial mode, without any further
             # reordering.
             if return_directly:
-                return gaussian_descriptors_np_test
+                return gaussian_descriptors_np
             else:
                 # Here, we want to do something else with the atomic density,
                 # and thus have to properly reorder it.

diff --git a/mala/targets/density.py b/mala/targets/density.py
@@ -1047,9 +1047,6 @@ def __setup_total_energy_module(
                     "Number of atoms is inconsistent between MALA "
                     "and Quantum Espresso."
                 )
-        test_positions = te.get_positions(
-            atoms_Angstrom.get_global_number_of_atoms()
-        )
 
         # We need to find out if the grid dimensions are consistent.
         # That depends on the form of the density data we received.
@@ -1105,7 +1102,8 @@ def __setup_total_energy_module(
                 "Using 1D density to calculate the total energy"
                 " requires reshaping of this data. "
                 "This is unproblematic, as long as you provided t"
-                "he correct grid_dimensions."
+                "he correct grid_dimensions.",
+                min_verbosity=2,
             )
             density_for_qe = self.get_density(
                 density_data, convert_to_threedimensional=True