Feat distrarray memoryimprov

pylops_mpi/optimization/cls_basic.py

@@ -93,10 +93,12 @@ def setup(
         self.tol = tol
 
         x = x0.copy()
-        self.r = self.y - self.Op.matvec(x)
+        # residual computed in-place into the output of the matvec
+        self.r = self.Op.matvec(x)
+        self.y.subtract(self.r, out=self.r)
         self.rank = x.rank
         self.c = self.r.copy()
-        self.kold = float(np.abs(self.r.dot(self.r.conj())))
+        self.kold = float(np.abs(self.r.dot(self.r, vdot=True)))
 
         # create variables to track the residual norm and iterations
         self.cost: List = []
@@ -132,11 +134,15 @@ def step(
         Opc = self.Op.matvec(self.c)
         cOpc = np.abs(self.c.dot(Opc.conj()))
         a = float(self.kold / cOpc)
-        x += a * self.c
-        self.r -= a * Opc
-        k = float(np.abs(self.r.dot(self.r.conj())))
+        # all updates are done in-place, reusing Opc as scratch space
+        Opc *= a
+        self.r -= Opc
+        self.c.multiply(a, out=Opc)
+        x += Opc
+        k = float(np.abs(self.r.dot(self.r, vdot=True)))
         b = float(k / self.kold)
-        self.c = self.r + b * self.c
+        self.c *= b
+        self.c += self.r
         self.kold = k
         self.iiter += 1
         self.cost.append(float(np.sqrt(self.kold)))
@@ -352,21 +358,29 @@ def setup(
         self.niter = niter
 
         x = x0.copy()
-        self.s = self.y - self.Op.matvec(x)
-        damped_x = x * damp
-        r = self.Op.rmatvec(self.s) - damped_x
+        # residual computed in-place into the output of the matvec
+        self.s = self.Op.matvec(x)
+        self.y.subtract(self.s, out=self.s)
+        r = self.Op.rmatvec(self.s)
+        if damp != 0.0:
+            r -= x * damp
         self.rank = x.rank
         self.c = r.copy()
         self.q = self.Op.matvec(self.c)
-        self.kold = float(np.abs(r.dot(r.conj())))
+        # r is kept as scratch space for the steps
+        self.r = r
+        self.kold = float(np.abs(r.dot(r, vdot=True)))
 
         # create variables to track the residual norm and iterations
         self.cost = []
         self.cost1 = []
         self.cost.append(float(self.s.norm()))
-        self.cost1.append(
-            np.sqrt(float(self.cost[0] ** 2 + damp * np.abs(x.dot(x.conj()))))
-        )
+        if damp == 0.0:
+            self.cost1.append(self.cost[0])
+        else:
+            self.cost1.append(
+                np.sqrt(float(self.cost[0] ** 2 + damp * np.abs(x.dot(x, vdot=True))))
+            )
         self.iiter = 0
 
         # print setup
@@ -397,28 +411,41 @@ def step(
 
         """
 
-        a = float(
-            np.abs(
-                self.kold
-                / (self.q.dot(self.q.conj()) + self.damp * self.c.dot(self.c.conj()))
-            )
-        )
-        x += a * self.c
-        self.s -= a * self.q
-        damped_x = self.damp * x
-        r = self.Op.rmatvec(self.s) - damped_x
-        k = float(np.abs(r.dot(r.conj())))
+        qq = self.q.dot(self.q, vdot=True)
+        if self.damp != 0.0:
+            qq += self.damp * self.c.dot(self.c, vdot=True)
+        a = float(np.abs(self.kold / qq))
+        # all updates are done in-place: the previous r is dead after the last
+        # update of c, so it is used as scratch space; q is replaced by the
+        # next matvec
+        self.c.multiply(a, out=self.r)
+        x += self.r
+        self.q *= a
+        self.s -= self.q
+        r = self.Op.rmatvec(self.s)
+        if self.damp != 0.0:
+            x.multiply(self.damp, out=self.r)
+            r -= self.r
+        k = float(np.abs(r.dot(r, vdot=True)))
         b = float(k / self.kold)
-        self.c = r + b * self.c
+        self.c *= b
+        self.c += r
+        self.r = r
         self.q = self.Op.matvec(self.c)
         self.kold = k
         self.iiter += 1
         self.cost.append(float(self.s.norm()))
-        self.cost1.append(
-            np.sqrt(
-                float(self.cost[self.iiter] ** 2 + self.damp * np.abs(x.dot(x.conj())))
+        if self.damp == 0.0:
+            self.cost1.append(self.cost[self.iiter])
+        else:
+            self.cost1.append(
+                np.sqrt(
+                    float(
+                        self.cost[self.iiter] ** 2
+                        + self.damp * np.abs(x.dot(x, vdot=True))
+                    )
+                )
             )
-        )
         if show and self.rank == 0:
             self._print_step(x)
         return x

tests/test_distributedarray.py

@@ -15,6 +15,7 @@
 
     backend = "numpy"
 from mpi4py import MPI
+import numpy as npp
 import pytest
 
 from pylops_mpi import DistributedArray, Partition
@@ -358,4 +359,131 @@ def test_distributed_maskednorm(par):
         arr.norm(ord=np.inf, axis=par['norm_axis']),
         np.linalg.norm(par['x'], ord=np.inf, axis=par['norm_axis']),
         rtol=1e-14
-    )
+    )
+
+
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("par1, par2", [(par4, par5), (par4j, par5j)])
+def test_distributed_out(par1, par2):
+    """Arithmetic with out= and in-place operators"""
+    x1 = DistributedArray.to_dist(x=par1['x'])
+    x2 = DistributedArray.to_dist(x=par2['x'])
+    for name, op in [('add', np.add), ('subtract', np.subtract),
+                     ('multiply', np.multiply)]:
+        expected = op(par1['x'], par2['x'])
+        # array operand
+        out = x1.empty_like()
+        res = getattr(x1, name)(x2, out=out)
+        assert res is out
+        assert_allclose(out.asarray(), expected, rtol=1e-14)
+        # scalar operand
+        out = x1.empty_like()
+        getattr(x1, name)(2.5, out=out)
+        assert_allclose(out.asarray(), op(par1['x'], 2.5), rtol=1e-14)
+        # operands untouched
+        assert_allclose(x1.asarray(), par1['x'], rtol=1e-14)
+        assert_allclose(x2.asarray(), par2['x'], rtol=1e-14)
+    # out aliasing the second operand
+    out = x2.copy()
+    x1.subtract(out, out=out)
+    assert_allclose(out.asarray(), par1['x'] - par2['x'], rtol=1e-14)
+    # negation
+    assert_allclose((-x1).asarray(), -par1['x'], rtol=1e-14)
+    # in-place operators keep the storage
+    for opname, op in [('__iadd__', np.add), ('__isub__', np.subtract),
+                       ('__imul__', np.multiply)]:
+        for other, other_g in [(x2, par2['x']), (2.5, 2.5)]:
+            y = x1.copy()
+            local = y.local_array
+            res = getattr(y, opname)(other)
+            assert res is y
+            assert y.local_array is local
+            assert_allclose(y.asarray(), op(par1['x'], other_g), rtol=1e-14)
+
+
+@pytest.mark.skipif(
+    int(os.environ.get("TEST_CUPY_PYLOPS", 0)) == 1, reason="Not CuPy enabled"
+)
+@pytest.mark.mpi(min_size=2)
+def test_distributed_out_cast():
+    """out= requiring an unsafe cast falls back to assignment (as before)"""
+    x = DistributedArray(global_shape=(10 * MPI.COMM_WORLD.Get_size(),),
+                         dtype=np.float64, engine=backend)
+    x[:] = 1.0
+    y = DistributedArray(global_shape=x.global_shape,
+                         dtype=np.complex128, engine=backend)
+    y[:] = 1.0 + 2.0j
+    with pytest.warns(Warning):
+        x.add(y, out=x)
+    assert_allclose(x.local_array, 2.0)
+
+
+@pytest.mark.mpi(min_size=2)
+def test_distributed_like_no_collectives():
+    """Result arrays reuse communicator and cached local shapes"""
+    size = MPI.COMM_WORLD.Get_size()
+    if size % 2 != 0:
+        pytest.skip("Number of processes is not divisible by 2")
+    mask = np.repeat(np.arange(2), size // 2)
+    x = DistributedArray(global_shape=(4 * size, 3), mask=mask,
+                         engine=backend)
+    x[:] = 1.0
+    _ = x.local_shapes
+    for arr in [x + x, x - x, x * x, x * 2.0, -x, x.copy(), x.conj(),
+                x.zeros_like(), x.empty_like()]:
+        assert arr.sub_comm is x.sub_comm
+        assert arr.mask is x.mask
+        assert arr.local_shapes is x.local_shapes
+        assert arr.local_shape == x.local_shape
+
+
+@pytest.mark.mpi(min_size=2)
+def test_distributed_broadcast_math():
+    """Arithmetic on BROADCAST arrays is equal on every rank"""
+    xg = np.asarray(npp.random.default_rng(42).standard_normal((20, 7)))
+    for part in [Partition.BROADCAST, Partition.UNSAFE_BROADCAST]:
+        x = DistributedArray.to_dist(x=xg, partition=part)
+        y = DistributedArray.to_dist(x=2 * xg, partition=part)
+        for res, expected in [(x + y, 3 * xg), (y - x, xg), (x * y, 2 * xg * xg),
+                              (-x, -xg)]:
+            assert res.partition is part
+            assert_allclose(res.local_array, expected, rtol=1e-14)
+        x += y
+        assert_allclose(x.local_array, 3 * xg, rtol=1e-14)
+        assert_allclose(x.dot(y), np.sum(3 * xg * 2 * xg), rtol=1e-12)
+
+
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("par", [(par4), (par4j)])
+def test_distributed_norm_dot_views(par):
+    """2-norm and dot do not modify (or depend on copies of) local arrays"""
+    x = DistributedArray.to_dist(x=par['x'])
+    before = x.local_array.copy()
+    assert_allclose(x.norm(), np.linalg.norm(par['x'].ravel()), rtol=1e-12)
+    assert_allclose(x.dot(x, vdot=True), np.vdot(par['x'], par['x']), rtol=1e-12)
+    assert_allclose(x.local_array, before)
+
+
+@pytest.mark.mpi(min_size=2)
+def test_distributed_norm_overflow():
+    """2-norm of overflowing complex data is inf (not nan)"""
+    x = DistributedArray(global_shape=(4 * MPI.COMM_WORLD.Get_size(),),
+                         dtype=np.complex128, engine=backend)
+    x[:] = (1.0 + 1.0j) * 1e160
+    assert np.isinf(x.norm())
+
+
+@pytest.mark.mpi(min_size=2)
+def test_distributed_like_checks():
+    """_like uses the dtype of local_array and checks its shape"""
+    x = DistributedArray(global_shape=(4 * MPI.COMM_WORLD.Get_size(), 3),
+                         dtype=np.float64, engine=backend)
+    # dtype of local_array wins over dtype
+    y = x._like(local_array=np.zeros(x.local_shape, dtype=np.complex64),
+                dtype=np.float64)
+    assert y.dtype == np.complex64
+    # dtype is used when allocating
+    assert x._like(dtype=np.float32).local_array.dtype == np.float32
+    assert x._like().local_array.dtype == np.float64
+    with pytest.raises(ValueError, match="expected"):
+        x._like(local_array=np.zeros((x.local_shape[0] + 1, 3)))

tests/test_solver.py

@@ -425,3 +425,89 @@ def test_cgls_stacked(par):
         y1 = StackedBDiag * x_global
         xinv1 = pylops.cgls(StackedBDiag, y1, x0=x0, niter=par["nx"], tol=1e-5, show=True)[0]
         assert_allclose(xinv_array, xinv1, rtol=1e-12)
+
+
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("dtype", ["float64", "complex128"])
+@pytest.mark.parametrize("damp", [0.0, 0.5])
+@pytest.mark.parametrize("stacked", [False, True])
+def test_cgls_damp_inplace(stacked, damp, dtype):
+    """CGLS with and without damping (in-place updates): compare with serial
+    pylops and check that data and initial guess are not modified"""
+    rng = np.random.default_rng(10)
+    nx, ny, niter = 8, 12, 6
+    imag = 1j if dtype == "complex128" else 0.0
+
+    def rand(*shape):
+        return rng.standard_normal(shape) + imag * rng.standard_normal(shape)
+
+    Amats = [[rand(ny, nx) for _ in range(size)] for _ in range(2 if stacked else 1)]
+    x_global = [rand(size * nx) for _ in Amats]
+    x0_global = [rand(size * nx) for _ in Amats]
+
+    def dist(vec):
+        d = DistributedArray(global_shape=size * nx, dtype=dtype, engine=backend)
+        d[:] = np.asarray(vec[rank * nx:(rank + 1) * nx])
+        return d
+
+    ops = [MPIBlockDiag(ops=[MatrixMult(np.asarray(A[rank]), dtype=dtype)], dtype=dtype)
+           for A in Amats]
+    if stacked:
+        Op = MPIStackedBlockDiag(ops=ops)
+        x = StackedDistributedArray([dist(v) for v in x_global])
+        x0 = StackedDistributedArray([dist(v) for v in x0_global])
+    else:
+        Op = ops[0]
+        x, x0 = dist(x_global[0]), dist(x0_global[0])
+    y = Op @ x
+    y_before, x0_before = y.asarray().copy(), x0.asarray().copy()
+
+    xinv, _, iiter, _, _, cost = cgls(Op, y, x0=x0, niter=niter, damp=damp, tol=0)
+    assert iiter == niter
+    assert_allclose(y.asarray(), y_before, rtol=1e-14)
+    assert_allclose(x0.asarray(), x0_before, rtol=1e-14)
+
+    # serial reference
+    serial = BlockDiag(
+        ops=[BlockDiag(ops=[MatrixMult(np.asarray(A[r]), dtype=dtype) for r in range(size)])
+             for A in Amats], forceflat=True) if stacked else \
+        BlockDiag(ops=[MatrixMult(np.asarray(Amats[0][r]), dtype=dtype) for r in range(size)])
+    xg = np.concatenate([np.asarray(v) for v in x_global])
+    x0g = np.concatenate([np.asarray(v) for v in x0_global])
+    xinv1, _, _, _, _, cost1 = pylops.cgls(serial, serial @ xg, x0=x0g, niter=niter,
+                                           damp=damp, tol=0)
+    assert_allclose(xinv.asarray(), xinv1, rtol=1e-8)
+    assert_allclose(cost, cost1, rtol=1e-8)
+
+
+@pytest.mark.mpi(min_size=2)
+@pytest.mark.parametrize("stacked", [False, True])
+def test_cg_cgls_allreduces(stacked):
+    """CGLS (no damping) needs 3 allreduces per iteration, CG 2"""
+    import pylops_mpi.Distributed as D
+
+    nx, niter = 6, 5
+    mats = [np.random.normal(0, 1, (nx, nx)) + 5 * np.eye(nx) for _ in range(size)]
+    op1 = MPIBlockDiag(ops=[MatrixMult(np.asarray(mats[rank]))])
+    Op = MPIStackedBlockDiag(ops=[op1, op1]) if stacked else op1
+
+    def dist():
+        d = DistributedArray(global_shape=size * nx, engine=backend)
+        d[:] = np.random.normal(0, 1, nx)
+        return d
+    y = StackedDistributedArray([dist(), dist()]) if stacked else dist()
+
+    orig, n = D.mpi_allreduce, [0]
+
+    def counting(*args, **kwargs):
+        n[0] += 1
+        return orig(*args, **kwargs)
+    D.mpi_allreduce = counting
+    try:
+        for solver, per_iter, nsetup in [(cgls, 3, 2), (cg, 2, 1)]:
+            n[0] = 0
+            x0 = y.zeros_like()
+            solver(Op, y, x0=x0, niter=niter, tol=0)
+            assert n[0] == per_iter * niter + nsetup
+    finally:
+        D.mpi_allreduce = orig