Merge pull request #257 from carterbox/assorted-upgrades

NEW: Assorted upgrades

src/tike/cluster.py

@@ -223,14 +223,16 @@ def compact(population, num_cluster, max_iter=500):
     Parameters
     ----------
     population : (M, N) array_like
-        The M samples of an N dimensional population that needs to be clustered.
+        The M samples of an N dimensional population that needs to be
+        clustered.
     num_cluster : int (0..M]
         The number of clusters in which to divide M samples.
 
     Returns
     -------
     indicies : (num_cluster,) list of array of integer
-        The indicies of population that belong to each cluster.
+        The indicies of population that belong to each cluster. Clusters are
+        sorted from largest to smallest.
 
     Raises
     ------
@@ -371,7 +373,9 @@ def compact(population, num_cluster, max_iter=500):
     if __debug__:
         for c in range(num_cluster):
             _assert_cluster_is_full(labels, c, max_size[c])
-    return [np.flatnonzero(labels == c) for c in range(num_cluster)]
+    indices = [np.flatnonzero(labels == c) for c in range(num_cluster)]
+    indices.sort(key=len, reverse=True)
+    return indices
 
 
 def _k_means_objective(population, labels, num_cluster):

src/tike/opt.py

@@ -6,11 +6,12 @@
 library.
 
 """
-
+import typing
 import logging
 import warnings
 
 import numpy as np
+import numpy.typing as npt
 
 logger = logging.getLogger(__name__)
 randomizer = np.random.default_rng()
@@ -80,7 +81,12 @@ def momentum(g, v, m, vdecay=None, mdecay=0.9):
     return m, None, m
 
 
-def adagrad(g, v=None, eps=1e-6):
+def adagrad(
+    g: npt.NDArray,
+    v: typing.Union[None, npt.NDArray] = None,
+    m: typing.Union[None, npt.NDArray] = None,
+    eps: float = 1e-6,
+):
     """Return the adaptive gradient algorithm direction.
 
     Used to provide a better search direction to stochastic gradient
@@ -109,10 +115,10 @@ def adagrad(g, v=None, eps=1e-6):
     learning research 12, no. 7 (2011).
     """
     if v is None:
-        return g, (g * g.conj()).real
+        return g, (g * g.conj()).real, m
     v += (g * g.conj()).real
     d = g / np.sqrt(v + eps)
-    return d, v
+    return d, v, m
 
 
 def adadelta(g, d0=None, v=None, m=None, decay=0.9, eps=1e-6):
@@ -155,7 +161,14 @@ def adadelta(g, d0=None, v=None, m=None, decay=0.9, eps=1e-6):
     return d, v, m
 
 
-def adam(g, v=None, m=None, vdecay=0.999, mdecay=0.9, eps=1e-8):
+def adam(
+    g: npt.NDArray,
+    v: typing.Union[None, npt.NDArray] = None,
+    m: typing.Union[None, npt.NDArray] = None,
+    vdecay: float = 0.999,
+    mdecay: float = 0.9,
+    eps: float = 1e-8,
+) -> typing.Tuple[npt.NDArray, npt.NDArray, npt.NDArray]:
     """Return the adaptive moment estimation direction.
 
     Used to provide a better search direction to stochastic gradient
@@ -190,8 +203,8 @@ def adam(g, v=None, m=None, vdecay=0.999, mdecay=0.9, eps=1e-8):
     optimization." arXiv preprint arXiv:1412.6980 (2014).
     """
     logger.info("ADAM decay m=%+.3e, v=%+.3e; eps=%+.3e", mdecay, vdecay, eps)
-    v = 0 if v is None else v
-    m = 0 if m is None else m
+    v = np.zeros_like(g.real) if v is None else v
+    m = np.zeros_like(g) if m is None else m
     m = mdecay * m + (1 - mdecay) * g
     v = vdecay * v + (1 - vdecay) * (g * g.conj()).real
     m_ = m / (1 - mdecay)
@@ -364,3 +377,23 @@ def conjugate_gradient(
         cost = cost_function(x)
 
     return x, cost
+
+
+def fit_line_least_squares(
+    y: typing.List[float],
+    x: typing.List[float],
+) -> typing.Tuple[float, float]:
+    """Return the `slope`, `intercept` pair that best fits `y`, `x` to a line.
+
+    y = slope * x + intercept
+
+    """
+    assert len(x) == len(y)
+    count = len(x)
+    assert count > 0
+    sum_x = np.sum(x)
+    sum_y = np.sum(y)
+    slope = (count * np.sum(x * y) - (sum_x * sum_y)) / (count * np.sum(x * x) -
+                                                         (sum_x * sum_x))
+    intercept = (sum_y - slope * sum_x) / count
+    return slope, intercept

src/tike/view.py

@@ -53,6 +53,8 @@
 
 import logging
 import warnings
+import typing
+import itertools
 
 from matplotlib.patches import Ellipse
 import matplotlib.transforms as transforms
@@ -91,13 +93,17 @@ def complexHSV_to_RGB(img0):
 
     hsv_img = np.ones((*sz, 3), 'float32')
 
-    hsv_img[ ..., 0 ] = np.angle( img0 )    # always scaled between +/- pi
-    hsv_img[ ..., 2 ] = np.abs( img0 )      # always scaled between 0 and +inf
+    hsv_img[..., 0] = np.angle(img0)  # always scaled between +/- pi
+    hsv_img[..., 2] = np.abs(img0)  # always scaled between 0 and +inf
+
+    if hsv_img[..., 2].max() > 1.0:
+        raise ValueError('The maximum amplitude of `img0` must be <= 1.0; '
+                         'rescale your image before converting to RGB.')
 
     #================================
     # Rescale hue to the range [0, 1]
 
-    hsv_img[ ..., 0 ] = ( hsv_img[ ..., 0 ] + np.pi ) / ( 2 * np.pi )
+    hsv_img[..., 0] = (hsv_img[..., 0] + np.pi) / (2 * np.pi)
 
     #==================================
     # convert HSV representation to RGB
@@ -157,11 +163,11 @@ def plot_probe_power(probe):
     probe : (..., 1, 1, SHARED, WIDE, HIGH) complex64
         The probes to be analyzed.
     """
-    power = np.square(tike.linalg.norm(
-        probe,
+    power = np.sum(
+        np.square(np.abs(probe)),
         axis=(-2, -1),
         keepdims=False,
-    )).flatten()
+    ).flatten()
     axes = plt.gca()
     axes.bar(
         range(len(power)),
@@ -573,7 +579,10 @@ def plot_trajectories(theta, v, h, t):
     return ax1a, ax1b
 
 
-def plot_cost_convergence(costs, times):
+def plot_cost_convergence(
+    costs: typing.Union[typing.List[float], typing.List[typing.List[float]]],
+    times: typing.List[float],
+):
     """Plot a twined plot of cost vs iteration/cumulative-time
 
     The plot is a semi-log line plot with two lines. One line shows cost as a
@@ -582,7 +591,7 @@ def plot_cost_convergence(costs, times):
 
     Parameters
     ----------
-    costs : (NUM_ITER, ) array-like
+    costs : (NUM_ITER, ), (NUM_ITER, NUM_BATCH, ) array-like
         The objective cost at each iteration.
     times : (NUM_ITER, ) array-like
         The wall-time for each iteration in seconds.
@@ -594,21 +603,30 @@ def plot_cost_convergence(costs, times):
     """
     ax1 = plt.subplot()
 
-    costs = np.asarray(costs)
-    alpha = 1.0 / costs.shape[1] if costs.ndim > 1 else 1.0
+    cost_summary = [np.mean(x) for x in costs]
+    num_iter = np.arange(1, len(times) + 1)
+    # Must handle the case in which the number of batches changes over time.
+    if isinstance(costs[0], list):
+        batches = list(itertools.zip_longest(*costs, fillvalue=None))
+    else:
+        batches = [costs]
+
+    alpha = max(0.05, 1.0 / len(batches[0]))
 
     color = 'black'
     ax1.semilogy()
     ax1.set_xlabel('iteration', color=color)
     ax1.set_ylabel('objective')
-    ax1.plot(costs, linestyle='--', color=color, alpha=alpha)
+    for batch in batches:
+        ax1.plot(num_iter, batch, linestyle='--', color=color, alpha=alpha)
     ax1.tick_params(axis='x', labelcolor=color)
+    ax1.set_xscale('log', base=10)
 
     ax2 = ax1.twiny()
 
     color = 'red'
     ax2.set_xlabel('wall-time [s]', color=color)
-    ax2.plot(np.cumsum(times), costs, color=color, alpha=alpha)
+    ax2.plot(np.cumsum(times), cost_summary, color=color, alpha=1.0)
     ax2.tick_params(axis='x', labelcolor=color)
 
     return ax1, ax2

tests/ptycho/io.py

@@ -49,26 +49,31 @@ def _save_probe(output_folder, probe, algorithm):
         probe.reshape((-1, *probe.shape[-2:])),
         axis=1,
     )
