velocity_modelling.scripts.plot_velocity_model

  1from enum import Enum
  2from pathlib import Path
  3from typing import Annotated, Optional
  4
  5import numpy as np
  6import pandas as pd
  7import pygmt
  8import typer
  9import yaml
 10
 11from pygmt_helper import plotting
 12from qcore import coordinates, grid
 13from velocity_modelling import bounding_box
 14from velocity_modelling.bounding_box import BoundingBox
 15
 16app = typer.Typer()
 17
 18
 19class VelocityModelComponent(str, Enum):
 20    """Enumeration for different components of a velocity model.
 21
 22    Attributes
 23    ----------
 24    p_wave : str
 25        P-wave (Primary wave) velocity component.
 26    s_wave : str
 27        S-wave (Secondary wave) velocity component.
 28    density : str
 29        Density component of the velocity model.
 30    """
 31
 32    p_wave = "p_wave"
 33    s_wave = "s_wave"
 34    density = "rho"
 35
 36
 37CB_LABEL_MAP = {
 38    VelocityModelComponent.s_wave: "S Wave Velocity",
 39    VelocityModelComponent.p_wave: "P Wave Velocity",
 40    VelocityModelComponent.density: "Density",
 41}
 42
 43
 44def plot_velocity_model(
 45    fig: pygmt.Figure,
 46    box: BoundingBox,
 47    velocity_model: np.memmap,
 48    resolution: float,
 49    slice: int,
 50    nx: int,
 51    ny: int,
 52    region: tuple[float, ...],
 53    **kwargs,
 54):
 55    """Plot the velocity model on a pygmt figure.
 56
 57    Parameters
 58    ----------
 59    fig : pygmt.Figure
 60        The figure to plot on.
 61    box : BoundingBox
 62        The bounding box the bounds of the velocity model.
 63    velocity_model : np.memmap
 64        The velocity model data as a memory-mapped array.
 65    resolution : float
 66        The resolution of the grid in meters.
 67    slice : int
 68        The z-level slice of the velocity model to plot.
 69    nx : int
 70        The number of grid points in the x-direction.
 71    ny : int
 72        The number of grid points in the y-direction.
 73    region : tuple of float
 74        The plotting region.
 75    **kwargs
 76        Additional keyword arguments passed to plotting.plot_grid.
 77    """
 78    corners = np.c_[box.corners, np.zeros_like(box.corners[:, 0])]
 79
 80    # The indexing here is to account for:
 81    # 1. The coordinate meshgrid generating along the boundary of
 82    # the domain (hence, 1:-1, 1:-1),
 83    # 2. The fact we don't need depth (hence :2)
 84    lat_lon_grid = grid.coordinate_meshgrid(
 85        corners[-1], corners[0], corners[-2], resolution * 1000
 86    )[1:-1, 1:-1, :2]
 87
 88    # The lat lon grid has shape (nx, ny, nz), so we flip that to
 89    # make the veloity model.
 90    lat_lon_grid = np.transpose(lat_lon_grid, (1, 0, 2))
 91    velocity_slice = velocity_model[:, slice, :].reshape((ny, nx))
 92    velocity_model_df = pd.DataFrame(
 93        {
 94            "lat": lat_lon_grid[:, :, 0].ravel(),
 95            "lon": lat_lon_grid[:, :, 1].ravel(),
 96            "value": velocity_slice.ravel(),
 97        }
 98    )
 99
