ENH accelarate dicodile

dicodile/update_z/coordinate_descent.py

@@ -10,7 +10,7 @@
 from dicodile.utils.csc import _dense_transpose_convolve, reconstruct
 from dicodile.utils import check_random_state
 from dicodile.utils import debug_flags as flags
-from dicodile.utils.segmentation import Segmentation
+from dicodile.utils.segmentation import WorkerSegmentation
 from dicodile.utils.csc import compute_ztz, compute_ztX
 from dicodile.utils.shape_helpers import get_valid_support
 from dicodile.utils.order_iterator import get_order_iterator
@@ -86,7 +86,7 @@ def coordinate_descent(X_i, D, reg, z0=None, DtD=None, n_seg='auto',
     if n_seg == 'auto':
         n_seg = np.array(valid_support) // (2 * np.array(atom_support) - 1)
         n_seg = tuple(np.maximum(1, n_seg))
-    segments = Segmentation(n_seg, signal_support=valid_support)
+    segments = WorkerSegmentation(n_seg, signal_support=valid_support)
 
     # Pre-compute constants for maintaining the auxillary variable beta and
     # compute the coordinate update values.
@@ -291,7 +291,8 @@ def _init_beta(X_i, D, reg, z_i=None, constants={}, z_positive=False,
     return beta, dz_opt, dE
 
 
-def _select_coordinate(dz_opt, dE, segments, i_seg, strategy, order=None):
+def _select_coordinate(dz_opt, dE, segments, i_seg, strategy, seg_bounds,
+                       order=None):
     """Pick a coordinate to update
 
     Parameters
@@ -322,21 +323,21 @@ def _select_coordinate(dz_opt, dE, segments, i_seg, strategy, order=None):
     if strategy in ['random', 'cyclic-r', 'cyclic']:
         k0, *pt0 = next(order)
     else:
-        if strategy in ['greedy', 'gs-r']:
-            seg_slice = segments.get_seg_slice(i_seg, inner=True)
-            dz_opt_seg = dz_opt[seg_slice]
-            i0 = abs(dz_opt_seg).argmax()
+        seg_slice = segments.get_seg_slice(seg_bounds)
 
+        if strategy in ['greedy', 'gs-r']:
+            d_seg = dz_opt[seg_slice]
         elif strategy == 'gs-q':
-            seg_slice = segments.get_seg_slice(i_seg, inner=True)
-            dE_seg = dE[seg_slice]
-            i0 = abs(dE_seg).argmax()
+            d_seg = dE[seg_slice]
+
+        i0 = abs(d_seg).argmax()
+
         # TODO: broken~~~!!!
-        k0, *pt0 = np.unravel_index(i0, dz_opt_seg.shape)
-        # k0, *pt0 = tuple(fast_unravel(i0, dz_opt_seg.shape))
-        pt0 = segments.get_global_coordinate(i_seg, pt0)
+        k0, *pt0 = np.unravel_index(i0, d_seg.shape)
+        pt0 = segments.get_global_coordinate(pt0, seg_bounds)
 
     dz = dz_opt[(k0, *pt0)]
+
     return k0, pt0, dz
 
 
@@ -435,7 +436,7 @@ def compute_dE(dz_opt, beta, z_hat, reg):
         dz_opt * (z_hat + .5 * dz_opt - beta)
         # l1 term
         + reg * (abs(z_hat) - abs(z_hat + dz_opt))
-        )
+    )
 
 
 def _check_convergence(segments, tol, iteration, dz_opt, n_coordinates,

dicodile/update_z/dicod.py

@@ -13,7 +13,7 @@
 from ..utils.csc import _is_rank1, compute_objective
 from ..utils.debugs import main_check_beta
 from .coordinate_descent import STRATEGIES
-from ..utils.segmentation import Segmentation
+from ..utils.segmentation import WorkerSegmentation
 from .coordinate_descent import coordinate_descent
 from ..utils.mpi import broadcast_array, recv_reduce_sum_array
 from ..utils.shape_helpers import get_valid_support, find_grid_size
@@ -256,14 +256,16 @@ def _send_signal(workers, w_world, atom_support, X, z0=None):
         X_info["workers_topology"] = w_world, n_workers // w_world
 
     # compute a segmentation for the image,
-    workers_segments = Segmentation(n_seg=X_info['workers_topology'],
-                                    signal_support=valid_support,
-                                    overlap=overlap)
+    workers_segments = WorkerSegmentation(n_seg=X_info['workers_topology'],
+                                          signal_support=valid_support,
+                                          overlap=overlap)
 
     # Make sure that each worker has at least a segment of twice the size of
     # the dictionary. If this is not the case, the algorithm is not valid as it
     # is possible to have interference with workers that are not neighbors.
-    worker_support = workers_segments.get_seg_support(0, inner=True)
+    worker_inner_bounds = workers_segments.get_seg_bounds(0, inner=True)
+    worker_support = workers_segments.get_seg_support(worker_inner_bounds)
+
     msg = ("The size of the support in each worker is smaller than twice the "
            "size of the atom support. The algorithm is does not converge in "
            "this condition. Reduce the number of cores.\n"
@@ -279,10 +281,10 @@ def _send_signal(workers, w_world, atom_support, X, z0=None):
     X = np.array(X, dtype='d')
 
     for i_seg in range(n_workers):
+        seg_bounds = workers_segments.get_seg_bounds(i_seg)
         if z0 is not None:
-            worker_slice = workers_segments.get_seg_slice(i_seg)
+            worker_slice = workers_segments.get_seg_slice(seg_bounds)
             _send_array(workers.comm, i_seg, z0[worker_slice])
-        seg_bounds = workers_segments.get_seg_bounds(i_seg)
         X_worker_slice = (Ellipsis,) + tuple([
             slice(start, end + size_atom_ax - 1)
             for (start, end), size_atom_ax in zip(seg_bounds, atom_support)
@@ -361,13 +363,12 @@ def recv_z_hat(comm, n_atoms, workers_segments):
     inner = not flags.GET_OVERLAP_Z_HAT
     z_hat = np.empty((n_atoms, *valid_support), dtype='d')
     for i_seg in range(workers_segments.effective_n_seg):
-        worker_support = workers_segments.get_seg_support(
-            i_seg, inner=inner)
+        bounds = workers_segments.get_seg_bounds(i_seg, inner=inner)
+        worker_support = workers_segments.get_seg_support(bounds)
         z_worker = np.zeros((n_atoms,) + worker_support, 'd')
         comm.Recv([z_worker.ravel(), MPI.DOUBLE], source=i_seg,
                   tag=constants.TAG_ROOT + i_seg)
-        worker_slice = workers_segments.get_seg_slice(
-            i_seg, inner=inner)
+        worker_slice = workers_segments.get_seg_slice(bounds)
         z_hat[worker_slice] = z_worker
 
     return z_hat

dicodile/utils/debugs.py

@@ -26,9 +26,11 @@ def main_check_beta(comm, workers_segments):
         sum_beta = np.empty(1, 'd')
         value = []
         for i_worker in range(workers_segments.effective_n_seg):
-
-            pt = workers_segments.get_local_coordinate(i_worker, pt_global)
-            if workers_segments.is_contained_coordinate(i_worker, pt):
+            worker_bounds = workers_segments.get_seg_bounds(i_worker)
+            pt = workers_segments.get_local_coordinate(pt_global,
+                                                       worker_bounds)
+            if workers_segments.is_contained_coordinate(pt,
+                                                        worker_bounds):
                 comm.Recv([sum_beta, MPI.DOUBLE], source=i_worker,
                           tag=constants.TAG_ROOT + i_probe)
                 value.append(sum_beta[0])
@@ -44,8 +46,10 @@ def worker_check_beta(rank, workers_segments, beta, D_shape):
 
     global_test_points = get_global_test_points(workers_segments)
     for i_probe, pt_global in enumerate(global_test_points):
-        pt = workers_segments.get_local_coordinate(rank, pt_global)
-        if workers_segments.is_contained_coordinate(rank, pt):
+        worker_bounds = workers_segments.get_seg_bounds(rank)
+        pt = workers_segments.get_local_coordinate(pt_global,
+                                                   worker_bounds)
+        if workers_segments.is_contained_coordinate(pt, worker_bounds):
             beta_slice = (Ellipsis,) + pt
             sum_beta = np.array(beta[beta_slice].sum(), dtype='d')