My First Salvus Inversion

import os
import numpy as np
import xarray as xr

import salvus.namespace as sn

SALVUS_FLOW_SITE_NAME = os.environ.get("SITE_NAME", "local")

# Generate a ring with 5 scalar sources.
sources = sn.simple_config.source.cartesian.collections.ScalarPoint2DRing(
    x=0, y=0, radius=0.09, count=5, f=1.0
)

events = []
# Use the same receiver ring for each but mask out the receivers close to each
# source.
for i, s in enumerate(sources._sources):
    all_receivers = (
        sn.simple_config.receiver.cartesian.collections.RingPoint2D(
            x=0, y=0, radius=0.09, count=100, fields=["phi"]
        )
    )
    receivers = [
        r
        for r in all_receivers._receivers
        if np.sqrt(
            (s.location[0] - r.location[0]) ** 2
            + (s.location[1] - r.location[1]) ** 2
        )
        > 0.03
    ]
    events.append(
        sn.Event(event_name=f"event_{i}", sources=s, receivers=receivers)
    )

def get_spherical_inclusion() -> xr.Dataset:
    nx, ny = 200, 200
    x = np.linspace(-0.1, +0.1, nx)
    y = np.linspace(-0.1, +0.1, ny)
    xx, yy = np.meshgrid(x, y, indexing="ij")

    # Add 3 spherical inclusions
    vp = 1500.0 * np.ones_like(xx)
    rho = 980.0 * np.ones_like(xx)
    mask = np.sqrt(xx**2 + yy**2) < 0.05
    vp[mask] = 1480.0
    rho[mask] = 1000.0

    mask = np.sqrt(xx**2 + (yy - 0.025) ** 2) < 0.015
    vp[mask] = 1550.0
    rho[mask] = 1040.0

    mask = np.sqrt(xx**2 + (yy + 0.025) ** 2) < 0.015
    vp[mask] = 1460.0
    rho[mask] = 1010.0

    ds = xr.Dataset(
        data_vars={
            "vp": (["x", "y"], vp),
            "rho": (["x", "y"], rho),
        },
        coords={"x": x, "y": y},
    )

    return ds


true_model = get_spherical_inclusion()
ax = true_model.vp.T.plot(figsize=(10, 6))
ax.axes.set_aspect("equal")

# Use the volume model to set up the project.
volume_model = sn.model.volume.cartesian.GenericModel(
    name="true_model", data=true_model
)
p = sn.Project.from_volume_model(
    path="project_full_waveform_inversion",
    volume_model=volume_model,
    load_if_exists=True,
)

# Add the events.
for event in events:
    p.add_to_project(event)

p.visualizations.notebook.domain()

# Generate the data for the target model.
wsc = sn.WaveformSimulationConfiguration(end_time_in_seconds=0.00015)
wsc.physics.wave_equation.time_step_in_seconds = 4.0e-7

ec = sn.EventConfiguration(
    waveform_simulation_configuration=wsc,
    wavelet=sn.simple_config.stf.Ricker(center_frequency=50000.0),
)
p += sn.SimulationConfiguration(
    name="true_model_100kHz",
    elements_per_wavelength=2,
    tensor_order=4,
    max_frequency_in_hertz=100000.0,
    model_configuration=sn.ModelConfiguration(
        background_model=None, volume_models="true_model"
    ),
    event_configuration=ec,
)
p.simulations.launch(
    simulation_configuration="true_model_100kHz",
    events=p.events.list(),
    site_name=SALVUS_FLOW_SITE_NAME,
    ranks_per_job=1,
)
p.simulations.query(block=True)

bm = sn.model.background.homogeneous.IsotropicAcoustic(vp=1500.0, rho=980.0)
mc = sn.ModelConfiguration(background_model=bm)
wsc = sn.WaveformSimulationConfiguration(end_time_in_seconds=0.00015)

p += sn.SimulationConfiguration(
    name="initial_model",
    elements_per_wavelength=2,
    tensor_order=2,
    max_frequency_in_hertz=100000.0,
    model_configuration=mc,
    event_configuration=sn.EventConfiguration(
        waveform_simulation_configuration=wsc,
        wavelet=sn.simple_config.stf.Ricker(center_frequency=50000.0),
    ),
)
p.simulations.launch(
    simulation_configuration="initial_model",
    events=p.events.list(),
    site_name=SALVUS_FLOW_SITE_NAME,
    ranks_per_job=1,
)
p.simulations.query(block=True)

# The misfit configuration defines how synthetics are compared to observed data.
p += sn.MisfitConfiguration(
    name="L2",
    # Could be observed data. Here we compare to the synthetic target.
    observed_data="true_model_100kHz",
    # Salvus comes with a variety of misfit functions. You can
    # also define your own.
    misfit_function="L2",
    # This is an acoustic simulation so we'll use recordings of phi.
    receiver_field="phi",
)

# Now we define the actual inverse problem
p += sn.InverseProblemConfiguration(
    name="my_inversion",
    # Starting model.
    prior_model="initial_model",
    # The events to use.
    events=p.events.list(),
    # What parameters to invert for.
    mapping=sn.Mapping(scaling="absolute", inversion_parameters=["VP", "RHO"]),
    # The smoothing.
    preconditioner=sn.ConstantSmoothing({"VP": 0.01, "RHO": 0.01}),
    # The inversion method.
    method=sn.TrustRegion(initial_trust_region_linf=10.0),
    # The misfit configuration we defined above.
    misfit_configuration="L2",
    # Compress the forward wavefield by subsampling in time.
    wavefield_compression=sn.WavefieldCompression(
        forward_wavefield_sampling_interval=10
    ),
    # Job submission settings.
    job_submission=sn.SiteConfig(
        site_name=SALVUS_FLOW_SITE_NAME, ranks_per_job=1
    ),
)

# Lastly we perform two iterations, and have a look at the results.
for i in range(2):
    p.inversions.iterate(
        inverse_problem_configuration="my_inversion",
        timeout_in_seconds=360,
        ping_interval_in_seconds=2,
        delete_disposable_files="all",
    )