Seismological Gradients in SalvusProject

# 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")
PROJECT_DIR = "project"

import time
import pathlib

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

from salvus import namespace as sn

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

# Define a spherical chunk domain.
lat_c, lon_c = -27.5, 135
lat_e, lon_e = 25.5, 42.0

# Uncomment the following line to start from scratch
# !rm -rf proj
if pathlib.Path(PROJECT_DIR).exists():
    print("Opening existing project.")
    p = sn.Project(path=PROJECT_DIR)
else:
    print("Creating new project.")
    d = sn.domain.dim3.SphericalChunkDomain(
        lat_center=lat_c,
        lat_extent=lat_e,
        lon_center=lon_c,
        lon_extent=lon_e,
        radius_in_meter=6371e3,
    )
    p = sn.Project.from_domain(path=PROJECT_DIR, domain=d)

# Add an event by parsing some prepared data.
e = sn.Event(
    sources=sn.simple_config.source.seismology.parse(
        "./event.txt", dimensions=3
    ),
    receivers=sn.simple_config.receiver.seismology.parse(
        obspy.read_inventory("stations.txt"),
        dimensions=3,
        fields=["displacement"],
    ),
)

p += e

lat = np.linspace(lat_c - lat_e / 2.0, lat_c + lat_e / 2.0, 50)
lon = np.linspace(lon_c - lon_e / 2.0, lon_c + lon_e / 2.0, 50)
depth = np.linspace(0, 1000.0, 40)

_, lon_grid, lat_grid = np.meshgrid(depth, lon, lat, indexing="ij")

# Distance from center in degree.
d = ((lon_grid - lon_c) ** 2 + (lat_grid - lat_c) ** 2) ** 0.5

# Apply Gaussian
sigma = 5.0
d = 1.0 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-0.5 * (d / sigma) ** 2)

# Normalize to 10 % deviation
d /= d.max()
d *= 10.0

ds = xr.Dataset(
    data_vars={
        "vs": (["depth", "longitude", "latitude"], d),
        "vp": (["depth", "longitude", "latitude"], d),
    },
    coords={"depth": depth, "latitude": lat, "longitude": lon},
)

ds.vs.attrs["units"] = "%"
ds.vp.attrs["units"] = "%"

# Same for the coordinate axes.
ds.depth.attrs["units"] = "km"
ds.latitude.attrs["units"] = "degrees_north"
ds.longitude.attrs["units"] = "degrees_east"
ds.attrs["geospatial_lat_units"] = "degrees_north"
ds.attrs["geospatial_lon_units"] = "degrees_east"
ds.attrs["geospatial_vertical_units"] = "km"

ds.vs.isel(depth=0).T.plot()
plt.show()

ec = sn.EventConfiguration(
    wavelet=sn.simple_config.stf.GaussianRate(half_duration_in_seconds=100.0),
    waveform_simulation_configuration=sn.WaveformSimulationConfiguration(
        end_time_in_seconds=1800.0
    ),
)

mc = sn.ModelConfiguration(background_model="prem_iso_no_crust")

p += sn.SimulationConfiguration(
    tensor_order=1,
    name="initial_model",
    elements_per_wavelength=1.0,
    min_period_in_seconds=200,
    max_depth_in_meters=500e3,
    model_configuration=sn.ModelConfiguration(
        background_model="prem_iso_no_crust",
    ),
    event_configuration=ec,
)

p += sn.model.volume.seismology.GenericModel(name="vs_blob", data=ds)

p += sn.SimulationConfiguration(
    tensor_order=1,
    name="target_model",
    elements_per_wavelength=1.0,
    min_period_in_seconds=200,
    max_depth_in_meters=500e3,
    model_configuration=sn.ModelConfiguration(
        background_model="prem_iso_no_crust", volume_models=["vs_blob"]
    ),
    event_configuration=ec,
)

p.viz.nb.simulation_setup("target_model", events=p.events.list())

for simulation in ["initial_model", "target_model"]:
    p.simulations.launch(
        ranks_per_job=4,
        site_name=SALVUS_FLOW_SITE_NAME,
        events=p.events.list()[0],
        wall_time_in_seconds_per_job=3600,
        simulation_configuration=simulation,
    )

    while not p.simulations.query():
        time.sleep(1.0)

