atomic_dpc.py 16.5 KB
Newer Older
Mukherjee, Debangshu's avatar
Mukherjee, Debangshu committed
1
2
import scipy.ndimage as scnd
import scipy.optimize as sio
Mukherjee, Debangshu's avatar
Mukherjee, Debangshu committed
3
import numpy as np
4
import stemtool as st
5
6
7
8
9
10
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import matplotlib_scalebar.scalebar as mpss
import matplotlib.offsetbox as mploff
import matplotlib.gridspec as mpgs
import matplotlib as mpl
Mukherjee, Debangshu's avatar
Mukherjee, Debangshu committed
11

12

13
class atomic_dpc(object):
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
    """
    Atomic Resolution DPC estimation

    Parameters
    ----------
    Data_4D:  ndarray
              Four-dimensional dataset where the first two
              dimensions are real space scanning dimensions,
              while the last two dimenions are the Fourier 
              space electron diffraction patterns
    Data_ADF: ndarray
              Simultaneously collected two-dimensional ADF-STEM
              image
    calib_pm: float
              Real space pixel calibration in picometers
    voltage:  float
30
              Microscope accelerating voltage in kV
31
32
33
34
35
36
37
38
39
40
41
42
43
    aperture: float
              The probe forming condenser aperture in milliradians

    Notes
    -----
    This class function takes in a 4D-STEM image, and a simultaneously
    collected atomic resolution ADF-STEM image. Based on the accelerating
    voltage and the condenser aperture this calculates the center of mass 
    (C.O.M.) shifts in the central undiffracted beam. Using the idea that 
    the curl of the beam shift vectors, should be minimized at the correct
    Fourier rotation angles, this class also corrects for rotation of the 
    collceted 4D-STEM data with respect to the optic axis. Using these, a 
    correct potential accumulation and charge accumulation maps could be 
44
    built. To prevent errors, we convert everything to SI units first.
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
    
    Examples
    --------
    Run as:
    
    >>> DPC = st.dpc.atomic_dpc(Data_4D, DataADF, calibration, voltage, aper)
    
    Once the data is loaded, the ADF-STEM and the BF-STEM images could be 
    visualized as:
    
    >>> DPC.show_ADF_BF()
    
    Then the following call generates the mean CBED image, and if the show_image
    call is True, shows the mean image.

    >>> DPC.get_cbed(show_image = True)
    
    The initial uncorrected DPC shifts are generated as:

    >>> DPC.initial_dpc()

    The corrected DPC shifts are generated:

    >>> DPC.correct_dpc()

    The charge map is generated through:

    >>> DPC.show_charge()

    While the potential map is generated though:

    >>> DPC.show_potential()

    If a section of the image needs to be observed, to visualize the beam shifts,
    call the following:

    >>> DPC.plot_color_dpc()

    References
    ----------
    .. [1] Müller, K. et al. "Atomic electric fields revealed by a quantum mechanical 
86
        approach to electron picodiffraction". Nat. Commun. 5:565303 doi: 10.1038/ncomms6653 (2014)
87
88
    """

89
    def __init__(self, Data_4D, Data_ADF, calib_pm, voltage, aperture):
90
91
92
93
94
95
96
97
        """
        Load the user defined values.
        It also calculates the wavelength based on the accelerating voltage
        This also loads several SI constants as the foolowing attributes:
        `planck`:   The Planck's constant
        `epsilon0`: The dielectric permittivity of free space
        `e_charge`: The charge of an electron in Coulombs
        """
98
99
100
        self.data_adf = Data_ADF
        self.data_4D = Data_4D
        self.calib = calib_pm
101
102
103
104
105
106
107
108
        self.voltage = voltage * 1000  # convert to volts
        self.wavelength = st.sim.wavelength_ang(voltage) * (
            10 ** (-10)
        )  # convert to meters
        self.aperture = aperture / 1000  # convert to radians
        self.planck = 6.62607004 * (10 ** (-34))
        self.epsilon0 = 8.85418782 * (10 ** (-12))
        self.e_charge = (-1) * 1.60217662 * (10 ** (-19))
