From d145d550d08f6989413f93fe142891b8aaa44dde Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Fri, 15 Nov 2024 09:50:11 +0100
Subject: [PATCH 01/15] Some small changes for the twoelement case

---
 mala/descriptors/in.bgrid.twoelements.python | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/mala/descriptors/in.bgrid.twoelements.python b/mala/descriptors/in.bgrid.twoelements.python
index b216c05f2..941948acd 100644
--- a/mala/descriptors/in.bgrid.twoelements.python
+++ b/mala/descriptors/in.bgrid.twoelements.python
@@ -11,10 +11,12 @@ read_data ${atom_config_fname}
 
 mass * 1.0
 
+# Needs to be defined for Kokkos
+run_style verlet
+
 # define grid compute and atom compute
 
 group 		snapgroup type 1 2
-#variable 	rcutfac equal 4.67637
 variable 	rfac0 equal 0.99363
 variable 	rmin0 equal 0
 
@@ -41,9 +43,9 @@ compute bgrid all sna/grid grid ${ngridx} ${ngridy} ${ngridz} ${rcutfac} ${rfac0
 
 
 # is this important? or does it just need to be big enough?
-#variable rcutneigh equal 2.0*${rcutfac}*${radelem}
+variable rcutneigh equal 2.0*${rcutfac}*${radelem}
 # for water
-variable rcutneigh equal 4.0*${rcutfac}*${radelem1}
+#variable rcutneigh equal 4.0*${rcutfac}*${radelem1}
 
 pair_style zero ${rcutneigh}
 pair_coeff * *

From ffe13c502e5c288efdc0515d60cefbabe3e54e2f Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Fri, 15 Nov 2024 11:17:31 +0100
Subject: [PATCH 02/15] Bispectrum descriptor calculation for two elements in
 principle works

---
 mala/descriptors/bispectrum.py                  | 17 ++++++++++++++++-
 .../{in.bgrid.python => in.bgrid_n1.python}     |  0
 ...id.twoelements.python => in.bgrid_n2.python} |  2 +-
 3 files changed, 17 insertions(+), 2 deletions(-)
 rename mala/descriptors/{in.bgrid.python => in.bgrid_n1.python} (100%)
 rename mala/descriptors/{in.bgrid.twoelements.python => in.bgrid_n2.python} (96%)

diff --git a/mala/descriptors/bispectrum.py b/mala/descriptors/bispectrum.py
index 207fac341..691012cff 100755
--- a/mala/descriptors/bispectrum.py
+++ b/mala/descriptors/bispectrum.py
@@ -137,6 +137,12 @@ def __calculate_lammps(self, outdir, **kwargs):
         # general LAMMPS import.
         from lammps import constants as lammps_constants
 
+        if len(set(self.atoms.numbers)) > 2:
+            raise ValueError(
+                "MALA can only compute bispectrum descriptors for 1- or "
+                "2-element systems currently."
+            )
+
         use_fp64 = kwargs.get("use_fp64", False)
         keep_logs = kwargs.get("keep_logs", False)
 
@@ -173,7 +179,10 @@ def __calculate_lammps(self, outdir, **kwargs):
                     )
             else:
                 self.parameters.lammps_compute_file = os.path.join(
-                    filepath, "in.bgrid.python"
+                    filepath,
+                    "in.bgrid_n{0}.python".format(
+                        len(set(self.atoms.numbers))
+                    ),
                 )
 
         # Do the LAMMPS calculation and clean up.
@@ -279,6 +288,12 @@ def __calculate_python(self, **kwargs):
             "large systems."
         )
 
+        if len(set(self.atoms.numbers)) > 1:
+            raise ValueError(
+                " MALA cannot compute bispectrum descriptors for "
+                "multi-element systems with python currently."
+            )
+
         # The entire bispectrum calculation may be extensively profiled.
         profile_calculation = kwargs.get("profile_calculation", False)
         if profile_calculation:
diff --git a/mala/descriptors/in.bgrid.python b/mala/descriptors/in.bgrid_n1.python
similarity index 100%
rename from mala/descriptors/in.bgrid.python
rename to mala/descriptors/in.bgrid_n1.python
diff --git a/mala/descriptors/in.bgrid.twoelements.python b/mala/descriptors/in.bgrid_n2.python
similarity index 96%
rename from mala/descriptors/in.bgrid.twoelements.python
rename to mala/descriptors/in.bgrid_n2.python
index 941948acd..82d22b7ab 100644
--- a/mala/descriptors/in.bgrid.twoelements.python
+++ b/mala/descriptors/in.bgrid_n2.python
@@ -43,7 +43,7 @@ compute bgrid all sna/grid grid ${ngridx} ${ngridy} ${ngridz} ${rcutfac} ${rfac0
 
 
 # is this important? or does it just need to be big enough?
-variable rcutneigh equal 2.0*${rcutfac}*${radelem}
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
 # for water
 #variable rcutneigh equal 4.0*${rcutfac}*${radelem1}
 

From a0ff2581eec248e4cb8187823ce23b60fdb4ec78 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Tue, 19 Nov 2024 09:23:42 +0100
Subject: [PATCH 03/15] Parallel processing works in principle

---
 mala/descriptors/bispectrum.py                | 11 +++--
 ...on => in.bgridlocal_defaultproc_n1.python} |  0
 .../in.bgridlocal_defaultproc_n2.python       | 41 +++++++++++++++++
 ...idlocal.python => in.bgridlocal_n1.python} |  0
 mala/descriptors/in.bgridlocal_n2.python      | 46 +++++++++++++++++++
 5 files changed, 95 insertions(+), 3 deletions(-)
 rename mala/descriptors/{in.bgridlocal_defaultproc.python => in.bgridlocal_defaultproc_n1.python} (100%)
 create mode 100644 mala/descriptors/in.bgridlocal_defaultproc_n2.python
 rename mala/descriptors/{in.bgridlocal.python => in.bgridlocal_n1.python} (100%)
 create mode 100644 mala/descriptors/in.bgridlocal_n2.python

diff --git a/mala/descriptors/bispectrum.py b/mala/descriptors/bispectrum.py
index 691012cff..3c5ca514a 100755
--- a/mala/descriptors/bispectrum.py
+++ b/mala/descriptors/bispectrum.py
@@ -171,11 +171,17 @@ def __calculate_lammps(self, outdir, **kwargs):
             if self.parameters._configuration["mpi"]:
                 if self.parameters.use_z_splitting:
                     self.parameters.lammps_compute_file = os.path.join(
-                        filepath, "in.bgridlocal.python"
+                        filepath,
+                        "in.bgridlocal_n{0}.python".format(
+                            len(set(self.atoms.numbers))
+                        ),
                     )
                 else:
                     self.parameters.lammps_compute_file = os.path.join(
-                        filepath, "in.bgridlocal_defaultproc.python"
+                        filepath,
+                        "in.bgridlocal_defaultproc_n{0}.python".format(
+                            len(set(self.atoms.numbers))
+                        ),
                     )
             else:
                 self.parameters.lammps_compute_file = os.path.join(
@@ -184,7 +190,6 @@ def __calculate_lammps(self, outdir, **kwargs):
                         len(set(self.atoms.numbers))
                     ),
                 )
