Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Multiple sources are not supported #102

Open
jinshanmu opened this issue Jan 16, 2024 · 5 comments
Open

Multiple sources are not supported #102

jinshanmu opened this issue Jan 16, 2024 · 5 comments
Labels
enhancement New feature or request

Comments

@jinshanmu
Copy link

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.
res1

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:
res2

Clearly only the first source (0, 0,25) is included, with the shots saved are exactly the same for the three sources:
Screenshot 2024-01-16 at 11 31 58 AM

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
Copy link
Author

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

@jinshanmu
Copy link
Author

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.
res

@Olender
Copy link
Collaborator

Olender commented Feb 5, 2024

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
Copy link
Collaborator

Olender commented Feb 5, 2024

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
Copy link
Author

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!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

When branches are created from issues, their pull requests are automatically linked.

2 participants