Lamb's problem

%matplotlib inline

# This notebook will use this variable to determine which
# remote site to run on.
import os

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

import numpy as np
import xarray as xr

import salvus.namespace as sn

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

p = sn.Project.from_domain(
    path="project_lambs_problem_volumetric_model",
    domain=d,
    load_if_exists=True,
)

src = sn.simple_config.source.cartesian.MomentTensorPoint2D(
    x=500.0, y=500.0, mxx=1.0, myy=1.0, mxy=0.0
)

recs = [
    sn.simple_config.receiver.cartesian.Point2D(
        y=800.0,
        x=x,
        network_code="REC",
        station_code=f"{_i:05}",
        fields=["displacement"],
    )
    for _i, x in enumerate(np.linspace(1010.0, 1410.0, 5))
]

p.add_to_project(sn.Event(event_name="event_0", sources=src, receivers=recs))

import numpy as np
import xarray as xr

# It helps to define the model larger than your domain.
nx, ny = 600, 400
x = np.linspace(-1000, 3000, nx)
y = np.linspace(-1000, 2000, ny)
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 function.
r = ((xx - 1000.0) / 200) ** 2 + ((yy - 500) / 200) ** 2
vp -= (np.sin(r) / r) * 1000

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

ds.VP.T.plot(figsize=(10, 6))

model = sn.model.volume.cartesian.GenericModel(name="volume", data=ds)
p.add_to_project(model, overwrite=True)

p.add_to_project(
    sn.SimulationConfiguration(
        name="volumetric_model",
        max_frequency_in_hertz=30.0,
        elements_per_wavelength=1.0,
        tensor_order=1,
        model_configuration=sn.ModelConfiguration(
            background_model=sn.model.background.from_volume_model(
                model=model
            ),
            volume_models=["volume"],
        ),
        event_configuration=sn.EventConfiguration(
            wavelet=sn.simple_config.stf.Ricker(center_frequency=15.0),
            waveform_simulation_configuration=sn.WaveformSimulationConfiguration(
                end_time_in_seconds=0.6
            ),
        ),
    ),
    overwrite=True,
)

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

p.simulations.launch(
    ranks_per_job=2,
    site_name=SALVUS_FLOW_SITE_NAME,
    events=p.events.list(),
    simulation_configuration="volumetric_model",
    extra_output_configuration={
        "volume_data": {
            "sampling_interval_in_time_steps": 10,
            "fields": ["displacement"],
        },
    },
)
p.simulations.query(block=True)

xdmf_path = (
    p.simulations.get_simulation_output_directory(
        simulation_configuration="volumetric_model", event=p.events.list()[0]
    )
    / "volume_data_output_elastic_volume.xdmf"
)

# OSX:
# !open $xdmf_path

# Linux:
# !xdg-open $xdmf_path

# Might need to install imageio.
# !pip install imageio

import imageio

image = np.require(
    imageio.imread("./data/salvus_image.png").T[:, ::-1], dtype=np.float32
)

ones = np.ones((2000, 1000), dtype=np.float32)

ds_image = xr.Dataset(
    data_vars={
        "VP": (["x", "y"], ones * 3000.0 + 2 * image),
        "VS": (["x", "y"], ones * 1847.5),
        "RHO": (["x", "y"], ones * 2200.0),
    },
    coords={"x": np.linspace(0, 2000, 2000), "y": np.linspace(0, 1000, 1000)},
)

ds_image.VP.T.plot(figsize=(10, 6))

model_image = sn.model.volume.cartesian.GenericModel(
    name="volume_image", data=ds_image
)
p.add_to_project(model_image, overwrite=True)

p.add_to_project(
    sn.SimulationConfiguration(
        name="volumetric_model_image",
        max_frequency_in_hertz=30.0,
        elements_per_wavelength=2.0,
        tensor_order=4,
        model_configuration=sn.ModelConfiguration(
            background_model=sn.model.background.from_volume_model(
                model=model_image
            ),
            volume_models=["volume_image"],
        ),
        event_configuration=sn.EventConfiguration(
            wavelet=sn.simple_config.stf.Ricker(center_frequency=15.0),
            waveform_simulation_configuration=sn.WaveformSimulationConfiguration(
                end_time_in_seconds=0.6
            ),
        ),
    ),
    overwrite=True,
)

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

p.simulations.launch(
    ranks_per_job=2,
    site_name=SALVUS_FLOW_SITE_NAME,
    events=p.events.list(),
    simulation_configuration="volumetric_model_image",
    extra_output_configuration={
        "volume_data": {
            "sampling_interval_in_time_steps": 10,
            "fields": ["displacement"],
        },
    },
)
p.simulations.query(block=True)

xdmf_path = (
    p.simulations.get_simulation_output_directory(
        simulation_configuration="volumetric_model_image",
        event=p.events.list()[0],
    )
    / "volume_data_output_elastic_volume.xdmf"
)

# OSX:
# !open $xdmf_path

# Linux:
# !xdg-open $xdmf_path