From 99e0936de0f234d7cb3bee66f558507d809fbf61 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Sat, 24 Sep 2022 14:11:15 -0700
Subject: [PATCH 01/57] origin fitting plot bug fix

---
 py4DSTEM/io/datastructure/py4dstem/braggvectors_fns.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/py4DSTEM/io/datastructure/py4dstem/braggvectors_fns.py b/py4DSTEM/io/datastructure/py4dstem/braggvectors_fns.py
index 076976f0d..66b8b08da 100644
--- a/py4DSTEM/io/datastructure/py4dstem/braggvectors_fns.py
+++ b/py4DSTEM/io/datastructure/py4dstem/braggvectors_fns.py
@@ -181,7 +181,7 @@ def fit_origin(
                 'H':2,
                 'W':3,
                 'cmap':'RdBu',
-                'clipvals':'manual',
+                'intensity_range':'absolute',
                 'vmin':-1*plot_range,
                 'vmax':1*plot_range,
                 'axsize':(6,2),

From ce22f75a5b1dee02eeb9c3b0fc5486f4b9b0b48b Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Thu, 13 Apr 2023 05:09:19 -0700
Subject: [PATCH 02/57] Inital polar peaks class

---
 py4DSTEM/process/calibration/origin.py        |  29 +-
 py4DSTEM/process/wholepatternfit/__init__.py  |   1 +
 .../process/wholepatternfit/polar_peaks.py    | 331 ++++++++++++++++++
 3 files changed, 351 insertions(+), 10 deletions(-)
 create mode 100644 py4DSTEM/process/wholepatternfit/polar_peaks.py

diff --git a/py4DSTEM/process/calibration/origin.py b/py4DSTEM/process/calibration/origin.py
index 47ac61b2f..0389ba0df 100644
--- a/py4DSTEM/process/calibration/origin.py
+++ b/py4DSTEM/process/calibration/origin.py
@@ -110,10 +110,11 @@ def get_origin_single_dp(dp, r, rscale=1.2):
 
 def get_origin(
     datacube,
-    r=None,
-    rscale=1.2,
-    dp_max=None,
-    mask=None
+    r = None,
+    rscale = 1.2,
+    dp_max = None,
+    mask = None,
+    fast_center = False,
     ):
     """
     Find the origin for all diffraction patterns in a datacube, assuming (a) there is no
@@ -139,6 +140,8 @@ def get_origin(
         mask (ndarray or None): if not None, should be an (R_Nx,R_Ny) shaped
                     boolean array. Origin is found only where mask==True, and masked
                     arrays are returned for qx0,qy0
+        fast_center: (bool)
+            Skip the center of mass refinement step.
 
     Returns:
         (2-tuple of (R_Nx,R_Ny)-shaped ndarrays): the origin, (x,y) at each scan position
@@ -164,10 +167,13 @@ def get_origin(
         ):
             dp = datacube.data[rx, ry, :, :]
             _qx0, _qy0 = np.unravel_index(
-                np.argmax(gaussian_filter(dp, r)), (datacube.Q_Nx, datacube.Q_Ny)
+                np.argmax(gaussian_filter(dp, r, mode='nearest')), (datacube.Q_Nx, datacube.Q_Ny)
             )
-            _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
-            qx0[rx, ry], qy0[rx, ry] = get_CoM(dp * _mask)
+            if fast_center:
+                qx0[rx, ry], qy0[rx, ry] = _qx0, _qy0
+            else:
+                _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
+                qx0[rx, ry], qy0[rx, ry] = get_CoM(dp * _mask)
 
     else:
         assert mask.shape == (datacube.R_Nx, datacube.R_Ny)
@@ -188,10 +194,13 @@ def get_origin(
             if mask[rx, ry]:
                 dp = datacube.data[rx, ry, :, :]
                 _qx0, _qy0 = np.unravel_index(
-                    np.argmax(gaussian_filter(dp, r)), (datacube.Q_Nx, datacube.Q_Ny)
+                    np.argmax(gaussian_filter(dp, r, mode='nearest')), (datacube.Q_Nx, datacube.Q_Ny)
                 )
-                _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
-                qx0.data[rx, ry], qy0.data[rx, ry] = get_CoM(dp * _mask)
+                if fast_center:
+                    qx0[rx, ry], qy0[rx, ry] = _qx0, _qy0
+                else:
+                    _mask = np.hypot(qxx - _qx0, qyy - _qy0) < r * rscale
+                    qx0.data[rx, ry], qy0.data[rx, ry] = get_CoM(dp * _mask)
             else:
                 qx0.mask, qy0.mask = True, True
 
diff --git a/py4DSTEM/process/wholepatternfit/__init__.py b/py4DSTEM/process/wholepatternfit/__init__.py
index 8fb0e351e..97d2e083e 100644
--- a/py4DSTEM/process/wholepatternfit/__init__.py
+++ b/py4DSTEM/process/wholepatternfit/__init__.py
@@ -1,2 +1,3 @@
 from .wp_models import *
 from .wpf import *
+from .polar_peaks import *
diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
new file mode 100644
index 000000000..bb2ebbc9e
--- /dev/null
+++ b/py4DSTEM/process/wholepatternfit/polar_peaks.py
@@ -0,0 +1,331 @@
+"""
+This sub-module contains functions for polar transform peak detection of amorphous / semicrystalline datasets.
+
+"""
+
+from py4DSTEM import tqdmnd
+from itertools import product
+from typing import Optional
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+
+class PolarPeaks:
+    """
+    Primary class for polar transform peak detection.
+    """
+
+
+    def __init__(
+        self,
+        datacube,
+        radial_min = 0.0,
+        radial_max = None,
+        radial_step = 1.0,
+        num_annular_bins = 60,
+        progress_bar = True,
+        ):
+        """
+        Initialize class by performing an intensity-preserving polar transformation.
+
+        Parameters
+        --------
+        datacube: py4DSTEM.io.DataCube
+            4D-STEM dataset, requires origin calibration
+        radial_min: float
+            Minimum radius of polar transformation.
+        radial_max: float
+            Maximum radius of polar transformation.
+        radial_step: float
+            Width of radial bins of polar transformation.
+        num_annular_bins: int
+            Number of bins in annular direction. Note that we fold data over 180 degrees periodically, 
+            so setting this value to 60 gives bin widths of 180/60 = 3.0 degrees.
+        progress_bar: bool
+            Turns on the progress bar for the polar transformation
+    
+        Returns
+        --------
+        
+
+        """
+
+        # radial bin coordinates
+        if radial_max is None:
+            radial_max = np.min(datacube.Qshape) / np.sqrt(2)
+        self.radial_bins = np.arange(
+            radial_min,
+            radial_max,
+            radial_step,
+            )
+        self.radial_step = np.array(radial_step)
+
+        # annular bin coordinates
+        self.annular_bins = np.linspace(
+            0,
+            np.pi,
+            num_annular_bins,
+            endpoint = False,
+            )
+        self.annular_step = self.annular_bins[1] - self.annular_bins[0]
+
+        # init polar transformation array
+        self.polar_shape = np.array((self.annular_bins.shape[0], self.radial_bins.shape[0]))
+        self.polar_size = np.prod(self.polar_shape)
+        self.data_polar = np.zeros((
+            datacube.R_Nx,
+            datacube.R_Ny,
+            self.polar_shape[0],
+            self.polar_shape[1],
+            ))
+
+        # init coordinates
+        xa, ya = np.meshgrid(
+            np.arange(datacube.Q_Nx),
+            np.arange(datacube.Q_Ny),
+            indexing = 'ij',
+        )
+
+        # polar transformation
+        for rx, ry in tqdmnd(
+            range(20,21),
+            range(110,111),
+            # range(datacube.R_Nx),
+            # range(datacube.R_Ny),
+            desc="polar transformation",
+            unit=" images",
+            disable=not progress_bar,
+            ):
+
+            # shifted coordinates
+            x = xa - datacube.calibration.get_qx0(rx,ry)
+            y = ya - datacube.calibration.get_qy0(rx,ry)
+
+            # polar coordinate indices
+            r_ind = (np.sqrt(x**2 + y**2) - self.radial_bins[0]) / self.radial_step
+            t_ind = np.arctan2(y, x) / self.annular_step
+            r_ind_floor = np.floor(r_ind).astype('int')
+            t_ind_floor = np.floor(t_ind).astype('int')
+            dr = r_ind - r_ind_floor
+            dt = t_ind - t_ind_floor
+            # t_ind_floor = np.mod(t_ind_floor, self.num_annular_bins)
+
+            # polar transformation
+            sub = np.logical_and(r_ind_floor >= 0, r_ind_floor < self.polar_shape[1])
+            im = np.bincount(
+                r_ind_floor[sub] + \
+                np.mod(t_ind_floor[sub],self.polar_shape[0]) * self.polar_shape[1],
+                weights = datacube.data[rx,ry][sub] * (1 - dr[sub]) * (1 - dt[sub]),
+                minlength = self.polar_size,
+            )
+            im += np.bincount(
+                r_ind_floor[sub] + \
+                np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]) * self.polar_shape[1],
+                weights = datacube.data[rx,ry][sub] * (1 - dr[sub]) * (    dt[sub]),
+                minlength = self.polar_size,
+            )
+            sub = np.logical_and(r_ind_floor >= -1, r_ind_floor < self.polar_shape[1]-1)
+            im += np.bincount(
+                r_ind_floor[sub] + 1 + \
+                np.mod(t_ind_floor[sub],self.polar_shape[0]) * self.polar_shape[1],
+                weights = datacube.data[rx,ry][sub] * (    dr[sub]) * (1 - dt[sub]),
+                minlength = self.polar_size,
+            )
+            im += np.bincount(
+                r_ind_floor[sub] + 1 + \
+                np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]) * self.polar_shape[1],
+                weights = datacube.data[rx,ry][sub] * (    dr[sub]) * (    dt[sub]),
+                minlength = self.polar_size,
+            )
+
+            # sub = np.logical_and(r_ind_floor >= 0, r_ind_floor < self.polar_shape[1])
+            # im = accumarray(
+            #     np.vstack((
+            #         np.mod(t_ind_floor[sub],self.polar_shape[0]),
+            #         r_ind_floor[sub],
+            #     )).T, 
+            #     datacube.data[rx,ry][sub] * (1 - dr[sub]) * (1 - dt[sub]),
+            #     size=self.polar_shape,
+            # )
+            # im += accumarray(
+            #     np.vstack((
+            #         np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]),
+            #         r_ind_floor[sub],
+            #     )).T, 
+            #     datacube.data[rx,ry][sub] * (1 - dr[sub]) * (    dt[sub]),
+            #     size=self.polar_shape,
+            # )
+            # sub = np.logical_and(r_ind_floor >= -1, r_ind_floor < self.polar_shape[1]-1)
+            # im += accumarray(
+            #     np.vstack((
+            #         np.mod(t_ind_floor[sub],self.polar_shape[0]),
+            #         r_ind_floor[sub] + 1,
+            #     )).T, 
+            #     datacube.data[rx,ry][sub] * (    dr[sub]) * (1 - dt[sub]),
+            #     size=self.polar_shape,
+            # )            
+            # im += accumarray(
+            #     np.vstack((
+            #         np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]),
+            #         r_ind_floor[sub] + 1,
+            #     )).T, 
+            #     datacube.data[rx,ry][sub] * (    dr[sub]) * (    dt[sub]),
+            #     size=self.polar_shape,
+            # )
+
+            # Output
+            # self.data_polar[rx,ry] = im
+
+
+
+            fig, ax = plt.subplots(figsize=(6,6))
+            ax.imshow(
+                np.reshape(im, self.polar_shape),
+                vmin = 0,
+                vmax = 10,
+                )
+
+
+
+
+
+    # fig, ax = plt.subplots(figsize=(6,6))
+    # ax.imshow(
+    #     im,
+    #     vmin = 0,
+    #     vmax = 10,
+    #     )
+
+
+
+# def accumarray(subs, vals, size=None, fun=np.sum):
+
+#     if len(subs.shape) == 1:
+#         if size is None:
+#             size = [subs.values.max() + 1, 0]
+
+#         acc = val.groupby(subs).agg(fun)
+#     else:
+#         if size is None:
+#             size = [subs.values.max()+1, subs.shape[1]]
+
+#         # subs = subs.copy().reset_index()
+#         subs = subs.copy()
+#         by = subs.columns.tolist()[1:]
+#         acc = subs.groupby(by=by)['index'].agg(list).apply(lambda x: val[x].agg(fun))
+#         acc = acc.to_frame().reset_index().pivot_table(index=0, columns=1, aggfunc='first')
+#         acc.columns = range(acc.shape[1])
+#         acc = acc.reindex(range(size[1]), axis=1).fillna(0)
+
+#     id_x = range(size[0])
+#     acc = acc.reindex(id_x).fillna(0)
+
+#     return acc
+
+
+# def accum(accmap, a, func=None, size=None, fill_value=0, dtype=None):
+#     """
+#     An accumulation function similar to Matlab's `accumarray` function.
+
+#     Parameters
+#     ----------
+#     accmap : ndarray
+#         This is the "accumulation map".  It maps input (i.e. indices into
+#         `a`) to their destination in the output array.  The first `a.ndim`
+#         dimensions of `accmap` must be the same as `a.shape`.  That is,
+#         `accmap.shape[:a.ndim]` must equal `a.shape`.  For example, if `a`
+#         has shape (15,4), then `accmap.shape[:2]` must equal (15,4).  In this
+#         case `accmap[i,j]` gives the index into the output array where
+#         element (i,j) of `a` is to be accumulated.  If the output is, say,
+#         a 2D, then `accmap` must have shape (15,4,2).  The value in the
+#         last dimension give indices into the output array. If the output is
+#         1D, then the shape of `accmap` can be either (15,4) or (15,4,1) 
+#     a : ndarray
+#         The input data to be accumulated.
+#     func : callable or None
+#         The accumulation function.  The function will be passed a list
+#         of values from `a` to be accumulated.
+#         If None, numpy.sum is assumed.
+#     size : ndarray or None
+#         The size of the output array.  If None, the size will be determined
+#         from `accmap`.
+#     fill_value : scalar
+#         The default value for elements of the output array. 
+#     dtype : numpy data type, or None
+#         The data type of the output array.  If None, the data type of
+#         `a` is used.
+
+#     Returns
+#     -------
+#     out : ndarray
+#         The accumulated results.
+
+#         The shape of `out` is `size` if `size` is given.  Otherwise the
+#         shape is determined by the (lexicographically) largest indices of
+#         the output found in `accmap`.
+
+
+#     Examples
+#     --------
+#     >>> from numpy import array, prod
+#     >>> a = array([[1,2,3],[4,-1,6],[-1,8,9]])
+#     >>> a
+#     array([[ 1,  2,  3],
+#            [ 4, -1,  6],
+#            [-1,  8,  9]])
+#     >>> # Sum the diagonals.
+#     >>> accmap = array([[0,1,2],[2,0,1],[1,2,0]])
+#     >>> s = accum(accmap, a)
+#     array([9, 7, 15])
+#     >>> # A 2D output, from sub-arrays with shapes and positions like this:
+#     >>> # [ (2,2) (2,1)]
+#     >>> # [ (1,2) (1,1)]
+#     >>> accmap = array([
+#             [[0,0],[0,0],[0,1]],
+#             [[0,0],[0,0],[0,1]],
+#             [[1,0],[1,0],[1,1]],
+#         ])
+#     >>> # Accumulate using a product.
+#     >>> accum(accmap, a, func=prod, dtype=float)
+#     array([[ -8.,  18.],
+#            [ -8.,   9.]])
+#     >>> # Same accmap, but create an array of lists of values.
+#     >>> accum(accmap, a, func=lambda x: x, dtype='O')
+#     array([[[1, 2, 4, -1], [3, 6]],
+#            [[-1, 8], [9]]], dtype=object)
+#     """
+
+#     # Check for bad arguments and handle the defaults.
+#     if accmap.shape[:a.ndim] != a.shape:
+#         raise ValueError("The initial dimensions of accmap must be the same as a.shape")
+#     if func is None:
+#         func = np.sum
+#     if dtype is None:
+#         dtype = a.dtype
+#     if accmap.shape == a.shape:
+#         accmap = np.expand_dims(accmap, -1)
+#     adims = tuple(range(a.ndim))
+#     if size is None:
+#         size = 1 + np.squeeze(np.apply_over_axes(np.max, accmap, axes=adims))
+#     size = np.atleast_1d(size)
+
+#     # Create an array of python lists of values.
+#     vals = np.empty(size, dtype='O')
+#     for s in product(*[range(k) for k in size]):
+#         vals[s] = []
+#     for s in product(*[range(k) for k in a.shape]):
+#         indx = tuple(accmap[s])
+#         val = a[s]
+#         vals[indx].append(val)
+
+#     # Create the output array.
+#     out = np.empty(size, dtype=dtype)
+#     for s in product(*[range(k) for k in size]):
+#         if vals[s] == []:
+#             out[s] = fill_value
+#         else:
+#             out[s] = func(vals[s])
+
+#     return out

From 82b591dba41f1db40440a47d0720056f6c30f10f Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Thu, 13 Apr 2023 06:31:39 -0700
Subject: [PATCH 03/57] Adding simple peak fitting function

---
 .../process/wholepatternfit/polar_peaks.py    | 307 +++++++-----------
 1 file changed, 111 insertions(+), 196 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
index bb2ebbc9e..889e46eab 100644
--- a/py4DSTEM/process/wholepatternfit/polar_peaks.py
+++ b/py4DSTEM/process/wholepatternfit/polar_peaks.py
@@ -4,6 +4,7 @@
 """
 
 from py4DSTEM import tqdmnd
+from scipy.ndimage import gaussian_filter
 from itertools import product
 from typing import Optional
 
@@ -24,7 +25,7 @@ def __init__(
         radial_min = 0.0,
         radial_max = None,
         radial_step = 1.0,
-        num_annular_bins = 60,
+        num_annular_bins = 180,
         progress_bar = True,
         ):
         """
@@ -41,8 +42,9 @@ def __init__(
         radial_step: float
             Width of radial bins of polar transformation.
         num_annular_bins: int
-            Number of bins in annular direction. Note that we fold data over 180 degrees periodically, 
-            so setting this value to 60 gives bin widths of 180/60 = 3.0 degrees.
+            Number of bins in annular direction. Note that we fold data over 
+            180 degrees periodically, so setting this value to 60 gives bin 
+            widths of 180/60 = 3.0 degrees.
         progress_bar: bool
             Turns on the progress bar for the polar transformation
     
@@ -90,10 +92,8 @@ def __init__(
 
         # polar transformation
         for rx, ry in tqdmnd(
-            range(20,21),
-            range(110,111),
-            # range(datacube.R_Nx),
-            # range(datacube.R_Ny),
+            range(datacube.R_Nx),
+            range(datacube.R_Ny),
             desc="polar transformation",
             unit=" images",
             disable=not progress_bar,
@@ -140,192 +140,107 @@ def __init__(
                 minlength = self.polar_size,
             )
 
-            # sub = np.logical_and(r_ind_floor >= 0, r_ind_floor < self.polar_shape[1])
-            # im = accumarray(
-            #     np.vstack((
-            #         np.mod(t_ind_floor[sub],self.polar_shape[0]),
-            #         r_ind_floor[sub],
-            #     )).T, 
-            #     datacube.data[rx,ry][sub] * (1 - dr[sub]) * (1 - dt[sub]),
-            #     size=self.polar_shape,
-            # )
-            # im += accumarray(
-            #     np.vstack((
-            #         np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]),
-            #         r_ind_floor[sub],
-            #     )).T, 
-            #     datacube.data[rx,ry][sub] * (1 - dr[sub]) * (    dt[sub]),
-            #     size=self.polar_shape,
-            # )
-            # sub = np.logical_and(r_ind_floor >= -1, r_ind_floor < self.polar_shape[1]-1)
-            # im += accumarray(
-            #     np.vstack((
-            #         np.mod(t_ind_floor[sub],self.polar_shape[0]),
-            #         r_ind_floor[sub] + 1,
-            #     )).T, 
-            #     datacube.data[rx,ry][sub] * (    dr[sub]) * (1 - dt[sub]),
-            #     size=self.polar_shape,
-            # )            
-            # im += accumarray(
-            #     np.vstack((
-            #         np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]),
-            #         r_ind_floor[sub] + 1,
-            #     )).T, 
-            #     datacube.data[rx,ry][sub] * (    dr[sub]) * (    dt[sub]),
-            #     size=self.polar_shape,
-            # )
-
-            # Output
-            # self.data_polar[rx,ry] = im
-
-
-
-            fig, ax = plt.subplots(figsize=(6,6))
-            ax.imshow(
-                np.reshape(im, self.polar_shape),
-                vmin = 0,
-                vmax = 10,
-                )
-
-
-
-
-
-    # fig, ax = plt.subplots(figsize=(6,6))
-    # ax.imshow(
-    #     im,
-    #     vmin = 0,
-    #     vmax = 10,
-    #     )
-
-
-
-# def accumarray(subs, vals, size=None, fun=np.sum):
-
-#     if len(subs.shape) == 1:
-#         if size is None:
-#             size = [subs.values.max() + 1, 0]
-
-#         acc = val.groupby(subs).agg(fun)
-#     else:
-#         if size is None:
-#             size = [subs.values.max()+1, subs.shape[1]]
-
-#         # subs = subs.copy().reset_index()
-#         subs = subs.copy()
-#         by = subs.columns.tolist()[1:]
-#         acc = subs.groupby(by=by)['index'].agg(list).apply(lambda x: val[x].agg(fun))
-#         acc = acc.to_frame().reset_index().pivot_table(index=0, columns=1, aggfunc='first')
-#         acc.columns = range(acc.shape[1])
-#         acc = acc.reindex(range(size[1]), axis=1).fillna(0)
-
-#     id_x = range(size[0])
-#     acc = acc.reindex(id_x).fillna(0)
-
-#     return acc
-
-
-# def accum(accmap, a, func=None, size=None, fill_value=0, dtype=None):
-#     """
-#     An accumulation function similar to Matlab's `accumarray` function.
-
-#     Parameters
-#     ----------
-#     accmap : ndarray
-#         This is the "accumulation map".  It maps input (i.e. indices into
-#         `a`) to their destination in the output array.  The first `a.ndim`
-#         dimensions of `accmap` must be the same as `a.shape`.  That is,
-#         `accmap.shape[:a.ndim]` must equal `a.shape`.  For example, if `a`
-#         has shape (15,4), then `accmap.shape[:2]` must equal (15,4).  In this
-#         case `accmap[i,j]` gives the index into the output array where
-#         element (i,j) of `a` is to be accumulated.  If the output is, say,
-#         a 2D, then `accmap` must have shape (15,4,2).  The value in the
-#         last dimension give indices into the output array. If the output is
-#         1D, then the shape of `accmap` can be either (15,4) or (15,4,1) 
-#     a : ndarray
-#         The input data to be accumulated.
-#     func : callable or None
-#         The accumulation function.  The function will be passed a list
-#         of values from `a` to be accumulated.
-#         If None, numpy.sum is assumed.
-#     size : ndarray or None
-#         The size of the output array.  If None, the size will be determined
-#         from `accmap`.
-#     fill_value : scalar
-#         The default value for elements of the output array. 
-#     dtype : numpy data type, or None
-#         The data type of the output array.  If None, the data type of
-#         `a` is used.
-
-#     Returns
-#     -------
-#     out : ndarray
-#         The accumulated results.
-
-#         The shape of `out` is `size` if `size` is given.  Otherwise the
-#         shape is determined by the (lexicographically) largest indices of
-#         the output found in `accmap`.
-
-
-#     Examples
-#     --------
-#     >>> from numpy import array, prod
-#     >>> a = array([[1,2,3],[4,-1,6],[-1,8,9]])
-#     >>> a
-#     array([[ 1,  2,  3],
-#            [ 4, -1,  6],
-#            [-1,  8,  9]])
-#     >>> # Sum the diagonals.
-#     >>> accmap = array([[0,1,2],[2,0,1],[1,2,0]])
-#     >>> s = accum(accmap, a)
-#     array([9, 7, 15])
-#     >>> # A 2D output, from sub-arrays with shapes and positions like this:
-#     >>> # [ (2,2) (2,1)]
-#     >>> # [ (1,2) (1,1)]
-#     >>> accmap = array([
-#             [[0,0],[0,0],[0,1]],
-#             [[0,0],[0,0],[0,1]],
-#             [[1,0],[1,0],[1,1]],
-#         ])
-#     >>> # Accumulate using a product.
-#     >>> accum(accmap, a, func=prod, dtype=float)
-#     array([[ -8.,  18.],
-#            [ -8.,   9.]])
-#     >>> # Same accmap, but create an array of lists of values.
-#     >>> accum(accmap, a, func=lambda x: x, dtype='O')
-#     array([[[1, 2, 4, -1], [3, 6]],
-#            [[-1, 8], [9]]], dtype=object)
-#     """
-
-#     # Check for bad arguments and handle the defaults.
-#     if accmap.shape[:a.ndim] != a.shape:
-#         raise ValueError("The initial dimensions of accmap must be the same as a.shape")
-#     if func is None:
-#         func = np.sum
-#     if dtype is None:
-#         dtype = a.dtype
-#     if accmap.shape == a.shape:
-#         accmap = np.expand_dims(accmap, -1)
-#     adims = tuple(range(a.ndim))
-#     if size is None:
-#         size = 1 + np.squeeze(np.apply_over_axes(np.max, accmap, axes=adims))
-#     size = np.atleast_1d(size)
-
-#     # Create an array of python lists of values.
-#     vals = np.empty(size, dtype='O')
-#     for s in product(*[range(k) for k in size]):
-#         vals[s] = []
-#     for s in product(*[range(k) for k in a.shape]):
-#         indx = tuple(accmap[s])
-#         val = a[s]
-#         vals[indx].append(val)
-
-#     # Create the output array.
-#     out = np.empty(size, dtype=dtype)
-#     for s in product(*[range(k) for k in size]):
-#         if vals[s] == []:
-#             out[s] = fill_value
-#         else:
-#             out[s] = func(vals[s])
-
-#     return out
+            # output
+            self.data_polar[rx,ry] = np.reshape(im, self.polar_shape)
+
+
+
+
+    def fit_peaks(
+        self,
+        num_peaks_fit = 1,
+        sigma_radial_pixels = 0.0,
+        sigma_annular_degrees = 1.0,
+        progress_bar = True,
+        ):
+        """
+        Fit both background signal and peak positions and intensities for each radial bin.
+
+        Parameters
+        --------
+        progress_bar: bool
+            Turns on the progress bar for the polar transformation
+    
+        Returns
+        --------
+        
+
+        """
+
+        # sigma in pixels
+        self._sigma_radial_px = sigma_radial_pixels / self.radial_step
+        self._sigma_annular_px = np.deg2rad(sigma_annular_degrees) / self.annular_step
+
+        # init
+        self.radial_median = np.zeros((
+            self.data_polar.shape[0],
+            self.data_polar.shape[1],
+            self.polar_shape[1],
+            ))
+        self.radial_peaks = np.zeros((
+            self.data_polar.shape[0],
+            self.data_polar.shape[1],
+            self.polar_shape[1],
+            num_peaks_fit,
+            2,
+            ))
+
+
+        # loop over probe positions
+        for rx, ry in tqdmnd(
+            self.data_polar.shape[0],
+            self.data_polar.shape[1],
+            desc="polar transformation",
+            unit=" positions",
+            disable=not progress_bar,
+            ):
+
+            im = gaussian_filter(
+                self.data_polar[rx,ry],
+                sigma = (self._sigma_annular_px, self._sigma_radial_px),
+                mode = ('wrap', 'nearest'),
+                truncate = 3.0,
+            )
+
+            # background signal
+            self.radial_median[rx,ry] = np.median(im,axis=0)
+
+            # local maxima
+            sub_peaks = np.logical_and(
+                im > np.roll(im,-1,axis=0),
+                im > np.roll(im, 1,axis=0),
+            )
+
+            for a0 in range(self.polar_shape[1]):
+                inds = np.squeeze(np.argwhere(sub_peaks[:,a0]))
+                vals = im[inds,a0]
+
+                inds_sort = np.argsort(vals)[::-1]
+                inds_keep = inds_sort[:num_peaks_fit]
+
+                peaks_val = np.maximum(vals[inds_keep] - self.radial_median[rx,ry,a0], 0)
+                peaks_ind = inds[inds_keep]
+                peaks_angle = self.annular_bins[peaks_ind]
+
+                # TODO - add subpixel peak fitting?
+
+                # output
+                self.radial_peaks[rx,ry,a0,:,0] = peaks_angle
+                self.radial_peaks[rx,ry,a0,:,1] = peaks_val
+
+
+
+
+
+        # fig, ax = plt.subplots(figsize=(12,8))
+        # ax.imshow(
+        #     np.hstack((
+        #         im,
+        #         sub_peaks * im,
+        #     )),
+        #     vmin = 0,
+        #     vmax = 5,
+        #     )
+        # ax.plot(
+        #     im[:,6]
+        #     )

From 9764297e8c725e668dfd7145a306fb51c98389ad Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Thu, 13 Apr 2023 10:30:49 -0700
Subject: [PATCH 04/57] More plotting functions

---
 .../process/wholepatternfit/polar_peaks.py    | 138 ++++++++++++++++--
 1 file changed, 127 insertions(+), 11 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
index 889e46eab..295de7068 100644
--- a/py4DSTEM/process/wholepatternfit/polar_peaks.py
+++ b/py4DSTEM/process/wholepatternfit/polar_peaks.py
@@ -5,8 +5,11 @@
 
 from py4DSTEM import tqdmnd
 from scipy.ndimage import gaussian_filter
+from sklearn.decomposition import PCA
 from itertools import product
 from typing import Optional
+from matplotlib.colors import hsv_to_rgb
+
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -225,22 +228,135 @@ def fit_peaks(
                 # TODO - add subpixel peak fitting?
 
                 # output
-                self.radial_peaks[rx,ry,a0,:,0] = peaks_angle
-                self.radial_peaks[rx,ry,a0,:,1] = peaks_val
+                num_peaks = peaks_val.shape[0]
+                self.radial_peaks[rx,ry,a0,:num_peaks,0] = peaks_angle
+                self.radial_peaks[rx,ry,a0,:num_peaks,1] = peaks_val
 
 
 
 
 
-        # fig, ax = plt.subplots(figsize=(12,8))
-        # ax.imshow(
-        #     np.hstack((
-        #         im,
-        #         sub_peaks * im,
-        #     )),
-        #     vmin = 0,
-        #     vmax = 5,
-        #     )
+    def orientation_map(
+        self,
+        radial_index = 0,
+        peak_index = 0,
+        intensity_range = (0,1),
+        plot_result = True,
+        ):
+        """
+        Create an RGB orientation map from a given peak bin
+
+        Parameters
+        --------
+        progress_bar: bool
+            Turns on the progress bar for the polar transformation
+    
+        Returns
+        --------
+        im_orientation: np.array
+            rgb image array
+
+        """
+
+        # intensity mask
+        val = np.squeeze(self.radial_peaks[:,:,radial_index,peak_index,1]).copy()
+        val -= intensity_range[0]
+        val /= intensity_range[1] - intensity_range[0]
+        val = np.clip(val,0,1)
+
+        # orientation
+        hue = np.squeeze(self.radial_peaks[:,:,radial_index,peak_index,0]).copy()
+        hue = np.mod(2*hue,1)
+
+
+        # generate image
+        im_orientation = np.ones((
+            self.data_polar.shape[0],
+            self.data_polar.shape[1],
+            3))
+        im_orientation[:,:,0] = hue
+        im_orientation[:,:,2] = val
+        im_orientation = hsv_to_rgb(im_orientation)
+
+        if plot_result:
+            fig, ax = plt.subplots(figsize=(8,8))
+            ax.imshow(
+                im_orientation,
+                vmin = 0,
+                vmax = 5,
+            )
         # ax.plot(
         #     im[:,6]
         #     )
+
+        return im_orientation
+
+
+
+    def background_pca(
+        self,
+        pca_index = 0,
+        intensity_range = (0,1),
+        normalize_mean = True,
+        normalize_std = True,
+        plot_result = True,
+        plot_coef = False,
+        ):
+        """
+        Generate PCA decompositions of the background signal
+
+        Parameters
+        --------
+        progress_bar: bool
+            Turns on the progress bar for the polar transformation
+    
+        Returns
+        --------
+        im_pca: np,array
+            rgb image array
+        coef_pca: np.array
+            radial PCA component selected
+
+        """
+
+        # PCA decomposition
+        shape = self.radial_median.shape
+        A = np.reshape(self.radial_median, (shape[0]*shape[1],shape[2]))
+        if normalize_mean:
+            A -= np.mean(A,axis=0)
+        if normalize_std:
+            A /= np.std(A,axis=0)
+        pca = PCA(n_components=np.maximum(pca_index+1,2))
+        pca.fit(A)
+
+        components = pca.components_
+        loadings = pca.transform(A)
+
+        # output image data
+        sig_pca = np.reshape(loadings[:,pca_index], shape[0:2])
+        sig_pca -= intensity_range[0]
+        sig_pca /= intensity_range[1] - intensity_range[0]
+        sig_pca = np.clip(sig_pca,0,1)
+        im_pca = np.tile(sig_pca[:,:,None],(1,1,3))
+
+        # output PCA coefficient
+        coef_pca = np.vstack((
+            self.radial_bins,
+            components[pca_index,:]
+        )).T
+
+        if plot_result:
+            fig, ax = plt.subplots(figsize=(8,8))
+            ax.imshow(
+                im_pca,
+                vmin = 0,
+                vmax = 5,
+            )
+        if plot_coef:
+            fig, ax = plt.subplots(figsize=(8,4))
+            ax.plot(
+                coef_pca[:,0],
+                coef_pca[:,1]
+                )
+
+        return im_pca, coef_pca
\ No newline at end of file

From dd5c97a0683856a2ef73092f2b98c729a9619606 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Thu, 13 Apr 2023 13:00:01 -0700
Subject: [PATCH 05/57] fixing bug

---
 .../process/wholepatternfit/polar_peaks.py    | 25 ++++++++++---------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
index 295de7068..0a8c1026c 100644
--- a/py4DSTEM/process/wholepatternfit/polar_peaks.py
+++ b/py4DSTEM/process/wholepatternfit/polar_peaks.py
@@ -215,22 +215,23 @@ def fit_peaks(
             )
 
             for a0 in range(self.polar_shape[1]):
-                inds = np.squeeze(np.argwhere(sub_peaks[:,a0]))
-                vals = im[inds,a0]
+                inds = np.atleast_1d(np.squeeze(np.argwhere(sub_peaks[:,a0])))
+                if inds.size > 0:
+                    vals = im[inds,a0]
 
-                inds_sort = np.argsort(vals)[::-1]
-                inds_keep = inds_sort[:num_peaks_fit]
+                    inds_sort = np.argsort(vals)[::-1]
+                    inds_keep = inds_sort[:num_peaks_fit]
 
-                peaks_val = np.maximum(vals[inds_keep] - self.radial_median[rx,ry,a0], 0)
-                peaks_ind = inds[inds_keep]
-                peaks_angle = self.annular_bins[peaks_ind]
+                    peaks_val = np.maximum(vals[inds_keep] - self.radial_median[rx,ry,a0], 0)
+                    peaks_ind = inds[inds_keep]
+                    peaks_angle = self.annular_bins[peaks_ind]
 
-                # TODO - add subpixel peak fitting?
+                    # TODO - add subpixel peak fitting?
 
-                # output
-                num_peaks = peaks_val.shape[0]
-                self.radial_peaks[rx,ry,a0,:num_peaks,0] = peaks_angle
-                self.radial_peaks[rx,ry,a0,:num_peaks,1] = peaks_val
+                    # output
+                    num_peaks = peaks_val.shape[0]
+                    self.radial_peaks[rx,ry,a0,:num_peaks,0] = peaks_angle
+                    self.radial_peaks[rx,ry,a0,:num_peaks,1] = peaks_val
 
 
 

From 4f3390e908cf34413a0547fb2ad01669476b715c Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Sat, 15 Apr 2023 14:21:58 -0700
Subject: [PATCH 06/57] Adding grain clustering + plotting

---
 .../process/wholepatternfit/polar_peaks.py    | 288 +++++++++++++++++-
 1 file changed, 286 insertions(+), 2 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
index 0a8c1026c..9a11c7e37 100644
--- a/py4DSTEM/process/wholepatternfit/polar_peaks.py
+++ b/py4DSTEM/process/wholepatternfit/polar_peaks.py
@@ -10,7 +10,7 @@
 from typing import Optional
 from matplotlib.colors import hsv_to_rgb
 
-
+import time, sys
 import numpy as np
 import matplotlib.pyplot as plt
 
@@ -233,6 +233,79 @@ def fit_peaks(
                     self.radial_peaks[rx,ry,a0,:num_peaks,0] = peaks_angle
                     self.radial_peaks[rx,ry,a0,:num_peaks,1] = peaks_val
 
+    def plot_peak_signal(
+        self,
+        figsize = (8,4),
+        ):
+        """
+        Plot the mean peak signal from all radial bins.
+
+        """
+
+        sig = np.mean(
+            np.sum(
+                self.radial_peaks[:,:,:,:,1], axis = 3,
+                ), axis = (0,1),
+            )
+
+        # plot
+        fig = plt.figure(figsize = figsize)
+        ax1 = fig.add_subplot(111)
+        ax2 = ax1.twiny()
+
+        ax1.plot(
+            self.radial_bins,
+            sig,
+            )
+        ax1.set_xlabel('Radial Bin [pixels]')
+        ax1.set_ylabel('Mean Peak Signal [counts]')
+
+        def tick_function(x):
+            v = (x - self.radial_bins[0]) / self.radial_step
+            return ["%.0f" % z for z in v]
+        tick_locations = ax1.get_xticks()
+        tick_locations += self.radial_bins[0]
+        ax2.set_xticks(tick_locations)
+        ax2.set_xticklabels(tick_function(tick_locations))
+        ax2.set_xlim(ax1.get_xlim())
+        ax2.set_xlabel('Radial Bin Index')
+        ax2.minorticks_on()
+        ax2.grid()
+
+        plt.show()
+
+
+
+        # fig,ax = plt.subplots(figsize = figsize)
+        # ax.plot(
+        #     self.radial_bins,
+        #     sig,
+        #     )
+
+        
+
+
+        #         fig = plt.figure()
+        # ax1 = fig.add_subplot(111)
+        # ax2 = ax1.twiny()
+
+        # X = np.linspace(0,1,1000)
+        # Y = np.cos(X*20)
+
+        # ax1.plot(X,Y)
+        # ax1.set_xlabel(r"Original x-axis: $X$")
+
+        # new_tick_locations = np.array([.2, .5, .9])
+
+        # def tick_function(X):
+        #     V = 1/(1+X)
+        #     return ["%.3f" % z for z in V]
+
+        # ax2.set_xlim(ax1.get_xlim())
+        # ax2.set_xticks(new_tick_locations)
+        # ax2.set_xticklabels(tick_function(new_tick_locations))
+        # ax2.set_xlabel(r"Modified x-axis: $1/(1+X)$")
+        # plt.show()
 
 
 
@@ -293,6 +366,193 @@ def orientation_map(
         return im_orientation
 
 
+    def cluster_grains(
+        self,
+        radial_index = 0,
+        threshold_add = 1.0,
+        threshold_grow = 0.0,
+        angle_tolerance_deg = 5.0,
+        progress_bar = False,
+        ):
+        """
+        Cluster grains from a specific radial bin
+
+        Parameters
+        --------
+        radial_index: int
+            Which radial bin to perform the clustering over.
+        threshold_add: float
+            Minimum signal required for a probe position to initialize a cluster.
+        threshold_grow: float
+            Minimum signal required for a probe position to be added to a cluster.
+        angle_tolerance_deg: float
+            Rotation rolerance for clustering grains.
+        progress_bar: bool
+            Turns on the progress bar for the polar transformation
+    
+        Returns
+        --------
+        
+
+        """
+
+        # Get data
+        phi = np.squeeze(self.radial_peaks[:,:,radial_index,:,0]).copy()
+        sig = np.squeeze(self.radial_peaks[:,:,radial_index,:,1]).copy()
+        sig_init = sig.copy()
+        mark = sig >= threshold_grow
+        sig[np.logical_not(mark)] = 0
+        
+        # init
+        self.cluster_sizes = np.array((), dtype='int')
+        self.cluster_sig = np.array(())
+        self.cluster_inds = []
+        inds_all = np.zeros_like(sig, dtype='int')
+        inds_all.ravel()[:] = np.arange(inds_all.size)
+
+        # Tolerance
+        tol = np.deg2rad(angle_tolerance_deg)
+
+        # Main loop
+        search = True
+        while search is True:
+            inds_grain = np.argmax(sig)
+            val = sig.ravel()[inds_grain]
+
+            if val < threshold_add:
+                search = False
+            
+            else:
+                # progressbar
+                comp = 1 - np.mean(np.max(mark,axis = 2))
+                update_progress(comp)
+
+                # Start cluster
+                x,y,z = np.unravel_index(inds_grain, sig.shape)
+                mark[x,y,z] = False
+                sig[x,y,z] = 0
+                phi_cluster = phi[x,y,z]
+
+                # Neighbors to search
+                xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1)
+                yr = np.clip(y + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1)
+                inds_cand = inds_all[xr,yr,:].ravel()
+                inds_cand = np.delete(inds_cand, mark.ravel()[inds_cand] == False)
+                # [mark[xr,yr,:].ravel()]
+
+                if inds_cand.size == 0:
+                    grow = False
+                else:
+                    grow = True
+
+                # grow the cluster
+                while grow is True:
+                    inds_new = np.array((),dtype='int')
+
+                    keep = np.zeros(inds_cand.size, dtype='bool')
+                    for a0 in range(inds_cand.size):
+                        xc,yc,zc = np.unravel_index(inds_cand[a0], sig.shape)
+
+                        phi_test = phi[xc,yc,zc]
+                        dphi = np.mod(phi_cluster - phi_test + np.pi/2.0, np.pi) - np.pi/2.0
+
+                        if np.abs(dphi) < tol:
+                            keep[a0] = True
+
+                            sig[xc,yc,zc] = 0
+                            mark[xc,yc,zc] = False 
+
+                            xr = np.clip(xc + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1)
+                            yr = np.clip(yc + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1)
+                            inds_add = inds_all[xr,yr,:].ravel()
+                            inds_new = np.append(inds_new, inds_add)
+
+
+                    inds_grain = np.append(inds_grain, inds_cand[keep])
+                    inds_cand = np.unique(np.delete(inds_new, mark.ravel()[inds_new] == False))
+
+                    if inds_cand.size == 0:
+                        grow = False
+
+            # convert grain to x,y coordinates, add = list
+            xg,yg,zg = np.unravel_index(inds_grain, sig.shape)
+            xyg = np.unique(np.vstack((xg,yg)), axis = 1)
+            sig_mean = np.mean(sig_init.ravel()[inds_grain])
+            self.cluster_sizes = np.append(self.cluster_sizes, xyg.shape[1])
+            self.cluster_sig = np.append(self.cluster_sig, sig_mean)
+            self.cluster_inds.append(xyg)
+
+        # finish progressbar
+        update_progress(1)
+
+    def cluster_plot_size(
+        self,
+        area_max = None,
+        weight_intensity = False,
+        pixel_area = 1.0,
+        pixel_area_units = 'px^2',
+        figsize = (8,6),
+        returnfig = False,
+        ):
+        """
+        Plot the cluster sizes
+
+        Parameters
+        --------
+        area_max: int (optional)
+            Max area bin in pixels
+        weight_intensity: bool
+            Weight histogram by the peak intensity.
+        pixel_area: float
+            Size of pixel area unit square
+        pixel_area_units: string
+            Units of the pixel area 
+        figsize: tuple
+            Size of the figure panel
+        returnfig: bool
+            Setting this to true returns the figure and axis handles
+    
+        Returns
+        --------
+        fig, ax (optional)
+            Figure and axes handles
+
+        """
+
+        if area_max is None:
+            area_max = np.max(self.cluster_sizes)
+        area = np.arange(area_max)
+        sub = self.cluster_sizes.astype('int') < area_max
+        if weight_intensity:
+            hist = np.bincount(
+                self.cluster_sizes[sub],
+                weights = self.cluster_sig[sub],
+                minlength = area_max,
+                )
+        else:
+            hist = np.bincount(
+                self.cluster_sizes[sub],
+                minlength = area_max,
+                )
+
+
+        # plotting
+        fig,ax = plt.subplots(figsize = figsize)
+        ax.bar(
+            area * pixel_area,
+            hist,
+            width = 0.8 * pixel_area,
+            )
+
+        ax.set_xlabel('Grain Area [' + pixel_area_units + ']')
+        if weight_intensity:
+            ax.set_ylabel('Total Signal [arb. units]')
+        else:
+            ax.set_ylabel('Number of Grains')
+
+        if returnfig:
+            return fig,ax
+
 
     def background_pca(
         self,
@@ -360,4 +620,28 @@ def background_pca(
                 coef_pca[:,1]
                 )
 
-        return im_pca, coef_pca
\ No newline at end of file
+        return im_pca, coef_pca
+
+
+# Progressbar taken from stackexchange:
+# https://stackoverflow.com/questions/3160699/python-progress-bar
+def update_progress(progress):
+    barLength = 60 # Modify this to change the length of the progress bar
+    status = ""
+    if isinstance(progress, int):
+        progress = float(progress)
+    if not isinstance(progress, float):
+        progress = 0
+        status = "error: progress var must be float\r\n"
+    if progress < 0:
+        progress = 0
+        status = "Halt...\r\n"
+    if progress >= 1:
+        progress = 1
+        status = "Done...\r\n"
+    block = int(round(barLength*progress))
+    text = "\rPercent: [{0}] {1}% {2}".format( "#"*block + "-"*(barLength-block), 
+        np.round(progress*100,4), 
+        status)
+    sys.stdout.write(text)
+    sys.stdout.flush()

From 81ccd097fb17e5e42b5052b9e4fbbbb543d19a68 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 12 May 2023 18:42:13 -0700
Subject: [PATCH 07/57] Fixing clustering

---
 py4DSTEM/process/wholepatternfit/polar_peaks.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
index 9a11c7e37..e65c81532 100644
--- a/py4DSTEM/process/wholepatternfit/polar_peaks.py
+++ b/py4DSTEM/process/wholepatternfit/polar_peaks.py
@@ -436,9 +436,8 @@ def cluster_grains(
                 # Neighbors to search
                 xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1)
                 yr = np.clip(y + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1)
-                inds_cand = inds_all[xr,yr,:].ravel()
+                inds_cand = inds_all[xr[:,None],yr[None],:].ravel()
                 inds_cand = np.delete(inds_cand, mark.ravel()[inds_cand] == False)
-                # [mark[xr,yr,:].ravel()]
 
                 if inds_cand.size == 0:
                     grow = False
@@ -464,7 +463,7 @@ def cluster_grains(
 
                             xr = np.clip(xc + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1)
                             yr = np.clip(yc + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1)
-                            inds_add = inds_all[xr,yr,:].ravel()
+                            inds_add = inds_all[xr[:,None],yr[None],:].ravel()
                             inds_new = np.append(inds_new, inds_add)
 
 

From 10d47f85596e32cc299d8d4050479f440d754d20 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Sat, 13 May 2023 10:21:41 -0700
Subject: [PATCH 08/57] fixing major bug in clustering

---
 .../process/wholepatternfit/polar_peaks.py    | 28 +++++++++++++------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
index e65c81532..fad76c4a6 100644
--- a/py4DSTEM/process/wholepatternfit/polar_peaks.py
+++ b/py4DSTEM/process/wholepatternfit/polar_peaks.py
@@ -486,7 +486,9 @@ def cluster_grains(
 
     def cluster_plot_size(
         self,
+        area_min = None,
         area_max = None,
+        area_step = 1,
         weight_intensity = False,
         pixel_area = 1.0,
         pixel_area_units = 'px^2',
@@ -498,8 +500,12 @@ def cluster_plot_size(
 
         Parameters
         --------
+        area_min: int (optional)
+            Min area bin in pixels
         area_max: int (optional)
             Max area bin in pixels
+        area_step: int (optional)
+            Step size of the histogram bin
         weight_intensity: bool
             Weight histogram by the peak intensity.
         pixel_area: float
@@ -520,18 +526,24 @@ def cluster_plot_size(
 
         if area_max is None:
             area_max = np.max(self.cluster_sizes)
-        area = np.arange(area_max)
-        sub = self.cluster_sizes.astype('int') < area_max
+        area = np.arange(0,area_max,area_step)
+        if area_min is None:
+            sub = self.cluster_sizes.astype('int') < area_max
+        else:
+            sub = np.logical_and(
+                self.cluster_sizes.astype('int') >= area_min,
+                self.cluster_sizes.astype('int') < area_max
+                )
         if weight_intensity:
             hist = np.bincount(
-                self.cluster_sizes[sub],
+                self.cluster_sizes[sub] // area_step,
                 weights = self.cluster_sig[sub],
-                minlength = area_max,
+                minlength = area.shape[0],
                 )
         else:
             hist = np.bincount(
-                self.cluster_sizes[sub],
-                minlength = area_max,
+                self.cluster_sizes[sub] // area_step,
+                minlength = area.shape[0],
                 )
 
 
@@ -540,9 +552,9 @@ def cluster_plot_size(
         ax.bar(
             area * pixel_area,
             hist,
-            width = 0.8 * pixel_area,
+            width = 0.8 * pixel_area * area_step,
             )
-
+        ax.set_xlim((0,area_max*pixel_area))
         ax.set_xlabel('Grain Area [' + pixel_area_units + ']')
         if weight_intensity:
             ax.set_ylabel('Total Signal [arb. units]')

From 78fc2a696ea2cf00c5caef6580ccf5f5b32068fb Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Thu, 8 Jun 2023 09:02:04 -0700
Subject: [PATCH 09/57] Fixing add beamstop bug

---
 py4DSTEM/classes/methods/datacube_methods.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/py4DSTEM/classes/methods/datacube_methods.py b/py4DSTEM/classes/methods/datacube_methods.py
index bf3a4de4d..f6269c555 100644
--- a/py4DSTEM/classes/methods/datacube_methods.py
+++ b/py4DSTEM/classes/methods/datacube_methods.py
@@ -1266,16 +1266,16 @@ def get_beamstop_mask(
             name = 'gen_params',
             data = {
                 #'gen_func' : 
-                'threshold' : 0.25,
-                'distance_edge' : 4.0,
-                'include_edges' : True,
+                'threshold' : threshold,
+                'distance_edge' : distance_edge,
+                'include_edges' : include_edges,
                 'name' : "mask_beamstop",
-                'returncalc' : True,
+                'returncalc' : returncalc,
             }
         )
 
         # Add to tree
-        self.tree( mask_beamstop )
+        self.tree(x)
 
         # return
         if returncalc:

From ff38927caf2510ac9a214520faf46c99dd12808d Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Sun, 25 Jun 2023 14:21:29 -0400
Subject: [PATCH 10/57] starting to switch to externally managed offset for WPF
 parameters

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 45 ++++++++++---------
 1 file changed, 25 insertions(+), 20 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index b19ffc902..221334a81 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -74,6 +74,11 @@ def __init__(
         else:
             self.set_params(initial_value, lower_bound, upper_bound)
 
+        # Store a dummy offset. This must be set by WPF during setup
+        # This stores the index in the master Jacobian array corresponding to
+        # this parameter
+        self.offset = np.nan
+
     def set_params(
         self, 
         initial_value, 
@@ -100,8 +105,8 @@ def __init__(self, background_value=0.0, name="DC Background"):
     def func(self, DP: np.ndarray, level, **kwargs) -> None:
         DP += level
 
-    def jacobian(self, J: np.ndarray, *args, offset: int, **kwargs):
-        J[:, offset] = 1
+    def jacobian(self, J: np.ndarray, *args, **kwargs):
+        J[:, self.params['DC Level'].offset] = 1
 
 
 class GaussianBackground(WPFModelPrototype):
@@ -130,7 +135,7 @@ def global_center_func(self, DP: np.ndarray, sigma, level, **kwargs) -> None:
         DP += level * np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
 
     def global_center_jacobian(
-        self, J: np.ndarray, sigma, level, offset: int, **kwargs
+        self, J: np.ndarray, sigma, level, **kwargs
     ) -> None:
 
         exp_expr = np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
@@ -146,10 +151,10 @@ def global_center_jacobian(
         ).ravel()
 
         # dF/s(sigma)
-        J[:, offset] = (level * kwargs["global_r"] ** 2 * exp_expr / sigma**3).ravel()
+        J[:, self.params['sigma'].offset] = (level * kwargs["global_r"] ** 2 * exp_expr / sigma**3).ravel()
 
         # dF/d(level)
-        J[:, offset + 1] = exp_expr.ravel()
+        J[:, self.params['intensity'].offset] = exp_expr.ravel()
 
     def local_center_func(self, DP: np.ndarray, sigma, level, x0, y0, **kwargs) -> None:
         DP += level * np.exp(
@@ -158,11 +163,11 @@ def local_center_func(self, DP: np.ndarray, sigma, level, x0, y0, **kwargs) -> N
         )
 
     def local_center_jacobian(
-        self, J: np.ndarray, sigma, level, x0, y0, offset: int, **kwargs
+        self, J: np.ndarray, sigma, level, x0, y0, **kwargs
     ) -> None:
 
         # dF/s(sigma)
-        J[:, offset] = (
+        J[:, self.params['sigma'].offset] = (
             level
             * ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             * np.exp(
@@ -173,13 +178,13 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(level)
-        J[:, offset + 1] = np.exp(
+        J[:, self.params['intensity'].offset] = np.exp(
             ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             / (-2 * sigma**2)
         ).ravel()
 
         # dF/d(x0)
-        J[:, offset + 2] = (
+        J[:, self.params['x center'].offset] = (
             level
             * (kwargs["xArray"] - x0)
             * np.exp(
@@ -190,7 +195,7 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(y0)
-        J[:, offset + 3] = (
+        J[:, self.params['y center'].offset] = (
             level
             * (kwargs["yArray"] - y0)
             * np.exp(
@@ -234,7 +239,7 @@ def global_center_func(
         DP += level * np.exp((kwargs["global_r"] - radius) ** 2 / (-2 * sigma**2))
 
     def global_center_jacobian(
-        self, J: np.ndarray, radius, sigma, level, offset: int, **kwargs
+        self, J: np.ndarray, radius, sigma, level, **kwargs
     ) -> None:
 
         local_r = radius - kwargs["global_r"]
@@ -261,15 +266,15 @@ def global_center_jacobian(
         ).ravel()
 
         # dF/d(radius)
-        J[:, offset] += (-1.0 * level * exp_expr * local_r / (sigma**2)).ravel()
+        J[:, self.params['radius'].offset] += (-1.0 * level * exp_expr * local_r / (sigma**2)).ravel()
 
         # dF/d(sigma)
-        J[:, offset + 1] = (
+        J[:, self.params['sigma'].offset] = (
             level * local_r ** 2 * exp_expr / sigma**3
         ).ravel()
 
-        # dF/d(level)
-        J[:, offset + 2] = exp_expr.ravel()
+        # dF/d(intensity)
+        J[:, self.params['intensity'].offset] = exp_expr.ravel()
 
     def local_center_func(
         self, DP: np.ndarray, radius, sigma, level, x0, y0, **kwargs
@@ -278,13 +283,13 @@ def local_center_func(
         DP += level * np.exp((local_r - radius) ** 2 / (-2 * sigma**2))
 
     def local_center_jacobian(
-        self, J: np.ndarray, radius, sigma, level, x0, y0, offset: int, **kwargs
+        self, J: np.ndarray, radius, sigma, level, x0, y0, **kwargs
     ) -> None:
         return NotImplementedError()
         # dF/d(radius)
 
         # dF/s(sigma)
-        J[:, offset] = (
+        J[:, self.params['sigma'].offset] = (
             level
             * ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             * np.exp(
@@ -295,13 +300,13 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(level)
-        J[:, offset + 1] = np.exp(
+        J[:, self.params['intensity'].offset] = np.exp(
             ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             / (-2 * sigma**2)
         ).ravel()
 
         # dF/d(x0)
-        J[:, offset + 2] = (
+        J[:, self.params['x center'].offset] = (
             level
             * (kwargs["xArray"] - x0)
             * np.exp(
@@ -312,7 +317,7 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(y0)
-        J[:, offset + 3] = (
+        J[:, self.params['y center'].offset] = (
             level
             * (kwargs["yArray"] - y0)
             * np.exp(

From ca785d6e145a32024df4f1a309e9c56fca9890a9 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 28 Jul 2023 10:28:07 -0700
Subject: [PATCH 11/57] Visualization tweaks

---
 py4DSTEM/visualize/overlay.py | 38 ++++++++++++++++++++++++++---------
 py4DSTEM/visualize/show.py    | 14 +++++++++++++
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/py4DSTEM/visualize/overlay.py b/py4DSTEM/visualize/overlay.py
index 7e7147a15..250d63d77 100644
--- a/py4DSTEM/visualize/overlay.py
+++ b/py4DSTEM/visualize/overlay.py
@@ -147,21 +147,32 @@ def add_annuli(ax,d):
     """
     Adds one or more annuli to Axis ax using the parameters in dictionary d.
     """
-    # Handle inputs
+    
+    # Check that all required inputs are present
     assert isinstance(ax,Axes)
-    # center
     assert('center' in d.keys())
+    assert('radii' in d.keys())
+
+    # Get user-provided center and radii
     center = d['center']
+    radii = d['radii']
+
+    # Determine number of annuli being plotted
+    if isinstance(center,list):
+        N = len(center)
+    elif isinstance(radii,list):
+        N = len(radii)
+    else:
+        N = 1
+
+    # center
     if isinstance(center,tuple):
         assert(len(center)==2)
-        center = [center]
-    assert(isinstance(center,list))
-    N = len(center)
+        center = [center]*N
+    # assert(isinstance(center,list))
     assert(all([isinstance(x,tuple) for x in center]))
     assert(all([len(x)==2 for x in center]))
     # radii
-    assert('radii' in d.keys())
-    radii = d['radii']
     if isinstance(radii,tuple):
         assert(len(radii)==2)
         ri = [radii[0] for i in range(N)]
@@ -183,7 +194,7 @@ def add_annuli(ax,d):
         assert is_color_like(color)
         color = [color for i in range(N)]
     # fill
-    fill = d['fill'] if 'fill' in d.keys() else False
+    fill = d['fill'] if 'fill' in d.keys() else True
     if isinstance(fill,bool):
         fill = [fill for i in range(N)]
     else:
@@ -694,8 +705,15 @@ def add_scalebar(ax,d):
         labeltext = f'{np.round(length_units,3)}'+' '+pixelunits
         if xshiftdir>0: va='top'
         else: va='bottom'
-        ax.text(labelpos_y,labelpos_x,labeltext,size=labelsize,
-                color=labelcolor,alpha=alpha,ha='center',va=va)
+        ax.text(
+            labelpos_y,
+            labelpos_x,
+            labeltext,
+            size = labelsize,
+            color = labelcolor,
+            alpha = alpha,
+            ha = 'center',
+            va = va)
 
     # if not ticks:
     #     ax.set_xticks([])
diff --git a/py4DSTEM/visualize/show.py b/py4DSTEM/visualize/show.py
index 3b9d99e43..2bbe563ac 100644
--- a/py4DSTEM/visualize/show.py
+++ b/py4DSTEM/visualize/show.py
@@ -679,6 +679,20 @@ def show(
             scalebar['pixelsize'] = pixelsize
             scalebar['pixelunits'] = pixelunits
             scalebar['space'] = space
+            # determine good default scale bar fontsize
+            if figax is not None:
+                # print(figax[0].figsize)
+                # size = figax[1].get_size_inches()
+                # ax_h = figax[1].bbox.transformed(figax[0].gca().transAxes).height
+                bbox = figax[1].get_window_extent()
+                dpi = figax[0].dpi
+                size = (
+                    bbox.width / dpi , 
+                    bbox.height / dpi 
+                    )
+                scalebar['labelsize'] = np.min(np.array(size)) * 3.0
+            if 'labelsize' not in scalebar.keys():
+                scalebar['labelsize'] = np.min(np.array(figsize)) * 2.0
             add_scalebar(ax,scalebar)
 
 

From 81c708735497284847605f44cf15cef035b16f49 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 28 Jul 2023 14:44:57 -0700
Subject: [PATCH 12/57] Updating viz to have white/black text

---
 py4DSTEM/process/wholepatternfit/wpf_viz.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index da972ca33..2820d1056 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -38,7 +38,26 @@ def show_model_grid(self, x=None, **plot_kwargs):
         m.func(DP, *x[ind : ind + m.nParams].tolist(), **shared_data)
 
         a.matshow(DP, cmap="turbo")
-        a.text(0.5, 0.92, m.name, transform=a.transAxes, ha="center", va="center")
+
+        # Determine if text color should be white or black
+        int_range = np.array((np.min(DP), np.max(DP)))
+        if int_range[0] != int_range[1]:
+            r = (np.mean(DP[:DP.shape[0]//10,:]) - int_range[0]) / (int_range[1] - int_range[0])
+            if r < 0.5:
+                color = 'w'
+            else:
+                color = 'k'
+        else:
+            color = 'w'
+
+        a.text(
+            0.5, 
+            0.92, 
+            m.name, 
+            transform = a.transAxes, 
+            ha = "center", 
+            va = "center",
+            color = color)
     for a in ax.flat:
         a.axis("off")
 

From 0a0f5eb86971bc60ddf73238a97947671ed7dcb7 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 28 Jul 2023 14:51:25 -0700
Subject: [PATCH 13/57] Updates

---
 py4DSTEM/process/wholepatternfit/wpf.py | 235 +++++++++++++++++++-----
 1 file changed, 191 insertions(+), 44 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 542fb1b14..54391b494 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -11,6 +11,8 @@
 from matplotlib.gridspec import GridSpec
 import warnings
 
+from multiprocessing import Pool
+
 class WholePatternFit:
 
     from py4DSTEM.process.wholepatternfit.wpf_viz import (
@@ -147,6 +149,7 @@ def fit_to_mean_CBED(self, **fit_opts):
         default_opts = {
             "method": "trf",
             "verbose": 1,
+            "x_scale": "jac",
         }
         default_opts.update(fit_opts)
 
@@ -208,9 +211,108 @@ def fit_to_mean_CBED(self, **fit_opts):
 
         return opt
 
+
+    def fit_single_pattern(
+        self, 
+        rx,
+        ry,
+        resume = False, 
+        **fit_opts
+        ):
+        """
+        Apply model fitting to one pattern.
+
+        Parameters
+        ----------
+        resume: bool (optional)
+            Set to true to continue a previous fit with more iterations.
+        rx: int
+            probe x coordinate
+        ry: int
+            probe y coordinate
+        fit_opts: args (optional)
+            args passed to scipy.optimize.least_squares
+
+        Returns
+        --------
+        fit_coefs: np.array
+            Fitted coefficients 
+        fit_metrics: np.array
+            Fitting metrics
+
+        """
+
+        # make sure we have the latest parameters
+        self._scrape_model_params()
+
+        # set tracking off
+        self._track = False
+        self._fevals = []
+
+        if resume:
+            assert hasattr(self, "fit_data"), "No existing data resuming fit!"
+
+        # init
+        fit_coefs = np.zeros(self.x0.shape[0])
+        fit_metric = np.zeros(4)
+
+        # Default fitting options
+        default_opts = {
+            "method": "trf",
+            "verbose": 1,
+            "x_scale": "jac",
+        }
+        default_opts.update(fit_opts)
+
+        # Loop over probe positions
+        current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
+        shared_data = self.static_data.copy()
+        self._cost_history = (
+            []
+        )  # clear this so it doesn't grow: TODO make this not stupid
+
+        try:
+            x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0
+
+            if self.hasJacobian & self.use_jacobian:
+                opt = least_squares(
+                    self._pattern_error,
+                    x0,
+                    jac=self._jacobian,
+                    bounds=(self.lower_bound, self.upper_bound),
+                    args=(current_pattern, shared_data),
+                    **default_opts,
+                    # **fit_opts,
+                )
+            else:
+                opt = least_squares(
+                    self._pattern_error,
+                    x0,
+                    bounds=(self.lower_bound, self.upper_bound),
+                    args=(current_pattern, shared_data),
+                    **default_opts,
+                    # **fit_opts,
+                )
+
+            fit_coefs = opt.x
+            fit_metrics_single = [
+                opt.cost,
+                opt.optimality,
+                opt.nfev,
+                opt.status,
+            ]
+        except:
+            fit_coefs = x0
+            fit_metrics_single = [0,0,0,0]
+
+        return fit_coefs, fit_metrics_single
+
+
     def fit_all_patterns(
         self, 
         resume = False, 
+        real_space_mask = None,
+        multiprocessing_num_threads = None,
         **fit_opts
         ):
         """
@@ -220,6 +322,11 @@ def fit_all_patterns(
         ----------
         resume: bool (optional)
             Set to true to continue a previous fit with more iterations.
+        real_space_mask: np.array() of bools (optional)
+            Only perform the fitting on a subset of the probe positions,
+            where real_space_mask[rx,ry] == True.
+        multiprocessing_num_threads: int (optional)
+            Set to an integer value of threads to parallelize over probe positions.
         fit_opts: args (optional)
             args passed to scipy.optimize.least_squares
 
@@ -242,52 +349,88 @@ def fit_all_patterns(
         if resume:
             assert hasattr(self, "fit_data"), "No existing data resuming fit!"
 
-        fit_data = np.zeros((self.datacube.R_Nx, self.datacube.R_Ny, self.x0.shape[0]))
-        fit_metrics = np.zeros((self.datacube.R_Nx, self.datacube.R_Ny, 4))
+        # init
+        fit_data = np.zeros((self.x0.shape[0], self.datacube.R_Nx, self.datacube.R_Ny))
+        fit_metrics = np.zeros((4, self.datacube.R_Nx, self.datacube.R_Ny))
+
+        # Default fitting options
+        default_opts = {
+            "method": "trf",
+            "verbose": 1,
+            "x_scale": "jac",
+        }
+        default_opts.update(fit_opts)
 
+        # Loop over probe positions
+        # if multiprocessing_num_threads is None:
         for rx, ry in tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny):
-            current_pattern = self.datacube.data[rx, ry, :, :] * self.intensity_scale
-            shared_data = self.static_data.copy()
-            self._cost_history = (
-                []
-            )  # clear this so it doesn't grow: TODO make this not stupid
-
-            try:
-                x0 = self.fit_data.data[rx, ry] if resume else self.x0
-
-                if self.hasJacobian & self.use_jacobian:
-                    opt = least_squares(
-                        self._pattern_error,
-                        x0,
-                        jac=self._jacobian,
-                        bounds=(self.lower_bound, self.upper_bound),
-                        args=(current_pattern, shared_data),
-                        **fit_opts,
-                    )
-                else:
-                    opt = least_squares(
-                        self._pattern_error,
-                        x0,
-                        bounds=(self.lower_bound, self.upper_bound),
-                        args=(current_pattern, shared_data),
-                        **fit_opts,
+            if real_space_mask is not None and real_space_mask[rx,ry] == True:
+                try:
+                    fit_coefs, fit_metrics_single = WPF.fit_single_pattern(
+                        rx,
+                        ry,
+                        **fitopts,
                     )
-
-                fit_data[rx, ry, :] = opt.x
-                fit_metrics[rx, ry, :] = [
-                    opt.cost,
-                    opt.optimality,
-                    opt.nfev,
-                    opt.status,
-                ]
-            # except LinAlgError as err:
-           # added so that sending an interupt or keyboard interupt breaks out of the for loop rather than just the probe
-            except InterruptedError:
-                break
-            except KeyboardInterrupt:
-                break
-            except:
-                warnings.warn(f'Fit on position ({rx,ry}) failed with error')
+                    fit_data[:, rx, ry] = fit_coefs
+                    fit_metrics[:, rx, ry] = fit_metrics_single
+
+                # except LinAlgError as err:
+                # added so that sending an interupt or keyboard interupt breaks out of the for loop rather than just the probe
+                except InterruptedError:
+                    break
+                except KeyboardInterrupt:
+                    break
+                except:
+                    warnings.warn(f'Fit on position ({rx,ry}) failed with error')
+        # else:
+            # def fit_single_probe(
+            #     rx,
+            #     ry,
+            #     ):
+            #     current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
+            #     shared_data = self.static_data.copy()
+            #     x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0.copy()
+
+            #     try:
+            #         if self.hasJacobian & self.use_jacobian:
+            #             opt = least_squares(
+            #                 self._pattern_error,
+            #                 x0,
+            #                 jac=self._jacobian,
+            #                 bounds=(self.lower_bound, self.upper_bound),
+            #                 args=(current_pattern, shared_data),
+            #                 **default_opts,
+            #                 # **fit_opts,
+            #             )
+            #         else:
+            #             opt = least_squares(
+            #                 self._pattern_error,
+            #                 x0,
+            #                 bounds=(self.lower_bound, self.upper_bound),
+            #                 args=(current_pattern, shared_data),
+            #                 **default_opts,
+            #                 # **fit_opts,
+            #             )
+
+            #         fit_data_single = opt.x
+            #         fit_metrics_single = [
+            #             opt.cost,
+            #             opt.optimality,
+            #             opt.nfev,
+            #             opt.status,
+            #         ]
+            #     except:
+            #         fit_data_single = x0
+            #         fit_metrics_single = [
+            #             0.0,
+            #             0.0,
+            #             0.0,
+            #             0.0,
+            #         ]
+
+            # with Pool(multiprocessing_num_threads) as p:
+            #     print(p.map(fit_single_probe, 
+            #         [(0,0), 1,1]))
 
 
         # Convert to RealSlices
@@ -314,7 +457,11 @@ def fit_all_patterns(
             slicelabels=["cost", "optimality", "nfev", "status"],
         )
 
-        self.show_fit_metrics()
+        # Adding try for testing
+        try:
+            self.show_fit_metrics()
+        except:
+            pass
 
         return self.fit_data, self.fit_metrics
 

From f1d9304de2b4c5e6182f7f5151b0cf6c1b3c048c Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 28 Jul 2023 16:10:18 -0700
Subject: [PATCH 14/57] Parallel fitting working

---
 py4DSTEM/process/wholepatternfit/wpf.py | 258 +++++++++++++++++-------
 1 file changed, 184 insertions(+), 74 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 54391b494..c5e9fd232 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -11,7 +11,10 @@
 from matplotlib.gridspec import GridSpec
 import warnings
 
-from multiprocessing import Pool
+# from multiprocessing import Pool
+# import multiprocess as mp
+from mpire import WorkerPool
+
 
 class WholePatternFit:
 
@@ -259,7 +262,7 @@ def fit_single_pattern(
         # Default fitting options
         default_opts = {
             "method": "trf",
-            "verbose": 1,
+            "verbose": 0,
             "x_scale": "jac",
         }
         default_opts.update(fit_opts)
@@ -312,7 +315,8 @@ def fit_all_patterns(
         self, 
         resume = False, 
         real_space_mask = None,
-        multiprocessing_num_threads = None,
+        num_jobs = None,
+        show_fit_metrics = True,
         **fit_opts
         ):
         """
@@ -325,7 +329,7 @@ def fit_all_patterns(
         real_space_mask: np.array() of bools (optional)
             Only perform the fitting on a subset of the probe positions,
             where real_space_mask[rx,ry] == True.
-        multiprocessing_num_threads: int (optional)
+        num_jobs: int (optional)
             Set to an integer value of threads to parallelize over probe positions.
         fit_opts: args (optional)
             args passed to scipy.optimize.least_squares
@@ -356,81 +360,190 @@ def fit_all_patterns(
         # Default fitting options
         default_opts = {
             "method": "trf",
-            "verbose": 1,
+            "verbose": 0,
             "x_scale": "jac",
         }
         default_opts.update(fit_opts)
 
         # Loop over probe positions
-        # if multiprocessing_num_threads is None:
-        for rx, ry in tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny):
-            if real_space_mask is not None and real_space_mask[rx,ry] == True:
-                try:
-                    fit_coefs, fit_metrics_single = WPF.fit_single_pattern(
-                        rx,
-                        ry,
-                        **fitopts,
-                    )
-                    fit_data[:, rx, ry] = fit_coefs
+        if num_jobs is None:
+            for rx, ry in tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny):
+                if real_space_mask is not None and real_space_mask[rx,ry] == True:
+                    current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
+                    shared_data = self.static_data.copy()
+                    x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0.copy()
+
+                    try:
+                        if self.hasJacobian & self.use_jacobian:
+                            opt = least_squares(
+                                self._pattern_error,
+                                x0,
+                                jac=self._jacobian,
+                                bounds=(self.lower_bound, self.upper_bound),
+                                args=(current_pattern, shared_data),
+                                **default_opts,
+                                # **fit_opts,
+                            )
+                        else:
+                            opt = least_squares(
+                                self._pattern_error,
+                                x0,
+                                bounds=(self.lower_bound, self.upper_bound),
+                                args=(current_pattern, shared_data),
+                                **default_opts,
+                                # **fit_opts,
+                            )
+
+                        fit_data_single = opt.x
+                        fit_metrics_single = [
+                            opt.cost,
+                            opt.optimality,
+                            opt.nfev,
+                            opt.status,
+                        ]
+                    except:
+                        fit_data_single = x0
+                        fit_metrics_single = [
+                            0.0,
+                            0.0,
+                            0.0,
+                            0.0,
+                        ],
+
+                    fit_data[:, rx, ry] = fit_data_single
                     fit_metrics[:, rx, ry] = fit_metrics_single
 
-                # except LinAlgError as err:
-                # added so that sending an interupt or keyboard interupt breaks out of the for loop rather than just the probe
-                except InterruptedError:
-                    break
-                except KeyboardInterrupt:
-                    break
-                except:
-                    warnings.warn(f'Fit on position ({rx,ry}) failed with error')
-        # else:
-            # def fit_single_probe(
+        else:
+            # Get list of probe positions
+            if real_space_mask is not None:
+                xa,ya = np.where(real_space_mask)
+            else:
+                xa,ya = np.meshgrid(
+                    np.arange(self.datacube.Rshape[0]),
+                    np.arange(self.datacube.Rshape[1]),
+                    indexing = 'ij',                    
+                    )
+                xa = xa.ravel()
+                ya = ya.ravel()
+            xy = np.vstack((xa,ya))
+
+            fit_inputs = [default_opts]*xy.shape[1]
+            for a0 in range(xy.shape[1]):
+                fit_inputs[a0]['rx'] = xy[0,a0]
+                fit_inputs[a0]['ry'] = xy[1,a0]
+
+            with WorkerPool(n_jobs = num_jobs) as pool:
+                results = pool.map(
+                    self.fit_single_pattern, 
+                    fit_inputs,
+                    progress_bar=True,
+                    )
+
+            for a0 in range(xy.shape[1]):
+                fit_data[:, xy[0,a0], xy[1,a0]] = results[a0][0]
+                fit_metrics[:, xy[0,a0], xy[1,a0]] = results[a0][1]
+
+
+
+            # print(results[0][0].shape)
+            # print(results[0][1].shape)
+
+            # fit_inputs = xy.tolist() + [[default_opts]*xy.shape[0]]
+            # # fits = xy.tolist()
+
+            # print(str(*default_opts))
+
+
+            # Parallel fitting
+            # def fun_wrapper(dict_args):
+            #     return self.fit_single_pattern(**dict_args)
+
+
+            # pool = Pool(processes=num_jobs)
+            # pool.starmap(
+            #     fun_wrapper, 
+            #     fit_inputs,
+            # )
+            # pool.close()
+            # with Pool(4) as p:
+            #     print(p.map(
+            #         self.fit_single_pattern, 
+            #         **tuple(xy)))
+            # rx,ry = 31,31
+            # fit_coefs, fit_metrics_single = self.fit_single_pattern(
             #     rx,
             #     ry,
-            #     ):
-            #     current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
-            #     shared_data = self.static_data.copy()
-            #     x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0.copy()
-
-            #     try:
-            #         if self.hasJacobian & self.use_jacobian:
-            #             opt = least_squares(
-            #                 self._pattern_error,
-            #                 x0,
-            #                 jac=self._jacobian,
-            #                 bounds=(self.lower_bound, self.upper_bound),
-            #                 args=(current_pattern, shared_data),
-            #                 **default_opts,
-            #                 # **fit_opts,
-            #             )
-            #         else:
-            #             opt = least_squares(
-            #                 self._pattern_error,
-            #                 x0,
-            #                 bounds=(self.lower_bound, self.upper_bound),
-            #                 args=(current_pattern, shared_data),
-            #                 **default_opts,
-            #                 # **fit_opts,
-            #             )
-
-            #         fit_data_single = opt.x
-            #         fit_metrics_single = [
-            #             opt.cost,
-            #             opt.optimality,
-            #             opt.nfev,
-            #             opt.status,
-            #         ]
-            #     except:
-            #         fit_data_single = x0
-            #         fit_metrics_single = [
-            #             0.0,
-            #             0.0,
-            #             0.0,
-            #             0.0,
-            #         ]
-
-            # with Pool(multiprocessing_num_threads) as p:
-            #     print(p.map(fit_single_probe, 
-            #         [(0,0), 1,1]))
+            #     **default_opts,
+            # )
+            # fit_data[:, rx, ry] = fit_coefs
+            # fit_metrics[:, rx, ry] = fit_metrics_single
+
+                    # try:
+                    #     fit_coefs, fit_metrics_single = WPF.fit_single_pattern(
+                    #         rx,
+                    #         ry,
+                    #         **fitopts,
+                    #     )
+                    #     fit_data[:, rx, ry] = fit_coefs
+                    #     fit_metrics[:, rx, ry] = fit_metrics_single
+
+                    # # except LinAlgError as err:
+                    # # added so that sending an interupt or keyboard interupt breaks out of the for loop rather than just the probe
+                    # except InterruptedError:
+                    #     break
+                    # except KeyboardInterrupt:
+                    #     break
+                    # except:
+                    #     warnings.warn(f'Fit on position ({rx,ry}) failed with error')
+        # # else:
+        #     # def fit_single_probe(
+        #     #     rx,
+        #     #     ry,
+        #     #     ):
+        #     #     current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
+        #     #     shared_data = self.static_data.copy()
+        #     #     x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0.copy()
+
+        #     #     try:
+        #     #         if self.hasJacobian & self.use_jacobian:
+        #     #             opt = least_squares(
+        #     #                 self._pattern_error,
+        #     #                 x0,
+        #     #                 jac=self._jacobian,
+        #     #                 bounds=(self.lower_bound, self.upper_bound),
+        #     #                 args=(current_pattern, shared_data),
+        #     #                 **default_opts,
+        #     #                 # **fit_opts,
+        #     #             )
+        #     #         else:
+        #     #             opt = least_squares(
+        #     #                 self._pattern_error,
+        #     #                 x0,
+        #     #                 bounds=(self.lower_bound, self.upper_bound),
+        #     #                 args=(current_pattern, shared_data),
+        #     #                 **default_opts,
+        #     #                 # **fit_opts,
+        #     #             )
+
+        #     #         fit_data_single = opt.x
+        #     #         fit_metrics_single = [
+        #     #             opt.cost,
+        #     #             opt.optimality,
+        #     #             opt.nfev,
+        #     #             opt.status,
+        #     #         ]
+        #     #     except:
+        #     #         fit_data_single = x0
+        #     #         fit_metrics_single = [
+        #     #             0.0,
+        #     #             0.0,
+        #     #             0.0,
+        #     #             0.0,
+        #     #         ]
+
+        #     # with Pool(num_jobs) as p:
+        #     #     print(p.map(fit_single_probe, 
+        #     #         [(0,0), 1,1]))
 
 
         # Convert to RealSlices
@@ -457,11 +570,8 @@ def fit_all_patterns(
             slicelabels=["cost", "optimality", "nfev", "status"],
         )
 
-        # Adding try for testing
-        try:
+        if show_fit_metrics:
             self.show_fit_metrics()
-        except:
-            pass
 
         return self.fit_data, self.fit_metrics
 

From f654d7e5fc7024912c81413ce61c40f474824271 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 28 Jul 2023 16:11:05 -0700
Subject: [PATCH 15/57] Cleaning up

---
 py4DSTEM/process/wholepatternfit/wpf.py | 105 +-----------------------
 1 file changed, 1 insertion(+), 104 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index c5e9fd232..dfc641990 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -330,7 +330,7 @@ def fit_all_patterns(
             Only perform the fitting on a subset of the probe positions,
             where real_space_mask[rx,ry] == True.
         num_jobs: int (optional)
-            Set to an integer value of threads to parallelize over probe positions.
+            Set to an integer value giving the number of jobs to parallelize over probe positions.
         fit_opts: args (optional)
             args passed to scipy.optimize.least_squares
 
@@ -443,109 +443,6 @@ def fit_all_patterns(
                 fit_data[:, xy[0,a0], xy[1,a0]] = results[a0][0]
                 fit_metrics[:, xy[0,a0], xy[1,a0]] = results[a0][1]
 
-
-
-            # print(results[0][0].shape)
-            # print(results[0][1].shape)
-
-            # fit_inputs = xy.tolist() + [[default_opts]*xy.shape[0]]
-            # # fits = xy.tolist()
-
-            # print(str(*default_opts))
-
-
-            # Parallel fitting
-            # def fun_wrapper(dict_args):
-            #     return self.fit_single_pattern(**dict_args)
-
-
-            # pool = Pool(processes=num_jobs)
-            # pool.starmap(
-            #     fun_wrapper, 
-            #     fit_inputs,
-            # )
-            # pool.close()
-            # with Pool(4) as p:
-            #     print(p.map(
-            #         self.fit_single_pattern, 
-            #         **tuple(xy)))
-            # rx,ry = 31,31
-            # fit_coefs, fit_metrics_single = self.fit_single_pattern(
-            #     rx,
-            #     ry,
-            #     **default_opts,
-            # )
-            # fit_data[:, rx, ry] = fit_coefs
-            # fit_metrics[:, rx, ry] = fit_metrics_single
-
-                    # try:
-                    #     fit_coefs, fit_metrics_single = WPF.fit_single_pattern(
-                    #         rx,
-                    #         ry,
-                    #         **fitopts,
-                    #     )
-                    #     fit_data[:, rx, ry] = fit_coefs
-                    #     fit_metrics[:, rx, ry] = fit_metrics_single
-
-                    # # except LinAlgError as err:
-                    # # added so that sending an interupt or keyboard interupt breaks out of the for loop rather than just the probe
-                    # except InterruptedError:
-                    #     break
-                    # except KeyboardInterrupt:
-                    #     break
-                    # except:
-                    #     warnings.warn(f'Fit on position ({rx,ry}) failed with error')
-        # # else:
-        #     # def fit_single_probe(
-        #     #     rx,
-        #     #     ry,
-        #     #     ):
-        #     #     current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
-        #     #     shared_data = self.static_data.copy()
-        #     #     x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0.copy()
-
-        #     #     try:
-        #     #         if self.hasJacobian & self.use_jacobian:
-        #     #             opt = least_squares(
-        #     #                 self._pattern_error,
-        #     #                 x0,
-        #     #                 jac=self._jacobian,
-        #     #                 bounds=(self.lower_bound, self.upper_bound),
-        #     #                 args=(current_pattern, shared_data),
-        #     #                 **default_opts,
-        #     #                 # **fit_opts,
-        #     #             )
-        #     #         else:
-        #     #             opt = least_squares(
-        #     #                 self._pattern_error,
-        #     #                 x0,
-        #     #                 bounds=(self.lower_bound, self.upper_bound),
-        #     #                 args=(current_pattern, shared_data),
-        #     #                 **default_opts,
-        #     #                 # **fit_opts,
-        #     #             )
-
-        #     #         fit_data_single = opt.x
-        #     #         fit_metrics_single = [
-        #     #             opt.cost,
-        #     #             opt.optimality,
-        #     #             opt.nfev,
-        #     #             opt.status,
-        #     #         ]
-        #     #     except:
-        #     #         fit_data_single = x0
-        #     #         fit_metrics_single = [
-        #     #             0.0,
-        #     #             0.0,
-        #     #             0.0,
-        #     #             0.0,
-        #     #         ]
-
-        #     # with Pool(num_jobs) as p:
-        #     #     print(p.map(fit_single_probe, 
-        #     #         [(0,0), 1,1]))
-
-
         # Convert to RealSlices
         model_names = []
         for m in self.model:

From e8e1f3a754e97af860ab6fd7563b4b4d73625e34 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 28 Jul 2023 17:53:42 -0700
Subject: [PATCH 16/57] minor cleanup

---
 py4DSTEM/process/wholepatternfit/wpf.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index dfc641990..b9035a736 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -403,12 +403,7 @@ def fit_all_patterns(
                         ]
                     except:
                         fit_data_single = x0
-                        fit_metrics_single = [
-                            0.0,
-                            0.0,
-                            0.0,
-                            0.0,
-                        ],
+                        fit_metrics_single = [0,0,0,0]
 
                     fit_data[:, rx, ry] = fit_data_single
                     fit_metrics[:, rx, ry] = fit_metrics_single

From 5b89bce2a118f668492a0bd205536757809600b7 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Fri, 28 Jul 2023 17:56:07 -0700
Subject: [PATCH 17/57] Removing old polar peaks

---
 .../process/wholepatternfit/polar_peaks.py    | 658 ------------------
 1 file changed, 658 deletions(-)
 delete mode 100644 py4DSTEM/process/wholepatternfit/polar_peaks.py

diff --git a/py4DSTEM/process/wholepatternfit/polar_peaks.py b/py4DSTEM/process/wholepatternfit/polar_peaks.py
deleted file mode 100644
index fad76c4a6..000000000
--- a/py4DSTEM/process/wholepatternfit/polar_peaks.py
+++ /dev/null
@@ -1,658 +0,0 @@
-"""
-This sub-module contains functions for polar transform peak detection of amorphous / semicrystalline datasets.
-
-"""
-
-from py4DSTEM import tqdmnd
-from scipy.ndimage import gaussian_filter
-from sklearn.decomposition import PCA
-from itertools import product
-from typing import Optional
-from matplotlib.colors import hsv_to_rgb
-
-import time, sys
-import numpy as np
-import matplotlib.pyplot as plt
-
-
-
-class PolarPeaks:
-    """
-    Primary class for polar transform peak detection.
-    """
-
-
-    def __init__(
-        self,
-        datacube,
-        radial_min = 0.0,
-        radial_max = None,
-        radial_step = 1.0,
-        num_annular_bins = 180,
-        progress_bar = True,
-        ):
-        """
-        Initialize class by performing an intensity-preserving polar transformation.
-
-        Parameters
-        --------
-        datacube: py4DSTEM.io.DataCube
-            4D-STEM dataset, requires origin calibration
-        radial_min: float
-            Minimum radius of polar transformation.
-        radial_max: float
-            Maximum radius of polar transformation.
-        radial_step: float
-            Width of radial bins of polar transformation.
-        num_annular_bins: int
-            Number of bins in annular direction. Note that we fold data over 
-            180 degrees periodically, so setting this value to 60 gives bin 
-            widths of 180/60 = 3.0 degrees.
-        progress_bar: bool
-            Turns on the progress bar for the polar transformation
-    
-        Returns
-        --------
-        
-
-        """
-
-        # radial bin coordinates
-        if radial_max is None:
-            radial_max = np.min(datacube.Qshape) / np.sqrt(2)
-        self.radial_bins = np.arange(
-            radial_min,
-            radial_max,
-            radial_step,
-            )
-        self.radial_step = np.array(radial_step)
-
-        # annular bin coordinates
-        self.annular_bins = np.linspace(
-            0,
-            np.pi,
-            num_annular_bins,
-            endpoint = False,
-            )
-        self.annular_step = self.annular_bins[1] - self.annular_bins[0]
-
-        # init polar transformation array
-        self.polar_shape = np.array((self.annular_bins.shape[0], self.radial_bins.shape[0]))
-        self.polar_size = np.prod(self.polar_shape)
-        self.data_polar = np.zeros((
-            datacube.R_Nx,
-            datacube.R_Ny,
-            self.polar_shape[0],
-            self.polar_shape[1],
-            ))
-
-        # init coordinates
-        xa, ya = np.meshgrid(
-            np.arange(datacube.Q_Nx),
-            np.arange(datacube.Q_Ny),
-            indexing = 'ij',
-        )
-
-        # polar transformation
-        for rx, ry in tqdmnd(
-            range(datacube.R_Nx),
-            range(datacube.R_Ny),
-            desc="polar transformation",
-            unit=" images",
-            disable=not progress_bar,
-            ):
-
-            # shifted coordinates
-            x = xa - datacube.calibration.get_qx0(rx,ry)
-            y = ya - datacube.calibration.get_qy0(rx,ry)
-
-            # polar coordinate indices
-            r_ind = (np.sqrt(x**2 + y**2) - self.radial_bins[0]) / self.radial_step
-            t_ind = np.arctan2(y, x) / self.annular_step
-            r_ind_floor = np.floor(r_ind).astype('int')
-            t_ind_floor = np.floor(t_ind).astype('int')
-            dr = r_ind - r_ind_floor
-            dt = t_ind - t_ind_floor
-            # t_ind_floor = np.mod(t_ind_floor, self.num_annular_bins)
-
-            # polar transformation
-            sub = np.logical_and(r_ind_floor >= 0, r_ind_floor < self.polar_shape[1])
-            im = np.bincount(
-                r_ind_floor[sub] + \
-                np.mod(t_ind_floor[sub],self.polar_shape[0]) * self.polar_shape[1],
-                weights = datacube.data[rx,ry][sub] * (1 - dr[sub]) * (1 - dt[sub]),
-                minlength = self.polar_size,
-            )
-            im += np.bincount(
-                r_ind_floor[sub] + \
-                np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]) * self.polar_shape[1],
-                weights = datacube.data[rx,ry][sub] * (1 - dr[sub]) * (    dt[sub]),
-                minlength = self.polar_size,
-            )
-            sub = np.logical_and(r_ind_floor >= -1, r_ind_floor < self.polar_shape[1]-1)
-            im += np.bincount(
-                r_ind_floor[sub] + 1 + \
-                np.mod(t_ind_floor[sub],self.polar_shape[0]) * self.polar_shape[1],
-                weights = datacube.data[rx,ry][sub] * (    dr[sub]) * (1 - dt[sub]),
-                minlength = self.polar_size,
-            )
-            im += np.bincount(
-                r_ind_floor[sub] + 1 + \
-                np.mod(t_ind_floor[sub] + 1,self.polar_shape[0]) * self.polar_shape[1],
-                weights = datacube.data[rx,ry][sub] * (    dr[sub]) * (    dt[sub]),
-                minlength = self.polar_size,
-            )
-
-            # output
-            self.data_polar[rx,ry] = np.reshape(im, self.polar_shape)
-
-
-
-
-    def fit_peaks(
-        self,
-        num_peaks_fit = 1,
-        sigma_radial_pixels = 0.0,
-        sigma_annular_degrees = 1.0,
-        progress_bar = True,
-        ):
-        """
-        Fit both background signal and peak positions and intensities for each radial bin.
-
-        Parameters
-        --------
-        progress_bar: bool
-            Turns on the progress bar for the polar transformation
-    
-        Returns
-        --------
-        
-
-        """
-
-        # sigma in pixels
-        self._sigma_radial_px = sigma_radial_pixels / self.radial_step
-        self._sigma_annular_px = np.deg2rad(sigma_annular_degrees) / self.annular_step
-
-        # init
-        self.radial_median = np.zeros((
-            self.data_polar.shape[0],
-            self.data_polar.shape[1],
-            self.polar_shape[1],
-            ))
-        self.radial_peaks = np.zeros((
-            self.data_polar.shape[0],
-            self.data_polar.shape[1],
-            self.polar_shape[1],
-            num_peaks_fit,
-            2,
-            ))
-
-
-        # loop over probe positions
-        for rx, ry in tqdmnd(
-            self.data_polar.shape[0],
-            self.data_polar.shape[1],
-            desc="polar transformation",
-            unit=" positions",
-            disable=not progress_bar,
-            ):
-
-            im = gaussian_filter(
-                self.data_polar[rx,ry],
-                sigma = (self._sigma_annular_px, self._sigma_radial_px),
-                mode = ('wrap', 'nearest'),
-                truncate = 3.0,
-            )
-
-            # background signal
-            self.radial_median[rx,ry] = np.median(im,axis=0)
-
-            # local maxima
-            sub_peaks = np.logical_and(
-                im > np.roll(im,-1,axis=0),
-                im > np.roll(im, 1,axis=0),
-            )
-
-            for a0 in range(self.polar_shape[1]):
-                inds = np.atleast_1d(np.squeeze(np.argwhere(sub_peaks[:,a0])))
-                if inds.size > 0:
-                    vals = im[inds,a0]
-
-                    inds_sort = np.argsort(vals)[::-1]
-                    inds_keep = inds_sort[:num_peaks_fit]
-
-                    peaks_val = np.maximum(vals[inds_keep] - self.radial_median[rx,ry,a0], 0)
-                    peaks_ind = inds[inds_keep]
-                    peaks_angle = self.annular_bins[peaks_ind]
-
-                    # TODO - add subpixel peak fitting?
-
-                    # output
-                    num_peaks = peaks_val.shape[0]
-                    self.radial_peaks[rx,ry,a0,:num_peaks,0] = peaks_angle
-                    self.radial_peaks[rx,ry,a0,:num_peaks,1] = peaks_val
-
-    def plot_peak_signal(
-        self,
-        figsize = (8,4),
-        ):
-        """
-        Plot the mean peak signal from all radial bins.
-
-        """
-
-        sig = np.mean(
-            np.sum(
-                self.radial_peaks[:,:,:,:,1], axis = 3,
-                ), axis = (0,1),
-            )
-
-        # plot
-        fig = plt.figure(figsize = figsize)
-        ax1 = fig.add_subplot(111)
-        ax2 = ax1.twiny()
-
-        ax1.plot(
-            self.radial_bins,
-            sig,
-            )
-        ax1.set_xlabel('Radial Bin [pixels]')
-        ax1.set_ylabel('Mean Peak Signal [counts]')
-
-        def tick_function(x):
-            v = (x - self.radial_bins[0]) / self.radial_step
-            return ["%.0f" % z for z in v]
-        tick_locations = ax1.get_xticks()
-        tick_locations += self.radial_bins[0]
-        ax2.set_xticks(tick_locations)
-        ax2.set_xticklabels(tick_function(tick_locations))
-        ax2.set_xlim(ax1.get_xlim())
-        ax2.set_xlabel('Radial Bin Index')
-        ax2.minorticks_on()
-        ax2.grid()
-
-        plt.show()
-
-
-
-        # fig,ax = plt.subplots(figsize = figsize)
-        # ax.plot(
-        #     self.radial_bins,
-        #     sig,
-        #     )
-
-        
-
-
-        #         fig = plt.figure()
-        # ax1 = fig.add_subplot(111)
-        # ax2 = ax1.twiny()
-
-        # X = np.linspace(0,1,1000)
-        # Y = np.cos(X*20)
-
-        # ax1.plot(X,Y)
-        # ax1.set_xlabel(r"Original x-axis: $X$")
-
-        # new_tick_locations = np.array([.2, .5, .9])
-
-        # def tick_function(X):
-        #     V = 1/(1+X)
-        #     return ["%.3f" % z for z in V]
-
-        # ax2.set_xlim(ax1.get_xlim())
-        # ax2.set_xticks(new_tick_locations)
-        # ax2.set_xticklabels(tick_function(new_tick_locations))
-        # ax2.set_xlabel(r"Modified x-axis: $1/(1+X)$")
-        # plt.show()
-
-
-
-
-    def orientation_map(
-        self,
-        radial_index = 0,
-        peak_index = 0,
-        intensity_range = (0,1),
-        plot_result = True,
-        ):
-        """
-        Create an RGB orientation map from a given peak bin
-
-        Parameters
-        --------
-        progress_bar: bool
-            Turns on the progress bar for the polar transformation
-    
-        Returns
-        --------
-        im_orientation: np.array
-            rgb image array
-
-        """
-
-        # intensity mask
-        val = np.squeeze(self.radial_peaks[:,:,radial_index,peak_index,1]).copy()
-        val -= intensity_range[0]
-        val /= intensity_range[1] - intensity_range[0]
-        val = np.clip(val,0,1)
-
-        # orientation
-        hue = np.squeeze(self.radial_peaks[:,:,radial_index,peak_index,0]).copy()
-        hue = np.mod(2*hue,1)
-
-
-        # generate image
-        im_orientation = np.ones((
-            self.data_polar.shape[0],
-            self.data_polar.shape[1],
-            3))
-        im_orientation[:,:,0] = hue
-        im_orientation[:,:,2] = val
-        im_orientation = hsv_to_rgb(im_orientation)
-
-        if plot_result:
-            fig, ax = plt.subplots(figsize=(8,8))
-            ax.imshow(
-                im_orientation,
-                vmin = 0,
-                vmax = 5,
-            )
-        # ax.plot(
-        #     im[:,6]
-        #     )
-
-        return im_orientation
-
-
-    def cluster_grains(
-        self,
-        radial_index = 0,
-        threshold_add = 1.0,
-        threshold_grow = 0.0,
-        angle_tolerance_deg = 5.0,
-        progress_bar = False,
-        ):
-        """
-        Cluster grains from a specific radial bin
-
-        Parameters
-        --------
-        radial_index: int
-            Which radial bin to perform the clustering over.
-        threshold_add: float
-            Minimum signal required for a probe position to initialize a cluster.
-        threshold_grow: float
-            Minimum signal required for a probe position to be added to a cluster.
-        angle_tolerance_deg: float
-            Rotation rolerance for clustering grains.
-        progress_bar: bool
-            Turns on the progress bar for the polar transformation
-    
-        Returns
-        --------
-        
-
-        """
-
-        # Get data
-        phi = np.squeeze(self.radial_peaks[:,:,radial_index,:,0]).copy()
-        sig = np.squeeze(self.radial_peaks[:,:,radial_index,:,1]).copy()
-        sig_init = sig.copy()
-        mark = sig >= threshold_grow
-        sig[np.logical_not(mark)] = 0
-        
-        # init
-        self.cluster_sizes = np.array((), dtype='int')
-        self.cluster_sig = np.array(())
-        self.cluster_inds = []
-        inds_all = np.zeros_like(sig, dtype='int')
-        inds_all.ravel()[:] = np.arange(inds_all.size)
-
-        # Tolerance
-        tol = np.deg2rad(angle_tolerance_deg)
-
-        # Main loop
-        search = True
-        while search is True:
-            inds_grain = np.argmax(sig)
-            val = sig.ravel()[inds_grain]
-
-            if val < threshold_add:
-                search = False
-            
-            else:
-                # progressbar
-                comp = 1 - np.mean(np.max(mark,axis = 2))
-                update_progress(comp)
-
-                # Start cluster
-                x,y,z = np.unravel_index(inds_grain, sig.shape)
-                mark[x,y,z] = False
-                sig[x,y,z] = 0
-                phi_cluster = phi[x,y,z]
-
-                # Neighbors to search
-                xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1)
-                yr = np.clip(y + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1)
-                inds_cand = inds_all[xr[:,None],yr[None],:].ravel()
-                inds_cand = np.delete(inds_cand, mark.ravel()[inds_cand] == False)
-
-                if inds_cand.size == 0:
-                    grow = False
-                else:
-                    grow = True
-
-                # grow the cluster
-                while grow is True:
-                    inds_new = np.array((),dtype='int')
-
-                    keep = np.zeros(inds_cand.size, dtype='bool')
-                    for a0 in range(inds_cand.size):
-                        xc,yc,zc = np.unravel_index(inds_cand[a0], sig.shape)
-
-                        phi_test = phi[xc,yc,zc]
-                        dphi = np.mod(phi_cluster - phi_test + np.pi/2.0, np.pi) - np.pi/2.0
-
-                        if np.abs(dphi) < tol:
-                            keep[a0] = True
-
-                            sig[xc,yc,zc] = 0
-                            mark[xc,yc,zc] = False 
-
-                            xr = np.clip(xc + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1)
-                            yr = np.clip(yc + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1)
-                            inds_add = inds_all[xr[:,None],yr[None],:].ravel()
-                            inds_new = np.append(inds_new, inds_add)
-
-
-                    inds_grain = np.append(inds_grain, inds_cand[keep])
-                    inds_cand = np.unique(np.delete(inds_new, mark.ravel()[inds_new] == False))
-
-                    if inds_cand.size == 0:
-                        grow = False
-
-            # convert grain to x,y coordinates, add = list
-            xg,yg,zg = np.unravel_index(inds_grain, sig.shape)
-            xyg = np.unique(np.vstack((xg,yg)), axis = 1)
-            sig_mean = np.mean(sig_init.ravel()[inds_grain])
-            self.cluster_sizes = np.append(self.cluster_sizes, xyg.shape[1])
-            self.cluster_sig = np.append(self.cluster_sig, sig_mean)
-            self.cluster_inds.append(xyg)
-
-        # finish progressbar
-        update_progress(1)
-
-    def cluster_plot_size(
-        self,
-        area_min = None,
-        area_max = None,
-        area_step = 1,
-        weight_intensity = False,
-        pixel_area = 1.0,
-        pixel_area_units = 'px^2',
-        figsize = (8,6),
-        returnfig = False,
-        ):
-        """
-        Plot the cluster sizes
-
-        Parameters
-        --------
-        area_min: int (optional)
-            Min area bin in pixels
-        area_max: int (optional)
-            Max area bin in pixels
-        area_step: int (optional)
-            Step size of the histogram bin
-        weight_intensity: bool
-            Weight histogram by the peak intensity.
-        pixel_area: float
-            Size of pixel area unit square
-        pixel_area_units: string
-            Units of the pixel area 
-        figsize: tuple
-            Size of the figure panel
-        returnfig: bool
-            Setting this to true returns the figure and axis handles
-    
-        Returns
-        --------
-        fig, ax (optional)
-            Figure and axes handles
-
-        """
-
-        if area_max is None:
-            area_max = np.max(self.cluster_sizes)
-        area = np.arange(0,area_max,area_step)
-        if area_min is None:
-            sub = self.cluster_sizes.astype('int') < area_max
-        else:
-            sub = np.logical_and(
-                self.cluster_sizes.astype('int') >= area_min,
-                self.cluster_sizes.astype('int') < area_max
-                )
-        if weight_intensity:
-            hist = np.bincount(
-                self.cluster_sizes[sub] // area_step,
-                weights = self.cluster_sig[sub],
-                minlength = area.shape[0],
-                )
-        else:
-            hist = np.bincount(
-                self.cluster_sizes[sub] // area_step,
-                minlength = area.shape[0],
-                )
-
-
-        # plotting
-        fig,ax = plt.subplots(figsize = figsize)
-        ax.bar(
-            area * pixel_area,
-            hist,
-            width = 0.8 * pixel_area * area_step,
-            )
-        ax.set_xlim((0,area_max*pixel_area))
-        ax.set_xlabel('Grain Area [' + pixel_area_units + ']')
-        if weight_intensity:
-            ax.set_ylabel('Total Signal [arb. units]')
-        else:
-            ax.set_ylabel('Number of Grains')
-
-        if returnfig:
-            return fig,ax
-
-
-    def background_pca(
-        self,
-        pca_index = 0,
-        intensity_range = (0,1),
-        normalize_mean = True,
-        normalize_std = True,
-        plot_result = True,
-        plot_coef = False,
-        ):
-        """
-        Generate PCA decompositions of the background signal
-
-        Parameters
-        --------
-        progress_bar: bool
-            Turns on the progress bar for the polar transformation
-    
-        Returns
-        --------
-        im_pca: np,array
-            rgb image array
-        coef_pca: np.array
-            radial PCA component selected
-
-        """
-
-        # PCA decomposition
-        shape = self.radial_median.shape
-        A = np.reshape(self.radial_median, (shape[0]*shape[1],shape[2]))
-        if normalize_mean:
-            A -= np.mean(A,axis=0)
-        if normalize_std:
-            A /= np.std(A,axis=0)
-        pca = PCA(n_components=np.maximum(pca_index+1,2))
-        pca.fit(A)
-
-        components = pca.components_
-        loadings = pca.transform(A)
-
-        # output image data
-        sig_pca = np.reshape(loadings[:,pca_index], shape[0:2])
-        sig_pca -= intensity_range[0]
-        sig_pca /= intensity_range[1] - intensity_range[0]
-        sig_pca = np.clip(sig_pca,0,1)
-        im_pca = np.tile(sig_pca[:,:,None],(1,1,3))
-
-        # output PCA coefficient
-        coef_pca = np.vstack((
-            self.radial_bins,
-            components[pca_index,:]
-        )).T
-
-        if plot_result:
-            fig, ax = plt.subplots(figsize=(8,8))
-            ax.imshow(
-                im_pca,
-                vmin = 0,
-                vmax = 5,
-            )
-        if plot_coef:
-            fig, ax = plt.subplots(figsize=(8,4))
-            ax.plot(
-                coef_pca[:,0],
-                coef_pca[:,1]
-                )
-
-        return im_pca, coef_pca
-
-
-# Progressbar taken from stackexchange:
-# https://stackoverflow.com/questions/3160699/python-progress-bar
-def update_progress(progress):
-    barLength = 60 # Modify this to change the length of the progress bar
-    status = ""
-    if isinstance(progress, int):
-        progress = float(progress)
-    if not isinstance(progress, float):
-        progress = 0
-        status = "error: progress var must be float\r\n"
-    if progress < 0:
-        progress = 0
-        status = "Halt...\r\n"
-    if progress >= 1:
-        progress = 1
-        status = "Done...\r\n"
-    block = int(round(barLength*progress))
-    text = "\rPercent: [{0}] {1}% {2}".format( "#"*block + "-"*(barLength-block), 
-        np.round(progress*100,4), 
-        status)
-    sys.stdout.write(text)
-    sys.stdout.flush()

From 32a4a41213cebb568338d71dc17dd2c3d530382e Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Sun, 30 Jul 2023 12:03:54 -0700
Subject: [PATCH 18/57] Adding mpire to requirements

---
 setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.py b/setup.py
index b0c7fa081..a806c0131 100644
--- a/setup.py
+++ b/setup.py
@@ -38,6 +38,7 @@
         'dask >= 2.3.0',
         'distributed >= 2.3.0',
         'emdfile >= 0.0.10',
+        'mpire >= 2.7.1',
         ],
     extras_require={
         'ipyparallel': ['ipyparallel >= 6.2.4', 'dill >= 0.3.3'],

From 907489e328c8335374f8d331d680a2f7ad1b5254 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Sun, 30 Jul 2023 12:21:23 -0700
Subject: [PATCH 19/57] Removing polar peaks py file

---
 py4DSTEM/process/wholepatternfit/__init__.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/py4DSTEM/process/wholepatternfit/__init__.py b/py4DSTEM/process/wholepatternfit/__init__.py
index 97d2e083e..8fb0e351e 100644
--- a/py4DSTEM/process/wholepatternfit/__init__.py
+++ b/py4DSTEM/process/wholepatternfit/__init__.py
@@ -1,3 +1,2 @@
 from .wp_models import *
 from .wpf import *
-from .polar_peaks import *

From 6e87332ca6f01a553620b23cddc14376000c1dba Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Wed, 2 Aug 2023 15:42:59 -0700
Subject: [PATCH 20/57] Adding get_strained_crystal method

---
 py4DSTEM/process/diffraction/crystal.py | 53 +++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index 1c4ac9073..c73358dc1 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -139,6 +139,59 @@ def calculate_lattice(self):
             self.pymatgen_available = True            
         else:
             self.pymatgen_available = False
+
+    def get_strained_crystal(
+        self,
+        exx = 0.0,
+        eyy = 0.0,
+        ezz = 0.0,
+        exy = 0.0,
+        exz = 0.0,
+        eyz = 0.0,
+        deformation_matrix = None,
+        return_deformation_matrix = False,
+        ):
+        """
+        This method returns new Crystal class with strain applied. The directions of (x,y,z)
+        are with respect to the default Crystal orientation, which can be checked with
+        print(Crystal.lat_real)
+        """
+
+        # deformation matrix
+        if deformation_matrix is None:
+            deformation_matrix = np.array([
+                [1+exx,     exy,    exz     ],
+                [exy,       1+eyy,  eyz     ],
+                [exz,       eyz,    1+ezz   ],
+            ])
+
+        # new unit cell
+        lat_new = self.lat_real @ deformation_matrix
+        a_new = np.linalg.norm(lat_new[0,:])
+        b_new = np.linalg.norm(lat_new[1,:])
+        c_new = np.linalg.norm(lat_new[2,:])
+        alpha_new = np.rad2deg(np.arccos(np.clip(np.sum(
+            lat_new[1,:]*lat_new[2,:])/b_new/c_new,-1,1)))
+        beta_new  = np.rad2deg(np.arccos(np.clip(np.sum(
+            lat_new[0,:]*lat_new[2,:])/a_new/c_new,-1,1)))
+        gamma_new = np.rad2deg(np.arccos(np.clip(np.sum(
+            lat_new[0,:]*lat_new[1,:])/a_new/b_new,-1,1)))
+        cell_new = np.array(
+            (a_new,b_new,c_new,alpha_new,beta_new,gamma_new)
+        )
+
+        # Make new crystal
+        from py4DSTEM.process.diffraction import Crystal
+        crystal_strained = Crystal(
+            self.positions, 
+            self.numbers, 
+            cell_new,
+        )
+
+        if return_deformation_matrix:
+            return crystal_strained, deformation_matrix
+        else:
+            return crystal_strained
         
 
     def from_CIF(CIF, conventional_standard_structure=True):

From dc6a11fa184ebb6b4a6c9cc258b586e29c5fa32b Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Wed, 2 Aug 2023 16:17:48 -0700
Subject: [PATCH 21/57] Adding better strain method

---
 py4DSTEM/process/diffraction/crystal.py | 62 ++++++++++++++++++++++---
 1 file changed, 56 insertions(+), 6 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index c73358dc1..2676a23f5 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -4,7 +4,8 @@
 import matplotlib.pyplot as plt
 from fractions import Fraction
 from typing import Union, Optional
-from copy import deepcopy
+# import copy
+# from copy import deepcopy
 from scipy.optimize import curve_fit
 import sys
 
@@ -108,7 +109,16 @@ def __init__(
         
         # Calculate lattice parameters
         self.calculate_lattice()
-        
+    
+    # TODO check for and copy all important attributes
+    def copy(self):
+        crystal = Crystal(
+            self.positions,
+            self.numbers,
+            self.cell,
+        )
+        return crystal
+
     def calculate_lattice(self):
         # calculate unit cell lattice vectors
         a = self.cell[0]
@@ -439,13 +449,36 @@ def calculate_structure_factors(
         k_max: float = 2.0,
         tol_structure_factor: float = 1e-4,
         return_intensities: bool = False,
-    ):
+        exx = 0.0,
+        eyy = 0.0,
+        ezz = 0.0,
+        exy = 0.0,
+        exz = 0.0,
+        eyz = 0.0,
+        deformation_matrix = None,
+        ):
+
+
         """
         Calculate structure factors for all hkl indices up to max scattering vector k_max
 
-        Args:
-            k_max (numpy float):                max scattering vector to include (1/Angstroms)
-            tol_structure_factor (numpy float): tolerance for removing low-valued structure factors
+        Parameters
+        --------
+        
+        k_max: float                
+            max scattering vector to include (1/Angstroms)
+        tol_structure_factor: float
+            tolerance for removing low-valued structure factors
+        return_intensities: bool
+            return the intensities and positions of all structure factor peaks.
+        exx 
+
+        Returns
+        --------
+        (q_SF, I_SF)
+            Tuple of the q vectors and intensities of each structure factor.
+        
+
         """
 
         # Store k_max
@@ -479,6 +512,23 @@ def calculate_structure_factors(
         # g_vec_all = self.lat_inv @ hkl
         g_vec_all =  (hkl.T @ self.lat_inv).T
 
+        # strain
+        if (exx != 0.0 or \
+            eyy != 0.0 or \
+            ezz != 0.0 or \
+            exy != 0.0 or \
+            exz != 0.0 or \
+            eyz != 0.0) and \
+            deformation_matrix is None:
+                deformation_matrix = np.array([
+                    [1+exx,     exy,    exz     ],
+                    [exy,       1+eyy,  eyz     ],
+                    [exz,       eyz,    1+ezz   ],
+                ])
+        if deformation_matrix is not None:
+            # Note that we need to take the matrix inverse, since we're operating in reciprocal space.
+            g_vec_all = np.linalg.inv(deformation_matrix) @ g_vec_all
+                
         # Delete lattice vectors outside of k_max
         keep = np.linalg.norm(g_vec_all, axis=0) <= self.k_max
         self.hkl = hkl[:, keep]

From 9d676ee177c45343a09ec335ab1c0ced071ff274 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Wed, 2 Aug 2023 16:33:17 -0700
Subject: [PATCH 22/57] Attempting to add strain to the bloch wave calcs

---
 py4DSTEM/process/diffraction/crystal.py       | 15 ++++++-
 py4DSTEM/process/diffraction/crystal_bloch.py | 40 +++++++++++++++++++
 2 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index 2676a23f5..a4ba23422 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -471,7 +471,20 @@ def calculate_structure_factors(
             tolerance for removing low-valued structure factors
         return_intensities: bool
             return the intensities and positions of all structure factor peaks.
-        exx 
+        exx: float
+            Strain in the x direction
+        eyy: float
+            Strain in the y direction
+        ezz: float
+            Strain in the z direction
+        exy: float
+            Shear in the x,y direction
+        exz: float
+            Shear in the x,z direction
+        eyz: float
+            Shear in the y,z direction
+        deformation_matrix: np.array
+            3x3 deformation matrix in real space. 
 
         Returns
         --------
diff --git a/py4DSTEM/process/diffraction/crystal_bloch.py b/py4DSTEM/process/diffraction/crystal_bloch.py
index 6a3c9b1ac..b297a0371 100644
--- a/py4DSTEM/process/diffraction/crystal_bloch.py
+++ b/py4DSTEM/process/diffraction/crystal_bloch.py
@@ -28,6 +28,13 @@ def calculate_dynamical_structure_factors(
     recompute_kinematic_structure_factors=True,
     g_vec_precision=None,
     verbose=True,
+    exx = 0.0,
+    eyy = 0.0,
+    ezz = 0.0,
+    exy = 0.0,
+    exz = 0.0,
+    eyz = 0.0,
+    deformation_matrix = None,
 ):
     """
     Calculate and store the relativistic corrected structure factors used for Bloch computations
@@ -74,6 +81,22 @@ def calculate_dynamical_structure_factors(
                                         substantial speedup at the cost of some reduced accuracy
 
         See WK_scattering_factors.py for details on the Weickenmeier-Kohl form factors.
+
+    exx: float
+        Strain in the x direction
+    eyy: float
+        Strain in the y direction
+    ezz: float
+        Strain in the z direction
+    exy: float
+        Shear in the x,y direction
+    exz: float
+        Shear in the x,z direction
+    eyz: float
+        Shear in the y,z direction
+    deformation_matrix: np.array
+        3x3 deformation matrix in real space.
+
     """
 
     assert method in (
@@ -133,6 +156,23 @@ def calculate_dynamical_structure_factors(
     hkl = np.vstack([xa.ravel(), ya.ravel(), za.ravel()])
     g_vec_all = lat_inv @ hkl
 
+    # strain
+    if (exx != 0.0 or \
+        eyy != 0.0 or \
+        ezz != 0.0 or \
+        exy != 0.0 or \
+        exz != 0.0 or \
+        eyz != 0.0) and \
+        deformation_matrix is None:
+            deformation_matrix = np.array([
+                [1+exx,     exy,    exz     ],
+                [exy,       1+eyy,  eyz     ],
+                [exz,       eyz,    1+ezz   ],
+            ])
+    if deformation_matrix is not None:
+        # Note that we need to take the matrix inverse, since we're operating in reciprocal space.
+        g_vec_all = np.linalg.inv(deformation_matrix) @ g_vec_all
+
     # Delete lattice vectors outside of k_max
     keep = np.linalg.norm(g_vec_all, axis=0) <= k_max
     hkl = hkl[:, keep]

From 1e4046512134fea71d9d16e9c29770eb616db999 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Wed, 2 Aug 2023 16:35:43 -0700
Subject: [PATCH 23/57] Commenting out strain in the dynamical diff SF function
 for now

---
 py4DSTEM/process/diffraction/crystal_bloch.py | 46 +++++++++----------
 1 file changed, 23 insertions(+), 23 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal_bloch.py b/py4DSTEM/process/diffraction/crystal_bloch.py
index b297a0371..48918a91b 100644
--- a/py4DSTEM/process/diffraction/crystal_bloch.py
+++ b/py4DSTEM/process/diffraction/crystal_bloch.py
@@ -28,13 +28,13 @@ def calculate_dynamical_structure_factors(
     recompute_kinematic_structure_factors=True,
     g_vec_precision=None,
     verbose=True,
-    exx = 0.0,
-    eyy = 0.0,
-    ezz = 0.0,
-    exy = 0.0,
-    exz = 0.0,
-    eyz = 0.0,
-    deformation_matrix = None,
+    # exx = 0.0,
+    # eyy = 0.0,
+    # ezz = 0.0,
+    # exy = 0.0,
+    # exz = 0.0,
+    # eyz = 0.0,
+    # deformation_matrix = None,
 ):
     """
     Calculate and store the relativistic corrected structure factors used for Bloch computations
@@ -156,22 +156,22 @@ def calculate_dynamical_structure_factors(
     hkl = np.vstack([xa.ravel(), ya.ravel(), za.ravel()])
     g_vec_all = lat_inv @ hkl
 
-    # strain
-    if (exx != 0.0 or \
-        eyy != 0.0 or \
-        ezz != 0.0 or \
-        exy != 0.0 or \
-        exz != 0.0 or \
-        eyz != 0.0) and \
-        deformation_matrix is None:
-            deformation_matrix = np.array([
-                [1+exx,     exy,    exz     ],
-                [exy,       1+eyy,  eyz     ],
-                [exz,       eyz,    1+ezz   ],
-            ])
-    if deformation_matrix is not None:
-        # Note that we need to take the matrix inverse, since we're operating in reciprocal space.
-        g_vec_all = np.linalg.inv(deformation_matrix) @ g_vec_all
+    # # strain
+    # if (exx != 0.0 or \
+    #     eyy != 0.0 or \
+    #     ezz != 0.0 or \
+    #     exy != 0.0 or \
+    #     exz != 0.0 or \
+    #     eyz != 0.0) and \
+    #     deformation_matrix is None:
+    #         deformation_matrix = np.array([
+    #             [1+exx,     exy,    exz     ],
+    #             [exy,       1+eyy,  eyz     ],
+    #             [exz,       eyz,    1+ezz   ],
+    #         ])
+    # if deformation_matrix is not None:
+    #     # Note that we need to take the matrix inverse, since we're operating in reciprocal space.
+    #     g_vec_all = np.linalg.inv(deformation_matrix) @ g_vec_all
 
     # Delete lattice vectors outside of k_max
     keep = np.linalg.norm(g_vec_all, axis=0) <= k_max

From 12be089db2965d6f22d2e3f009ae52be734127a5 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Wed, 2 Aug 2023 16:59:13 -0700
Subject: [PATCH 24/57] Making strain simpler

---
 py4DSTEM/process/diffraction/crystal.py       | 64 +++++--------------
 py4DSTEM/process/diffraction/crystal_bloch.py | 38 -----------
 2 files changed, 17 insertions(+), 85 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index a4ba23422..ee6d0fc1d 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -175,6 +175,14 @@ def get_strained_crystal(
                 [exz,       eyz,    1+ezz   ],
             ])
 
+        # copy crystal
+        crystal_strained = self.copy()
+        # crystal_strained = Crystal(
+        #     positions = self.positions.copy(),
+        #     numbers = self.numbers.copy(),
+        #     cell = self.cell.copy(),
+        #     )
+
         # new unit cell
         lat_new = self.lat_real @ deformation_matrix
         a_new = np.linalg.norm(lat_new[0,:])
@@ -186,17 +194,17 @@ def get_strained_crystal(
             lat_new[0,:]*lat_new[2,:])/a_new/c_new,-1,1)))
         gamma_new = np.rad2deg(np.arccos(np.clip(np.sum(
             lat_new[0,:]*lat_new[1,:])/a_new/b_new,-1,1)))
-        cell_new = np.array(
+        crystal_strained.cell = np.array(
             (a_new,b_new,c_new,alpha_new,beta_new,gamma_new)
         )
 
-        # Make new crystal
-        from py4DSTEM.process.diffraction import Crystal
-        crystal_strained = Crystal(
-            self.positions, 
-            self.numbers, 
-            cell_new,
-        )
+        # Update lattice
+        crystal_strained.lat_real = crystal_strained.lat_real @ deformation_matrix
+
+        # Inverse lattice, metric tensors
+        crystal_strained.metric_real = crystal_strained.lat_real @ crystal_strained.lat_real.T
+        crystal_strained.metric_inv = np.linalg.inv(crystal_strained.metric_real)
+        crystal_strained.lat_inv = crystal_strained.metric_inv @ crystal_strained.lat_real
 
         if return_deformation_matrix:
             return crystal_strained, deformation_matrix
@@ -449,13 +457,6 @@ def calculate_structure_factors(
         k_max: float = 2.0,
         tol_structure_factor: float = 1e-4,
         return_intensities: bool = False,
-        exx = 0.0,
-        eyy = 0.0,
-        ezz = 0.0,
-        exy = 0.0,
-        exz = 0.0,
-        eyz = 0.0,
-        deformation_matrix = None,
         ):
 
 
@@ -471,21 +472,7 @@ def calculate_structure_factors(
             tolerance for removing low-valued structure factors
         return_intensities: bool
             return the intensities and positions of all structure factor peaks.
-        exx: float
-            Strain in the x direction
-        eyy: float
-            Strain in the y direction
-        ezz: float
-            Strain in the z direction
-        exy: float
-            Shear in the x,y direction
-        exz: float
-            Shear in the x,z direction
-        eyz: float
-            Shear in the y,z direction
-        deformation_matrix: np.array
-            3x3 deformation matrix in real space. 
-
+        
         Returns
         --------
         (q_SF, I_SF)
@@ -524,23 +511,6 @@ def calculate_structure_factors(
         hkl = np.vstack([xa.ravel(), ya.ravel(), za.ravel()])
         # g_vec_all = self.lat_inv @ hkl
         g_vec_all =  (hkl.T @ self.lat_inv).T
-
-        # strain
-        if (exx != 0.0 or \
-            eyy != 0.0 or \
-            ezz != 0.0 or \
-            exy != 0.0 or \
-            exz != 0.0 or \
-            eyz != 0.0) and \
-            deformation_matrix is None:
-                deformation_matrix = np.array([
-                    [1+exx,     exy,    exz     ],
-                    [exy,       1+eyy,  eyz     ],
-                    [exz,       eyz,    1+ezz   ],
-                ])
-        if deformation_matrix is not None:
-            # Note that we need to take the matrix inverse, since we're operating in reciprocal space.
-            g_vec_all = np.linalg.inv(deformation_matrix) @ g_vec_all
                 
         # Delete lattice vectors outside of k_max
         keep = np.linalg.norm(g_vec_all, axis=0) <= self.k_max
diff --git a/py4DSTEM/process/diffraction/crystal_bloch.py b/py4DSTEM/process/diffraction/crystal_bloch.py
index 48918a91b..008135461 100644
--- a/py4DSTEM/process/diffraction/crystal_bloch.py
+++ b/py4DSTEM/process/diffraction/crystal_bloch.py
@@ -28,13 +28,6 @@ def calculate_dynamical_structure_factors(
     recompute_kinematic_structure_factors=True,
     g_vec_precision=None,
     verbose=True,
-    # exx = 0.0,
-    # eyy = 0.0,
-    # ezz = 0.0,
-    # exy = 0.0,
-    # exz = 0.0,
-    # eyz = 0.0,
-    # deformation_matrix = None,
 ):
     """
     Calculate and store the relativistic corrected structure factors used for Bloch computations
@@ -82,20 +75,6 @@ def calculate_dynamical_structure_factors(
 
         See WK_scattering_factors.py for details on the Weickenmeier-Kohl form factors.
 
-    exx: float
-        Strain in the x direction
-    eyy: float
-        Strain in the y direction
-    ezz: float
-        Strain in the z direction
-    exy: float
-        Shear in the x,y direction
-    exz: float
-        Shear in the x,z direction
-    eyz: float
-        Shear in the y,z direction
-    deformation_matrix: np.array
-        3x3 deformation matrix in real space.
 
     """
 
@@ -156,23 +135,6 @@ def calculate_dynamical_structure_factors(
     hkl = np.vstack([xa.ravel(), ya.ravel(), za.ravel()])
     g_vec_all = lat_inv @ hkl
 
-    # # strain
-    # if (exx != 0.0 or \
-    #     eyy != 0.0 or \
-    #     ezz != 0.0 or \
-    #     exy != 0.0 or \
-    #     exz != 0.0 or \
-    #     eyz != 0.0) and \
-    #     deformation_matrix is None:
-    #         deformation_matrix = np.array([
-    #             [1+exx,     exy,    exz     ],
-    #             [exy,       1+eyy,  eyz     ],
-    #             [exz,       eyz,    1+ezz   ],
-    #         ])
-    # if deformation_matrix is not None:
-    #     # Note that we need to take the matrix inverse, since we're operating in reciprocal space.
-    #     g_vec_all = np.linalg.inv(deformation_matrix) @ g_vec_all
-
     # Delete lattice vectors outside of k_max
     keep = np.linalg.norm(g_vec_all, axis=0) <= k_max
     hkl = hkl[:, keep]

From 295177878a3d25507bcacee3e77a0843ea0751b5 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Wed, 2 Aug 2023 17:03:20 -0700
Subject: [PATCH 25/57] simplifying code

---
 py4DSTEM/process/diffraction/crystal.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index ee6d0fc1d..03d709192 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -185,6 +185,8 @@ def get_strained_crystal(
 
         # new unit cell
         lat_new = self.lat_real @ deformation_matrix
+
+        # update cell params
         a_new = np.linalg.norm(lat_new[0,:])
         b_new = np.linalg.norm(lat_new[1,:])
         c_new = np.linalg.norm(lat_new[2,:])
@@ -199,7 +201,7 @@ def get_strained_crystal(
         )
 
         # Update lattice
-        crystal_strained.lat_real = crystal_strained.lat_real @ deformation_matrix
+        crystal_strained.lat_real = lat_new
 
         # Inverse lattice, metric tensors
         crystal_strained.metric_real = crystal_strained.lat_real @ crystal_strained.lat_real.T

From 63c4f15e2af642572c00b11468c3162c508f5487 Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Wed, 2 Aug 2023 17:21:43 -0700
Subject: [PATCH 26/57] Fixing the convention for strain

---
 py4DSTEM/process/diffraction/crystal.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index 03d709192..c4ae8c0db 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -170,9 +170,9 @@ def get_strained_crystal(
         # deformation matrix
         if deformation_matrix is None:
             deformation_matrix = np.array([
-                [1+exx,     exy,    exz     ],
-                [exy,       1+eyy,  eyz     ],
-                [exz,       eyz,    1+ezz   ],
+                [1.0+exx,   1.0*exy,    1.0*exz],
+                [1.0*exy,   1.0+eyy,    1.0*eyz],
+                [1.0*exz,   1.0*eyz,    1.0+ezz],
             ])
 
         # copy crystal

From f616f5f6f4febd037a36e95bb45834671d17f4bf Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Thu, 3 Aug 2023 13:24:28 -0700
Subject: [PATCH 27/57] Updating Crystal.__init__ to accept 3x3 cells

---
 py4DSTEM/process/diffraction/crystal.py | 96 ++++++++++++-------------
 1 file changed, 46 insertions(+), 50 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index c4ae8c0db..e2c6a1170 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -79,6 +79,7 @@ def __init__(
                 1 number: the lattice parameter for a cubic cell
                 3 numbers: the three lattice parameters for an orthorhombic cell
                 6 numbers: the a,b,c lattice parameters and ɑ,β,ɣ angles for any cell
+                3x3 array: row vectors containing the (u,v,w) lattice vectors.
 
         """
         # Initialize Crystal
@@ -93,7 +94,10 @@ def __init__(
         else:
             raise Exception("Number of positions and atomic numbers do not match")
 
-        # unit cell, as either [a a a 90 90 90], [a b c 90 90 90], or [a b c alpha beta gamma]
+        # unit cell, as one of:
+        # [a a a 90 90 90]
+        # [a b c 90 90 90]
+        # [a b c alpha beta gamma]
         cell = np.asarray(cell, dtype="float_")
         if np.size(cell) == 1:
             self.cell = np.hstack([cell, cell, cell, 90, 90, 90])
@@ -101,8 +105,20 @@ def __init__(
             self.cell = np.hstack([cell, 90, 90, 90])
         elif np.size(cell) == 6:
             self.cell = cell
+        elif np.shape(cell)[0] == 3 and np.shape(cell)[1] == 3:
+            self.lat_real = np.array(cell)
+            a = np.linalg.norm(self.lat_real[0,:])
+            b = np.linalg.norm(self.lat_real[1,:])
+            c = np.linalg.norm(self.lat_real[2,:])
+            alpha = np.rad2deg(np.arccos(np.clip(np.sum(
+                self.lat_real[1,:]*self.lat_real[2,:])/b/c,-1,1)))
+            beta  = np.rad2deg(np.arccos(np.clip(np.sum(
+                self.lat_real[0,:]*self.lat_real[2,:])/a/c,-1,1)))
+            gamma = np.rad2deg(np.arccos(np.clip(np.sum(
+                self.lat_real[0,:]*self.lat_real[1,:])/a/b,-1,1)))
+            self.cell = (a,b,c,alpha,beta,gamma)
         else:
-            raise Exception("Cell cannot contain " + np.size(cell) + " elements")
+            raise Exception("Cell cannot contain " + np.size(cell) + " entries")
         
         # pymatgen flag
         self.pymatgen_available = False
@@ -120,24 +136,26 @@ def copy(self):
         return crystal
 
     def calculate_lattice(self):
-        # calculate unit cell lattice vectors
-        a = self.cell[0]
-        b = self.cell[1]
-        c = self.cell[2]
-        alpha = np.deg2rad(self.cell[3])
-        beta = np.deg2rad(self.cell[4])
-        gamma = np.deg2rad(self.cell[5])
-        f = np.cos(beta) * np.cos(gamma) - np.cos(alpha)
-        vol = a*b*c*np.sqrt(1 \
-            + 2*np.cos(alpha)*np.cos(beta)*np.cos(gamma) \
-            - np.cos(alpha)**2 - np.cos(beta)**2 - np.cos(gamma)**2)
-        self.lat_real = np.array(
-            [
-                [a,               0,                 0],
-                [b*np.cos(gamma), b*np.sin(gamma),   0],
-                [c*np.cos(beta), -c*f/np.sin(gamma), vol/(a*b*np.sin(gamma))],
-            ]
-        )
+        
+        if not hasattr(self, 'lat_real'):
+            # calculate unit cell lattice vectors
+            a = self.cell[0]
+            b = self.cell[1]
+            c = self.cell[2]
+            alpha = np.deg2rad(self.cell[3])
+            beta = np.deg2rad(self.cell[4])
+            gamma = np.deg2rad(self.cell[5])
+            f = np.cos(beta) * np.cos(gamma) - np.cos(alpha)
+            vol = a*b*c*np.sqrt(1 \
+                + 2*np.cos(alpha)*np.cos(beta)*np.cos(gamma) \
+                - np.cos(alpha)**2 - np.cos(beta)**2 - np.cos(gamma)**2)
+            self.lat_real = np.array(
+                [
+                    [a,               0,                 0],
+                    [b*np.cos(gamma), b*np.sin(gamma),   0],
+                    [c*np.cos(beta), -c*f/np.sin(gamma), vol/(a*b*np.sin(gamma))],
+                ]
+            )
 
         # Inverse lattice, metric tensors
         self.metric_real = self.lat_real @ self.lat_real.T
@@ -164,7 +182,7 @@ def get_strained_crystal(
         """
         This method returns new Crystal class with strain applied. The directions of (x,y,z)
         are with respect to the default Crystal orientation, which can be checked with
-        print(Crystal.lat_real)
+        print(Crystal.lat_real) applied to the original Crystal.
         """
 
         # deformation matrix
@@ -175,38 +193,16 @@ def get_strained_crystal(
                 [1.0*exz,   1.0*eyz,    1.0+ezz],
             ])
 
-        # copy crystal
-        crystal_strained = self.copy()
-        # crystal_strained = Crystal(
-        #     positions = self.positions.copy(),
-        #     numbers = self.numbers.copy(),
-        #     cell = self.cell.copy(),
-        #     )
-
         # new unit cell
         lat_new = self.lat_real @ deformation_matrix
 
-        # update cell params
-        a_new = np.linalg.norm(lat_new[0,:])
-        b_new = np.linalg.norm(lat_new[1,:])
-        c_new = np.linalg.norm(lat_new[2,:])
-        alpha_new = np.rad2deg(np.arccos(np.clip(np.sum(
-            lat_new[1,:]*lat_new[2,:])/b_new/c_new,-1,1)))
-        beta_new  = np.rad2deg(np.arccos(np.clip(np.sum(
-            lat_new[0,:]*lat_new[2,:])/a_new/c_new,-1,1)))
-        gamma_new = np.rad2deg(np.arccos(np.clip(np.sum(
-            lat_new[0,:]*lat_new[1,:])/a_new/b_new,-1,1)))
-        crystal_strained.cell = np.array(
-            (a_new,b_new,c_new,alpha_new,beta_new,gamma_new)
-        )
-
-        # Update lattice
-        crystal_strained.lat_real = lat_new
-
-        # Inverse lattice, metric tensors
-        crystal_strained.metric_real = crystal_strained.lat_real @ crystal_strained.lat_real.T
-        crystal_strained.metric_inv = np.linalg.inv(crystal_strained.metric_real)
-        crystal_strained.lat_inv = crystal_strained.metric_inv @ crystal_strained.lat_real
+        # make new crystal class
+        from py4DSTEM.process.diffraction import Crystal
+        crystal_strained = Crystal(
+            positions = self.positions.copy(),
+            numbers = self.numbers.copy(),
+            cell = lat_new,
+            )
 
         if return_deformation_matrix:
             return crystal_strained, deformation_matrix

From 8a96817e012bdb9a6ac6ebb062b80c6db3dce31d Mon Sep 17 00:00:00 2001
From: Colin <cophus@gmail.com>
Date: Thu, 3 Aug 2023 13:26:52 -0700
Subject: [PATCH 28/57] Adding Alex's docstring suggestion

---
 py4DSTEM/process/diffraction/crystal.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index e2c6a1170..36ccf3c07 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -183,6 +183,8 @@ def get_strained_crystal(
         This method returns new Crystal class with strain applied. The directions of (x,y,z)
         are with respect to the default Crystal orientation, which can be checked with
         print(Crystal.lat_real) applied to the original Crystal.
+
+        Strains are given in fractional values, so exx = 0.01 is 1% strain along the x direction.
         """
 
         # deformation matrix

From 12d3ea20ecba3b04fa36dfb3f0de0c410b7d0360 Mon Sep 17 00:00:00 2001
From: alex-rakowski <alexrakowski90@gmail.com>
Date: Thu, 3 Aug 2023 14:09:25 -0700
Subject: [PATCH 29/57] removing white space from bloch

---
 py4DSTEM/process/diffraction/crystal_bloch.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal_bloch.py b/py4DSTEM/process/diffraction/crystal_bloch.py
index 008135461..6a3c9b1ac 100644
--- a/py4DSTEM/process/diffraction/crystal_bloch.py
+++ b/py4DSTEM/process/diffraction/crystal_bloch.py
@@ -74,8 +74,6 @@ def calculate_dynamical_structure_factors(
                                         substantial speedup at the cost of some reduced accuracy
 
         See WK_scattering_factors.py for details on the Weickenmeier-Kohl form factors.
-
-
     """
 
     assert method in (

From 177f92e70a8c1096153c3b9700e7d8c0e181b971 Mon Sep 17 00:00:00 2001
From: Steve Zeltmann <37132012+sezelt@users.noreply.github.com>
Date: Thu, 3 Aug 2023 17:17:05 -0400
Subject: [PATCH 30/57] format and remove crystal copy

---
 py4DSTEM/process/diffraction/crystal.py | 13 +------------
 1 file changed, 1 insertion(+), 12 deletions(-)

diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py
index 36ccf3c07..4d4d4a248 100644
--- a/py4DSTEM/process/diffraction/crystal.py
+++ b/py4DSTEM/process/diffraction/crystal.py
@@ -4,8 +4,6 @@
 import matplotlib.pyplot as plt
 from fractions import Fraction
 from typing import Union, Optional
-# import copy
-# from copy import deepcopy
 from scipy.optimize import curve_fit
 import sys
 
@@ -125,15 +123,6 @@ def __init__(
         
         # Calculate lattice parameters
         self.calculate_lattice()
-    
-    # TODO check for and copy all important attributes
-    def copy(self):
-        crystal = Crystal(
-            self.positions,
-            self.numbers,
-            self.cell,
-        )
-        return crystal
 
     def calculate_lattice(self):
         
@@ -204,7 +193,7 @@ def get_strained_crystal(
             positions = self.positions.copy(),
             numbers = self.numbers.copy(),
             cell = lat_new,
-            )
+        )
 
         if return_deformation_matrix:
             return crystal_strained, deformation_matrix

From 46541ee547dc892777736b78c25d1326081c64bd Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Sun, 6 Aug 2023 12:34:48 -0400
Subject: [PATCH 31/57] building out parameter sharing infra

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 81 +++++++++++++++----
 py4DSTEM/process/wholepatternfit/wpf.py       | 35 ++++----
 2 files changed, 88 insertions(+), 28 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 221334a81..fe1f8558a 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -46,7 +46,7 @@ def __init__(
     #         len(signature(self.func).parameters) == len(params) + 2
     #     ), f"The model function has the wrong number of arguments in its signature. It must be written as func(DP, param1, param2, ..., **kwargs). The current signature is {str(signature(self.func))}"
 
-    def func(self, DP: np.ndarray, *args, **kwargs) -> None:
+    def func(self, DP: np.ndarray, x, **kwargs) -> None:
         raise NotImplementedError()
 
     # Required signature for the Jacobian:
@@ -75,8 +75,8 @@ def __init__(
             self.set_params(initial_value, lower_bound, upper_bound)
 
         # Store a dummy offset. This must be set by WPF during setup
-        # This stores the index in the master Jacobian array corresponding to
-        # this parameter
+        # This stores the index in the master parameter and Jacobian arrays 
+        # corresponding to this parameter
         self.offset = np.nan
 
     def set_params(
@@ -95,6 +95,23 @@ def __str__(self):
     def __repr__(self):
         return f"Value: {self.initial_value} (Range: {self.lower_bound},{self.upper_bound})"
 
+class _BaseModel(WPFModelPrototype):
+    """
+    Model object used by the WPF class as a container for the global Parameters
+    """
+    def __init__(self, x0, y0, name="Global"):
+        params = {
+            "x center": Parameter(x0),
+            "y center": Parameter(y0)
+        }
+
+        super.__init__(name, params)
+
+    def func(self, DP: np.ndarray, x, **kwargs) -> None:
+        pass
+
+    def jacobian(self, J: np.ndarray, *args, **kwargs) -> None:
+        pass
 
 class DCBackground(WPFModelPrototype):
     def __init__(self, background_value=0.0, name="DC Background"):
@@ -102,8 +119,8 @@ def __init__(self, background_value=0.0, name="DC Background"):
 
         super().__init__(name, params)
 
-    def func(self, DP: np.ndarray, level, **kwargs) -> None:
-        DP += level
+    def func(self, DP: np.ndarray, x, **kwargs) -> None:
+        DP += x[self.params['DC Level'].offset]
 
     def jacobian(self, J: np.ndarray, *args, **kwargs):
         J[:, self.params['DC Level'].offset] = 1
@@ -131,13 +148,19 @@ def __init__(
 
         super().__init__(name, params)
 
-    def global_center_func(self, DP: np.ndarray, sigma, level, **kwargs) -> None:
+    def global_center_func(self, DP: np.ndarray, x:np.ndarray, **kwargs) -> None:
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['intensity'].offset]
+
         DP += level * np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
 
     def global_center_jacobian(
-        self, J: np.ndarray, sigma, level, **kwargs
+        self, J: np.ndarray, x:np.ndarray, **kwargs
     ) -> None:
 
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['intensity'].offset]
+
         exp_expr = np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
 
         # dF/d(global_x0)
@@ -156,16 +179,25 @@ def global_center_jacobian(
         # dF/d(level)
         J[:, self.params['intensity'].offset] = exp_expr.ravel()
 
-    def local_center_func(self, DP: np.ndarray, sigma, level, x0, y0, **kwargs) -> None:
+    def local_center_func(self, DP: np.ndarray, x:np.ndarray, **kwargs) -> None:
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['intensity'].offset]
+        x0 = x[self.params['x center'].offset]
+        y0 = x[self.params['y center'].offset]
         DP += level * np.exp(
             ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             / (-2 * sigma**2)
         )
 
     def local_center_jacobian(
-        self, J: np.ndarray, sigma, level, x0, y0, **kwargs
+        self, J: np.ndarray, x:np.ndarray, **kwargs
     ) -> None:
 
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['intensity'].offset]
+        x0 = x[self.params['x center'].offset]
+        y0 = x[self.params['y center'].offset]
+
         # dF/s(sigma)
         J[:, self.params['sigma'].offset] = (
             level
@@ -234,14 +266,21 @@ def __init__(
         super().__init__(name, params)
 
     def global_center_func(
-        self, DP: np.ndarray, radius, sigma, level, **kwargs
+        self, DP: np.ndarray, x: np.ndarray, **kwargs
     ) -> None:
+        radius = x[self.params['radius'].offset]
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['level'].offset]
         DP += level * np.exp((kwargs["global_r"] - radius) ** 2 / (-2 * sigma**2))
 
     def global_center_jacobian(
-        self, J: np.ndarray, radius, sigma, level, **kwargs
+        self, J: np.ndarray, x:np.ndarray, **kwargs
     ) -> None:
 
+        radius = x[self.params['radius'].offset]
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['level'].offset]
+
         local_r = radius - kwargs["global_r"]
         clipped_r = np.maximum(local_r, 0.1)
 
@@ -277,15 +316,29 @@ def global_center_jacobian(
         J[:, self.params['intensity'].offset] = exp_expr.ravel()
 
     def local_center_func(
-        self, DP: np.ndarray, radius, sigma, level, x0, y0, **kwargs
+        self, DP: np.ndarray, x:np.ndarray, **kwargs
     ) -> None:
+
+        radius = x[self.params['radius'].offset]
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['level'].offset]
+        x0 = x[self.params['x center'].offset]
+        y0 = x[self.params['y center'].offset]
+
         local_r = np.hypot(kwargs["xArray"] - x0, kwargs["yArray"] - y0)
         DP += level * np.exp((local_r - radius) ** 2 / (-2 * sigma**2))
 
     def local_center_jacobian(
-        self, J: np.ndarray, radius, sigma, level, x0, y0, **kwargs
+        self, J: np.ndarray, x:np.ndarray, **kwargs
     ) -> None:
         return NotImplementedError()
+
+        radius = x[self.params['radius'].offset]
+        sigma = x[self.params['sigma'].offset]
+        level = x[self.params['level'].offset]
+        x0 = x[self.params['x center'].offset]
+        y0 = x[self.params['y center'].offset]
+
         # dF/d(radius)
 
         # dF/s(sigma)
@@ -447,7 +500,7 @@ def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         for i, (u, v) in enumerate(zip(self.u_inds, self.v_inds)):
             x = x0 + (u * ux) + (v * vx)
             y = y0 + (u * uy) + (v * vy)
-            # if (x > 0) & (x < kwargs["Q_Nx"]) & (y > 0) & (y < kwargs["Q_Nx"]):
+
             DP += args[i + 6] / (
                 1.0
                 + np.exp(
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index b9035a736..a6b83d5d2 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -1,6 +1,6 @@
 from py4DSTEM import DataCube, RealSlice
 from emdfile import tqdmnd
-from py4DSTEM.process.wholepatternfit.wp_models import WPFModelPrototype
+from py4DSTEM.process.wholepatternfit.wp_models import WPFModelPrototype, _BaseModel
 
 from typing import Optional
 import numpy as np
@@ -11,10 +11,9 @@
 from matplotlib.gridspec import GridSpec
 import warnings
 
-# from multiprocessing import Pool
-# import multiprocess as mp
 from mpire import WorkerPool
 
+__all__ = ["WholePatternFit"]
 
 class WholePatternFit:
 
@@ -91,20 +90,28 @@ def __init__(
             x0 = np.array(x0)
             y0 = np.array(y0)
             if x0.size == 2:
-                self.global_xy0_lb = np.array([x0[0] - x0[1], y0[0] - y0[1]])
-                self.global_xy0_ub = np.array([x0[0] + x0[1], y0[0] + y0[1]])
+                global_xy0_lb = np.array([x0[0] - x0[1], y0[0] - y0[1]])
+                global_xy0_ub = np.array([x0[0] + x0[1], y0[0] + y0[1]])
             elif x0.size == 3:
-                self.global_xy0_lb = np.array([x0[1], y0[1]])
-                self.global_xy0_ub = np.array([x0[2], y0[2]])
+                global_xy0_lb = np.array([x0[1], y0[1]])
+                global_xy0_ub = np.array([x0[2], y0[2]])
             else:
-                self.global_xy0_lb = np.array([0.0, 0.0])
-                self.global_xy0_ub = np.array([datacube.Q_Nx, datacube.Q_Ny])
+                global_xy0_lb = np.array([0.0, 0.0])
+                global_xy0_ub = np.array([datacube.Q_Nx, datacube.Q_Ny])
             x0 = x0[0]
             y0 = y0[0]
 
         else:
-            self.global_xy0_lb = np.array([0.0, 0.0])
-            self.global_xy0_ub = np.array([datacube.Q_Nx, datacube.Q_Ny])
+            global_xy0_lb = np.array([0.0, 0.0])
+            global_xy0_ub = np.array([datacube.Q_Nx, datacube.Q_Ny])
+
+        # The WPF object holds a special Model that manages the shareable center coordinates
+        self.global_params = _BaseModel(
+            x0 = (x0, global_xy0_lb[0], global_xy0_ub[0]),
+            y0 = (y0, global_xy0_lb[1], global_xy0_ub[1])
+        )
+        # TODO: remove special cases for global/local center in the Models
+        # Needs an efficient way to handle calculation of q_r
 
         # set up the global arguments
         self._setup_static_data(x0,y0)
@@ -217,9 +224,9 @@ def fit_to_mean_CBED(self, **fit_opts):
 
     def fit_single_pattern(
         self, 
-        rx,
-        ry,
-        resume = False, 
+        data: np.ndarray,
+        resume:bool = False, 
+        restart_data:np.ndarray = None,
         **fit_opts
         ):
         """

From fa90c5409dd1cdd95a7adae62a721effcd876a1f Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Sun, 6 Aug 2023 12:47:57 -0400
Subject: [PATCH 32/57] add model type flags and format

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 205 +++++++++---------
 py4DSTEM/process/wholepatternfit/wpf.py       |  85 ++++----
 py4DSTEM/process/wholepatternfit/wpf_viz.py   |  61 +++---
 3 files changed, 170 insertions(+), 181 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index fe1f8558a..3b966ce6a 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -1,10 +1,24 @@
 from inspect import signature
 from typing import Optional
+from enum import Flag, auto
 import numpy as np
 
 from pdb import set_trace
 
 
+class WPFModelType(Flag):
+    """
+    Flags to signify capabilities and other semantics of a Model
+    """
+
+    BACKGROUND = auto()
+
+    AMORPHOUS = auto()
+    LATTICE = auto()
+
+    DUMMY = auto()
+
+
 class WPFModelPrototype:
     """
     Prototype class for a compent of a whole-pattern model.
@@ -29,11 +43,7 @@ class WPFModelPrototype:
                         • keyword arguments. this is to provide some pre-computed information for convenience
     """
 
-    def __init__(
-        self,
-        name: str,
-        params: dict,
-    ):
+    def __init__(self, name: str, params: dict, model_type=WPFModelType.DUMMY):
         self.name = name
         self.params = params
 
@@ -41,6 +51,8 @@ def __init__(
 
         self.hasJacobian = getattr(self, "jacobian", None) is not None
 
+        self.model_type = model_type
+
     #     # check the function obeys the spec
     #     assert (
     #         len(signature(self.func).parameters) == len(params) + 2
@@ -61,33 +73,32 @@ def __init__(
         initial_value,
         lower_bound: Optional[float] = None,
         upper_bound: Optional[float] = None,
-        ):
-
+    ):
         if hasattr(initial_value, "__iter__"):
             if len(initial_value) == 2:
                 initial_value = (
                     initial_value[0],
-                    initial_value[0]-initial_value[1],
-                    initial_value[0]+initial_value[1],
-                    )
+                    initial_value[0] - initial_value[1],
+                    initial_value[0] + initial_value[1],
+                )
             self.set_params(*initial_value)
         else:
             self.set_params(initial_value, lower_bound, upper_bound)
 
         # Store a dummy offset. This must be set by WPF during setup
-        # This stores the index in the master parameter and Jacobian arrays 
+        # This stores the index in the master parameter and Jacobian arrays
         # corresponding to this parameter
         self.offset = np.nan
 
     def set_params(
-        self, 
-        initial_value, 
-        lower_bound, 
+        self,
+        initial_value,
+        lower_bound,
         upper_bound,
-        ):
+    ):
         self.initial_value = initial_value
         self.lower_bound = lower_bound if lower_bound is not None else -np.inf
-        self.upper_bound = upper_bound if upper_bound is not None else  np.inf
+        self.upper_bound = upper_bound if upper_bound is not None else np.inf
 
     def __str__(self):
         return f"Value: {self.initial_value} (Range: {self.lower_bound},{self.upper_bound})"
@@ -95,17 +106,16 @@ def __str__(self):
     def __repr__(self):
         return f"Value: {self.initial_value} (Range: {self.lower_bound},{self.upper_bound})"
 
+
 class _BaseModel(WPFModelPrototype):
     """
     Model object used by the WPF class as a container for the global Parameters
     """
-    def __init__(self, x0, y0, name="Global"):
-        params = {
-            "x center": Parameter(x0),
-            "y center": Parameter(y0)
-        }
 
-        super.__init__(name, params)
+    def __init__(self, x0, y0, name="Globals"):
+        params = {"x center": Parameter(x0), "y center": Parameter(y0)}
+
+        super().__init__(name, params, model_type=WPFModelType.DUMMY)
 
     def func(self, DP: np.ndarray, x, **kwargs) -> None:
         pass
@@ -113,17 +123,18 @@ def func(self, DP: np.ndarray, x, **kwargs) -> None:
     def jacobian(self, J: np.ndarray, *args, **kwargs) -> None:
         pass
 
+
 class DCBackground(WPFModelPrototype):
     def __init__(self, background_value=0.0, name="DC Background"):
         params = {"DC Level": Parameter(background_value)}
 
-        super().__init__(name, params)
+        super().__init__(name, params, model_type=WPFModelType.BACKGROUND)
 
     def func(self, DP: np.ndarray, x, **kwargs) -> None:
-        DP += x[self.params['DC Level'].offset]
+        DP += x[self.params["DC Level"].offset]
 
     def jacobian(self, J: np.ndarray, *args, **kwargs):
-        J[:, self.params['DC Level'].offset] = 1
+        J[:, self.params["DC Level"].offset] = 1
 
 
 class GaussianBackground(WPFModelPrototype):
@@ -146,20 +157,17 @@ def __init__(
             self.func = self.local_center_func
             self.jacobian = self.local_center_jacobian
 
-        super().__init__(name, params)
+        super().__init__(name, params, model_type=WPFModelType.BACKGROUND)
 
-    def global_center_func(self, DP: np.ndarray, x:np.ndarray, **kwargs) -> None:
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['intensity'].offset]
+    def global_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["intensity"].offset]
 
         DP += level * np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
 
-    def global_center_jacobian(
-        self, J: np.ndarray, x:np.ndarray, **kwargs
-    ) -> None:
-
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['intensity'].offset]
+    def global_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["intensity"].offset]
 
         exp_expr = np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
 
@@ -174,32 +182,31 @@ def global_center_jacobian(
         ).ravel()
 
         # dF/s(sigma)
-        J[:, self.params['sigma'].offset] = (level * kwargs["global_r"] ** 2 * exp_expr / sigma**3).ravel()
+        J[:, self.params["sigma"].offset] = (
+            level * kwargs["global_r"] ** 2 * exp_expr / sigma**3
+        ).ravel()
 
         # dF/d(level)
-        J[:, self.params['intensity'].offset] = exp_expr.ravel()
+        J[:, self.params["intensity"].offset] = exp_expr.ravel()
 
-    def local_center_func(self, DP: np.ndarray, x:np.ndarray, **kwargs) -> None:
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['intensity'].offset]
-        x0 = x[self.params['x center'].offset]
-        y0 = x[self.params['y center'].offset]
+    def local_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["intensity"].offset]
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
         DP += level * np.exp(
             ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             / (-2 * sigma**2)
         )
 
-    def local_center_jacobian(
-        self, J: np.ndarray, x:np.ndarray, **kwargs
-    ) -> None:
-
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['intensity'].offset]
-        x0 = x[self.params['x center'].offset]
-        y0 = x[self.params['y center'].offset]
+    def local_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["intensity"].offset]
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
 
         # dF/s(sigma)
-        J[:, self.params['sigma'].offset] = (
+        J[:, self.params["sigma"].offset] = (
             level
             * ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             * np.exp(
@@ -210,13 +217,13 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(level)
-        J[:, self.params['intensity'].offset] = np.exp(
+        J[:, self.params["intensity"].offset] = np.exp(
             ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             / (-2 * sigma**2)
         ).ravel()
 
         # dF/d(x0)
-        J[:, self.params['x center'].offset] = (
+        J[:, self.params["x center"].offset] = (
             level
             * (kwargs["xArray"] - x0)
             * np.exp(
@@ -227,7 +234,7 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(y0)
-        J[:, self.params['y center'].offset] = (
+        J[:, self.params["y center"].offset] = (
             level
             * (kwargs["yArray"] - y0)
             * np.exp(
@@ -263,23 +270,18 @@ def __init__(
             self.func = self.local_center_func
             self.jacobian = self.local_center_jacobian
 
-        super().__init__(name, params)
+        super().__init__(name, params, model_type=WPFModelType.AMORPHOUS)
 
-    def global_center_func(
-        self, DP: np.ndarray, x: np.ndarray, **kwargs
-    ) -> None:
-        radius = x[self.params['radius'].offset]
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['level'].offset]
+    def global_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
+        radius = x[self.params["radius"].offset]
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["level"].offset]
         DP += level * np.exp((kwargs["global_r"] - radius) ** 2 / (-2 * sigma**2))
 
-    def global_center_jacobian(
-        self, J: np.ndarray, x:np.ndarray, **kwargs
-    ) -> None:
-
-        radius = x[self.params['radius'].offset]
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['level'].offset]
+    def global_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
+        radius = x[self.params["radius"].offset]
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["level"].offset]
 
         local_r = radius - kwargs["global_r"]
         clipped_r = np.maximum(local_r, 0.1)
@@ -305,44 +307,41 @@ def global_center_jacobian(
         ).ravel()
 
         # dF/d(radius)
-        J[:, self.params['radius'].offset] += (-1.0 * level * exp_expr * local_r / (sigma**2)).ravel()
+        J[:, self.params["radius"].offset] += (
+            -1.0 * level * exp_expr * local_r / (sigma**2)
+        ).ravel()
 
         # dF/d(sigma)
-        J[:, self.params['sigma'].offset] = (
-            level * local_r ** 2 * exp_expr / sigma**3
+        J[:, self.params["sigma"].offset] = (
+            level * local_r**2 * exp_expr / sigma**3
         ).ravel()
 
         # dF/d(intensity)
-        J[:, self.params['intensity'].offset] = exp_expr.ravel()
-
-    def local_center_func(
-        self, DP: np.ndarray, x:np.ndarray, **kwargs
-    ) -> None:
+        J[:, self.params["intensity"].offset] = exp_expr.ravel()
 
-        radius = x[self.params['radius'].offset]
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['level'].offset]
-        x0 = x[self.params['x center'].offset]
-        y0 = x[self.params['y center'].offset]
+    def local_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
+        radius = x[self.params["radius"].offset]
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["level"].offset]
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
 
         local_r = np.hypot(kwargs["xArray"] - x0, kwargs["yArray"] - y0)
         DP += level * np.exp((local_r - radius) ** 2 / (-2 * sigma**2))
 
-    def local_center_jacobian(
-        self, J: np.ndarray, x:np.ndarray, **kwargs
-    ) -> None:
+    def local_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
         return NotImplementedError()
 
-        radius = x[self.params['radius'].offset]
-        sigma = x[self.params['sigma'].offset]
-        level = x[self.params['level'].offset]
-        x0 = x[self.params['x center'].offset]
-        y0 = x[self.params['y center'].offset]
+        radius = x[self.params["radius"].offset]
+        sigma = x[self.params["sigma"].offset]
+        level = x[self.params["level"].offset]
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
 
         # dF/d(radius)
 
         # dF/s(sigma)
-        J[:, self.params['sigma'].offset] = (
+        J[:, self.params["sigma"].offset] = (
             level
             * ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             * np.exp(
@@ -353,13 +352,13 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(level)
-        J[:, self.params['intensity'].offset] = np.exp(
+        J[:, self.params["intensity"].offset] = np.exp(
             ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
             / (-2 * sigma**2)
         ).ravel()
 
         # dF/d(x0)
-        J[:, self.params['x center'].offset] = (
+        J[:, self.params["x center"].offset] = (
             level
             * (kwargs["xArray"] - x0)
             * np.exp(
@@ -370,7 +369,7 @@ def local_center_jacobian(
         ).ravel()
 
         # dF/d(y0)
-        J[:, self.params['y center'].offset] = (
+        J[:, self.params["y center"].offset] = (
             level
             * (kwargs["yArray"] - y0)
             * np.exp(
@@ -472,7 +471,7 @@ def __init__(
         if refine_width:
             params["edge width"] = Parameter(disk_width)
 
-        super().__init__(name, params)
+        super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
     def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         # copy the global centers in the right place for the local center generator
@@ -481,7 +480,6 @@ def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         )
 
     def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-
         x0 = args[0]
         y0 = args[1]
         ux = args[2]
@@ -508,19 +506,20 @@ def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
                         4
                         * (
                             np.sqrt(
-                                (kwargs["xArray"] - x) ** 2 + (kwargs["yArray"] - y) ** 2
+                                (kwargs["xArray"] - x) ** 2
+                                + (kwargs["yArray"] - y) ** 2
                             )
                             - disk_radius
                         )
-                        / disk_width, 
-                    20)
+                        / disk_width,
+                        20,
+                    )
                 )
             )
 
     def global_center_jacobian(
         self, J: np.ndarray, *args, offset: int, **kwargs
     ) -> None:
-
         x0 = kwargs["global_x0"]
         y0 = kwargs["global_y0"]
         r = np.maximum(5e-1, kwargs["global_r"])
@@ -631,7 +630,6 @@ def __init__(
         name="Complex Overlapped Disk Lattice",
         verbose=False,
     ):
-
         params = {}
 
         # if global_center:
@@ -696,7 +694,7 @@ def __init__(
 
         self.func = self.global_center_func
 
-        super().__init__(name, params)
+        super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
     def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         # copy the global centers in the right place for the local center generator
@@ -705,7 +703,6 @@ def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         )
 
     def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-
         x0 = args[0]
         y0 = args[1]
         ux = args[2]
@@ -749,7 +746,6 @@ def __init__(
         name="Custom Kernel Disk Lattice",
         verbose=False,
     ):
-
         params = {}
 
         # if global_center:
@@ -808,7 +804,7 @@ def __init__(
 
         self.func = self.global_center_func
 
-        super().__init__(name, params)
+        super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
     def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         # copy the global centers in the right place for the local center generator
@@ -817,7 +813,6 @@ def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         )
 
     def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-
         x0 = args[0]
         y0 = args[1]
         ux = args[2]
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index a6b83d5d2..9651748f0 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -15,8 +15,8 @@
 
 __all__ = ["WholePatternFit"]
 
-class WholePatternFit:
 
+class WholePatternFit:
     from py4DSTEM.process.wholepatternfit.wpf_viz import (
         show_model_grid,
         show_lattice_points,
@@ -76,7 +76,7 @@ def __init__(
             meanCBED if meanCBED is not None else np.mean(datacube.data, axis=(0, 1))
         )
         # Global scaling parameter
-        self.intensity_scale = 1/np.mean(self.meanCBED)
+        self.intensity_scale = 1 / np.mean(self.meanCBED)
 
         self.mask = mask if mask is not None else np.ones_like(self.meanCBED)
 
@@ -106,15 +106,15 @@ def __init__(
             global_xy0_ub = np.array([datacube.Q_Nx, datacube.Q_Ny])
 
         # The WPF object holds a special Model that manages the shareable center coordinates
-        self.global_params = _BaseModel(
-            x0 = (x0, global_xy0_lb[0], global_xy0_ub[0]),
-            y0 = (y0, global_xy0_lb[1], global_xy0_ub[1])
+        self.coordinate_model = _BaseModel(
+            x0=(x0, global_xy0_lb[0], global_xy0_ub[0]),
+            y0=(y0, global_xy0_lb[1], global_xy0_ub[1]),
         )
         # TODO: remove special cases for global/local center in the Models
         # Needs an efficient way to handle calculation of q_r
 
         # set up the global arguments
-        self._setup_static_data(x0,y0)
+        self._setup_static_data(x0, y0)
 
         self.fit_power = fit_power
 
@@ -138,13 +138,11 @@ def add_model_list(self, model_list):
             self.add_model(m)
 
     def generate_initial_pattern(self):
-
         # update parameters:
         self._scrape_model_params()
         return self._pattern(self.x0, self.static_data.copy()) / self.intensity_scale
 
     def fit_to_mean_CBED(self, **fit_opts):
-
         # first make sure we have the latest parameters
         self._scrape_model_params()
 
@@ -221,14 +219,13 @@ def fit_to_mean_CBED(self, **fit_opts):
 
         return opt
 
-
     def fit_single_pattern(
-        self, 
+        self,
         data: np.ndarray,
-        resume:bool = False, 
-        restart_data:np.ndarray = None,
-        **fit_opts
-        ):
+        resume: bool = False,
+        restart_data: np.ndarray = None,
+        **fit_opts,
+    ):
         """
         Apply model fitting to one pattern.
 
@@ -246,7 +243,7 @@ def fit_single_pattern(
         Returns
         --------
         fit_coefs: np.array
-            Fitted coefficients 
+            Fitted coefficients
         fit_metrics: np.array
             Fitting metrics
 
@@ -313,19 +310,18 @@ def fit_single_pattern(
             ]
         except:
             fit_coefs = x0
-            fit_metrics_single = [0,0,0,0]
+            fit_metrics_single = [0, 0, 0, 0]
 
         return fit_coefs, fit_metrics_single
 
-
     def fit_all_patterns(
-        self, 
-        resume = False, 
-        real_space_mask = None,
-        num_jobs = None,
-        show_fit_metrics = True,
-        **fit_opts
-        ):
+        self,
+        resume=False,
+        real_space_mask=None,
+        num_jobs=None,
+        show_fit_metrics=True,
+        **fit_opts,
+    ):
         """
         Apply model fitting to all patterns.
 
@@ -375,8 +371,10 @@ def fit_all_patterns(
         # Loop over probe positions
         if num_jobs is None:
             for rx, ry in tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny):
-                if real_space_mask is not None and real_space_mask[rx,ry] == True:
-                    current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
+                if real_space_mask is not None and real_space_mask[rx, ry] == True:
+                    current_pattern = (
+                        self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
+                    )
                     shared_data = self.static_data.copy()
                     x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0.copy()
 
@@ -410,7 +408,7 @@ def fit_all_patterns(
                         ]
                     except:
                         fit_data_single = x0
-                        fit_metrics_single = [0,0,0,0]
+                        fit_metrics_single = [0, 0, 0, 0]
 
                     fit_data[:, rx, ry] = fit_data_single
                     fit_metrics[:, rx, ry] = fit_metrics_single
@@ -418,32 +416,32 @@ def fit_all_patterns(
         else:
             # Get list of probe positions
             if real_space_mask is not None:
-                xa,ya = np.where(real_space_mask)
+                xa, ya = np.where(real_space_mask)
             else:
-                xa,ya = np.meshgrid(
+                xa, ya = np.meshgrid(
                     np.arange(self.datacube.Rshape[0]),
                     np.arange(self.datacube.Rshape[1]),
-                    indexing = 'ij',                    
-                    )
+                    indexing="ij",
+                )
                 xa = xa.ravel()
                 ya = ya.ravel()
-            xy = np.vstack((xa,ya))
+            xy = np.vstack((xa, ya))
 
-            fit_inputs = [default_opts]*xy.shape[1]
+            fit_inputs = [default_opts] * xy.shape[1]
             for a0 in range(xy.shape[1]):
-                fit_inputs[a0]['rx'] = xy[0,a0]
-                fit_inputs[a0]['ry'] = xy[1,a0]
+                fit_inputs[a0]["rx"] = xy[0, a0]
+                fit_inputs[a0]["ry"] = xy[1, a0]
 
-            with WorkerPool(n_jobs = num_jobs) as pool:
+            with WorkerPool(n_jobs=num_jobs) as pool:
                 results = pool.map(
-                    self.fit_single_pattern, 
+                    self.fit_single_pattern,
                     fit_inputs,
                     progress_bar=True,
-                    )
+                )
 
             for a0 in range(xy.shape[1]):
-                fit_data[:, xy[0,a0], xy[1,a0]] = results[a0][0]
-                fit_metrics[:, xy[0,a0], xy[1,a0]] = results[a0][1]
+                fit_data[:, xy[0, a0], xy[1, a0]] = results[a0][0]
+                fit_metrics[:, xy[0, a0], xy[1, a0]] = results[a0][1]
 
         # Convert to RealSlices
         model_names = []
@@ -498,7 +496,7 @@ def get_lattice_maps(self):
         ]
 
         g_maps = []
-        for (i, l) in lattices:
+        for i, l in lattices:
             param_list = list(l.params.keys())
             lattice_offset = param_list.index("ux")
             data_offset = self.model_param_inds[i] + 2 + lattice_offset
@@ -515,7 +513,7 @@ def get_lattice_maps(self):
 
         return g_maps
 
-    def _setup_static_data(self,x0,y0):
+    def _setup_static_data(self, x0, y0):
         self.static_data = {}
 
         xArray, yArray = np.mgrid[0 : self.datacube.Q_Nx, 0 : self.datacube.Q_Ny]
@@ -533,7 +531,6 @@ def _setup_static_data(self,x0,y0):
         )
 
     def _pattern_error(self, x, current_pattern, shared_data):
-
         DP = np.zeros((self.datacube.Q_Nx, self.datacube.Q_Ny))
 
         shared_data["global_x0"] = x[0]
@@ -557,7 +554,6 @@ def _pattern_error(self, x, current_pattern, shared_data):
         return DP.ravel()
 
     def _pattern(self, x, shared_data):
-
         DP = np.zeros((self.datacube.Q_Nx, self.datacube.Q_Ny))
 
         shared_data["global_x0"] = x[0]
@@ -592,7 +588,6 @@ def _jacobian(self, x, current_pattern, shared_data):
         return J * self.mask.ravel()[:, np.newaxis]
 
     def _scrape_model_params(self):
-
         self.x0 = np.zeros((self.nParams + 2,))
         self.upper_bound = np.zeros_like(self.x0)
         self.lower_bound = np.zeros_like(self.x0)
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index 2820d1056..c02560653 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -5,6 +5,7 @@
 import matplotlib.colors as mpl_c
 from matplotlib.gridspec import GridSpec
 
+
 def show_model_grid(self, x=None, **plot_kwargs):
     if x is None:
         x = self.mean_CBED_fit.x
@@ -42,37 +43,34 @@ def show_model_grid(self, x=None, **plot_kwargs):
         # Determine if text color should be white or black
         int_range = np.array((np.min(DP), np.max(DP)))
         if int_range[0] != int_range[1]:
-            r = (np.mean(DP[:DP.shape[0]//10,:]) - int_range[0]) / (int_range[1] - int_range[0])
+            r = (np.mean(DP[: DP.shape[0] // 10, :]) - int_range[0]) / (
+                int_range[1] - int_range[0]
+            )
             if r < 0.5:
-                color = 'w'
+                color = "w"
             else:
-                color = 'k'
+                color = "k"
         else:
-            color = 'w'
+            color = "w"
 
         a.text(
-            0.5, 
-            0.92, 
-            m.name, 
-            transform = a.transAxes, 
-            ha = "center", 
-            va = "center",
-            color = color)
+            0.5,
+            0.92,
+            m.name,
+            transform=a.transAxes,
+            ha="center",
+            va="center",
+            color=color,
+        )
     for a in ax.flat:
         a.axis("off")
 
     plt.show()
 
+
 def show_lattice_points(
-    self, 
-    im = None,
-    vmin = None,
-    vmax = None,
-    power = None,
-    returnfig=False, 
-    *args, 
-    **kwargs
-    ):
+    self, im=None, vmin=None, vmax=None, power=None, returnfig=False, *args, **kwargs
+):
     """
     Plotting utility to show the initial lattice points.
     """
@@ -85,16 +83,16 @@ def show_lattice_points(
     fig, ax = plt.subplots(*args, **kwargs)
     if vmin is None and vmax is None:
         ax.matshow(
-            im**power, 
+            im**power,
             cmap="gray",
-            )
+        )
     else:
         ax.matshow(
             im**power,
-            vmin = vmin,
-            vmax = vmax, 
+            vmin=vmin,
+            vmax=vmax,
             cmap="gray",
-            )
+        )
 
     for m in self.model:
         if "Lattice" in m.name:
@@ -109,17 +107,18 @@ def show_lattice_points(
             spots[:, 1] += self.static_data["global_y0"]
 
             ax.scatter(
-                spots[:, 1], 
-                spots[:, 0], 
-                s = 100,
-                marker="x", 
+                spots[:, 1],
+                spots[:, 0],
+                s=100,
+                marker="x",
                 label=m.name,
-                )
+            )
 
     ax.legend()
 
     return (fig, ax) if returnfig else plt.show()
 
+
 def show_fit_metrics(self, returnfig=False, **subplots_kwargs):
     assert hasattr(self, "fit_metrics"), "Please run fitting first!"
 
@@ -167,4 +166,4 @@ def show_fit_metrics(self, returnfig=False, **subplots_kwargs):
 
     fig.set_facecolor("w")
 
-    return (fig, ax) if returnfig else plt.show()
\ No newline at end of file
+    return (fig, ax) if returnfig else plt.show()

From 4e7a7d8b6b8e0cc946299629abd90f4c2558d4f2 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Sun, 6 Aug 2023 13:27:19 -0400
Subject: [PATCH 33/57] starting moire model

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 24 ++++++++++---------
 py4DSTEM/process/wholepatternfit/wpf.py       | 19 +++++++--------
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 3b966ce6a..4802ced19 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -16,7 +16,8 @@ class WPFModelType(Flag):
     AMORPHOUS = auto()
     LATTICE = auto()
 
-    DUMMY = auto()
+    DUMMY = auto() # Model has no direct contribution to pattern
+    META = auto() # Model depends on multiple sub-Models (triggers extra linking on inclusion)
 
 
 class WPFModelPrototype:
@@ -140,6 +141,7 @@ def jacobian(self, J: np.ndarray, *args, **kwargs):
 class GaussianBackground(WPFModelPrototype):
     def __init__(
         self,
+        WPF,
         sigma,
         intensity,
         global_center=True,
@@ -149,6 +151,8 @@ def __init__(
     ):
         params = {"sigma": Parameter(sigma), "intensity": Parameter(intensity)}
         if global_center:
+            params['x center'] = WPF.coordinate_model.params['x center']
+            params['y center'] = WPF.coordinate_model.params['y center']
             self.func = self.global_center_func
             self.jacobian = self.global_center_jacobian
         else:
@@ -544,7 +548,6 @@ def global_center_jacobian(
 
             disk_intensity = args[i + 6]
 
-            # if (x > 0) & (x < kwargs["Q_Nx"]) & (y > 0) & (y < kwargs["Q_Nx"]):
             r_disk = np.maximum(
                 5e-1,
                 np.sqrt((kwargs["xArray"] - x) ** 2 + (kwargs["yArray"] - y) ** 2),
@@ -572,11 +575,6 @@ def global_center_jacobian(
                 / ((1.0 + top_exp) ** 2 * disk_width * r)
             ).ravel()
 
-            # because... reasons, sometimes we get NaN
-            # very far from the disk center. let's zero those:
-            # dx[np.isnan(dx)] = 0.0
-            # dy[np.isnan(dy)] = 0.0
-
             # insert global positional derivatives
             J[:, 0] += disk_intensity * dx
             J[:, 1] += disk_intensity * dy
@@ -589,7 +587,6 @@ def global_center_jacobian(
 
             # insert intensity derivative
             dI = (mask * (1.0 / (1.0 + top_exp))).ravel()
-            # dI[np.isnan(dI)] = 0.0
             J[:, offset + i + 4] = dI
 
             # insert disk radius derivative
@@ -597,7 +594,6 @@ def global_center_jacobian(
                 dR = (
                     4.0 * args[i + 4] * top_exp / (disk_width * (1.0 + top_exp) ** 2)
                 ).ravel()
-                # dR[np.isnan(dR)] = 0.0
                 J[:, offset + len(args) + radius_ind] += dR
 
             if self.refine_width:
@@ -608,11 +604,17 @@ def global_center_jacobian(
                     * (r_disk - disk_radius)
                     / (disk_width**2 * (1.0 + top_exp) ** 2)
                 ).ravel()
-                # dW[np.isnan(dW)] = 0.0
                 J[:, offset + len(args) - 1] += dW
 
-            # set_trace()
 
+class SyntheticDiskMoire(WPFModelPrototype):
+    def __init__(
+        self, 
+        lattice_a:SyntheticDiskLattice, 
+        lattice_b:SyntheticDiskLattice, 
+        decorated_peaks:list=None
+    ):
+        super().__init__(name, params, model_type=WPFModelType.META | WPFModelType.LATTICE)
 
 class ComplexOverlapKernelDiskLattice(WPFModelPrototype):
     def __init__(
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 9651748f0..426874f43 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -1,6 +1,6 @@
 from py4DSTEM import DataCube, RealSlice
 from emdfile import tqdmnd
-from py4DSTEM.process.wholepatternfit.wp_models import WPFModelPrototype, _BaseModel
+from py4DSTEM.process.wholepatternfit.wp_models import WPFModelPrototype, _BaseModel, WPFModelType
 
 from typing import Optional
 import numpy as np
@@ -80,11 +80,6 @@ def __init__(
 
         self.mask = mask if mask is not None else np.ones_like(self.meanCBED)
 
-        self.model = []
-        self.model_param_inds = []
-
-        self.nParams = 0
-        self.use_jacobian = use_jacobian
 
         if hasattr(x0, "__iter__") and hasattr(y0, "__iter__"):
             x0 = np.array(x0)
@@ -113,6 +108,11 @@ def __init__(
         # TODO: remove special cases for global/local center in the Models
         # Needs an efficient way to handle calculation of q_r
 
+        self.model = [self.coordinate_model, ]
+
+        self.nParams = 0
+        self.use_jacobian = use_jacobian
+
         # set up the global arguments
         self._setup_static_data(x0, y0)
 
@@ -127,8 +127,6 @@ def __init__(
     def add_model(self, model: WPFModelPrototype):
         self.model.append(model)
 
-        # keep track of where each model's parameter list begins
-        self.model_param_inds.append(self.nParams)
         self.nParams += len(model.params.keys())
 
         self._scrape_model_params()
@@ -492,11 +490,12 @@ def get_lattice_maps(self):
         lattices = [
             (i, m)
             for i, m in enumerate(self.model)
-            if "lattice" in type(m).__name__.lower()
+            if WPFModelType.LATTICE in m.model_type
         ]
 
         g_maps = []
         for i, l in lattices:
+            # TODO: use parameter object offsets to get indices
             param_list = list(l.params.keys())
             lattice_offset = param_list.index("ux")
             data_offset = self.model_param_inds[i] + 2 + lattice_offset
@@ -523,8 +522,6 @@ def _setup_static_data(self, x0, y0):
         self.static_data["Q_Nx"] = self.datacube.Q_Nx
         self.static_data["Q_Ny"] = self.datacube.Q_Ny
 
-        self.static_data["global_x0"] = x0
-        self.static_data["global_y0"] = y0
         self.static_data["global_r"] = np.hypot(
             (self.static_data["xArray"] - x0),
             (self.static_data["yArray"] - y0),

From e3d9b4e48a4bb80e24be03c0d53001d44a446c2a Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Sun, 6 Aug 2023 13:42:23 -0400
Subject: [PATCH 34/57] add linking function

---
 py4DSTEM/process/wholepatternfit/wpf.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 426874f43..71c381dac 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -135,6 +135,19 @@ def add_model_list(self, model_list):
         for m in model_list:
             self.add_model(m)
 
+    def link_parameters(self, parent_model, child_model, parameters):
+        """
+        Link parameters of separate models together. The parameters of
+        the child_model are replaced with the parameters of the parent_model.
+        Note, this does not add the models to the WPF object, that must
+        be performed separately. 
+        """
+        # Make sure parameters is iterable
+        parameters = [parameters,] if not hasattr(parameters,"__iter__")
+
+        for par in parameters:
+            child_model.params[par] = parent_model.params[par]
+
     def generate_initial_pattern(self):
         # update parameters:
         self._scrape_model_params()

From 9b5fa0475c7edfd6618e1bf156cf9f3a0e5a3f0b Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Sun, 6 Aug 2023 13:49:53 -0400
Subject: [PATCH 35/57] new offset assignment logic

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 23 ++++---
 py4DSTEM/process/wholepatternfit/wpf.py       | 64 +++++++++++--------
 2 files changed, 50 insertions(+), 37 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 4802ced19..b89de856e 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -16,8 +16,10 @@ class WPFModelType(Flag):
     AMORPHOUS = auto()
     LATTICE = auto()
 
-    DUMMY = auto() # Model has no direct contribution to pattern
-    META = auto() # Model depends on multiple sub-Models (triggers extra linking on inclusion)
+    DUMMY = auto()  # Model has no direct contribution to pattern
+    META = (
+        auto()
+    )  # Model depends on multiple sub-Models (triggers extra linking on inclusion)
 
 
 class WPFModelPrototype:
@@ -151,8 +153,8 @@ def __init__(
     ):
         params = {"sigma": Parameter(sigma), "intensity": Parameter(intensity)}
         if global_center:
-            params['x center'] = WPF.coordinate_model.params['x center']
-            params['y center'] = WPF.coordinate_model.params['y center']
+            params["x center"] = WPF.coordinate_model.params["x center"]
+            params["y center"] = WPF.coordinate_model.params["y center"]
             self.func = self.global_center_func
             self.jacobian = self.global_center_jacobian
         else:
@@ -609,12 +611,15 @@ def global_center_jacobian(
 
 class SyntheticDiskMoire(WPFModelPrototype):
     def __init__(
-        self, 
-        lattice_a:SyntheticDiskLattice, 
-        lattice_b:SyntheticDiskLattice, 
-        decorated_peaks:list=None
+        self,
+        lattice_a: SyntheticDiskLattice,
+        lattice_b: SyntheticDiskLattice,
+        decorated_peaks: list = None,
     ):
-        super().__init__(name, params, model_type=WPFModelType.META | WPFModelType.LATTICE)
+        super().__init__(
+            name, params, model_type=WPFModelType.META | WPFModelType.LATTICE
+        )
+
 
 class ComplexOverlapKernelDiskLattice(WPFModelPrototype):
     def __init__(
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 71c381dac..455ebdec8 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -1,6 +1,10 @@
 from py4DSTEM import DataCube, RealSlice
 from emdfile import tqdmnd
-from py4DSTEM.process.wholepatternfit.wp_models import WPFModelPrototype, _BaseModel, WPFModelType
+from py4DSTEM.process.wholepatternfit.wp_models import (
+    WPFModelPrototype,
+    _BaseModel,
+    WPFModelType,
+)
 
 from typing import Optional
 import numpy as np
@@ -80,7 +84,6 @@ def __init__(
 
         self.mask = mask if mask is not None else np.ones_like(self.meanCBED)
 
-
         if hasattr(x0, "__iter__") and hasattr(y0, "__iter__"):
             x0 = np.array(x0)
             y0 = np.array(y0)
@@ -108,7 +111,9 @@ def __init__(
         # TODO: remove special cases for global/local center in the Models
         # Needs an efficient way to handle calculation of q_r
 
-        self.model = [self.coordinate_model, ]
+        self.model = [
+            self.coordinate_model,
+        ]
 
         self.nParams = 0
         self.use_jacobian = use_jacobian
@@ -129,7 +134,7 @@ def add_model(self, model: WPFModelPrototype):
 
         self.nParams += len(model.params.keys())
 
-        self._scrape_model_params()
+        self._finalize_model()
 
     def add_model_list(self, model_list):
         for m in model_list:
@@ -140,22 +145,28 @@ def link_parameters(self, parent_model, child_model, parameters):
         Link parameters of separate models together. The parameters of
         the child_model are replaced with the parameters of the parent_model.
         Note, this does not add the models to the WPF object, that must
-        be performed separately. 
+        be performed separately.
         """
         # Make sure parameters is iterable
-        parameters = [parameters,] if not hasattr(parameters,"__iter__")
+        parameters = (
+            [
+                parameters,
+            ]
+            if not hasattr(parameters, "__iter__")
+            else parameters
+        )
 
         for par in parameters:
             child_model.params[par] = parent_model.params[par]
 
     def generate_initial_pattern(self):
         # update parameters:
-        self._scrape_model_params()
+        self._finalize_model()
         return self._pattern(self.x0, self.static_data.copy()) / self.intensity_scale
 
     def fit_to_mean_CBED(self, **fit_opts):
         # first make sure we have the latest parameters
-        self._scrape_model_params()
+        self._finalize_model()
 
         # set the current active pattern to the mean CBED:
         current_pattern = self.meanCBED * self.intensity_scale
@@ -261,7 +272,7 @@ def fit_single_pattern(
         """
 
         # make sure we have the latest parameters
-        self._scrape_model_params()
+        self._finalize_model()
 
         # set tracking off
         self._track = False
@@ -358,7 +369,7 @@ def fit_all_patterns(
         """
 
         # make sure we have the latest parameters
-        self._scrape_model_params()
+        self._finalize_model()
 
         # set tracking off
         self._track = False
@@ -597,23 +608,20 @@ def _jacobian(self, x, current_pattern, shared_data):
 
         return J * self.mask.ravel()[:, np.newaxis]
 
-    def _scrape_model_params(self):
-        self.x0 = np.zeros((self.nParams + 2,))
-        self.upper_bound = np.zeros_like(self.x0)
-        self.lower_bound = np.zeros_like(self.x0)
-
-        self.x0[0:2] = np.array(
-            [self.static_data["global_x0"], self.static_data["global_y0"]]
-        )
-        self.upper_bound[0:2] = self.global_xy0_ub
-        self.lower_bound[0:2] = self.global_xy0_lb
-
-        for i, m in enumerate(self.model):
-            ind = self.model_param_inds[i] + 2
-
-            for j, v in enumerate(m.params.values()):
-                self.x0[ind + j] = v.initial_value
-                self.upper_bound[ind + j] = v.upper_bound
-                self.lower_bound[ind + j] = v.lower_bound
+    def _finalize_model(self):
+        # iterate over all models and assign indices, accumulate list
+        # of unique parameters. then, accumulate initial value and bounds vectors
+        unique_params = []
+        idx = 0
+        for model in self.model:
+            for param in model.params:
+                if param not in unique_params:
+                    unique_params.append(param)
+                    param.offset = idx
+                    idx += 1
+
+        self.x0 = np.array([param.initial_value for param in unique_params])
+        self.upper_bound = np.array([param.upper_bound for param in unique_params])
+        self.lower_bound = np.array([param.lower_bound for param in unique_params])
 
         self.hasJacobian = all([m.hasJacobian for m in self.model])

From f77552f4c33facc00feb116f8c219978c11d7ecc Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Tue, 8 Aug 2023 10:56:15 -0400
Subject: [PATCH 36/57] remove global special case in gaussian bkgd

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 105 +++++-------------
 py4DSTEM/process/wholepatternfit/wpf.py       |  38 ++++---
 2 files changed, 48 insertions(+), 95 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index b89de856e..fa2ebfbf1 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -17,9 +17,7 @@ class WPFModelType(Flag):
     LATTICE = auto()
 
     DUMMY = auto()  # Model has no direct contribution to pattern
-    META = (
-        auto()
-    )  # Model depends on multiple sub-Models (triggers extra linking on inclusion)
+    META = auto()  # Model depends on multiple sub-Models
 
 
 class WPFModelPrototype:
@@ -56,11 +54,6 @@ def __init__(self, name: str, params: dict, model_type=WPFModelType.DUMMY):
 
         self.model_type = model_type
 
-    #     # check the function obeys the spec
-    #     assert (
-    #         len(signature(self.func).parameters) == len(params) + 2
-    #     ), f"The model function has the wrong number of arguments in its signature. It must be written as func(DP, param1, param2, ..., **kwargs). The current signature is {str(signature(self.func))}"
-
     def func(self, DP: np.ndarray, x, **kwargs) -> None:
         raise NotImplementedError()
 
@@ -155,100 +148,48 @@ def __init__(
         if global_center:
             params["x center"] = WPF.coordinate_model.params["x center"]
             params["y center"] = WPF.coordinate_model.params["y center"]
-            self.func = self.global_center_func
-            self.jacobian = self.global_center_jacobian
         else:
             params["x center"] = Parameter(x0)
             params["y center"] = Parameter(y0)
-            self.func = self.local_center_func
-            self.jacobian = self.local_center_jacobian
 
         super().__init__(name, params, model_type=WPFModelType.BACKGROUND)
 
-    def global_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
+    def func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["intensity"].offset]
 
-        DP += level * np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
-
-    def global_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
-        sigma = x[self.params["sigma"].offset]
-        level = x[self.params["intensity"].offset]
+        r = WPF._get_distance(x, self.params["x center"], self.params["y center"])
 
-        exp_expr = np.exp(kwargs["global_r"] ** 2 / (-2 * sigma**2))
+        DP += level * np.exp(r**2 / (-2 * sigma**2))
 
-        # dF/d(global_x0)
-        J[:, 0] += (
-            level * (kwargs["xArray"] - kwargs["global_x0"]) * exp_expr / sigma**2
-        ).ravel()
-
-        # dF/d(global_y0)
-        J[:, 1] += (
-            level * (kwargs["yArray"] - kwargs["global_y0"]) * exp_expr / sigma**2
-        ).ravel()
-
-        # dF/s(sigma)
-        J[:, self.params["sigma"].offset] = (
-            level * kwargs["global_r"] ** 2 * exp_expr / sigma**3
-        ).ravel()
-
-        # dF/d(level)
-        J[:, self.params["intensity"].offset] = exp_expr.ravel()
-
-    def local_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
+    def jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["intensity"].offset]
         x0 = x[self.params["x center"].offset]
         y0 = x[self.params["y center"].offset]
-        DP += level * np.exp(
-            ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-            / (-2 * sigma**2)
-        )
 
-    def local_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
-        sigma = x[self.params["sigma"].offset]
-        level = x[self.params["intensity"].offset]
-        x0 = x[self.params["x center"].offset]
-        y0 = x[self.params["y center"].offset]
+        r = kwargs["parent"]._get_distance(
+            x, self.params["x center"], self.params["y center"]
+        )
+        exp_expr = np.exp(r**2 / (-2 * sigma**2))
 
-        # dF/s(sigma)
-        J[:, self.params["sigma"].offset] = (
-            level
-            * ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-            * np.exp(
-                ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-                / (-2 * sigma**2)
-            )
-            / sigma**3
+        # dF/d(x0)
+        J[:, self.params["x center"].offset] += (
+            level * (kwargs["xArray"] - x0) * exp_expr / sigma**2
         ).ravel()
 
-        # dF/d(level)
-        J[:, self.params["intensity"].offset] = np.exp(
-            ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-            / (-2 * sigma**2)
+        # dF/d(y0)
+        J[:, self.params["y center"].offset] += (
+            level * (kwargs["yArray"] - y0) * exp_expr / sigma**2
         ).ravel()
 
-        # dF/d(x0)
-        J[:, self.params["x center"].offset] = (
-            level
-            * (kwargs["xArray"] - x0)
-            * np.exp(
-                ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-                / (-2 * sigma**2)
-            )
-            / sigma**2
+        # dF/s(sigma)
+        J[:, self.params["sigma"].offset] += (
+            level * r**2 * exp_expr / sigma**3
         ).ravel()
 
-        # dF/d(y0)
-        J[:, self.params["y center"].offset] = (
-            level
-            * (kwargs["yArray"] - y0)
-            * np.exp(
-                ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-                / (-2 * sigma**2)
-            )
-            / sigma**2
-        ).ravel()
+        # dF/d(level)
+        J[:, self.params["intensity"].offset] += exp_expr.ravel()
 
 
 class GaussianRing(WPFModelPrototype):
@@ -620,6 +561,12 @@ def __init__(
             name, params, model_type=WPFModelType.META | WPFModelType.LATTICE
         )
 
+        # ensure both models share the same center coordinate
+
+        # pick the right pair of lattice vectors for generating the moire reciprocal lattice
+
+        #
+
 
 class ComplexOverlapKernelDiskLattice(WPFModelPrototype):
     def __init__(
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 455ebdec8..8ba61751f 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -4,6 +4,7 @@
     WPFModelPrototype,
     _BaseModel,
     WPFModelType,
+    Parameter,
 )
 
 from typing import Optional
@@ -546,26 +547,31 @@ def _setup_static_data(self, x0, y0):
         self.static_data["Q_Nx"] = self.datacube.Q_Nx
         self.static_data["Q_Ny"] = self.datacube.Q_Ny
 
-        self.static_data["global_r"] = np.hypot(
-            (self.static_data["xArray"] - x0),
-            (self.static_data["yArray"] - y0),
-        )
-
-    def _pattern_error(self, x, current_pattern, shared_data):
-        DP = np.zeros((self.datacube.Q_Nx, self.datacube.Q_Ny))
+        # todo: does this create a problematic cicular reference?
+        self.static_data["parent"] = self
 
-        shared_data["global_x0"] = x[0]
-        shared_data["global_y0"] = x[1]
-        shared_data["global_r"] = np.hypot(
-            (shared_data["xArray"] - x[0]),
-            (shared_data["yArray"] - x[1]),
+    def _get_distance(self, params: np.ndarray, x: Parameter, y: Parameter):
+        """
+        Return the distance from a point in pixel coordinates specified
+        by two Parameter objects.
+        This method caches the result from the _BaseModel for performance
+        """
+        if (
+            x is self.model[0].params["x center"]
+            and y is self.model[0].params["y center"]
+        ):
+            # TODO: actually implement caching
+            pass
+
+        return np.hypot(
+            self.static_data["xArray"] - params[x.offset],
+            self.static_data["yArray"] - params[y.offset],
         )
 
-        for i, m in enumerate(self.model):
-            ind = self.model_param_inds[i] + 2
-            m.func(DP, *x[ind : ind + m.nParams].tolist(), **shared_data)
+    def _pattern_error(self, x, current_pattern, shared_data):
+        DP = self._pattern(x, shared_data)
 
-        DP = (DP**self.fit_power - current_pattern**self.fit_power) * self.mask
+        DP = (DP - current_pattern**self.fit_power) * self.mask
 
         if self._track:
             self._fevals.append(DP)

From 54ff82fba5ec83998d49cc0bc11b349d41507783 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Tue, 8 Aug 2023 11:08:54 -0400
Subject: [PATCH 37/57] remove global special case from gaussian ring

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 102 +++++-------------
 1 file changed, 24 insertions(+), 78 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index fa2ebfbf1..8e33d56a7 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -195,6 +195,7 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
 class GaussianRing(WPFModelPrototype):
     def __init__(
         self,
+        WPF,
         radius,
         sigma,
         intensity,
@@ -207,15 +208,13 @@ def __init__(
             "radius": Parameter(radius),
             "sigma": Parameter(sigma),
             "intensity": Parameter(intensity),
+            "x center": WPF.coordinate_model.params["x center"]
+            if global_center
+            else Parameter(x0),
+            "y center": WPF.coordinate_model.params["y center"]
+            if global_center
+            else Parameter(y0),
         }
-        if global_center:
-            self.func = self.global_center_func
-            self.jacobian = self.global_center_jacobian
-        else:
-            params["x center"] = Parameter(x0)
-            params["y center"] = Parameter(y0)
-            self.func = self.local_center_func
-            self.jacobian = self.local_center_jacobian
 
         super().__init__(name, params, model_type=WPFModelType.AMORPHOUS)
 
@@ -223,32 +222,39 @@ def global_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
         radius = x[self.params["radius"].offset]
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["level"].offset]
-        DP += level * np.exp((kwargs["global_r"] - radius) ** 2 / (-2 * sigma**2))
+
+        r = WPF._get_distance(x, self.params["x center"], self.params["y center"])
+
+        DP += level * np.exp((r - radius) ** 2 / (-2 * sigma**2))
 
     def global_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
         radius = x[self.params["radius"].offset]
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["level"].offset]
 
-        local_r = radius - kwargs["global_r"]
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
+        r = WPF._get_distance(x, self.params["x center"], self.params["y center"])
+
+        local_r = radius - r
         clipped_r = np.maximum(local_r, 0.1)
 
         exp_expr = np.exp(local_r**2 / (-2 * sigma**2))
 
-        # dF/d(global_x0)
-        J[:, 0] += (
+        # dF/d(x0)
+        J[:, self.params["x center"].offset] += (
             level
             * exp_expr
-            * (kwargs["xArray"] - kwargs["global_x0"])
+            * (kwargs["xArray"] - x0)
             * local_r
             / (sigma**2 * clipped_r)
         ).ravel()
 
-        # dF/d(global_y0)
-        J[:, 1] += (
+        # dF/d(y0)
+        J[:, self.parans["y center"].offset] += (
             level
             * exp_expr
-            * (kwargs["yArray"] - kwargs["global_y0"])
+            * (kwargs["yArray"] - x0)
             * local_r
             / (sigma**2 * clipped_r)
         ).ravel()
@@ -259,72 +265,12 @@ def global_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None
         ).ravel()
 
         # dF/d(sigma)
-        J[:, self.params["sigma"].offset] = (
+        J[:, self.params["sigma"].offset] += (
             level * local_r**2 * exp_expr / sigma**3
         ).ravel()
 
         # dF/d(intensity)
-        J[:, self.params["intensity"].offset] = exp_expr.ravel()
-
-    def local_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
-        radius = x[self.params["radius"].offset]
-        sigma = x[self.params["sigma"].offset]
-        level = x[self.params["level"].offset]
-        x0 = x[self.params["x center"].offset]
-        y0 = x[self.params["y center"].offset]
-
-        local_r = np.hypot(kwargs["xArray"] - x0, kwargs["yArray"] - y0)
-        DP += level * np.exp((local_r - radius) ** 2 / (-2 * sigma**2))
-
-    def local_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
-        return NotImplementedError()
-
-        radius = x[self.params["radius"].offset]
-        sigma = x[self.params["sigma"].offset]
-        level = x[self.params["level"].offset]
-        x0 = x[self.params["x center"].offset]
-        y0 = x[self.params["y center"].offset]
-
-        # dF/d(radius)
-
-        # dF/s(sigma)
-        J[:, self.params["sigma"].offset] = (
-            level
-            * ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-            * np.exp(
-                ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-                / (-2 * sigma**2)
-            )
-            / sigma**3
-        ).ravel()
-
-        # dF/d(level)
-        J[:, self.params["intensity"].offset] = np.exp(
-            ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-            / (-2 * sigma**2)
-        ).ravel()
-
-        # dF/d(x0)
-        J[:, self.params["x center"].offset] = (
-            level
-            * (kwargs["xArray"] - x0)
-            * np.exp(
-                ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-                / (-2 * sigma**2)
-            )
-            / sigma**2
-        ).ravel()
-
-        # dF/d(y0)
-        J[:, self.params["y center"].offset] = (
-            level
-            * (kwargs["yArray"] - y0)
-            * np.exp(
-                ((kwargs["xArray"] - x0) ** 2 + (kwargs["yArray"] - y0) ** 2)
-                / (-2 * sigma**2)
-            )
-            / sigma**2
-        ).ravel()
+        J[:, self.params["intensity"].offset] += exp_expr.ravel()
 
 
 class SyntheticDiskLattice(WPFModelPrototype):

From 022b941a34be3762586041593a34fb202ebc9203 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Wed, 9 Aug 2023 10:46:45 -0400
Subject: [PATCH 38/57] fully convert synthetic disk model and squash bugs

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 177 ++++++++++--------
 py4DSTEM/process/wholepatternfit/wpf.py       |  31 +--
 py4DSTEM/process/wholepatternfit/wpf_viz.py   |  35 ++--
 3 files changed, 128 insertions(+), 115 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 8e33d56a7..5db679b59 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -158,7 +158,7 @@ def func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["intensity"].offset]
 
-        r = WPF._get_distance(x, self.params["x center"], self.params["y center"])
+        r = kwargs['parent']._get_distance(x, self.params["x center"], self.params["y center"])
 
         DP += level * np.exp(r**2 / (-2 * sigma**2))
 
@@ -302,15 +302,14 @@ def __init__(
         params = {}
 
         if global_center:
-            self.func = self.global_center_func
-            self.jacobian = self.global_center_jacobian
-
-            x0 = WPF.static_data["global_x0"]
-            y0 = WPF.static_data["global_y0"]
+            params["x center"] = WPF.coordinate_model.params["x center"]
+            params["y center"] = WPF.coordinate_model.params["y center"]
         else:
             params["x center"] = Parameter(x0)
             params["y center"] = Parameter(y0)
-            self.func = self.local_center_func
+
+        x0 = params["x center"].initial_value
+        y0 = params["y center"].initial_value
 
         params["ux"] = Parameter(ux)
         params["uy"] = Parameter(uy)
@@ -366,41 +365,39 @@ def __init__(
 
         super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
-    def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-        # copy the global centers in the right place for the local center generator
-        self.local_center_func(
-            DP, kwargs["global_x0"], kwargs["global_y0"], *args, **kwargs
+    def func(self, DP: np.ndarray, x: np.ndarray, **static_data) -> None:
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
+        ux = x[self.params["ux"].offset]
+        uy = x[self.params["uy"].offset]
+        vx = x[self.params["vx"].offset]
+        vy = x[self.params["vy"].offset]
+
+        disk_radius = (
+            x[self.params["disk radius"].offset]
+            if self.refine_radius
+            else self.disk_radius
         )
 
-    def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-        x0 = args[0]
-        y0 = args[1]
-        ux = args[2]
-        uy = args[3]
-        vx = args[4]
-        vy = args[5]
-
-        if self.refine_radius & self.refine_width:
-            disk_radius = args[-2]
-        elif self.refine_radius:
-            disk_radius = args[-1]
-        else:
-            disk_radius = self.disk_radius
-        disk_width = args[-1] if self.refine_width else self.disk_width
+        disk_width = (
+            x[self.params["edge width"].offset]
+            if self.refine_width
+            else self.disk_width
+        )
 
         for i, (u, v) in enumerate(zip(self.u_inds, self.v_inds)):
-            x = x0 + (u * ux) + (v * vx)
-            y = y0 + (u * uy) + (v * vy)
+            x_pos = x0 + (u * ux) + (v * vx)
+            y_pos = y0 + (u * uy) + (v * vy)
 
-            DP += args[i + 6] / (
+            DP += x[self.params[f"[{u},{v}] Intensity"].offset] / (
                 1.0
                 + np.exp(
                     np.minimum(
                         4
                         * (
                             np.sqrt(
-                                (kwargs["xArray"] - x) ** 2
-                                + (kwargs["yArray"] - y) ** 2
+                                (static_data["xArray"] - x_pos) ** 2
+                                + (static_data["yArray"] - y_pos) ** 2
                             )
                             - disk_radius
                         )
@@ -410,109 +407,137 @@ def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
                 )
             )
 
-    def global_center_jacobian(
-        self, J: np.ndarray, *args, offset: int, **kwargs
-    ) -> None:
-        x0 = kwargs["global_x0"]
-        y0 = kwargs["global_y0"]
-        r = np.maximum(5e-1, kwargs["global_r"])
-        ux = args[0]
-        uy = args[1]
-        vx = args[2]
-        vy = args[3]
-
-        if self.refine_radius & self.refine_width:
-            disk_radius = args[-2]
-            radius_ind = -2
-        elif self.refine_radius:
-            disk_radius = args[-1]
-            radius_ind = -1
-        else:
-            disk_radius = self.disk_radius
-        disk_width = args[-1] if self.refine_width else self.disk_width
+    def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data) -> None:
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
+        ux = x[self.params["ux"].offset]
+        uy = x[self.params["uy"].offset]
+        vx = x[self.params["vx"].offset]
+        vy = x[self.params["vy"].offset]
+        WPF = static_data['parent']
+
+        r = np.maximum(
+            5e-1, WPF._get_distance(x, self.params["x center"], self.params["y center"])
+        )
+
+        disk_radius = (
+            x[self.params["disk radius"].offset]
+            if self.refine_radius
+            else self.disk_radius
+        )
+
+        disk_width = (
+            x[self.params["edge width"].offset]
+            if self.refine_width
+            else self.disk_width
+        )
 
         for i, (u, v) in enumerate(zip(self.u_inds, self.v_inds)):
-            x = x0 + (u * ux) + (v * vx)
-            y = y0 + (u * uy) + (v * vy)
+            x_pos = x0 + (u * ux) + (v * vx)
+            y_pos = y0 + (u * uy) + (v * vy)
 
-            disk_intensity = args[i + 6]
+            disk_intensity = x[self.params[f"[{u},{v}] Intensity"].offset]
 
             r_disk = np.maximum(
                 5e-1,
-                np.sqrt((kwargs["xArray"] - x) ** 2 + (kwargs["yArray"] - y) ** 2),
+                np.sqrt(
+                    (static_data["xArray"] - x_pos) ** 2 + (static_data["yArray"] - y_pos) ** 2
+                ),
             )
 
             mask = r_disk < (2 * disk_radius)
 
             top_exp = mask * np.exp(4 * ((mask * r_disk) - disk_radius) / disk_width)
 
-            # dF/d(global_x0)
+            # dF/d(x0)
             dx = (
                 4
-                * args[i + 4]
-                * (kwargs["xArray"] - x)
+                * disk_intensity
+                * (static_data["xArray"] - x_pos)
                 * top_exp
                 / ((1.0 + top_exp) ** 2 * disk_width * r)
             ).ravel()
 
-            # dF/d(global_y0)
+            # dF/d(y0)
             dy = (
                 4
-                * args[i + 4]
-                * (kwargs["yArray"] - y)
+                * disk_intensity
+                * (static_data["yArray"] - y_pos)
                 * top_exp
                 / ((1.0 + top_exp) ** 2 * disk_width * r)
             ).ravel()
 
-            # insert global positional derivatives
-            J[:, 0] += disk_intensity * dx
-            J[:, 1] += disk_intensity * dy
+            # insert center position derivatives
+            J[:, self.params["x center"].offset] += disk_intensity * dx
+            J[:, self.params["y center"].offset] += disk_intensity * dy
 
             # insert lattice vector derivatives
-            J[:, offset] += disk_intensity * u * dx
-            J[:, offset + 1] += disk_intensity * u * dy
-            J[:, offset + 2] += disk_intensity * v * dx
-            J[:, offset + 3] += disk_intensity * v * dy
+            J[:, self.params["ux"].offset] += disk_intensity * u * dx
+            J[:, self.params["uy"].offset] += disk_intensity * u * dy
+            J[:, self.params["vx"].offset] += disk_intensity * v * dx
+            J[:, self.params["vy"].offset] += disk_intensity * v * dy
 
             # insert intensity derivative
             dI = (mask * (1.0 / (1.0 + top_exp))).ravel()
-            J[:, offset + i + 4] = dI
+            J[:, self.params[f"[{u},{v}] Intensity"].offset] += dI
 
             # insert disk radius derivative
             if self.refine_radius:
                 dR = (
-                    4.0 * args[i + 4] * top_exp / (disk_width * (1.0 + top_exp) ** 2)
+                    4.0 * disk_intensity * top_exp / (disk_width * (1.0 + top_exp) ** 2)
                 ).ravel()
-                J[:, offset + len(args) + radius_ind] += dR
+                J[:, self.params["disk radius"].offset] += dR
 
             if self.refine_width:
                 dW = (
                     4.0
-                    * args[i + 4]
+                    * disk_intensity
                     * top_exp
                     * (r_disk - disk_radius)
                     / (disk_width**2 * (1.0 + top_exp) ** 2)
                 ).ravel()
-                J[:, offset + len(args) - 1] += dW
+                J[:, self.params["edge width"].offset] += dW
 
 
 class SyntheticDiskMoire(WPFModelPrototype):
+    """
+    Add Moire peaks arising from two SyntheticDiskLattice lattices.
+    The positions of the Moire peaks are derived from the lattice
+    vectors of the parent lattices. This model object adds only the intensity of
+    each Moire peak as parameters, all other attributes are inherited from the parents
+    """
+
     def __init__(
         self,
         lattice_a: SyntheticDiskLattice,
         lattice_b: SyntheticDiskLattice,
         decorated_peaks: list = None,
+        name: str = "Moire Lattice",
     ):
-        super().__init__(
-            name, params, model_type=WPFModelType.META | WPFModelType.LATTICE
-        )
+        """
+        Parameters
+        ----------
 
+        lattice_a, lattice_b: SyntheticDiskLattice
+        """
         # ensure both models share the same center coordinate
 
         # pick the right pair of lattice vectors for generating the moire reciprocal lattice
 
         #
 
+        super().__init__(
+            name,
+            params,
+            model_type=WPFModelType.META,
+        )
+
+    def func(self, DP, x, **static_data):
+        pass
+
+    def jacobian(self, J, x, **static_data):
+        pass
+
 
 class ComplexOverlapKernelDiskLattice(WPFModelPrototype):
     def __init__(
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 8ba61751f..b0333974a 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -547,7 +547,6 @@ def _setup_static_data(self, x0, y0):
         self.static_data["Q_Nx"] = self.datacube.Q_Nx
         self.static_data["Q_Ny"] = self.datacube.Q_Ny
 
-        # todo: does this create a problematic cicular reference?
         self.static_data["parent"] = self
 
     def _get_distance(self, params: np.ndarray, x: Parameter, y: Parameter):
@@ -583,34 +582,18 @@ def _pattern_error(self, x, current_pattern, shared_data):
     def _pattern(self, x, shared_data):
         DP = np.zeros((self.datacube.Q_Nx, self.datacube.Q_Ny))
 
-        shared_data["global_x0"] = x[0]
-        shared_data["global_y0"] = x[1]
-        shared_data["global_r"] = np.hypot(
-            (shared_data["xArray"] - x[0]),
-            (shared_data["yArray"] - x[1]),
-        )
-
-        for i, m in enumerate(self.model):
-            ind = self.model_param_inds[i] + 2
-            m.func(DP, *x[ind : ind + m.nParams].tolist(), **shared_data)
+        for m in self.model:
+            m.func(DP, x, **shared_data)
 
         return (DP**self.fit_power) * self.mask
 
     def _jacobian(self, x, current_pattern, shared_data):
         # TODO: automatic mixed analytic/finite difference
 
-        J = np.zeros(((self.datacube.Q_Nx * self.datacube.Q_Ny), self.nParams + 2))
+        J = np.zeros(((self.datacube.Q_Nx * self.datacube.Q_Ny), self.nParams))
 
-        shared_data["global_x0"] = x[0]
-        shared_data["global_y0"] = x[1]
-        shared_data["global_r"] = np.hypot(
-            (shared_data["xArray"] - x[0]),
-            (shared_data["yArray"] - x[1]),
-        )
-
-        for i, m in enumerate(self.model):
-            ind = self.model_param_inds[i] + 2
-            m.jacobian(J, *x[ind : ind + m.nParams].tolist(), offset=ind, **shared_data)
+        for m in self.model:
+            m.jacobian(J, x, **shared_data)
 
         return J * self.mask.ravel()[:, np.newaxis]
 
@@ -620,7 +603,7 @@ def _finalize_model(self):
         unique_params = []
         idx = 0
         for model in self.model:
-            for param in model.params:
+            for param in model.params.values():
                 if param not in unique_params:
                     unique_params.append(param)
                     param.offset = idx
@@ -631,3 +614,5 @@ def _finalize_model(self):
         self.lower_bound = np.array([param.lower_bound for param in unique_params])
 
         self.hasJacobian = all([m.hasJacobian for m in self.model])
+
+        self.nParams = self.x0.shape[0]
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index c02560653..eb7efd450 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -5,6 +5,8 @@
 import matplotlib.colors as mpl_c
 from matplotlib.gridspec import GridSpec
 
+from py4DSTEM.process.wholepatternfit.wp_models import WPFModelType
+
 
 def show_model_grid(self, x=None, **plot_kwargs):
     if x is None:
@@ -94,25 +96,26 @@ def show_lattice_points(
             cmap="gray",
         )
 
-    for m in self.model:
-        if "Lattice" in m.name:
-            ux, uy = m.params["ux"].initial_value, m.params["uy"].initial_value
-            vx, vy = m.params["vx"].initial_value, m.params["vy"].initial_value
+    lattices = [m for m in self.model if WPFModelType.LATTICE in m.model_type]
 
-            lat = np.array([[ux, uy], [vx, vy]])
-            inds = np.stack([m.u_inds, m.v_inds], axis=1)
+    for m in lattices:
+        ux, uy = m.params["ux"].initial_value, m.params["uy"].initial_value
+        vx, vy = m.params["vx"].initial_value, m.params["vy"].initial_value
 
-            spots = inds @ lat
-            spots[:, 0] += self.static_data["global_x0"]
-            spots[:, 1] += self.static_data["global_y0"]
+        lat = np.array([[ux, uy], [vx, vy]])
+        inds = np.stack([m.u_inds, m.v_inds], axis=1)
 
-            ax.scatter(
-                spots[:, 1],
-                spots[:, 0],
-                s=100,
-                marker="x",
-                label=m.name,
-            )
+        spots = inds @ lat
+        spots[:, 0] += self.coordinate_model.params["x center"].initial_value
+        spots[:, 1] += self.coordinate_model.params["y center"].initial_value
+
+        ax.scatter(
+            spots[:, 1],
+            spots[:, 0],
+            s=100,
+            marker="x",
+            label=m.name,
+        )
 
     ax.legend()
 

From d5aa018af1a9dd743e4dabd1b51edd0257a3c6a1 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Wed, 9 Aug 2023 13:08:26 -0400
Subject: [PATCH 39/57] link multiple models at the same time and fix plots

---
 py4DSTEM/process/wholepatternfit/wpf.py     | 28 +++++++++++----------
 py4DSTEM/process/wholepatternfit/wpf_viz.py | 25 ++++--------------
 2 files changed, 20 insertions(+), 33 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index b0333974a..c69ce0a54 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -148,7 +148,15 @@ def link_parameters(self, parent_model, child_model, parameters):
         Note, this does not add the models to the WPF object, that must
         be performed separately.
         """
-        # Make sure parameters is iterable
+        # Make sure child_model and parameters are iterable
+        child_model = (
+            [
+                child_model,
+            ]
+            if not hasattr(child_model, "__iter__")
+            else child_model
+        )
+
         parameters = (
             [
                 parameters,
@@ -157,8 +165,9 @@ def link_parameters(self, parent_model, child_model, parameters):
             else parameters
         )
 
-        for par in parameters:
-            child_model.params[par] = parent_model.params[par]
+        for child in child_model:
+            for par in parameters:
+                child.params[par] = parent_model.params[par]
 
     def generate_initial_pattern(self):
         # update parameters:
@@ -497,17 +506,10 @@ def fit_all_patterns(
 
     def accept_mean_CBED_fit(self):
         x = self.mean_CBED_fit.x
-        self.static_data["global_x0"] = x[0]
-        self.static_data["global_y0"] = x[1]
 
-        self.static_data["global_r"] = np.hypot(
-            (self.static_data["xArray"] - x[0]), (self.static_data["yArray"] - x[1])
-        )
-
-        for i, m in enumerate(self.model):
-            ind = self.model_param_inds[i] + 2
-            for j, k in enumerate(m.params.keys()):
-                m.params[k].initial_value = x[ind + j]
+        for model in self.model:
+            for param in model.params.values():
+                param.initial_value = x[param.offset]
 
     def get_lattice_maps(self):
         assert hasattr(self, "fit_data"), "Please run fitting first!"
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index eb7efd450..d6201a15c 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -9,25 +9,11 @@
 
 
 def show_model_grid(self, x=None, **plot_kwargs):
-    if x is None:
-        x = self.mean_CBED_fit.x
-
-    shared_data = self.static_data.copy()
-    shared_data["global_x0"] = x[0]
-    shared_data["global_y0"] = x[1]
-    shared_data["global_r"] = np.hypot(
-        (shared_data["xArray"] - x[0]),
-        (shared_data["yArray"] - x[1]),
-    )
+    x = x or self.mean_CBED_fit.x
 
-    shared_data["global_x0"] = x[0]
-    shared_data["global_y0"] = x[1]
-    shared_data["global_r"] = np.hypot(
-        (shared_data["xArray"] - x[0]),
-        (shared_data["yArray"] - x[1]),
-    )
+    model = [m for m in self.model if WPFModelType.DUMMY not in m.model_type]
 
-    N = len(self.model)
+    N = len(model)
     cols = int(np.ceil(np.sqrt(N)))
     rows = (N + 1) // cols
 
@@ -35,10 +21,9 @@ def show_model_grid(self, x=None, **plot_kwargs):
     kwargs.update(plot_kwargs)
     fig, ax = plt.subplots(rows, cols, **kwargs)
 
-    for i, (a, m) in enumerate(zip(ax.flat, self.model)):
+    for (a, m) in zip(ax.flat, model):
         DP = np.zeros((self.datacube.Q_Nx, self.datacube.Q_Ny))
-        ind = self.model_param_inds[i] + 2
-        m.func(DP, *x[ind : ind + m.nParams].tolist(), **shared_data)
+        m.func(DP, x, **self.static_data)
 
         a.matshow(DP, cmap="turbo")
 

From d143cd187e1eef0ac928cc452e7786f4921cc9d9 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Wed, 9 Aug 2023 16:21:17 -0400
Subject: [PATCH 40/57] changes to distributed WPF

---
 py4DSTEM/process/wholepatternfit/wp_models.py |   9 +-
 py4DSTEM/process/wholepatternfit/wpf.py       | 255 +++++++++---------
 py4DSTEM/process/wholepatternfit/wpf_viz.py   |   2 +-
 3 files changed, 134 insertions(+), 132 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 5db679b59..5f1207ed0 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -158,7 +158,9 @@ def func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["intensity"].offset]
 
-        r = kwargs['parent']._get_distance(x, self.params["x center"], self.params["y center"])
+        r = kwargs["parent"]._get_distance(
+            x, self.params["x center"], self.params["y center"]
+        )
 
         DP += level * np.exp(r**2 / (-2 * sigma**2))
 
@@ -414,7 +416,7 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data) -> None:
         uy = x[self.params["uy"].offset]
         vx = x[self.params["vx"].offset]
         vy = x[self.params["vy"].offset]
-        WPF = static_data['parent']
+        WPF = static_data["parent"]
 
         r = np.maximum(
             5e-1, WPF._get_distance(x, self.params["x center"], self.params["y center"])
@@ -441,7 +443,8 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data) -> None:
             r_disk = np.maximum(
                 5e-1,
                 np.sqrt(
-                    (static_data["xArray"] - x_pos) ** 2 + (static_data["yArray"] - y_pos) ** 2
+                    (static_data["xArray"] - x_pos) ** 2
+                    + (static_data["yArray"] - y_pos) ** 2
                 ),
             )
 
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index c69ce0a54..f1d511f28 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -16,8 +16,6 @@
 from matplotlib.gridspec import GridSpec
 import warnings
 
-from mpire import WorkerPool
-
 __all__ = ["WholePatternFit"]
 
 
@@ -251,107 +249,13 @@ def fit_to_mean_CBED(self, **fit_opts):
 
         return opt
 
-    def fit_single_pattern(
-        self,
-        data: np.ndarray,
-        resume: bool = False,
-        restart_data: np.ndarray = None,
-        **fit_opts,
-    ):
-        """
-        Apply model fitting to one pattern.
-
-        Parameters
-        ----------
-        resume: bool (optional)
-            Set to true to continue a previous fit with more iterations.
-        rx: int
-            probe x coordinate
-        ry: int
-            probe y coordinate
-        fit_opts: args (optional)
-            args passed to scipy.optimize.least_squares
-
-        Returns
-        --------
-        fit_coefs: np.array
-            Fitted coefficients
-        fit_metrics: np.array
-            Fitting metrics
-
-        """
-
-        # make sure we have the latest parameters
-        self._finalize_model()
-
-        # set tracking off
-        self._track = False
-        self._fevals = []
-
-        if resume:
-            assert hasattr(self, "fit_data"), "No existing data resuming fit!"
-
-        # init
-        fit_coefs = np.zeros(self.x0.shape[0])
-        fit_metric = np.zeros(4)
-
-        # Default fitting options
-        default_opts = {
-            "method": "trf",
-            "verbose": 0,
-            "x_scale": "jac",
-        }
-        default_opts.update(fit_opts)
-
-        # Loop over probe positions
-        current_pattern = self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
-        shared_data = self.static_data.copy()
-        self._cost_history = (
-            []
-        )  # clear this so it doesn't grow: TODO make this not stupid
-
-        try:
-            x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0
-
-            if self.hasJacobian & self.use_jacobian:
-                opt = least_squares(
-                    self._pattern_error,
-                    x0,
-                    jac=self._jacobian,
-                    bounds=(self.lower_bound, self.upper_bound),
-                    args=(current_pattern, shared_data),
-                    **default_opts,
-                    # **fit_opts,
-                )
-            else:
-                opt = least_squares(
-                    self._pattern_error,
-                    x0,
-                    bounds=(self.lower_bound, self.upper_bound),
-                    args=(current_pattern, shared_data),
-                    **default_opts,
-                    # **fit_opts,
-                )
-
-            fit_coefs = opt.x
-            fit_metrics_single = [
-                opt.cost,
-                opt.optimality,
-                opt.nfev,
-                opt.status,
-            ]
-        except:
-            fit_coefs = x0
-            fit_metrics_single = [0, 0, 0, 0]
-
-        return fit_coefs, fit_metrics_single
-
     def fit_all_patterns(
         self,
         resume=False,
         real_space_mask=None,
-        num_jobs=None,
         show_fit_metrics=True,
+        distributed=True,
+        num_jobs=None,
         **fit_opts,
     ):
         """
@@ -364,6 +268,8 @@ def fit_all_patterns(
         real_space_mask: np.array() of bools (optional)
             Only perform the fitting on a subset of the probe positions,
             where real_space_mask[rx,ry] == True.
+        distributed: bool (optional)
+            Whether to evaluate using a pool of worker threads
         num_jobs: int (optional)
             Set to an integer value giving the number of jobs to parallelize over probe positions.
         fit_opts: args (optional)
@@ -401,7 +307,7 @@ def fit_all_patterns(
         default_opts.update(fit_opts)
 
         # Loop over probe positions
-        if num_jobs is None:
+        if not distributed:
             for rx, ry in tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny):
                 if real_space_mask is not None and real_space_mask[rx, ry] == True:
                     current_pattern = (
@@ -446,34 +352,14 @@ def fit_all_patterns(
                     fit_metrics[:, rx, ry] = fit_metrics_single
 
         else:
-            # Get list of probe positions
-            if real_space_mask is not None:
-                xa, ya = np.where(real_space_mask)
-            else:
-                xa, ya = np.meshgrid(
-                    np.arange(self.datacube.Rshape[0]),
-                    np.arange(self.datacube.Rshape[1]),
-                    indexing="ij",
-                )
-                xa = xa.ravel()
-                ya = ya.ravel()
-            xy = np.vstack((xa, ya))
-
-            fit_inputs = [default_opts] * xy.shape[1]
-            for a0 in range(xy.shape[1]):
-                fit_inputs[a0]["rx"] = xy[0, a0]
-                fit_inputs[a0]["ry"] = xy[1, a0]
-
-            with WorkerPool(n_jobs=num_jobs) as pool:
-                results = pool.map(
-                    self.fit_single_pattern,
-                    fit_inputs,
-                    progress_bar=True,
-                )
-
-            for a0 in range(xy.shape[1]):
-                fit_data[:, xy[0, a0], xy[1, a0]] = results[a0][0]
-                fit_metrics[:, xy[0, a0], xy[1, a0]] = results[a0][1]
+            # distributed evaluation
+            self._fit_distributed(
+                resume=resume,
+                num_jobs=num_jobs,
+                fit_opts=default_opts,
+                fit_data=fit_data,
+                fit_metrics=fit_metrics,
+            )
 
         # Convert to RealSlices
         model_names = []
@@ -486,7 +372,8 @@ def fit_all_patterns(
                     i += 1
             model_names.append(n)
 
-        param_names = ["global_x0", "global_y0"] + [
+        # TODO: this produces duplicate entries for linked params- is this good or not?
+        param_names = [
             n + "/" + k
             for m, n in zip(self.model, model_names)
             for k in m.params.keys()
@@ -618,3 +505,115 @@ def _finalize_model(self):
         self.hasJacobian = all([m.hasJacobian for m in self.model])
 
         self.nParams = self.x0.shape[0]
+
+    def _fit_single_pattern(
+        self,
+        data: np.ndarray,
+        initial_guess: np.ndarray,
+        fit_opts,
+    ):
+        """
+        Apply model fitting to one pattern.
+
+        Parameters
+        ----------
+        data: np.ndarray
+            Diffraction pattern
+        initial_guess: np.ndarray
+            starting guess for fitting
+        fit_opts:
+            args passed to scipy.optimize.least_squares
+
+        Returns
+        --------
+        fit_coefs: np.array
+            Fitted coefficients
+        fit_metrics: np.array
+            Fitting metrics
+
+        """
+
+        try:
+            if self.hasJacobian & self.use_jacobian:
+                opt = least_squares(
+                    self._pattern_error,
+                    initial_guess,
+                    jac=self._jacobian,
+                    bounds=(self.lower_bound, self.upper_bound),
+                    args=(data, self.static_data),
+                    **fit_opts,
+                )
+            else:
+                opt = least_squares(
+                    self._pattern_error,
+                    initial_guess,
+                    bounds=(self.lower_bound, self.upper_bound),
+                    args=(data, self.static_data),
+                    **fit_opts,
+                )
+
+            fit_coefs = opt.x
+            fit_metrics_single = [
+                opt.cost,
+                opt.optimality,
+                opt.nfev,
+                opt.status,
+            ]
+        except Exception as err:
+            print(err)
+            fit_coefs = initial_guess
+            fit_metrics_single = [0, 0, 0, 0]
+
+        return fit_coefs, fit_metrics_single
+
+    def _fit_distributed(
+        self,
+        fit_opts: dict,
+        fit_data:np.ndarray,
+        fit_metrics:np.ndarray,
+        resume=False,
+        num_jobs=None,
+    ):
+        from mpire import WorkerPool
+        from threadpoolctl import threadpool_limits
+
+        def f(shared_data, args):
+            with threadpool_limits(limits=1):
+                return self._fit_single_pattern(**args, fit_opts=shared_data)
+
+        fit_inputs = [
+            (
+                {
+                    "data": self.datacube[rx, ry],
+                    "initial_guess": self.fit_data[rx, ry] if resume else self.x0,
+                },
+            )
+            for rx in range(self.datacube.R_Nx)
+            for ry in range(self.datacube.R_Ny)
+        ]
+
+        with WorkerPool(
+            n_jobs=num_jobs,
+            shared_objects=fit_opts,
+        ) as pool:
+            results = pool.map(
+                f,
+                fit_inputs,
+                progress_bar=True,
+            )
+
+        for (rx,ry), res in zip(np.ndindex((self.datacube.R_Nx, self.datacube.R_Ny)), results):
+            fit_data[:,rx,ry] = res[0]
+            fit_metrics[:,rx,ry] = res[1]
+
+
+    def __getstate__(self):
+        # Prevent pickling from copying the datacube, so that distributed
+        # evaluation does not balloon memory usage.
+        # Copy the object's state from self.__dict__ which contains
+        # all our instance attributes. Always use the dict.copy()
+        # method to avoid modifying the original state.
+        state = self.__dict__.copy()
+        # Remove the unpicklable entries.
+        del state["datacube"]
+        return state
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index d6201a15c..9d6649e1b 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -21,7 +21,7 @@ def show_model_grid(self, x=None, **plot_kwargs):
     kwargs.update(plot_kwargs)
     fig, ax = plt.subplots(rows, cols, **kwargs)
 
-    for (a, m) in zip(ax.flat, model):
+    for a, m in zip(ax.flat, model):
         DP = np.zeros((self.datacube.Q_Nx, self.datacube.Q_Ny))
         m.func(DP, x, **self.static_data)
 

From c377e57bb6c5f9d8ddb699dc2b8d1f14229a931d Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Wed, 9 Aug 2023 16:23:17 -0400
Subject: [PATCH 41/57] linter fixes

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 5f1207ed0..a0cffae5f 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -225,7 +225,7 @@ def global_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["level"].offset]
 
-        r = WPF._get_distance(x, self.params["x center"], self.params["y center"])
+        r = kwargs['parent']._get_distance(x, self.params["x center"], self.params["y center"])
 
         DP += level * np.exp((r - radius) ** 2 / (-2 * sigma**2))
 
@@ -236,7 +236,7 @@ def global_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None
 
         x0 = x[self.params["x center"].offset]
         y0 = x[self.params["y center"].offset]
-        r = WPF._get_distance(x, self.params["x center"], self.params["y center"])
+        r = kwargs['parent']._get_distance(x, self.params["x center"], self.params["y center"])
 
         local_r = radius - r
         clipped_r = np.maximum(local_r, 0.1)
@@ -529,6 +529,8 @@ def __init__(
 
         #
 
+        params = {}
+
         super().__init__(
             name,
             params,

From da829ebadd6eaa91539a5da01f72f2cd301bf474 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Thu, 10 Aug 2023 09:28:30 -0400
Subject: [PATCH 42/57] fix how fit data is stored, improve distributed

---
 py4DSTEM/process/wholepatternfit/wp_models.py |  12 +-
 py4DSTEM/process/wholepatternfit/wpf.py       | 114 ++++++++++--------
 2 files changed, 69 insertions(+), 57 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index a0cffae5f..ac7ae447f 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -220,23 +220,27 @@ def __init__(
 
         super().__init__(name, params, model_type=WPFModelType.AMORPHOUS)
 
-    def global_center_func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
+    def func(self, DP: np.ndarray, x: np.ndarray, **kwargs) -> None:
         radius = x[self.params["radius"].offset]
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["level"].offset]
 
-        r = kwargs['parent']._get_distance(x, self.params["x center"], self.params["y center"])
+        r = kwargs["parent"]._get_distance(
+            x, self.params["x center"], self.params["y center"]
+        )
 
         DP += level * np.exp((r - radius) ** 2 / (-2 * sigma**2))
 
-    def global_center_jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
+    def jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
         radius = x[self.params["radius"].offset]
         sigma = x[self.params["sigma"].offset]
         level = x[self.params["level"].offset]
 
         x0 = x[self.params["x center"].offset]
         y0 = x[self.params["y center"].offset]
-        r = kwargs['parent']._get_distance(x, self.params["x center"], self.params["y center"])
+        r = kwargs["parent"]._get_distance(
+            x, self.params["x center"], self.params["y center"]
+        )
 
         local_r = radius - r
         clipped_r = np.maximum(local_r, 0.1)
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index f1d511f28..6979d1f6c 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -256,6 +256,7 @@ def fit_all_patterns(
         show_fit_metrics=True,
         distributed=True,
         num_jobs=None,
+        threads_per_job=1,
         **fit_opts,
     ):
         """
@@ -285,7 +286,7 @@ def fit_all_patterns(
         """
 
         # make sure we have the latest parameters
-        self._finalize_model()
+        unique_params, unique_names = self._finalize_model()
 
         # set tracking off
         self._track = False
@@ -306,15 +307,16 @@ def fit_all_patterns(
         }
         default_opts.update(fit_opts)
 
+        mask = real_space_mask or np.ones((self.datacube.R_Nx, self.datacube.R_Ny), dtype=bool)
+
         # Loop over probe positions
         if not distributed:
             for rx, ry in tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny):
-                if real_space_mask is not None and real_space_mask[rx, ry] == True:
+                if mask[rx,ry]:
                     current_pattern = (
-                        self.datacube.data[rx, ry, :, :].copy() * self.intensity_scale
+                        self.datacube.data[rx, ry, :, :] * self.intensity_scale
                     )
-                    shared_data = self.static_data.copy()
-                    x0 = self.fit_data.data[rx, ry].copy() if resume else self.x0.copy()
+                    x0 = self.fit_data.data[rx, ry] if resume else self.x0
 
                     try:
                         if self.hasJacobian & self.use_jacobian:
@@ -323,18 +325,16 @@ def fit_all_patterns(
                                 x0,
                                 jac=self._jacobian,
                                 bounds=(self.lower_bound, self.upper_bound),
-                                args=(current_pattern, shared_data),
+                                args=(current_pattern, self.static_data),
                                 **default_opts,
-                                # **fit_opts,
                             )
                         else:
                             opt = least_squares(
                                 self._pattern_error,
                                 x0,
                                 bounds=(self.lower_bound, self.upper_bound),
-                                args=(current_pattern, shared_data),
+                                args=(current_pattern, self.static_data),
                                 **default_opts,
-                                # **fit_opts,
                             )
 
                         fit_data_single = opt.x
@@ -344,7 +344,8 @@ def fit_all_patterns(
                             opt.nfev,
                             opt.status,
                         ]
-                    except:
+                    except Exception as err:
+                        # print(err)
                         fit_data_single = x0
                         fit_metrics_single = [0, 0, 0, 0]
 
@@ -356,30 +357,13 @@ def fit_all_patterns(
             self._fit_distributed(
                 resume=resume,
                 num_jobs=num_jobs,
+                threads_per_job=threads_per_job,
                 fit_opts=default_opts,
                 fit_data=fit_data,
                 fit_metrics=fit_metrics,
             )
 
-        # Convert to RealSlices
-        model_names = []
-        for m in self.model:
-            n = m.name
-            if n in model_names:
-                i = 1
-                while n in model_names:
-                    n = m.name + "_" + str(i)
-                    i += 1
-            model_names.append(n)
-
-        # TODO: this produces duplicate entries for linked params- is this good or not?
-        param_names = [
-            n + "/" + k
-            for m, n in zip(self.model, model_names)
-            for k in m.params.keys()
-        ]
-
-        self.fit_data = RealSlice(fit_data, name="Fit Data", slicelabels=param_names)
+        self.fit_data = RealSlice(fit_data, name="Fit Data", slicelabels=unique_names)
         self.fit_metrics = RealSlice(
             fit_metrics,
             name="Fit Metrics",
@@ -402,25 +386,26 @@ def get_lattice_maps(self):
         assert hasattr(self, "fit_data"), "Please run fitting first!"
 
         lattices = [
-            (i, m)
-            for i, m in enumerate(self.model)
+            m
+            for m in self.model
             if WPFModelType.LATTICE in m.model_type
         ]
 
         g_maps = []
-        for i, l in lattices:
-            # TODO: use parameter object offsets to get indices
-            param_list = list(l.params.keys())
-            lattice_offset = param_list.index("ux")
-            data_offset = self.model_param_inds[i] + 2 + lattice_offset
-
-            # TODO: Use proper RealSlice semantics for access
-            data = self.fit_data.data[:, :, data_offset : data_offset + 4]
+        for lat in lattices:
+            data = np.stack(
+                [
+                    self.fit_data.data[lat.params['ux'].offset],
+                    self.fit_data.data[lat.params['uy'].offset],
+                    self.fit_data.data[lat.params['vx'].offset],
+                    self.fit_data.data[lat.params['vy'].offset],
+                    np.ones(self.fit_data.data.shape[1:], dtype=np.bool_),
+                ],axis=0)
 
             g_map = RealSlice(
-                np.dstack((data, np.ones(data.shape[:2], dtype=np.bool_))),
+                data,
                 slicelabels=["g1x", "g1y", "g2x", "g2y", "mask"],
-                name=l.name,
+                name=lat.name,
             )
             g_maps.append(g_map)
 
@@ -489,12 +474,26 @@ def _jacobian(self, x, current_pattern, shared_data):
     def _finalize_model(self):
         # iterate over all models and assign indices, accumulate list
         # of unique parameters. then, accumulate initial value and bounds vectors
+
+        # get unique names for each model
+        model_names = []
+        for m in self.model:
+            n = m.name
+            if n in model_names:
+                i = 1
+                while n in model_names:
+                    n = m.name + "_" + str(i)
+                    i += 1
+            model_names.append(n)
+
         unique_params = []
+        unique_names = []
         idx = 0
-        for model in self.model:
-            for param in model.params.values():
+        for model, model_name in zip(self.model, model_names):
+            for param_name, param in model.params.items():
                 if param not in unique_params:
                     unique_params.append(param)
+                    unique_names.append(model_name + "/" + param_name)
                     param.offset = idx
                     idx += 1
 
@@ -506,6 +505,8 @@ def _finalize_model(self):
 
         self.nParams = self.x0.shape[0]
 
+        return unique_params, unique_names
+
     def _fit_single_pattern(
         self,
         data: np.ndarray,
@@ -540,7 +541,7 @@ def _fit_single_pattern(
                     initial_guess,
                     jac=self._jacobian,
                     bounds=(self.lower_bound, self.upper_bound),
-                    args=(data, self.static_data),
+                    args=(data*self.intensity_scale, self.static_data),
                     **fit_opts,
                 )
             else:
@@ -548,7 +549,7 @@ def _fit_single_pattern(
                     self._pattern_error,
                     initial_guess,
                     bounds=(self.lower_bound, self.upper_bound),
-                    args=(data, self.static_data),
+                    args=(data*self.intensity_scale, self.static_data),
                     **fit_opts,
                 )
 
@@ -569,18 +570,23 @@ def _fit_single_pattern(
     def _fit_distributed(
         self,
         fit_opts: dict,
-        fit_data:np.ndarray,
-        fit_metrics:np.ndarray,
+        fit_data: np.ndarray,
+        fit_metrics: np.ndarray,
         resume=False,
         num_jobs=None,
+        threads_per_job=1,
     ):
-        from mpire import WorkerPool
+        from mpire import WorkerPool, cpu_count
         from threadpoolctl import threadpool_limits
 
+        # prevent oversubscription when using multiple threads per job
+        num_jobs = num_jobs or cpu_count() // threads_per_job
+
         def f(shared_data, args):
-            with threadpool_limits(limits=1):
+            with threadpool_limits(limits=threads_per_job):
                 return self._fit_single_pattern(**args, fit_opts=shared_data)
 
+        # hopefully the data entries remain as views until dispatch time...
         fit_inputs = [
             (
                 {
@@ -592,6 +598,7 @@ def f(shared_data, args):
             for ry in range(self.datacube.R_Ny)
         ]
 
+        # TODO: auto set n_jobs when using multi threads each
         with WorkerPool(
             n_jobs=num_jobs,
             shared_objects=fit_opts,
@@ -602,10 +609,11 @@ def f(shared_data, args):
                 progress_bar=True,
             )
 
-        for (rx,ry), res in zip(np.ndindex((self.datacube.R_Nx, self.datacube.R_Ny)), results):
-            fit_data[:,rx,ry] = res[0]
-            fit_metrics[:,rx,ry] = res[1]
-
+        for (rx, ry), res in zip(
+            np.ndindex((self.datacube.R_Nx, self.datacube.R_Ny)), results
+        ):
+            fit_data[:, rx, ry] = res[0]
+            fit_metrics[:, rx, ry] = res[1]
 
     def __getstate__(self):
         # Prevent pickling from copying the datacube, so that distributed

From 5dc18dc2e781cff8cc632f4de3c33e4e4a696d45 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Thu, 10 Aug 2023 13:36:53 -0400
Subject: [PATCH 43/57] initialization of moire object built out

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 173 +++++++++++++++++-
 py4DSTEM/process/wholepatternfit/wpf.py       |  28 +--
 2 files changed, 181 insertions(+), 20 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index ac7ae447f..9addd0b13 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -516,8 +516,10 @@ class SyntheticDiskMoire(WPFModelPrototype):
 
     def __init__(
         self,
+        WPF,
         lattice_a: SyntheticDiskLattice,
         lattice_b: SyntheticDiskLattice,
+        intensity_0: float,
         decorated_peaks: list = None,
         name: str = "Moire Lattice",
     ):
@@ -528,12 +530,140 @@ def __init__(
         lattice_a, lattice_b: SyntheticDiskLattice
         """
         # ensure both models share the same center coordinate
+        if (lattice_a.params["x center"] is not lattice_b.params["x center"]) or (
+            lattice_a.params["y center"] is not lattice_b.params["y center"]
+        ):
+            raise ValueError(
+                "The center coordinates for each model must be linked, "
+                "either by passing global_center=True or linking after instantiation."
+            )
 
-        # pick the right pair of lattice vectors for generating the moire reciprocal lattice
+        self.lattice_a = lattice_a
+        self.lattice_b = lattice_b
+
+        # construct a 2x4 matrix "M" that transforms the parent lattices into
+        # the moire lattice vectors
+        lat_a = np.array(
+            [
+                [
+                    lattice_a.params["ux"].initial_value,
+                    lattice_a.params["uy"].initial_value,
+                ],
+                [
+                    lattice_a.params["vx"].initial_value,
+                    lattice_a.params["vy"].initial_value,
+                ],
+            ]
+        )
 
-        #
+        lat_b = np.array(
+            [
+                [
+                    lattice_b.params["ux"].initial_value,
+                    lattice_b.params["uy"].initial_value,
+                ],
+                [
+                    lattice_b.params["vx"].initial_value,
+                    lattice_b.params["vy"].initial_value,
+                ],
+            ]
+        )
 
-        params = {}
+        lat_ab = np.vstack((lat_a, lat_b))
+
+        # pick the pairing that gives the smallest unit cell
+        M_test = np.hstack((np.eye(2), -np.eye(2)))
+        M_test_flip = np.hstack((np.eye(2), -np.flipud(np.eye(2))))
+
+        M = (
+            M_test
+            if np.max(np.linalg.norm(M_test @ lat_ab, axis=1))
+            < np.max(np.linalg.norm(M_test_flip, axis=1))
+            else M_test_flip
+        )
+
+        # ensure the moire vectors are less 90 deg apart
+        if np.arccos(
+            ((M @ lat_ab)[0] @ (M @ lat_ab)[1])
+            / (np.linalg.norm((M @ lat_ab)[0]) * np.linalg.norm((M @ lat_ab)[1]))
+        ) > np.radians(90):
+            M[1] *= -1.0
+
+        # ensure they are right-handed
+        if np.cross(*(M @ lat_ab)) < 0.0:
+            M = np.flipud(np.eye(2)) @ M
+
+        # store moire construction
+        self.moire_matrix = M
+
+        # generate the indices of each peak, then find unique peaks
+        if decorated_peaks is not None:
+            decorated_peaks = np.array(decorated_peaks)
+            parent_peaks = np.vstack(
+                (
+                    np.concatenate(
+                        (decorated_peaks, np.zeros_like(decorated_peaks)), axis=1
+                    ),
+                    np.concatenate(
+                        (np.zeros_like(decorated_peaks), decorated_peaks), axis=1
+                    ),
+                )
+            )
+        else:
+            parent_peaks = np.vstack(
+                (
+                    np.concatenate(
+                        (
+                            np.stack((lattice_a.u_inds, lattice_a.v_inds), axis=1),
+                            np.zeros((lattice_a.u_inds.shape[0], 2)),
+                        ),
+                        axis=1,
+                    ),
+                    np.concatenate(
+                        (
+                            np.zeros((lattice_b.u_inds.shape[0], 2)),
+                            np.stack((lattice_b.u_inds, lattice_b.v_inds), axis=1),
+                        ),
+                        axis=1,
+                    ),
+                )
+            )
+
+        # trial indices for moire peaks
+        mx, my = np.mgrid[-1:2, -1:2]
+        moire_peaks = np.stack([mx.ravel(), my.ravel()], axis=1)[1:-1]
+
+        # construct a giant index array with columns a_h a_k b_h b_k m_h m_k
+        parent_expanded = np.zeros((parent_peaks.shape[0], 6))
+        parent_expanded[:, :4] = parent_peaks
+        moire_expanded = np.zeros((moire_peaks.shape[0], 6))
+        moire_expanded[:, 4:] = moire_peaks
+
+        all_indices = (
+            parent_expanded[:, None, :] + moire_expanded[None, :, :]
+        ).reshape(-1, 6)
+
+        lat_abm = np.vstack((lat_ab, M @ lat_ab))
+
+        all_peaks = all_indices @ lat_abm
+
+        _, idx_unique = np.unique(all_peaks, axis=0, return_index=True)
+
+        all_indices = all_indices[idx_unique]
+
+        # remove spots that coincide with primary peaks
+        parent_spots = parent_peaks @ lat_ab
+        self.moire_indices_uvm = np.array(
+            [idx for idx in all_indices if (idx @ lat_abm) not in parent_spots]
+        )
+
+        # each order of parent reflection has a separate moire intensity
+        max_order = np.max(np.abs(self.moire_indices_uvm[:, :4]))
+
+        params = {
+            f"Order {n} Moire Intensity": Parameter(intensity_0)
+            for n in range(max_order)
+        }
 
         super().__init__(
             name,
@@ -541,10 +671,41 @@ def __init__(
             model_type=WPFModelType.META,
         )
 
-    def func(self, DP, x, **static_data):
-        pass
+    def func(self, DP: np.ndarray, x: np.ndarray, **static_data):
+        # construct the moire unit cell from the current vectors
+        # of the two parent lattices
+        lat_a = np.array(
+            [
+                [
+                    x[self.lattice_a.params["ux"].offset],
+                    x[self.lattice_a.params["uy"].offset],
+                ],
+                [
+                    x[self.lattice_a.params["vx"].offset],
+                    x[self.lattice_a.params["vy"].offset],
+                ],
+            ]
+        )
+
+        lat_b = np.array(
+            [
+                [
+                    x[self.lattice_b.params["ux"].offset],
+                    x[self.lattice_b.params["uy"].offset],
+                ],
+                [
+                    x[self.lattice_b.params["vx"].offset],
+                    x[self.lattice_b.params["vy"].offset],
+                ],
+            ]
+        )
+
+        lat_ab = np.vstack((lat_a, lat_b))
+        lat_abm = np.vstack((lat_ab, self.moire_matrix @ lat_ab))
+
+        # moire = self.moire_matrix @ lat_ab
 
-    def jacobian(self, J, x, **static_data):
+    def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
         pass
 
 
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 6979d1f6c..af6ac3d15 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -307,12 +307,14 @@ def fit_all_patterns(
         }
         default_opts.update(fit_opts)
 
-        mask = real_space_mask or np.ones((self.datacube.R_Nx, self.datacube.R_Ny), dtype=bool)
+        mask = real_space_mask or np.ones(
+            (self.datacube.R_Nx, self.datacube.R_Ny), dtype=bool
+        )
 
         # Loop over probe positions
         if not distributed:
             for rx, ry in tqdmnd(self.datacube.R_Nx, self.datacube.R_Ny):
-                if mask[rx,ry]:
+                if mask[rx, ry]:
                     current_pattern = (
                         self.datacube.data[rx, ry, :, :] * self.intensity_scale
                     )
@@ -385,22 +387,20 @@ def accept_mean_CBED_fit(self):
     def get_lattice_maps(self):
         assert hasattr(self, "fit_data"), "Please run fitting first!"
 
-        lattices = [
-            m
-            for m in self.model
-            if WPFModelType.LATTICE in m.model_type
-        ]
+        lattices = [m for m in self.model if WPFModelType.LATTICE in m.model_type]
 
         g_maps = []
         for lat in lattices:
             data = np.stack(
                 [
-                    self.fit_data.data[lat.params['ux'].offset],
-                    self.fit_data.data[lat.params['uy'].offset],
-                    self.fit_data.data[lat.params['vx'].offset],
-                    self.fit_data.data[lat.params['vy'].offset],
+                    self.fit_data.data[lat.params["ux"].offset],
+                    self.fit_data.data[lat.params["uy"].offset],
+                    self.fit_data.data[lat.params["vx"].offset],
+                    self.fit_data.data[lat.params["vy"].offset],
                     np.ones(self.fit_data.data.shape[1:], dtype=np.bool_),
-                ],axis=0)
+                ],
+                axis=0,
+            )
 
             g_map = RealSlice(
                 data,
@@ -541,7 +541,7 @@ def _fit_single_pattern(
                     initial_guess,
                     jac=self._jacobian,
                     bounds=(self.lower_bound, self.upper_bound),
-                    args=(data*self.intensity_scale, self.static_data),
+                    args=(data * self.intensity_scale, self.static_data),
                     **fit_opts,
                 )
             else:
@@ -549,7 +549,7 @@ def _fit_single_pattern(
                     self._pattern_error,
                     initial_guess,
                     bounds=(self.lower_bound, self.upper_bound),
-                    args=(data*self.intensity_scale, self.static_data),
+                    args=(data * self.intensity_scale, self.static_data),
                     **fit_opts,
                 )
 

From c88f7e312655d0122a9fddd406ed2ade4b0f87de Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Fri, 11 Aug 2023 16:59:29 -0400
Subject: [PATCH 44/57] moire generates pattern but derivatives are not
 complete

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 244 ++++++++++++++----
 py4DSTEM/process/wholepatternfit/wpf_viz.py   |  27 +-
 2 files changed, 217 insertions(+), 54 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 9addd0b13..0069783c4 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -15,6 +15,7 @@ class WPFModelType(Flag):
 
     AMORPHOUS = auto()
     LATTICE = auto()
+    MOIRE = auto()
 
     DUMMY = auto()  # Model has no direct contribution to pattern
     META = auto()  # Model depends on multiple sub-Models
@@ -521,6 +522,11 @@ def __init__(
         lattice_b: SyntheticDiskLattice,
         intensity_0: float,
         decorated_peaks: list = None,
+        link_disk_parameters: bool = True,
+        refine_width: bool = True,
+        edge_width: list = None,
+        refine_radius: bool = True,
+        disk_radius: list = None,
         name: str = "Moire Lattice",
     ):
         """
@@ -543,44 +549,23 @@ def __init__(
 
         # construct a 2x4 matrix "M" that transforms the parent lattices into
         # the moire lattice vectors
-        lat_a = np.array(
-            [
-                [
-                    lattice_a.params["ux"].initial_value,
-                    lattice_a.params["uy"].initial_value,
-                ],
-                [
-                    lattice_a.params["vx"].initial_value,
-                    lattice_a.params["vy"].initial_value,
-                ],
-            ]
-        )
 
-        lat_b = np.array(
+        lat_ab = self._get_parent_lattices(lattice_a, lattice_b)
+        
+        # pick the pairing that gives the smallest unit cell
+        mx, my = np.mgrid[-1:2, -1:2]
+        test_peaks = np.stack([mx.ravel(), my.ravel()], axis=1)
+        tests = np.stack(
             [
-                [
-                    lattice_b.params["ux"].initial_value,
-                    lattice_b.params["uy"].initial_value,
-                ],
-                [
-                    lattice_b.params["vx"].initial_value,
-                    lattice_b.params["vy"].initial_value,
-                ],
-            ]
+                np.hstack((np.eye(2), np.vstack((b1, b2))))
+                for b1 in test_peaks
+                for b2 in test_peaks
+                if not np.allclose(b1, b2)
+            ],
+            axis=0,
         )
 
-        lat_ab = np.vstack((lat_a, lat_b))
-
-        # pick the pairing that gives the smallest unit cell
-        M_test = np.hstack((np.eye(2), -np.eye(2)))
-        M_test_flip = np.hstack((np.eye(2), -np.flipud(np.eye(2))))
-
-        M = (
-            M_test
-            if np.max(np.linalg.norm(M_test @ lat_ab, axis=1))
-            < np.max(np.linalg.norm(M_test_flip, axis=1))
-            else M_test_flip
-        )
+        M = tests[np.argmin(np.max(np.linalg.norm(tests @ lat_ab, axis=-1), axis=-1))]
 
         # ensure the moire vectors are less 90 deg apart
         if np.arccos(
@@ -658,31 +643,50 @@ def __init__(
         )
 
         # each order of parent reflection has a separate moire intensity
-        max_order = np.max(np.abs(self.moire_indices_uvm[:, :4]))
+        max_order = int(np.max(np.abs(self.moire_indices_uvm[:, :4])))
 
         params = {
             f"Order {n} Moire Intensity": Parameter(intensity_0)
-            for n in range(max_order)
+            for n in range(max_order + 1)
         }
 
+        params["x center"] = lattice_a.params["x center"]
+        params["y center"] = lattice_a.params["y center"]
+
+        # add disk edge and width parameters if needed
+        if link_disk_parameters:
+            if (lattice_a.refine_width) and (lattice_b.refine_width):
+                self.refine_width = True
+                params["edge width"] = lattice_a.params["edge width"]
+            if (lattice_a.refine_radius) and (lattice_b.refine_radius):
+                self.refine_radius = True
+                params["disk radius"] = lattice_a.params["disk radius"]
+        else:
+            self.refine_width = refine_width
+            if self.refine_width:
+                params["edge width"] = Parameter(edge_width)
+
+            self.refine_radius = refine_radius
+            if self.refine_radius:
+                params["disk radius"] = Parameter(disk_radius)
+
         super().__init__(
             name,
             params,
-            model_type=WPFModelType.META,
+            model_type=WPFModelType.META | WPFModelType.MOIRE,
         )
 
-    def func(self, DP: np.ndarray, x: np.ndarray, **static_data):
-        # construct the moire unit cell from the current vectors
-        # of the two parent lattices
+    def _get_parent_lattices(self, lattice_a, lattice_b):
+
         lat_a = np.array(
             [
                 [
-                    x[self.lattice_a.params["ux"].offset],
-                    x[self.lattice_a.params["uy"].offset],
+                    lattice_a.params["ux"].initial_value,
+                    lattice_a.params["uy"].initial_value,
                 ],
                 [
-                    x[self.lattice_a.params["vx"].offset],
-                    x[self.lattice_a.params["vy"].offset],
+                    lattice_a.params["vx"].initial_value,
+                    lattice_a.params["vy"].initial_value,
                 ],
             ]
         )
@@ -690,23 +694,161 @@ def func(self, DP: np.ndarray, x: np.ndarray, **static_data):
         lat_b = np.array(
             [
                 [
-                    x[self.lattice_b.params["ux"].offset],
-                    x[self.lattice_b.params["uy"].offset],
+                    lattice_b.params["ux"].initial_value,
+                    lattice_b.params["uy"].initial_value,
                 ],
                 [
-                    x[self.lattice_b.params["vx"].offset],
-                    x[self.lattice_b.params["vy"].offset],
+                    lattice_b.params["vx"].initial_value,
+                    lattice_b.params["vy"].initial_value,
                 ],
             ]
         )
 
-        lat_ab = np.vstack((lat_a, lat_b))
+        return np.vstack((lat_a, lat_b))
+
+    def func(self, DP: np.ndarray, x: np.ndarray, **static_data):
+        # construct the moire unit cell from the current vectors
+        # of the two parent lattices
+
+        lat_ab = self._get_parent_lattices(self.lattice_a, self.lattice_b)
         lat_abm = np.vstack((lat_ab, self.moire_matrix @ lat_ab))
 
-        # moire = self.moire_matrix @ lat_ab
+        # grab shared parameters
+        disk_radius = (
+            x[self.params["disk radius"].offset]
+            if self.refine_radius
+            else self.disk_radius
+        )
+
+        disk_width = (
+            x[self.params["edge width"].offset]
+            if self.refine_width
+            else self.disk_width
+        )
+
+        # compute positions of each moire peak
+        positions = self.moire_indices_uvm @ lat_abm
+        positions += np.array(
+            [x[self.params["x center"].offset], x[self.params["y center"].offset]]
+        )
+
+        for (x_pos, y_pos), indices in zip(positions, self.moire_indices_uvm):
+            # Each peak has an intensity based on the max index of parent lattice
+            # which it decorates
+            order = int(np.max(np.abs(indices[:4])))
+            intensity = x[self.params[f"Order {order} Moire Intensity"].offset]
+
+            DP += intensity / (
+                1.0
+                + np.exp(
+                    np.minimum(
+                        4
+                        * (
+                            np.sqrt(
+                                (static_data["xArray"] - x_pos) ** 2
+                                + (static_data["yArray"] - y_pos) ** 2
+                            )
+                            - disk_radius
+                        )
+                        / disk_width,
+                        20,
+                    )
+                )
+            )
 
     def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
-        pass
+        # construct the moire unit cell from the current vectors
+        # of the two parent lattices
+        lat_ab = self._get_parent_lattices(self.lattice_a, self.lattice_b)
+        lat_abm = np.vstack((lat_ab, self.moire_matrix @ lat_ab))
+
+        # grab shared parameters
+        disk_radius = (
+            x[self.params["disk radius"].offset]
+            if self.refine_radius
+            else self.disk_radius
+        )
+
+        disk_width = (
+            x[self.params["edge width"].offset]
+            if self.refine_width
+            else self.disk_width
+        )
+
+        # compute positions of each moire peak
+        positions = self.moire_indices_uvm @ lat_abm
+        positions += np.array(
+            [x[self.params["x center"].offset], x[self.params["y center"].offset]]
+        )
+
+        for (x_pos, y_pos), indices in zip(positions, self.moire_indices_uvm):
+            # Each peak has an intensity based on the max index of parent lattice
+            # which it decorates
+            order = int(np.max(np.abs(indices[:4])))
+            intensity_idx = self.params[f"Order {order} Moire Intensity"].offset
+            intensity = x[intensity_idx]
+
+            r_disk = np.maximum(
+                5e-1,
+                np.sqrt(
+                    (static_data["xArray"] - x_pos) ** 2
+                    + (static_data["yArray"] - y_pos) ** 2
+                ),
+            )
+
+            mask = r_disk < (2 * disk_radius)
+
+            top_exp = mask * np.exp(4 * ((mask * r_disk) - disk_radius) / disk_width)
+
+            # dF/d(x0)
+            dx = (
+                4
+                * disk_intensity
+                * (static_data["xArray"] - x_pos)
+                * top_exp
+                / ((1.0 + top_exp) ** 2 * disk_width * r)
+            ).ravel()
+
+            # dF/d(y0)
+            dy = (
+                4
+                * disk_intensity
+                * (static_data["yArray"] - y_pos)
+                * top_exp
+                / ((1.0 + top_exp) ** 2 * disk_width * r)
+            ).ravel()
+
+            # insert center position derivatives
+            J[:, self.params["x center"].offset] += disk_intensity * dx
+            J[:, self.params["y center"].offset] += disk_intensity * dy
+
+            # insert lattice vector derivatives
+            # TODO: these need to be scaled for chain rule!!!
+            J[:, self.params["ux"].offset] += disk_intensity * u * dx
+            J[:, self.params["uy"].offset] += disk_intensity * u * dy
+            J[:, self.params["vx"].offset] += disk_intensity * v * dx
+            J[:, self.params["vy"].offset] += disk_intensity * v * dy
+
+            # insert intensity derivative
+            dI = (mask * (1.0 / (1.0 + top_exp))).ravel()
+            J[:, intensity_idx] += dI
+
+            # insert disk radius derivative
+            if self.refine_radius:
+                dR = (
+                    4.0 * disk_intensity * top_exp / (disk_width * (1.0 + top_exp) ** 2)
+                ).ravel()
+                J[:, self.params["disk radius"].offset] += dR
+
+            if self.refine_width:
+                dW = (
+                    4.0
+                    * disk_intensity
+                    * top_exp
+                    * (r_disk - disk_radius)
+                    / (disk_width**2 * (1.0 + top_exp) ** 2)
+                ).ravel()
+                J[:, self.params["edge width"].offset] += dW
 
 
 class ComplexOverlapKernelDiskLattice(WPFModelPrototype):
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index 9d6649e1b..b853b8f90 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -9,7 +9,7 @@
 
 
 def show_model_grid(self, x=None, **plot_kwargs):
-    x = x or self.mean_CBED_fit.x
+    x = self.mean_CBED_fit.x if x is None else x
 
     model = [m for m in self.model if WPFModelType.DUMMY not in m.model_type]
 
@@ -91,8 +91,8 @@ def show_lattice_points(
         inds = np.stack([m.u_inds, m.v_inds], axis=1)
 
         spots = inds @ lat
-        spots[:, 0] += self.coordinate_model.params["x center"].initial_value
-        spots[:, 1] += self.coordinate_model.params["y center"].initial_value
+        spots[:, 0] += m.params["x center"].initial_value
+        spots[:, 1] += m.params["y center"].initial_value
 
         ax.scatter(
             spots[:, 1],
@@ -102,8 +102,29 @@ def show_lattice_points(
             label=m.name,
         )
 
+    moires = [m for m in self.model if WPFModelType.MOIRE in m.model_type]
+
+    for m in moires:
+        lat_ab = m._get_parent_lattices(m.lattice_a, m.lattice_b)
+        lat_abm = np.vstack((lat_ab, m.moire_matrix @ lat_ab))
+
+        spots = m.moire_indices_uvm @ lat_abm
+        spots[:, 0] += m.params["x center"].initial_value
+        spots[:, 1] += m.params["y center"].initial_value
+
+        ax.scatter(
+            spots[:, 1],
+            spots[:, 0],
+            s=100,
+            marker="+",
+            label=m.name,
+        )
+
     ax.legend()
 
+    ax.set_xlim(0,im.shape[1]-1)
+    ax.set_ylim(im.shape[0]-1,0)
+
     return (fig, ax) if returnfig else plt.show()
 
 

From 64dde4dba6e1b9493aa69a71cd1b004f9dc18143 Mon Sep 17 00:00:00 2001
From: cophus <cophus@gmail.com>
Date: Sun, 13 Aug 2023 15:32:49 -0700
Subject: [PATCH 45/57] Fix for the mask in real space

---
 py4DSTEM/process/wholepatternfit/wpf.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index af6ac3d15..d5b3b25aa 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -307,9 +307,14 @@ def fit_all_patterns(
         }
         default_opts.update(fit_opts)
 
-        mask = real_space_mask or np.ones(
-            (self.datacube.R_Nx, self.datacube.R_Ny), dtype=bool
-        )
+        # Masking function
+        if mask is None:
+            mask = np.ones(
+                (self.datacube.R_Nx, self.datacube.R_Ny), 
+                dtype=bool,
+            )
+        else:
+            mask = real_space_mask
 
         # Loop over probe positions
         if not distributed:

From f4f3f42bbe8b2631138c9607d39f7a28d6788e5b Mon Sep 17 00:00:00 2001
From: cophus <cophus@gmail.com>
Date: Sun, 13 Aug 2023 15:34:37 -0700
Subject: [PATCH 46/57] Typo

---
 py4DSTEM/process/wholepatternfit/wpf.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index d5b3b25aa..ae9b017a1 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -308,7 +308,7 @@ def fit_all_patterns(
         default_opts.update(fit_opts)
 
         # Masking function
-        if mask is None:
+        if real_space_mask is None:
             mask = np.ones(
                 (self.datacube.R_Nx, self.datacube.R_Ny), 
                 dtype=bool,

From 4441083f62236ba95d70ce8d878ff17ef5aae04b Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Mon, 14 Aug 2023 11:09:04 -0400
Subject: [PATCH 47/57] fixes for moire, derivative now working

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 53 ++++++++++++++++---
 py4DSTEM/process/wholepatternfit/wpf_viz.py   | 15 ++++--
 2 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 0069783c4..18b6f2aee 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -551,7 +551,7 @@ def __init__(
         # the moire lattice vectors
 
         lat_ab = self._get_parent_lattices(lattice_a, lattice_b)
-        
+
         # pick the pairing that gives the smallest unit cell
         mx, my = np.mgrid[-1:2, -1:2]
         test_peaks = np.stack([mx.ravel(), my.ravel()], axis=1)
@@ -636,6 +636,22 @@ def __init__(
 
         all_indices = all_indices[idx_unique]
 
+        # remove peaks outside of pattern
+        Q_Nx = WPF.static_data["Q_Nx"]
+        Q_Ny = WPF.static_data["Q_Ny"]
+        all_peaks = all_indices @ lat_abm
+        all_peaks[:,0] += lattice_a.params['x center'].initial_value
+        all_peaks[:,1] += lattice_a.params['y center'].initial_value
+        delete_mask = np.logical_or.reduce(
+            [
+                all_peaks[:, 0] < 0.0,
+                all_peaks[:, 0] >= Q_Nx,
+                all_peaks[:, 1] < 0.0,
+                all_peaks[:, 1] >= Q_Ny,
+            ]
+        )
+        all_indices = all_indices[~delete_mask]
+
         # remove spots that coincide with primary peaks
         parent_spots = parent_peaks @ lat_ab
         self.moire_indices_uvm = np.array(
@@ -670,6 +686,22 @@ def __init__(
             if self.refine_radius:
                 params["disk radius"] = Parameter(disk_radius)
 
+        # store some data that helps compute the derivatives
+        selector_matrices = np.eye(8).reshape(-1, 4, 2)
+        selector_parameters = [
+            self.lattice_a.params["ux"],
+            self.lattice_a.params["uy"],
+            self.lattice_a.params["vx"],
+            self.lattice_a.params["vy"],
+            self.lattice_b.params["ux"],
+            self.lattice_b.params["uy"],
+            self.lattice_b.params["vx"],
+            self.lattice_b.params["vy"],
+        ]
+        self.parent_vector_selectors = [
+            (p, m) for p, m in zip(selector_parameters, selector_matrices)
+        ]
+
         super().__init__(
             name,
             params,
@@ -677,7 +709,6 @@ def __init__(
         )
 
     def _get_parent_lattices(self, lattice_a, lattice_b):
-
         lat_a = np.array(
             [
                 [
@@ -775,6 +806,11 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
             else self.disk_width
         )
 
+        # distance from center coordinate
+        r = np.maximum(
+            5e-1, static_data['parent']._get_distance(x, self.params["x center"], self.params["y center"])
+        )
+
         # compute positions of each moire peak
         positions = self.moire_indices_uvm @ lat_abm
         positions += np.array(
@@ -786,7 +822,7 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
             # which it decorates
             order = int(np.max(np.abs(indices[:4])))
             intensity_idx = self.params[f"Order {order} Moire Intensity"].offset
-            intensity = x[intensity_idx]
+            disk_intensity = x[intensity_idx]
 
             r_disk = np.maximum(
                 5e-1,
@@ -823,11 +859,12 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
             J[:, self.params["y center"].offset] += disk_intensity * dy
 
             # insert lattice vector derivatives
-            # TODO: these need to be scaled for chain rule!!!
-            J[:, self.params["ux"].offset] += disk_intensity * u * dx
-            J[:, self.params["uy"].offset] += disk_intensity * u * dy
-            J[:, self.params["vx"].offset] += disk_intensity * v * dx
-            J[:, self.params["vy"].offset] += disk_intensity * v * dy
+            for par, mat in self.parent_vector_selectors:
+                # find the x and y derivatives of the position of this
+                # disk in terms of each of the parent lattice vectors
+                d_abm = np.vstack((mat, self.moire_matrix @ mat))
+                d_param = indices @ d_abm
+                J[:, par.offset] += disk_intensity * (d_param[0] * dx + d_param[1] * dy)
 
             # insert intensity derivative
             dI = (mask * (1.0 / (1.0 + top_exp))).ravel()
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index b853b8f90..f510ca804 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -56,7 +56,15 @@ def show_model_grid(self, x=None, **plot_kwargs):
 
 
 def show_lattice_points(
-    self, im=None, vmin=None, vmax=None, power=None, returnfig=False, *args, **kwargs
+    self,
+    im=None,
+    vmin=None,
+    vmax=None,
+    power=None,
+    crop_to_pattern=False,
+    returnfig=False,
+    *args,
+    **kwargs,
 ):
     """
     Plotting utility to show the initial lattice points.
@@ -122,8 +130,9 @@ def show_lattice_points(
 
     ax.legend()
 
-    ax.set_xlim(0,im.shape[1]-1)
-    ax.set_ylim(im.shape[0]-1,0)
+    if crop_to_pattern:
+        ax.set_xlim(0, im.shape[1] - 1)
+        ax.set_ylim(im.shape[0] - 1, 0)
 
     return (fig, ax) if returnfig else plt.show()
 

From 087cd38120949786e982a85f25f86815815a8cd1 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Mon, 14 Aug 2023 12:17:01 -0400
Subject: [PATCH 48/57] allow each moire disk to have separate intensity

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 60 +++++++++++++++----
 py4DSTEM/process/wholepatternfit/wpf.py       |  2 +-
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 18b6f2aee..11fa274da 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -522,11 +522,13 @@ def __init__(
         lattice_b: SyntheticDiskLattice,
         intensity_0: float,
         decorated_peaks: list = None,
+        link_moire_disk_intensities: bool = False,
         link_disk_parameters: bool = True,
         refine_width: bool = True,
         edge_width: list = None,
         refine_radius: bool = True,
         disk_radius: list = None,
+        lattice_b_search_range: int = 1,
         name: str = "Moire Lattice",
     ):
         """
@@ -553,7 +555,8 @@ def __init__(
         lat_ab = self._get_parent_lattices(lattice_a, lattice_b)
 
         # pick the pairing that gives the smallest unit cell
-        mx, my = np.mgrid[-1:2, -1:2]
+        imax = lattice_b_search_range
+        mx, my = np.mgrid[-imax : imax + 1, -imax : imax + 1]
         test_peaks = np.stack([mx.ravel(), my.ravel()], axis=1)
         tests = np.stack(
             [
@@ -565,6 +568,10 @@ def __init__(
             axis=0,
         )
 
+        # pick the combination that gives the smallest cell
+        # (it might be better to choose the cell that gives the smallest volume
+        # but then we have to handle the case of nearly-parallel vectors in
+        # the search space)
         M = tests[np.argmin(np.max(np.linalg.norm(tests @ lat_ab, axis=-1), axis=-1))]
 
         # ensure the moire vectors are less 90 deg apart
@@ -640,8 +647,8 @@ def __init__(
         Q_Nx = WPF.static_data["Q_Nx"]
         Q_Ny = WPF.static_data["Q_Ny"]
         all_peaks = all_indices @ lat_abm
-        all_peaks[:,0] += lattice_a.params['x center'].initial_value
-        all_peaks[:,1] += lattice_a.params['y center'].initial_value
+        all_peaks[:, 0] += lattice_a.params["x center"].initial_value
+        all_peaks[:, 1] += lattice_a.params["y center"].initial_value
         delete_mask = np.logical_or.reduce(
             [
                 all_peaks[:, 0] < 0.0,
@@ -658,13 +665,22 @@ def __init__(
             [idx for idx in all_indices if (idx @ lat_abm) not in parent_spots]
         )
 
-        # each order of parent reflection has a separate moire intensity
-        max_order = int(np.max(np.abs(self.moire_indices_uvm[:, :4])))
+        self.link_moire_disk_intensities = link_moire_disk_intensities
+        if link_moire_disk_intensities:
+            # each order of parent reflection has a separate moire intensity
+            max_order = int(np.max(np.abs(self.moire_indices_uvm[:, :4])))
 
-        params = {
-            f"Order {n} Moire Intensity": Parameter(intensity_0)
-            for n in range(max_order + 1)
-        }
+            params = {
+                f"Order {n} Moire Intensity": Parameter(intensity_0)
+                for n in range(max_order + 1)
+            }
+        else:
+            params = {
+                f"a ({ax},{ay}), b ({bx},{by}), moire ({mx},{my}) Intensity": Parameter(
+                    intensity_0
+                )
+                for ax, ay, bx, by, mx, my in self.moire_indices_uvm
+            }
 
         params["x center"] = lattice_a.params["x center"]
         params["y center"] = lattice_a.params["y center"]
@@ -767,7 +783,16 @@ def func(self, DP: np.ndarray, x: np.ndarray, **static_data):
             # Each peak has an intensity based on the max index of parent lattice
             # which it decorates
             order = int(np.max(np.abs(indices[:4])))
-            intensity = x[self.params[f"Order {order} Moire Intensity"].offset]
+
+            if self.link_moire_disk_intensities:
+                intensity = x[self.params[f"Order {order} Moire Intensity"].offset]
+            else:
+                ax, ay, bx, by, mx, my = indices
+                intensity = x[
+                    self.params[
+                        f"a ({ax},{ay}), b ({bx},{by}), moire ({mx},{my}) Intensity"
+                    ].offset
+                ]
 
             DP += intensity / (
                 1.0
@@ -808,7 +833,10 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
 
         # distance from center coordinate
         r = np.maximum(
-            5e-1, static_data['parent']._get_distance(x, self.params["x center"], self.params["y center"])
+            5e-1,
+            static_data["parent"]._get_distance(
+                x, self.params["x center"], self.params["y center"]
+            ),
         )
 
         # compute positions of each moire peak
@@ -820,8 +848,14 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
         for (x_pos, y_pos), indices in zip(positions, self.moire_indices_uvm):
             # Each peak has an intensity based on the max index of parent lattice
             # which it decorates
-            order = int(np.max(np.abs(indices[:4])))
-            intensity_idx = self.params[f"Order {order} Moire Intensity"].offset
+            if self.link_moire_disk_intensities:
+                order = int(np.max(np.abs(indices[:4])))
+                intensity_idx = self.params[f"Order {order} Moire Intensity"].offset
+            else:
+                ax, ay, bx, by, mx, my = indices
+                intensity_idx = self.params[
+                    f"a ({ax},{ay}), b ({bx},{by}), moire ({mx},{my}) Intensity"
+                ].offset
             disk_intensity = x[intensity_idx]
 
             r_disk = np.maximum(
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index ae9b017a1..94bbea7b8 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -310,7 +310,7 @@ def fit_all_patterns(
         # Masking function
         if real_space_mask is None:
             mask = np.ones(
-                (self.datacube.R_Nx, self.datacube.R_Ny), 
+                (self.datacube.R_Nx, self.datacube.R_Ny),
                 dtype=bool,
             )
         else:

From fef118c9a2b6eab82e08915c75f50d67020ec2db Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Mon, 14 Aug 2023 12:20:07 -0400
Subject: [PATCH 49/57] remove fit gamma scaling

---
 py4DSTEM/process/wholepatternfit/wpf.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 94bbea7b8..518647bdb 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -34,7 +34,6 @@ def __init__(
         mask: Optional[np.ndarray] = None,
         use_jacobian: bool = True,
         meanCBED: Optional[np.ndarray] = None,
-        fit_power: float = 1,
     ):
         """
         Perform pixelwise fits using composable models and numerical optimization.
@@ -70,8 +69,6 @@ def __init__(
         meanCBED: Optional np.ndarray, used to specify the diffraction pattern used
             for initial refinement of the parameters. If not specified, the average across
             all scan positions is computed
-        fit_power: float, diffraction patterns are raised to this power, sets the gamma
-            level at which the patterns are compared
 
         """
         self.datacube = datacube
@@ -120,8 +117,6 @@ def __init__(
         # set up the global arguments
         self._setup_static_data(x0, y0)
 
-        self.fit_power = fit_power
-
         # for debugging: tracks all function evals
         self._track = False
         self._fevals = []
@@ -449,7 +444,7 @@ def _get_distance(self, params: np.ndarray, x: Parameter, y: Parameter):
     def _pattern_error(self, x, current_pattern, shared_data):
         DP = self._pattern(x, shared_data)
 
-        DP = (DP - current_pattern**self.fit_power) * self.mask
+        DP = (DP - current_pattern) * self.mask
 
         if self._track:
             self._fevals.append(DP)
@@ -464,7 +459,7 @@ def _pattern(self, x, shared_data):
         for m in self.model:
             m.func(DP, x, **shared_data)
 
-        return (DP**self.fit_power) * self.mask
+        return DP * self.mask
 
     def _jacobian(self, x, current_pattern, shared_data):
         # TODO: automatic mixed analytic/finite difference

From fa6934a99ee61c9b324e455251b8ab34dafbe1cf Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Mon, 14 Aug 2023 13:17:46 -0400
Subject: [PATCH 50/57] add docstrings all around

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 214 ++++++++++++++++--
 py4DSTEM/process/wholepatternfit/wpf.py       | 196 +++++++++++-----
 setup.py                                      |   1 +
 3 files changed, 330 insertions(+), 81 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 11fa274da..3dfdfca2a 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -21,7 +21,7 @@ class WPFModelType(Flag):
     META = auto()  # Model depends on multiple sub-Models
 
 
-class WPFModelPrototype:
+class WPFModel:
     """
     Prototype class for a compent of a whole-pattern model.
     Holds the following:
@@ -71,6 +71,14 @@ def __init__(
         lower_bound: Optional[float] = None,
         upper_bound: Optional[float] = None,
     ):
+        """
+        Object representing a fitting parameter with bounds.
+
+        Can be specified three ways:
+        Parameter(initial_value) - Unbounded, with an initial guess
+        Parameter(initial_value, deviation) - Bounded within deviation of initial_guess
+        Parameter(initial_value, lower_bound, upper_bound) - Both bounds specified
+        """
         if hasattr(initial_value, "__iter__"):
             if len(initial_value) == 2:
                 initial_value = (
@@ -104,9 +112,11 @@ def __repr__(self):
         return f"Value: {self.initial_value} (Range: {self.lower_bound},{self.upper_bound})"
 
 
-class _BaseModel(WPFModelPrototype):
+class _BaseModel(WPFModel):
     """
-    Model object used by the WPF class as a container for the global Parameters
+    Model object used by the WPF class as a container for the global Parameters.
+
+    **This object should not be instantiated directly.**
     """
 
     def __init__(self, x0, y0, name="Globals"):
@@ -121,7 +131,19 @@ def jacobian(self, J: np.ndarray, *args, **kwargs) -> None:
         pass
 
 
-class DCBackground(WPFModelPrototype):
+class DCBackground(WPFModel):
+    """
+    Model representing constant background intensity.
+
+    Parameters
+    ----------
+    background_value
+        Background intensity value.
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    """
+
     def __init__(self, background_value=0.0, name="DC Background"):
         params = {"DC Level": Parameter(background_value)}
 
@@ -134,7 +156,34 @@ def jacobian(self, J: np.ndarray, *args, **kwargs):
         J[:, self.params["DC Level"].offset] = 1
 
 
-class GaussianBackground(WPFModelPrototype):
+class GaussianBackground(WPFModel):
+    """
+    Model representing a 2D Gaussian intensity distribution
+
+    Parameters
+    ----------
+    WPF: WholePatternFit
+        Parent WPF object
+    sigma
+        parameter specifying width of the Gaussian
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    intensity
+        parameter specifying intensity of the Gaussian
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    global_center: bool
+        If True, uses same center coordinate as the global model
+        If False, uses an independent center
+    x0, y0:
+        Center coordinates of model for local origin
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    """
+
     def __init__(
         self,
         WPF,
@@ -195,7 +244,39 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
         J[:, self.params["intensity"].offset] += exp_expr.ravel()
 
 
-class GaussianRing(WPFModelPrototype):
+class GaussianRing(WPFModel):
+    """
+    Model representing a halo with Gaussian falloff
+
+    Parameters
+    ----------
+    WPF: WholePatternFit
+        parent fitting object
+    radius:
+        radius of halo
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    sigma:
+        width of Gaussian falloff
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    intensity:
+        Intensity of the halo
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    global_center: bool
+        If True, uses same center coordinate as the global model
+        If False, uses an independent center
+    x0, y0:
+        Center coordinates of model for local origin
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    """
+
     def __init__(
         self,
         WPF,
@@ -280,7 +361,57 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
         J[:, self.params["intensity"].offset] += exp_expr.ravel()
 
 
-class SyntheticDiskLattice(WPFModelPrototype):
+class SyntheticDiskLattice(WPFModel):
+    """
+    Model representing a lattice of diffraction disks with a soft edge
+
+    Parameters
+    ----------
+
+    WPF: WholePatternFit
+        parent fitting object
+    ux,uy,vx,vy
+        x and y components of the lattice vectors u and v.
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    disk_radius
+        Radius of each diffraction disk.
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    disk_width
+        Width of the smooth falloff at the edge of the disk
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    u_max, v_max
+        Maximum lattice indices to include in the pattern.
+        Disks outside the pattern are automatically clipped.
+    intensity_0
+        Initial intensity for each diffraction disk.
+        Each disk intensity is an independent fit variable in the final model
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    refine_radius: bool
+        Flag whether disk radius is made a fitting parameter
+    refine_width: bool
+        Flag whether disk edge width is made a fitting parameter
+    global_center: bool
+        If True, uses same center coordinate as the global model
+        If False, uses an independent center
+    x0, y0:
+        Center coordinates of model for local origin
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    exclude_indices: list
+        Indices to exclude from the pattern
+    include_indices: list
+        If specified, only the indices in the list are added to the pattern
+    """
+
     def __init__(
         self,
         WPF,
@@ -507,12 +638,48 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data) -> None:
                 J[:, self.params["edge width"].offset] += dW
 
 
-class SyntheticDiskMoire(WPFModelPrototype):
+class SyntheticDiskMoire(WPFModel):
     """
-    Add Moire peaks arising from two SyntheticDiskLattice lattices.
-    The positions of the Moire peaks are derived from the lattice
-    vectors of the parent lattices. This model object adds only the intensity of
-    each Moire peak as parameters, all other attributes are inherited from the parents
+    Model of diffraction disks arising from interference between two lattices.
+
+    The Moire unit cell is determined automatically using the two input lattices.
+
+    Parameters
+    ----------
+    WPF: WholePatternFit
+        parent fitting object
+    lattice_a, lattice_b: SyntheticDiskLattice
+        parent lattices for the Moire
+    intensity_0
+        Initial guess of Moire disk intensity
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    decorated_peaks: list
+        When specified, only the reflections in the list are decorated with Moire spots
+        If not specified, all peaks are decorated
+    link_moire_disk_intensities: bool
+        When False, each Moire disk has an independently fit intensity
+        When True, Moire disks arising from the same order of parent reflection share
+        the same intensity
+    link_disk_parameters: bool
+        When True, edge_width and disk_radius are inherited from lattice_a
+    refine_width: bool
+        Flag whether disk edge width is a fit variable
+    edge_width
+        Width of the soft edge of the diffraction disk.
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    refine_radius: bool
+        Flag whether disk radius is a fit variable
+    disk radius
+        Radius of the diffraction disks
+        Specified as initial_value, (initial_value, deviation), or
+            (initial_value, lower_bound, upper_bound). See
+            Parameter documentation for details.
+    lattice_b_search_range: int
+        Range of index values for lattice_b to test when finding the Moire unit cell.
     """
 
     def __init__(
@@ -922,7 +1089,7 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **static_data):
                 J[:, self.params["edge width"].offset] += dW
 
 
-class ComplexOverlapKernelDiskLattice(WPFModelPrototype):
+class ComplexOverlapKernelDiskLattice(WPFModel):
     def __init__(
         self,
         WPF,
@@ -938,18 +1105,11 @@ def __init__(
         name="Complex Overlapped Disk Lattice",
         verbose=False,
     ):
-        params = {}
-
-        # if global_center:
-        #     self.func = self.global_center_func
-        #     self.jacobian = self.global_center_jacobian
+        return NotImplementedError(
+            "This model type has not been updated for use with the new architecture."
+        )
 
-        #     x0 = WPF.static_data["global_x0"]
-        #     y0 = WPF.static_data["global_y0"]
-        # else:
-        #     params["x center"] = Parameter(x0)
-        #     params["y center"] = Parameter(y0)
-        #     self.func = self.local_center_func
+        params = {}
 
         self.probe_kernelFT = np.fft.fft2(probe_kernel)
 
@@ -1038,7 +1198,7 @@ def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
         DP += np.abs(localDP) ** 2
 
 
-class KernelDiskLattice(WPFModelPrototype):
+class KernelDiskLattice(WPFModel):
     def __init__(
         self,
         WPF,
@@ -1054,6 +1214,10 @@ def __init__(
         name="Custom Kernel Disk Lattice",
         verbose=False,
     ):
+        return NotImplementedError(
+            "This model type has not been updated for use with the new architecture."
+        )
+
         params = {}
 
         # if global_center:
diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 518647bdb..89ad3640c 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -1,7 +1,7 @@
 from py4DSTEM import DataCube, RealSlice
 from emdfile import tqdmnd
 from py4DSTEM.process.wholepatternfit.wp_models import (
-    WPFModelPrototype,
+    WPFModel,
     _BaseModel,
     WPFModelType,
     Parameter,
@@ -104,8 +104,6 @@ def __init__(
             x0=(x0, global_xy0_lb[0], global_xy0_ub[0]),
             y0=(y0, global_xy0_lb[1], global_xy0_ub[1]),
         )
-        # TODO: remove special cases for global/local center in the Models
-        # Needs an efficient way to handle calculation of q_r
 
         self.model = [
             self.coordinate_model,
@@ -115,7 +113,7 @@ def __init__(
         self.use_jacobian = use_jacobian
 
         # set up the global arguments
-        self._setup_static_data(x0, y0)
+        self._setup_static_data()
 
         # for debugging: tracks all function evals
         self._track = False
@@ -123,23 +121,54 @@ def __init__(
         self._xevals = []
         # self._cost_history = []
 
-    def add_model(self, model: WPFModelPrototype):
+    def add_model(self, model: WPFModel):
+        """
+        Add a WPFModel to the current model
+
+        Parameters
+        ----------
+        model: WPFModel
+            model to add to the fitting routine
+        """
         self.model.append(model)
 
         self.nParams += len(model.params.keys())
 
         self._finalize_model()
 
-    def add_model_list(self, model_list):
+    def add_model_list(self, model_list: list[WPFModel]):
+        """
+        Add multiple WPFModel objects to the current model
+
+        Parameters
+        ----------
+        model: list[WPFModel]
+            models to add to the fitting routine
+        """
         for m in model_list:
             self.add_model(m)
 
-    def link_parameters(self, parent_model, child_model, parameters):
+    def link_parameters(
+        self,
+        parent_model: WPFModel,
+        child_model: WPFModel | list[WPFModel],
+        parameters: str | list[str],
+    ):
         """
         Link parameters of separate models together. The parameters of
         the child_model are replaced with the parameters of the parent_model.
         Note, this does not add the models to the WPF object, that must
         be performed separately.
+
+        Parameters
+        ----------
+        parent_model: WPFModel
+            model from which parameters will be copied
+        child_model: WPFModel or list of WPFModels
+            model(s) whose independent parameters are to be linked
+            with those of the parent_model
+        parameters: str or list of str
+            names of parameters to be linked
         """
         # Make sure child_model and parameters are iterable
         child_model = (
@@ -162,18 +191,41 @@ def link_parameters(self, parent_model, child_model, parameters):
             for par in parameters:
                 child.params[par] = parent_model.params[par]
 
-    def generate_initial_pattern(self):
+    def generate_initial_pattern(self) -> np.ndarray:
+        """
+        Generate a diffraction pattern using the initial parameter
+        guesses for each model component
+
+        Returns
+        -------
+        initial_pattern: np.ndarray
+
+        """
+
         # update parameters:
         self._finalize_model()
         return self._pattern(self.x0, self.static_data.copy()) / self.intensity_scale
 
     def fit_to_mean_CBED(self, **fit_opts):
+        """
+        Fit model parameters to the mean CBED pattern
+
+        Parameters
+        ----------
+        fit_opts: keyword arguments passed to scipy.optimize.least_squares
+
+        Returns
+        -------
+        optimizer_result: dict
+            Output of scipy.optimize.least_squares
+            (also stored in self.mean_CBED_fit)
+
+        """
         # first make sure we have the latest parameters
         self._finalize_model()
 
         # set the current active pattern to the mean CBED:
         current_pattern = self.meanCBED * self.intensity_scale
-        shared_data = self.static_data.copy()
 
         self._fevals = []
         self._xevals = []
@@ -192,7 +244,7 @@ def fit_to_mean_CBED(self, **fit_opts):
                 self.x0,
                 jac=self._jacobian,
                 bounds=(self.lower_bound, self.upper_bound),
-                args=(current_pattern, shared_data),
+                args=(current_pattern, self.static_data),
                 **default_opts,
             )
         else:
@@ -200,7 +252,7 @@ def fit_to_mean_CBED(self, **fit_opts):
                 self._pattern_error,
                 self.x0,
                 bounds=(self.lower_bound, self.upper_bound),
-                args=(current_pattern, shared_data),
+                args=(current_pattern, self.static_data),
                 **default_opts,
             )
 
@@ -217,7 +269,9 @@ def fit_to_mean_CBED(self, **fit_opts):
         ax.set_xlabel("Iterations")
         ax.set_yscale("log")
 
-        DP = self._pattern(self.mean_CBED_fit.x, shared_data) / self.intensity_scale
+        DP = (
+            self._pattern(self.mean_CBED_fit.x, self.static_data) / self.intensity_scale
+        )
         ax = fig.add_subplot(gs[0, 1])
         CyRd = mpl_c.LinearSegmentedColormap.from_list(
             "CyRd", ["#00ccff", "#ffffff", "#ff0000"]
@@ -246,12 +300,12 @@ def fit_to_mean_CBED(self, **fit_opts):
 
     def fit_all_patterns(
         self,
-        resume=False,
-        real_space_mask=None,
-        show_fit_metrics=True,
-        distributed=True,
-        num_jobs=None,
-        threads_per_job=1,
+        resume: bool = False,
+        real_space_mask: Optional[np.ndarray] = None,
+        show_fit_metrics: bool = True,
+        distributed: bool = True,
+        num_jobs: int = None,
+        threads_per_job: int = 1,
         **fit_opts,
     ):
         """
@@ -261,13 +315,18 @@ def fit_all_patterns(
         ----------
         resume: bool (optional)
             Set to true to continue a previous fit with more iterations.
-        real_space_mask: np.array() of bools (optional)
+        real_space_mask: np.ndarray of bools (optional)
             Only perform the fitting on a subset of the probe positions,
             where real_space_mask[rx,ry] == True.
         distributed: bool (optional)
             Whether to evaluate using a pool of worker threads
         num_jobs: int (optional)
-            Set to an integer value giving the number of jobs to parallelize over probe positions.
+            number of parallel worker threads to launch if distributed=True
+            Defaults to number of CPU cores
+        threads_per_job: int (optional)
+            number of threads for each parallel job. If num_jobs is not specified,
+            the number of workers is automatically chosen so as to not oversubscribe
+            the cores (num_jobs = CPU_count // threads_per_job)
         fit_opts: args (optional)
             args passed to scipy.optimize.least_squares
 
@@ -276,7 +335,7 @@ def fit_all_patterns(
         fit_data: RealSlice
             Fitted coefficients for all probe positions
         fit_metrics: RealSlice
-            Fitting metrixs for all probe positions
+            Fitting metrics for all probe positions
 
         """
 
@@ -358,6 +417,7 @@ def fit_all_patterns(
             # distributed evaluation
             self._fit_distributed(
                 resume=resume,
+                real_space_mask=mask,
                 num_jobs=num_jobs,
                 threads_per_job=threads_per_job,
                 fit_opts=default_opts,
@@ -378,13 +438,25 @@ def fit_all_patterns(
         return self.fit_data, self.fit_metrics
 
     def accept_mean_CBED_fit(self):
+        """
+        Sets the parameters optimized by fitting to mean CBED
+        as the initial guess for each of the component models.
+        """
         x = self.mean_CBED_fit.x
 
         for model in self.model:
             for param in model.params.values():
                 param.initial_value = x[param.offset]
 
-    def get_lattice_maps(self):
+    def get_lattice_maps(self) -> list[RealSlice]:
+        """
+        Get the fitted reciprical lattice vectors refined at each scan point.
+
+        Returns
+        -------
+        g_maps: list[RealSlice]
+            RealSlice objects containing the lattice data for each scan position
+        """
         assert hasattr(self, "fit_data"), "Please run fitting first!"
 
         lattices = [m for m in self.model if WPFModelType.LATTICE in m.model_type]
@@ -411,7 +483,10 @@ def get_lattice_maps(self):
 
         return g_maps
 
-    def _setup_static_data(self, x0, y0):
+    def _setup_static_data(self):
+        """
+        Generate basic data that each model can access during the fitting routine
+        """
         self.static_data = {}
 
         xArray, yArray = np.mgrid[0 : self.datacube.Q_Nx, 0 : self.datacube.Q_Ny]
@@ -511,6 +586,7 @@ def _fit_single_pattern(
         self,
         data: np.ndarray,
         initial_guess: np.ndarray,
+        mask: bool,
         fit_opts,
     ):
         """
@@ -522,6 +598,8 @@ def _fit_single_pattern(
             Diffraction pattern
         initial_guess: np.ndarray
             starting guess for fitting
+        mask: bool
+            Fitting is skipped if mask is False, and default values are returned
         fit_opts:
             args passed to scipy.optimize.least_squares
 
@@ -533,49 +611,55 @@ def _fit_single_pattern(
             Fitting metrics
 
         """
+        if mask:
+            try:
+                if self.hasJacobian & self.use_jacobian:
+                    opt = least_squares(
+                        self._pattern_error,
+                        initial_guess,
+                        jac=self._jacobian,
+                        bounds=(self.lower_bound, self.upper_bound),
+                        args=(data * self.intensity_scale, self.static_data),
+                        **fit_opts,
+                    )
+                else:
+                    opt = least_squares(
+                        self._pattern_error,
+                        initial_guess,
+                        bounds=(self.lower_bound, self.upper_bound),
+                        args=(data * self.intensity_scale, self.static_data),
+                        **fit_opts,
+                    )
 
-        try:
-            if self.hasJacobian & self.use_jacobian:
-                opt = least_squares(
-                    self._pattern_error,
-                    initial_guess,
-                    jac=self._jacobian,
-                    bounds=(self.lower_bound, self.upper_bound),
-                    args=(data * self.intensity_scale, self.static_data),
-                    **fit_opts,
-                )
-            else:
-                opt = least_squares(
-                    self._pattern_error,
-                    initial_guess,
-                    bounds=(self.lower_bound, self.upper_bound),
-                    args=(data * self.intensity_scale, self.static_data),
-                    **fit_opts,
-                )
-
-            fit_coefs = opt.x
-            fit_metrics_single = [
-                opt.cost,
-                opt.optimality,
-                opt.nfev,
-                opt.status,
-            ]
-        except Exception as err:
-            print(err)
-            fit_coefs = initial_guess
-            fit_metrics_single = [0, 0, 0, 0]
-
-        return fit_coefs, fit_metrics_single
+                fit_coefs = opt.x
+                fit_metrics_single = [
+                    opt.cost,
+                    opt.optimality,
+                    opt.nfev,
+                    opt.status,
+                ]
+            except Exception as err:
+                # print(err)
+                fit_coefs = initial_guess
+                fit_metrics_single = [0, 0, 0, 0]
+
+            return fit_coefs, fit_metrics_single
+        else:
+            return np.zeros_like(initial_guess), [0, 0, 0, 0]
 
     def _fit_distributed(
         self,
         fit_opts: dict,
         fit_data: np.ndarray,
         fit_metrics: np.ndarray,
+        real_space_mask: np.ndarray,
         resume=False,
         num_jobs=None,
         threads_per_job=1,
     ):
+        """
+        Run fitting using multiprocessing to fit several patterns in parallel
+        """
         from mpire import WorkerPool, cpu_count
         from threadpoolctl import threadpool_limits
 
@@ -592,13 +676,13 @@ def f(shared_data, args):
                 {
                     "data": self.datacube[rx, ry],
                     "initial_guess": self.fit_data[rx, ry] if resume else self.x0,
+                    "mask": real_space_mask[rx, ry],
                 },
             )
             for rx in range(self.datacube.R_Nx)
             for ry in range(self.datacube.R_Ny)
         ]
 
-        # TODO: auto set n_jobs when using multi threads each
         with WorkerPool(
             n_jobs=num_jobs,
             shared_objects=fit_opts,
diff --git a/setup.py b/setup.py
index a806c0131..45a243be2 100644
--- a/setup.py
+++ b/setup.py
@@ -39,6 +39,7 @@
         'distributed >= 2.3.0',
         'emdfile >= 0.0.10',
         'mpire >= 2.7.1',
+        'threadpoolctl >= 3.1.0'
         ],
     extras_require={
         'ipyparallel': ['ipyparallel >= 6.2.4', 'dill >= 0.3.3'],

From c825e0d948958a96a42b8908616123a0bfcc2692 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Mon, 14 Aug 2023 13:26:08 -0400
Subject: [PATCH 51/57] update the weird lattice models to maybe work with new
 architecture

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 101 ++++++++----------
 1 file changed, 44 insertions(+), 57 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 3dfdfca2a..0618a42eb 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -1102,6 +1102,9 @@ def __init__(
         v_max: int,
         intensity_0: float,
         exclude_indices: list = [],
+        global_center: bool = True,
+        x0=0.0,
+        y0=0.0,
         name="Complex Overlapped Disk Lattice",
         verbose=False,
     ):
@@ -1113,6 +1116,16 @@ def __init__(
 
         self.probe_kernelFT = np.fft.fft2(probe_kernel)
 
+        if global_center:
+            params["x center"] = WPF.coordinate_model.params["x center"]
+            params["y center"] = WPF.coordinate_model.params["y center"]
+        else:
+            params["x center"] = Parameter(x0)
+            params["y center"] = Parameter(y0)
+
+        x0 = params["x center"].initial_value
+        y0 = params["y center"].initial_value
+
         params["ux"] = Parameter(ux)
         params["uy"] = Parameter(uy)
         params["vx"] = Parameter(vx)
@@ -1160,23 +1173,15 @@ def __init__(
         self.u_inds = self.u_inds[~delete_mask]
         self.v_inds = self.v_inds[~delete_mask]
 
-        self.func = self.global_center_func
-
         super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
-    def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-        # copy the global centers in the right place for the local center generator
-        self.local_center_func(
-            DP, kwargs["global_x0"], kwargs["global_y0"], *args, **kwargs
-        )
-
-    def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-        x0 = args[0]
-        y0 = args[1]
-        ux = args[2]
-        uy = args[3]
-        vx = args[4]
-        vy = args[5]
+    def func(self, DP: np.ndarray, x_fit, **kwargs) -> None:
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
+        ux = x[self.params["ux"].offset]
+        uy = x[self.params["uy"].offset]
+        vx = x[self.params["vx"].offset]
+        vy = x[self.params["vy"].offset]
 
         localDP = np.zeros_like(DP, dtype=np.complex64)
 
@@ -1185,8 +1190,8 @@ def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
             y = y0 + (u * uy) + (v * vy)
 
             localDP += (
-                args[2 * i + 6]
-                * np.exp(1j * args[2 * i + 7])
+                x_fit[self.params[f"[{u},{v}] Intensity"].offset]
+                * np.exp(1j * x_fit[self.params[f"[{u},{v}] Phase"].offset])
                 * np.abs(
                     np.fft.ifft2(
                         self.probe_kernelFT
@@ -1211,27 +1216,25 @@ def __init__(
         v_max: int,
         intensity_0: float,
         exclude_indices: list = [],
+        global_center: bool = True,
+        x0=0.0,
+        y0=0.0,
         name="Custom Kernel Disk Lattice",
         verbose=False,
     ):
-        return NotImplementedError(
-            "This model type has not been updated for use with the new architecture."
-        )
-
         params = {}
 
-        # if global_center:
-        #     self.func = self.global_center_func
-        #     self.jacobian = self.global_center_jacobian
+        self.probe_kernelFT = np.fft.fft2(probe_kernel)
 
-        #     x0 = WPF.static_data["global_x0"]
-        #     y0 = WPF.static_data["global_y0"]
-        # else:
-        #     params["x center"] = Parameter(x0)
-        #     params["y center"] = Parameter(y0)
-        #     self.func = self.local_center_func
+        if global_center:
+            params["x center"] = WPF.coordinate_model.params["x center"]
+            params["y center"] = WPF.coordinate_model.params["y center"]
+        else:
+            params["x center"] = Parameter(x0)
+            params["y center"] = Parameter(y0)
 
-        self.probe_kernelFT = np.fft.fft2(probe_kernel)
+        x0 = params["x center"].initial_value
+        y0 = params["y center"].initial_value
 
         params["ux"] = Parameter(ux)
         params["uy"] = Parameter(uy)
@@ -1250,16 +1253,8 @@ def __init__(
         self.xqArray = np.tile(np.fft.fftfreq(Q_Nx)[:, np.newaxis], (1, Q_Ny))
 
         for i, (u, v) in enumerate(zip(u_inds.ravel(), v_inds.ravel())):
-            x = (
-                WPF.static_data["global_x0"]
-                + (u * params["ux"].initial_value)
-                + (v * params["vx"].initial_value)
-            )
-            y = (
-                WPF.static_data["global_y0"]
-                + (u * params["uy"].initial_value)
-                + (v * params["vy"].initial_value)
-            )
+            x = x0 + (u * params["ux"].initial_value) + (v * params["vx"].initial_value)
+            y = y0 + (u * params["uy"].initial_value) + (v * params["vy"].initial_value)
             if [u, v] in exclude_indices:
                 delete_mask[i] = True
             elif (x < 0) or (x > Q_Nx) or (y < 0) or (y > Q_Ny):
@@ -1274,30 +1269,22 @@ def __init__(
         self.u_inds = self.u_inds[~delete_mask]
         self.v_inds = self.v_inds[~delete_mask]
 
-        self.func = self.global_center_func
-
         super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
-    def global_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-        # copy the global centers in the right place for the local center generator
-        self.local_center_func(
-            DP, kwargs["global_x0"], kwargs["global_y0"], *args, **kwargs
-        )
-
-    def local_center_func(self, DP: np.ndarray, *args, **kwargs) -> None:
-        x0 = args[0]
-        y0 = args[1]
-        ux = args[2]
-        uy = args[3]
-        vx = args[4]
-        vy = args[5]
+    def func(self, DP: np.ndarray, x_fit: np.ndarray, **static_data) -> None:
+        x0 = x[self.params["x center"].offset]
+        y0 = x[self.params["y center"].offset]
+        ux = x[self.params["ux"].offset]
+        uy = x[self.params["uy"].offset]
+        vx = x[self.params["vx"].offset]
+        vy = x[self.params["vy"].offset]
 
         for i, (u, v) in enumerate(zip(self.u_inds, self.v_inds)):
             x = x0 + (u * ux) + (v * vx)
             y = y0 + (u * uy) + (v * vy)
 
             DP += (
-                args[i + 6]
+                x_fit[params[f"[{u},{v}] Intensity"].offset]
                 * np.abs(
                     np.fft.ifft2(
                         self.probe_kernelFT

From 9b79ef5e3a1d900a8ecf823868e3f432df8d6831 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Mon, 14 Aug 2023 13:29:38 -0400
Subject: [PATCH 52/57] linter fixes

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 0618a42eb..a5a3c0142 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -1176,12 +1176,12 @@ def __init__(
         super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
     def func(self, DP: np.ndarray, x_fit, **kwargs) -> None:
-        x0 = x[self.params["x center"].offset]
-        y0 = x[self.params["y center"].offset]
-        ux = x[self.params["ux"].offset]
-        uy = x[self.params["uy"].offset]
-        vx = x[self.params["vx"].offset]
-        vy = x[self.params["vy"].offset]
+        x0 = x_fit[self.params["x center"].offset]
+        y0 = x_fit[self.params["y center"].offset]
+        ux = x_fit[self.params["ux"].offset]
+        uy = x_fit[self.params["uy"].offset]
+        vx = x_fit[self.params["vx"].offset]
+        vy = x_fit[self.params["vy"].offset]
 
         localDP = np.zeros_like(DP, dtype=np.complex64)
 
@@ -1272,12 +1272,12 @@ def __init__(
         super().__init__(name, params, model_type=WPFModelType.LATTICE)
 
     def func(self, DP: np.ndarray, x_fit: np.ndarray, **static_data) -> None:
-        x0 = x[self.params["x center"].offset]
-        y0 = x[self.params["y center"].offset]
-        ux = x[self.params["ux"].offset]
-        uy = x[self.params["uy"].offset]
-        vx = x[self.params["vx"].offset]
-        vy = x[self.params["vy"].offset]
+        x0 = x_fit[self.params["x center"].offset]
+        y0 = x_fit[self.params["y center"].offset]
+        ux = x_fit[self.params["ux"].offset]
+        uy = x_fit[self.params["uy"].offset]
+        vx = x_fit[self.params["vx"].offset]
+        vy = x_fit[self.params["vy"].offset]
 
         for i, (u, v) in enumerate(zip(self.u_inds, self.v_inds)):
             x = x0 + (u * ux) + (v * vx)

From 51f6c1eff7a135541c829cf60304d7044ab807db Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Mon, 14 Aug 2023 13:30:58 -0400
Subject: [PATCH 53/57] another linter fix

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index a5a3c0142..d4d190365 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -1284,7 +1284,7 @@ def func(self, DP: np.ndarray, x_fit: np.ndarray, **static_data) -> None:
             y = y0 + (u * uy) + (v * vy)
 
             DP += (
-                x_fit[params[f"[{u},{v}] Intensity"].offset]
+                x_fit[self.params[f"[{u},{v}] Intensity"].offset]
                 * np.abs(
                     np.fft.ifft2(
                         self.probe_kernelFT

From 4fe16dd79ac9419e6539c0f05d2eba2e3d930bd6 Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Tue, 15 Aug 2023 16:43:12 -0400
Subject: [PATCH 54/57] improvements to moire

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 19 ++---
 py4DSTEM/process/wholepatternfit/wpf_viz.py   | 77 ++++++++++++++++++-
 2 files changed, 80 insertions(+), 16 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index d4d190365..2ef87119b 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -678,8 +678,6 @@ class SyntheticDiskMoire(WPFModel):
         Specified as initial_value, (initial_value, deviation), or
             (initial_value, lower_bound, upper_bound). See
             Parameter documentation for details.
-    lattice_b_search_range: int
-        Range of index values for lattice_b to test when finding the Moire unit cell.
     """
 
     def __init__(
@@ -695,15 +693,8 @@ def __init__(
         edge_width: list = None,
         refine_radius: bool = True,
         disk_radius: list = None,
-        lattice_b_search_range: int = 1,
         name: str = "Moire Lattice",
     ):
-        """
-        Parameters
-        ----------
-
-        lattice_a, lattice_b: SyntheticDiskLattice
-        """
         # ensure both models share the same center coordinate
         if (lattice_a.params["x center"] is not lattice_b.params["x center"]) or (
             lattice_a.params["y center"] is not lattice_b.params["y center"]
@@ -722,14 +713,14 @@ def __init__(
         lat_ab = self._get_parent_lattices(lattice_a, lattice_b)
 
         # pick the pairing that gives the smallest unit cell
-        imax = lattice_b_search_range
-        mx, my = np.mgrid[-imax : imax + 1, -imax : imax + 1]
-        test_peaks = np.stack([mx.ravel(), my.ravel()], axis=1)
+        mx, my = np.mgrid[-1 : 2, -1 : 2]
+        test_peaks_a = np.stack([mx.ravel(), my.ravel()], axis=1)
+        test_peaks_b = np.stack((lattice_b.u_inds, lattice_b.v_inds), axis=1)
         tests = np.stack(
             [
                 np.hstack((np.eye(2), np.vstack((b1, b2))))
-                for b1 in test_peaks
-                for b2 in test_peaks
+                for b1 in test_peaks_a
+                for b2 in test_peaks_b
                 if not np.allclose(b1, b2)
             ],
             axis=0,
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index f510ca804..79ea1e0a3 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -61,13 +61,40 @@ def show_lattice_points(
     vmin=None,
     vmax=None,
     power=None,
+    show_vectors=True,
     crop_to_pattern=False,
     returnfig=False,
+    moire_origin_idx=[0,0,0,0],
     *args,
     **kwargs,
 ):
     """
     Plotting utility to show the initial lattice points.
+
+    Parameters
+    ----------
+    im: np.ndarray
+        Optional: Image to show, defaults to mean CBED
+    vmin, vmax: float
+        Intensity ranges for plotting im
+    power: float
+        Gamma level for showing im
+    show_vectors: bool
+        Flag to plot the lattice vectors
+    crop_to_pattern: bool
+        Flag to limit the field of view to the pattern area. If False,
+        spots outside the pattern are shown
+    returnfig: bool
+        If True, (fig,ax) are returned and plt.show() is not called
+    moire_origin_idx: list of length 4
+        Indices of peak on which to draw Moire vectors, written as
+        [a_u, a_v, b_u, b_v]
+    args, kwargs
+        Passed to plt.subplots
+
+    Returns
+    -------
+    fig,ax: If returnfig=True
     """
 
     if im is None:
@@ -102,7 +129,7 @@ def show_lattice_points(
         spots[:, 0] += m.params["x center"].initial_value
         spots[:, 1] += m.params["y center"].initial_value
 
-        ax.scatter(
+        axpts = ax.scatter(
             spots[:, 1],
             spots[:, 0],
             s=100,
@@ -110,6 +137,27 @@ def show_lattice_points(
             label=m.name,
         )
 
+        if show_vectors:
+            ax.arrow(
+                m.params["y center"].initial_value, 
+                m.params["x center"].initial_value,
+                m.params["uy"].initial_value,
+                m.params["ux"].initial_value,
+                length_includes_head=True,
+                color=axpts.get_facecolor(),
+                width=1., 
+                )
+
+            ax.arrow(
+                m.params["y center"].initial_value, 
+                m.params["x center"].initial_value,
+                m.params["vy"].initial_value,
+                m.params["vx"].initial_value,
+                length_includes_head=True,
+                color=axpts.get_facecolor(),
+                width=1., 
+                )
+
     moires = [m for m in self.model if WPFModelType.MOIRE in m.model_type]
 
     for m in moires:
@@ -120,7 +168,7 @@ def show_lattice_points(
         spots[:, 0] += m.params["x center"].initial_value
         spots[:, 1] += m.params["y center"].initial_value
 
-        ax.scatter(
+        axpts = ax.scatter(
             spots[:, 1],
             spots[:, 0],
             s=100,
@@ -128,6 +176,31 @@ def show_lattice_points(
             label=m.name,
         )
 
+        if show_vectors:
+            arrow_origin = np.array(moire_origin_idx) @ lat_ab
+            arrow_origin[0] += m.params["x center"].initial_value
+            arrow_origin[1] += m.params["y center"].initial_value
+
+            ax.arrow(
+                arrow_origin[1], 
+                arrow_origin[0],
+                lat_abm[4,1],
+                lat_abm[4,0],
+                length_includes_head=True,
+                color=axpts.get_facecolor(),
+                width=1., 
+                )
+
+            ax.arrow(
+                arrow_origin[1], 
+                arrow_origin[0],
+                lat_abm[5,1],
+                lat_abm[5,0],
+                length_includes_head=True,
+                color=axpts.get_facecolor(),
+                width=1., 
+                )
+
     ax.legend()
 
     if crop_to_pattern:

From a5785786cf5245f2d9616a4b5c794b7c7465f119 Mon Sep 17 00:00:00 2001
From: Stephanie Ribet <stephanie.ribet@gmail.com>
Date: Fri, 18 Aug 2023 09:19:27 -0700
Subject: [PATCH 55/57] small updates

---
 py4DSTEM/process/wholepatternfit/wp_models.py |  7 ++--
 py4DSTEM/process/wholepatternfit/wpf_viz.py   | 34 +++++++++----------
 py4DSTEM/visualize/show.py                    |  3 --
 3 files changed, 19 insertions(+), 25 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index 2ef87119b..f28c65147 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -1,10 +1,7 @@
-from inspect import signature
 from typing import Optional
 from enum import Flag, auto
 import numpy as np
 
-from pdb import set_trace
-
 
 class WPFModelType(Flag):
     """
@@ -342,7 +339,7 @@ def jacobian(self, J: np.ndarray, x: np.ndarray, **kwargs) -> None:
         J[:, self.parans["y center"].offset] += (
             level
             * exp_expr
-            * (kwargs["yArray"] - x0)
+            * (kwargs["yArray"] - y0)
             * local_r
             / (sigma**2 * clipped_r)
         ).ravel()
@@ -713,7 +710,7 @@ def __init__(
         lat_ab = self._get_parent_lattices(lattice_a, lattice_b)
 
         # pick the pairing that gives the smallest unit cell
-        mx, my = np.mgrid[-1 : 2, -1 : 2]
+        mx, my = np.mgrid[-1:2, -1:2]
         test_peaks_a = np.stack([mx.ravel(), my.ravel()], axis=1)
         test_peaks_b = np.stack((lattice_b.u_inds, lattice_b.v_inds), axis=1)
         tests = np.stack(
diff --git a/py4DSTEM/process/wholepatternfit/wpf_viz.py b/py4DSTEM/process/wholepatternfit/wpf_viz.py
index 79ea1e0a3..06c55edfb 100644
--- a/py4DSTEM/process/wholepatternfit/wpf_viz.py
+++ b/py4DSTEM/process/wholepatternfit/wpf_viz.py
@@ -64,7 +64,7 @@ def show_lattice_points(
     show_vectors=True,
     crop_to_pattern=False,
     returnfig=False,
-    moire_origin_idx=[0,0,0,0],
+    moire_origin_idx=[0, 0, 0, 0],
     *args,
     **kwargs,
 ):
@@ -139,24 +139,24 @@ def show_lattice_points(
 
         if show_vectors:
             ax.arrow(
-                m.params["y center"].initial_value, 
+                m.params["y center"].initial_value,
                 m.params["x center"].initial_value,
                 m.params["uy"].initial_value,
                 m.params["ux"].initial_value,
                 length_includes_head=True,
                 color=axpts.get_facecolor(),
-                width=1., 
-                )
+                width=1.0,
+            )
 
             ax.arrow(
-                m.params["y center"].initial_value, 
+                m.params["y center"].initial_value,
                 m.params["x center"].initial_value,
                 m.params["vy"].initial_value,
                 m.params["vx"].initial_value,
                 length_includes_head=True,
                 color=axpts.get_facecolor(),
-                width=1., 
-                )
+                width=1.0,
+            )
 
     moires = [m for m in self.model if WPFModelType.MOIRE in m.model_type]
 
@@ -182,24 +182,24 @@ def show_lattice_points(
             arrow_origin[1] += m.params["y center"].initial_value
 
             ax.arrow(
-                arrow_origin[1], 
+                arrow_origin[1],
                 arrow_origin[0],
-                lat_abm[4,1],
-                lat_abm[4,0],
+                lat_abm[4, 1],
+                lat_abm[4, 0],
                 length_includes_head=True,
                 color=axpts.get_facecolor(),
-                width=1., 
-                )
+                width=1.0,
+            )
 
             ax.arrow(
-                arrow_origin[1], 
+                arrow_origin[1],
                 arrow_origin[0],
-                lat_abm[5,1],
-                lat_abm[5,0],
+                lat_abm[5, 1],
+                lat_abm[5, 0],
                 length_includes_head=True,
                 color=axpts.get_facecolor(),
-                width=1., 
-                )
+                width=1.0,
+            )
 
     ax.legend()
 
diff --git a/py4DSTEM/visualize/show.py b/py4DSTEM/visualize/show.py
index 2bbe563ac..c9785be5b 100644
--- a/py4DSTEM/visualize/show.py
+++ b/py4DSTEM/visualize/show.py
@@ -681,9 +681,6 @@ def show(
             scalebar['space'] = space
             # determine good default scale bar fontsize
             if figax is not None:
-                # print(figax[0].figsize)
-                # size = figax[1].get_size_inches()
-                # ax_h = figax[1].bbox.transformed(figax[0].gca().transAxes).height
                 bbox = figax[1].get_window_extent()
                 dpi = figax[0].dpi
                 size = (

From d5652e3fb8da59e8c2c072d22adf880ea8ac3c9b Mon Sep 17 00:00:00 2001
From: Stephanie Ribet <stephanie.ribet@gmail.com>
Date: Fri, 18 Aug 2023 09:50:58 -0700
Subject: [PATCH 56/57] wpf imports

---
 py4DSTEM/process/wholepatternfit/wpf.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wpf.py b/py4DSTEM/process/wholepatternfit/wpf.py
index 89ad3640c..5851e1f09 100644
--- a/py4DSTEM/process/wholepatternfit/wpf.py
+++ b/py4DSTEM/process/wholepatternfit/wpf.py
@@ -10,11 +10,10 @@
 from typing import Optional
 import numpy as np
 
-from scipy.optimize import least_squares, minimize
+from scipy.optimize import least_squares
 import matplotlib.pyplot as plt
 import matplotlib.colors as mpl_c
 from matplotlib.gridspec import GridSpec
-import warnings
 
 __all__ = ["WholePatternFit"]
 
@@ -276,13 +275,12 @@ def fit_to_mean_CBED(self, **fit_opts):
         CyRd = mpl_c.LinearSegmentedColormap.from_list(
             "CyRd", ["#00ccff", "#ffffff", "#ff0000"]
         )
-        im = ax.matshow(
+        ax.matshow(
             err_im := -(DP - self.meanCBED),
             cmap=CyRd,
             vmin=-np.abs(err_im).max() / 4,
             vmax=np.abs(err_im).max() / 4,
         )
-        # fig.colorbar(im)
         ax.set_title("Error")
         ax.axis("off")
 

From c6eaaac18fb2e0c172aab08c929ae9ee8f24e6ab Mon Sep 17 00:00:00 2001
From: Steven Zeltmann <steven.zeltmann@cornell.edu>
Date: Fri, 18 Aug 2023 14:04:33 -0400
Subject: [PATCH 57/57] fix for moire cell generation

---
 py4DSTEM/process/wholepatternfit/wp_models.py | 36 +++++++++++++------
 1 file changed, 25 insertions(+), 11 deletions(-)

diff --git a/py4DSTEM/process/wholepatternfit/wp_models.py b/py4DSTEM/process/wholepatternfit/wp_models.py
index f28c65147..8a10303b9 100644
--- a/py4DSTEM/process/wholepatternfit/wp_models.py
+++ b/py4DSTEM/process/wholepatternfit/wp_models.py
@@ -710,24 +710,38 @@ def __init__(
         lat_ab = self._get_parent_lattices(lattice_a, lattice_b)
 
         # pick the pairing that gives the smallest unit cell
-        mx, my = np.mgrid[-1:2, -1:2]
-        test_peaks_a = np.stack([mx.ravel(), my.ravel()], axis=1)
-        test_peaks_b = np.stack((lattice_b.u_inds, lattice_b.v_inds), axis=1)
+        test_peaks = np.stack((lattice_b.u_inds, lattice_b.v_inds), axis=1)
         tests = np.stack(
             [
                 np.hstack((np.eye(2), np.vstack((b1, b2))))
-                for b1 in test_peaks_a
-                for b2 in test_peaks_b
+                for b1 in test_peaks
+                for b2 in test_peaks
                 if not np.allclose(b1, b2)
             ],
             axis=0,
         )
-
-        # pick the combination that gives the smallest cell
-        # (it might be better to choose the cell that gives the smallest volume
-        # but then we have to handle the case of nearly-parallel vectors in
-        # the search space)
-        M = tests[np.argmin(np.max(np.linalg.norm(tests @ lat_ab, axis=-1), axis=-1))]
+        # choose only cells where the two unit vectors are not nearly parallel,
+        # and penalize cells with large discrepancy in lattce vector length
+        lat_m = tests @ lat_ab
+        a_dot_b = (
+            np.sum(lat_m[:, 0] * lat_m[:, 1], axis=1)
+            / np.minimum(
+                np.linalg.norm(lat_m[:, 0], axis=1), np.linalg.norm(lat_m[:, 1], axis=1)
+            )
+            ** 2
+        )
+        tests = tests[np.abs(a_dot_b) < 0.9] # this factor of 0.9 sets the parallel cutoff
+        # with the parallel vectors filtered, pick the cell with the smallest volume
+        lat_m = tests @ lat_ab
+        V = np.sum(
+            lat_m[:, 0]
+            * np.cross(
+                np.hstack((lat_m[:, 1], np.zeros((lat_m.shape[0],))[:, None])),
+                [0, 0, 1],
+            )[:, :2],
+            axis=1,
+        )
+        M = tests[np.argmin(np.abs(V))]
 
         # ensure the moire vectors are less 90 deg apart
         if np.arccos(