Implement Atomic forces

examples/in.numdiff

@@ -0,0 +1,89 @@
+# example with unique grid type
+
+variable         nfreq equal 1 # how frequently to run fix (should be every step)
+variable         ngridx equal 18 # how many grid points
+variable         ngridy equal 18 # how many grid points
+variable         ngridz equal 27 # how many grid points
+variable         nthermo equal 1 #frequency for printing thermodynamic data
+variable         nrep index 3 # number of repeated unit cells for Ca
+variable         a index 4.1 # fake lattice constant
+variable         fdelta equal 5.0e-7
+variable         fdeltam equal -1.0e-6
+
+units            metal
+atom_modify      map hash
+
+
+boundary         p p p
+#boundary         f f f
+
+read_data        Be_snapshot1.data
+#make a fictitious system to test grid force calculations
+mass             1 9.0
+
+
+pair_style       zero 6.0
+pair_coeff       * *
+
+group            acegroup type 1
+group            first id 1
+group            second id 2
+variable         rcutfac equal 6.0 # define rcutfac for pairstyle_zero (must be bigger than radial cutoffs in coupling_coefficients_Be.yace
+
+variable         ace_options string "coupling_coefficients_Be.yace ugridtype 1"
+
+pair_style       zero ${rcutfac}
+pair_coeff       * * 
+#
+timestep         0.5e-3
+neighbor         0.3 bin
+neigh_modify     every 1 delay 0 check yes
+
+compute          mygrid all pace/grid  grid ${ngridx} ${ngridy} ${ngridz} ${ace_options}
+
+#define lammps python command
+python           pre_force_callback file mala_betas.py
+#python          pre_force_callback file dummy_betas.py
+
+fix             4 all python/acegridforce ${nfreq} pre_force pre_force_callback grid ${ngridx} ${ngridy} ${ngridz} ${ace_options} 
+
+thermo          ${nthermo}
+variable        fsq atom fx^2+fy^2+fz^2
+compute         favsq all reduce ave v_fsq
+
+thermo_style    custom step temp pe etotal press c_favsq
+
+
+dump            1 all custom 1 dump.*.mala id type x y z fx fy fz
+run             0
+
+displace_atoms  first move ${fdelta} 0.0 0.0 units box
+undump 1
+reset_timestep  1
+dump 1 all custom 1 dump.*.mala id type x y z fx fy fz
+run             0
+displace_atoms  first move ${fdeltam} 0.0 0.0 units box
+undump 1
+reset_timestep  2
+dump 1 all custom 1 dump.*.mala id type x y z fx fy fz
+run             0
+displace_atoms  first move ${fdelta} ${fdelta} 0.0 units box
+undump 1
+reset_timestep  3
+dump 1 all custom 1 dump.*.mala id type x y z fx fy fz
+run             0
+displace_atoms  first move 0.0 ${fdeltam} 0.0 units box
+undump 1
+reset_timestep  4
+dump 1 all custom 1 dump.*.mala id type x y z fx fy fz
+run             0
+displace_atoms  first move 0.0 ${fdelta} ${fdelta} units box
+undump 1
+reset_timestep  5
+dump 1 all custom 1 dump.*.mala id type x y z fx fy fz
+run             0
+displace_atoms  first move 0.0 0.0 ${fdeltam} units box
+undump 1
+reset_timestep  6
+dump 1 all custom 1 dump.*.mala id type x y z fx fy fz
+run             0

examples/mala_betas.py

