Source code for jarvis.analysis.stm.tersoff_hamann

"""Module to simulate STM with Tershoff-Hamann approach."""
# Reference: https://www.nature.com/articles/s41597-021-00824-y
import matplotlib.pyplot as plt
from jarvis.io.vasp.outputs import Chgcar
import numpy as np
import matplotlib.transforms as mtransforms
import scipy
from jarvis.core.image import Image


[docs]class TersoffHamannSTM(object): """Generate constant height and constant current STM images.""" def __init__( self, chg_name="PARCHG", min_size=50.0, height_tol=2, zcut=None, use_interpolated=True, crop_from_center=False, crop_mult=3, min_pixel=288, interp_step=1.0, skew=True, extend=0, ft_image_path=None, ft_image_zoom_factor=5, dpi=240, cmap="gray", ): """Initialize class with pathe of PARCHG and other input params.""" # In original paper, extend used as 1 chgcar = Chgcar(filename=chg_name) self.atoms = chgcar.atoms self.dim = chgcar.dim self.zmaxp = self.atoms.cart_coords[:, 2].max() self.nz = self.dim[2] volume = self.atoms.volume tmp = chgcar.chg[-1] * volume chg = tmp.reshape(self.dim[::-1]).T self.chg = chg self.cmap = cmap self.ft_image_path = ft_image_path self.ft_image_zoom_factor = ft_image_zoom_factor self.height_tol = height_tol self.a = self.atoms.lattice.a self.b = self.atoms.lattice.b self.c = self.atoms.lattice.c self.skew = skew self.zcut = zcut self.extend = extend z_frac_coords = self.atoms.frac_coords[:, 2] z_frac_coords_moved = [] for i in z_frac_coords: if i > 0.5: i = i - 1 elif i < -0.5: i = i + 1 z_frac_coords_moved.append(i) self.crop_mult = crop_mult self.crop_from_center = crop_from_center if self.crop_from_center: min_size = min_size * self.crop_mult self.min_size = min_size self.zmaxp = max(np.array(z_frac_coords_moved) * self.c) rep_x = int(min_size / self.a) + self.extend rep_y = int(min_size / self.b) + self.extend self.repeat = [rep_x, rep_y] self.scell = self.atoms.make_supercell_matrix([rep_x, rep_y, 1]) self.use_interpolated = use_interpolated self.interp_step = interp_step self.min_pixel = min_pixel self.dpi = dpi
[docs] def constant_height( self, filename="testh.png", ): """Get iso-height image.""" tol = self.height_tol if not self.zcut: self.zcut = int((self.zmaxp + tol) / self.c * self.nz) # print("zcut", self.zcut, self.repeat) info = {} img_ext = np.tile(self.chg[:, :, self.zcut], self.repeat) if not self.use_interpolated: exts = ( 0, self.a * self.repeat[0], 0, self.b * (self.repeat[1] - 1), ) plt.close() fig, ax = plt.subplots() plt.xticks([]) plt.yticks([]) if self.skew: tmp = 90 - self.atoms.lattice.angles[2] else: tmp = 0 data = self.get_plot( ax, img_ext, exts, mtransforms.Affine2D().skew_deg(tmp, 0), ) info["data"] = data info["img_ext"] = img_ext fig.subplots_adjust(bottom=0, top=1, left=0.0, right=1) plt.savefig( filename, dpi=self.dpi ) # , bbox_inches="tight", pad_inches=0.0, dpi=240) plt.close() else: img_ext = self.get_interpolated_data(img_data=img_ext) # print ('img_ext',img_ext) plt.close() # fig, ax = plt.subplots() plt.imshow(img_ext, interpolation="none", cmap=self.cmap) # ax.set_aspect('equal') plt.axis("off") plt.savefig(filename, dpi=self.dpi) plt.close() plt.gca().set_axis_off() plt.subplots_adjust( top=1, bottom=0, right=1, left=0, hspace=0, wspace=0 ) plt.margins(0, 0) plt.gca().xaxis.set_major_locator(plt.NullLocator()) plt.gca().yaxis.set_major_locator(plt.NullLocator()) # print ('img_ext',img_ext) plt.close() info["scell"] = self.scell info["zcut"] = self.zcut info["img_ext"] = img_ext if self.crop_from_center: im = Image.from_file(filename) p_x = int(self.min_pixel / self.crop_mult) # print('shape',im.values.shape,p_x) im = Image.crop_from_center( image_path=filename, target_left=p_x, target_right=p_x ) info["img_ext"] = im plt.imshow(im, cmap=self.cmap) plt.tight_layout() plt.axis("off") plt.savefig(filename, dpi=self.dpi) plt.close() if self.ft_image_path is not None: im = Image.from_file(filename).fourier_transform2D( zoom_factor=self.ft_image_zoom_factor ) im.save(self.ft_image_path) return info
[docs] def constant_current( self, pc=None, ext=0.15, filename="testc.png", use_interpolated=True, ): """Return the constant-current cut the charge density.""" zmax_ind = int(self.zmaxp / self.c * self.nz) + 1 info = {} tol = self.height_tol # Find what z value is near the current, and take avergae if not self.zcut: self.zcut = int((self.zmaxp + tol) / self.c * self.nz) zext = int(self.nz * ext) zcut_min = self.zcut - zext if zcut_min < zmax_ind: zcut_min = zmax_ind zcut_max = self.zcut + zext if pc is None: c = np.average(self.chg[:, :, self.zcut]) else: tmp = np.zeros(zcut_max - zcut_min) for ii in range(tmp.size): tmp[ii] = np.average(self.chg[:, :, zcut_min + ii]) c = np.linspace(tmp.min(), tmp.max(), 100)[pc] # height of iso-current img = np.argmin(np.abs(self.chg[:, :, zcut_min:zcut_max] - c), axis=2) img_ext = np.tile(img, self.repeat[::-1]) + self.zcut - zext if not use_interpolated: fig, ax = plt.subplots() exts = ( 0, self.a * self.repeat[0], 0, self.b * (self.repeat[1] - 1), ) plt.xticks([]) plt.yticks([]) if self.skew: tmp = 90 - self.atoms.lattice.angles[2] else: tmp = 0 data = self.get_plot( ax, img_ext, exts, mtransforms.Affine2D().skew_deg(tmp, 0), ) info["data"] = data info["img_ext"] = img_ext plt.savefig(filename, dpi=self.dpi) plt.close() else: img_ext = self.get_interpolated_data(img_data=img_ext) plt.close() # fig, ax = plt.subplots() plt.imshow(img_ext, interpolation="none", cmap=self.cmap) # ax.set_aspect('equal') plt.axis("off") plt.savefig(filename, dpi=self.dpi) plt.close() if self.crop_from_center: im = Image.from_file(filename) p_x = int(self.min_pixel / self.crop_mult) # print('shape',im.values.shape,p_x) im = Image.crop_from_center( image_path=filename, target_left=p_x, target_right=p_x ) plt.imshow(im, cmap=self.cmap) plt.axis("off") plt.savefig(filename, dpi=self.dpi) plt.close() info["scell"] = self.scell info["zcut"] = self.zcut return info
[docs] def get_plot(self, ax, Z, extent, transform): """Apply affine transformation.""" ax.axis("off") im = ax.imshow( Z, interpolation="none", origin="lower", extent=extent, clip_on=True, aspect="equal", ) # ,cmap=plt.get_cmap('gray') trans_data = transform + ax.transData im.set_transform(trans_data) # display intended extent of the image x1, x2, y1, y2 = im.get_extent() data = ax.plot( [x1, x2, x2, x1, x1], [y1, y1, y2, y2, y1], "y--", transform=trans_data, ) print("min Z and maxZ", np.min(Z), np.max(Z)) return data
[docs] def get_interpolated_data(self, img_data=[]): """Get interpolated data after moving pixels.""" x = [] y = [] zz = [] xy = [] step = self.interp_step atoms = self.atoms for i in range(img_data.shape[0]): for j in range(img_data.shape[1]): z = img_data[i][j] xyz = i * atoms.lattice_mat[0] + j * atoms.lattice_mat[1] x.append(xyz[0]) y.append(xyz[1]) zz.append(z) xy.append([xyz[0], xyz[1]]) grid_x, grid_y = np.mgrid[ min(x) : max(x) : step, min(y) : max(y) : step ] # stepx=(max(x)-min(x))/bins # stepy=(max(y)-min(y))/bins # grid_x, grid_y = np.mgrid[min(x):max(x):stepx, min(y):max(y):stepy] interp = scipy.interpolate.griddata(xy, zz, (grid_x, grid_y)).T # plt.scatter(x, y, 1, zz) # plt.savefig("ok.png") # plt.close() return interp
""" if __name__ == "__main__": plt.switch_backend("agg") f = "PARCHG" t = TersoffHamannSTM(chg_name=f).constant_height() t = TersoffHamannSTM(chg_name=f).constant_current() """