(1) now using DCT unwrap implementation even without weights, which allows full Rytov calculation on the GPU and is much faster. (2) renamed _to_cpu() to to_cpu(). (3) other minor improvements

Author: ptbrown1729

mcsim/analysis/optimize.py

@@ -2,6 +2,7 @@
 Tools for solving inverse problems using accelerated proximal gradient methods
 """
 
+import warnings
 import numpy as np
 import time
 import random
@@ -25,7 +26,7 @@
 array = Union[np.ndarray, cp.ndarray]
 
 
-def _to_cpu(m):
+def to_cpu(m):
     """
     Ensure array is CPU/NumPy
     :param m:
@@ -204,7 +205,7 @@ def run(self,
                    "niterations": max_iterations,
                    "use_fista": use_fista,
                    "use_gpu": use_gpu,
-                   "x_init": _to_cpu(xp.array(x_start, copy=True)),
+                   "x_init": to_cpu(xp.array(x_start, copy=True)),
                    "prox_parameters": self.prox_parameters,
                    "stop_condition": "ok"
                    }
@@ -252,9 +253,9 @@ def run(self,
 
                 if compute_cost:
                     if compute_all_costs:
-                        costs[ii] = _to_cpu(self.cost(x))
+                        costs[ii] = to_cpu(self.cost(x))
                     else:
-                        costs[ii, inds] = _to_cpu(self.cost(x, inds=inds))
+                        costs[ii, inds] = to_cpu(self.cost(x, inds=inds))
 
                 timing["cost"] = np.concatenate((timing["cost"], np.array([time.perf_counter() - tstart_err])))
 
@@ -282,11 +283,11 @@ def run(self,
 
                 if compute_all_costs:
                     c_all = self.cost(x)
-                    costs[ii] = _to_cpu(c_all)
+                    costs[ii] = to_cpu(c_all)
                     cx = xp.mean(c_all[inds], axis=0)
                 else:
                     c_now = self.cost(x, inds=inds)
-                    costs[ii, inds] = _to_cpu(c_now)
+                    costs[ii, inds] = to_cpu(c_now)
                     cx = xp.mean(c_now, axis=0)
 
                 timing["cost"] = np.concatenate((timing["cost"], np.array([time.perf_counter() - tstart_err])))
@@ -347,15 +348,19 @@ def lipschitz_condition_violated(y, cx, gx):
 
             # print information
             if verbose:
-                status = f"iteration {ii + 1:d}/{max_iterations:d}," \
-                         f" cost={np.nanmean(costs[ii]):.3g}," \
-                         f" step={steps[ii]:.3g}," \
-                         f" line search iters={line_search_iters[ii]:d}," \
-                         f" grad={timing['grad'][ii]:.3f}s," \
-                         f" prox={timing['prox'][ii]:.3f}s," \
-                         f" cost={timing['cost'][ii]:.3f}s," \
-                         f" iter={timing['iteration'][ii]:.3f}s," \
-                         f" total={time.perf_counter() - tstart:.3f}s"
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore", category=RuntimeWarning)
+
+                    status = f"iteration {ii + 1:d}/{max_iterations:d}," \
+                             f" cost={np.nanmean(costs[ii]):.3g}," \
+                             f" step={steps[ii]:.3g}," \
+                             f" line search iters={line_search_iters[ii]:d}," \
+                             f" grad={timing['grad'][ii]:.3f}s," \
+                             f" prox={timing['prox'][ii]:.3f}s," \
+                             f" cost={timing['cost'][ii]:.3f}s," \
+                             f" iter={timing['iteration'][ii]:.3f}s," \
+                             f" total={time.perf_counter() - tstart:.3f}s"
+
                 if use_gpu:
                     status += f", GPU={mempool.used_bytes()/1e9:.3}GB"
 
@@ -369,9 +374,9 @@ def lipschitz_condition_violated(y, cx, gx):
         # compute final cost
         if compute_cost:
             if compute_all_costs:
-                costs[ii + 1] = _to_cpu(self.cost(x))
+                costs[ii + 1] = to_cpu(self.cost(x))
             else:
-                costs[ii + 1, inds] = _to_cpu(self.cost(x, inds=inds))
+                costs[ii + 1, inds] = to_cpu(self.cost(x, inds=inds))
 
         # store results
         results.update({"timing": timing,

mcsim/analysis/sim_reconstruction.py

@@ -46,7 +46,7 @@
 from matplotlib.patches import Circle, Rectangle
 # code from our projects
 import mcsim.analysis.analysis_tools as tools
-from mcsim.analysis.optimize import Optimizer, soft_threshold, tv_prox, _to_cpu
+from mcsim.analysis.optimize import Optimizer, soft_threshold, tv_prox, to_cpu
 from mcsim.analysis.fft import ft2, ift2
 from localize_psf.rois import get_centered_rois, cut_roi
 from localize_psf.fit_psf import circ_aperture_otf, blur_img_psf, oversample_voxel
@@ -1332,11 +1332,11 @@ def proc_bands(bands, phases, amps, frqs, dy, dx, upsample_fact, use_gpu):
 
                 # if cupy array, move off GPU
                 if isinstance(attr, cp.ndarray):
-                    setattr(self, attr_name, _to_cpu(attr))
+                    setattr(self, attr_name, to_cpu(attr))
 
                 # if dask array, move off GPU delayed
                 if isinstance(attr, da.core.Array):
-                    on_cpu = da.map_blocks(_to_cpu, attr, dtype=attr.dtype)
+                    on_cpu = da.map_blocks(to_cpu, attr, dtype=attr.dtype)
                     setattr(self, attr_name, on_cpu)
 
         self.print_log(f"reconstruction took {time.perf_counter() - tstart_recon:.2f}s")
@@ -2370,13 +2370,28 @@ def show_sim_napari(fname_zarr: str,
 
     dxy = imgz.attrs["dx"]
     dxy_sim = dxy / imgz.attrs["upsample_factor"]
+
+    # translate to put FFT zero coordinates at origin
     translate_wf = (-(wf.shape[-2] // 2) * dxy, -(wf.shape[-1] // 2) * dxy)
     translate_sim = (-((2 * wf.shape[-2]) // 2) * dxy_sim, -((2 * wf.shape[-1]) // 2) * dxy_sim)
+    translate_pattern_2x = [a - 0.25 * dxy for a in translate_wf]
 
     if viewer is None:
         viewer = napari.Viewer()
 
-    # translate to put FFT zero coordinates at origin
+    if hasattr(imgz, "patterns"):
+        viewer.add_image(imgz.patterns,
+                         scale=(dxy, dxy),
+                         translate=translate_wf,
+                         name="patterns")
+
+    if hasattr(imgz, "patterns_2x"):
+        viewer.add_image(imgz.patterns_2x,
+                         scale=(dxy_sim, dxy_sim),
+                         # translate=translate_sim,
+                         translate=translate_pattern_2x,
+                         name="patterns upsampled")
+
     if hasattr(imgz, "sim_os"):
         sim_os = np.expand_dims(imgz.sim_os, axis=-3)
 
@@ -2392,6 +2407,18 @@ def show_sim_napari(fname_zarr: str,
                          name="wf deconvolved",
                          visible=False)
 
+    if hasattr(imgz, "sim_fista_forward_model"):
+        viewer.add_image(imgz.sim_fista_forward_model,
+                         scale=(dxy, dxy),
+                         translate=translate_wf,
+                         name="FISTA forward model")
+
+    if hasattr(imgz, "sim_sr_fista"):
+        viewer.add_image(np.expand_dims(imgz.sim_sr_fista, axis=-3),
+                         scale=(dxy_sim, dxy_sim),
+                         translate=translate_pattern_2x,
+                         name="SIM-SR FISTA")
+
     if hasattr(imgz, "sim_sr"):
         viewer.add_image(np.expand_dims(imgz.sim_sr, axis=-3),
                          scale=(dxy_sim, dxy_sim),
@@ -2420,19 +2447,6 @@ def show_sim_napari(fname_zarr: str,
                          translate=translate_wf,
                          name="raw images")
 
-    if hasattr(imgz, "patterns"):
-        viewer.add_image(imgz.patterns,
-                         scale=(dxy, dxy),
-                         translate=translate_wf,
-                         name="patterns")
-
-    if hasattr(imgz, "patterns_2x"):
-        viewer.add_image(imgz.patterns_2x,
-                         scale=(dxy_sim, dxy_sim),
-                         # translate=translate_sim,
-                         translate=[a - 0.25 * dxy for a in translate_wf],
-                         name="patterns upsampled")
-
     viewer.show(block=block)
 
     return viewer
@@ -2582,10 +2596,10 @@ def fit_modulation_frq(mft1: np.ndarray,
         raise ValueError("must have ft1.shape = ft2.shape")
 
     # must be on CPU for this function to work
-    mft1 = _to_cpu(mft1)
-    mft2 = _to_cpu(mft2)
+    mft1 = to_cpu(mft1)
+    mft2 = to_cpu(mft2)
     if otf is not None:
-        otf = _to_cpu(otf)
+        otf = to_cpu(otf)
 
     # mask
     if mask is None: