Merge branch 'dev' into occupancy

.github/scripts/update_version.py

@@ -8,7 +8,7 @@
     lines = f.readlines()
 
 line_split = lines[0].split(".")
-patch_number = line_split[2].split("'")[0]
+patch_number = line_split[2].split("'")[0].split('"')[0]
 
 # Increment patch number
 patch_number = str(int(patch_number) + 1) + "'"

.github/workflows/check_install_dev.yml

@@ -16,7 +16,7 @@ jobs:
                 allow_failure: [false]
                 runs-on: [ubuntu-latest]
                 architecture: [x86_64]
-                python-version: ["3.9", "3.10", "3.11",]
+                python-version: ["3.9", "3.10", "3.11", "3.12"]
                 # include:
                 #     - python-version: "3.12.0-beta.4"
                 #       runs-on: ubuntu-latest

.github/workflows/check_install_main.yml

@@ -16,7 +16,7 @@ jobs:
                 allow_failure: [false]
                 runs-on: [ubuntu-latest, windows-latest, macos-latest]
                 architecture: [x86_64]
-                python-version: ["3.9", "3.10", "3.11",]
+                python-version: ["3.9", "3.10", "3.11", "3.12"]
                 #include:
                 #    - python-version: "3.12.0-beta.4"
                 #      runs-on: ubuntu-latest

.github/workflows/check_install_quick.yml

@@ -20,7 +20,7 @@ jobs:
                 allow_failure: [false]
                 runs-on: [ubuntu-latest] 
                 architecture: [x86_64]
-                python-version: ["3.10"]
+                python-version: ["3.9", "3.12"]
                 # Currently no public runners available for this but this or arm64 should work next time
                 # include:
                 #     - python-version: "3.10"

.gitignore

@@ -4,6 +4,7 @@
 *.swp
 *.ipynb_checkpoints*
 .vscode/
+pyrightconfig.json
 
 # Folders #
 .idea/

.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+  # Using this mirror lets us use mypyc-compiled black, which is about 2x faster
+  - repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 23.11.0
+    hooks:
+      - id: black
+        language_version: python3.9
+  - repo: https://github.com/pycqa/flake8
+    rev: '6.1.0'  
+    hooks:
+      - id: flake8

py4DSTEM/__init__.py

@@ -52,17 +52,25 @@
     BraggVectorMap,
 )
 
-# strain
-from py4DSTEM.process import StrainMap
+from py4DSTEM.process import classification
 
-# TODO - crystal
-# TODO - ptycho
-# TODO - others
 
-# TODO - where
-from py4DSTEM.process import (
-    PolarDatacube,
-)
+# diffraction
+from py4DSTEM.process.diffraction import Crystal, Orientation
+
+
+# ptycho
+from py4DSTEM.process import phase
+
+
+# polar
+from py4DSTEM.process.polar import PolarDatacube
+
+
+# strain
+from py4DSTEM.process.strain.strain import StrainMap
+
+from py4DSTEM.process import wholepatternfit
 
 
 ### more submodules

py4DSTEM/braggvectors/braggvector_methods.py

@@ -1,12 +1,14 @@
 # BraggVectors methods
 
-import numpy as np
-from scipy.ndimage import gaussian_filter
-from warnings import warn
 import inspect
+from warnings import warn
 
-from emdfile import Array, Metadata, tqdmnd, _read_metadata
+import matplotlib.pyplot as plt
+import numpy as np
+from emdfile import Array, Metadata, _read_metadata, tqdmnd
+from py4DSTEM import show
 from py4DSTEM.datacube import VirtualImage
+from scipy.ndimage import gaussian_filter
 
 
 class BraggVectorMethods:
