Multiple sources are not supported

[jinshanmu]

Here is my code for an FWI of a simple two-layer model:

import sys
sys.path.append("/home/spyro")

import os
os.environ["OMP_NUM_THREADS"] = "1"

from firedrake import *
import numpy as np
import finat
from ROL.firedrake_vector import FiredrakeVector as FeVector
import ROL
from mpi4py import MPI
import psutil

import spyro

sou_pos = [(0, 0.5)]
rec_pos = spyro.create_transect((-1.0, 0.25), (-1.0, 0.75), 10)

model = {}
model["opts"] = {
    "method": "KMV", 
    "quadrature": "KMV", 
    "degree": 1, 
    "dimension": 2, 
}
model["parallelism"] = {
    "type": "spatial", 
}
model["mesh"] = {
    "Lz": 1.0, 
    "Lx": 1.0, 
    "Ly": 0.0, 
    "meshfile": "not_used.msh", 
    "initmodel": "not_used.hdf5", 
    "truemodel": "not_used.hdf5", 
}
model["BCs"] = {
    "status": True, 
    "outer_bc": "non-reflective", 
    "damping_type": "polynomial", 
    "exponent": 2, 
    "cmax": 4.5, 
    "R": 1e-6, 
    "lz": 0.25, 
    "lx": 0.25,  
    "ly": 0.0, 
}
model["acquisition"] = {
    "source_type": "Ricker", 
    "source_pos": sou_pos, 
    "frequency": 5.0, 
    "delay": 0.1, 
    "receiver_locations": rec_pos, 
}
model["timeaxis"] = {
    "t0": 0.0, 
    "tf": 2.0, 
    "dt": 5e-4, 
    "amplitude": 1, 
    "nspool": 100, 
    "fspool": 100, 
}

mesh = RectangleMesh(150, 150, 1.5, 1.5)
mesh.coordinates.dat.data[:, 0] -= 1.25
mesh.coordinates.dat.data[:, 1] -= 0.25

comm = spyro.utils.mpi_init(model)
element = spyro.domains.space.FE_method(mesh, model["opts"]["method"], model["opts"]["degree"])
V = FunctionSpace(mesh, element)

x, y = SpatialCoordinate(mesh)
velocity = conditional(x > -0.5, 1, 2)
vp = Function(V, name="velocity").interpolate(velocity)
if comm.ensemble_comm.rank == 0:
    File("simple_velocity_model.pvd", comm=comm.comm).write(vp)

sources = spyro.Sources(model, mesh, V, comm)
receivers = spyro.Receivers(model, mesh, V, comm)
wavelet = spyro.full_ricker_wavelet(
    dt=model["timeaxis"]["dt"],
    tf=model["timeaxis"]["tf"],
    freq=model["acquisition"]["frequency"],
)

p_field, p_at_recv = spyro.solvers.forward(model, mesh, comm, vp, sources, wavelet, receivers)
spyro.plots.plot_shots(model, comm, p_at_recv)
spyro.io.save_shots(model, comm, p_at_recv)

outdir = "out/"
if not os.path.exists(outdir):
    os.mkdir(outdir)

velocity = conditional(x < 0, 1.5, 1.5)
vp = Function(V, name="velocity").interpolate(velocity)
if comm.ensemble_comm.rank == 0:
    File("simple_velocity_guess.pvd", comm=comm.comm).write(vp)

if comm.ensemble_comm.rank == 0:
    control_file = File(outdir + "control.pvd", comm=comm.comm)
    grad_file = File(outdir + "grad.pvd", comm=comm.comm)
quad_rule = finat.quadrature.make_quadrature(V.finat_element.cell, V.ufl_element().degree(), "KMV")
dxlump = dx(scheme=quad_rule)

class L2Inner(object):
    def __init__(self):
        self.A = assemble(TrialFunction(V) * TestFunction(V) * dxlump, mat_type="matfree")
        self.Ap = as_backend_type(self.A).mat()

    def eval(self, _u, _v):
        upet = as_backend_type(_u).vec()
        vpet = as_backend_type(_v).vec()
        A_u = self.Ap.createVecLeft()
        self.Ap.mult(upet, A_u)
        return vpet.dot(A_u)

def regularize_gradient(vp, dJ):
    m_u = TrialFunction(V)
    m_v = TestFunction(V)
    mgrad = m_u * m_v * dx(scheme=quad_rule)
    ffG = dot(grad(vp), grad(m_v)) * dx(scheme=quad_rule)
    G = mgrad - ffG
    lhsG, rhsG = lhs(G), rhs(G)
    gradreg = Function(V)
    grad_prob = LinearVariationalProblem(lhsG, rhsG, gradreg)
    grad_solver = LinearVariationalSolver(
        grad_prob,
        solver_parameters={
            "ksp_type": "preonly",
            "pc_type": "jacobi",
            "mat_type": "matfree",
        },
    )
    grad_solver.solve()
    dJ += gradreg
    return dJ

