unroll change

Author: astroC86

pylops_mpi/StackedLinearOperator.py

@@ -70,9 +70,7 @@ def matvec(self, x: Union[DistributedArray, StackedDistributedArray]) -> Union[D
 
         M, N = self.shape
         if isinstance(x, StackedDistributedArray):
-            stacked_shape = (
-                np.sum(np.fromiter((a.global_shape[0] for a in x.distarrays), dtype=np.int64)),
-            )
+            stacked_shape = (np.sum([a.global_shape for a in x.distarrays]), )
             if stacked_shape != (N, ):
                 raise ValueError("dimension mismatch")
         if isinstance(x, DistributedArray) and x.global_shape != (N,):
@@ -105,9 +103,7 @@ def rmatvec(self, x: Union[DistributedArray, StackedDistributedArray]) -> Union[
 
         M, N = self.shape
         if isinstance(x, StackedDistributedArray):
-            stacked_shape = (
-                np.sum(np.fromiter((a.global_shape[0] for a in x.distarrays), dtype=np.int64)),
-            )
+            stacked_shape = (np.sum([a.global_shape for a in x.distarrays]), )
             if stacked_shape != (M, ):
                 raise ValueError("dimension mismatch")
         if isinstance(x, DistributedArray) and x.global_shape != (M,):

pylops_mpi/basicoperators/BlockDiag.py

@@ -170,10 +170,8 @@ def __init__(self, ops: Sequence[MPILinearOperator], base_comm: MPI.Comm = MPI.C
                  dtype: Optional[DTypeLike] = None):
         self.ops = ops
         dtype = _get_dtype(self.ops) if dtype is None else np.dtype(dtype)
-        shape = (
-            int(np.sum(np.fromiter((op.shape[0] for op in ops), dtype=np.int64))),
-            int(np.sum(np.fromiter((op.shape[1] for op in ops), dtype=np.int64))),
-        )
+        shape = (int(np.sum(op.shape[0] for op in ops)),
+                 int(np.sum(op.shape[1] for op in ops)))
         super().__init__(shape=shape, dtype=dtype, base_comm=base_comm)
 
     def _matvec(self, x: StackedDistributedArray) -> StackedDistributedArray:

pylops_mpi/basicoperators/VStack.py

@@ -183,10 +183,7 @@ def __init__(self, ops: Sequence[MPILinearOperator],
         self.ops = ops
         if len(set(op.shape[1] for op in ops)) > 1:
             raise ValueError("Operators have different number of columns")
-        shape = (
-            int(np.sum(np.fromiter((op.shape[0] for op in ops), dtype=np.int64))),
-            ops[0].shape[1],
-        )
+        shape = (int(np.sum(op.shape[0] for op in ops)), ops[0].shape[1])
         dtype = _get_dtype(self.ops) if dtype is None else np.dtype(dtype)
         super().__init__(shape=shape, dtype=dtype, base_comm=base_comm)