NEW: Solve the lamino-alignment problem

.gitignore

@@ -118,5 +118,7 @@ core
 .DS_Store
 .nfs*
 
+*.tiff
+
 # MPI host files
 host_list

src/tike/admm/__init__.py

@@ -0,0 +1,4 @@
+from .admm import *
+from .al import *
+from .pal import *
+from .pl import *

src/tike/admm/admm.py

@@ -0,0 +1,69 @@
+import logging
+
+import numpy as np
+
+import tike.align
+import tike.lamino
+import tike.ptycho
+
+
+def simulate(
+    u,
+    scan,
+    probe,
+    flow,
+    angle,
+    tilt,
+    theta,
+    padded_shape,
+    detector_shape,
+):
+    phi = np.exp(1j * tike.lamino.simulate(
+        obj=u,
+        tilt=tilt,
+        theta=theta,
+    ))
+    psi = tike.align.simulate(
+        original=phi,
+        flow=flow,
+        padded_shape=padded_shape,
+        angle=angle,
+        cval=1.0,
+    )
+    data = tike.ptycho.simulate(
+        psi=psi,
+        probe=probe,
+        detector_shape=detector_shape,
+        scan=scan,
+    )
+    return data, psi, phi
+
+
+def print_log_line(**kwargs):
+    """Print keyword arguments and values on a single comma-separated line.
+
+    The format of the line is as follows:
+
+    ```
+    foo: 003, bar: +1.234e+02, hello: world\n
+    ```
+
+    Parameters
+    ----------
+    line: dictionary
+        The key value pairs to be printed.
+
+    """
+    line = []
+    for k, v in kwargs.items():
+        # Use special formatting for float and integers
+        if isinstance(v, (float, np.floating)):
+            line.append(f'"{k}": {v:6.3e}')
+        elif isinstance(v, (int, np.integer)):
+            line.append(f'"{k}": {v:3d}')
+        elif isinstance(v, str):
+            line.append(f'"{k}": "{v}"')
+        else:
+            line.append(f'"{k}": {v}')
+    # Combine all the strings and strip the last comma
+    print("{", ", ".join(line), "}", flush=True)

src/tike/admm/al.py

@@ -0,0 +1,156 @@
+import logging
+
+import numpy as np
+import cupy as cp
+
+import tike.admm.subproblem
+import tike.communicator
+
+from .admm import print_log_line
+
+logger = logging.getLogger(__name__)
+
+
+def ptycho__align_lamino(
+    data,
+    psi,
+    scan,
+    probe,
+    theta,
+    tilt,
+    angle,
+    w,
+    flow=None,
+    shift=None,
+    niter=1,
+    interval=8,
+    folder=None,
+    cg_iter=4,
+    align_method=False,
+    skip_ptycho=False,
+):
+    """Solve the joint ptycho-lamino problem using ADMM."""
+    presult = {
+        'psi': psi,
+        'scan': scan,
+        'probe': probe,
+    }
+
+    u = np.zeros((w, w, w), dtype='complex64')
+    Hu = np.ones((len(theta), w, w), dtype='complex64')
+    phi = Hu
+    Aφ = np.ones(psi.shape, dtype='complex64')
+
+    λ_l = np.zeros([len(theta), w, w], dtype='complex64')
+    ρ_l = 0.5
+
+    comm = tike.communicator.MPICommunicator()
+
+    with cp.cuda.Device(comm.rank if comm.size > 1 else None):
+
+        if not skip_ptycho:
+            presult, _ = tike.admm.subproblem.ptycho(
+                # constants
+                comm=comm,
+                data=data,
+                λ=None,
+                ρ=None,
+                Aφ=None,
+                # updated
+                presult=presult,
+                # parameters
+                num_iter=4 * niter,
+                cg_iter=cg_iter,
+                folder=folder,
+                save_result=niter + 1,
+                rescale=True,
+                rtol=1e-6,
+            )
+
+        for k in range(1, niter + 1):
+            logger.info(f"Start ADMM iteration {k}.")
+            save_result = k if k % interval == 0 else False
+
+            (
+                phi,
+                _,
+                _,
+                flow,
+                shift,
+                Aφ,
+                align_cost,
+                winsize,
+            ) = tike.admm.subproblem.align(
+                # constants
+                comm=comm,
+                psi=presult['psi'],
+                angle=angle,
+                Hu=Hu,
+                λ_l=λ_l,
+                ρ_l=ρ_l,
+                # updated
+                phi=phi,
+                λ_p=None,
+                ρ_p=1,
+                flow=flow,
+                shift=shift,
+                Aφ0=None,
+                # parameters
+                align_method=align_method,
+                cg_iter=cg_iter,
+                num_iter=4,
+                folder=folder,
+                save_result=save_result,
+                winsize=winsize if k > 1 else 129,
+            )
+
+            (
+                u,
+                λ_l,
+                ρ_l,
+                Hu,
+                lamino_cost,
+            ) = tike.admm.subproblem.lamino(
+                # constants
+                comm=comm,
+                phi=phi,
+                theta=theta,
+                tilt=tilt,
+                # updated
+                u=u,
+                λ_l=λ_l,
+                ρ_l=ρ_l,
+                Hu0=Hu,
+                # parameters
+                num_iter=4,
+                cg_iter=cg_iter,
+                folder=folder,
+                save_result=save_result,
+            )
+
+            # Record metrics for each subproblem
+            ψAφ = presult['psi'] - Aφ
+            φHu = phi - Hu
+            lagrangian = (
+                [np.mean(np.real(ψAφ.conj() * ψAφ))],
+                [
+                    2 * np.mean(np.real(λ_l.conj() * φHu)) +
+                    ρ_l * np.mean(np.real(φHu.conj() * φHu))
+                ],
+                [align_cost],
+            )
+            lagrangian = [comm.gather(x) for x in lagrangian]
+
+            if comm.rank == 0:
+                lagrangian = [np.sum(x) for x in lagrangian]
+                print_log_line(
+                    k=k,
+                    ρ_l=ρ_l,
+                    winsize=winsize,
+                    align_method=align_method,
+                    Lagrangian=np.sum(lagrangian[:2]),
+                    ψAφ=lagrangian[0],
+                    φHu=lagrangian[1],
+                    align=lagrangian[2],
+                    lamino=float(lamino_cost),
+                )