Lamb's Problem

%matplotlib inline

# This line helps with tab-completion of the simple_config objects.
# The IPython/Jupyter project default to a differnet inference based
# tab completion engine which unfortunately does not yet fully work
# with SalvusFlow. This is completely optional and a convenience
# option.
%config Completer.use_jedi = False

# Standard Python packages
import os
import toml
import numpy as np
import matplotlib.pyplot as plt

# Workflow management.
from salvus.flow import api

# Specific objects to aid in setting up simulations.
from pyasdf import ASDFDataSet
from salvus.mesh.simple_mesh import basic_mesh
import salvus.flow.simple_config as sc

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

# Domain setup (m).
max_x = 2000.0  # Distance in meters.
max_y = 1000.0  # Distance in meters.
max_frequency = 25.0  # Frequency in Hz.

# Material properties.
vp = 3200.0  # P-wave velocity in m/s.
vs = 1847.5  # S-wave velocity in m/s.
rho = 2200.0  # Density in kg / m^3.

# Generate a mesh using the "simplified" meshing interface.
mesh = basic_mesh.CartesianHomogeneousIsotropicElastic2D(
    vp=vp,
    vs=vs,
    rho=rho,
    x_max=max_x,
    y_max=max_y,
    max_frequency=max_frequency,
).create_mesh()

# Visualize the mesh in the notebook.
mesh

# Write the mesh to a file which can be visualized in Paraview.
# Open the associated lambs_problem_mesh.xdmf to visualize.
mesh.write_h5("lambs_problem_mesh.h5")

# Sources.
fx, fy = 0.0, -1e10  # Source components (Nm)
sx, sy = 1000.0, 500.0  # Source position (m)
stf = sc.source.stf.Ricker(center_frequency=14.5)  # Center frequency in Hz.
source = sc.source.cartesian.VectorPoint2D(
    x=sx, y=sy, fx=fx, fy=fy, source_time_function=stf
)

# Receivers.
nr = 5  # Number of receivers.
ry0 = 800.0  # Receiver y-value.
rx0 = 1010.0  # x-value of first receiver.
rx1 = 1410.0  # x-value of last receiver.
receivers = [
    sc.receiver.cartesian.Point2D(
        x=x, y=ry0, station_code=f"{_i:03d}", fields=["displacement"]
    )
    for _i, x in enumerate(np.linspace(rx0, rx1, nr))
]

# Set up a waveform simulation object.
sim = sc.simulation.Waveform(mesh=mesh, sources=source, receivers=receivers)

# Modify the start-time, end-time, and time-step of the simulation.
sim.physics.wave_equation.end_time_in_seconds = 0.52
sim.physics.wave_equation.time_step_in_seconds = 1e-3
sim.physics.wave_equation.start_time_in_seconds = -0.08

# sim.output.volume_data.filename = "wavefield.h5"
# sim.output.volume_data.format = "hdf5"
# sim.output.volume_data.fields = ["displacement"]
# sim.output.volume_data.sampling_interval_in_time_steps = 10


# Make sure that Salvus _Compute_ will accept the current options.
sim.validate()

api.run(
    input_file=sim,
    site_name=SALVUS_FLOW_SITE_NAME,
    ranks=4,
    wall_time_in_seconds=60,
    output_folder="output",
    get_all=True,
)

def plot_asdf_seismograms(asdf_file, axes):
    """
    A simple utility function to plot traces side-by side with a
    semi-analytic solution to Lamb's problem.
    :param asdf_file: ASDF datafile containing receivers.
    :param axes: Axes on which to plot.
    """
    for i, waveform in enumerate(asdf_file.waveforms):
        for j, cmp in enumerate(["X", "Y"]):
            trace = waveform.displacement.select(component=cmp)[0]
            axes[j, i].plot(trace.data)

    for a in axes[:, 0]:
        a.set_ylabel("Displacement (m)")
    for a in axes[1, :]:
        a.set_xlabel("Time sample")

fig, axes = plt.subplots(2, 5, sharex="col", sharey="row", figsize=(20, 5))
with ASDFDataSet("data/force_verticale.h5", mode="r") as ex2ddir, ASDFDataSet(
    "output/receivers.h5", mode="r"
) as simdata:
    plot_asdf_seismograms(simdata, axes)
    plot_asdf_seismograms(ex2ddir, axes)

# Initialize first-order absorbing boundary conditions.
boundaries = sc.boundary.Absorbing(
    side_sets=["x0", "x1", "y0"], taper_amplitude=0.0, width_in_meters=0.0
)

# Associate boundaries with our simulation.
sim.add_boundary_conditions(boundaries)
sim.validate()

# Run a job.
api.run(
    input_file=sim,
    site_name=SALVUS_FLOW_SITE_NAME,
    ranks=2,
    wall_time_in_seconds=60,
    output_folder="output",
    get_all=True,
    overwrite=True,
)

# Plot results.
fig, axes = plt.subplots(2, 5, sharex="col", sharey="row", figsize=(20, 5))
with ASDFDataSet("output/receivers.h5", mode="r") as simdata, ASDFDataSet(
    "data/force_verticale.h5", mode="r"
) as analytic:
    # Plot.
    plot_asdf_seismograms(analytic, axes)
    plot_asdf_seismograms(simdata, axes)

    # Ensure seismograms are similar.
    for ana, sal in zip(analytic.waveforms, simdata.waveforms):
        for cmp in ["X", "Y"]:
            ana_trace = ana.displacement.select(component=cmp)[0].data[250:450]
            sal_trace = sal.displacement.select(component=cmp)[0].data[250:450]
            assert np.linalg.norm(ana_trace - sal_trace) < 1e-3

# Write the input file as a toml.
with open("sample_toml_file.toml", "w") as fh:
    toml.dump(sim.get_dictionary(), fh)

# Look at the file we've written.
!cat "sample_toml_file.toml"