Add density cutoff option to fix FDBM for all-electron densities.

doc/sphinx/source/tutorials/afmulator-tutorial.md

@@ -222,6 +222,7 @@ The full-density-based model (FDBM) can provide a better approximation of the Pa
 In order to use the FDBM, more input files and parameters are required:
 ```python
 from ppafm.ocl.AFMulator import AFMulator, HartreePotential, ElectronDensity, TipDensity
+
 pot, xyzs, Zs = HartreePotential.from_file("LOCPOT.xsf", scale=-1) # Sample Hartree potential
 rho_sample, _, _ = ElectronDensity.from_file("CHGCAR.xsf") # Sample electron density
 rho_tip, xyzs_tip, Zs_tip = TipDensity.from_file("density_CO.xsf") # Tip electron density
@@ -242,11 +243,22 @@ Above, we first load the input files.
 In addition to the Hartree potential we also load the sample electron density `rho_sample`, and two densities for the tip, the full electron density `rho_tip`, and a delta density `rho_tip_delta`.
 When constructing the `AFMulator`, we use both of the densities.
 In this case the `rho=rho_tip` is used for the Pauli density overlap intergral, and the `rho_delta=rho_tip_delta` is used in the electrostatics calculation as in the previous section.
+
+````{warning}
+All-electron densities from codes such as FHI-aims can sometimes produce severe artifacts when used with the FDBM.
+A simple work around is to simply cutoff extremely high values of the electron density around the nuclei, which can be done using the {meth}`.clamp` method for the electron density grid:
+```python
+rho_sample = rho_sample.clamp(maximum=100)
+```
+A value of 100 for the cutoff is usually safe to use, although sometimes when both the prefactor ``'A_pauli'`` and exponent ``'B_pauli'`` are large, a lower cutoff may be required.
+````
+
 The delta density can be separately calculated and loaded from a file as above, or it can be estimated by subtracting the valence electrons from the full electron density:
 ```python
 rho_tip_delta = rho_tip.subCores(xyzs_tip, Zs_tip, Rcore=0.7)
 ```
 
+
 We also specify the two parameters of the Pauli overlap integral, the prefactor `A_pauli` and the exponent `B_pauli`, as well as the DFT-D3 parameters based on the functional name via `d3_params`.
 In this case we suppose that the electron densities were calculated using the `PBE` functional, so we use the corresponding parameters.
 Other predefined parameter sets can be found in the description of the method {meth}`.add_dftd3`.
@@ -257,7 +269,3 @@ afm_images = afmulator(xyzs, Zs, qs=pot, rho_sample=rho_sample)
 ```
 
 For examples of using the FDBM, see the [pyridine example](https://github.com/Probe-Particle/ppafm/blob/main/examples/pyridineDensOverlap/run_gpu.py), as well as another [more advanced example](https://github.com/Probe-Particle/ppafm/blob/main/examples/paper_figure/run_simulation.py) where multiple simulations are performed using all of the different force field models described here.
-
-```{warning}
-Electron densities from FHI-aims are known to not work well with the FDBM. This simulation model is generally less explored compared to the other ones, so expect more problems. Known working configurations have used densities calculated with VASP.
-```

ppafm/cli/conv_rho.py

@@ -27,6 +27,7 @@ def main():
     parser.add_argument('-o', '--output', action='store', default='pauli', help='Name of output energy/force files.')
     parser.add_argument('--saveDebugXsfs', action='store_true', help='Save auxiliary xsf files for debugging.')
     parser.add_argument('--no_negative_check', action='store_true', help='Input density files may contain negative voxels. This is handled by default by setting negative values to zero and rescaling the density so that the total charge is conserved. Setting this option disables the check.' )
+    parser.add_argument("--density_cutoff", action="store", default=None, type=float, help="Apply a cutoff to the electron densities to cut out high values. Helpful when using all-electron densities where extremely high values at the nuclei positions can cause artifacts in the resulting simulations. In these cases, 100 is usually a safe value to use, although sometimes when both the prefactor 'Apauli' and exponent 'Bpauli' are large, a lower cutoff may be required.")
     parser.add_arguments(['output_format', 'energy', 'Apauli', 'Bpauli'])
     # fmt: on
 
@@ -57,6 +58,11 @@ def main():
             io.save_scal_field("sample_density_pow_%03.3f.xsf" % args.Bpauli, rho_sample, lvec_sample, data_format=args.output_format, head=head_sample)
             io.save_scal_field("tip_density_pow_%03.3f.xsf" % args.Bpauli, rho_tip, lvec_tip, data_format=args.output_format, head=head_tip)
 
+    if args.density_cutoff:
+        print(f">>> Applying a density cutoff of {args.density_cutoff} to sample and tip electron densities.")
+        rho_sample[rho_sample > args.density_cutoff] = args.density_cutoff
+        rho_tip[rho_tip > args.density_cutoff] = args.density_cutoff
+
     print(">>> Evaluating convolution E(R) = A*Integral_r ( rho_tip^B(r-R) * rho_sample^B(r) ) using FFT ... ")
     f_x, f_y, f_z, energy = fieldFFT.potential2forces_mem(rho_sample, lvec_sample, n_dim_sample, rho=rho_tip, doForce=True, doPot=True, deleteV=True)
 

ppafm/cli/gui/ppafm_gui.py

@@ -74,6 +74,7 @@
     "Rotation": "Rotation: Set sample counter-clockwise rotation angle around center of atom coordinates.",
     "fdbm_V0": "FDBM V0: Prefactor in Pauli interaction integral in the full-density based model.",
     "fdbm_alpha": "FDBM alpha: Exponent in Pauli interaction integral in the full-density based model.",
+    "fdbm_cutoff": "Density cutoff: cut out high values of the electron density. Useful when working with all-electron densities where extremely high values can arise near nuclei, causing artifacts in the images.",
     "PBCz": "z periodicity: When checked, the lattice is also periodic in z direction. This is usually not required, since the scan is aligned with the xy direction.",
     "PBC": "Periodic Boundaries: Lattice vectors for periodic images of atoms. Does not affect electrostatics calculated from a Hartree potential file, which is always assumed to be periodic.",
     "k": "k: Cantilever spring constant. Only appears as a scaling constant.",
@@ -123,11 +124,13 @@ def __init__(self, input_files=None, device=0, verbose=0):
         self.Zs = None
         self.qs = None
         self.rho_sample = None
+        self.rho_sample_original = None  # Stores the original density when a cutoff is applied
         self.rho_tip = None
         self.rho_tip_delta = None
         self.sample_lvec = None
         self.rot = np.eye(3)
         self.df_points = []
+        self.previous_cutoff = None
 
         # Initialize OpenCL environment on chosen device and create an afmulator instance to use for simulations
         oclu.init_env(device)
@@ -141,6 +144,7 @@ def __init__(self, input_files=None, device=0, verbose=0):
         self.afmulator.forcefield.verbose = verbose
         self.afmulator.scanner.verbose = verbose
         FFcl.bRuntime = verbose > 1
+        self.afmulator.bRuntime = verbose > 1
 
         # --- init QtMain
         QtWidgets.QMainWindow.__init__(self)
@@ -302,17 +306,42 @@ def updateParams(self, preset_none=True):
         if self.verbose > 0:
             print("updateParams", Q, sigma, multipole, tipStiffness, tipR0, use_point_charge, A_pauli, B_pauli, type(self.qs))
 
-        self.afmulator.iZPP = int(self.bxZPP.value())
         if self.rho_sample is not None:  # Use FDBM
-            if self.afmulator.B_pauli != B_pauli:
-                self.afmulator.setRho(self.rho_tip, B_pauli=B_pauli)
+
+            if self.bxCutoffEnable.isChecked():
+                self.bxCutoffValue.setDisabled(False)
+                cutoff = float(self.bxCutoffValue.value())
+            else:
+                self.bxCutoffValue.setDisabled(True)
+                cutoff = None
+
+            if cutoff is None or self.previous_cutoff is None:
+                cutoff_changed = cutoff != self.previous_cutoff
+            else:
+                cutoff_changed = not np.isclose(cutoff, self.previous_cutoff)
+            self.previous_cutoff = cutoff
+
+            if cutoff_changed:
+                if cutoff is not None:
+                    self.rho_sample = self.rho_sample_original.clamp(maximum=cutoff, in_place=False)
+                else:
+                    self.rho_sample = self.rho_sample_original
+
+            if self.afmulator.B_pauli != B_pauli or cutoff_changed:
+                if cutoff is not None:
+                    rho_tip = self.rho_tip.clamp(maximum=cutoff, in_place=False)
+                else:
+                    rho_tip = self.rho_tip
+                self.afmulator.setRho(rho_tip, B_pauli=B_pauli)
         else:
             if self.rho_tip is None:  # else we just keep the tip density loaded from file
                 if use_point_charge:
                     self.afmulator.setQs(Qs, QZs)
                     self.afmulator.setRho(None, sigma)
                 elif isinstance(self.qs, HartreePotential):
                     self.afmulator.setRho({multipole: Q}, sigma)
+
+        self.afmulator.iZPP = int(self.bxZPP.value())
         self.afmulator.scanner.stiffness = np.array(tipStiffness, dtype=np.float32) / -PPU.eVA_Nm
         self.afmulator.tipR0 = tipR0
         self.afmulator.A_pauli = A_pauli
@@ -471,6 +500,8 @@ def loadInput(self, main_input, rho_sample=None, rho_tip=None, rho_tip_delta=Non
             self.qs.release()
         if self.rho_sample is not None:
             self.rho_sample.release()
+        if self.rho_sample_original is not None:
+            self.rho_sample_original.release()
         if self.rho_tip is not None:
             self.rho_tip.release()
         if self.rho_tip_delta is not None:
@@ -502,6 +533,7 @@ def loadInput(self, main_input, rho_sample=None, rho_tip=None, rho_tip_delta=Non
         # Load auxiliary files (if any)
         if rho_sample:
             self.rho_sample, _, _ = ElectronDensity.from_file(rho_sample)
+            self.rho_sample_original = self.rho_sample
         else:
             self.rho_sample = None
         if rho_tip:
@@ -577,9 +609,14 @@ def loadInput(self, main_input, rho_sample=None, rho_tip=None, rho_tip_delta=Non
         if self.rho_sample is None:
             self.bxV0.setDisabled(True)
             self.bxAlpha.setDisabled(True)
+            self.bxCutoffEnable.setDisabled(True)
+            self.bxCutoffEnable.setChecked(False)
+            self.bxCutoffValue.setDisabled(True)
         else:
             self.bxV0.setDisabled(False)
             self.bxAlpha.setDisabled(False)
+            self.bxCutoffEnable.setDisabled(False)
+            self.previous_cutoff = None
 
         # Create geometry editor widget
         self.createGeomEditor()
@@ -1000,6 +1037,18 @@ def _create_probe_settings_ui(self):
         l.addWidget(lb)
         self.bxAlpha = _spin_box((0.1, 9.9), 1.0, 0.02, self.updateParams, TTips["fdbm_alpha"], l)
 
+        vb = QtWidgets.QHBoxLayout()
+        self.l0.addLayout(vb)
+        self.bxCutoffEnable = QtWidgets.QCheckBox()
+        self.bxCutoffEnable.setChecked(False)
+        self.bxCutoffEnable.toggled.connect(self.updateParams)
+        vb.addWidget(self.bxCutoffEnable)
+        lb = QtWidgets.QLabel("Density cutoff")
+        lb.setToolTip(TTips["fdbm_cutoff"])
+        vb.addWidget(lb)
+        self.bxCutoffValue = _spin_box((0.1, 10000.0), 100.0, 10.0, self.updateParams, TTips["fdbm_cutoff"], vb)
+        self.bxCutoffValue.setDisabled(True)
+
     def _create_scan_settings_ui(self):
         # Title
         vb = QtWidgets.QHBoxLayout()

ppafm/ml/Generator.py

@@ -342,6 +342,9 @@ class GeneratorAFMtrainer:
             is FDBM, where it is used for calculating the electrostatic interaction.
         ignore_elements: list of int. Atomic numbers of elements to ignore in scan window distance and position calculations.
             Useful for example for ignoring surface slab atoms in centering the scan window.
+        density_cutoff: float or None. If not None, apply a cutoff to electron densities in the FDBM Pauli integral. Useful when
+            working with all-electron densities where large density values near nuclei can cause artifacts in the resulting images.
+            Ignored when sim_type is not ``'FDBM'``.
     """
 
     bRuntime = False
@@ -361,6 +364,7 @@ def __init__(
         rhos=None,
         rho_deltas=None,
         ignore_elements=[],
+        density_cutoff=None,
     ):
         self.afmulator = afmulator
         self.aux_maps = aux_maps
@@ -370,9 +374,11 @@ def __init__(
         self.distAbove = distAbove
         self.distAboveActive = distAbove
         self.iZPPs = iZPPs
-        self._prepareBuffers(rhos, rho_deltas)
         self.ignore_elements = ignore_elements
 
+        self.density_cutoff = density_cutoff if self.sim_type == "fdbm" else None
+        self._prepare_tip_buffers(rhos, rho_deltas, self.density_cutoff)
+
         if Qs is None or QZs is None:
             if self.sim_type == "lj+pc":
                 raise ValueError("Both Qs and QZs should be specified when using point-charge electrostatics.")
@@ -389,29 +395,33 @@ def __init__(
         self.df_steps = self.afmulator.df_steps
         self.z_size = self.scan_dim[2] - self.df_steps + 1
 
-    def _prepareBuffers(self, rhos=None, rho_deltas=None):
+    def _prepare_tip_buffers(self, rhos=None, rho_deltas=None, density_cutoff=None):
         if rhos is None:
-            self.rhos = self.ffts = [(None, None)] * len(self.iZPPs)
+            self.rhos_tip = self.rhos_tip_delta = self.ffts_tip = self.ffts_tip_delta = [None] * len(self.iZPPs)
             return
         elif len(rhos) != len(self.iZPPs):
             raise ValueError(f"The length of rhos ({len(rhos)}) does not match the length of iZPPs ({len(self.iZPPs)})")
         if rho_deltas is not None and (len(rhos) != len(rho_deltas)):
             raise ValueError(f"The length of rhos ({len(rhos)}) does not match the length of rho_deltas ({len(rho_deltas)})")
-        self.rhos = []
-        self.ffts = []
-        for i in range(len(rhos)):
-            self.afmulator.setRho(rhos[i], sigma=self.afmulator.sigma, B_pauli=self.afmulator.B_pauli)
-            rhos_ = [self.afmulator.forcefield.rho]
-            ffts_ = [self.afmulator.forcefield.fft_corr]
-            if rho_deltas is not None:
-                self.afmulator.setRhoDelta(rho_deltas[i])
-                rhos_.append(self.afmulator.forcefield.rho_delta)
-                ffts_.append(self.afmulator.forcefield.fft_corr_delta)
-            else:
-                rhos_.append(None)
-                ffts_.append(None)
-            self.rhos.append(rhos_)
-            self.ffts.append(ffts_)
+        self.rhos_tip = []
+        self.ffts_tip = []
+        self._rhos_original = []
+        for rho in rhos:
+            if density_cutoff is not None:
+                rho = rho.clamp(maximum=density_cutoff, in_place=False)
+            self.afmulator.setRho(rho, sigma=self.afmulator.sigma, B_pauli=self.afmulator.B_pauli)
+            self.rhos_tip.append(self.afmulator.forcefield.rho)
+            self.ffts_tip.append(self.afmulator.forcefield.fft_corr)
+            self._rhos_original.append(rho)
+        if rho_deltas is not None:
+            self.rhos_tip_delta = []
+            self.ffts_tip_delta = []
+            for rho_delta in rho_deltas:
+                self.afmulator.setRhoDelta(rho_delta)
+                self.rhos_tip_delta.append(self.afmulator.forcefield.rho_delta)
+                self.ffts_tip_delta.append(self.afmulator.forcefield.fft_corr_delta)
+        else:
+            self.rhos_tip_delta = self.ffts_tip_delta = [None] * len(self.iZPPs)
 
     def __iter__(self):
         self.sample_dict = {}
@@ -471,14 +481,16 @@ def __next__(self):
                 print(f"Sample {s} preparation time [s]: {time.perf_counter() - sample_start}")
 
             # Get AFM
-            for i, (iZPP, rho, fft, Qs, QZs) in enumerate(zip(self.iZPPs, self.rhos, self.ffts, self.Qs, self.QZs)):  # Loop over different tips
+            for i, (iZPP, rho_tip, rho_tip_delta, fft_tip, fft_tip_delta, Qs, QZs) in enumerate(
+                zip(self.iZPPs, self.rhos_tip, self.rhos_tip_delta, self.ffts_tip, self.ffts_tip_delta, self.Qs, self.QZs)
+            ):  # Loop over different tips
                 # Set interaction parameters
                 self.afmulator.iZPP = iZPP
                 self.afmulator.setQs(Qs, QZs)
-                self.afmulator.forcefield.rho = rho[0]
-                self.afmulator.forcefield.fft_corr = fft[0]
-                self.afmulator.forcefield.rho_delta = rho[1]
-                self.afmulator.forcefield.fft_corr_delta = fft[1]
+                self.afmulator.forcefield.rho = rho_tip
+                self.afmulator.forcefield.fft_corr = fft_tip
+                self.afmulator.forcefield.rho_delta = rho_tip_delta
+                self.afmulator.forcefield.fft_corr_delta = fft_tip_delta
                 self.sample_dict["REAs"] = PPU.getAtomsREA(self.afmulator.iZPP, self.sample_dict["Zs"], self.afmulator.typeParams, alphaFac=-1.0)
 
                 # Make sure tip-sample distance is right
@@ -563,6 +575,9 @@ def _load_next_sample(self):
         if "rot" not in sample_dict:
             sample_dict["rot"] = np.eye(3)
 
+        if self.density_cutoff is not None:
+            sample_dict["rho_sample"] = sample_dict["rho_sample"].clamp(maximum=self.density_cutoff, in_place=False)
+
         return sample_dict
 
     def __len__(self):
@@ -615,6 +630,26 @@ def handle_distance(self):
             (sw[1][0], sw[1][1], z_min + self.scan_size[2]),
         )
 
+    def set_fdbm_parameters(self, A_pauli, B_pauli):
+        """
+        Set the Pauli integral parameters in an FDBM simulation. If set simulation type is not FDBM, does nothing.
+
+        Arguments:
+            A_pauli: float. Integral prefactor.
+            B_pauli: float. Integrant exponent.
+        """
+        if self.sim_type != "fdbm":
+            return
+        self.afmulator.A_pauli = A_pauli
+        new_rhos = []
+        new_ffts = []
+        for rho in self._rhos_original:
+            self.afmulator.setRho(rho, sigma=self.afmulator.sigma, B_pauli=B_pauli)
+            new_rhos.append(self.afmulator.forcefield.rho)
+            new_ffts.append(self.afmulator.forcefield.fft_corr)
+        self.rhos_tip = new_rhos
+        self.ffts_tip = new_ffts
+
     def randomize_df_steps(self, minimum=4, maximum=20):
         """Randomize oscillation amplitude by randomizing the number of steps in df convolution.