Manual Gradient Computation

# Variable used in the notebook to determine which site
# is used to run the simulations.
import os

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

%matplotlib inline

import matplotlib.pyplot as plt
import numpy as np
import pathlib
import xarray as xr

import salvus.namespace as sn

# Use xarray to define the blob in vs.
def get_target_model():
    x = np.linspace(-500, 500, 200)
    y = np.linspace(-500, 500, 200)
    xx, yy = np.meshgrid(x, y, indexing="ij")

    # Simple gaussian blob around the center of the domain.
    d = (xx**2 + yy**2) ** 0.5
    vs = np.exp(-0.5 * (d / 100.0) ** 2) * 250.0 + 1000.0

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

    return ds


target_model = get_target_model()
target_model.vs.plot()

# Build mesh for the initial model.
initial_model_mesh = sn.simple_mesh.CartesianHomogeneousIsotropicElastic2D(
    vp=1500.0,
    vs=1000.0,
    rho=2000.0,
    x_max=float(target_model.x.data.ptp()),
    y_max=float(target_model.y.data.ptp()),
    max_frequency=10.0,
    tensor_order=1,
).create_mesh()

# Shift.
initial_model_mesh.points[:, 0] += target_model.x.data.min()
initial_model_mesh.points[:, 1] += target_model.y.data.min()

target_model_mesh = initial_model_mesh.copy()

target_model_mesh = sn.toolbox.interpolate_cartesian(
    data=target_model, mesh=target_model_mesh, extrapolate=True
)
target_model_mesh

# Explosive source in one corner of the mesh.
src = sn.simple_config.source.cartesian.MomentTensorPoint2D(
    mxx=10000.0,
    myy=10000.0,
    mxy=0.0,
    x=300.0,
    y=300.0,
    source_time_function=sn.simple_config.stf.Ricker(center_frequency=5.0),
)

# A centered ring of receivers.
recs = sn.simple_config.receiver.cartesian.collections.RingPoint2D(
    x=0, y=0, radius=200.0, count=50, fields=["displacement"]
)

# Create two simulation objects.
# (1) The target will serve as observed data.
w_target = sn.simple_config.simulation.Waveform(
    mesh=target_model_mesh, sources=src, receivers=recs
)
w_target.physics.wave_equation.end_time_in_seconds = 1.0
w_target


# (2) The initial model will be used to compute synthetics.
w_initial = sn.simple_config.simulation.Waveform(
    mesh=initial_model_mesh, sources=src, receivers=recs
)
w_initial.physics.wave_equation.end_time_in_seconds = 1.0

# Must store the checkpoints for the subsequent adjoint run.
w_initial.output.volume_data.format = "hdf5"
w_initial.output.volume_data.filename = "output.h5"
w_initial.output.volume_data.fields = ["adjoint-checkpoint"]
w_initial.output.volume_data.sampling_interval_in_time_steps = (
    "auto-for-checkpointing"
)

# Create the "observed data".
j_target = sn.api.run(
    input_file=w_target,
    site_name=SALVUS_FLOW_SITE_NAME,
    output_folder="output_target",
    overwrite=True,
)

# Forward run for the synthetics.
j_initial = sn.api.run(
    input_file=w_initial,
    site_name=SALVUS_FLOW_SITE_NAME,
    output_folder="output_initial",
    overwrite=True,
    # Don't delete the remote files as they are needed
    # for the adjoint run!
    delete_remote_files=False,
)

# Load the data from both runs, and visualize.
#
# Note that this is a vectorial receiver and thus the component
# has to be specified.
target_event = j_target.get_as_event()
initial_event = j_initial.get_as_event()

target_event.plot(receiver_field="displacement", component="X")
initial_event.plot(receiver_field="displacement", component="X")

import typing

import numpy as np
import obspy

from salvus.project.tools.data_selection import compute_window


def compute_temporal_weights(
    st: obspy.Stream,
    receiver: sn.simple_config.receiver._Base,
    sources: typing.List[sn.simple_config.source._Base],
) -> typing.Dict[str, typing.List[np.ndarray]]:
    """
    Function computing temporal weights.

    This example function will taper to zero after the first arrival
    which is just computed geometrically in this case.

    Args:
        st: The input data, one trace per component. Don't
            modify it!
        receiver: The receiver for the data.
        sources: All sources for this event. Just one source
            in many cases.

    Returns:
        A dictionary with a list of weight sets per component.
    """
    # This project only has a single source per event.
    # Does not yet account for the source time function.
    distance = np.linalg.norm(
        np.array(receiver.location) - np.array(sources[0].location)
    )

    first_arrival = distance / 1000.0 - 0.1  # Account for STF
    taper_length = 0.1

    weights = {}
    for tr in st:
        all_weights_for_components = []
        component = tr.stats.channel[-1]
        times = tr.times() + tr.stats.starttime.timestamp

        # Only a single window for each trace.
        weights[component] = [
            {
                "values": compute_window(
                    t=times,
                    window_center=first_arrival + 0.05,
                    window_width=0.2,
                    ramp_width=0.1,
                )
            }
        ]

    return weights


# Can be hooked up to the event that will then process
# the data on the fly.
initial_event.register_temporal_weights_function(compute_temporal_weights)
target_event.register_temporal_weights_function(compute_temporal_weights)

# Now the processing is always applied for any subsequent action.
initial_event.plot(receiver_field="displacement", component="X")
target_event.plot(receiver_field="displacement", component="X")

# Now construct a misfit object.
event_misfit = sn.EventMisfit(
    # For the purpose of this gradient we treat the simulation
    # with the target model as observed data and the simulation
    # with the intitial model as the synthetic data.
    observed_event=target_event,
    synthetic_event=initial_event,
    # Choose misfit function - can also be a function.
    misfit_function="L2",
    # Some misfit function take extra keyword arguments.
    extra_kwargs_misfit_function={},
    # Let's measure on the displacement.
    receiver_field="displacement",
)

print("Total misfit for event:", event_misfit.misfit_value)

from salvus.flow.simple_config.simulation_generator import (
    create_adjoint_waveform_simulation,
)

# Use the adjoint source to generate a simulation object for the
# adjoint simulation. It is aware of the parameters of the forward
# simulation and can thus guarantee everything is compatible.
adjoint_source_filename = pathlib.Path("adjoint_source.h5")
event_misfit.write(filename=adjoint_source_filename)

w_adjoint = create_adjoint_waveform_simulation(
    meta_json_forward_run=event_misfit.synthetic_event.meta_json_contents,
    adjoint_source_file=adjoint_source_filename,
    gradient_parameterization="rho-vp-vs",
)

# Last but not least actually run the adjoint simulation to compute
# the gradient.
sn.api.run(
    input_file=w_adjoint,
    site_name=SALVUS_FLOW_SITE_NAME,
    output_folder="output_gradient",
    overwrite=True,
)

# Time to clean up the forward run.
j_initial.delete()

# Finally visualize the gradient.
from salvus.mesh.unstructured_mesh import UnstructuredMesh

UnstructuredMesh.from_h5("./output_gradient/gradient.h5")