Develop (#183)

* Image augmentation.

* Augment images.

* Augment images.

* Specie update.

* Add Latt2D, STM image (b-1) fix, image augmentation fix.

* Add Latt2D, STM image (b-1) fix, image augmentation fix.

* Update conf.py

* Update conf.py

* Multi-output graph bacthing.

* Add EDOS dataset.

* Temp.

* Add circuit maker.

* Add circuit maker.

* NELECT update.

* Version update, more DBs added.

* Fix CHGCAR vasp.

* Added volumetric reshape for CHGCAR.

* Tmp

* Tershoff Hamman update, specie update.

* Add crop from center in STM.

* Add Fourier transfor in STM.

* Update STM pytest.

* Add DPI to STM.

* Zeo++ added, Atoms cif update, STM update, random vacancy maker added.

* Atoms tempfile fix, Potcar from atoms module added.

* Test for docs.

* C2DB link update, docs Atoms update.

* C2DB link update, docs Atoms update.

* Version update, COD DB, QM9 JCTC DB added.

jarvis/__init__.py

@@ -1,2 +1,2 @@
 """Version number."""
-__version__ = "2021.06.18"
+__version__ = "2021.07.04"

jarvis/analysis/defects/vacancy.py

@@ -4,6 +4,8 @@
 from jarvis.analysis.structure.spacegroup import Spacegroup3D
 from jarvis.core.utils import rand_select
 from jarvis.core.atoms import Atoms
+# import numpy as np
+import random
 
 
 class Vacancy(object):
@@ -102,6 +104,38 @@ def __repr__(self, indent=2):
         return pprint.pformat(self.to_dict(), indent=indent)
 
 
+def generate_random_defects(n_vacs=10, atoms=None, element=None, seed=123):
+    """Generate random defects for an element."""
+    # natoms = atoms.num_atoms
+    atoms = atoms.to_dict()
+    elements = atoms["elements"]
+    if element is None:
+        element = elements[0]
+    coords = atoms["coords"]
+    lattice_mat = atoms["lattice_mat"]
+    # ids = np.arange(natoms)
+    new_elements = []
+    new_coords = []
+    options = []
+    for ii, i in enumerate(elements):
+        if i == element:
+            options.append(ii)
+    random.seed(seed)
+    random.shuffle(options)
+    to_delete = options[0:n_vacs]
+    for ii, i in enumerate(elements):
+        if ii not in to_delete:
+            new_elements.append(i)
+            new_coords.append(coords[ii])
+    new_atoms = Atoms(
+        coords=new_coords,
+        lattice_mat=lattice_mat,
+        elements=new_elements,
+        cartesian=False,
+    )
+    return new_atoms
+
+
 """
 if __name__ == "__main__":
     from jarvis.io.vasp.inputs import Poscar

jarvis/analysis/stm/tersoff_hamann.py

@@ -4,13 +4,30 @@
 from jarvis.io.vasp.outputs import Chgcar
 import numpy as np
 import matplotlib.transforms as mtransforms
+import scipy
+from jarvis.core.image import Image
 
 
 class TersoffHamannSTM(object):
     """Generate constant height and constant current STM images."""
 
     def __init__(
-        self, chg_name="PARCHG", min_size=50.0, skew=True, zcut=None, extend=0
+        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
@@ -23,6 +40,10 @@ def __init__(
         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
@@ -37,48 +58,120 @@ def __init__(
             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])
 
-    def constant_height(self, tol=2, filename="testh.png"):
+        self.use_interpolated = use_interpolated
+        self.interp_step = interp_step
+
+        self.min_pixel = min_pixel
+        self.dpi = dpi
+
+    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)
-        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]
+        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:
-            tmp = 0
-        data = self.get_plot(
-            ax,
-            img_ext,
-            exts,
-            mtransforms.Affine2D().skew_deg(tmp, 0),
-        )
-        info = {}
-        info["img_ext"] = img_ext
-        info["data"] = data
+            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
-        fig.subplots_adjust(bottom=0, top=1, left=0.0, right=1)
-        plt.savefig(
-            filename
-        )  # , bbox_inches="tight", pad_inches=0.0, dpi=240)
-        plt.close()
+        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
 
-    def constant_current(self, tol=2, pc=None, ext=0.15, filename="testc.png"):
+    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)
@@ -98,26 +191,53 @@ def constant_current(self, tol=2, pc=None, ext=0.15, filename="testc.png"):
         # 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
-        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]
+        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:
-            tmp = 0
+            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()
 
-        data = self.get_plot(
-            ax,
-            img_ext,
-            exts,
-            mtransforms.Affine2D().skew_deg(tmp, 0),
-        )
-        plt.savefig(filename)
-        plt.close()
-        info = {}
-        info["img_ext"] = img_ext
-        info["data"] = data
         info["scell"] = self.scell
         info["zcut"] = self.zcut
         return info
@@ -148,6 +268,35 @@ def get_plot(self, ax, Z, extent, transform):
         print("min Z and maxZ", np.min(Z), np.max(Z))
         return data
 
+    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__":