Merge pull request #668 from py4dstem/lite

I like my coffee black, my py4DSTEM lite, and my tea in the river

py4DSTEM/__init__.py

@@ -1,6 +1,9 @@
 from py4DSTEM.version import __version__
 from emdfile import tqdmnd
 
+from importlib.metadata import packages_distributions
+
+is_package_lite = "py4DSTEM-lite" in packages_distributions()["py4DSTEM"]
 
 ### io
 
@@ -52,8 +55,11 @@
     BraggVectorMap,
 )
 
-from py4DSTEM.process import classification
-
+try:
+    from py4DSTEM.process import classification
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 
 # diffraction
 from py4DSTEM.process.diffraction import Crystal, Orientation
@@ -70,7 +76,11 @@
 # strain
 from py4DSTEM.process.strain.strain import StrainMap
 
-from py4DSTEM.process import wholepatternfit
+try:
+    from py4DSTEM.process import wholepatternfit
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 
 
 ### more submodules

py4DSTEM/braggvectors/diskdetection.py

@@ -5,12 +5,18 @@
 from scipy.ndimage import gaussian_filter
 
 from emdfile import tqdmnd
+from py4DSTEM import is_package_lite
 from py4DSTEM.braggvectors.braggvectors import BraggVectors
 from py4DSTEM.data import QPoints
 from py4DSTEM.datacube import DataCube
 from py4DSTEM.preprocess.utils import get_maxima_2D
 from py4DSTEM.process.utils.cross_correlate import get_cross_correlation_FT
