diff --git a/examples/ex19_atomic_forces.py b/examples/ex19_atomic_forces.py new file mode 100644 index 000000000..bb0fd7242 --- /dev/null +++ b/examples/ex19_atomic_forces.py @@ -0,0 +1,500 @@ +import os + +from ase.io import read +import mala +from mala import printout +import numpy as np +import torch + +from mala.datahandling.data_repo import data_repo_path + +data_path = os.path.join(data_repo_path, "Be2") + +""" +ex19_at.py: Show how a prediction can be made using MALA. +Using nothing more then the trained network and atomic configurations, +predictions can be made. +""" + + +# Uses a network to make a prediction. +assert os.path.exists("Be_model.zip"), "Be model missing, run ex01 first." + + +def run_prediction(backprop=False): + """ + This just runs a regular MALA prediction for a two-atom Beryllium model. + """ + parameters, network, data_handler, predictor = mala.Predictor.load_run( + "Be_model" + ) + + 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 = "Bispectrum" + parameters.descriptors.bispectrum_twojmax = 10 + parameters.descriptors.bispectrum_cutoff = 4.67637 + parameters.targets.pseudopotential_path = data_path + + atoms = read(os.path.join(data_path, "Be_snapshot3.out")) + ldos = predictor.predict_for_atoms(atoms, save_grads=backprop) + ldos_calculator: mala.LDOS = predictor.target_calculator + ldos_calculator.read_from_array(ldos) + return ldos, ldos_calculator, parameters, predictor + + +def band_energy_contribution(): + """ + Test the band energy contribution to the forces by calculating it + via finite differences and analytically. This should return arrays of 1s. + Note that the LDOS is only shifted at one point in space, but since the + band energy is a not a spatially resolved quantity, the forces should + be the same for all grid points. + """ + ldos, ldos_calculator, parameters, predictor = run_prediction() + + # Only compute a specific part of the forces. + ldos_calculator.debug_forces_flag = "band_energy" + mala_forces = ldos_calculator.atomic_forces.copy() + + delta_numerical = 1e-6 + numerical_forces = [] + + for i in range(0, parameters.targets.ldos_gridsize): + ldos_plus = ldos.copy() + ldos_plus[0, i] += delta_numerical * 0.5 + ldos_calculator.read_from_array(ldos_plus) + derivative_plus = ldos_calculator.band_energy + + ldos_minus = ldos.copy() + ldos_minus[0, i] -= delta_numerical * 0.5 + ldos_calculator.read_from_array(ldos_minus) + derivative_minus = ldos_calculator.band_energy + + numerical_forces.append( + (derivative_plus - derivative_minus) / delta_numerical + ) + + print("FORCE TEST: Band energy contribution.") + print(mala_forces[0, :] / np.array(numerical_forces)) + print(mala_forces[2000, :] / np.array(numerical_forces)) + print(mala_forces[4000, :] / np.array(numerical_forces)) + + +def entropy_contribution(): + """ + Test the entropy contribution to the forces by calculating it + via finite differences and analytically. This should return arrays of 1s, + it will most likely not for many test settings due to the fact that the + entropy contribution is VERY small. + Note that the LDOS is only shifted at one point in space, but since the + entropy is a not a spatially resolved quantity, the forces should + be the same for all grid points. + """ + ldos, ldos_calculator, parameters, predictor = run_prediction() + + # Only compute a specific part of the forces. + ldos_calculator.debug_forces_flag = "entropy_contribution" + mala_forces = ldos_calculator.atomic_forces.copy() + + delta_numerical = 1e-8 + numerical_forces = [] + + for i in range(0, parameters.targets.ldos_gridsize): + ldos_plus = ldos.copy() + ldos_plus[0, i] += delta_numerical * 0.5 + ldos_calculator.read_from_array(ldos_plus) + derivative_plus = ldos_calculator.entropy_contribution + + ldos_minus = ldos.copy() + ldos_minus[0, i] -= delta_numerical * 0.5 + ldos_calculator.read_from_array(ldos_minus) + derivative_minus = ldos_calculator.entropy_contribution + + numerical_forces.append( + (derivative_plus - derivative_minus) / delta_numerical + ) + + print("FORCE TEST: Entropy contribution.") + print(mala_forces[0, :] / np.array(numerical_forces)) + print(mala_forces[2000, :] / np.array(numerical_forces)) + print(mala_forces[4000, :] / np.array(numerical_forces)) + + +def hartree_contribution(): + """ + Test the Hartree contribution to the forces by calculating it + via finite differences and analytically. This should return arrays of 1s. + Since the Hartree energy is dependent on the electronic density, the + derivative will vary across space, and we have to compute the forces + at multiple points. + """ + ldos, ldos_calculator, parameters, predictor = run_prediction() + + # Only compute a specific part of the forces. + ldos_calculator.debug_forces_flag = "hartree" + mala_forces = ldos_calculator.atomic_forces.copy() + + delta_numerical = 1e-6 + points = [0, 2000, 4000] + + print("FORCE TEST: Hartree contribution.") + for point in points: + numerical_forces = [] + for i in range(0, parameters.targets.ldos_gridsize): + ldos_plus = ldos.copy() + ldos_plus[point, i] += delta_numerical * 0.5 + ldos_calculator.read_from_array(ldos_plus) + _, derivative_plus = ldos_calculator.get_total_energy( + return_energy_contributions=True + ) + derivative_plus = derivative_plus["e_hartree"] + + ldos_minus = ldos.copy() + ldos_minus[point, i] -= delta_numerical * 0.5 + ldos_calculator.read_from_array(ldos_minus) + _, derivative_minus = ldos_calculator.get_total_energy( + return_energy_contributions=True + ) + derivative_minus = derivative_minus["e_hartree"] + numerical_forces.append( + (derivative_plus - derivative_minus) / delta_numerical + ) + + print(mala_forces[point, :] / np.array(numerical_forces)) + + +def check_input_gradient_scaling(): + # Load a model, set paramters. + predictor: mala.Predictor + parameters: mala.Parameters + 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.targets.restrict_targets = None + parameters.descriptors.descriptors_contain_xyz = False + parameters.descriptors.descriptor_type = "ACE" + + # Compute descriptors, only do further computations for one point. + atoms1 = read("/home/fiedlerl/data/mala_data_repo/Be2/Be_snapshot1.out") + descriptors, ngrid = ( + predictor.data.descriptor_calculator.calculate_from_atoms( + atoms1, [18, 18, 27] + ) + ) + snap_descriptors = torch.from_numpy(np.array([descriptors[0, 0, 0]])) + snap_descriptors_work = snap_descriptors.clone() + snap_descriptors = predictor.data.input_data_scaler.transform( + snap_descriptors + ) + snap_descriptors.requires_grad = True + + # Forward pass through the network. + ldos = network(snap_descriptors) + d_d_d_B = [] + + # Compute "gradient". This is ACTUALLY the Jacobian matrix, i.e., the one + # we can compare with finite differences. + for i in range(0, 11): + if snap_descriptors.grad is not None: + snap_descriptors.grad.zero_() + + ldos[0, i].backward(retain_graph=True) + # d_d_d_B = snap_descriptors.grad.clone() + d_d_d_B.append( + predictor.data.input_data_scaler.inverse_transform_backpropagation( + snap_descriptors.grad.clone() + ) + ) + + # Just for debugging purposes, not part of the actual test. + # if False: + # # This would be the direct application of the autograd. + # # Autograd computes the vector Jacobian product, see + # # https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html + # # Here: compare same_as_sum with d_d_d_B_sum. + # ldos = network(snap_descriptors) + # snap_descriptors.grad.zero_() + # # print(ldos / ldos) + # ldos.backward(ldos / ldos, retain_graph=True) + # same_as_sum = ( + # predictor.data.input_data_scaler.inverse_transform_backpropagation( + # snap_descriptors.grad.clone() + # ) + # ) + # d_d_d_B_sum = d_d_d_B[0] + # for i in range(1, 11): + # d_d_d_B_sum += d_d_d_B[i] + + with torch.no_grad(): + # In case we want to compare different levels of finite differences. + for diff in [ + # 5.0e-1, + # 5.0e-2, + # 5.0e-3, + 5.0e-4, + # 5.0e-5, + # 5.0e-6, + ]: + + # Compute finite differences. j theoretically goes up to + # 36, but for a simple check this is completely enough. + for j in range(0, 5): + snap_descriptors_work[0, j] += diff + descriptors_scaled1 = snap_descriptors_work.clone() + ldos_1 = network( + predictor.data.input_data_scaler.transform( + descriptors_scaled1 + ) + ).clone() + + snap_descriptors_work[0, j] -= 2.0 * diff + descriptors_scaled2 = snap_descriptors_work.clone() + ldos_2 = network( + predictor.data.input_data_scaler.transform( + descriptors_scaled2 + ) + ).clone() + + # Comparison - we only compare the first component of the + # LDOS for brevity, but this works for all components. + snap_descriptors_work[0, j] += diff + force = -1.0 * ((ldos_2[0, 0] - ldos_1[0, 0]) / (2 * diff)) + print( + diff, + force.double().numpy(), + d_d_d_B[0][0, j], + force.double().numpy() / d_d_d_B[0][0, j], + ) + + +def check_output_gradient_scaling(): + # Load a model, set paramters. + predictor: mala.Predictor + parameters: mala.Parameters + 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.targets.restrict_targets = None + parameters.descriptors.descriptors_contain_xyz = False + parameters.descriptors.descriptor_type = "ACE" + + # Compute descriptors, only do further computations for one point. + atoms1 = read("/home/fiedlerl/data/mala_data_repo/Be2/Be_snapshot1.out") + descriptors, ngrid = ( + predictor.data.descriptor_calculator.calculate_from_atoms( + atoms1, [18, 18, 27] + ) + ) + snap_descriptors = torch.from_numpy(np.array([descriptors[0, 0, 0]])) + snap_descriptors_work = snap_descriptors.clone() + snap_descriptors = predictor.data.input_data_scaler.transform( + snap_descriptors + ) + snap_descriptors.requires_grad = True + + # Forward pass through the network. + # ldos = network(snap_descriptors) + # d_d_d_B = [] + + # Compute "gradient". This is ACTUALLY the Jacobian matrix, i.e., the one + # we can compare with finite differences. + # for i in range(0, 11): + # if snap_descriptors.grad is not None: + # snap_descriptors.grad.zero_() + # + # ldos[0, i].backward(retain_graph=True) + # # d_d_d_B = snap_descriptors.grad.clone() + # d_d_d_B.append( + # predictor.data.input_data_scaler.inverse_transform_backpropagation( + # snap_descriptors.grad.clone() + # ) + # ) + + if True: + # This would be the direct application of the autograd. + # Autograd computes the vector Jacobian product, see + # https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html + ldos = network(snap_descriptors) + ldos_unscaled = predictor.data.output_data_scaler.inverse_transform( + ldos, copy=True + ) + if snap_descriptors.grad is not None: + snap_descriptors.grad.zero_() + ldos.backward( + predictor.data.output_data_scaler.transform_backpropagation( + 2 * ldos_unscaled + ), + retain_graph=True, + ) + d_d_d_B = ( + predictor.data.input_data_scaler.inverse_transform_backpropagation( + snap_descriptors.grad.clone() + ) + ) + + with torch.no_grad(): + # In case we want to compare different levels of finite differences. + for diff in [ + # 5.0e-1, + # 5.0e-2, + # 5.0e-3, + 5.0e-4, + # 5.0e-5, + # 5.0e-6, + ]: + + # Compute finite differences. j theoretically goes up to + # 36, but for a simple check this is completely enough. + for j in range(0, 5): + snap_descriptors_work[0, j] += diff + descriptors_scaled1 = snap_descriptors_work.clone() + ldos_1 = predictor.data.output_data_scaler.inverse_transform( + network( + predictor.data.input_data_scaler.transform( + descriptors_scaled1 + ) + ).clone() + ) + energy_1 = 0 + for i in range(0, 11): + energy_1 += ldos_1[0, i] * ldos_1[0, i] + + snap_descriptors_work[0, j] -= 2.0 * diff + descriptors_scaled2 = snap_descriptors_work.clone() + ldos_2 = predictor.data.output_data_scaler.inverse_transform( + network( + predictor.data.input_data_scaler.transform( + descriptors_scaled2 + ) + ).clone() + ) + energy_2 = 0 + for i in range(0, 11): + energy_2 += ldos_2[0, i] * ldos_2[0, i] + + # Comparison - we only compare the first component of the + # LDOS for brevity, but this works for all components. + snap_descriptors_work[0, j] += diff + force = -1.0 * ((energy_2 - energy_1) / (2 * diff)) + print( + diff, + force.double().numpy(), + d_d_d_B[0, j], + force.double().numpy() / d_d_d_B[0, j], + ) + + +def check_backpropagation(): + predictor: mala.Predictor + parameters: mala.Parameters + 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.targets.restrict_targets = None + parameters.descriptors.descriptors_contain_xyz = False + parameters.descriptors.descriptor_type = "ACE" + + # Compute descriptors, only do further computations for one point. + atoms1 = read("/home/fiedlerl/data/mala_data_repo/Be2/Be_snapshot1.out") + descriptors, ngrid = ( + predictor.data.descriptor_calculator.calculate_from_atoms( + atoms1, [18, 18, 27] + ) + ) + snap_descriptors = torch.from_numpy( + descriptors.reshape([ngrid, descriptors.shape[-1]]) + ) + snap_descriptors_work = snap_descriptors.clone() + + ldos = predictor.predict_for_atoms( + atoms1, + save_grads=True, + pass_descriptors=[ + descriptors, + ngrid, + descriptors.shape[-1], + parameters.descriptors.descriptors_contain_xyz, + ], + ) + 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") + ) + ldos_calculator.debug_forces_flag = "band_energy" + ldos_calculator.setup_for_forces(predictor) + mala_forces = ldos_calculator.atomic_forces.copy() + + for point in [0, 2000, 4000]: + for diff in [ + # 5.0e-1, + # 5.0e-2, + # 5.0e-3, + 5.0e-4, + # 5.0e-5, + # 5.0e-6, + ]: + + # Compute finite differences. j theoretically goes up to + # 36, but for a simple check this is completely enough. + for j in range(0, 5): + snap_descriptors_work[point, j] += diff + descriptors_scaled1 = snap_descriptors_work.clone() + ldos_1 = predictor.data.output_data_scaler.inverse_transform( + network( + predictor.data.input_data_scaler.transform( + descriptors_scaled1 + ) + ), + as_numpy=True, + ).copy() + ldos_calculator.read_from_array(ldos_1) + energy_1 = ldos_calculator.band_energy + + snap_descriptors_work[point, j] -= 2.0 * diff + descriptors_scaled2 = snap_descriptors_work.clone() + ldos_2 = predictor.data.output_data_scaler.inverse_transform( + network( + predictor.data.input_data_scaler.transform( + descriptors_scaled2 + ) + ), + as_numpy=True, + ).copy() + ldos_calculator.read_from_array(ldos_2) + energy_2 = ldos_calculator.band_energy + + # Comparison - we only compare the first component of the + # LDOS for brevity, but this works for all components. + snap_descriptors_work[point, j] += diff + force = -1.0 * ((energy_2 - energy_1) / (2 * diff)) + print( + diff, + force, + mala_forces[point, j], + force / mala_forces[point, j], + ) + + +# check_input_gradient_scaling() +# check_output_gradient_scaling() +check_backpropagation() +# band_energy_contribution() +# entropy_contribution() +# hartree_contribution() diff --git a/examples/in.numdiff b/examples/in.numdiff new file mode 100644 index 000000000..f04583a34 --- /dev/null +++ b/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 diff --git a/examples/mala_betas.py b/examples/mala_betas.py new file mode 100644 index 000000000..7303609e8 --- /dev/null +++ b/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) diff --git a/external_modules/total_energy_module/total_energy.f90 b/external_modules/total_energy_module/total_energy.f90 index f1165e01e..3b6f8c63e 100644 --- a/external_modules/total_energy_module/total_energy.f90 +++ b/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 diff --git a/mala/common/parameters.py b/mala/common/parameters.py index 5fd2d551d..d77457e27 100644 --- a/mala/common/parameters.py +++ b/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 diff --git a/mala/datahandling/data_scaler.py b/mala/datahandling/data_scaler.py index 652a3b7c4..e9ef0804b 100644 --- a/mala/datahandling/data_scaler.py +++ b/mala/datahandling/data_scaler.py @@ -404,12 +404,16 @@ def transform(self, unscaled, copy=False): if self.scale_standard: scaled -= self.means scaled /= self.stds - unscaled.nan_to_num_(nan = 0.0) #account for possible 0 valued ACE descriptors + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors if self.scale_minmax: scaled -= self.mins scaled /= self.maxs - self.mins - unscaled.nan_to_num_(nan = 0.0) #account for possible 0 valued ACE descriptors + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors else: @@ -420,12 +424,102 @@ def transform(self, unscaled, copy=False): if self.scale_standard: scaled -= self.total_mean scaled /= self.total_std - scaled.nan_to_num_(nan = 0.0) #account for possible 0 valued ACE descriptors + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors if self.scale_minmax: scaled -= self.total_min scaled /= self.total_max - self.total_min - scaled.nan_to_num_(nan = 0.0) #account for possible 0 valued ACE descriptors + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors + + return scaled + + def transform_backpropagation(self, unscaled, copy=False): + """ + Transform data from unscaled to scaled during backpropagation. + + Unscaled means real world data, scaled means data as is used in + the network. Data is transformed in-place. + + Parameters + ---------- + unscaled : torch.Tensor + Real world data. + + copy : bool + If False, data is modified in-place. If True, a copy of the + data is modified. Default is False. + + Returns + ------- + scaled : torch.Tensor + Scaled data. + """ + if len(unscaled.size()) != 2: + raise ValueError( + "MALA DataScaler are designed for 2D-arrays, " + "while a {0}D-array has been provided.".format( + len(unscaled.size()) + ) + ) + + # Backward compatability. + if not hasattr(self, "scale_minmax") and hasattr(self, "scale_normal"): + self.scale_minmax = self.scale_normal + + # First we need to find out if we even have to do anything. + if self.scale_standard is False and self.scale_minmax is False: + pass + + elif self.cantransform is False: + raise Exception( + "Transformation cannot be done, this DataScaler " + "was never initialized" + ) + + # Perform the actual scaling, but use no_grad to make sure + # that the next couple of iterations stay untracked. + scaled = unscaled.clone() if copy else unscaled + + with torch.no_grad(): + if self.feature_wise: + + ########################## + # Feature-wise-scaling + ########################## + + if self.scale_standard: + scaled *= self.stds + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors + + if self.scale_minmax: + scaled *= self.maxs - self.mins + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors + + else: + + ########################## + # Total scaling + ########################## + + if self.scale_standard: + scaled *= self.total_std + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors + + if self.scale_minmax: + scaled *= self.total_max - self.total_min + scaled.nan_to_num_( + nan=0.0 + ) # account for possible 0 valued ACE descriptors return scaled @@ -517,6 +611,94 @@ def inverse_transform(self, scaled, copy=False, as_numpy=False): else: return unscaled + def inverse_transform_backpropagation( + self, scaled, copy=False, as_numpy=False + ): + """ + Transform data from scaled to unscaled during backpropagation. + + Unscaled means real world data, scaled means data as is used in + the network. + + Parameters + ---------- + scaled : torch.Tensor + Scaled data. + + as_numpy : bool + If True, a numpy array is returned, otherwise a torch tensor. + + copy : bool + If False, data is modified in-place. If True, a copy of the + data is modified. Default is False. + + Returns + ------- + unscaled : torch.Tensor + Real world data. + + """ + if len(scaled.size()) != 2: + raise ValueError( + "MALA DataScaler are designed for 2D-arrays, " + "while a {0}D-array has been provided.".format( + len(scaled.size()) + ) + ) + + # Backward compatability. + if not hasattr(self, "scale_minmax") and hasattr(self, "scale_normal"): + self.scale_minmax = self.scale_normal + + # Perform the actual scaling, but use no_grad to make sure + # that the next couple of iterations stay untracked. + unscaled = scaled.clone() if copy else scaled + + # First we need to find out if we even have to do anything. + if self.scale_standard is False and self.scale_minmax is False: + pass + + else: + if self.cantransform is False: + raise Exception( + "Backtransformation cannot be done, this " + "DataScaler was never initialized" + ) + + # Perform the actual scaling, but use no_grad to make sure + # that the next couple of iterations stay untracked. + with torch.no_grad(): + if self.feature_wise: + + ########################## + # Feature-wise-scaling + ########################## + + if self.scale_standard: + unscaled /= self.stds + # unscaled += self.means + + if self.scale_minmax: + unscaled /= self.maxs - self.mins + # unscaled += self.mins + + else: + + ########################## + # Total scaling + ########################## + + if self.scale_standard: + unscaled /= self.total_std + + if self.scale_minmax: + unscaled /= self.total_max - self.total_min + + if as_numpy: + return unscaled.detach().numpy().astype(np.float64) + else: + return unscaled + def save(self, filename, save_format="pickle"): """ Save the Scaler object so that it can be accessed again later. diff --git a/mala/network/predictor.py b/mala/network/predictor.py index 7db376a95..9f5e654f1 100644 --- a/mala/network/predictor.py +++ b/mala/network/predictor.py @@ -47,6 +47,9 @@ def __init__(self, params, network, data): self._number_of_batches_per_snapshot = 0 self.target_calculator = data.target_calculator + self.input_data = None + self.output_data = None + def predict_from_qeout(self, path_to_file, gather_ldos=False): """ Get predicted LDOS for the atomic configuration of a QE.out file. @@ -74,7 +77,14 @@ def predict_from_qeout(self, path_to_file, gather_ldos=False): self.data.target_calculator.atoms, gather_ldos=gather_ldos ) - def predict_for_atoms(self, atoms, gather_ldos=False, temperature=None): + def predict_for_atoms( + self, + atoms, + gather_ldos=False, + temperature=None, + save_grads=False, + pass_descriptors=None, + ): """ Get predicted LDOS for an atomic configuration. @@ -128,24 +138,32 @@ def predict_for_atoms(self, atoms, gather_ldos=False, temperature=None): self.data.target_calculator.invalidate_target() # Calculate descriptors. - time_before = perf_counter() - snap_descriptors, local_size = ( - self.data.descriptor_calculator.calculate_from_atoms( - atoms, self.data.grid_dimension - ) - ) - printout( - "Time for descriptor calculation: {:.8f}s".format( - perf_counter() - time_before - ), - min_verbosity=2, - ) + if pass_descriptors is None: + time_before = perf_counter() + snap_descriptors, local_size = ( + self.data.descriptor_calculator.calculate_from_atoms( + atoms, self.data.grid_dimension + ) + ) + printout( + "Time for descriptor calculation: {:.8f}s".format( + perf_counter() - time_before + ), + min_verbosity=2, + ) + feature_length = self.data.descriptor_calculator.feature_size + descs_with_xyz = ( + self.data.descriptor_calculator.descriptors_contain_xyz + ) + else: + snap_descriptors, local_size, feature_length, descs_with_xyz = ( + pass_descriptors + ) # Provide info from current snapshot to target calculator. self.data.target_calculator.read_additional_calculation_data( [atoms, self.data.grid_dimension], "atoms+grid" ) - feature_length = self.data.descriptor_calculator.feature_size # The actual calculation of the LDOS from the descriptors depends # on whether we run in parallel or serial. In the former case, @@ -166,7 +184,7 @@ def predict_for_atoms(self, atoms, gather_ldos=False, temperature=None): return None else: - if self.data.descriptor_calculator.descriptors_contain_xyz: + if descs_with_xyz: self.data.target_calculator.local_grid = snap_descriptors[ :, 0:3 ].copy() @@ -191,7 +209,8 @@ def predict_for_atoms(self, atoms, gather_ldos=False, temperature=None): ) if get_rank() == 0: - if self.data.descriptor_calculator.descriptors_contain_xyz: + # if self.data.descriptor_calculator.descriptors_contain_xyz: + if descs_with_xyz: snap_descriptors = snap_descriptors[:, :, :, 3:] feature_length -= 3 @@ -200,10 +219,17 @@ def predict_for_atoms(self, atoms, gather_ldos=False, temperature=None): ) snap_descriptors = torch.from_numpy(snap_descriptors).float() self.data.input_data_scaler.transform(snap_descriptors) - return self._forward_snap_descriptors(snap_descriptors) + if save_grads is True: + self.input_data = snap_descriptors + self.input_data.requires_grad = True + return self._forward_snap_descriptors( + snap_descriptors, save_torch_outputs=True + ) + else: + return self._forward_snap_descriptors(snap_descriptors) def _forward_snap_descriptors( - self, snap_descriptors, local_data_size=None + self, snap_descriptors, local_data_size=None, save_torch_outputs=False ): """ Forward a scaled tensor of descriptors through the neural network. @@ -238,6 +264,10 @@ def _forward_snap_descriptors( predicted_outputs = np.zeros( (local_data_size, self.data.target_calculator.feature_size) ) + if save_torch_outputs: + self.output_data = torch.zeros( + (local_data_size, self.data.target_calculator.feature_size) + ) # Only predict if there is something to predict. # Elsewise, we just wait at the barrier down below. @@ -267,6 +297,8 @@ def _forward_snap_descriptors( inputs = snap_descriptors[sl].to( self.parameters._configuration["device"] ) + if save_torch_outputs: + self.output_data[sl] = self.network(inputs) predicted_outputs[sl] = ( self.data.output_data_scaler.inverse_transform( self.network(inputs).to("cpu"), as_numpy=True diff --git a/mala/targets/calculation_helpers.py b/mala/targets/calculation_helpers.py index 5cb2b09d2..c39162954 100644 --- a/mala/targets/calculation_helpers.py +++ b/mala/targets/calculation_helpers.py @@ -6,6 +6,7 @@ from scipy import integrate + def integrate_values_on_spacing(values, spacing, method, axis=0): """ Integrate values assuming a uniform grid with a provided spacing. @@ -399,6 +400,108 @@ def analytical_integration(D, I0, I1, fermi_energy, energy_grid, temperature): return integral_value +def analytical_integration_weights(I0, I1, fermi_energy, energy_grid, + temperature): + """ + Constructs the weights used in analytical integration as described in Eq. 32 of [1]. + + Parameters + ---------- + I0 : string + I_0 from Eq. 33. The user only needs to provide which function should + be used. Arguments are: + + - F0 + - F1 + - F2 + - S0 + - S1 + + I1 : string + I_1 from Eq. 33. The user only needs to provide which function should + be used. Arguments are: + + - F0 + - F1 + - F2 + - S0 + - S1 + + fermi_energy : float + The fermi energy in eV. + + energy_grid : numpy.array + Energy grid on which the integration should be performed. + + temperature : float + Temperature in K. + + Returns + ------- + weights_vector : numpy.array + Weights to be used in the integral. + + """ + + # Mappings for the functions further down. + function_mappings = { + "F0": get_f0_value, + "F1": get_f1_value, + "F2": get_f2_value, + "S0": get_s0_value, + "S1": get_s1_value, + } + + # Check if everything makes sense. + if I0 not in list(function_mappings.keys()) or I1 not in\ + list(function_mappings.keys()): + raise Exception("Could not calculate analytical intergal, " + "wrong choice of auxiliary functions.") + + # Construct the weight vector. + weights_vector = np.zeros(energy_grid.shape, dtype=np.float64) + gridsize = energy_grid.shape[0] + energy_grid_edges = np.zeros(energy_grid.shape[0]+2, dtype=np.float64) + energy_grid_edges[1:-1] = energy_grid + spacing = (energy_grid[1]-energy_grid[0]) + energy_grid_edges[0] = energy_grid[0] - spacing + energy_grid_edges[-1] = energy_grid[-1] + spacing + + # Calculate the weights. + # It is not really possible to express this as a vector operation, + # since mp.polylog (which is called in function_mappings) does not support + # that. + beta = 1 / (kB * temperature) + for i in range(0, gridsize): + # Some aliases for readibility + ei = energy_grid_edges[i+1] + ei_plus = energy_grid_edges[i+2] + ei_minus = energy_grid_edges[i] + + # Calculate x + x = beta*(ei - fermi_energy) + x_plus = beta*(ei_plus - fermi_energy) + x_minus = beta*(ei_minus - fermi_energy) + + # Calculate the I0 value + i0 = function_mappings[I0](x, beta) + i0_plus = function_mappings[I0](x_plus, beta) + i0_minus = function_mappings[I0](x_minus, beta) + + # Calculate the I1 value + i1 = function_mappings[I1](x, beta) + i1_plus = function_mappings[I1](x_plus, beta) + i1_minus = function_mappings[I1](x_minus, beta) + + weights_vector[i] = (i0_plus-i0) * (1 + + ((ei - fermi_energy) / (ei_plus - ei))) \ + + (i0-i0_minus) * (1 - ((ei - fermi_energy) / (ei - ei_minus))) - \ + ((i1_plus-i1) / (ei_plus-ei)) + ((i1 - i1_minus) + / (ei - ei_minus)) + + return weights_vector + + # Define Gaussian def gaussians(grid, centers, sigma): r""" diff --git a/mala/targets/density.py b/mala/targets/density.py index 39df84a6e..fda9e8ac3 100644 --- a/mala/targets/density.py +++ b/mala/targets/density.py @@ -329,6 +329,19 @@ def total_energy_contributions(self): "No cached density available to calculate this property." ) + @cached_property + def force_contributions(self): + """ + All density based contributions to the forces. + + Calculated via the cached density. + """ + if self.density is not None: + return self.get_force_contributions() + else: + raise Exception("No cached density available to " + "calculate this property.") + def uncache_properties(self): """Uncache all cached properties of this calculator.""" if self._is_property_cached("number_of_electrons"): @@ -337,6 +350,9 @@ def uncache_properties(self): if self._is_property_cached("total_energy_contributions"): del self.total_energy_contributions + if self._is_property_cached("force_contributions"): + del self.force_contributions + ############################## # Methods ############################## @@ -836,54 +852,9 @@ def get_energy_contributions( } return energies_dict - def get_atomic_forces( - self, - density_data=None, - create_file=True, - atoms_Angstrom=None, - qe_input_data=None, - qe_pseudopotentials=None, - ): - """ - Calculate the atomic forces. - - This function uses an interface to QE. The atomic forces are - calculated via the Hellman-Feynman theorem, although only the local - contributions are calculated. The non-local contributions, as well - as the SCF correction (so anything wavefunction dependent) are ignored. - Therefore, this function is best used for data that was created using - local pseudopotentials. - - Parameters - ---------- - density_data : numpy.ndarray - Density data on a grid. If None, then the cached - density will be used for the calculation. - - create_file : bool - If False, the last mala.pw.scf.in file will be used as input for - Quantum Espresso. If True (recommended), MALA will create this - file according to calculation parameters. - - atoms_Angstrom : ase.Atoms - ASE atoms object for the current system. If None, MALA will - create one. - - qe_input_data : dict - Quantum Espresso parameters dictionary for the ASE<->QE interface. - If None (recommended), MALA will create one. - - qe_pseudopotentials : dict - Quantum Espresso pseudopotential dictionaty for the ASE<->QE - interface. If None (recommended), MALA will create one. - - Returns - ------- - atomic_forces : numpy.ndarray - An array of the form (natoms, 3), containing the atomic forces - in eV/Ang. - - """ + def get_force_contributions(self, density_data=None, create_file=True, + atoms_Angstrom=None, qe_input_data=None, + qe_pseudopotentials=None): if density_data is None: density_data = self.density if density_data is None: @@ -892,27 +863,48 @@ def get_atomic_forces( " this quantity." ) - # First, set up the total energy module for calculation. if atoms_Angstrom is None: atoms_Angstrom = self.atoms - self.__setup_total_energy_module( - density_data, - atoms_Angstrom, - create_file=create_file, - qe_input_data=qe_input_data, - qe_pseudopotentials=qe_pseudopotentials, - ) - # Now calculate the forces. - atomic_forces = np.array( - te.calc_forces(len(atoms_Angstrom)) - ).transpose() + # If the total energy has been calculated, we don't need to + # re-execute this particular code. + if self._is_property_cached("total_energy_contributions"): + pass + else: + self.__setup_total_energy_module(density_data, atoms_Angstrom, + create_file=create_file, + qe_input_data=qe_input_data, + qe_pseudopotentials= + qe_pseudopotentials) + + # Get the number of gridpoints in QE and MALA. + number_of_gridpoints = te.get_nnr() + if len(density_data.shape) == 4: + number_of_gridpoints_mala = density_data.shape[0] * \ + density_data.shape[1] * \ + density_data.shape[2] + elif len(density_data.shape) == 2: + number_of_gridpoints_mala = density_data.shape[0] + else: + raise Exception("Density data has wrong dimensions. ") - # QE returns the forces in Ry/Bohr. - atomic_forces = AtomicForce.convert_units( - atomic_forces, in_units="Ry/Bohr" - ) - return atomic_forces + spin = te.get_nspin() + + # Putting the force contributions together. + # @Normand: If you add more things to be extracted from QE, + # I would suggest you do it here. + # Hartree potential + hartree_potential = np.zeros([number_of_gridpoints, 1]) + te.v_h_wrapper(hartree_potential, number_of_gridpoints, spin) + hartree_potential = np.reshape(hartree_potential, + self.grid_dimensions+[1], order="F") + force_contribution = {"hartree": np.reshape(hartree_potential, + [number_of_gridpoints_mala, + 1], + order="C") * Rydberg * \ + self.voxel.volume} + + return force_contribution @staticmethod def get_scaled_positions_for_qe(atoms): diff --git a/mala/targets/dos.py b/mala/targets/dos.py index 4dad3e2f2..3898bc1bd 100644 --- a/mala/targets/dos.py +++ b/mala/targets/dos.py @@ -17,6 +17,7 @@ analytical_integration, get_beta, entropy_multiplicator, + analytical_integration_weights, ) @@ -292,6 +293,15 @@ def band_energy(self): "No cached DOS available to calculate this property." ) + @cached_property + def d_band_energy_d_dos(self): + """Derivative of band energy, calculated via cached DOS.""" + if self.density_of_states is not None: + return self.get_d_band_energy_d_dos() + else: + raise Exception("No cached DOS available to calculate this " + "property.") + @cached_property def number_of_electrons(self): """ @@ -343,6 +353,24 @@ def entropy_contribution(self): "No cached DOS available to calculate this property." ) + @cached_property + def d_entropy_contribution_d_dos(self): + """Derivative of band energy, calculated via cached DOS.""" + if self.density_of_states is not None: + return self.get_d_entropy_contribution_d_dos() + else: + raise Exception("No cached DOS available to calculate this " + "property.") + + @cached_property + def d_number_of_electrons_d_mu(self): + """ """ + if self.density_of_states is not None: + return self.get_d_number_of_electrons_d_mu() + else: + raise Exception("No cached DOS available to calculate this " + "property.") + def uncache_properties(self): """Uncache all cached properties of this calculator.""" if self._is_property_cached("number_of_electrons"): @@ -351,8 +379,16 @@ def uncache_properties(self): del self.energy_grid if self._is_property_cached("band_energy"): del self.band_energy + if self._is_property_cached("entropy_contribution"): + del self.entropy_contribution if self._is_property_cached("fermi_energy"): del self.fermi_energy + if self._is_property_cached("d_band_energy_dos"): + del self.d_band_energy_d_dos + if self._is_property_cached("d_entropy_contribution_d_dos"): + del self.d_entropy_contribution_d_dos + if self._is_property_cached("d_number_of_electrons_d_mu"): + del self.d_number_of_electrons_d_mu ############################## # Methods @@ -659,7 +695,6 @@ def get_band_energy( Band energy in eV. """ # Parse the parameters. - # Parse the parameters. if dos_data is None and self.density_of_states is None: raise Exception( "No DOS data provided, cannot calculate this quantity." @@ -709,13 +744,81 @@ def get_band_energy( integration_method, ) - return self.__band_energy_from_dos( - dos_data, - energy_grid, - fermi_energy, - temperature, - integration_method, - ) + def get_d_band_energy_d_dos(self, dos_data=None, fermi_energy=None, + temperature=None, broadcast_derivative=True, + number_of_electrons=None): + """ + Calculate the derivative of the band energy with respect to the DOS. + + Note that the derivative is taken at a constant number of electrons + with the assumption that the change in the number of electrons is + compensated by a change in the Fermi energy. This makes sense in + metals, but probably does not make sense in insulators or + semiconductors. + + Parameters + ---------- + dos_data : numpy.array + DOS data with dimension [energygrid]. + + fermi_energy : float + Fermi energy level in eV. + + temperature : float + Temperature in K. + + broadcast_derivative : bool + If True then the band energy will only be calculated on one + rank and thereafter be distributed to all other ranks. + + Returns + ------- + d_band_energy_d_DOS : numpy.array with dimension [energygrid]. + Derivative of the band energy (in eV) with respect to the DOS. + """ + # Parse the parameters. + if dos_data is None and self.density_of_states is None: + raise Exception("No DOS data provided, cannot calculate" + " this quantity.") + + # Here we check whether we will use our internal, cached + # DOS, or calculate everything from scratch. + if dos_data is not None: + if fermi_energy is None: + printout("Warning: No fermi energy was provided or could be " + "calculated from electronic structure data. " + "Using the DFT fermi energy, this may " + "yield unexpected results", min_verbosity=1) + fermi_energy = self.fermi_energy_dft + number_of_electrons = self.number_of_electrons_exact + else: + dos_data = self.density_of_states + fermi_energy = self.fermi_energy + number_of_electrons = self.number_of_electrons + + if temperature is None: + temperature = self.temperature + + if self.parameters._configuration["mpi"] and broadcast_derivative: + if get_rank() == 0: + energy_grid = self.energy_grid + d_band_energy_d_dos = self.\ + __d_band_energy_d_dos_from_dos(dos_data, energy_grid, + fermi_energy, temperature, + number_of_electrons) + else: + d_band_energy_d_dos = None + + d_band_energy_d_dos = get_comm().bcast(d_band_energy_d_dos, + root=0) + barrier() + return d_band_energy_d_dos + else: + energy_grid = self.energy_grid + return self.\ + __d_band_energy_d_dos_from_dos(dos_data, energy_grid, + fermi_energy, temperature, + number_of_electrons) def get_number_of_electrons( self, @@ -873,6 +976,132 @@ def get_entropy_contribution( integration_method, ) + def get_d_entropy_contribution_d_dos(self, dos_data=None, + fermi_energy=None, + temperature=None, + broadcast_derivative=True): + """ + Calculate the derivative of the entropy contribution w.r.t the DOS. + + Note that the derivative is taken at a constant number of electrons + with the assumption that the change in the number of electrons is + compensated by a change in the Fermi energy. This makes sense in + metals, but probably does not make sense in insulators or + semiconductors. + + Parameters + ---------- + dos_data : numpy.array + DOS data with dimension [energygrid]. If None, then the cached + DOS will be used for the calculation. + + fermi_energy : float + Fermi energy level in eV. + + temperature : float + Temperature in K. + + broadcast_derivative : bool + If True then the derivative will only be calculated on one + rank and thereafter be distributed to all other ranks. + + Returns + ------- + entropy_contribution : float + S/beta in eV. + """ + # Parse the parameters. + if dos_data is None and self.density_of_states is None: + raise Exception("No DOS data provided, cannot calculate" + " this quantity.") + + # Here we check whether we will use our internal, cached + # DOS, or calculate everything from scratch. + if dos_data is not None: + if fermi_energy is None: + printout("Warning: No fermi energy was provided or could be " + "calculated from electronic structure data. " + "Using the DFT fermi energy, this may " + "yield unexpected results", min_verbosity=1) + fermi_energy = self.fermi_energy_dft + else: + dos_data = self.density_of_states + fermi_energy = self.fermi_energy + if temperature is None: + temperature = self.temperature + + if self.parameters._configuration["mpi"] and broadcast_derivative: + if get_rank() == 0: + energy_grid = self.energy_grid + entropy = self. \ + __d_entropy_contribution_d_dos_from_dos(dos_data, energy_grid, + fermi_energy, temperature) + else: + entropy = None + + entropy = get_comm().bcast(entropy, root=0) + barrier() + return entropy + else: + energy_grid = self.energy_grid + return self. \ + __d_entropy_contribution_d_dos_from_dos(dos_data, energy_grid, + fermi_energy, temperature) + + def get_d_number_of_electrons_d_mu(self, dos_data=None, + fermi_energy=None, + temperature=None, + broadcast_derivative=True, + delta=None): + """ + + Returns + ------- + + """ + # Parse the parameters. + if dos_data is None and self.density_of_states is None: + raise Exception("No DOS data provided, cannot calculate" + " this quantity.") + + # Here we check whether we will use our internal, cached + # DOS, or calculate everything from scratch. + if dos_data is not None: + if fermi_energy is None: + printout("Warning: No fermi energy was provided or could be " + "calculated from electronic structure data. " + "Using the DFT fermi energy, this may " + "yield unexpected results", min_verbosity=1) + fermi_energy = self.fermi_energy_dft + else: + dos_data = self.density_of_states + fermi_energy = self.fermi_energy + if temperature is None: + temperature = self.temperature + + if delta is None: + delta = self.parameters.delta_forces + + if self.parameters._configuration["mpi"] and broadcast_derivative: + if get_rank() == 0: + energy_grid = self.energy_grid + derivative = self. \ + __d_number_of_electrons_d_mu(dos_data, energy_grid, + fermi_energy, temperature, + delta) + else: + derivative = None + + derivative = get_comm().bcast(derivative, root=0) + barrier() + return derivative + else: + energy_grid = self.energy_grid + return self. \ + __d_number_of_electrons_d_mu(dos_data, energy_grid, + fermi_energy, temperature, + delta) + def get_self_consistent_fermi_energy( self, dos_data=None, @@ -1157,6 +1386,76 @@ def __band_energy_from_dos( return band_energy + @staticmethod + def __d_band_energy_d_dos_from_dos(dos_data, energy_grid, fermi_energy, + temperature, number_of_electrons): + """Calculate the derivative of the band energy w.r.t. the DOS.""" + d_band_energy_minus_uN_d_dos = \ + analytical_integration_weights("F1", "F2", fermi_energy, + energy_grid, temperature) + d_number_of_electrons_d_dos = \ + analytical_integration_weights("F0", "F1", fermi_energy, + energy_grid, temperature) + + # The following crude numerical derivatives could be replaced with + # analytic derivatives with some work.Or they could be replaced with + # more sophisticated numerical derivatives. + delta = 0.001 + d_number_of_electrons_d_mu = (analytical_integration(dos_data, "F0", + "F1", + fermi_energy + + delta, + energy_grid, + temperature) + - analytical_integration(dos_data, "F0", + "F1", + fermi_energy - + delta, + energy_grid, + temperature)) /\ + (2.0*delta) + d_band_energy_minus_uN_d_mu = (analytical_integration(dos_data, "F1", + "F2", + fermi_energy + + delta, + energy_grid, + temperature) + - analytical_integration(dos_data, "F1", + "F2", + fermi_energy + - delta, + energy_grid, + temperature))/\ + (2.0*delta) + d_band_energy_d_mu = d_band_energy_minus_uN_d_mu + number_of_electrons + + # Since the number of electrons is constant + # \frac{dN}{dDOS} = (\partial{dN}{dDOS})_{\mu} + + # \frac{dN}{d\mu}*\frac{d\mu}{dDOS} = 0 + # So, + # \frac{d\mu}{dDOS} = - (\partial{dN}{dDOS})_{\mu} / \frac{dN}{d\mu} + # + # Also, + # \frac{d(E - \mu*N)}{dDOS} = (\partial{d(E - \mu*N)}{dDOS})_{\mu} + + # \frac{d(E - \mu*N)}{d\mu}*\frac{d\mu}{dDOS} + # \frac{dE}{dDOS} - \frac{d\mu}{dDOS}*N = (\partial{d(E - + # \mu*N)}{dDOS})_{\mu} + \frac{dE}{d\mu}*\frac{d\mu}{dDOS} - + # N*\frac{d\mu}{dDOS} + # \frac{dE}{dDOS} = (\partial{d(E - \mu*N)}{dDOS})_{\mu} + + # \frac{dE}{d\mu}*\frac{d\mu}{dDOS} + # + # Combining and rearranging, + # \frac{dE}{dDOS} = (\partial{d(E - \mu*N)}{dDOS})_{\mu} - + # \frac{dE}{d\mu}*(\partial{dN}{dDOS})_{\mu} / \frac{dN}{d\mu} + # = (\partial{d(E - \mu*N)}{dDOS})_{\mu} - + # (\partial{dN}{dDOS})_{\mu}*\frac{dE}{d\mu}/\frac{dN}{d\mu} + + d_band_energy_d_dos = d_band_energy_minus_uN_d_dos \ + - d_number_of_electrons_d_dos * \ + d_band_energy_d_mu / d_number_of_electrons_d_mu + + return d_band_energy_d_dos + @staticmethod def __entropy_contribution_from_dos( dos_data, energy_grid, fermi_energy, temperature, integration_method @@ -1228,3 +1527,76 @@ def __entropy_contribution_from_dos( raise Exception("Unknown integration method.") return entropy_contribution + + + @staticmethod + def __d_entropy_contribution_d_dos_from_dos(dos_data, energy_grid, fermi_energy, + temperature): + """Calculate the derivative of the entropy w.r.t the DOS.""" + d_entropy_contribution_d_dos = \ + analytical_integration_weights("S0", "S1", fermi_energy, + energy_grid, temperature) + d_number_of_electrons_d_dos = \ + analytical_integration_weights("F0", "F1", fermi_energy, + energy_grid, temperature) + # The following crude numerical derivatives could be replaced with + # analytic derivatives with some work. + # Or they could be replaced with more sophisticated numerical + # derivatives. + delta = 1e-10 + d_number_of_electrons_d_mu = (analytical_integration(dos_data, "F0", "F1", + fermi_energy + delta, + energy_grid, + temperature) \ + - analytical_integration(dos_data, "F0", "F1", + fermi_energy - delta, + energy_grid, + temperature)) / ( + 2.0 * delta) + d_entropy_contribution_d_mu = (analytical_integration(dos_data, "S0", "S1", + fermi_energy + delta, + energy_grid, + temperature) \ + - analytical_integration(dos_data, "S0", "S1", + fermi_energy - delta, + energy_grid, + temperature)) / ( + 2.0 * delta) + + # Since the number of electrons is constant + # \frac{dN}{dDOS} = (\partial{dN}{dDOS})_{\mu} + + # \frac{dN}{d\mu}*\frac{d\mu}{dDOS} = 0 + # So, + # \frac{d\mu}{dDOS} = - (\partial{dN}{dDOS})_{\mu} / + # \frac{dN}{d\mu} + # + # Also, + # \frac{dTS}{dDOS} = (\partial{dTS}{dDOS})_{\mu} + + # \frac{dTS}{d\mu}*\frac{d\mu}{dDOS} + # = (\partial{dTS}{dDOS})_{\mu} - + # \frac{dTS}{d\mu}*(\partial{dN}{dDOS})_{\mu} / \frac{dN}{d\mu} + # = (\partial{dTS}{dDOS})_{\mu} - + # (\partial{dN}{dDOS})_{\mu} * \frac{dTS}{d\mu}/\frac{dN}{d\mu} + + d_entropy_contribution_d_dos = d_entropy_contribution_d_dos \ + - d_number_of_electrons_d_dos * d_entropy_contribution_d_mu / d_number_of_electrons_d_mu + return d_entropy_contribution_d_dos + + @staticmethod + def __d_number_of_electrons_d_mu(dos_data, energy_grid, fermi_energy, + temperature, delta): + d_number_of_electrons_d_mu = (analytical_integration(dos_data, "F0", + "F1", + fermi_energy + + delta, + energy_grid, + temperature) + - analytical_integration(dos_data, "F0", + "F1", + fermi_energy - + delta, + energy_grid, + temperature)) /\ + (2.0*delta) + return d_number_of_electrons_d_mu + diff --git a/mala/targets/ldos.py b/mala/targets/ldos.py index 5d8253520..1ef198475 100644 --- a/mala/targets/ldos.py +++ b/mala/targets/ldos.py @@ -2,10 +2,12 @@ from functools import cached_property +import mala.datahandling.data_scaler from ase.units import Rydberg, Bohr, J, m import math import numpy as np from scipy import integrate +import torch from mala.common.parallelizer import ( get_comm, @@ -22,6 +24,7 @@ fermi_function, analytical_integration, integrate_values_on_spacing, + analytical_integration_weights, ) from mala.targets.dos import DOS from mala.targets.density import Density @@ -44,6 +47,13 @@ def __init__(self, params): super(LDOS, self).__init__(params) # Consistency check for parameters describing energy grid. self.local_density_of_states = None + self.debug_forces_flag = None + self.input_data_derivative = None + self.output_data_torch = None + self.input_data_scaler: mala.datahandling.data_scaler.DataScaler = None + self.output_data_scaler: mala.datahandling.data_scaler.DataScaler = ( + None + ) @classmethod def from_numpy_file(cls, params, path, units="1/(eV*A^3)"): @@ -291,6 +301,8 @@ def uncache_properties(self): del self._density_of_states_calculator if self._is_property_cached("entropy_contribution"): del self.entropy_contribution + if self._is_property_cached("atomic_forces"): + del self.atomic_forces @cached_property def energy_grid(self): @@ -307,6 +319,16 @@ def total_energy(self): "No cached LDOS available to calculate this property." ) + @cached_property + def atomic_forces(self): + """Atomic forces of the system, calculated via cached LDOS.""" + if self.local_density_of_states is not None: + return self.get_atomic_forces() + else: + raise Exception( + "No cached LDOS available to calculate this " "property." + ) + @cached_property def band_energy(self): """Band energy of the system, calculated via cached LDOS.""" @@ -1404,7 +1426,21 @@ def get_density_of_states( return dos_values def get_atomic_forces( - self, ldos_data, dE_dd, used_data_handler, snapshot_number=0 + self, + ldos_data=None, + dos_data=None, + density_data=None, + fermi_energy=None, + temperature=None, + voxel=None, + grid_integration_method="summation", + energy_integration_method="analytical", + atoms_Angstrom=None, + qe_input_data=None, + qe_pseudopotentials=None, + create_qe_file=True, + return_energy_contributions=False, + delta=None, ): r""" Get the atomic forces, currently work in progress. @@ -1438,11 +1474,169 @@ def get_atomic_forces( the descriptors. """ + if voxel is None: + voxel = self.voxel + if fermi_energy is None: + if ldos_data is None: + fermi_energy = self.fermi_energy + if fermi_energy is None: + printout( + "Warning: No fermi energy was provided or could be " + "calculated from electronic structure data. " + "Using the DFT fermi energy, this may " + "yield unexpected results", + min_verbosity=1, + ) + fermi_energy = self.fermi_energy_dft + if temperature is None: + temperature = self.temperature + if delta is None: + delta = self.parameters.delta_forces + + energy_grid = self._get_energy_grid() + + # Here we check whether we will use our internal, cached + # LDOS, or calculate everything from scratch. + if ldos_data is not None or ( + dos_data is not None and density_data is not None + ): + # Not using cached properties is currently not implemented. + # We can implement this later as needed. + raise Exception( + "Forces can currently only be calculated using " + "cached properties!" + ) + else: + # In this case, we use cached propeties wherever possible. + ldos_data = self.local_density_of_states + if ldos_data is None: + raise Exception( + "No input data provided to caculate " + "total energy. Provide EITHER LDOS" + " OR DOS and density." + ) + + d_E_d_d = np.zeros_like(self.local_density_of_states) + + ############################# + # DOS dependent energy terms. + ############################# + + # Band energy and entropy contribution can be calculated separately + # (if the debug flag asks for it) or as a singular DOS + # contribution. + if self.debug_forces_flag == "band_energy": + d_E_d_D = ( + self._density_of_states_calculator.d_band_energy_d_dos + ) * voxel.volume + + if self.debug_forces_flag == "entropy_contribution": + d_E_d_D = ( + self._density_of_states_calculator.d_entropy_contribution_d_dos + ) * voxel.volume + + # This is the entire DOS contribution. + if self.debug_forces_flag is None: + d_E_d_D = ( + self._density_of_states_calculator.d_band_energy_d_dos + + self._density_of_states_calculator.d_entropy_contribution_d_dos + ) * voxel.volume + + # The DOS contribution may need reshaping. + if ( + self.debug_forces_flag == "band_energy" + or self.debug_forces_flag == "entropy_contribution" + or self.debug_forces_flag is None + ): + if len(np.shape(self.local_density_of_states)) == 4: + d_E_d_d[:, :, :] = d_E_d_D + elif len(np.shape(self.local_density_of_states)) == 2: + d_E_d_d[:] = d_E_d_D + else: + raise Exception("Invalid LDOS shape provided.") + + ################################# + # Density dependent energy terms. + ################################# + + # Currently only the Hartree energy is computded. + if ( + self.debug_forces_flag is None + or self.debug_forces_flag == "hartree" + ): + + # The derivative of the Hartree energy w.r.t the density + # is simply the Hartree potential. + # Then we have to apply the chain rule. + d_E_h_d_n = self._density_calculator.force_contributions[ + "hartree" + ] + d_n_d_d = analytical_integration_weights( + "F0", "F1", fermi_energy, energy_grid, temperature + ) + + # The second portion of the derivative is the derivative of the + # Hartree energy w.r.t. the chemical potential. + # That's what happens here. + d_f01_d_mu = ( + analytical_integration_weights( + "F0", + "F1", + fermi_energy + delta, + energy_grid, + temperature, + ) + - analytical_integration_weights( + "F0", + "F1", + fermi_energy - delta, + energy_grid, + temperature, + ) + ) / (2.0 * delta) + d_n_d_mu = np.dot(ldos_data, d_f01_d_mu) + d_E_h_d_mu = np.dot(d_E_h_d_n[:, 0], d_n_d_mu) + d_mu_d_d = np.zeros_like(ldos_data) + d_mu_d_d[:] = d_n_d_d * self.voxel.volume + d_mu_d_d /= (-1.0) * ( + self._density_of_states_calculator.d_number_of_electrons_d_mu + ) + + # Add both portions together. + d_E_d_n = d_E_h_d_n * d_n_d_d + d_E_h_d_mu * d_mu_d_d + d_E_d_d += d_E_d_n + + ################################# + # Backpropagation of LDOS terms. + ################################# + + if self.input_data_derivative is not None: + # Scale energy derivative. + d_E_d_d = torch.from_numpy(d_E_d_d) + d_E_d_d = self.output_data_scaler.transform_backpropagation( + d_E_d_d + ) + + # Backprop it through the network. + self.output_data_torch.backward(d_E_d_d, retain_graph=True) + + # Get new energy derivative and return it. + d_E_d_B = self.input_data_derivative.grad.clone() + d_E_d_B = ( + self.input_data_scaler.inverse_transform_backpropagation( + d_E_d_B + ) + ) + + return d_E_d_B.detach().numpy().astype(np.float64) + + else: + return d_E_d_d # For now this only works with ML generated LDOS. # Gradient of the LDOS respect to the descriptors. - ldos_data.backward(dE_dd) - dd_dB = used_data_handler.get_test_input_gradient(snapshot_number) - return dd_dB + # ldos_data.backward(dE_dd) + # dd_dB = used_data_handler.get_test_input_gradient(snapshot_number) + # return dd_dB # Private methods ################# @@ -1809,3 +2003,18 @@ def _read_from_qe_files( else: self.local_density_of_states = ldos_data return ldos_data + + def setup_for_forces(self, predictor): + """ + Set calculator up for forces prediction by copying data from Predictor. + + + Parameters + ---------- + predictor : mala.Predictor + Predictor class used for LDOS prediction. + """ + self.input_data_derivative = predictor.input_data + self.output_data_torch = predictor.output_data + self.input_data_scaler = predictor.data.input_data_scaler + self.output_data_scaler = predictor.data.output_data_scaler