Conical Pitting

import os

import build123d
import matplotlib.pyplot as plt
import numpy as np
import salvus.namespace as sn
import wscan

# Salvus site on which simulations are to be run
SALVUS_FLOW_SITE_NAME = os.environ.get("SITE_NAME", "local")

# Path where the project should be saved
PATH_TO_PROJECT = "project_cone_pitting_tutorial"

# Maximum frequency of the source in Hz; also controls how finely the mesh
# is to be discretized
MAX_FREQUENCY = 10.0e6
CENTER_FREQUENCY = 3 / 4 * MAX_FREQUENCY

# Number of elements per wavelength; one of the key parameters that controls
# how fine the mesh is.  Refer to
# https://docs.mondaic.com/knowledge_base/spectral_element_modelling/hp_refinement
# for more details.
ELEMENTS_PER_WAVELENGTH = 2.0

# Simulation end time in seconds
SIMULATION_END_TIME = 5.0e-6

SPECIMEN_MATERIAL = wscan.library.materials.materials["steel"].material

# Horizontal extents of the plate in meters.
X_MIN = -10.0e-3
X_MAX = 10.0e-3

# Plate thickness in meters
THICKNESS = 8.0e-3

# Diameter of the pit mouth in meters
PIT_DIAMETER = 1.0e-3

# Maximum depth of the pit in meters
PIT_DEPTH = 1.0e-3

pit = wscan.geometry.dim2.defects.ConePit(
    diameter=PIT_DIAMETER,
    depth=PIT_DEPTH,
    origin=np.array([0.0, -THICKNESS]),
    normal=np.array([0.0, -1.0]),
    n_samples=9,
    lateral_variability=0.05,
    random_seed=42,
)
pit.plot()

# Create the plate geometry; get the build123d face as a `Face` object for the
# subsequent Boolean subtraction
plate = build123d.Polygon(
    [
        (X_MIN, 0.0),
        (X_MIN, -THICKNESS),
        (X_MAX, -THICKNESS),
        (X_MAX, 0.0),
    ]
).face()

# Add the pitting to the plate.  The pit's `geometry` property returns a
# `build123d.Face` object that we can use to perform a Boolean subtraction with
# the plate.
specimen = plate - pit.geometry

# Plot the entity
specimen.label = "specimen"
wscan.geometry.dim2.utils.plot_build123d_entity(specimen)
plt.show()

mesh = sn.UnstructuredMesh.from_h5("data/plate_with_pitting.h5")

mesh.find_side_sets()
mesh

p = sn.Project.from_mesh(PATH_TO_PROJECT, mesh)

probe = wscan.instruments.transducer.ContactTransducer2D(
    center=np.array([0.0, 0.0]),
    normal=np.array([0.0, 1.0]),
    width=10.0e-3,
    maximum_frequency=MAX_FREQUENCY,
    reference_velocity=SPECIMEN_MATERIAL.VS.p,
    sampling_factor=5.0,
)

p.add_to_project(
    sn.Event(
        sources=probe.get_sources(
            excitation_mode="longitudinal",
            source_type=sn.simple_config.source.cartesian.SideSetVectorPoint2D,
            side_set_name="y1",
        ),
    )
)

stf = sn.simple_config.stf.ToneBurst(
    center_frequency=CENTER_FREQUENCY,
    num_cycles=4,
)
stf.plot()

p.add_to_project(
    sn.UnstructuredMeshSimulationConfiguration(
        name="pit",
        unstructured_mesh=mesh,
        event_configuration=sn.EventConfiguration(
            wavelet=stf,
            waveform_simulation_configuration=sn.WaveformSimulationConfiguration(
                end_time_in_seconds=SIMULATION_END_TIME,
                boundary_conditions=sn.simple_config.boundary.Absorbing(
                    taper_amplitude=MAX_FREQUENCY,
                    side_sets=["x0", "x1"],
                    width_in_meters=0.0,
                ),
            ),
        ),
    )
)

ax = wscan.geometry.dim2.utils.plot_build123d_entity(
    specimen, legend=True, figsize=(12, 12)
)

for source in p.events.get(p.events.list()[0]).sources:
    plt.plot(source.point[0], source.point[1], "k.")

plt.show()

time_step, _ = mesh.compute_time_step(
    max_velocity=np.max(mesh.elemental_fields["VP"], axis=1)
)

print(f"Time step: {time_step} s")

p.simulations.launch(
    simulation_configuration="pit",
    events=p.events.list(),
    site_name=SALVUS_FLOW_SITE_NAME,
    ranks_per_job=sn.api.get_site(SALVUS_FLOW_SITE_NAME).config[
        "default_ranks"
    ],
    extra_output_configuration={
        "volume_data": {
            "sampling_interval_in_time_steps": int(
                (SIMULATION_END_TIME - stf.get_auto_start_time())
                / time_step
                / 300
            ),
            "fields": ["displacement"],
        },
    },
)
p.simulations.query(block=True)

event_data = p.waveforms.get("pit", p.events.list()[0])[0]

wavefield = sn.wavefield_output_to_xarray(
    wo=event_data.get_wavefield_output(
        output_type="volume",
        field="displacement",
    ),
    points=[
        np.linspace(-probe.width / 2, probe.width / 2, 301),
        np.linspace(-THICKNESS, 0.0, 301),
    ],
)

# Approximate time at which the wave arrives at the pit
t_plot = np.linspace(0.5e-6, 3.0e-6, 6)

# Fixed color range for the displacement field
vrange = 1.0e-11

# Create the plot
_, axs = plt.subplots(
    nrows=2, ncols=3, figsize=(10, 6), sharex=True, sharey=True
)

for i, ax in enumerate(axs.flatten()):
    # Plot the outline of the geometry
    ax = wscan.geometry.dim2.utils.plot_build123d_entity(
        specimen,
        face_alpha=0.0,
        ax=ax,
    )

    # Get the time index that is closest to `t_plot`
    t_ind = np.argmin(np.abs(wavefield.t.to_numpy() - t_plot[i]))

    # Plot the wavefield; `c=1` is the y-component of displacement, `c=0` is
    # the x-component
    wavefield.sel(c=1, t=wavefield.t[t_ind]).T.plot(
        shading="gouraud",
        infer_intervals=False,
        ax=ax,
        cmap="coolwarm",
        zorder=-1,
        add_colorbar=False,
        vmin=-vrange,
        vmax=vrange,
    )

    # Formatting
    ax.set_aspect("equal")
    ax.set_xlabel(None)
    ax.set_ylabel(None)

# Add x- and y-labels
for ax in axs[:, 0]:
    ax.set_ylabel("y [m]")
for ax in axs[-1, :]:
    ax.set_xlabel("x [m]")

plt.tight_layout()
plt.show()

print(
    p.simulations.get_simulation_output_directory(
        simulation_configuration="pit", event=p.events.list()[0]
    )
)