@@ -0,0 +1,150 @@
+import os
+
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+os.environ["OMP_NUM_THREADS"] = "1"
+from ase.io import read, write
+from ase import Atoms
+import mala
+from lammps import lammps
+import torch
+import numpy as np
+
+torch.set_num_threads(1)
+from mala.descriptors.lammps_utils import extract_compute_np
+from mala.datahandling.data_repo import data_repo_path
+
+data_path = os.path.join(data_repo_path, "Be2")
+
+
+def run_prediction(backprop=False, atoms=None, pass_descriptors=None):
+    """
+    This just runs a regular MALA prediction for a two-atom Beryllium model.
+    """
+    parameters, network, data_handler, predictor = mala.Predictor.load_run(
+        "Be_ACE_model_FULLSCALING"
+    )
+
+    parameters.targets.target_type = "LDOS"
+    parameters.targets.ldos_gridsize = 11
+    parameters.targets.ldos_gridspacing_ev = 2.5
+    parameters.targets.ldos_gridoffset_ev = -5
+
+    parameters.descriptors.descriptor_type = "ACE"
+
+    grd_loc = [18, 18, 27]
+    if type(atoms) == str:
+        atoms = read(atoms)
+    else:
+        atoms = atoms
+    predictor.parameters.inference_data_grid = grd_loc
+    assert atoms != None, "need to supply ase atoms object or file"
+    if type(atoms) == str:
+        atoms = read(atoms)
+    else:
+        atoms = atoms
+    ldos = predictor.predict_for_atoms(
+        atoms, save_grads=backprop, pass_descriptors=pass_descriptors
+    )
+    ldos_calculator: mala.LDOS = predictor.target_calculator
+    ldos_calculator.read_from_array(ldos)
+    ldos_calculator.read_additional_calculation_data(
+        os.path.join(data_path, "Be_snapshot1.out")
+    )
+    return ldos, ldos_calculator, parameters, predictor
+
+
+# boxlo, boxhi, xy, yz, xz, periodicity, box_change
+def lammps_box_2_ASE_cell(lmpbox):
+    Lx = lmpbox[1][0] - lmpbox[0][0]
+    Ly = lmpbox[1][1] - lmpbox[0][1]
+    Lz = lmpbox[1][2] - lmpbox[0][2]
+    xy = lmpbox[2]
+    yz = lmpbox[3]
+    xz = lmpbox[4]
+    a = [Lx, 0, 0]
+    b = [xy, Ly, 0]
+    c = [xz, yz, Lz]
+    cel = [a, b, c]
+    return cel
+
+
+def lammps_2_ase_atoms(lmp, typ_map):
+    cell = lammps_box_2_ASE_cell(lmp.extract_box())
+    x = lmp.extract_atom("x")
+    natoms = lmp.get_natoms()
+    pos = np.array([[x[i][0], x[i][1], x[i][2]] for i in range(natoms)])
+    # Extract atom types
+    atom_types = lmp.extract_atom("type")
+    # Convert atom types to NumPy array
+    atom_types_lst = [atom_types[i] for i in range(natoms)]
+    atom_syms = [typ_map[typi] for typi in atom_types_lst]
+    atoms = Atoms(atom_syms)
+    atoms.positions = pos
+    atoms.set_cell(cell)
+    atoms.set_pbc(True)  # assume pbc
+    return atoms
+
+
+def pre_force_callback(lmp):
+    # gc.collect()
+    L = lammps(ptr=lmp)
+    """
+    Test whether backpropagation works. To this end, the entire forces are
+    computed, and then backpropagated through the network.
+    """
+    # Only compute a specific part of the forces.
+    atoms = lammps_2_ase_atoms(L, typ_map={1: "Be"})
+    write("test_mala_lammps.vasp", atoms)
+    local_size = (18, 18, 27)
+    nx, ny, nz = (18, 18, 27)
+    feature_length = 36  # 5
+    fingerprint_length = feature_length + 3
+    ace_descriptors_np = extract_compute_np(
+        L,
+        "mygrid",
+        0,
+        2,
+        (nz, ny, nx, fingerprint_length),
+        use_fp64=False,
+    )
+    ace_descriptors_np = ace_descriptors_np.transpose([2, 1, 0, 3])
+    print(fingerprint_length, np.shape(ace_descriptors_np))
+    pass_descriptors = (
+        ace_descriptors_np,
+        local_size,
+        fingerprint_length,
+        True,
+    )
+    print("ace descs", ace_descriptors_np[0][0][0])
+    print("positions", atoms.positions)
+    ldos, ldos_calculator, parameters, predictor = run_prediction(
+        backprop=True, atoms=atoms, pass_descriptors=pass_descriptors
+    )
+    ldos_calculator.debug_forces_flag = "band_energy"
+    ldos_calculator.setup_for_forces(predictor)
+    ldos_calculator.read_from_array(ldos)
+    ldos_calculator.read_additional_calculation_data(
+        os.path.join(data_path, "Be_snapshot1.out")
+    )
+    mala_forces = ldos_calculator.atomic_forces.copy()
+    # energy attempt
+    eng = ldos_calculator.band_energy
+    # L.fix_external_set_energy_global('5', eng)
+    # end energy attempt
+    mala_forces = np.nan_to_num(mala_forces)
+    mala_test = mala_forces.reshape(27, 18, 18, feature_length)
+    # mala_test = mala_forces.reshape(18,18,27,feature_length)
+    mala_test = mala_test.transpose([2, 1, 0, 3])
+    print("mala_betas", mala_test[0][0][0])
+    print(
+        "mala force coeffs info:",
+        mala_forces.shape,
+        mala_test.shape,
+        np.amax(mala_forces),
+        np.mean(mala_forces),
+    )
+    print('mala "energy"', eng)
+    mala_2_lammps = mala_test.flatten()
+    L.close()
+    return np.ascontiguousarray(mala_2_lammps)