def process(st, inv, receiver, sources):
    # This particular function adds some random noise to
    # the data. To make it reproducible let's create some
    # deterministic seed - the seed is based on the receiver
    # number and the point-to-point amplitude of the first
    # trace.
    np.random.seed(
        int.from_bytes((receiver.name + str(st[0].data.ptp())).encode(), "big")
        % 2 ** 32
    )
    # Add 5% of normal distributed noise.
    for tr in st:
        tr.data += np.random.normal(
            scale=0.05 * tr.data.ptp(), size=tr.data.shape
        )

    # Filter it again to not have really high-frequencies.
    st.filter("lowpass", freq=1.0 / 80.0, zerophase=True)

    return st


p += sn.processing.seismology.SeismologyProcessingConfiguration(
    name="target_model_with_noise",
    data_source_name="SYNTHETIC_DATA:target_model",
    processing_function=process,
)

if not p.entities.has(
    "data_selection_configuration", "initial_model_100_to_300s"
):
    p.actions.seismology.pick_windows(
        data_selection_configuration="initial_model_100_to_300s",
        observed_data_name="PROCESSED_DATA:target_model_with_noise",
        synthetic_data_name="initial_model",
        events=p.events.list(),
        receiver_field="displacement",
        window_taper_width_in_seconds=50.0,
        window_picking_function="built-in",
        window_picking_function_kwargs={
            "minimum_period_in_seconds": 100.0,
            "maximum_period_in_seconds": 300.0,
        },
    )

p.viz.nb.waveforms(
    data=[
        "SYNTHETIC_DATA:initial_model",
        "PROCESSED_DATA:target_model_with_noise",
    ],
    receiver_field="displacement",
    data_selection_configuration="initial_model_100_to_300s",
)

p.viz.seismology.windows("initial_model_100_to_300s")

def l2_misfit_and_adjoint_source(
    data_synthetic: np.ndarray,
    data_observed: np.ndarray,
    sampling_rate_in_hertz: float,
):
    # Necessary to not run into crazy small numbers.
    f = 1e10
    assert data_synthetic.shape == data_observed.shape

    # the adjoint source is the negative derivative w.r.t.
    # synthetic data, i.e., we flip the sign here
    diff = data_observed - data_synthetic

    misfit = 0.5 * (diff ** 2).sum() / sampling_rate_in_hertz * f
    adj_src = diff / sampling_rate_in_hertz * f

    return misfit, adj_src


# Add a misfit configuration to the project.
# Note that this is independet of the synthetic data of the forward run
# but it only defines how to compare two sets of data.
p += sn.MisfitConfiguration(
    name="L2-misfit-to-target-model",
    observed_data="PROCESSED_DATA:target_model_with_noise",
    misfit_function=l2_misfit_and_adjoint_source,
    receiver_field="displacement",
    data_selection_configuration="initial_model_100_to_300s",
)

# Misfits can be computed before running an adjoint simulation.
misfits = None
while not misfits:
    misfits = p.actions.inversion.compute_misfits(
        simulation_configuration="initial_model",
        misfit_configuration="L2-misfit-to-target-model",
        store_checkpoints=True,
        events=p.events.list()[0],
        ranks_per_job=4,
        site_name=SALVUS_FLOW_SITE_NAME,
    )
    time.sleep(5.0)

print(misfits)

# Actually compute the gradients.
while not p.actions.inversion.compute_gradients(
    simulation_configuration="initial_model",
    misfit_configuration="L2-misfit-to-target-model",
    events=p.events.list()[0],
    ranks_per_job=4,
    site_name=SALVUS_FLOW_SITE_NAME,
):
    time.sleep(5.0)
p.viz.nb.gradients(
    simulation_configuration="initial_model",
    misfit_configuration="L2-misfit-to-target-model",
    events=p.events.list(),
)

# Look at the misfits from the different receivers.
p.viz.nb.misfit_comparison(
    reference_data="initial_model",
    misfit_configuration="L2-misfit-to-target-model",
    event=p.events.list()[0],
)

# Can also be shown on a map.
p.viz.seismology.misfit_map(
    reference_data="initial_model",
    misfit_configuration="L2-misfit-to-target-model",
)

# Or have a detailed look at how the individual misfits are computed.
p.viz.nb.misfits(
    simulation_configuration="initial_model",
    misfit_configuration="L2-misfit-to-target-model",
)

p.viz.misfit_histogram(
    simulation_configuration_a="initial_model",
    misfit_configuration="L2-misfit-to-target-model",
    events=p.events.list(),
)