109

110
    def show_ADF_BF(self, imsize=(20, 10)):
111
112
113
114
115
116
        """
        The ADF-STEM image is already loaded, while the `data_bf`
        attribute is obtained by summing up the 4D-STEM dataset along it's 
        Fourier dimensions. This is also a great checkpoint to see whether
        the ADF-STEM and the BF-STEM images are the inverse of each other.
        """
117
        self.data_bf = np.sum(self.data_4D, axis=(-1, -2))
118
119
120
        fontsize = int(np.amax(np.asarray(imsize)))
        plt.figure(figsize=imsize)
        plt.subplot(1, 2, 1)
121
        plt.imshow(self.data_adf, cmap="inferno")
122
123
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
124
        scalebar.box_alpha = 0
125
        scalebar.color = "w"
126
        plt.gca().add_artist(scalebar)
127
128
129
130
        plt.axis("off")
        at = mploff.AnchoredText(
            "ADF-STEM", prop=dict(size=fontsize), frameon=True, loc="lower left"
        )
131
132
133
134
        at.patch.set_boxstyle("round, pad=0., rounding_size=0.2")
        plt.gca().add_artist(at)

        plt.subplot(1, 2, 2)
135
        plt.imshow(self.data_bf, cmap="inferno")
136
137
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
138
        scalebar.box_alpha = 0
139
        scalebar.color = "w"
140
        plt.gca().add_artist(scalebar)
141
142
143
144
        plt.axis("off")
        at = mploff.AnchoredText(
            "Summed 4D-STEM", prop=dict(size=fontsize), frameon=True, loc="lower left"
        )
145
146
147
        at.patch.set_boxstyle("round, pad=0., rounding_size=0.2")
        plt.gca().add_artist(at)
        plt.tight_layout()
148

149
    def get_cbed(self, imsize=(15, 15), show_image=False):
150
151
152
153
154
155
156
157
        """
        We calculate the mean CBED pattern by averaging the Fourier data, to
        get the object attribute `cbed`. We fit this with a circle function to
        obtain the object attributes `beam_x`, `beam_y` and `beam_r`, which
        are respectively the x coordinate, y coordinate and the radius of the 
        circle which best represents the CBED disk. We use the calculated radius
        and the known aperture size to get the Fourier space calibration.
        """
158
        self.cbed = np.mean(self.data_4D, axis=(0, 1))
159
        self.beam_x, self.beam_y, self.beam_r = st.util.sobel_circle(self.cbed)
160
        self.inverse = self.aperture / (self.beam_r * self.wavelength)
161
162
163
        if show_image:
            plt.figure(figsize=imsize)
            plt.imshow(self.cbed)
164
            scalebar = mpss.ScaleBar(self.inverse, "1/m", mpss.SI_LENGTH_RECIPROCAL)
165
166
167
168
169
            scalebar.location = "lower right"
            scalebar.box_alpha = 1
            scalebar.color = "k"
            plt.gca().add_artist(scalebar)
            plt.axis("off")
170
171

    def initial_dpc(self, imsize=(30, 15)):
172
173
174
        """
        
        """
175
176
        qq, pp = np.mgrid[0 : self.data_4D.shape[-1], 0 : self.data_4D.shape[-2]]
        yy, xx = np.mgrid[0 : self.data_4D.shape[0], 0 : self.data_4D.shape[1]]
177
178
179
180
181
182
        yy = np.ravel(yy)
        xx = np.ravel(xx)
        self.YCom = np.empty(self.data_4D.shape[0:2], dtype=np.float)
        self.XCom = np.empty(self.data_4D.shape[0:2], dtype=np.float)
        for ii in range(len(yy)):
            pattern = self.data_4D[yy[ii], xx[ii], :, :]
183
184
185
186
187
188
189
            self.YCom[yy[ii], xx[ii]] = self.inverse * (
                (np.sum(np.multiply(qq, pattern)) / np.sum(pattern)) - self.beam_y
            )
            self.XCom[yy[ii], xx[ii]] = self.inverse * (
                (np.sum(np.multiply(pp, pattern)) / np.sum(pattern)) - self.beam_x
            )

