Develop (#239)

* 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. * Compostion bug fix, elemental descriptor added. * Develop (#186) * Update outputs.py I added the calculation of the Raman intensities inside parse_raman_dat * Update outputs.py * Update outputs.py * Update outputs.py * Update cfid.py * Delete __init__.py * stylecss added. * stylecss added. * Adding extra Makefile/ * Remove examples from docs. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Tutorials update. * Tutorials docs update. * Docs update,pdb reader updated. * Update action_build.yml * Update action_build.yml * Remove pytraj strong dependencies. * Update docs, Added PDBBind and HPOV datasets. * Docs update. * Add thcikness to surface builder. * Surface builder update, Chemical only magpie descriptors added, pdb_core dataset added, zeopp tempfile bugfix. * Typo fix. * Add names to chem descs. * Lessen hermsolver pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * No DFT3D * Exclude dft_3d dataset for memory issue. * Update figshare test. * Update figshare test. * Exclude db from coverage. * Exclude db from coverage. * Add magpie.json. * Add magpie.json. * Wien2k bands bug fix. * Wien2k bands bug fix. * Update JARVIS-FF,Elastictensor,LAMMPS parse folder, VASP bandstructure plot code. * JFF update. * Add JQE_TB3 and hMOF dataset. * Update LAMMPS module. * Update LAMMPS module. * Fix elastic tensor module. * Figshare update, docs db name update. * Substitutions. * Update figshare dft_3d, cfid_3d. * Docs data update. * Generate substitutions. * Lint fix. * Update DOS. * Update DOS. * Adding folders for nexus setup. * Update QMOF and hMOF info. * Fixing auto klength error. * Adding zeopp surface area. * Vacancy bug fix, added hmof to docs. * vacancy update. * QE inputs update. * Pyhon 3.9 test, added QE el-ph, VASP XANES, modified hmof db details. * Update Contribution.rst * Minor lint fix. * Update action_build.yml * Update action_build.yml * Fix qiskit requirements. * Fix phonopy requirements. * Fix all requirements. * Fix phonopy requirements. * QE test. * QE test. * Update action_build. * Try other python versions. * Try other python versions. * README updates. * Adding nexus. (#197) * Develop (#196) * 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. * Compostion bug fix, elemental descriptor added. * Develop (#186) * Update outputs.py I added the calculation of the Raman intensities inside parse_raman_dat * Update outputs.py * Update outputs.py * Update outputs.py * Update cfid.py * Delete __init__.py * stylecss added. * stylecss added. * Adding extra Makefile/ * Remove examples from docs. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Tutorials update. * Tutorials docs update. * Docs update,pdb reader updated. * Update action_build.yml * Update action_build.yml * Remove pytraj strong dependencies. * Update docs, Added PDBBind and HPOV datasets. * Docs update. * Add thcikness to surface builder. * Surface builder update, Chemical only magpie descriptors added, pdb_core dataset added, zeopp tempfile bugfix. * Typo fix. * Add names to chem descs. * Lessen hermsolver pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * No DFT3D * Exclude dft_3d dataset for memory issue. * Update figshare test. * Update figshare test. * Exclude db from coverage. * Exclude db from coverage. * Add magpie.json. * Add magpie.json. * Wien2k bands bug fix. * Wien2k bands bug fix. * Update JARVIS-FF,Elastictensor,LAMMPS parse folder, VASP bandstructure plot code. * JFF update. * Add JQE_TB3 and hMOF dataset. * Update LAMMPS module. * Update LAMMPS module. * Fix elastic tensor module. * Figshare update, docs db name update. * Substitutions. * Update figshare dft_3d, cfid_3d. * Docs data update. * Generate substitutions. * Lint fix. * Update DOS. * Update DOS. Co-authored-by: tavazza <tavazza@gmail.com> Co-authored-by: knc6 <kamal.choudhary@nist.gov> Co-authored-by: KAMAL CHOUDHARY <kamalch@login1.frontera.tacc.utexas.edu> * First input. * added black * Example folder. Co-authored-by: Kamal Choudhary <writetokamal.1989@gmail.com> Co-authored-by: tavazza <tavazza@gmail.com> Co-authored-by: knc6 <kamal.choudhary@nist.gov> Co-authored-by: KAMAL CHOUDHARY <kamalch@login1.frontera.tacc.utexas.edu> * Revert back. * Update .readthedocs.yaml * Update dev-requirements.txt * Docs requirements update. * Update .readthedocs.yaml * Update requirements.txt * Update .readthedocs.yaml * Update requirements.txt * Update .readthedocs.yaml * Update requirements.txt * Update requirements.txt * Update requirements.txt * Update requirements.txt * Update requirements.txt * Update atoms.py * Fixe pbc in ase_to_Atoms. * Add installation tests (#214) * QE inputs, XANES, GHAction updates. (#210) * 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. * Compostion bug fix, elemental descriptor added. * Develop (#186) * Update outputs.py I added the calculation of the Raman intensities inside parse_raman_dat * Update outputs.py * Update outputs.py * Update outputs.py * Update cfid.py * Delete __init__.py * stylecss added. * stylecss added. * Adding extra Makefile/ * Remove examples from docs. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Tutorials update. * Tutorials docs update. * Docs update,pdb reader updated. * Update action_build.yml * Update action_build.yml * Remove pytraj strong dependencies. * Update docs, Added PDBBind and HPOV datasets. * Docs update. * Add thcikness to surface builder. * Surface builder update, Chemical only magpie descriptors added, pdb_core dataset added, zeopp tempfile bugfix. * Typo fix. * Add names to chem descs. * Lessen hermsolver pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * No DFT3D * Exclude dft_3d dataset for memory issue. * Update figshare test. * Update figshare test. * Exclude db from coverage. * Exclude db from coverage. * Add magpie.json. * Add magpie.json. * Wien2k bands bug fix. * Wien2k bands bug fix. * Update JARVIS-FF,Elastictensor,LAMMPS parse folder, VASP bandstructure plot code. * JFF update. * Add JQE_TB3 and hMOF dataset. * Update LAMMPS module. * Update LAMMPS module. * Fix elastic tensor module. * Figshare update, docs db name update. * Substitutions. * Update figshare dft_3d, cfid_3d. * Docs data update. * Generate substitutions. * Lint fix. * Update DOS. * Update DOS. * Adding folders for nexus setup. * Update QMOF and hMOF info. * Fixing auto klength error. * Adding zeopp surface area. * Vacancy bug fix, added hmof to docs. * vacancy update. * QE inputs update. * Pyhon 3.9 test, added QE el-ph, VASP XANES, modified hmof db details. * Update Contribution.rst * Minor lint fix. * Update action_build.yml * Update action_build.yml * Fix qiskit requirements. * Fix phonopy requirements. * Fix all requirements. * Fix phonopy requirements. * QE test. * QE test. * Update action_build. * Try other python versions. * Try other python versions. * README updates. * Adding nexus. (#197) * Develop (#196) * 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. * Compostion bug fix, elemental descriptor added. * Develop (#186) * Update outputs.py I added the calculation of the Raman intensities inside parse_raman_dat * Update outputs.py * Update outputs.py * Update outputs.py * Update cfid.py * Delete __init__.py * stylecss added. * stylecss added. * Adding extra Makefile/ * Remove examples from docs. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Docs update. * Tutorials update. * Tutorials docs update. * Docs update,pdb reader updated. * Update action_build.yml * Update action_build.yml * Remove pytraj strong dependencies. * Update docs, Added PDBBind and HPOV datasets. * Docs update. * Add thcikness to surface builder. * Surface builder update, Chemical only magpie descriptors added, pdb_core dataset added, zeopp tempfile bugfix. * Typo fix. * Add names to chem descs. * Lessen hermsolver pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * Reduced pytest. * No DFT3D * Exclude dft_3d dataset for memory issue. * Update figshare test. * Update figshare test. * Exclude db from coverage. * Exclude db from coverage. * Add magpie.json. * Add magpie.json. * Wien2k bands bug fix. * Wien2k bands bug fix. * Update JARVIS-FF,Elastictensor,LAMMPS parse folder, VASP bandstructure plot code. * JFF update. * Add JQE_TB3 and hMOF dataset. * Update LAMMPS module. * Update LAMMPS module. * Fix elastic tensor module. * Figshare update, docs db name update. * Substitutions. * Update figshare dft_3d, cfid_3d. * Docs data update. * Generate substitutions. * Lint fix. * Update DOS. * Update DOS. Co-authored-by: tavazza <tavazza@gmail.com> Co-authored-by: knc6 <kamal.choudhary@nist.gov> Co-authored-by: KAMAL CHOUDHARY <kamalch@login1.frontera.tacc.utexas.edu> * First input. * added black * Example folder. Co-authored-by: Kamal Choudhary <writetokamal.1989@gmail.com> Co-authored-by: tavazza <tavazza@gmail.com> Co-authored-by: knc6 <kamal.choudhary@nist.gov> Co-authored-by: KAMAL CHOUDHARY <kamalch@login1.frontera.tacc.utexas.edu> * Revert back. * Update .readthedocs.yaml * Update dev-requirements.txt * Docs requirements update. * Update .readthedocs.yaml * Update requirements.txt * Update .readthedocs.yaml * Update requirements.txt * Update .readthedocs.yaml * Update requirements.txt * Update requirements.txt * Update requirements.txt * Update requirements.txt * Update requirements.txt * Update atoms.py * Fixe pbc in ase_to_Atoms. Co-authored-by: tavazza <tavazza@gmail.com> Co-authored-by: knc6 <kamal.choudhary@nist.gov> Co-authored-by: KAMAL CHOUDHARY <kamalch@login1.frontera.tacc.utexas.edu> Co-authored-by: wines1 <74620550+wines1@users.noreply.github.com> * Add installation tests * Fix codestyle * Fix codestyle with black * Add docstrings * Fix pydocstyle error * Update __init__.py * Update __init__.py Co-authored-by: Kamal Choudhary <writetokamal.1989@gmail.com> Co-authored-by: tavazza <tavazza@gmail.com> Co-authored-by: knc6 <kamal.choudhary@nist.gov> Co-authored-by: KAMAL CHOUDHARY <kamalch@login1.frontera.tacc.utexas.edu> Co-authored-by: wines1 <74620550+wines1@users.noreply.github.com> * Adding QE super. * Minor changes to QE module, atoms xyz fix. * Adding qe_tb info, and version update. * Update __init__.py * WIP super QE. * Working version of ScSi. * QE inputs and task update. * Add master super. * Add master super. * Lint fix. * Lint fix. * Minor fix. * ET update. * Fix ET test. * Update sanitize atoms. * Additonal checks on supercond. * Debye bug fix. * Pressure in QE Super. * Version fix, publication update, supercond workflow update. * Lint fix. * Tensorboard fix. * Tensorboard fix. * Tensorboard fix. * Melting temp fix. * Update vasp.py (#234) * Local tetra tmp. * Version update. * Lint fix. * HSE06 * Tmp. * Vacancy update, Optimade structureand other minor lint updates. * STEM pytest fix. * Minor lint fix. Co-authored-by: tavazza <tavazza@gmail.com> Co-authored-by: knc6 <kamal.choudhary@nist.gov> Co-authored-by: KAMAL CHOUDHARY <kamalch@login1.frontera.tacc.utexas.edu> Co-authored-by: wines1 <74620550+wines1@users.noreply.github.com> Co-authored-by: Saurav Maheshkar <sauravvmaheshkar@gmail.com>
usnistgov · May 21, 2022 · 1364444 · 1364444
1 parent 30a2c0d
commit 1364444
Show file tree

Hide file tree

Showing 13 changed files with 332 additions and 115 deletions.
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -63,7 +63,7 @@ mpmath==1.2.1
 multitasking==0.0.10
 networkx==2.6.3
 ntlm-auth==1.5.0
-numpy==1.21.0
+numpy==1.20.0
 oauthlib==3.1.1
 opencv-python==4.5.4.60
 opt-einsum==3.3.0

diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -64,7 +64,7 @@ mpmath==1.2.1
 multitasking==0.0.10
 networkx==2.6.3
 ntlm-auth==1.5.0
-numpy==1.21.0
+numpy==1.20.0
 oauthlib==3.1.1
 opencv-python==4.5.4.60
 opt-einsum==3.3.0

diff --git a/jarvis/__init__.py b/jarvis/__init__.py
@@ -1,5 +1,5 @@
 """Version number."""
-__version__ = "2022.05.05"
+__version__ = "2022.05.20"
 
 import os
 

diff --git a/jarvis/ai/descriptors/elemental.py b/jarvis/ai/descriptors/elemental.py
@@ -6,6 +6,7 @@
 
 
 def get_element_fraction_desc(formula="SiO2", max_nelements=103):
+    """Get element fraction."""
     x = np.zeros(max_nelements)
     fracs = Composition.from_string(formula).atomic_fraction
     for i, j in fracs.items():

diff --git a/jarvis/analysis/defects/vacancy.py b/jarvis/analysis/defects/vacancy.py
@@ -2,7 +2,7 @@
 import pprint
 from collections import OrderedDict
 from jarvis.analysis.structure.spacegroup import Spacegroup3D
-from jarvis.core.utils import rand_select
+# from jarvis.core.utils import rand_select
 from jarvis.core.atoms import Atoms
 
 # import numpy as np
@@ -26,13 +26,9 @@ def __init__(
         Arguments are given below.
         Args:
             atoms: jarvis.core.Atoms object.
-
             defect_index: atoms index for defect.
-
             defect_structure:  Atoms with defect.
-
             wyckoff_multiplicity: Wyckoff multiplicity.
-
             symbol: Elemenyt symbol.
         """
         self._atoms = atoms
@@ -53,7 +49,11 @@ def from_dict(self, d={}):
         )
 
     def generate_defects(
-        self, enforce_c_size=10.0, on_conventional_cell=False, extend=1
+        self,
+        enforce_c_size=10.0,
+        on_conventional_cell=False,
+        extend=1,
+        using_wyckoffs=True,
     ):
         """Provide function to generate defects."""
         atoms = self._atoms
@@ -74,13 +74,25 @@ def generate_defects(
             if dim3 == 0:
                 dim3 = 1
             supercell_size = [dim1, dim2, dim3]
-        spg = Spacegroup3D(atoms)
-        wyckoffs = spg._dataset["wyckoffs"]
-        atoms.props = wyckoffs
+        if using_wyckoffs:
+            spg = Spacegroup3D(atoms)
+            wyckoffs = spg._dataset["wyckoffs"]
+            atoms.props = wyckoffs
+        else:
+            wyckoffs = ["a" for i in range(atoms.num_atoms)]
+            atoms.props = wyckoffs
         supercell = atoms.make_supercell(supercell_size)
-        props = rand_select(supercell.props)
+        # props = rand_select(supercell.props)
+        new_props = {}
+        for ii, jj, kk in zip(
+            supercell.props, supercell.elements, range(supercell.num_atoms)
+        ):
+            if ii + "_" + jj not in new_props:
+                new_props[ii + "_" + jj] = kk
+        # print ('props',props)
+        # print ('newprops',new_props)
         vacs = []
-        for i, j in props.items():
+        for i, j in new_props.items():
             defect_strt = supercell.remove_site_by_index(j)
             vac = Vacancy(
                 atoms=supercell,
@@ -140,28 +152,3 @@ def generate_random_defects(n_vacs=10, atoms=None, element=None, seed=123):
         cartesian=False,
     )
     return new_atoms
-
-
-"""
-if __name__ == "__main__":
-    from jarvis.io.vasp.inputs import Poscar
-
-    box = [[2.715, 2.715, 0], [0, 2.715, 2.715], [2.715, 0, 2.715]]
-    coords = [[0, 0, 0], [0.25, 0.25, 0.25]]
-    elements = ["Si", "Si"]
-    Si = Atoms(lattice_mat=box, coords=coords, elements=elements)
-    Si = Poscar.from_file(
-        "/rk2/knc6/JARVIS-DFT/2D-bulk/mp-1143_bulk_LDA/MAIN-ELASTIC-bulk@mp_1143/POSCAR"
-    ).atoms
-    vacs = Vacancy(atoms=Si).generate_defects()
-    for i in vacs:
-        print(i)
-    spg = Spacegroup3D(Si)
-    cvn = spg.conventional_standard_structure
-    spg = Spacegroup3D(cvn)
-    props = spg._dataset["wyckoffs"]
-    Si.props = props
-    ss = Si.make_supercell([2, 2, 2])
-    props = ss.props
-    # print (rand_select(props))
-"""
diff --git a/jarvis/analysis/stem/convolution_apprx.py b/jarvis/analysis/stem/convolution_apprx.py
@@ -6,6 +6,8 @@
 # from numbers import Number
 from jarvis.core.utils import gaussian
 from jarvis.core.utils import lorentzian2 as lorentzian
+from jarvis.core.atoms import Atoms  # , get_supercell_dims, crop_square
+from typing import List
 
 
 class STEMConv(object):
@@ -15,16 +17,15 @@ def __init__(
         self,
         atoms=None,
         output_size=[50, 50],
-        power_factor=1.4,
+        power_factor=1.7,
         gaussian_width=0.5,
         lorentzian_width=0.5,
         intensity_ratio=0.5,
         nbins=100,
         tol=0.5,
+        crop=False,
     ):
-        """
-        Intitialize the class.
-        """
+        """Intitialize the class."""
         self.atoms = atoms
         self.output_size = output_size
         self.power_factor = power_factor
@@ -33,9 +34,9 @@ def __init__(
         self.intensity_ratio = intensity_ratio
         self.nbins = nbins
         self.tol = tol
+        self.crop = crop
 
-    @staticmethod
-    def superpose_deltas(positions, array):
+    def superpose_deltas(self, positions, array):
         """Superpose deltas."""
         z = 0
         shape = array.shape[-2:]
@@ -54,59 +55,140 @@ def superpose_deltas(positions, array):
         array[z, (rows + 1) % shape[0], (cols + 1) % shape[1]] += (
             rows - positions[:, 0]
         ) * (cols - positions[:, 1])
+        return array
+
+    def simulate_surface(
+        self,
+        atoms: Atoms,
+        px_scale: float = 0.2,
+        eps: float = 0.6,
+        rot: float = 0,
+        shift: List = [0, 0],
+    ):
+        """Simulate a STEM image.
 
-    def simulate_surface(self):
-        """Simulate a STEM image."""
+        atoms: jarvis.core.Atoms material slab
+        px_scale: pixel size in angstroms/px
+        eps: tolerance factor (angstroms)
+        for rendering atoms outside the field of view
+        rot: rotation about the image center (degrees)
+        shift: rigid translation of field of view [dx, dy] (angstroms)
 
-        extent = np.diag(self.atoms.lattice_mat)[:2]
+        """
+        shift = np.squeeze(shift)
+        output_px = np.squeeze(self.output_size)  # px
 
-        # shape = 1
-        sampling = (
-            extent[0] / self.output_size[0],
-            extent[1] / self.output_size[1],
-        )
+        # field of view size in angstroms
+        view_size = px_scale * (output_px - 1)
 
-        margin = int(np.ceil(5 / min(sampling)))  # int like 20
-        shape_w_margin = (
-            self.output_size[0] + 2 * margin,
-            self.output_size[1] + 2 * margin,
-        )
-        # Set up a grid
-        x = np.fft.fftfreq(shape_w_margin[0]) * shape_w_margin[1] * sampling[0]
-        y = np.fft.fftfreq(shape_w_margin[1]) * shape_w_margin[1] * sampling[1]
-        r = np.sqrt(x[:, None] ** 2 + y[None] ** 2)
+        # construct a supercell grid big enough to fill the field of view
+        cell_extent = atoms.lattice.abc[0:2]  # np.diag(atoms.lattice_mat)[:2]
 
-        # proble profile
+        cells = ((view_size // cell_extent) + 1).astype(int)
+        # print ('cells',cells)
+        atoms = atoms.make_supercell_matrix((3 * cells[0], 3 * cells[1], 1))
 
+        # Set up real-space grid (in angstroms)
+        # construct the probe array with the output target size
+        # fftshift, pad, un-fftshift
+        x = np.fft.fftfreq(output_px[0]) * output_px[0] * px_scale
+        y = np.fft.fftfreq(output_px[1]) * output_px[1] * px_scale
+        r = np.sqrt(x[:, None] ** 2 + y[None] ** 2)
+
+        # construct the probe profile centered
+        # at (0,0) on the periodic spatial grid
         x = np.linspace(0, 4 * self.lorentzian_width, self.nbins)
         profile = gaussian(
             x, self.gaussian_width
         ) + self.intensity_ratio * lorentzian(x, self.lorentzian_width)
-
         profile /= profile.max()
         f = interp1d(x, profile, fill_value=0, bounds_error=False)
         intensity = f(r)
-        positions = self.atoms.cart_coords[:, :2] / sampling - self.tol
-        # Check if atoms are within the specified range
-        inside = (
-            (positions[:, 0] > -margin)
-            & (positions[:, 1] > -margin)
-            & (positions[:, 0] < self.output_size[0] + margin)
-            & (positions[:, 1] < self.output_size[1] + margin)
-        )
-        positions = positions[inside] + margin
-        numbers = np.array(self.atoms.atomic_numbers)[inside]
-
-        array = np.zeros((1,) + shape_w_margin)  # adding extra 1
-        for number in np.unique(np.array(self.atoms.atomic_numbers)):
-            temp = np.zeros((1,) + shape_w_margin)
-            self.superpose_deltas(positions[numbers == number], temp)
 
+        # shift the probe profile to the center
+        # apply zero-padding, and shift back to the origin
+        margin = int(np.ceil(5 / px_scale))  # int like 20
+        intensity = np.fft.fftshift(intensity)
+        intensity = np.pad(intensity, (margin, margin))
+        intensity = np.fft.fftshift(intensity)
+
+        # project atomic coordinates onto the image
+        # center them as well
+        centroid = np.mean(atoms.cart_coords[:, :2], axis=0)
+
+        # center atom positions around (0,0)
+        pos = atoms.cart_coords[:, :2] - centroid
+
+        # apply field of view rotation
+        # (actually rotate the lattice coordinates)
+        if rot != 0:
+            rot = np.radians(rot)
+            R = np.array(
+                [[np.cos(rot), -np.sin(rot)], [np.sin(rot), np.cos(rot)]]
+            )
+            pos = pos @ R
+
+        # shift to center of image
+        pos += view_size / 2
+
+        # apply rigid translation of atoms wrt image field of view
+        pos += shift
+
+        # select only atoms in field of view
+        in_view = (
+            (pos[:, 0] > -eps)
+            & (pos[:, 0] < view_size[0] + eps)
+            & (pos[:, 1] > -eps)
+            & (pos[:, 1] < view_size[1] + eps)
+        )
+        # pos = pos[in_view]
+        numbers = np.array(atoms.atomic_numbers)
+        # numbers = numbers[in_view]
+
+        atom_px = pos / px_scale  # AA / (AA/px) -> px
+
+        # atom_px = atom_px + margin
+
+        render = in_view
+        # render = (
+        #     (pos[:, 0] > 0)
+        #     & (pos[:, 0] < view_size[0])
+        #     & (pos[:, 1] > 0)
+        #     & (pos[:, 1] < view_size[1])
+        # )
+
+        numbers_render = numbers[render]
+        # # shift atomic positions to offset zero padding
+        atom_px_render = atom_px[render] + margin
+
+        # initialize arrays with zero padding
+        array = np.zeros((1,) + intensity.shape)  # adding extra 1
+        mask = np.zeros((1,) + intensity.shape)
+        # print(f"intensity: {array.shape}")
+        for number in np.unique(np.array(atoms.atomic_numbers)):
+
+            temp = np.zeros((1,) + intensity.shape)
+            temp = self.superpose_deltas(
+                atom_px_render[numbers_render == number], temp
+            )
             array += temp * number ** self.power_factor
+            temp = np.where(temp > 0, number, temp)
+            mask += temp[0]
 
+        # FFT convolution of beam profile and atom position delta functions
         array = np.fft.ifft2(np.fft.fft2(array) * np.fft.fft2(intensity)).real
-        array = array[0, margin:-margin, margin:-margin]
-        return array
+
+        # crop the FFT padding and fix atom coordinates relative to
+        # the image field of view
+        sel = slice(margin, -margin)
+        array = array[0, sel, sel]
+        mask = mask[0, sel, sel]
+        # atom_px = atom_px - margin
+
+        atom_px = pos[in_view] / px_scale
+        numbers = numbers[in_view]
+
+        return array, mask, atom_px, numbers
 
 
 """
@@ -119,10 +201,14 @@ def simulate_surface(self):
     plt.switch_backend("agg")
 
     a = Atoms.from_dict(get_jid_data("JVASP-667")["atoms"])
-    c = crop_square(a)
-    # c = a.make_supercell_matrix([2, 2, 1])
-    p = STEMConv(atoms=c).simulate_surface()
-    plt.imshow(p, interpolation="gaussian", cmap="plasma")
+    stem = STEMConv(output_size=(256, 256))
+    p, mask, atom_x, nb = stem.simulate_surface(
+        a, px_scale=0.15, eps=0.6, rot=2, shift=[-0.2, 0.3]
+    )
+    plt.imshow(p, origin="lower", cmap="plasma")
+    plt.scatter(atom_x[:, 1], atom_x[:, 0], marker="o",
+    facecolors="none", edgecolors="r"
+    )
     plt.savefig("stem_example.png")
     plt.close()
 """