-
         # Do the LAMMPS calculation and clean up.
         lmp.file(self.parameters.lammps_compute_file)
 
diff --git a/mala/descriptors/in.bgridlocal_defaultproc.python b/mala/descriptors/in.bgridlocal_defaultproc_n1.python
similarity index 100%
rename from mala/descriptors/in.bgridlocal_defaultproc.python
rename to mala/descriptors/in.bgridlocal_defaultproc_n1.python
diff --git a/mala/descriptors/in.bgridlocal_defaultproc_n2.python b/mala/descriptors/in.bgridlocal_defaultproc_n2.python
new file mode 100644
index 000000000..00d06456e
--- /dev/null
+++ b/mala/descriptors/in.bgridlocal_defaultproc_n2.python
@@ -0,0 +1,41 @@
+# Calculate bispectrum descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, twojmax, rcutfac, atom_config_fname
+# using command-line -var option
+
+# Initialize simulation
+
+units		metal
+
+read_data ${atom_config_fname}
+
+mass * 1.0
+
+# define grid compute and atom compute
+
+group 		snapgroup type 1
+variable 	rfac0 equal 0.99363
+variable 	rmin0 equal 0
+variable 	wj1 equal 1
+variable 	wj2 equal 1
+variable 	wj1local equal ${wj1}-1.0e-15 # inject a bit of fuzz
+variable 	wj2local equal ${wj2}-1.0e-15 # inject a bit of fuzz
+variable 	radelem1 equal 0.5
+variable 	radelem2 equal 0.5
+variable 	bzero equal 0
+variable 	quadratic equal 0
+
+compute bgridlocal all sna/grid/local grid ${ngridx} ${ngridy} ${ngridz} ${rcutfac} ${rfac0} ${twojmax} ${radelem1} ${radelem2} ${wj1local} ${wj2local}  rmin0 ${rmin0} bzeroflag ${bzero} quadraticflag ${quadratic} switchflag ${switch}
+
+# is this important? or does it just need to be big enough?
+
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+# define output
+
+thermo_modify norm yes
+
+run 0
diff --git a/mala/descriptors/in.bgridlocal.python b/mala/descriptors/in.bgridlocal_n1.python
similarity index 100%
rename from mala/descriptors/in.bgridlocal.python
rename to mala/descriptors/in.bgridlocal_n1.python
diff --git a/mala/descriptors/in.bgridlocal_n2.python b/mala/descriptors/in.bgridlocal_n2.python
new file mode 100644
index 000000000..11e75fbe8
--- /dev/null
+++ b/mala/descriptors/in.bgridlocal_n2.python
@@ -0,0 +1,46 @@
+# Calculate bispectrum descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, twojmax, rcutfac, atom_config_fname 
+# using command-line -var option
+
+# set up processor grid to match QE
+${lammps_procs}
+
+# Initialize simulation
+
+units		metal
+
+read_data ${atom_config_fname}
+
+mass * 1.0
+
+# is this important? or does it just need to be big enough?
+
+variable 	radelem1 equal 0.5
+variable 	radelem2 equal 0.5
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+# set up plane cutoff values to evenly distribute grid points
+
+${zbal}
+
+# define grid compute and atom compute
+
+group 		snapgroup type 1 2
+variable 	rfac0 equal 0.99363
+variable 	rmin0 equal 0
+variable 	wj1 equal 1
+variable 	wj2 equal 1
+variable 	wj1local equal ${wj1}-1.0e-15 # inject a bit of fuzz
+variable 	wj2local equal ${wj2}-1.0e-15 # inject a bit of fuzz
+variable 	bzero equal 0
+variable 	quadratic equal 0
+
+compute bgridlocal all sna/grid/local grid ${ngridx} ${ngridy} ${ngridz} ${rcutfac} ${rfac0} ${twojmax} ${radelem1} ${radelem2} ${wj1local} ${wj2local} rmin0 ${rmin0} bzeroflag ${bzero} quadraticflag ${quadratic} switchflag ${switch}
+
+thermo_modify norm yes
+
+run 0

From 6b81953ca215cec9a4d87d8808851d4d4e7936e7 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Fri, 29 Nov 2024 14:48:59 +0100
Subject: [PATCH 04/15] Adapted bispectrum file to new internal nomenclature

---
 mala/descriptors/bispectrum.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mala/descriptors/bispectrum.py b/mala/descriptors/bispectrum.py
index 7f354433e..2dceabf97 100755
--- a/mala/descriptors/bispectrum.py
+++ b/mala/descriptors/bispectrum.py
@@ -132,7 +132,7 @@ def __calculate_lammps(self, outdir, **kwargs):
         # general LAMMPS import.
         from lammps import constants as lammps_constants
 