190
191
192
        vm = (np.amax(np.abs(np.concatenate((self.XCom, self.YCom), axis=1)))) / (
            10 ** 9
        )
193
        fontsize = int(np.amax(np.asarray(imsize)))
194
        sc_font = {"weight": "bold", "size": fontsize}
195

196
197
198
199
        fig = plt.figure(figsize=imsize)
        gs = mpgs.GridSpec(1, 2)
        ax1 = plt.subplot(gs[0, 0])
        ax2 = plt.subplot(gs[0, 1])
200

201
        im = ax1.imshow(self.XCom / (10 ** 9), vmin=-vm, vmax=vm, cmap="RdBu_r")
202
203
204
205
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
        scalebar.box_alpha = 1
        scalebar.color = "k"
206
        ax1.add_artist(scalebar)
207
208
209
210
211
212
213
        at = mploff.AnchoredText(
            "Shift in X direction",
            prop=dict(size=fontsize),
            frameon=True,
            loc="upper left",
        )
        at.patch.set_boxstyle("round, pad= 0., rounding_size= 0.2")
214
        ax1.add_artist(at)
215
        ax1.axis("off")
216

217
        im = ax2.imshow(self.YCom / (10 ** 9), vmin=-vm, vmax=vm, cmap="RdBu_r")
218
219
220
221
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
        scalebar.box_alpha = 1
        scalebar.color = "k"
222
        ax2.add_artist(scalebar)
223
224
225
226
227
228
229
        at = mploff.AnchoredText(
            "Shift in Y direction",
            prop=dict(size=fontsize),
            frameon=True,
            loc="upper left",
        )
        at.patch.set_boxstyle("round, pad= 0., rounding_size= 0.2")
230
        ax2.add_artist(at)
231
232
233
234
        ax2.axis("off")

        p1 = ax1.get_position().get_points().flatten()
        p2 = ax2.get_position().get_points().flatten()
235

236
237
        ax_cbar = fig.add_axes([p1[0] - 0.075, -0.01, p2[2], 0.02])
        cbar = plt.colorbar(im, cax=ax_cbar, orientation="horizontal")
238
        cbar.set_label(r"$\mathrm{Beam\: Shift\: \left(nm^{-1}\right)}$", **sc_font)
239

240
    def correct_dpc(self, imsize=(30, 15)):
241
242
243
244
245
246
        flips = np.zeros(4, dtype=bool)
        flips[2:4] = True
        chg_sums = np.zeros(4, dtype=self.XCom.dtype)
        angles = np.zeros(4, dtype=self.YCom.dtype)
        x0 = 90
        for ii in range(2):
247
            to_flip = flips[2 * ii]
248
249
250
251
252
253
254
255
256
            if to_flip:
                xdpcf = np.flip(self.XCom)
            else:
                xdpcf = self.XCom
            rho_dpc, phi_dpc = st.dpc.cart2pol(self.XCom, self.YCom)
            x = sio.minimize(st.dpc.angle_fun, x0, args=(rho_dpc, phi_dpc))
            min_x = x.x
            sol1 = min_x - 90
            sol2 = min_x + 90
257
258
259
260
261
262
263
264
265
266
267
            chg_sums[int(2 * ii)] = np.sum(
                st.dpc.charge_dpc(xdpcf, self.YCom, sol1) * self.data_adf
            )
            chg_sums[int(2 * ii + 1)] = np.sum(
                st.dpc.charge_dpc(xdpcf, self.YCom, sol2) * self.data_adf
            )
            angles[int(2 * ii)] = sol1
            angles[int(2 * ii + 1)] = sol2
        self.angle = (-1) * angles[chg_sums == np.amin(chg_sums)][0]
        self.final_flip = flips[chg_sums == np.amin(chg_sums)][0]