class Objective(ROL.Objective):
    def __init__(self, inner_product):
        ROL.Objective.__init__(self)
        self.inner_product = inner_product
        self.p_guess = None
        self.misfit = 0.0
        self.p_exact_recv = spyro.io.load_shots(model, comm)

    def value(self, x, tol):
        J_total = np.zeros((1))
        self.p_guess, p_guess_recv = spyro.solvers.forward(
            model,
            mesh,
            comm,
            vp,
            sources,
            wavelet,
            receivers,
        )
        self.misfit = spyro.utils.evaluate_misfit(model, p_guess_recv, self.p_exact_recv)
        J_total[0] += spyro.utils.compute_functional(model, self.misfit, velocity=vp)
        J_total = COMM_WORLD.allreduce(J_total, op=MPI.SUM)
        J_total[0] /= comm.ensemble_comm.size
        if comm.comm.size > 1:
            J_total[0] /= comm.comm.size
        return J_total[0]

    def gradient(self, g, x, tol):
        dJ = Function(V, name="gradient")
        dJ_local = spyro.solvers.gradient(
            model,
            mesh,
            comm,
            vp,
            receivers,
            self.p_guess,
            self.misfit,
        )
        if comm.ensemble_comm.size > 1:
            comm.allreduce(dJ_local, dJ)
        else:
            dJ = dJ_local
        dJ /= comm.ensemble_comm.size
        if comm.comm.size > 1:
            dJ /= comm.comm.size
        if "regularization" in model["opts"] and model["opts"]["regularization"]:
            dJ = regularize_gradient(vp, dJ)
        if comm.ensemble_comm.rank == 0:
            grad_file.write(dJ)
        g.scale(0)
        g.vec += dJ

    def update(self, x, flag, iteration):
        vp.assign(Function(V, x.vec, name="velocity"))
        if iteration >= 0:
            if comm.ensemble_comm.rank == 0:
                control_file.write(vp)

paramsDict = {
    "General": {"Secant": {"Type": "Limited-Memory BFGS", "Maximum Storage": 10}},
    "Step": {
        "Type": "Augmented Lagrangian",
        "Augmented Lagrangian": {
            "Subproblem Step Type": "Line Search",
            "Subproblem Iteration Limit": 5,
        },
        "Line Search": {"Descent Method": {"Type": "Quasi-Newton Step"}},
    },
    "Status Test": {
        "Gradient Tolerance": 1e-15,
        "Iteration Limit": 100,
        "Step Tolerance": 1e-15,
    },
}

params = ROL.ParameterList(paramsDict, "Parameters")

inner_product = L2Inner()

obj = Objective(inner_product)

u = Function(V, name="velocity").assign(vp)

opt = FeVector(u.vector(), inner_product)

algo = ROL.Algorithm("Line Search", params)

algo.run(opt, obj)

if comm.ensemble_comm.rank == 0:
    File("res.pvd", comm=comm.comm).write(vp)

The result looks normal.

But when I increased the source number to 3:
sou_pos = [(0, 0.25), (0, 0.5), (0, 0.75)]

there is an unexpected result:

Clearly only the first source (0, 0,25) is included, with the shots saved are exactly the same for the three sources:

I suspect some bug exists for the setup of the sources, where only the first source in the array is implemented in the simulation.

Hope someone can help resolve this issue.

Thank you!

[jinshanmu]

@Olender I have tried setting self.current_source = range(self.num_receivers) like issue #6 said but nothing changed. Any suggestion🙏?

[jinshanmu]

I deleted the if condition in apply_source and it seems to work, although I'm not sure if there are other consequences to this change.

[Olender]

So, the current implementation is focused on using multiple cores, specifically with the core count being a multiple of the number of sources, with each source propagating independently. For example, a 120-core simulation with 20 sources would run each source independently with spatial parallelism of 6 cores per shot. This is the automatic parallelism strategy setup.

Nothing is wrong with your changes to run a "super shot." You shouldn't note any errors or bugs in a serial problem. However, if you were to try that in parallel, the "if" you excluded should be changed to "if self.is_local[source_id] and source_id in self.current_sources". Where the current sources are a list of the sources you wish to propagate. This should be enough if all your sources are propagated simultaneously in a parallel model. But remember to set the parallelism strategy in the dictionary to spatial. If you want to propagate sets of 2 or more shots independently in parallel, then you will need to alter the mpi method in utils accordingly

[Olender]

If you wish to make the changes so that a super shot becomes an option in the dictionary, please submit a PR because that would be a welcome addition.

[jinshanmu]

So, the current implementation is focused on using multiple cores, specifically with the core count being a multiple of the number of sources, with each source propagating independently. For example, a 120-core simulation with 20 sources would run each source independently with spatial parallelism of 6 cores per shot. This is the automatic parallelism strategy setup.

Nothing is wrong with your changes to run a "super shot." You shouldn't note any errors or bugs in a serial problem. However, if you were to try that in parallel, the "if" you excluded should be changed to "if self.is_local[source_id] and source_id in self.current_sources". Where the current sources are a list of the sources you wish to propagate. This should be enough if all your sources are propagated simultaneously in a parallel model. But remember to set the parallelism strategy in the dictionary to spatial. If you want to propagate sets of 2 or more shots independently in parallel, then you will need to alter the mpi method in utils accordingly

Thank you for the clarification!