MALA for multiple elements

external_modules/total_energy_module/total_energy.f90

@@ -663,8 +663,9 @@ 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))
-
-  CALL rho_r2g(dfftp, gaussian_descriptors, strf)
+  DO isp = 1, nsp
+     CALL rho_r2g(dfftp, gaussian_descriptors(:, isp:isp), strf(:,isp:isp))
+  ENDDO
 
   CALL rho_r2g(dfftp, reference_gaussian_descriptors, rgd_of_g)
 
@@ -807,6 +808,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 +818,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"

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
 

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
 
@@ -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.
@@ -197,8 +213,27 @@ def __calculate_lammps(self, outdir, **kwargs):
             array_shape=(nrows_ggrid, ncols_ggrid),
             use_fp64=use_fp64,
         )
+
         self._clean_calculation(lmp, keep_logs)
 
+        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
         # if we operate without MPI.
@@ -225,7 +260,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(

mala/descriptors/bispectrum.py

@@ -132,6 +132,12 @@ def __calculate_lammps(self, outdir, **kwargs):
         # general LAMMPS import.
         from lammps import constants as lammps_constants
 
+        if len(set(self._atoms.numbers)) > 3:
+            raise ValueError(
+                "MALA can only compute bispectrum descriptors up to "
+                "3-element systems currently."
+            )
+
         use_fp64 = kwargs.get("use_fp64", False)
         keep_logs = kwargs.get("keep_logs", False)
 
@@ -151,6 +157,21 @@ def __calculate_lammps(self, outdir, **kwargs):
             "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.
@@ -160,17 +181,25 @@ 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(
-                    filepath, "in.bgrid.python"
+                    filepath,
+                    "in.bgrid_n{0}.python".format(
+                        len(set(self._atoms.numbers))
+                    ),
                 )
-
         # Do the LAMMPS calculation and clean up.
         lmp.file(self.parameters.lammps_compute_file)
 
@@ -274,6 +303,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:

mala/descriptors/in.bgrid.twoelements.python → mala/descriptors/in.bgrid_n2.python

@@ -11,16 +11,18 @@ 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
 
 #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
@@ -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}*${radelem1}
 # for water
-variable rcutneigh equal 4.0*${rcutfac}*${radelem1}
+#variable rcutneigh equal 4.0*${rcutfac}*${radelem1}
 
 pair_style zero ${rcutneigh}
 pair_coeff * *

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

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 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 	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