@@ -97,12 +99,12 @@ def histogram(
         # then scale by the sampling factor
         else:
             # get pixel calibration
-            if self.calstate["pixel"] == True:
+            if self.calstate["pixel"] is True:
                 qpix = self.calibration.get_Q_pixel_size()
                 qx /= qpix
                 qy /= qpix
             # origin calibration
-            if self.calstate["center"] == True:
+            if self.calstate["center"] is True:
                 origin = self.calibration.get_origin_mean()
                 qx += origin[0]
                 qy += origin[1]
@@ -151,12 +153,12 @@ def histogram(
         ).reshape(Q_Nx, Q_Ny)
 
         # determine the resampled grid center and pixel size
-        if mode == "cal" and self.calstate["center"] == True:
+        if mode == "cal" and self.calstate["center"] is True:
             x0 = sampling * origin[0]
             y0 = sampling * origin[1]
         else:
             x0, y0 = 0, 0
-        if mode == "cal" and self.calstate["pixel"] == True:
+        if mode == "cal" and self.calstate["pixel"] is True:
             pixelsize = qpix / sampling
         else:
             pixelsize = 1 / sampling
@@ -518,6 +520,7 @@ def fit_origin(
         mask_check_data=True,
         plot=True,
         plot_range=None,
+        cmap="RdBu_r",
         returncalc=True,
         **kwargs,
     ):
@@ -537,6 +540,7 @@ def fit_origin(
             mask_check_data (bool):     Get mask from origin measurements equal to zero. (TODO - replace)
             plot (bool, optional): plot results
             plot_range (float):    min and max color range for plot (pixels)
+            cmap (colormap): plotting colormap
 
         Returns:
             (variable): Return value depends on returnfitp. If ``returnfitp==False``
@@ -552,84 +556,113 @@ def fit_origin(
         from py4DSTEM.process.calibration import fit_origin
 
         if mask_check_data is True:
-            # TODO - replace this bad hack for the mask for the origin fit
-            mask = np.logical_not(q_meas[0] == 0)
-            qx0_fit, qy0_fit, qx0_residuals, qy0_residuals = fit_origin(
-                tuple(q_meas),
-                mask=mask,
-            )
-        else:
-            qx0_fit, qy0_fit, qx0_residuals, qy0_residuals = fit_origin(tuple(q_meas))
+            data_mask = np.logical_not(q_meas[0] == 0)
+            if mask is None:
+                mask = data_mask
+            else:
+                mask = np.logical_and(mask, data_mask)
+
+        qx0_fit, qy0_fit, qx0_residuals, qy0_residuals = fit_origin(
+            tuple(q_meas),
+            mask=mask,
+            fitfunction=fitfunction,
+            robust=robust,
+            robust_steps=robust_steps,
+            robust_thresh=robust_thresh,
+        )
 
-        # try to add to calibration
+        # try to add update calibration metadata
         try:
-            self.calibration.set_origin([qx0_fit, qy0_fit])
+            self.calibration.set_origin((qx0_fit, qy0_fit))
+            self.setcal()
         except AttributeError:
             warn(
                 "No calibration found on this datacube - fit values are not being stored"
             )
             pass
-        if plot:
-            from py4DSTEM.visualize import show_image_grid
 
-            if mask is None:
-                qx0_meas, qy0_meas = q_meas
-                qx0_res_plot = qx0_residuals
-                qy0_res_plot = qy0_residuals
-            else:
-                qx0_meas = np.ma.masked_array(q_meas[0], mask=np.logical_not(mask))
-                qy0_meas = np.ma.masked_array(q_meas[1], mask=np.logical_not(mask))
-                qx0_res_plot = np.ma.masked_array(
-                    qx0_residuals, mask=np.logical_not(mask)
-                )
-                qy0_res_plot = np.ma.masked_array(
-                    qy0_residuals, mask=np.logical_not(mask)
-                )
-            qx0_mean = np.mean(qx0_fit)
-            qy0_mean = np.mean(qy0_fit)
-
-            if plot_range is None:
-                plot_range = 2 * np.max(qx0_fit - qx0_mean)
-
-            cmap = kwargs.get("cmap", "RdBu_r")
-            kwargs.pop("cmap", None)
-            axsize = kwargs.get("axsize", (6, 2))
-            kwargs.pop("axsize", None)
-
-            show_image_grid(
-                lambda i: [
-                    qx0_meas - qx0_mean,
-                    qx0_fit - qx0_mean,
-                    qx0_res_plot,
-                    qy0_meas - qy0_mean,
-                    qy0_fit - qy0_mean,
-                    qy0_res_plot,
-                ][i],
-                H=2,
-                W=3,
+        # show
+        if plot:
+            self.show_origin_fit(
+                q_meas[0],
+                q_meas[1],
+                qx0_fit,
+                qy0_fit,
+                qx0_residuals,
+                qy0_residuals,
+                mask=mask,
+                plot_range=plot_range,
                 cmap=cmap,
-                axsize=axsize,
-                title=[
-                    "measured origin, x",
-                    "fitorigin, x",
-                    "residuals, x",
-                    "measured origin, y",
-                    "fitorigin, y",
-                    "residuals, y",
-                ],
-                vmin=-1 * plot_range,
-                vmax=1 * plot_range,
-                intensity_range="absolute",
                 **kwargs,
             )
 
-        # update calibration metadata
-        self.calibration.set_origin((qx0_fit, qy0_fit))
-        self.setcal()
-
+        # return
         if returncalc:
             return qx0_fit, qy0_fit, qx0_residuals, qy0_residuals
 
+    def show_origin_fit(
+        self,
+        qx0_meas,
+        qy0_meas,
+        qx0_fit,
+        qy0_fit,
+        qx0_residuals,
+        qy0_residuals,
+        mask=None,
+        plot_range=None,
+        cmap="RdBu_r",
+        **kwargs,
+    ):
+        # apply mask
+        if mask is not None:
+            qx0_meas = np.ma.masked_array(qx0_meas, mask=np.logical_not(mask))
+            qy0_meas = np.ma.masked_array(qy0_meas, mask=np.logical_not(mask))
+            qx0_residuals = np.ma.masked_array(qx0_residuals, mask=np.logical_not(mask))
+            qy0_residuals = np.ma.masked_array(qy0_residuals, mask=np.logical_not(mask))
+        qx0_mean = np.mean(qx0_fit)
+        qy0_mean = np.mean(qy0_fit)
+
+        # set range
+        if plot_range is None:
+            plot_range = max(
+                (
+                    1.5 * np.max(np.abs(qx0_fit - qx0_mean)),
+                    1.5 * np.max(np.abs(qy0_fit - qy0_mean)),
+                )
+            )
+
+        # set figsize
+        imsize_ratio = np.sqrt(qx0_meas.shape[1] / qx0_meas.shape[0])
+        axsize = (3 * imsize_ratio, 3 / imsize_ratio)
+        axsize = kwargs.pop("axsize", axsize)
+
+        # plot
+        fig, ax = show(
+            [
+                [qx0_meas - qx0_mean, qx0_fit - qx0_mean, qx0_residuals],
+                [qy0_meas - qy0_mean, qy0_fit - qy0_mean, qy0_residuals],
+            ],
+            cmap=cmap,
+            axsize=axsize,
+            title=[
+                "measured origin, x",
+                "fitorigin, x",
+                "residuals, x",
+                "measured origin, y",
+                "fitorigin, y",
+                "residuals, y",
+            ],
+            vmin=-1 * plot_range,
+            vmax=1 * plot_range,
+            intensity_range="absolute",
+            show_cbar=True,
+            returnfig=True,
+            **kwargs,
+        )
+        plt.tight_layout()
+
+        return
+
     def fit_p_ellipse(
         self, bvm, center, fitradii, mask=None, returncalc=False, **kwargs
     ):
@@ -765,6 +798,21 @@ def mask_in_R(self, mask, update_inplace=False, returncalc=True):
         else:
             return
 
+    def to_strainmap(self, name: str = None):
+        """
+        Generate a StrainMap object from the BraggVectors
+        equivalent to py4DSTEM.StrainMap(braggvectors=braggvectors)
+
+        Args:
+            name (str, optional): The name of the strainmap. Defaults to None which reverts to default name 'strainmap'.
+
+        Returns:
+            py4DSTEM.StrainMap: A py4DSTEM StrainMap object generated from the BraggVectors
+        """
+        from py4DSTEM.process.strain import StrainMap
+
+        return StrainMap(self, name) if name else StrainMap(self)
+
 
 ######### END BraggVectorMethods CLASS ########