source_modelling.scripts.plot_srf

  1"""Plot multi-segment rupture with slip."""
  2
  3from pathlib import Path
  4from typing import Annotated, Optional
  5
  6import numpy as np
  7import typer
  8
  9from pygmt_helper import plotting
 10from qcore import coordinates
 11from source_modelling import srf
 12
 13app = typer.Typer()
 14
 15
 16@app.command(help="Plot multi-segment rupture with slip.")
 17def plot_srf(
 18    srf_ffp: Annotated[
 19        Path, typer.Argument(help="Path to SRF file to plot.", exists=True)
 20    ],
 21    output_ffp: Annotated[
 22        Path, typer.Argument(help="Output plot image.", dir_okay=False)
 23    ],
 24    dpi: Annotated[
 25        float, typer.Option(help="Plot output DPI (higher is better)")
 26    ] = 300,
 27    title: Annotated[Optional[str], typer.Option(help="Plot title to use")] = None,
 28    levels: Annotated[
 29        float,
 30        typer.Option(
 31            help="Plot time as contours of every LEVELS seconds", metavar="LEVELS"
 32        ),
 33    ] = 1,
 34    realisation_ffp: Annotated[
 35        Optional[Path],
 36        typer.Option(
 37            help="Path to realisation, used to mark jump points.", exists=True
 38        ),
 39    ] = None,
 40    latitude_pad: Annotated[
 41        float, typer.Option(help="Latitude padding to apply (degrees)")
 42    ] = 0,
 43    longitude_pad: Annotated[
 44        float, typer.Option(help="longitude padding to apply (degrees)")
 45    ] = 0,
 46    annotations: Annotated[bool, typer.Option(help="Label contours")] = True,
 47) -> None:
 48    """Plot multi-segment rupture with slip.
 49
 50    Parameters
 51    ----------
 52    srf_ffp : Path
 53        Path to SRF file to plot.
 54    output_ffp : Path
 55        Output plot image.
 56    dpi : float
 57        Plot output DPI (higher is better).
 58    title : Optional[str]
 59        Plot title to use
 60    levels : float
 61        Plot time as contours of every `levels` seconds.
 62    realisation_ffp : Optional[Path]
 63        Path to realisation, used to mark jump points.
 64    latitude_pad : float
 65        Latitude padding to apply (degrees).
 66    longitude_pad : float
 67        Longitude padding to apply (degrees).
 68    annotations : bool
 69        Label contours.
 70    """
 71    srf_data = srf.read_srf(srf_ffp)
 72
 73    region = (
 74        srf_data.points["lon"].min() - longitude_pad,
 75        srf_data.points["lon"].max() + longitude_pad,
 76        srf_data.points["lat"].min() - latitude_pad,
 77        srf_data.points["lat"].max() + latitude_pad,
 78    )
 79
 80    # Compute slip limits
 81    slip_quantile = srf_data.points["slip"].quantile(0.98)
 82    slip_cb_max = max(int(np.round(slip_quantile, -1)), 10)
 83    cmap_limits = (0, slip_cb_max, slip_cb_max / 10)
 84
 85    fig = plotting.gen_region_fig(title, region=region, map_data=None)
 86
 87    for (_, segment), segment_points in zip(
 88        srf_data.header.iterrows(), srf_data.segments
 89    ):
 90        nstk = segment["nstk"]
 91        ndip = segment["ndip"]
 92
 93        # Create standard slip heatmap.
 94        cur_grid = plotting.create_grid(
 95            segment_points,
 96            "slip",
 97            grid_spacing="5e/5e",
 98            region=(
 99                segment_points["lon"].min(),
100                segment_points["lon"].max(),
101                segment_points["lat"].min(),
102                segment_points["lat"].max(),
103            ),
104            set_water_to_nan=False,
105        )
106
107        plotting.plot_grid(
108            fig,
109            cur_grid,
110            "hot",
111            cmap_limits,
112            ("white", "black"),
113            transparency=0,
114            reverse_cmap=True,
115            plot_contours=False,
116            cb_label="slip",
117            continuous_cmap=True,
118        )
119
120        # Plot time contours
121        time_grid = plotting.create_grid(
122            segment_points,
123            "tinit",
124            grid_spacing="5e/5e",
125            region=(
126                segment_points["lon"].min(),
127                segment_points["lon"].max(),
128                segment_points["lat"].min(),
129                segment_points["lat"].max(),
130            ),
131            set_water_to_nan=False,
132        )
133        fig.grdcontour(levels=1, grid=time_grid, pen="0.1p")
134
135        # Plot bounds of the current segment.
136        corners = segment_points.iloc[[0, nstk - 1, -1, (ndip - 1) * nstk]]
137        fig.plot(
138            x=corners["lon"].iloc[list(range(len(corners))) + [0]].to_list(),
139            y=corners["lat"].iloc[list(range(len(corners))) + [0]].to_list(),
140            pen="0.5p,black,-",
141        )
142        fig.plot(
143            x=corners["lon"].iloc[:2].to_list(),
144            y=corners["lat"].iloc[:2].to_list(),
145            pen="0.8p,black",
146        )
147
148        if not annotations:
149            continue
150
151        # Custom annotation implementation. Rough algorithm is:
152        # 1. Compute the number of whole second increments in tinit
153        tinit_max = int(np.round(segment_points["tinit"].max()))
154        tinit_min = int(np.round(segment_points["tinit"].min()))
155        # 2. Provided we have a non-trivial number of increments
156        if tinit_max - tinit_min >= 1:
157            # 3. Iterate over every second between the two and
158            for j in range(tinit_min, tinit_max):
159                # 4. Find the segment point with the closest tinit value to the current second.
160                min_delta = (segment_points["tinit"] - j).abs().min()
161                # 5. Provided that this point has a tinit closest enough to the current second.
162                if min_delta < 0.1:
163                    # 6. Get that point and
164                    closest_point = segment_points[
165                        (segment_points["tinit"] - j).abs() == min_delta
166                    ].iloc[0]
167                    # Add an annotation label to the plot.
168                    fig.text(
169                        text=f"{j}",
170                        x=closest_point["lon"],
171                        y=closest_point["lat"],
172                        font="5p",
173                        fill="white",
174                    )
175
176    # Plot the hypocentre.
177    hypocentre = srf_data.points[
178        srf_data.points["tinit"] == srf_data.points["tinit"].min()
179    ].iloc[0]
180    fig.plot(
181        x=hypocentre["lon"],
182        y=hypocentre["lat"],
183        style="a0.4c",
184        pen="1p,black",
185        fill="white",
186    )
187
188    # If we are supplied a JSON realisation, we can add labels for jump points.
189    if realisation_ffp: # pragma: no cover
190        # NOTE: this import is here because the workflow is, as yet,
191        # not ready to be installed along-side source modelling.
192        from workflow.realisations import RupturePropagationConfig, SourceConfig
193
194        rupture_propagation_config = RupturePropagationConfig.read_from_realisation(
195            realisation_ffp
196        )
197        source_config = SourceConfig.read_from_realisation(realisation_ffp)
198        for fault_name, jump_point in rupture_propagation_config.jump_points.items():
199            parent_name = rupture_propagation_config.rupture_causality_tree[fault_name]
200            if not parent_name:
201                continue
202            source = source_config.source_geometries[fault_name]
203            parent = source_config.source_geometries[parent_name]
204
205            # Ruptures jump from_point --> to_point
206            from_point = parent.fault_coordinates_to_wgs_depth_coordinates(
207                jump_point.from_point
208            )
209
210            # Find the closest point to the theoretical jump point (so we can lookup the time).
211            closest_from_point_distance_idx = (
212                coordinates.distance_between_wgs_depth_coordinates(
213                    srf_data.points[["lat", "lon", "dep"]].to_numpy()
214                    * np.array([1, 1, 1000]),
215                    from_point,
216                )
217            ).argmin()
218            srf_jump_point = srf_data.points.iloc[closest_from_point_distance_idx]
219
220            to_point = source.fault_coordinates_to_wgs_depth_coordinates(
221                jump_point.to_point
222            )[:2]
223
224            fig.plot(
225                x=from_point[1],
226                y=from_point[0],
227                style="t0.4c",
228                pen="1p,black",
229                fill="white",
230            )
231            fig.text(
232                x=srf_jump_point["lon"],
233                y=srf_jump_point["lat"] - 0.01,
234                font="5p",
235                fill="white",
236                text=f"t_jump = {srf_jump_point['tinit']:.2f}",
237            )
238            fig.plot(
239                x=to_point[1],
240                y=to_point[0],
241                style="i0.4c",
242                pen="1p,black",
243                fill="white",
244            )
245
246    fig.savefig(
247        output_ffp,
248        dpi=dpi,
249        anti_alias=True,
250    )
251
252
253if __name__ == "__main__":
254    app()