Merge pull request #192 from rohanbabbar04/ista

ISTA/FISTA implementation in cls_sparsity

Author: mrava87

docs/source/api/index.rst

@@ -106,6 +106,24 @@ Basic
     cg
     cgls
 
+Sparsity
+~~~~~
+
+.. currentmodule:: pylops_mpi.optimization.cls_sparsity
+
+.. autosummary::
+   :toctree: generated/
+
+    ISTA
+    FISTA
+
+.. currentmodule:: pylops_mpi.optimization.sparsity
+
+.. autosummary::
+   :toctree: generated/
+
+    ista
+    fista
 
 Utils
 -----

environment-dev.yml

@@ -8,7 +8,7 @@ dependencies:
   - numpy>=1.15.0
   - scipy>=1.8.0
   - mpi4py
-  - pylops
+  - pylops>=2.0.0
   - matplotlib
   - ipython
   - pytest

environment.yml

@@ -6,6 +6,6 @@ dependencies:
   - python>=3.8.0
   - numpy>=1.15.0
   - scipy>=1.8.0
-  - pylops>=2.0
+  - pylops>=2.0.0
   - matplotlib
   - mpi4py

pylops_mpi/DistributedArray.py

@@ -205,12 +205,19 @@ def __setitem__(self, index, value):
             the specified index positions.
         """
         if self.partition is Partition.BROADCAST:
+            ncp = get_module(self.engine)
             view = self.local_array[index]
+            buf = ncp.empty(view.shape, dtype=self.dtype)
             if self.rank == 0:
-                view[...] = value
-            view = self._bcast(self.base_comm, self.base_comm_nccl,
-                               view, root=0, engine=self.engine)
-            self.local_array[index] = view
+                buf[...] = value
+            buf = self._bcast(
+                self.base_comm,
+                self.base_comm_nccl,
+                buf,
+                root=0,
+                engine=self.engine
+            )
+            self.local_array[index] = buf
         else:
             self.local_array[index] = value
 
@@ -644,8 +651,25 @@ def multiply(self, dist_array):
             ProductArray[:] = self.local_array * dist_array
         return ProductArray
 
-    def dot(self, dist_array):
-        """Distributed Dot Product
+    def dot(self, dist_array, vdot: bool = False):
+        """
+        Compute the distributed dot product between this array and another
+        distributed array.
+
+        Parameters
+        ----------
+        dist_array : :obj:`pylops_mpi.DistributedArray`
+            The distributed array with which to compute the dot product.
+            It must have a compatible shape and partitioning.
+        vdot : bool, optional, Defaults to `False`
+            If True, compute the complex conjugate dot product (like numpy.vdot),
+            where the first argument is conjugated. If False, compute the standard
+            dot product.
+
+        Returns
+        -------
+        result : float
+            The result of the dot product across all ranks. This is reduced across all processes.
         """
         self._check_partition_shape(dist_array)
         self._check_mask(dist_array)
@@ -656,8 +680,9 @@ def dot(self, dist_array):
         y = DistributedArray.to_dist(x=dist_array.local_array, base_comm=self.base_comm, base_comm_nccl=self.base_comm_nccl) \
             if self.partition in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST] else dist_array
         # Flatten the local arrays and calculate dot product
+        dot_func = ncp.vdot if vdot else ncp.dot
         return self._allreduce_subcomm(self.sub_comm, self.base_comm_nccl,
-                                       ncp.dot(x.local_array.flatten(), y.local_array.flatten()),
+                                       dot_func(x.local_array.flatten(), y.local_array.flatten()),
                                        engine=self.engine)
 
     def _compute_vector_norm(self, local_array: NDArray,
@@ -839,6 +864,15 @@ def ravel(self, order: Optional[str] = "C"):
         arr[:] = x
         return arr
 
+    def empty_like(self):
+        """Creates an empty like DistributedArray with uninitialized values
+        """
+        dist = DistributedArray(global_shape=self.global_shape, base_comm=self.base_comm,
+                                base_comm_nccl=self.base_comm_nccl, partition=self.partition,
+                                axis=self.axis, local_shapes=self.local_shapes, mask=self.mask,
+                                engine=self.engine, dtype=self.dtype)
+        return dist
+
     def add_ghost_cells(self, cells_front: Optional[int] = None,
                         cells_back: Optional[int] = None):
         """Add ghost cells to the DistributedArray along the axis
@@ -1038,13 +1072,30 @@ def multiply(self, stacked_array):
                 ProductArray[iarr][:] = (self[iarr] * stacked_array)[:]
         return ProductArray
 
-    def dot(self, stacked_array):
-        """Dot Product of Stacked Distributed Arrays
+    def dot(self, stacked_array, vdot: bool = False):
+        """
+        Compute the distributed dot product between this array and another
+        distributed array.
+
+        Parameters
+        ----------
+        stacked_array : :obj:`pylops_mpi.StackedDistributedArray`
+            The distributed array with which to compute the dot product.
+            It must have a compatible shape and partitioning.
+        vdot : bool, optional, Defaults to `False`
+            If True, compute the complex conjugate dot product (like numpy.vdot),
+            where the first argument is conjugated. If False, compute the standard
+            dot product.
+
+        Returns
+        -------
+        result : float
+            The result of the dot product across all ranks. This is reduced across all processes.
         """
         self._check_stacked_size(stacked_array)
         dotprod = 0.
         for iarr in range(self.narrays):
-            dotprod += self[iarr].dot(stacked_array[iarr])
+            dotprod += self[iarr].dot(stacked_array[iarr], vdot=vdot)
         return dotprod
 
     def norm(self, ord: Optional[int] = None):
@@ -1085,6 +1136,19 @@ def copy(self):
         arr = StackedDistributedArray([distarray.copy() for distarray in self.distarrays])
         return arr
 
+    def empty_like(self):
+        """Creates an empty like StackedDistributedArray with uninitialized values
+        """
+        dists = []
+        for iarr in range(self.narrays):
+            distarray = self.distarrays[iarr]
+            dist = DistributedArray(global_shape=distarray.global_shape, base_comm=distarray.base_comm,
+                                    base_comm_nccl=distarray.base_comm_nccl, partition=distarray.partition,
+                                    axis=distarray.axis, local_shapes=distarray.local_shapes, mask=distarray.mask,
+                                    engine=distarray.engine, dtype=distarray.dtype)
+            dists.append(dist)
+        return StackedDistributedArray(distarrays=dists)
+
     def __repr__(self):
         repr_dist = "\n".join([distarray.__repr__() for distarray in self.distarrays])
         return f"<StackedDistributedArray with {self.narrays} distributed arrays: \n" + repr_dist

pylops_mpi/__init__.py

@@ -11,6 +11,7 @@
 )
 from .plotting.plotting import *
 from .optimization.basic import *
+from .optimization.sparsity import *
 
 try:
     from .version import version as __version__

pylops_mpi/optimization/__init__.py

@@ -8,5 +8,6 @@
 
     cg                              Conjugate gradient
     cgls                            Conjugate gradient least-squares.
-
+    ista                            Iterative Shrinkage-Thresholding Algorithm
+    fista                           Fast Iterative Shrinkage-Thresholding Algorithm
 """