velocity_modelling.scripts.view_velocity_model

 1"""View a velocity model in 3-dimensions."""
 2
 3from enum import StrEnum
 4from pathlib import Path
 5from typing import Annotated
 6
 7import numpy as np
 8import pyvista as pv
 9import typer
10
11app = typer.Typer()
12
13
14class Slices(StrEnum):
15    "Slicing control options for the velocity model."
16
17    XY = "xy"
18    XZ = "xz"
19    YZ = "yz"
20
21
22@app.command(help="View a velocity model in 3D.")
23def view_velocity_model(
24    velocity_model_ffp: Annotated[
25        Path, typer.Argument(help="The velocity model file to plot.")
26    ],
27    nx: Annotated[int, typer.Argument(help="The number of x-slices")],
28    ny: Annotated[int, typer.Argument(help="The number of y-slices")],
29    nz: Annotated[int, typer.Argument(help="The number of z-slices")],
30    slices: Annotated[
31        Slices, typer.Option(help="Add slice controls for the given axes.")
32    ] = Slices.XY,
33) -> None:
34    """View a velocity model in 3 dimensions.
35
36    Parameters
37    ----------
38    velocity_model_ffp : Path
39        Path to the velocity model file to view.
40    nx : int
41        The number of x-slices.
42    ny : int
43        The number of y-slices.
44    nz : int
45        The number of z-slices.
46    slices : Slices
47        The slicing controls to add, defaults to x and y axis slicing.
48    """
49    data_vs = np.fromfile(velocity_model_ffp, "<f4")
50
51    # reshape the data
52    data_vs = data_vs.reshape([ny, nz, nx])
53    data_vs = np.swapaxes(data_vs, 0, 2)
54    data_vs = np.swapaxes(data_vs, 1, 2)
55    data_vs = np.flip(data_vs, 2)  # reverse z axis
56
57    vol = pv.ImageData()
58    vol.dimensions = tuple(np.array(data_vs.shape) + 1)
59    vol.cell_data["values"] = data_vs.flatten(order="F")
60
61    p = pv.Plotter()
62
63    p.add_mesh_clip_plane(vol, assign_to_axis=slices[0])
64    p.add_mesh_clip_plane(vol, assign_to_axis=slices[1], invert=(slices[1] == "z"))
65
66    p.show()