2-D Global Gradients without SalvusProject

# Global site name for all runs
import os

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

from salvus.mesh.simple_mesh import basic_mesh
from salvus.flow import simple_config
from salvus.flow import api
import salvus.namespace as sn

import json
import h5py
import matplotlib.pyplot as plt
import numpy as np
import obspy
import pathlib
import typing

# The most important settings are collected here.
PERIOD = 100.0

# A good idea as the gradients are always computed
# relative to the input parameterization. This will
# just result in a much smoother looking gradient.
# Planet will also be actually round.
TENSOR_ORDER = 4

# Probably up this a bit for any final results but works
# just fine for testing purposes.
ELEMENTS_PER_WAVELENGTH = 1.0

# No latitude in 2-D. Receiver always at the surface.
SRC_LONGITUDE = 0.0
SRC_DEPTH_KM = 100.0
REC_LONGITUDE = 135.0

# It will by default select a window around this phase.
PHASE = "Pdiff"

m = basic_mesh.Circular2D()

m.basic.model = "prem_ani_no_crust"
m.basic.min_period_in_seconds = PERIOD
m.basic.elements_per_wavelength = ELEMENTS_PER_WAVELENGTH
m.advanced.tensor_order = TENSOR_ORDER

m

src = simple_config.source.seismology.SideSetMomentTensorPoint2D(
    longitude=SRC_LONGITUDE,
    depth_in_m=SRC_DEPTH_KM * 1000.0,
    mrr=0.0,
    mpp=0.0,
    mrp=1e20,
    side_set_name="r1",
    source_time_function=simple_config.stf.GaussianRate(
        half_duration_in_seconds=PERIOD / 2.0
    ),
)

rec = simple_config.receiver.seismology.SideSetPoint2D(
    longitude=REC_LONGITUDE,
    station_code="AA",
    side_set_name="r1",
    fields=["velocity"],
)

# Create the simulation object.
w = simple_config.simulation.Waveform(
    mesh=m.create_mesh(),
    sources=src,
    receivers=rec,
    end_time_in_seconds=2000.0,
    # Necessary to be able to later run an adjoint
    # simulation.
    store_adjoint_checkpoints=True,
)

w.validate()
w

job_forward = api.run(
    input_file=w,
    site_name=SALVUS_FLOW_SITE_NAME,
    output_folder="output",
    ranks=4,
    overwrite=True,
    delete_remote_files=False,
)

event = job_forward.get_as_event()
event.get_receiver_data(
    receiver_name="XX.AA.", receiver_field="velocity"
).plot()

from obspy.taup import TauPyModel
from obspy.geodetics import locations2degrees

# Same model as in the Salvus meshing constructor.
m = TauPyModel(model="prem")

distance = locations2degrees(0.0, SRC_LONGITUDE, 0.0, REC_LONGITUDE)

arrivals = m.get_ray_paths(
    distance_in_degree=distance,
    source_depth_in_km=SRC_DEPTH_KM,
    phase_list=[PHASE],
)

arrivals.plot_rays()

window_center = arrivals[0].time

print(f"Distance: {distance} degrees")
print(f"Window center: {window_center:.2f} seconds")

from salvus.project.tools.data_selection import compute_window


def temporal_weights_fct(
    st: obspy.Stream,
    receiver: simple_config.receiver._Base,
    sources: typing.List[simple_config.source._Base],
) -> typing.Dict[str, typing.List[typing.Dict]]:
    # Absolute time array.
    t = st[0].times() + st[0].stats.starttime.timestamp

    weights = {}
    for tr in st:
        component = tr.stats.channel[-1]
        # Only keep a window on the vertical component.
        if component != "Z":
            continue
        weights[component] = [
            {
                "values": compute_window(
                    t=t,
                    window_center=window_center,
                    window_width=2.0 * PERIOD,
                    ramp_width=PERIOD,
                )
            }
        ]
    return weights


event.register_temporal_weights_function(temporal_weights_fct)
event.get_receiver_data(
    receiver_name="XX.AA.", receiver_field="velocity"
).plot()

def cc_misfit(data_synthetic: np.ndarray, sampling_rate_in_hertz: float):
    dt = 1.0 / sampling_rate_in_hertz

    s_dt = np.gradient(data_synthetic, dt)
    s_dt_2 = np.gradient(s_dt, dt)

    N = (data_synthetic * s_dt_2).sum() * dt

    adj_src = -1.0 * (1.0 / N * s_dt)

    # The first value is the misfit which does not really have a meaning here.
    return 1.0, adj_src


# Finally construct the misfit object.
event_misfit = sn.EventMisfit(
    # This particular misfit function does
    # not require data.
    observed_event=None,
    synthetic_event=event,
    # Passed as a function here for demonstration purposes.
    # One could also pass a string to use Salvus' built-in
    # misfit functions.
    misfit_function=cc_misfit,
    receiver_field="velocity",
)

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="tti",
    gradient_output_format="hdf5",
)

# Finally run it to produce the gradient.
job_adjoint = api.run(
    input_file=w_adjoint,
    site_name=SALVUS_FLOW_SITE_NAME,
    output_folder="output_adjoint",
    overwrite=True,
    ranks=4,
)

# Visualize the gradients. Make sure to play with the color
# scale to get a nicer result.

from salvus.mesh.unstructured_mesh import UnstructuredMesh

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

# The forward job can now be deleted.
job_forward.delete()