My First Salvus Simulation

# Three standard Python imports
import os

import numpy as np
import xarray as xr

# The Salvus package.
import salvus.namespace as sn

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

domain = sn.domain.dim2.BoxDomain(x0=0.0, x1=2000.0, y0=0.0, y1=1000.0)
domain.plot()

# Once we have a domain we can initialize a project with a path on disk.
p = sn.Project.from_domain(
    path="my_project", domain=domain, load_if_exists=True
)

# Source and receivers exist for multiple dimensions and coordinate systems.
source = sn.simple_config.source.cartesian.VectorPoint2D(
    x=600.0, y=700.0, fx=0.0, fy=-1e10
)

receivers = [
    sn.simple_config.receiver.cartesian.Point2D(
        y=800.0,
        x=x,
        station_code=f"REC{i + 1}",
        # Note that one can specify the desired recording field here.
        fields=["displacement"],
    )
    for i, x in enumerate(np.linspace(1000.0, 1400.0, 5))
]

# The source and receivers object are quite similar to Python dictionaries.
source, receivers

# Sources and receivers are tied together in an event. An event can contain an
# arbitrary number of sources and receivers and each event will be one simulation.
event = sn.Event(event_name="event_a", sources=source, receivers=receivers)

# Lastly add the event to the already existing project and visualize it.
p.add_to_project(event)
p.visualizations.notebook.domain()

# It helps to define the model larger than your domain to avoid
# edge effects during the model interpolation step.
nx = 500
x = np.linspace(-100, 2100, nx)
y = np.linspace(-100, 1100, nx // 2)
xx, yy = np.meshgrid(x, y, indexing="ij")

# Background model.
vp = 3000.0 * np.ones_like(xx)
vs = 1847.5 * np.ones_like(xx)
rho = 2200.0 * np.ones_like(xx)

# Sinc perturbations.
r = ((xx - 1000.0) / 200) ** 2 + ((yy - 500) / 200) ** 2
vp -= (np.sin(r) / r) * 1000

# Assemble into one dataset.
ds = xr.Dataset(
    data_vars={
        "VP": (["x", "y"], vp),
        "VS": (["x", "y"], vs),
        "RHO": (["x", "y"], rho),
    },
    coords={"x": x, "y": y},
)

# Plot it.
ax = ds.VP.T.plot(figsize=(10, 6))
ax.axes.set_aspect("equal")

# Once we have an xarray data set, we can add it to the project.
model = sn.model.volume.cartesian.GenericModel(name="my_volume_model", data=ds)
p.add_to_project(model, overwrite=True)

sc = sn.SimulationConfiguration(
    name="my_first_simulation",
    # Mesh resolution settings.
    max_frequency_in_hertz=30.0,
    elements_per_wavelength=2.0,
    tensor_order=2,
    # Using the volumetric model we specified before.
    model_configuration=sn.ModelConfiguration(
        volume_models=["my_volume_model"],
    ),
    # Event specific configuration.
    event_configuration=sn.EventConfiguration(
        # Source wavelet.
        wavelet=sn.simple_config.stf.Ricker(center_frequency=15.0),
        waveform_simulation_configuration=sn.WaveformSimulationConfiguration(
            end_time_in_seconds=0.6
        ),
    ),
    # Absorbing boundaries everywhere but the top.
    absorbing_boundaries=sn.AbsorbingBoundaryParameters(
        reference_frequency=15.0,
        number_of_wavelengths=0.0,
        reference_velocity=1900.0,
        free_surface=["y1"],
    ),
)

# As always, we'll add it to the project, and visualize the full setup.
p.add_to_project(
    sc,
    overwrite=True,
)
p.visualizations.notebook.simulation_setup("my_first_simulation")

# It is finally time to launch the simulation.
p.simulations.launch(
    simulation_configuration="my_first_simulation",
    events=p.events.list(),
    # Determines where and how to run. Depending on the site
    # it will either directly launch, use a job queuing system, use GPUs, ...
    # The simulations will be automatically load balanced.
    site_name=SALVUS_FLOW_SITE_NAME,
    ranks_per_job=2,
    # For many applications, volumetric output is not necessary, for
    # illustrative purposes we use it here.
    extra_output_configuration={
        "volume_data": {
            "sampling_interval_in_time_steps": 10,
            "fields": ["displacement"],
        },
    },
)
p.simulations.query(block=True)

# Results can be queried per event from the project.
event_data = p.waveforms.get("my_first_simulation", p.events.list()[0])[0]
event_data

# Plot receiver data.
event_data.plot(receiver_field="displacement", component="X")

# It is of course possible to visualize the volumetric wavefield in Paraview
# and we do this in other tutorials. Here we interpolate the wavefield
# to a regular gridded xarray which is oftentimes helpful for future analysis.
ds_wavefield = sn.wavefield_output_to_xarray(
    wo=event_data.get_wavefield_output(
        output_type="volume",
        field="displacement",
        # We could access all time steps here but for larger examples this would be
        # excessive.
        time_steps=slice(30, 31),
    ),
    points=[np.linspace(0, 2000, 200), np.linspace(0, 1000, 100)],
)

# Standard xarray functionality.
ax = ds_wavefield.sel(c=0, t=ds_wavefield.t[0]).T.plot(
    shading="gouraud", infer_intervals=False, figsize=(10, 6)
)
ax.axes.set_aspect("equal")