100    cur_grid = plotting.create_grid(
101        velocity_model_df,
102        "value",
103        grid_spacing="100e/100e",
104        region=region,
105        set_water_to_nan=False,
106    )
107    cmap_limits = (
108        velocity_model_df["value"].min().round(1),
109        velocity_model_df["value"].max().round(1),
110        (
111            (velocity_model_df["value"].max() - velocity_model_df["value"].min()) / 10
112        ).round(1),
113    )
114
115    plotting.plot_grid(
116        fig,
117        cur_grid,
118        "hot",
119        cmap_limits,
120        ("white", "black"),
121        reverse_cmap=True,
122        plot_contours=False,
123        continuous_cmap=True,
124        **kwargs,
125    )
126
127
128def load_velocity_model_file(
129    velocity_model_ffp: Path,
130    component: VelocityModelComponent,
131    nx: int,
132    ny: int,
133    nz: int,
134) -> np.memmap:
135    """Load a velocity model from a file as a memory-mapped array.
136
137    Parameters
138    ----------
139    velocity_model_ffp : Path
140        The file path to the velocity model data.
141    component : VelocityModelComponent
142        The component of the velocity model to load.
143    nx : int
144        The number of grid points in the x-direction.
145    ny : int
146        The number of grid points in the y-direction.
147    nz : int
148        The number of grid points in the z-direction.
149
150    Returns
151    -------
152    np.memmap
153        The loaded velocity model as a memory-mapped array.
154    """
155    filepath_map = {
156        VelocityModelComponent.p_wave: "vp3dfile.p",
157        VelocityModelComponent.s_wave: "vs3dfile.s",
158        VelocityModelComponent.density: "rho3dfile.d",
159    }
160
161    return np.memmap(
162        velocity_model_ffp / filepath_map[component],
163        shape=(ny, nz, nx),
164        # NOTE: This may require tweaks in the future to account
165        # for differing endianness across machines. For now, this
166        # should be fine.
167        dtype="<f4",
168    )
169
170
171@app.command(
172    help="Plot a velocity model from a velocity model params yaml file (vm_params.yaml).",
173    name="vm-params",
174)
175def plot_vm_params(
176    vm_params_ffp: Annotated[
177        Path, typer.Argument(help="Path to velocity model parameters")
178    ],
179    output_plot_path: Annotated[
180        Path,
181        typer.Argument(
182            help="Path to output velocity model plot.", writable=True, dir_okay=False
183        ),
184    ],
185    velocity_model_ffp: Annotated[
186        Path | None,
187        typer.Option(
188            help="Path to velocity model directory", file_okay=False, exists=True
189        ),
190    ] = None,
191    component: Annotated[
192        VelocityModelComponent,
193        typer.Option(help="Velocity model component to overlay."),
194    ] = VelocityModelComponent.p_wave,
195    slice: Annotated[
196        int, typer.Option(help="z-level slice of velocity model to plot.")
197    ] = 0,
198    transparency: Annotated[
199        int, typer.Option(help="Velocity model overlay transparency, (0 = opaque)")
200    ] = 80,
201    title: Annotated[Optional[str], typer.Option(help="Figure title")] = None,
202    latitude_pad: Annotated[
203        float, typer.Option(help="Latitude padding for figure (in degrees)", min=0)
204    ] = 0,
205    longitude_pad: Annotated[
206        float, typer.Option(help="Longitude padding for figure (in degrees)", min=0)
207    ] = 0,
208    dpi: Annotated[
209        float, typer.Option(help="Plot output DPI (higher is better)")
210    ] = 300,
211) -> None:
212    """Plot a velocity model from a velocity model params yaml file (vm_params.yaml).
213
214    Parameters
215    ----------
216    vm_params_ffp : Path
217        Path to the YAML file containing velocity model parameters.
218    output_plot_path : Path
219        Path to save the generated plot.
220    velocity_model_ffp : Path | None
221        Path to the directory containing velocity model files. If None, no velocity model is overlaid.
222    component : VelocityModelComponent, default = `VelocityModelComponent.p_wave`
223        The velocity model component to overlay.
224    slice : int, optional, default = 0
225        The z-level slice of the velocity model to plot.
226    transparency : int, optional, default = 80
227        Transparency level for the velocity model overlay.
228    title : str, optional
229        Title of the figure.
230    latitude_pad : float, optional, default = 0
231        Padding for latitude (in degrees) for the figure.
232    longitude_pad : float, optional, default = 0
233        Padding for longitude (in degrees) for the figure.
234    dpi : float, optional, default = 300
235        DPI for the output plot.
236    """
237    with open(vm_params_ffp) as vm_params_handle:
238        vm_params = yaml.safe_load(vm_params_handle)
239
240        # Extent x and extent y are swapped in meaning between the old
241        # vm params and the new bounding box class. So we swap them
242        # here when we load them.
243        origin = np.array([vm_params["MODEL_LAT"], vm_params["MODEL_LON"]])
244        great_circle_bearing = vm_params["MODEL_ROT"]
245        extent_x = vm_params["extent_x"]
246        extent_y = vm_params["extent_y"]
247        nztm_bearing = coordinates.great_circle_bearing_to_nztm_bearing(
248            origin, extent_x / 2, great_circle_bearing
249        )
250        box = bounding_box.BoundingBox.from_centroid_bearing_extents(
251            origin, nztm_bearing, extent_x, extent_y
252        )
253
254        resolution = vm_params.get("hh")
255        nx = vm_params["nx"]
256        ny = vm_params["ny"]
257        nz = vm_params["nz"]
258
259    region = (
260        box.corners[:, 1].min() - longitude_pad,
261        box.corners[:, 1].max() + longitude_pad,
262        box.corners[:, 0].min() - latitude_pad,
263        box.corners[:, 0].max() + latitude_pad,
264    )
265    fig = plotting.gen_region_fig(title, region)
266    if velocity_model_ffp:
267        velocity_model = load_velocity_model_file(
268            velocity_model_ffp, component, nx, ny, nz
269        )
270        plot_velocity_model(
271            fig,
272            box,
273            velocity_model,
274            resolution,
275            slice,
276            nx,
277            ny,
278            region,
279            cb_label=CB_LABEL_MAP[component],
280            transparency=transparency,
281        )
282    fig.plot(
283        x=box.corners[:, 1],
284        y=box.corners[:, 0],
285        close=True,
286        pen="1p,black,-",
287    )
288    fig.savefig(
289        output_plot_path,
290        dpi=dpi,
291        anti_alias=True,
292    )
293
294
295if __name__ == "__main__":
296    app()