-from py4DSTEM.braggvectors.diskdetection_aiml import find_Bragg_disks_aiml
+
+try:
+    from py4DSTEM.braggvectors.diskdetection_aiml import find_Bragg_disks_aiml
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 
 
 def find_Bragg_disks(

py4DSTEM/datacube/virtualimage.py

@@ -5,7 +5,6 @@
 # for bragg virtual imaging methods, goto diskdetection.virtualimage.py
 
 import numpy as np
-import dask.array as da
 from typing import Optional
 import inspect
 
@@ -220,7 +219,9 @@ def get_virtual_image(
                     virtual_image[rx, ry] = np.sum(self.data[rx, ry] * mask)
 
             # dask
-            if dask is True:
+            if dask:
+                import dask.array as da
+
                 # set up a generalized universal function for dask distribution
                 def _apply_mask_dask(self, mask):
                     virtual_image = np.sum(

py4DSTEM/io/filereaders/__init__.py

@@ -1,6 +1,12 @@
+from py4DSTEM import is_package_lite
+from py4DSTEM.io.filereaders.empad import read_empad
 from py4DSTEM.io.filereaders.read_dm import read_dm
 from py4DSTEM.io.filereaders.read_K2 import read_gatan_K2_bin
-from py4DSTEM.io.filereaders.empad import read_empad
 from py4DSTEM.io.filereaders.read_mib import load_mib
-from py4DSTEM.io.filereaders.read_arina import read_arina
+
+try:
+    from py4DSTEM.io.filereaders.read_arina import read_arina
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 from py4DSTEM.io.filereaders.read_abTEM import read_abTEM

py4DSTEM/io/importfile.py

@@ -7,7 +7,6 @@
 from py4DSTEM.io.filereaders import (
     load_mib,
     read_abTEM,
-    read_arina,
     read_dm,
     read_empad,
     read_gatan_K2_bin,
@@ -90,6 +89,8 @@ def import_file(
     elif filetype == "mib":
         data = load_mib(filepath, mem=mem, binfactor=binfactor, **kwargs)
     elif filetype == "arina":
+        from py4DSTEM.io.filereaders import read_arina
+
         data = read_arina(filepath, mem=mem, binfactor=binfactor, **kwargs)
     elif filetype == "abTEM":
         data = read_abTEM(filepath, mem=mem, binfactor=binfactor, **kwargs)

py4DSTEM/process/__init__.py

@@ -1,9 +1,21 @@
+from py4DSTEM import is_package_lite
 from py4DSTEM.process.polar import PolarDatacube
 from py4DSTEM.process.strain.strain import StrainMap
 
 from py4DSTEM.process import phase
 from py4DSTEM.process import calibration
 from py4DSTEM.process import utils
-from py4DSTEM.process import classification
+
+try:
+    from py4DSTEM.process import classification
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
+
 from py4DSTEM.process import diffraction
-from py4DSTEM.process import wholepatternfit
+
+try:
+    from py4DSTEM.process import wholepatternfit
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc

py4DSTEM/process/diffraction/crystal_viz.py

@@ -8,7 +8,6 @@
 from scipy.signal import medfilt
 from scipy.ndimage import gaussian_filter
 from scipy.ndimage import distance_transform_edt
-from skimage.morphology import dilation, erosion
 
 import warnings
 import numpy as np
@@ -1913,6 +1912,8 @@ def plot_clusters(
     for a0 in range(self.cluster_sizes.shape[0]):
         if self.cluster_sizes[a0] >= area_min:
             if outline_grains:
+                from skimage.morphology import erosion
+
                 im_grain[:] = False
                 im_grain[
                     self.cluster_inds[a0][0, :],

py4DSTEM/process/phase/__init__.py

@@ -2,6 +2,7 @@
 
 _emd_hook = True
 
+from py4DSTEM import is_package_lite
 from py4DSTEM.process.phase.dpc import DPC
 from py4DSTEM.process.phase.magnetic_ptychographic_tomography import MagneticPtychographicTomography
 from py4DSTEM.process.phase.magnetic_ptychography import MagneticPtychography
@@ -11,7 +12,11 @@
 from py4DSTEM.process.phase.parallax import Parallax
 from py4DSTEM.process.phase.ptychographic_tomography import PtychographicTomography
 from py4DSTEM.process.phase.singleslice_ptychography import SingleslicePtychography
-from py4DSTEM.process.phase.parameter_optimize import OptimizationParameter, PtychographyOptimizer
 from py4DSTEM.process.phase.xray_magnetic_ptychography import XRayMagneticPtychography
+try:
+    from py4DSTEM.process.phase.parameter_optimize import OptimizationParameter, PtychographyOptimizer
+except (ImportError, ModuleNotFoundError) as exc:
+    if not is_package_lite:
+        raise exc
 
 # fmt: on

py4DSTEM/process/phase/utils.py

@@ -22,7 +22,6 @@ def get_array_module(*args):
 from py4DSTEM.process.utils import get_CoM
 from py4DSTEM.process.utils.cross_correlate import align_and_shift_images
 from py4DSTEM.process.utils.utils import electron_wavelength_angstrom
-from skimage.restoration import unwrap_phase
 
 # fmt: off
 
@@ -1755,6 +1754,8 @@ def unwrap_phase_2d(array, weights=None, gauge=None, corner_centered=True, xp=np
 
 
 def unwrap_phase_2d_skimage(array, corner_centered=True, xp=np):
+    from skimage.restoration import unwrap_phase
+
     if xp is np:
         array = array.astype(np.float64)
         unwrapped_array = unwrap_phase(array, wrap_around=corner_centered).astype(

py4DSTEM/process/polar/polar_analysis.py

@@ -5,7 +5,6 @@
 import matplotlib.pyplot as plt
 from scipy.optimize import curve_fit
 from scipy.ndimage import gaussian_filter
-from sklearn.decomposition import PCA
 
 from emdfile import tqdmnd
 
@@ -980,6 +979,8 @@ def background_pca(
         radial PCA component selected
     """
 
+    from sklearn.decomposition import PCA
+
     # PCA decomposition
     shape = self.radial_all.shape
     A = np.reshape(self.radial_all, (shape[0] * shape[1], shape[2]))

py4DSTEM/process/polar/polar_peaks.py

@@ -3,7 +3,6 @@
 
 from scipy.ndimage import gaussian_filter, gaussian_filter1d
 from scipy.signal import peak_prominences
-from skimage.feature import peak_local_max
 from scipy.optimize import curve_fit, leastsq
 import warnings
 
@@ -105,6 +104,8 @@ def find_peaks_single_pattern(
 
     """
 
+    from skimage.feature import peak_local_max
+
     # if needed, generate mask from Bragg peaks
     if bragg_peaks is not None:
         mask_bragg = self._datacube.get_braggmask(

py4DSTEM/utils/configuration_checker.py

@@ -5,13 +5,19 @@
 import re
 from importlib.util import find_spec
 
+from py4DSTEM import is_package_lite
+
 # need a mapping of pypi/conda names to import names
-import_mapping_dict = {
-    "scikit-image": "skimage",
-    "scikit-learn": "sklearn",
-    "scikit-optimize": "skopt",
-    "mp-api": "mp_api",
-}
+import_mapping_dict = (
+    {}
+    if is_package_lite
+    else {
+        "scikit-image": "skimage",
+        "scikit-learn": "sklearn",
+        "scikit-optimize": "skopt",
+        "mp-api": "mp_api",
+    }
+)
 
 
 # programatically get all possible requirements in the import name style
@@ -88,7 +94,8 @@ def get_modules_dict():
 
 
 # module_depenencies = get_modules_dict()
-modules = get_modules_list()
+# modules = get_modules_list()
+modules = [] if is_package_lite else get_modules_list()
 
 
 #### Class and Functions to Create Coloured Strings ####
@@ -526,9 +533,13 @@ def print_no_extra_checks(m: str):
 
 
 # dict of extra check functions
-funcs_dict = {
-    "cupy": check_cupy_gpu,
-}
+funcs_dict = (
+    {}
+    if is_package_lite
+    else {
+        "cupy": check_cupy_gpu,
+    }
+)
 
 
 #### main function used to check the configuration of the installation

setup.py

@@ -20,7 +20,7 @@
     author="Benjamin H. Savitzky",
     author_email="[email protected]",
     license="GNU GPLv3",
-    keywords="STEM 4DSTEM",
+    keywords="STEM,4DSTEM",
     python_requires=">=3.10",
     install_requires=[
         "numpy >= 1.19",