-    with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", category=UserWarning)
-        plt.imsave(
-            f'{output_folder}/probe-phase.png',
-            np.angle(flattened),
-            # The output of np.angle is locked to (-pi, pi]
-            cmap=plt.cm.twilight,
-            vmin=-np.pi,
-            vmax=np.pi,
-        )
-        plt.imsave(
-            f'{output_folder}/probe-ampli.png',
-            np.abs(flattened),
-        )
+    flattened /= (np.abs(flattened).max() * 1.001)
+    plt.imsave(
+        f'{output_folder}/probe.png',
+        tike.view.complexHSV_to_RGB(flattened),
+    )
     f = plt.figure()
     tike.view.plot_probe_power(probe)
     plt.semilogy()
     plt.title(algorithm)
     plt.savefig(f'{output_folder}/probe-power.png')
     plt.close(f)
+    nmodes = probe.shape[-3]
+    probe_orthogonality_matrix = np.zeros((nmodes, nmodes))
+    for i in range(nmodes):
+        for j in range(nmodes):
+            probe_orthogonality_matrix[i, j] = np.abs(tike.linalg.inner(
+                probe[..., i, :, :],
+                probe[..., j, :, :]
+            ))
+    f = plt.figure()
+    plt.imshow(probe_orthogonality_matrix, interpolation='nearest')
+    plt.colorbar()
+    plt.tight_layout()
+    plt.savefig(f'{output_folder}/probe-orthogonality.png')
+    plt.close(f)
 
 
 def _save_ptycho_result(result, algorithm):
@@ -89,7 +94,6 @@ def _save_ptycho_result(result, algorithm):
             )
             ax2.set_xlim(0, 60)
             ax1.set_ylim(10**(-1), 10**2)
-            ax1.set_xscale('log', base=10)
             fig.suptitle(algorithm)
             fig.tight_layout()
             plt.savefig(os.path.join(fname, 'convergence.png'))

tests/test_opt.py

@@ -1,7 +1,6 @@
 import numpy as np
 import tike.opt
 
-
 class AlgorithmOptionsStub:
 
     def __init__(self, costs, window=5) -> None:
@@ -16,3 +15,11 @@ def test_is_converged():
         AlgorithmOptionsStub((np.zeros(11) / 1234).tolist(), 5))
     assert not tike.opt.is_converged(
         AlgorithmOptionsStub((-np.arange(11) / 1234).tolist(), 5))
+
+
+def test_fit_line():
+    result = np.around(tike.opt.fit_line_least_squares(
+        y=np.asarray([0, np.log(0.9573), np.log(0.8386)]),
+        x=np.asarray([0, 1, 2]),
+    ), 4)
+    np.testing.assert_array_equal((-0.0880, 0.0148), result)