Add scan and tilt plots

Author: hexane360

phaser/engines/common/output.py

@@ -17,9 +17,10 @@ def output_images(state: ReconsState, out_dir: Path, options: SaveOptions):
         if ty not in _SAVE_FUNCS:
             raise ValueError(f"Unknown image type '{ty}'")
 
+        ext = options.plot_ext if ty in _PLOT_FUNCS else 'tiff'
         try:
             out_name = options.img_fmt.format(
-                type=ty, iter=state.iter,
+                type=ty, iter=state.iter, ext=ext
             )
             out_path = out_dir / out_name
         except KeyError as e:
@@ -175,6 +176,60 @@ def _save_object_mag(state: ReconsState, out_path: Path, options: SaveOptions, s
         w.write(obj_mag, **write_opts)
 
 
+def _plot_scan(state: ReconsState, out_path: Path, options: SaveOptions):
+    from matplotlib import pyplot
+    fig, ax = pyplot.subplots(figsize=(4, 4), dpi=options.plot_dpi, constrained_layout=True)
+
+    ax.set_aspect(1.)
+    [l, r, b, t] = state.object.sampling.mpl_extent()
+    ax.set_xlim(l, r)
+    ax.set_ylim(b, t)
+
+    scan = to_numpy(state.scan)
+    i = numpy.arange(scan[..., 0].size)
+    ax.scatter(scan[..., 1].ravel(), scan[..., 0].ravel(), c=i, s=0.2, cmap='plasma')
+
+    fig.savefig(out_path)
+    pyplot.close(fig)
+
+
+def _plot_tilt(state: ReconsState, out_path: Path, options: SaveOptions):
+    from matplotlib import pyplot
+
+    fig, ax = pyplot.subplots(figsize=(4, 4), dpi=options.plot_dpi, constrained_layout=True)
+
+    ax.set_aspect(1.)
+    [l, r, b, t] = state.object.sampling.mpl_extent()
+    ax.set_xlim(l, r)
+    ax.set_ylim(b, t)
+
+    scan = to_numpy(state.scan)
+    tilt = to_numpy(state.tilt)
+    tilt = tilt[..., 1] + tilt[..., 0]*1.j
+    max_tilt = max(numpy.max(numpy.abs(tilt)), 1.0)  # at least 1 mrad
+    c = colorize_complex(tilt / max_tilt, amp=True, rescale=False)
+    ax.scatter(scan[..., 1].ravel(), scan[..., 0].ravel(), c=c, s=0.2)
+
+    fig.draw_without_rendering()
+    trans = ax.transAxes + fig.transFigure.inverted()
+    legend_ax_max = trans.transform([0.95, 0.02])
+    legend_ax_size = (0.1, 0.1)
+    legend_ax = fig.add_axes((legend_ax_max[0] - legend_ax_size[0], legend_ax_max[1], *legend_ax_size), projection='polar')
+
+    legend_ax.set_rmax(max_tilt)  # type: ignore
+    legend_ax.set_theta_direction(-1)  # type: ignore
+    legend_ax.set_axis_off()
+
+    thetas = numpy.linspace(0., 2*numpy.pi, 70)
+    rs = numpy.concatenate([[0.0], numpy.geomspace(0.1, 1.0, 30)])
+    rr, tt = numpy.meshgrid(rs, thetas, indexing='ij')
+    c2 = colorize_complex(rr * numpy.exp(1.j * tt), rescale=False)
+    legend_ax.pcolormesh(tt, rr * max_tilt, c2)
+    legend_ax.text(-numpy.pi/2., max_tilt * 1.05, f"{max_tilt:.1f} mrad", ha='center', va='bottom', size='small')
+
+    fig.savefig(out_path)
+    pyplot.close(fig)
+
 
 _SAVE_FUNCS: t.Dict[str, t.Callable[[ReconsState, Path, SaveOptions], t.Any]] = {
     'probe': _save_probe,
@@ -185,4 +240,8 @@ def _save_object_mag(state: ReconsState, out_path: Path, options: SaveOptions, s
     'object_phase_sum': partial(_save_object_phase, stack=False),
     'object_mag_stack': partial(_save_object_mag, stack=True),
     'object_mag_sum': partial(_save_object_mag, stack=False),
-}
+    'scan': _plot_scan,
+    'tilt': _plot_tilt,
+}
+# save functions with special handling of file extensions
+_PLOT_FUNCS: t.Set[str] = {'scan', 'tilt'}

phaser/plan.py

@@ -12,6 +12,7 @@
     'probe', 'probe_mag', 'probe_recip', 'probe_recip_mag',
     'object_phase_stack', 'object_phase_sum',
     'object_mag_stack', 'object_mag_sum',
+    'scan', 'tilt',
 ]
 
 
@@ -27,11 +28,15 @@ class InitPlan(Dataclass, kw_only=True):
 class SaveOptions(Dataclass, kw_only=True):
     images: t.Tuple[SaveType, ...] = ('probe', 'object_phase_stack')
     crop_roi: bool = True
-    unwrap_phase: bool = True
+    unwrap_phase: bool = False
     img_dtype: t.Literal['float', '8bit', '16bit', '32bit'] = '16bit'
 
+    plot_ext: str = "svg"
+    """Extension to use for matplotlib savefig"""
+    plot_dpi: int = 300
+
     out_dir: str = "{name}"
-    img_fmt: str = "{type}_iter{iter.total_iter}.tiff"
+    img_fmt: str = "{type}_iter{iter.total_iter}.{ext}"
     hdf5_fmt: str = "iter{iter.total_iter}.h5"
 
 

phaser/utils/image.py

@@ -54,23 +54,21 @@ def remove_linear_ramp(
     return output
 
 
-def colorize_complex(vals: ArrayLike, magnitude_only=False) -> NDArray[numpy.floating]:
+def colorize_complex(vals: ArrayLike, amp: bool = False, rescale: bool = True) -> NDArray[numpy.floating]:
     """Colorize a ndarray of complex values as rgb."""
     from matplotlib.colors import hsv_to_rgb
     xp = get_array_module(vals)
 
     vals = xp.asarray(vals, dtype=numpy.complexfloating)
-    mag = abs2(vals)
-    arg = xp.angle(vals) 
-    max_mag = xp.max(mag)
 
-    if magnitude_only:
-        return mag
+    v = xp.abs(vals) if amp else abs2(vals)
+    if rescale:
+        v /= xp.max(v)
+    arg = xp.angle(vals) 
 
     h = (arg + numpy.pi) / (2*numpy.pi)
-    s = 0.85 * xp.ones_like(mag)
-    v = mag / max_mag 
-    return hsv_to_rgb(to_numpy(xp.stack((h, s, v), axis=-1)))
+    s = 0.85 * xp.ones_like(v)
+    return xp.clip(hsv_to_rgb(to_numpy(xp.stack((h, s, v), axis=-1))), 0.0, 1.0)
 
 
 def scale_to_integral_type(