268
269
270
271
272
        if self.final_flip:
            xdpcf = np.fliplr(self.XCom)
        else:
            xdpcf = np.copy(self.XCom)
        rho_dpc, phi_dpc = st.dpc.cart2pol(xdpcf, self.YCom)
273
274
275
276
        self.XComC, self.YComC = st.dpc.pol2cart(
            rho_dpc, (phi_dpc - (self.angle * ((np.pi) / 180)))
        )

277
278
279
        vm = (np.amax(np.abs(np.concatenate((self.XComC, self.YComC), axis=1)))) / (
            10 ** 9
        )
280
        fontsize = int(np.amax(np.asarray(imsize)))
281
        sc_font = {"weight": "bold", "size": fontsize}
282

283
284
285
286
        fig = plt.figure(figsize=imsize)
        gs = mpgs.GridSpec(1, 2)
        ax1 = plt.subplot(gs[0, 0])
        ax2 = plt.subplot(gs[0, 1])
287

288
        im = ax1.imshow(self.XComC / (10 ** 9), vmin=-vm, vmax=vm, cmap="RdBu_r")
289
290
291
292
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
        scalebar.box_alpha = 1
        scalebar.color = "k"
293
        ax1.add_artist(scalebar)
294
295
296
297
298
299
300
        at = mploff.AnchoredText(
            "Corrected shift in X direction",
            prop=dict(size=fontsize),
            frameon=True,
            loc="upper left",
        )
        at.patch.set_boxstyle("round, pad= 0., rounding_size= 0.2")
301
        ax1.add_artist(at)
302
        ax1.axis("off")
303

304
        im = ax2.imshow(self.YComC / (10 ** 9), vmin=-vm, vmax=vm, cmap="RdBu_r")
305
306
307
308
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
        scalebar.box_alpha = 1
        scalebar.color = "k"
309
        ax2.add_artist(scalebar)
310
311
312
313
314
315
316
        at = mploff.AnchoredText(
            "Corrected shift in Y direction",
            prop=dict(size=fontsize),
            frameon=True,
            loc="upper left",
        )
        at.patch.set_boxstyle("round, pad= 0., rounding_size= 0.2")
317
        ax2.add_artist(at)
318
319
320
321
        ax2.axis("off")

        p1 = ax1.get_position().get_points().flatten()
        p2 = ax2.get_position().get_points().flatten()
322

323
324
        ax_cbar = fig.add_axes([p1[0] - 0.075, -0.01, p2[2], 0.02])
        cbar = plt.colorbar(im, cax=ax_cbar, orientation="horizontal")
325
326
327
328
329
330
331
        cbar.set_label(r"$\mathrm{Beam\: Shift\: \left(nm^{-1}\right)}$", **sc_font)

        self.MomentumX = self.planck * self.XComC
        self.MomentumY = self.planck * self.YComC
        # assuming infinitely thin sample
        self.e_fieldX = self.MomentumX / self.e_charge
        self.e_fieldY = self.MomentumY / self.e_charge
332

333
    def show_charge(self, imsize=(15, 15)):
334
        fontsize = int(np.amax(np.asarray(imsize)))
335
336
337
338
339
340
341
342
343
344

        # Use Poisson's equation
        self.charge = (
            (
                (np.gradient(self.e_fieldX)[1] + np.gradient(self.e_fieldY)[0])
                * (self.calib * (10 ** (-12)))
            )
            * self.epsilon0
            * 4
            * np.pi
345
        )
346
347
        cm = np.amax(np.abs(self.charge))
        plt.figure(figsize=imsize)
348
349
350
351
352
        plt.imshow(self.charge, vmin=-cm, vmax=cm, cmap="seismic")
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
        scalebar.box_alpha = 1
        scalebar.color = "k"
353
        plt.gca().add_artist(scalebar)
354
355
356
357
        plt.axis("off")
        at = mploff.AnchoredText(
            "Charge from DPC", prop=dict(size=fontsize), frameon=True, loc="lower left"
        )
358
359
360
        at.patch.set_boxstyle("round,pad=0.,rounding_size=0.2")
        plt.gca().add_artist(at)
        plt.tight_layout()
361
362

    def show_potential(self, imsize=(15, 15)):
363
        fontsize = int(np.amax(np.asarray(imsize)))
364
365
        self.potential = st.dpc.integrate_dpc(self.e_fieldX, self.e_fieldY)
        cm = np.amax(np.abs(self.potential))
366
        plt.figure(figsize=imsize)
367
        plt.imshow(self.potential, vmin=-cm, vmax=cm, cmap="RdBu_r")
368
369
370
371
        scalebar = mpss.ScaleBar(self.calib / 1000, "nm")
        scalebar.location = "lower right"
        scalebar.box_alpha = 1
        scalebar.color = "k"
372
        plt.gca().add_artist(scalebar)
373
374
375
376
377
378
379
        plt.axis("off")
        at = mploff.AnchoredText(
            "Measured potential",
            prop=dict(size=fontsize),
            frameon=True,
            loc="lower left",
        )
380
381
382
        at.patch.set_boxstyle("round, pad=0., rounding_size=0.2")
        plt.gca().add_artist(at)
        plt.tight_layout()
383
384

    def plot_color_dpc(self, skip=2, portion=7, imsize=(20, 10)):
385
        fontsize = int(np.amax(np.asarray(imsize)))
386
387
388
        sc_font = {"weight": "bold", "size": fontsize}
        mpl.rc("font", **sc_font)
        cc = self.XComC + ((1j) * self.YComC)
389
        cc_color = st.util.cp_image_val(cc)
390
391
392
393
394
395
396
        cutter = 1 / portion
        cutstart = (
            np.round(
                np.asarray(self.XComC.shape) - (cutter * np.asarray(self.XComC.shape))
            )
        ).astype(int)
        ypos, xpos = np.mgrid[0 : self.YComC.shape[0], 0 : self.XComC.shape[1]]
397
        ypos = ypos
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
        xcut = (
            xpos[cutstart[0] : self.XComC.shape[0], cutstart[1] : self.XComC.shape[1]]
            - cutstart[1]
        )
        ycut = (
            np.flipud(
                ypos[
                    cutstart[0] : self.XComC.shape[0], cutstart[1] : self.XComC.shape[1]
                ]
            )
            - cutstart[0]
        )
        dx = self.XComC[
            cutstart[0] : self.XComC.shape[0], cutstart[1] : self.XComC.shape[1]
        ]
        dy = self.YComC[
            cutstart[0] : self.XComC.shape[0], cutstart[1] : self.XComC.shape[1]
        ]
        cc_cut = cc_color[
            cutstart[0] : self.XComC.shape[0], cutstart[1] : self.XComC.shape[1], :
        ]
419
420
421
422

        plt.figure(figsize=imsize)
        plt.subplot(1, 2, 1)
        plt.imshow(cc_color)
423
424
        scalebar = mpss.ScaleBar(self.calib, "pm")
        scalebar.location = "lower right"
425
        scalebar.box_alpha = 0
426
        scalebar.color = "w"
427
        plt.gca().add_artist(scalebar)
428
429
430
431
432
433
434
        plt.axis("off")
        at = mploff.AnchoredText(
            "Center of Mass Shift",
            prop=dict(size=fontsize),
            frameon=True,
            loc="lower left",
        )
435
436
437
438
439
        at.patch.set_boxstyle("round, pad=0., rounding_size=0.2")
        plt.gca().add_artist(at)

        plt.subplot(1, 2, 2)
        plt.imshow(cc_cut)
440
441
442
443
444
445
446
447
448
449
        plt.quiver(
            xcut[::skip, ::skip],
            ycut[::skip, ::skip],
            dx[::skip, ::skip],
            dy[::skip, ::skip],
            pivot="mid",
            color="w",
        )
        scalebar = mpss.ScaleBar(self.calib, "pm")
        scalebar.location = "lower right"
450
        scalebar.box_alpha = 0
451
        scalebar.color = "w"
452
        plt.gca().add_artist(scalebar)
453
454
        plt.axis("off")
        plt.tight_layout()