external_modules/total_energy_module/total_energy.f90

@@ -913,7 +913,8 @@ SUBROUTINE v_xc_wrapper(rho_of_r, rho_of_g, rho_core, rhog_core,&
 
 END SUBROUTINE v_xc_wrapper
 
-SUBROUTINE v_h_wrapper(rhog, ehart, charge, v,  nnr_in, nspin_in, ngm_in)
+
+SUBROUTINE v_h_wrapper(v, nnr_in, nspin_in)
   !----------------------------------------------------------------------------
   !  Derived from Quantum Espresso code
   !! author: Paolo Giannozzi
@@ -931,20 +932,22 @@ SUBROUTINE v_h_wrapper(rhog, ehart, charge, v,  nnr_in, nspin_in, ngm_in)
   USE lsda_mod,    ONLY : nspin
   USE kinds,       ONLY : DP
   USE fft_base,    ONLY : dfftp
+!  USE fft_rho,       ONLY : rho_r2g
+  USE scf,           ONLY : rho
 
   IMPLICIT NONE
 
-  INTEGER,     INTENT(IN)  :: ngm_in, nnr_in, nspin_in
-  DOUBLE COMPLEX, INTENT(IN)  :: rhog(ngm_in)
+  INTEGER,     INTENT(IN)  :: nnr_in, nspin_in
   DOUBLE PRECISION,  INTENT(INOUT) :: v(nnr_in,nspin_in)
-  DOUBLE PRECISION,    INTENT(OUT) :: ehart, charge
 
-  ! Check consistency of dimensions
+  DOUBLE PRECISION :: ehart, charge
+
+    ! Check consistency of dimensions
   IF (nnr_in /= dfftp%nnr) STOP "*** nnr provided to v_h_wrapper() does not match dfftp%nnr"
   IF (nspin_in /= nspin) STOP "*** nspin provided to v_h_wrapper() does not mach lsda_mod%nspin"
-  IF (ngm_in /= ngm) STOP "*** ngm provided to v_h_wrapper() does not match gvect%ngm"
 
-  CALL v_h(rhog, ehart, charge, v)
+!  CALL rho_r2g(dfftp, rho_of_r, rho%of_g)
+  CALL v_h(rho%of_g(:,1), ehart, charge, v)
 
 END SUBROUTINE v_h_wrapper
 

mala/common/parameters.py

@@ -804,6 +804,7 @@ def __init__(self):
         self.rdf_parameters = {"number_of_bins": 500, "rMax": "mic"}
         self.tpcf_parameters = {"number_of_bins": 20, "rMax": "mic"}
         self.ssf_parameters = {"number_of_bins": 100, "kMax": 12.0}
+        self.delta_forces = 0.001
         self.assume_two_dimensional = False
 
     @property