Skip to content

Commit

Permalink
minor: fix linting
Browse files Browse the repository at this point in the history
  • Loading branch information
mrava87 committed Jun 16, 2024
1 parent 03fa3f2 commit 1a9549a
Show file tree
Hide file tree
Showing 7 changed files with 37 additions and 49 deletions.
26 changes: 12 additions & 14 deletions examples/plot_mdc.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
This example shows how to use the :py:class:`pylops_mpi.waveeqprocessing.MPIMDC` operator
to convolve a 3D kernel with an input seismic data in a distributed fashion (where
parallelism is harnessed over the frequency axis when performing repeated matrix-vector
or matrix-matrix multiplications).
or matrix-matrix multiplications).
"""
from matplotlib import pyplot as plt
Expand Down Expand Up @@ -84,12 +84,11 @@
m, mwav = m.T, mwav.T
Gwav_fft = Gwav_fft.transpose(2, 0, 1)


###############################################################################
# Now that we have created the kernel of our MDC operator in ``Gwav_fft``, we
# are ready to define a strategy on how to split it along the first
# (i.e., frequency) axis over different ranks. In practical applications, one
# would of course pre-compute the kernel and just load the relevant part in
# Now that we have created the kernel of our MDC operator in ``Gwav_fft``, we
# are ready to define a strategy on how to split it along the first
# (i.e., frequency) axis over different ranks. In practical applications, one
# would of course pre-compute the kernel and just load the relevant part in
# each rank from file.

# Choose how to split sources to ranks
Expand All @@ -104,21 +103,20 @@
G = Gwav_fft[ifin_rank:ifend_rank].astype(cdtype)
print(f'Rank: {rank}, G: {G.shape}')


###############################################################################
# We can finally create the MDC operator using
# We can finally create the MDC operator using
# :py:class:`pylops_mpi.waveeqprocessing.MPIMDC` so that the most
# demanding computations can be run in parallel.

# Define operator
Fop = pylops_mpi.waveeqprocessing.MPIMDC(
G, nt=2 * par["nt"] - 1, nv=1, nfreq=nf,
G, nt=2 * par["nt"] - 1, nv=1, nfreq=nf,
dt=0.004, dr=1.0, twosided=True)

# Apply forward
md = pylops_mpi.DistributedArray(global_shape=(2 * par["nt"] - 1) * par["nx"] * 1,
partition=pylops_mpi.Partition.BROADCAST,
dtype=dtype)
md = pylops_mpi.DistributedArray(global_shape=(2 * par["nt"] - 1) * par["nx"] * 1,
partition=pylops_mpi.Partition.BROADCAST,
dtype=dtype)
md[:] = m.astype(dtype).ravel()

dd = Fop @ md
Expand All @@ -132,7 +130,7 @@

###############################################################################
# Finally let's display input model, data and adjoint model

if rank == 0:
fig, axs = plt.subplots(1, 3, figsize=(9, 6))
axs[0].imshow(
Expand Down Expand Up @@ -171,4 +169,4 @@
axs[2].set_title(r"$m_{adj}$", fontsize=15)
axs[2].set_xlabel("s")
axs[2].set_ylabel("t")
fig.tight_layout()
fig.tight_layout()
2 changes: 1 addition & 1 deletion pylops_mpi/basicoperators/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
================================
The subpackage basicoperators extends some of the basic operations
provided by pylops.basicoperators providing forward and adjoint
provided by pylops.basicoperators providing forward and adjoint
functionalities using MPI.
A list of operators present in pylops_mpi.basicoperators:
Expand Down
15 changes: 6 additions & 9 deletions pylops_mpi/signalprocessing/Fredholm1.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@
Partition
)

from pylops_mpi.utils.decorators import reshaped


class MPIFredholm1(MPILinearOperator):
r"""Fredholm integral of first kind.
Expand Down Expand Up @@ -41,17 +39,17 @@ class MPIFredholm1(MPILinearOperator):
MPI Base Communicator. Defaults to ``mpi4py.MPI.COMM_WORLD``.
dtype : :obj:`str`, optional
Type of elements in input array.
Attributes
----------
shape : :obj:`tuple`
Operator shape
Raises
------
NotImplementedError
If the size of the first dimension of ``G`` is equal to 1 in any of the ranks
Notes
-----
A multi-dimensional Fredholm integral of first kind can be expressed as
Expand All @@ -68,7 +66,7 @@ class MPIFredholm1(MPILinearOperator):
m(k, y, z) = \int{G^*(k, y, x) d(k, x, z) \,\mathrm{d}x}
\quad \forall k=1,\ldots,n_{\text{slice}}
This integral is implemented in a distributed fashion, where ``G``
This integral is implemented in a distributed fashion, where ``G``
is split across ranks along its first dimension. The inputs
of both the forward and adjoint are distributed arrays with broadcast partion:
each rank takes a portion of such arrays, computes a partial integral, and
Expand Down Expand Up @@ -99,7 +97,7 @@ def __init__(
self.rank = base_comm.Get_rank()
self.dims = (nslstot, self.ny, self.nz)
self.dimsd = (nslstot, self.nx, self.nz)
shape = (np.prod(self.dimsd),
shape = (np.prod(self.dimsd),
np.prod(self.dims))
super().__init__(shape=shape, dtype=np.dtype(dtype), base_comm=base_comm)

Expand Down Expand Up @@ -157,8 +155,7 @@ def _rmatvec(self, x: NDArray) -> NDArray:
else:
for isl in range(self.nsl):
y1[isl] = ncp.dot(x[isl].T.conj(), self.G[isl]).T.conj()

# gather results
y[:] = np.vstack(self.base_comm.allgather(y1)).ravel()
return y

2 changes: 1 addition & 1 deletion pylops_mpi/signalprocessing/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
=====================================
The subpackage signalprocessing extends some of the signal processing
functionalities in pylops.signalprocessing providing forward and adjoint
functionalities in pylops.signalprocessing providing forward and adjoint
functionalities using MPI.
A list of operators present in pylops_mpi.signalprocessing :
Expand Down
30 changes: 12 additions & 18 deletions pylops_mpi/waveeqprocessing/MDC.py
Original file line number Diff line number Diff line change
@@ -1,26 +1,20 @@
import logging
import warnings
import numpy as np

from mpi4py import MPI
from pylops import Identity, Transpose
from pylops import Identity
from pylops.signalprocessing import FFT, Fredholm1
from pylops.waveeqprocessing.mdd import _MDC

from pylops_mpi import (
DistributedArray,
MPILinearOperator,
Partition
)
from pylops_mpi import MPILinearOperator
from pylops_mpi.signalprocessing.Fredholm1 import MPIFredholm1


def _MDC(G, nt, nv, nfmax, dt=1., dr=1., twosided=True,
saveGt=True, conj=False, prescaled=False,
base_comm=MPI.COMM_WORLD,
_Identity=Identity, _Transpose=Transpose, _FFT=FFT,
_Fredholm1=Fredholm1, args_Identity={}, args_Transpose={},
args_FFT={}, args_Identity1={}, args_Transpose1={},
_Identity=Identity, _FFT=FFT,
_Fredholm1=Fredholm1, args_Identity={},
args_FFT={}, args_Identity1={},
args_FFT1={}, args_Fredholm1={}):
r"""Multi-dimensional convolution.
Expand All @@ -40,11 +34,11 @@ def _MDC(G, nt, nv, nfmax, dt=1., dr=1., twosided=True,

# create Fredholm operator
if prescaled:
Frop = _Fredholm1(G, nv, saveGt=saveGt,
Frop = _Fredholm1(G, nv, saveGt=saveGt,
base_comm=base_comm,
dtype=dtype, **args_Fredholm1)
else:
Frop = _Fredholm1(dr * dt * np.sqrt(nt) * G, nv,
Frop = _Fredholm1(dr * dt * np.sqrt(nt) * G, nv,
saveGt=saveGt, base_comm=base_comm,
dtype=dtype, **args_Fredholm1)
if conj:
Expand Down Expand Up @@ -73,7 +67,7 @@ def _MDC(G, nt, nv, nfmax, dt=1., dr=1., twosided=True,

# create MDC operator
MDCop = F1opH * I1opH * Frop * Iop * Fop

# force dtype to be real (as FFT operators assume real inputs and outputs)
MDCop.dtype = rdtype

Expand All @@ -82,7 +76,7 @@ def _MDC(G, nt, nv, nfmax, dt=1., dr=1., twosided=True,

def MPIMDC(G, nt, nv, nfreq, dt=1., dr=1., twosided=True,
fftengine='numpy',
saveGt=True, conj=False,
saveGt=True, conj=False,
usematmul=False, prescaled=False,
base_comm: MPI.Comm = MPI.COMM_WORLD):
r"""Multi-dimensional convolution.
Expand All @@ -92,11 +86,11 @@ def MPIMDC(G, nt, nv, nfreq, dt=1., dr=1., twosided=True,
:math:`[n_t \times n_r (\times n_{vs})]` and
:math:`[n_t \times n_s (\times n_{vs})]` (or :math:`2*n_t-1` for
``twosided=True``), respectively.
Parameters
----------
G : :obj:`numpy.ndarray`
Multi-dimensional convolution kernel in frequency domain of size
Multi-dimensional convolution kernel in frequency domain of size
:math:`[n_{fmax} \times n_s \times n_r]`
nt : :obj:`int`
Number of samples along time axis for model and data (note that this
Expand Down Expand Up @@ -180,4 +174,4 @@ def MPIMDC(G, nt, nv, nfreq, dt=1., dr=1., twosided=True,
base_comm=base_comm,
_Fredholm1=MPIFredholm1,
args_FFT={'engine': fftengine},
args_Fredholm1={'usematmul': usematmul})
args_Fredholm1={'usematmul': usematmul})
2 changes: 1 addition & 1 deletion pylops_mpi/waveeqprocessing/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
=====================================
The subpackage waveeqprocessing extends some of the wave equation processing
functionalities in pylops.waveeqprocessing providing forward and adjoint
functionalities in pylops.waveeqprocessing providing forward and adjoint
functionalities using MPI.
A list of operators present in pylops_mpi.waveeqprocessing :
Expand Down
9 changes: 4 additions & 5 deletions tests/test_fredholm.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ def test_Gsize1(par):
)
"""


@pytest.mark.mpi(min_size=2)
@pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4), (par5), (par6)])
def test_Fredholm1(par):
Expand All @@ -118,8 +119,8 @@ def test_Fredholm1(par):
usematmul=par["usematmul"],
dtype=par["dtype"],
)
x = DistributedArray(global_shape=par["nsl"] * par["ny"] * par["nz"],

x = DistributedArray(global_shape=par["nsl"] * par["ny"] * par["nz"],
partition=pylops_mpi.Partition.BROADCAST,
dtype=par["dtype"])
x[:] = 1. + par["imag"] * 1.
Expand All @@ -130,7 +131,7 @@ def test_Fredholm1(par):
# Adjoint
y_adj_dist = Fop_MPI.H @ y_dist
y_adj = y_adj_dist.asarray()

if rank == 0:
Fop = pylops.signalprocessing.Fredholm1(
F,
Expand All @@ -145,5 +146,3 @@ def test_Fredholm1(par):
y_adj_np = Fop.H @ y_np
assert_allclose(y, y_np, rtol=1e-14)
assert_allclose(y_adj, y_adj_np, rtol=1e-14)


0 comments on commit 1a9549a

Please sign in to comment.