Merge pull request #254 from carterbox/cluster-compact

REF: Cluster compact

.github/workflows/unit-tests.yml

@@ -27,7 +27,7 @@ jobs:
       matrix:
         python-version:
           - "3.10"
-          - "3.7"
+          - "3.11"
           - "3.9"
 
     steps:

pyproject.toml

@@ -28,15 +28,14 @@ classifiers = [
 dependencies = [
   "cupy >=10.0, !=10.3.0",
   'importlib_resources; python_version<"3.9"',
-  'importlib_metadata; python_version<"3.8"',
   "matplotlib ==3.*",
   "numpy ~=1.17",
   "opencv-python >=3.4, <5.0",
   "scipy >=1.6.0",
 ]
 license = {text = "BSD 3-Clause License"}
 name = "tike"
-requires-python = "~=3.6"
+requires-python = "~=3.8"
 dynamic = ["version", "readme"]
 optional-dependencies = {mpi = [
       "mpi4py ==3.*"

requirements.txt

@@ -1,11 +1,10 @@
 cupy>=10,!=10.3.0
-importlib_metadata # backport for python<3.8
 importlib_resources # backport for python<3.9
 matplotlib-base=3.*
 mpi4py>=3.*
 numpy>=1.17
 py-opencv>=3.4,<5.0
-python>=3.6
+python>=3.8
 scipy>=1.6.0
 setuptools_scm>=7
 setuptools>=45

src/tike/communicators/pool.py

@@ -35,6 +35,7 @@ class ThreadPool(ThreadPoolExecutor):
         When invalid GPU device ids are provided.
         When the current CUDA device does not match the first GPU id in the
         list of workers.
+
     """
 
     Device = cp.cuda.Device
@@ -316,11 +317,12 @@ def reduce_mean(
     ) -> cp.ndarray:
         """Reduce x by addition to one GPU from all other GPUs."""
         worker = self.workers[0] if worker is None else worker
-        return cp.mean(
-            self.gather(x, worker=worker, axis=axis),
-            keepdims=x[0].ndim > 0,
-            axis=axis,
-        )
+        with self.Device(worker):
+            return cp.mean(
+                self.gather(x, worker=worker, axis=axis),
+                keepdims=x[0].ndim > 0,
+                axis=axis,
+            )
 
     def allreduce(
         self,

src/tike/ptycho/object.py

@@ -7,10 +7,12 @@
 from __future__ import annotations
 import dataclasses
 import logging
+import typing
 
 import cupy as cp
 import cupyx.scipy.ndimage
 import numpy as np
+import numpy.typing as npt
 import scipy.interpolate
 
 import tike.linalg
@@ -40,18 +42,28 @@ class ObjectOptions:
     mdecay: float = 0.9
     """The proportion of the first moment that is previous first moments."""
 
-    v: np.array = dataclasses.field(init=False, default_factory=lambda: None)
+    v: typing.Union[npt.NDArray, None] = dataclasses.field(
+        init=False,
+        default_factory=lambda: None,
+    )
     """The second moment for adaptive moment."""
 
-    m: np.array = dataclasses.field(init=False, default_factory=lambda: None)
+    m: typing.Union[npt.NDArray, None] = dataclasses.field(
+        init=False,
+        default_factory=lambda: None,
+    )
     """The first moment for adaptive moment."""
 
-    preconditioner: np.array = dataclasses.field(init=False,
-                                                 default_factory=lambda: None)
+    preconditioner: typing.Union[npt.NDArray, None] = dataclasses.field(
+        init=False,
+        default_factory=lambda: None,
+    )
     """The magnitude of the illumination used for conditioning the object updates."""
 
-    combined_update: np.array = dataclasses.field(init=False,
-                                                  default_factory=lambda: None)
+    combined_update: typing.Union[npt.NDArray, None] = dataclasses.field(
+        init=False,
+        default_factory=lambda: None,
+    )
     """Used for compact batch updates."""
 
     clip_magnitude: bool = True
@@ -174,7 +186,7 @@ def get_absorbtion_image(data, scan, *, rescale=1.0, method='cubic'):
     ----------
     data : (FRAME, WIDE, HIGH)
         The intensity (square of the absolute value) of the propagated
-        wavefront; i.e. what the detector records. FFT-shifted so the
+        wavefront i.e. what the detector records. FFT-shifted so the
         diffraction peak is at the corners.
     scan : (POSI, 2) float32
         Coordinates of the minimum corner of the probe grid for each
@@ -203,3 +215,17 @@ def get_absorbtion_image(data, scan, *, rescale=1.0, method='cubic'):
         fill_value=np.amax(values),
     )
     return np.reshape(absorption_image, coord0.shape)
+
+
+def remove_object_ambiguity(
+    psi: npt.NDArray[cp.complex64],
+    probe: npt.NDArray[cp.complex64],
+    preconditioner: npt.NDArray[cp.complex64],
+) -> typing.Tuple[npt.NDArray[cp.complex64], npt.NDArray[cp.complex64]]:
+    """Remove normalization ambiguity between probe and psi"""
+    W = preconditioner.real
+    W = W / tike.linalg.mnorm(W)
+    object_norm = 2 * np.sqrt(np.mean(np.square(np.abs(psi)) * W))
+    psi = psi / object_norm
+    probe = probe * object_norm
+    return psi, probe

src/tike/ptycho/probe.py

@@ -118,12 +118,25 @@ class ProbeOptions:
     hard constraint.
     """
 
-    def copy_to_device(self):
+    probe_update_sum: typing.Union[npt.NDArray, None] = dataclasses.field(
+        init=False,
+        default_factory=lambda: None,
+    )
+    """Used for momentum updates."""
+
+    preconditioner: typing.Union[npt.NDArray, None] = dataclasses.field(
+        init=False,
+        default_factory=lambda: None,
+    )
+
+    def copy_to_device(self, comm):
         """Copy to the current GPU memory."""
         if self.v is not None:
             self.v = cp.asarray(self.v)
         if self.m is not None:
             self.m = cp.asarray(self.m)
+        if self.preconditioner is not None:
+            self.preconditioner = comm.pool.bcast([self.preconditioner])
         return self
 
     def copy_to_host(self):
@@ -132,6 +145,8 @@ def copy_to_host(self):
             self.v = cp.asnumpy(self.v)
         if self.m is not None:
             self.m = cp.asnumpy(self.m)
+        if self.preconditioner is not None:
+            self.preconditioner = cp.asnumpy(self.preconditioner[0])
         return self
 
     def resample(self, factor: float) -> ProbeOptions:
@@ -580,7 +595,7 @@ def add_modes_cartesian_hermite(probe, nmodes: int):
             new_probes.append(basis)
 
             if len(new_probes) == nmodes:
-                return np.concatenate(new_probes, axis=-3)
+                return np.concatenate(new_probes, axis=-3)[..., :nmodes, :, :]
 
     raise RuntimeError(
         "`add_modes_cartesian_hermite` never reached a return statement."
@@ -695,15 +710,11 @@ def orthogonalize_eig(x):
     # We find the eigen vectors of x^H @ x in order to get v^H from SVD of x
     # without computing u, s.
     val, vectors = xp.linalg.eigh(A, UPLO='U')
-    # np.linalg.eigh guarantees that the eigen values are returned in ascending
-    # order, so we just reverse the order of modes to have them sorted in
-    # descending order.
-    vectors = vectors[..., ::-1].swapaxes(-1, -2)
-    result = (vectors @ x.reshape(*x.shape[:-2], -1)).reshape(*x.shape)
-    assert np.all(
-        np.diff(tike.linalg.norm(result, axis=(-2, -1), keepdims=False),
-                axis=-1) <= 0
-    ), f"Power of the orthogonalized probes should be monotonically decreasing! {val}"
+    result = (vectors.swapaxes(-1, -2) @ x.reshape(*x.shape[:-2], -1)).reshape(
+        *x.shape)
+    power = tike.linalg.norm(result, axis=(-2, -1), keepdims=False)
+    result = result[...,
+                    np.argsort(power, axis=None, kind='stable')[::-1], :, :]
     return result
 
 

src/tike/ptycho/ptycho.py

@@ -369,7 +369,7 @@ def __enter__(self):
 
         if self._device_parameters.probe_options is not None:
             self._device_parameters.probe_options = self._device_parameters.probe_options.copy_to_device(
-            )
+                self.comm,)
 
         if self._device_parameters.object_options is not None:
             self._device_parameters.object_options = self._device_parameters.object_options.copy_to_device(