-        if len(set(self.atoms.numbers)) > 2:
+        if len(set(self._atoms.numbers)) > 2:
             raise ValueError(
                 "MALA can only compute bispectrum descriptors for 1- or "
                 "2-element systems currently."
@@ -168,21 +168,21 @@ def __calculate_lammps(self, outdir, **kwargs):
                     self.parameters.lammps_compute_file = os.path.join(
                         filepath,
                         "in.bgridlocal_n{0}.python".format(
-                            len(set(self.atoms.numbers))
+                            len(set(self._atoms.numbers))
                         ),
                     )
                 else:
                     self.parameters.lammps_compute_file = os.path.join(
                         filepath,
                         "in.bgridlocal_defaultproc_n{0}.python".format(
-                            len(set(self.atoms.numbers))
+                            len(set(self._atoms.numbers))
                         ),
                     )
             else:
                 self.parameters.lammps_compute_file = os.path.join(
                     filepath,
                     "in.bgrid_n{0}.python".format(
-                        len(set(self.atoms.numbers))
+                        len(set(self._atoms.numbers))
                     ),
                 )
         # Do the LAMMPS calculation and clean up.

From 64956b8ef3b029c1a5e5f739ced9a6765e02679f Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Fri, 29 Nov 2024 15:16:08 +0100
Subject: [PATCH 05/15] Started implementing elemental weights

---
 mala/common/parameters.py      | 25 +++++++++++++++++++++++++
 mala/descriptors/bispectrum.py | 12 ++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/mala/common/parameters.py b/mala/common/parameters.py
index 1d2ba9d96..1900649c3 100644
--- a/mala/common/parameters.py
+++ b/mala/common/parameters.py
@@ -389,6 +389,7 @@ def __init__(self):
         self.bispectrum_twojmax = 10
         self.bispectrum_cutoff = 4.67637
         self.bispectrum_switchflag = 1
+        self.bispectrum_element_weights = None
 
         # Everything pertaining to the atomic density.
         # Seperate cutoff given here because bispectrum descriptors and
@@ -473,6 +474,30 @@ def bispectrum_switchflag(self, value):
         if _int_value > 0:
             self._snap_switchflag = 1
 
+    @property
+    def bispectrum_element_weights(self):
+        """
+        Element species weights for the bispectrum calculation.
+
+        They are provided as an ordered list, and will be assigned to the
+        elements alphabetically, i.e., the first entry will go to the element
+        coming first in the alphabet and so on. Weights are always relative, so
+        the list will be rescaled such that the largest value is 1 and all
+        the other ones are scaled accordingly.
+        """
+        return self._bispectrum_element_weights
+
+    @bispectrum_element_weights.setter
+    def bispectrum_element_weights(self, value):
+        if not isinstance(value, list) and value is not None:
+            raise ValueError("Bispectrum element weights must be list.")
+        if value is not None:
+            if np.max(value) != 1.0:
+                max = np.max(value)
+                for element in range(len(value)):
+                    value[element] /= max
+        self._bispectrum_element_weights = value
+
     def _update_mpi(self, new_mpi):
         self._configuration["mpi"] = new_mpi
 
diff --git a/mala/descriptors/bispectrum.py b/mala/descriptors/bispectrum.py
index 2dceabf97..3ecc478cf 100755
--- a/mala/descriptors/bispectrum.py
+++ b/mala/descriptors/bispectrum.py
@@ -153,6 +153,18 @@ def __calculate_lammps(self, outdir, **kwargs):
         nz = self.grid_dimensions[2]
 
         # Create LAMMPS instance.
+        if (
+            len(set(self._atoms.numbers)) > 1
+            and self.parameters.bispectrum_element_weights is None
+        ):
+            self.parameters.bispectrum_element_weights = [1] * len(
+                set(self._atoms.numbers)
+            )
+            printout(
+                "Multielement system selected without providing elemental "
+                "weights. Set weights to: ",
+                self.parameters.bispectrum_element_weights,
+            )
         lammps_dict = {
             "twojmax": self.parameters.bispectrum_twojmax,
             "rcutfac": self.parameters.bispectrum_cutoff,

From 62f3304109bf68178a7c7e8da4f7ea12266e6366 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 2 Dec 2024 09:50:02 +0100
Subject: [PATCH 06/15] Added test for atoms to LAMMPS data

---
 test/complete_interfaces_test.py | 40 +++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

diff --git a/test/complete_interfaces_test.py b/test/complete_interfaces_test.py
index 6a7956656..3a59eb558 100644
--- a/test/complete_interfaces_test.py
+++ b/test/complete_interfaces_test.py
@@ -1,7 +1,7 @@
 import importlib
 import os
 
-from ase.io import read
+from ase.io import read, write
 import mala
 from mala.common.parameters import ParametersBase
 import numpy as np
@@ -357,3 +357,41 @@ def test_additional_calculation_data_json(self):
                     new_ldos_calculator.atoms.get_positions()[i, j],
                     rtol=accuracy_fine,
                 )
+
+    def test_atoms_to_lammps_data_multielement(self):
+        """
+        Test if the atoms to LAMMPS interface gives alphabetical ordering.
+
+        If I am not misunderstanding anything than how the element ordering
+        is transferred from an atoms object to LAMMPS is not necessarily
+        well-defined. In the current interface, it seems to simply be
+        alphabetical ordering. We need to check whether this holds.
+        It could change, and since LAMMPS itself does not name the elements
+        in its data file, one would not notice.
+
+        Returns
+        -------
+
+        """
+        vasp_source_file = open("test_element.vasp", mode="w")
+        vasp_source_file.write(
+            "O  Al Te H\n 1.0000000000000000\n     6.0000000000000000    0.0000000000000000    0.0000000000000000\n     0.0000000000000000    6.0000000000000000    0.0000000000000000\n     0.0000000000000000    0.0000000000000000    6.0000000000000000\nO  Al Te H\n   1   1   1   1\nCartesian\n  0.0 0.0 1.0\n  0.0 0.0 2.0\n  0.0 0.0 3.0\n  0.0 0.0 4.0 \n"
+        )
+        vasp_source_file.close()
+        atoms = read("test_element.vasp", format="vasp")
+        write("test_element.lammps", atoms, format="lammps-data")
+
+        raw_file = open("test_element.lammps")
+        lines = raw_file.readlines()
+        found_atoms = 0
+        for idx, line in enumerate(lines):
+            if "Atoms" in line:
+                found_atoms = idx
+
+        alphabetical_order_elements = [3, 1, 4, 2]
+
+        for i in range(2, 6):
+            assert (
+                int(lines[found_atoms + i].split()[1])
+                == alphabetical_order_elements[i - 2]
+            )

From 240ca440f61d77a7121c7d5ff5063f2388504f17 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 2 Dec 2024 11:22:38 +0100
Subject: [PATCH 07/15] Made adjustable weights available

---
 mala/descriptors/bispectrum.py                | 27 ++++++++++---------
 mala/descriptors/in.bgrid_n2.python           |  4 +--
 .../in.bgridlocal_defaultproc_n2.python       |  4 +--
 mala/descriptors/in.bgridlocal_n2.python      |  4 +--
 test/complete_interfaces_test.py              |  4 ---
 5 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/mala/descriptors/bispectrum.py b/mala/descriptors/bispectrum.py
index 3ecc478cf..46fa68484 100755
--- a/mala/descriptors/bispectrum.py
+++ b/mala/descriptors/bispectrum.py
@@ -153,22 +153,25 @@ def __calculate_lammps(self, outdir, **kwargs):
         nz = self.grid_dimensions[2]
 
         # Create LAMMPS instance.
-        if (
-            len(set(self._atoms.numbers)) > 1
-            and self.parameters.bispectrum_element_weights is None
-        ):
-            self.parameters.bispectrum_element_weights = [1] * len(
-                set(self._atoms.numbers)
-            )
-            printout(
-                "Multielement system selected without providing elemental "
-                "weights. Set weights to: ",
-                self.parameters.bispectrum_element_weights,
-            )
         lammps_dict = {
             "twojmax": self.parameters.bispectrum_twojmax,
             "rcutfac": self.parameters.bispectrum_cutoff,
         }
+        if len(set(self._atoms.numbers)) > 1:
+
+            if self.parameters.bispectrum_element_weights is None:
+                self.parameters.bispectrum_element_weights = [1] * len(
+                    set(self._atoms.numbers)
+                )
+                printout(
+                    "Multielement system selected without providing elemental "
+                    "weights. Set weights to: ",
+                    self.parameters.bispectrum_element_weights,
+                )
+            for i in range(len(self.parameters.bispectrum_element_weights)):
+                lammps_dict["wj" + str(i + 1)] = (
+                    self.parameters.bispectrum_element_weights[i]
+                )
         lmp = self._setup_lammps(nx, ny, nz, lammps_dict)
 
         # An empty string means that the user wants to use the standard input.
diff --git a/mala/descriptors/in.bgrid_n2.python b/mala/descriptors/in.bgrid_n2.python
index 82d22b7ab..26f6514c4 100644
--- a/mala/descriptors/in.bgrid_n2.python
+++ b/mala/descriptors/in.bgrid_n2.python
@@ -21,8 +21,8 @@ variable 	rfac0 equal 0.99363
 variable 	rmin0 equal 0
 
 #variable 	wj equal 1
-variable 	wj1 equal 1
-variable 	wj2 equal 1
+#variable 	wj1 equal 1
+#variable 	wj2 equal 1
 
 #variable 	radelem equal 0.5
 variable 	radelem1 equal 0.5
diff --git a/mala/descriptors/in.bgridlocal_defaultproc_n2.python b/mala/descriptors/in.bgridlocal_defaultproc_n2.python
index 00d06456e..826f8a239 100644
--- a/mala/descriptors/in.bgridlocal_defaultproc_n2.python
+++ b/mala/descriptors/in.bgridlocal_defaultproc_n2.python
@@ -16,8 +16,8 @@ mass * 1.0
 group 		snapgroup type 1
 variable 	rfac0 equal 0.99363
 variable 	rmin0 equal 0
-variable 	wj1 equal 1
-variable 	wj2 equal 1
+# variable 	wj1 equal 1
+# variable 	wj2 equal 1
 variable 	wj1local equal ${wj1}-1.0e-15 # inject a bit of fuzz
 variable 	wj2local equal ${wj2}-1.0e-15 # inject a bit of fuzz
 variable 	radelem1 equal 0.5
diff --git a/mala/descriptors/in.bgridlocal_n2.python b/mala/descriptors/in.bgridlocal_n2.python
index 11e75fbe8..160283d71 100644
--- a/mala/descriptors/in.bgridlocal_n2.python
+++ b/mala/descriptors/in.bgridlocal_n2.python
@@ -32,8 +32,8 @@ ${zbal}
 group 		snapgroup type 1 2
 variable 	rfac0 equal 0.99363
 variable 	rmin0 equal 0
-variable 	wj1 equal 1
-variable 	wj2 equal 1
+# variable 	wj1 equal 1
+# variable 	wj2 equal 1
 variable 	wj1local equal ${wj1}-1.0e-15 # inject a bit of fuzz
 variable 	wj2local equal ${wj2}-1.0e-15 # inject a bit of fuzz
 variable 	bzero equal 0
diff --git a/test/complete_interfaces_test.py b/test/complete_interfaces_test.py
index 3a59eb558..c4ad25f28 100644
--- a/test/complete_interfaces_test.py
+++ b/test/complete_interfaces_test.py
@@ -368,10 +368,6 @@ def test_atoms_to_lammps_data_multielement(self):
         alphabetical ordering. We need to check whether this holds.
         It could change, and since LAMMPS itself does not name the elements
         in its data file, one would not notice.
-
-        Returns
-        -------
-
         """
         vasp_source_file = open("test_element.vasp", mode="w")
         vasp_source_file.write(

From a6ff4c31ec8ddba974202358e44952bc060f23d4 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 2 Dec 2024 12:26:43 +0100
Subject: [PATCH 08/15] Added runfiles for multielement Gaussian grid

---
 mala/descriptors/atomic_density.py            | 24 +++++++++++---
 ....python => in.ggrid_defaultproc_n1.python} |  0
 .../in.ggrid_defaultproc_n2.python            | 25 +++++++++++++++
 mala/descriptors/in.ggrid_n1.python           | 32 +++++++++++++++++++
 .../{in.ggrid.python => in.ggrid_n2.python}   |  5 +--
 5 files changed, 80 insertions(+), 6 deletions(-)
 rename mala/descriptors/{in.ggrid_defaultproc.python => in.ggrid_defaultproc_n1.python} (100%)
 create mode 100644 mala/descriptors/in.ggrid_defaultproc_n2.python
 create mode 100644 mala/descriptors/in.ggrid_n1.python
 rename mala/descriptors/{in.ggrid.python => in.ggrid_n2.python} (83%)

diff --git a/mala/descriptors/atomic_density.py b/mala/descriptors/atomic_density.py
index 4459c838b..558f66957 100755
--- a/mala/descriptors/atomic_density.py
+++ b/mala/descriptors/atomic_density.py
@@ -150,9 +150,16 @@ def __calculate_lammps(self, outdir, **kwargs):
 
         # Create LAMMPS instance.
         lammps_dict = {
-            "sigma": self.parameters.atomic_density_sigma,
             "rcutfac": self.parameters.atomic_density_cutoff,
         }
+        if len(set(self._atoms.numbers)) == 1:
+            lammps_dict["sigma"] = self.parameters.atomic_density_sigma
+        else:
+            for i in range(len(set(self._atoms.numbers))):
+                lammps_dict["sigma" + str(i + 1)] = (
+                    self.parameters.atomic_density_sigma
+                )
+
         lmp = self._setup_lammps(nx, ny, nz, lammps_dict)
 
         # For now the file is chosen automatically, because this is used
@@ -161,15 +168,24 @@ def __calculate_lammps(self, outdir, **kwargs):
         if self.parameters._configuration["mpi"]:
             if self.parameters.use_z_splitting:
                 self.parameters.lammps_compute_file = os.path.join(
-                    filepath, "in.ggrid.python"
+                    filepath,
+                    "in.ggrid_n{0}.python".format(
+                        len(set(self._atoms.numbers))
+                    ),
                 )
             else:
                 self.parameters.lammps_compute_file = os.path.join(
-                    filepath, "in.ggrid_defaultproc.python"
+                    filepath,
+                    "in.ggrid_defaultproc_n{0}.python".format(
+                        len(set(self._atoms.numbers))
+                    ),
                 )
         else:
             self.parameters.lammps_compute_file = os.path.join(
-                filepath, "in.ggrid_defaultproc.python"
+                filepath,
+                "in.ggrid_defaultproc_n{0}.python".format(
+                    len(set(self._atoms.numbers))
+                ),
             )
 
         # Do the LAMMPS calculation and clean up.
diff --git a/mala/descriptors/in.ggrid_defaultproc.python b/mala/descriptors/in.ggrid_defaultproc_n1.python
similarity index 100%
rename from mala/descriptors/in.ggrid_defaultproc.python
rename to mala/descriptors/in.ggrid_defaultproc_n1.python
diff --git a/mala/descriptors/in.ggrid_defaultproc_n2.python b/mala/descriptors/in.ggrid_defaultproc_n2.python
new file mode 100644
index 000000000..dbcd4884b
--- /dev/null
+++ b/mala/descriptors/in.ggrid_defaultproc_n2.python
@@ -0,0 +1,25 @@
+# Calculate Gaussian atomic density descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, sigma, atom_config_fname
+# using command-line -var option
+
+# Initialize simulation
+
+units		metal
+read_data ${atom_config_fname}
+mass * 1.0
+
+
+# is this important? or does it just need to be big enough?
+variable	radelem1 equal 0.5
+variable	radelem2 equal 0.5
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+compute ggrid all gaussian/grid/local  grid ${ngridx} ${ngridy} ${ngridz}  ${rcutfac}  ${radelem1}  ${radelem2}  ${sigma1} ${sigma2}
+
+thermo_modify norm yes
+
+run 0
diff --git a/mala/descriptors/in.ggrid_n1.python b/mala/descriptors/in.ggrid_n1.python
new file mode 100644
index 000000000..c7821beec
--- /dev/null
+++ b/mala/descriptors/in.ggrid_n1.python
@@ -0,0 +1,32 @@
+# Calculate Gaussian atomic density descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, sigma, atom_config_fname
+# using command-line -var option
+
+# set up processor grid to match QE
+${lammps_procs}
+
+# Initialize simulation
+
+units		metal
+read_data ${atom_config_fname}
+mass * 1.0
+
+
+# is this important? or does it just need to be big enough?
+variable	radelem1 equal 0.5
+variable	radelem2 equal 0.5
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+# set up plane cutoff values to evenly distribute grid points
+${zbal}
+
+
+compute ggrid all gaussian/grid/local  grid ${ngridx} ${ngridy} ${ngridz}  ${rcutfac}  ${radelem1}  ${radelem2}  ${sigma1} ${sigma2}
+
+thermo_modify norm yes
+
+run 0
diff --git a/mala/descriptors/in.ggrid.python b/mala/descriptors/in.ggrid_n2.python
similarity index 83%
rename from mala/descriptors/in.ggrid.python
rename to mala/descriptors/in.ggrid_n2.python
index 265eac8f8..ad362b50a 100644
--- a/mala/descriptors/in.ggrid.python
+++ b/mala/descriptors/in.ggrid_n2.python
@@ -14,7 +14,8 @@ mass * 1.0
 
 
 # is this important? or does it just need to be big enough?
-variable	radelem equal 0.5
+variable	radelem1 equal 0.5
+variable	radelem2 equal 0.5
 variable rcutneigh equal 2.0*${rcutfac}*${radelem}
 
 pair_style zero ${rcutneigh}
@@ -24,7 +25,7 @@ pair_coeff * *
 ${zbal}
 
 
-compute ggrid all gaussian/grid/local  grid ${ngridx} ${ngridy} ${ngridz}  ${rcutfac}  ${radelem}  ${sigma}
+compute ggrid all gaussian/grid/local  grid ${ngridx} ${ngridy} ${ngridz}  ${rcutfac}  ${radelem1}  ${radelem2}  ${sigma1} ${sigma2}
 
 thermo_modify norm yes
 

From ace21a949e8069665627864624b18af6cc156c01 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 2 Dec 2024 12:57:59 +0100
Subject: [PATCH 09/15] Atomic density formula working for two elements

---
 mala/descriptors/atomic_density.py | 3 ++-
 mala/targets/density.py            | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/mala/descriptors/atomic_density.py b/mala/descriptors/atomic_density.py
index 558f66957..45c889c30 100755
--- a/mala/descriptors/atomic_density.py
+++ b/mala/descriptors/atomic_density.py
@@ -213,6 +213,7 @@ def __calculate_lammps(self, outdir, **kwargs):
             array_shape=(nrows_ggrid, ncols_ggrid),
             use_fp64=use_fp64,
         )
+
         self._clean_calculation(lmp, keep_logs)
 
         # In comparison to bispectrum, the atomic density always returns
@@ -241,7 +242,7 @@ def __calculate_lammps(self, outdir, **kwargs):
                         self.grid_dimensions[2],
                         self.grid_dimensions[1],
                         self.grid_dimensions[0],
-                        7,
+                        6 + len(set(self._atoms.numbers)),
                     )
                 )
                 gaussian_descriptors_np = gaussian_descriptors_np.transpose(
diff --git a/mala/targets/density.py b/mala/targets/density.py
index 61623fe24..be04f6f5b 100644
--- a/mala/targets/density.py
+++ b/mala/targets/density.py
@@ -1214,7 +1214,7 @@ def __setup_total_energy_module(
             t0 = time.perf_counter()
             gaussian_descriptors = np.reshape(
                 gaussian_descriptors,
-                [number_of_gridpoints_mala, 1],
+                [number_of_gridpoints_mala, len(set(self.atoms.numbers))],
                 order="F",
             )
             reference_gaussian_descriptors = np.reshape(
@@ -1244,7 +1244,7 @@ def __setup_total_energy_module(
                 reference_gaussian_descriptors,
                 sigma,
                 number_of_gridpoints,
-                1,
+                len(set(self.atoms.numbers)),
             )
             barrier()
             printout(

From 907e254a51322ee565620e08f58baadad0095e54 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 2 Dec 2024 15:15:43 +0100
Subject: [PATCH 10/15] Three elements in principle work, but atomic density
 formula not for some reason

---
 mala/descriptors/bispectrum.py                |  6 +-
 mala/descriptors/in.bgrid_n3.python           | 57 +++++++++++++++++++
 .../in.bgridlocal_defaultproc_n2.python       |  2 +-
 .../in.bgridlocal_defaultproc_n3.python       | 41 +++++++++++++
 mala/descriptors/in.bgridlocal_n3.python      | 46 +++++++++++++++
 .../in.ggrid_defaultproc_n3.python            | 26 +++++++++
 mala/descriptors/in.ggrid_n3.python           | 33 +++++++++++
 7 files changed, 207 insertions(+), 4 deletions(-)
 create mode 100644 mala/descriptors/in.bgrid_n3.python
 create mode 100644 mala/descriptors/in.bgridlocal_defaultproc_n3.python
 create mode 100644 mala/descriptors/in.bgridlocal_n3.python
 create mode 100644 mala/descriptors/in.ggrid_defaultproc_n3.python
 create mode 100644 mala/descriptors/in.ggrid_n3.python

diff --git a/mala/descriptors/bispectrum.py b/mala/descriptors/bispectrum.py
index 46fa68484..3027b6bb5 100755
--- a/mala/descriptors/bispectrum.py
+++ b/mala/descriptors/bispectrum.py
@@ -132,10 +132,10 @@ def __calculate_lammps(self, outdir, **kwargs):
         # general LAMMPS import.
         from lammps import constants as lammps_constants
 
-        if len(set(self._atoms.numbers)) > 2:
+        if len(set(self._atoms.numbers)) > 3:
             raise ValueError(
-                "MALA can only compute bispectrum descriptors for 1- or "
-                "2-element systems currently."
+                "MALA can only compute bispectrum descriptors up to "
+                "3-element systems currently."
             )
 
         use_fp64 = kwargs.get("use_fp64", False)
diff --git a/mala/descriptors/in.bgrid_n3.python b/mala/descriptors/in.bgrid_n3.python
new file mode 100644
index 000000000..08ee0db70
--- /dev/null
+++ b/mala/descriptors/in.bgrid_n3.python
@@ -0,0 +1,57 @@
+# Calculate bispectrum descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, twojmax, rcutfac, atom_config_fname 
+# using command-line -var option
+
+# Initialize simulation
+
+units		metal
+
+read_data ${atom_config_fname} 
+
+mass * 1.0
+
+# Needs to be defined for Kokkos
+run_style verlet
+
+# define grid compute and atom compute
+
+group 		snapgroup type 1 2 3
+variable 	rfac0 equal 0.99363
+variable 	rmin0 equal 0
+
+#variable 	radelem equal 0.5
+variable 	radelem1 equal 0.5
+variable 	radelem2 equal 0.5
+variable 	radelem3 equal 0.5
+
+variable 	bzero equal 0
+variable 	quadratic equal 0
+
+#variable        snap_options string &
+#"${rcutfac} ${rfac0} ${twojmax} ${radelem} ${wj} rmin0 ${rmin0} quadraticflag ${quadratic} bzeroflag ${bzero} switchflag ${switch}"
+
+#compute bgrid all sna/grid grid ${ngridx} ${ngridy} ${ngridz} ${snap_options}
+
+compute bgrid all sna/grid grid ${ngridx} ${ngridy} ${ngridz} ${rcutfac} ${rfac0} ${twojmax} ${radelem1} ${radelem2} ${radelem3} ${wj1} ${wj2} ${wj3} rmin0 ${rmin0} bzeroflag ${bzero} quadraticflag ${quadratic} switchflag ${switch}
+
+
+# create dummy potential for neighbor list
+
+
+# is this important? or does it just need to be big enough?
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+# for water
+#variable rcutneigh equal 4.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+# define output
+
+thermo_style	custom step temp ke pe vol c_bgrid[1][1]
+thermo_modify norm yes
+
+# run
+
+run 0
diff --git a/mala/descriptors/in.bgridlocal_defaultproc_n2.python b/mala/descriptors/in.bgridlocal_defaultproc_n2.python
index 826f8a239..b315d1d5b 100644
--- a/mala/descriptors/in.bgridlocal_defaultproc_n2.python
+++ b/mala/descriptors/in.bgridlocal_defaultproc_n2.python
@@ -13,7 +13,7 @@ mass * 1.0
 
 # define grid compute and atom compute
 
-group 		snapgroup type 1
+group 		snapgroup type 1 2
 variable 	rfac0 equal 0.99363
 variable 	rmin0 equal 0
 # variable 	wj1 equal 1
diff --git a/mala/descriptors/in.bgridlocal_defaultproc_n3.python b/mala/descriptors/in.bgridlocal_defaultproc_n3.python
new file mode 100644
index 000000000..49d2a534f
--- /dev/null
+++ b/mala/descriptors/in.bgridlocal_defaultproc_n3.python
@@ -0,0 +1,41 @@
+# Calculate bispectrum descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, twojmax, rcutfac, atom_config_fname
+# using command-line -var option
+
+# Initialize simulation
+
+units		metal
+
+read_data ${atom_config_fname}
+
+mass * 1.0
+
+# define grid compute and atom compute
+
+group 		snapgroup type 1 2 3
+variable 	rfac0 equal 0.99363
+variable 	rmin0 equal 0
+variable 	wj1local equal ${wj1}-1.0e-15 # inject a bit of fuzz
+variable 	wj2local equal ${wj2}-1.0e-15 # inject a bit of fuzz
+variable 	wj3local equal ${wj2}-1.0e-15 # inject a bit of fuzz
+variable 	radelem1 equal 0.5
+variable 	radelem2 equal 0.5
+variable 	radelem3 equal 0.5
+variable 	bzero equal 0
+variable 	quadratic equal 0
+
+compute bgridlocal all sna/grid/local grid ${ngridx} ${ngridy} ${ngridz} ${rcutfac} ${rfac0} ${twojmax} ${radelem1} ${radelem2} ${radelem3} ${wj1local} ${wj2local} ${wj3local}  rmin0 ${rmin0} bzeroflag ${bzero} quadraticflag ${quadratic} switchflag ${switch}
+
+# is this important? or does it just need to be big enough?
+
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+# define output
+
+thermo_modify norm yes
+
+run 0
diff --git a/mala/descriptors/in.bgridlocal_n3.python b/mala/descriptors/in.bgridlocal_n3.python
new file mode 100644
index 000000000..d4cf2fcfb
--- /dev/null
+++ b/mala/descriptors/in.bgridlocal_n3.python
@@ -0,0 +1,46 @@
+# Calculate bispectrum descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, twojmax, rcutfac, atom_config_fname 
+# using command-line -var option
+
+# set up processor grid to match QE
+${lammps_procs}
+
+# Initialize simulation
+
+units		metal
+
+read_data ${atom_config_fname}
+
+mass * 1.0
+
+# is this important? or does it just need to be big enough?
+
+variable 	radelem1 equal 0.5
+variable 	radelem2 equal 0.5
+variable 	radelem3 equal 0.5
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+# set up plane cutoff values to evenly distribute grid points
+
+${zbal}
+
+# define grid compute and atom compute
+
+group 		snapgroup type 1 2 3
+variable 	rfac0 equal 0.99363
+variable 	rmin0 equal 0
+variable 	wj1local equal ${wj1}-1.0e-15 # inject a bit of fuzz
+variable 	wj2local equal ${wj2}-1.0e-15 # inject a bit of fuzz
+variable 	wj3local equal ${wj2}-1.0e-15 # inject a bit of fuzz
+variable 	bzero equal 0
+variable 	quadratic equal 0
+
+compute bgridlocal all sna/grid/local grid ${ngridx} ${ngridy} ${ngridz} ${rcutfac} ${rfac0} ${twojmax} ${radelem1} ${radelem2} ${radelem3} ${wj1local} ${wj2local} ${wj3local} rmin0 ${rmin0} bzeroflag ${bzero} quadraticflag ${quadratic} switchflag ${switch}
+
+thermo_modify norm yes
+
+run 0
diff --git a/mala/descriptors/in.ggrid_defaultproc_n3.python b/mala/descriptors/in.ggrid_defaultproc_n3.python
new file mode 100644
index 000000000..c3e34431e
--- /dev/null
+++ b/mala/descriptors/in.ggrid_defaultproc_n3.python
@@ -0,0 +1,26 @@
+# Calculate Gaussian atomic density descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, sigma, atom_config_fname
+# using command-line -var option
+
+# Initialize simulation
+
+units		metal
+read_data ${atom_config_fname}
+mass * 1.0
+
+
+# is this important? or does it just need to be big enough?
+variable	radelem1 equal 0.5
+variable	radelem2 equal 0.5
+variable	radelem3 equal 0.5
+variable rcutneigh equal 2.0*${rcutfac}*${radelem1}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+compute ggrid all gaussian/grid/local  grid ${ngridx} ${ngridy} ${ngridz}  ${rcutfac}  ${radelem1}  ${radelem2} ${radelem3} ${sigma1} ${sigma2} ${sigma3}
+
+thermo_modify norm yes
+
+run 0
diff --git a/mala/descriptors/in.ggrid_n3.python b/mala/descriptors/in.ggrid_n3.python
new file mode 100644
index 000000000..e0a582bcd
--- /dev/null
+++ b/mala/descriptors/in.ggrid_n3.python
@@ -0,0 +1,33 @@
+# Calculate Gaussian atomic density descriptors on a 3D grid
+
+# pass in values ngridx, ngridy, ngridz, sigma, atom_config_fname
+# using command-line -var option
+
+# set up processor grid to match QE
+${lammps_procs}
+
+# Initialize simulation
+
+units		metal
+read_data ${atom_config_fname}
+mass * 1.0
+
+
+# is this important? or does it just need to be big enough?
+variable	radelem1 equal 0.5
+variable	radelem2 equal 0.5
+variable	radelem3 equal 0.5
+variable rcutneigh equal 2.0*${rcutfac}*${radelem}
+
+pair_style zero ${rcutneigh}
+pair_coeff * *
+
+# set up plane cutoff values to evenly distribute grid points
+${zbal}
+
+
+compute ggrid all gaussian/grid/local  grid ${ngridx} ${ngridy} ${ngridz}  ${rcutfac}  ${radelem1}  ${radelem2}  ${radelem3}  ${sigma1} ${sigma2} ${sigma3}
+
+thermo_modify norm yes
+
+run 0

From 8823bda5f56c900e96a9e62a14b59b574b6f2538 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Thu, 5 Dec 2024 11:49:27 +0100
Subject: [PATCH 11/15] Testing some things

---
 mala/descriptors/atomic_density.py | 8 ++++++--
 mala/targets/density.py            | 3 +++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/mala/descriptors/atomic_density.py b/mala/descriptors/atomic_density.py
index 45c889c30..6b5a1349e 100755
--- a/mala/descriptors/atomic_density.py
+++ b/mala/descriptors/atomic_density.py
@@ -215,14 +215,18 @@ def __calculate_lammps(self, outdir, **kwargs):
         )
 
         self._clean_calculation(lmp, keep_logs)
-
+        # gaussian_descriptors_np_test = gaussian_descriptors_np.copy()
+        # if np.shape(gaussian_descriptors_np_test)[1] > 7:
+        #     gaussian_descriptors_np_test[:, 6] = gaussian_descriptors_np[:, 7]
+        #     gaussian_descriptors_np_test[:, 7] = gaussian_descriptors_np[:, 8]
+        #     gaussian_descriptors_np_test[:, 8] = gaussian_descriptors_np[:, 6]
         # In comparison to bispectrum, the atomic density always returns
         # in the "local mode". Thus we have to make some slight adjustments
         # if we operate without MPI.
         self.grid_dimensions = [nx, ny, nz]
         if self.parameters._configuration["mpi"]:
             if return_directly:
-                return gaussian_descriptors_np
+                return gaussian_descriptors_np_test
             else:
                 self.feature_size = 4
                 return gaussian_descriptors_np, nrows_ggrid
diff --git a/mala/targets/density.py b/mala/targets/density.py
index be04f6f5b..cc032cfb4 100644
--- a/mala/targets/density.py
+++ b/mala/targets/density.py
@@ -1047,6 +1047,9 @@ 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.

From 27aa15de4c28bc1f70004094c58453050fbd8cc7 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 9 Dec 2024 10:51:27 +0100
Subject: [PATCH 12/15] Fixed parser to allow for sampling LDOS with over 1000
 cube files

---
 mala/targets/ldos.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/mala/targets/ldos.py b/mala/targets/ldos.py
index c53245003..741b55dcc 100644
--- a/mala/targets/ldos.py
+++ b/mala/targets/ldos.py
@@ -1609,7 +1609,14 @@ def _read_from_qe_files(
 
         for i in range(start_index, end_index):
             tmp_file_name = path_scheme
-            tmp_file_name = tmp_file_name.replace("*", str(i).zfill(digits))
+            if digits < 4:
+                tmp_file_name = tmp_file_name.replace(
+                    "*", str(i).zfill(digits)
+                )
+            else:
+                # For some reason, there are no leading zeros above 3 digits
+                # in QE.
+                tmp_file_name = tmp_file_name.replace("*", str(i).zfill(3))
 
             # Open the cube file
             if file_type == ".cube":

From 5c0843748b35d3caaea90040a5cd2979aab3ce83 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Tue, 10 Dec 2024 16:57:59 +0100
Subject: [PATCH 13/15] More testing

---
 mala/descriptors/atomic_density.py | 8 ++++----
 mala/targets/density.py            | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/mala/descriptors/atomic_density.py b/mala/descriptors/atomic_density.py
index 6b5a1349e..6502e0595 100755
--- a/mala/descriptors/atomic_density.py
+++ b/mala/descriptors/atomic_density.py
@@ -214,8 +214,8 @@ def __calculate_lammps(self, outdir, **kwargs):
             use_fp64=use_fp64,
         )
 
-        self._clean_calculation(lmp, keep_logs)
-        # gaussian_descriptors_np_test = gaussian_descriptors_np.copy()
+        # self._clean_calculation(lmp, keep_logs)
+        gaussian_descriptors_np_test = gaussian_descriptors_np.copy()
         # if np.shape(gaussian_descriptors_np_test)[1] > 7:
         #     gaussian_descriptors_np_test[:, 6] = gaussian_descriptors_np[:, 7]
         #     gaussian_descriptors_np_test[:, 7] = gaussian_descriptors_np[:, 8]
@@ -226,7 +226,7 @@ def __calculate_lammps(self, outdir, **kwargs):
         self.grid_dimensions = [nx, ny, nz]
         if self.parameters._configuration["mpi"]:
             if return_directly:
-                return gaussian_descriptors_np_test
+                return gaussian_descriptors_np
             else:
                 self.feature_size = 4
                 return gaussian_descriptors_np, nrows_ggrid
@@ -236,7 +236,7 @@ def __calculate_lammps(self, outdir, **kwargs):
             # the descriptors, even in serial mode, without any further
             # reordering.
             if return_directly:
-                return gaussian_descriptors_np
+                return gaussian_descriptors_np_test
             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
index cc032cfb4..e42f69c7a 100644
--- a/mala/targets/density.py
+++ b/mala/targets/density.py
@@ -1014,8 +1014,8 @@ def __setup_total_energy_module(
             barrier()
 
             # Right after setup we can delete the file.
-            if get_rank() == 0:
-                os.remove(tem_input_name)
+            # if get_rank() == 0:
+            #     os.remove(tem_input_name)
 
             printout(
                 "Total energy module: Time used by total energy initialization: {:.8f}s".format(

From f84dd75eba678ac02044b93e714510e2b8d14c21 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Fri, 13 Dec 2024 16:45:34 +0100
Subject: [PATCH 14/15] Everything after second component was missing from
 structure factor calculation

---
 .../total_energy_module/total_energy.f90       | 15 +++++++++++++--
 mala/descriptors/atomic_density.py             | 18 +++++++++++++-----
 mala/targets/density.py                        | 10 +++++-----
 3 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/external_modules/total_energy_module/total_energy.f90 b/external_modules/total_energy_module/total_energy.f90
index f1165e01e..d3eb5f669 100644
--- a/external_modules/total_energy_module/total_energy.f90
+++ b/external_modules/total_energy_module/total_energy.f90
@@ -663,9 +663,13 @@ SUBROUTINE set_positions_gauss(verbose, gaussian_descriptors,reference_gaussian_
   CALL mp_barrier( intra_image_comm )
   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, gaussian_descriptors, strf)
+!
   CALL rho_r2g(dfftp, reference_gaussian_descriptors, rgd_of_g)
 
   DO isp = 1, nsp
@@ -807,6 +811,8 @@ SUBROUTINE set_rho_of_r(rho_of_r,nnr_in,nspin_in)
   USE cell_base,     ONLY : omega
   USE mp,            ONLY : mp_sum
   USE mp_bands,      ONLY : intra_bgrp_comm
+  USE vlocal,        ONLY : strf
+  USE ions_base,            ONLY : nsp
 
   IMPLICIT NONE
   INTEGER,     INTENT(IN)  :: nnr_in, nspin_in
@@ -815,6 +821,7 @@ SUBROUTINE set_rho_of_r(rho_of_r,nnr_in,nspin_in)
   REAL(DP) :: charge
   REAL(DP) :: etotefield
   INTEGER :: ir
+  INTEGER :: isp
 
   ! Check consistency of dimensions
   IF (nnr_in /= dfftp%nnr) STOP "*** nnr provided to set_rho_of_r() does not match dfftp%nnr"
@@ -861,6 +868,10 @@ 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
index 6502e0595..661fe3e20 100755
--- a/mala/descriptors/atomic_density.py
+++ b/mala/descriptors/atomic_density.py
@@ -214,12 +214,20 @@ def __calculate_lammps(self, outdir, **kwargs):
             use_fp64=use_fp64,
         )
 
-        # self._clean_calculation(lmp, keep_logs)
+        self._clean_calculation(lmp, keep_logs)
         gaussian_descriptors_np_test = gaussian_descriptors_np.copy()
-        # if np.shape(gaussian_descriptors_np_test)[1] > 7:
-        #     gaussian_descriptors_np_test[:, 6] = gaussian_descriptors_np[:, 7]
-        #     gaussian_descriptors_np_test[:, 7] = gaussian_descriptors_np[:, 8]
-        #     gaussian_descriptors_np_test[:, 8] = gaussian_descriptors_np[:, 6]
+
+        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
+                    ]
+                )
+
         # In comparison to bispectrum, the atomic density always returns
         # in the "local mode". Thus we have to make some slight adjustments
         # if we operate without MPI.
diff --git a/mala/targets/density.py b/mala/targets/density.py
index e42f69c7a..61b1ab334 100644
--- a/mala/targets/density.py
+++ b/mala/targets/density.py
@@ -1014,8 +1014,8 @@ def __setup_total_energy_module(
             barrier()
 
             # Right after setup we can delete the file.
-            # if get_rank() == 0:
-            #     os.remove(tem_input_name)
+            if get_rank() == 0:
+                os.remove(tem_input_name)
 
             printout(
                 "Total energy module: Time used by total energy initialization: {:.8f}s".format(
@@ -1047,9 +1047,9 @@ 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()
-        # )
+        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.

From 9dd9a34205a7e81728c1b4338f425d9473ea3ff0 Mon Sep 17 00:00:00 2001
From: Lenz Fiedler <l.fiedler@hzdr.de>
Date: Mon, 16 Dec 2024 11:57:26 +0100
Subject: [PATCH 15/15] Atomic density formula for multiple elements works now

---
 .../total_energy_module/total_energy.f90      |  7 ----
 mala/descriptors/atomic_density.py            | 34 +++++++++++--------
 mala/targets/density.py                       |  6 ++--
 3 files changed, 22 insertions(+), 25 deletions(-)

diff --git a/external_modules/total_energy_module/total_energy.f90 b/external_modules/total_energy_module/total_energy.f90
index d3eb5f669..faa6b741e 100644
--- 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
index 661fe3e20..4e5ee884d 100755
--- 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
index 61b1ab334..0c9e17536 100644
--- 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