From abe0d0e1c1114880cab660a50a1164685c2b4cf7 Mon Sep 17 00:00:00 2001 From: Kristjan Eimre Date: Sat, 2 Jul 2022 17:13:46 +0300 Subject: [PATCH] skip p-tip calc. if only s tip is specified --- cp2k_spm_tools/cp2k_stm_sts.py | 110 ++++++++---------- .../run_cube_from_wfn.sh | 8 +- examples/benzene_stm/run_stm_sts_from_wfn.sh | 6 +- stm_sts_from_wfn.py | 1 + 4 files changed, 55 insertions(+), 70 deletions(-) diff --git a/cp2k_spm_tools/cp2k_stm_sts.py b/cp2k_spm_tools/cp2k_stm_sts.py index 79aab26..959f75f 100644 --- a/cp2k_spm_tools/cp2k_stm_sts.py +++ b/cp2k_spm_tools/cp2k_stm_sts.py @@ -40,10 +40,13 @@ def __init__(self, mpi_comm, cp2k_grid_orb, p_tip_ratios): self.mpi_comm = mpi_comm self.p_tip_ratios = p_tip_ratios + # add a check if ptip needs to be calculated, as it can use a lot of memory + self.ptip_enabled = True + if all(r == 0.0 for r in p_tip_ratios): + self.ptip_enabled = False self.global_morb_energies_by_rank = None - self.z_arr = np.arange(0.0, self.cell_n[2]*self.dv[2], self.dv[2]) + self.origin[2] # to angstrom and WRT to topmost atom self.z_arr /= ang_2_bohr @@ -58,12 +61,6 @@ def __init__(self, mpi_comm, cp2k_grid_orb, p_tip_ratios): # p-tip contribution of the orbitals defined in local space for this mpi_rank self.local_orbital_ptip = None - -# # List of dictionaries containing everything to do with STM/STS maps -# self.stm_maps_data_list = [] -# -# # Dictionary containing everything to do with orbital maps -# self.orb_maps_data = None # Dictionary containing all the STM/STS/ORB output self.series_output = {} @@ -125,11 +122,6 @@ def divide_by_space(self): orbitals_per_rank = np.array([len(gme) for gme in self.global_morb_energies_by_rank[ispin]]) total_orb = sum(orbitals_per_rank) - ### Calculate and divide also the p-tip contribution, - ### as derivatives are hard to account for after dividing the orbitals in space - grad_x, grad_y = np.gradient(self.cgo.morb_grids[ispin], axis=(1, 2)) - p_tip_contrib = (grad_x/self.dv[0]) ** 2 + (grad_y/self.dv[1]) ** 2 - for rank in range(self.mpi_size): # which indexes to send? @@ -152,17 +144,24 @@ def divide_by_space(self): if self.mpi_rank == rank: self.local_orbitals.append(recvbuf.reshape(total_orb, self.local_cell_n[0], self.local_cell_n[1], self.local_cell_n[2])) - - ### Divide also the p-tip contribution! - sendbuf = p_tip_contrib[:, ix_start:ix_end, :, :].ravel() - if self.mpi_rank == rank: - recvbuf = np.empty(sum(orbitals_per_rank)*num_spatial_points, dtype=self.cgo.dtype) - else: - recvbuf = None - self.mpi_comm.Gatherv(sendbuf=sendbuf, - recvbuf=[recvbuf, orbitals_per_rank*num_spatial_points], root=rank) - if self.mpi_rank == rank: - self.local_orbital_ptip.append(recvbuf.reshape(total_orb, self.local_cell_n[0], self.local_cell_n[1], self.local_cell_n[2])) + + if self.ptip_enabled: + ### Calculate and divide also the p-tip contribution, + ### as derivatives are hard to account for after dividing the orbitals in space + p_tip_contrib = (np.gradient(self.cgo.morb_grids[ispin], axis=1)/self.dv[0])**2 + p_tip_contrib += (np.gradient(self.cgo.morb_grids[ispin], axis=2)/self.dv[1])**2 + + for rank in range(self.mpi_size): + ix_start, ix_end = self.x_ind_per_rank(rank) + if self.mpi_rank == rank: + recvbuf = np.empty(sum(orbitals_per_rank)*num_spatial_points, dtype=self.cgo.dtype) + else: + recvbuf = None + sendbuf = p_tip_contrib[:, ix_start:ix_end, :, :].ravel() + self.mpi_comm.Gatherv(sendbuf=sendbuf, + recvbuf=[recvbuf, orbitals_per_rank*num_spatial_points], root=rank) + if self.mpi_rank == rank: + self.local_orbital_ptip.append(recvbuf.reshape(total_orb, self.local_cell_n[0], self.local_cell_n[1], self.local_cell_n[2])) def gather_global_energies(self): @@ -261,8 +260,11 @@ def local_data_plane_above_atoms(self, local_data, height): plane_index = int(np.round(plane_z_wrt_orig/self.local_cell[2]*self.local_cell_n[2])) return local_data[:, :, plane_index] - def s_p_type_signal(self, orbital, orbital_ptip, p_tip_ratio): - return (1.0 - p_tip_ratio) * orbital**2 + p_tip_ratio * orbital_ptip + def s_p_type_signal(self, i_spin, i_mo, p_tip_ratio): + if p_tip_ratio == 0.0: + return self.local_orbitals[i_spin][i_mo]**2 + else: + return (1.0 - p_tip_ratio) * self.local_orbitals[i_spin][i_mo]**2 + p_tip_ratio * self.local_orbital_ptip[i_spin][i_mo] def build_stm_series(self, e_arr, fwhms, heights, isovalues, p_tip_ratio=0.0): @@ -292,35 +294,25 @@ def index_energy(inp): for i_e, e in enumerate(e_arr): # --------------------- # Contributing orbitals in the energy range since last energy value - close_energies = [] - close_grids = [] - close_grids_ptip = [] + close_indexes = [] for ispin in range(self.nspin): e1 = np.min([last_e, e]) e2 = np.max([last_e, e]) close_i1 = index_energy(self.global_morb_energies[ispin] > e1 - 2.0*fwhm) close_i2 = index_energy(self.global_morb_energies[ispin] > e2 + 2.0*fwhm) - #if close_i2 is not None: - # close_i2 += 1 - close_energies.append(self.global_morb_energies[ispin][close_i1:close_i2]) - close_grids.append(self.local_orbitals[ispin][close_i1:close_i2]) - close_grids_ptip.append(self.local_orbital_ptip[ispin][close_i1:close_i2]) - - #print("fwhm %.2f, energy range %.4f : %.4f" % (fwhm, last_e, e)) - #if close_i2 is not None: - # print("---- contrib %d:%d" % (close_i1, close_i2)) - #else: - # print("---- contrib %d:" % (close_i1)) - #print("---- ens:" + str(close_energies)) + if close_i1 is None: + close_i1 = 0 + if close_i2 is None: + close_i2 = len(self.global_morb_energies[ispin]) + close_indexes.append(np.arange(close_i1, close_i2)) # --------------------- # Update charge density for ispin in range(self.nspin): - for i_m, morb_en in enumerate(close_energies[ispin]): + for i_mo in close_indexes[ispin]: + morb_en = self.global_morb_energies[ispin][i_mo] broad_factor = self.gaussian_area(last_e, e, morb_en, fwhm) - cur_charge_dens += broad_factor*( - self.s_p_type_signal(close_grids[ispin][i_m], close_grids_ptip[ispin][i_m], p_tip_ratio) - ) + cur_charge_dens += broad_factor*self.s_p_type_signal(ispin, i_mo, p_tip_ratio) # --------------------- # find surfaces corresponding to isovalues @@ -329,23 +321,26 @@ def index_energy(inp): i_isosurf = self._get_isosurf_indexes(cur_charge_dens, isoval, True) cc_map[i_fwhm, i_iso, :, :, i_e] = self._index_with_interpolation(i_isosurf, self.z_arr) - for ispin in range(self.nspin): - for i_m, morb_en in enumerate(close_energies[ispin]): + for ispin in range(self.nspin): + for i_mo in close_indexes[ispin]: + morb_en = self.global_morb_energies[ispin][i_mo] morb_on_surf = self._index_with_interpolation_3d( i_isosurf, - self.s_p_type_signal(close_grids[ispin][i_m], close_grids_ptip[ispin][i_m], p_tip_ratio) + self.s_p_type_signal(ispin, i_mo, p_tip_ratio) ) cc_ldos[i_fwhm, i_iso, :, :, i_e] += self.gaussian(e - morb_en, fwhm) * morb_on_surf + # --------------------- # find constant height images for i_h, height in enumerate(heights): ch_map[i_fwhm, i_h, :, :, i_e] = self.local_data_plane_above_atoms(cur_charge_dens, height) - for ispin in range(self.nspin): - for i_m, morb_en in enumerate(close_energies[ispin]): + for ispin in range(self.nspin): + for i_mo in close_indexes[ispin]: + morb_en = self.global_morb_energies[ispin][i_mo] morb_on_plane = self.local_data_plane_above_atoms( - self.s_p_type_signal(close_grids[ispin][i_m], close_grids_ptip[ispin][i_m], p_tip_ratio), + self.s_p_type_signal(ispin, i_mo, p_tip_ratio), height ) ch_ldos[i_fwhm, i_h, :, :, i_e] += self.gaussian(e - morb_en, fwhm) * morb_on_plane @@ -559,14 +554,6 @@ def create_orb_series(self, orb_indexes, height_list=[], isoval_list=[], fwhm_li + len(isoval_list) * 2*len(fwhm_list) * len(self.p_tip_ratios) # p-type cc sts & stm signals ) - #number_of_series = ( - # len(height_list) # just s-type WFN LDOS - # + len(height_list) * len(self.p_tip_ratios) # p-type WFN ch-signals - # + len(isoval_list) * len(self.p_tip_ratios) # p-type WFN cc-signals - # + len(height_list) * 2*len(fwhm_list) # p-type ch sts & stm signals - # + len(isoval_list) * 2*len(fwhm_list) # p-type cc sts & stm signals - #) - # Orbital series for i_spin in range(self.nspin): label = 's%d_orb' % i_spin @@ -611,10 +598,9 @@ def create_orb_series(self, orb_indexes, height_list=[], isoval_list=[], fwhm_li # series data i_orb_count = 0 for i_mo in orb_indexes_wrt_data_start[i_spin]: - orb_plane = self.local_data_plane_above_atoms(self.local_orbitals[i_spin][i_mo], h) - orb_plane_ptip = self.local_data_plane_above_atoms(self.local_orbital_ptip[i_spin][i_mo], h) + s_p_data = self.s_p_type_signal(i_spin, i_mo, p_tip_ratio) self.series_output[label]['series_data'][i_series_counter, i_orb_count, :, :] = ( - self.s_p_type_signal(orb_plane, orb_plane_ptip, p_tip_ratio) + self.local_data_plane_above_atoms(s_p_data, h) ) i_orb_count += 1 i_series_counter += 1 @@ -635,7 +621,7 @@ def create_orb_series(self, orb_indexes, height_list=[], isoval_list=[], fwhm_li i_orb_count = 0 for i_mo in orb_indexes_wrt_data_start[i_spin]: i_isosurf = self._get_isosurf_indexes( - self.s_p_type_signal(self.local_orbitals[i_spin][i_mo], self.local_orbital_ptip[i_spin][i_mo], p_tip_ratio), + self.s_p_type_signal(i_spin, i_mo, p_tip_ratio), isov, True ) self.series_output[label]['series_data'][i_series_counter, i_orb_count, :, :] = ( diff --git a/examples/benzene_cube_from_wfn/run_cube_from_wfn.sh b/examples/benzene_cube_from_wfn/run_cube_from_wfn.sh index b32f912..0034e26 100644 --- a/examples/benzene_cube_from_wfn/run_cube_from_wfn.sh +++ b/examples/benzene_cube_from_wfn/run_cube_from_wfn.sh @@ -1,12 +1,10 @@ -DIR="../benzene_cp2k_scf/" +#!/bin/bash -l -# Path to cube_from_wfn.py -# Leave empty if already in $PATH -SCRIPT_PATH="" +DIR="../benzene_cp2k_scf/" mkdir cubes -"$SCRIPT_PATH"cube_from_wfn.py \ +../../cube_from_wfn.py \ --cp2k_input_file $DIR/cp2k.inp \ --basis_set_file ../BASIS_MOLOPT \ --xyz_file $DIR/geom.xyz \ diff --git a/examples/benzene_stm/run_stm_sts_from_wfn.sh b/examples/benzene_stm/run_stm_sts_from_wfn.sh index e722d2b..485d2e1 100644 --- a/examples/benzene_stm/run_stm_sts_from_wfn.sh +++ b/examples/benzene_stm/run_stm_sts_from_wfn.sh @@ -17,10 +17,10 @@ mpirun -n 2 python3 ../../stm_sts_from_wfn.py \ --dx 0.15 \ --eval_cutoff 14.0 \ --extrap_extent 2.0 \ - --p_tip_ratios 0.0 1.0\ + --p_tip_ratios 0.0 1.0 \ \ - --n_homo 5 \ - --n_lumo 5 \ + --n_homo 4 \ + --n_lumo 4 \ --orb_heights 3.0 5.0 \ --orb_isovalues 1e-7 \ --orb_fwhms 0.1 \ diff --git a/stm_sts_from_wfn.py b/stm_sts_from_wfn.py index c5121e8..c2d4b18 100644 --- a/stm_sts_from_wfn.py +++ b/stm_sts_from_wfn.py @@ -295,6 +295,7 @@ ]) ion_pot = cube_utils.find_vacuum_level_naive(hart_cube) - (homo_en + cp2k_grid_orb.ref_energy) print("IONIZATION POTENIAL (eV): %.6f (accurate only for isolated molecules)" % ion_pot) + sys.stdout.flush() ### ------------------------------------------------------ ### Set up STM object