Stoch file plotting tool

pyproject.toml

@@ -20,6 +20,7 @@ plot-srf-rise = "visualisation.sources.plot_rise:app"
 plot-mw-contributions = "visualisation.sources.plot_mw_contributions:app"
 plot-slip-rise-rake = "visualisation.sources.plot_slip_rise_rake:app"
 plot-srf-distribution = "visualisation.sources.plot_srf_distribution:app"
+plot-stoch = "visualisation.sources.plot_stoch:app"
 
 [tool.setuptools.package-dir]
 visualisation = "visualisation"

tests/stoch/realisation.stoch

@@ -0,0 +1,111 @@
+7
+  172.1811   -43.5095     4     4     2.00     2.00
+ 150   54  255     1.40     1.83     5.77
+    59   139   157   200
+   195   153   103   223
+   221   161   159   217
+   284   179   312   115
+  0.58  0.44  0.05  0.38
+  0.80  0.66  0.39  0.70
+  0.56  0.55  0.48  0.65
+  0.45  0.33  0.33  0.34
+  5.09  4.02  3.70  3.62
+  3.38  2.78  2.00  1.82
+  2.65  1.77  0.78  0.94
+  2.28  1.59  0.72  1.04
+  172.1255   -43.5728     4     5     2.00     2.00
+  35   70   39     0.97  -999.00  -999.00
+    73   159   165   224
+   101    92   114   118
+   253   200    83   281
+   204   140   208   177
+   270   177   276    94
+  0.32  0.60  0.73  0.61
+  0.68  0.66  0.61  0.71
+  0.72  0.66  0.64  0.68
+  0.39  0.38  0.44  0.47
+  0.22  0.28  0.39  0.30
+  7.63  6.63  5.88  5.19
+  6.16  5.20  4.37  3.63
+  4.98  4.32  3.64  2.31
+  4.71  4.26  3.43  2.91
+  4.73  4.45  3.76  4.02
+  172.0362   -43.5808     8     5     2.00     2.00
+ 303   75  127     0.48  -999.00  -999.00
+   171   151   160   151   191   321   237   177
+    29    81   223   278   179   295   281   254
+   109   164   126   199   169   333    89    56
+   292   234   145    97   236   223   176   348
+   268   401   129   129    74    92   227   163
+  0.85  0.75  0.50  0.65  0.93  0.95  0.83  0.89
+  0.71  0.55  0.92  0.98  0.78  0.67  0.88  0.86
+  0.60  0.88  0.84  1.03  0.93  1.10  0.77  0.42
+  0.50  0.48  0.37  0.51  0.65  0.59  0.32  0.55
+  0.49  0.67  0.46  0.33  0.40  0.52  0.46  0.38
+  9.24  9.95 10.67 11.54 12.46 13.01 13.92 14.58
+  7.25  7.89  8.64  9.33 10.12 10.90 11.95 12.50
+  6.08  6.93  8.08  8.61  9.43 10.12 11.26 11.93
+  6.40  7.11  7.73  8.51  8.93  9.86 10.31 10.62
+  7.23  7.21  8.22  8.67  9.31  9.85 10.21 10.97
+  172.1871   -43.5920     9     5     2.00     2.00
+  86   80   55     0.49  -999.00  -999.00
+    99   351   291   152   356   304   264   399   202
+   349   206   196   204   375   161   155   246   171
+   132   214    70   218    95   172   350   101   208
+   106   152   169   135    72   185   337   198   324
+   177   273   120    91   172   185   239   156   171
+  0.35  0.90  1.11  1.04  1.04  0.99  1.09  1.27  1.05
+  0.48  0.53  0.76  0.54  1.19  0.60  0.63  1.00  1.00
+  0.60  0.88  0.75  0.67  0.32  0.40  0.73  0.99  0.83
+  0.53  0.69  0.67  0.61  0.37  0.23  0.66  0.64  0.56
+  0.58  0.54  0.55  0.34  0.36  0.38  0.37  0.59  0.52
+ 10.20  9.92 10.66 11.88 12.64 13.78 14.69 15.37 16.56
+  7.87  8.11  9.07  9.74 10.69 12.18 13.03 13.80 14.22
+  7.79  7.51  8.66  9.41 10.74 11.36 11.62 12.91 13.38
+  8.61  8.18  8.60  9.25 10.34 10.84 11.29 12.11 12.59
+  8.67  8.65  9.32  9.62 10.22 10.82 11.37 11.94 12.53
+  172.3556   -43.5759     6     4     2.00     2.00
+  86   78   65     0.49  -999.00  -999.00
+   191   155   180   250   149   196
+   121   120   105   164   177   252
+   261   270   274   185   295   278
+   138   209   115   117   267   113
+  0.52  0.60  0.81  0.81  0.74  0.41
+  0.74  0.74  0.67  0.57  0.65  0.92
+  0.69  0.74  0.65  0.46  0.56  0.77
+  0.38  0.31  0.23  0.28  0.26  0.26
+ 17.58 17.87 18.29 18.65 19.42 19.96
+ 15.43 15.47 16.06 16.39 16.97 17.59
+ 13.71 14.21 14.50 15.07 15.69 16.66
+ 13.37 13.50 14.44 14.88 15.35 16.52
+  172.3115   -43.5509     6     5     2.00     2.00
+  40   80   64     1.49  -999.00  -999.00
+   100   226   271   142   169    96
+   292   174    51   161   312   205
+   124   213   144   176   288   173
+   151   195   192   158   110   256
+   146    37    68   203   141   213
+  0.77  0.65  0.68  0.60  0.54  0.72
+  1.12  0.78  0.62  0.55  0.83  1.01
+  0.89  0.60  0.48  0.55  0.76  0.59
+  0.40  0.49  0.46  0.38  0.41  0.48
+  0.32  0.32  0.27  0.33  0.22  0.27
+ 17.67 17.57 18.15 19.03 19.37 19.95
+ 16.13 16.37 17.20 17.57 17.69 18.52
+ 15.43 15.47 16.23 16.77 16.82 17.73
+ 14.60 15.14 15.56 16.25 16.78 17.23
+ 14.25 15.11 15.70 16.05 16.58 17.19
+  171.9397   -43.5815     3     4     2.00     2.00
+ 216   50   79     0.88  -999.00  -999.00
+   222   171   298
+   179   208   202
+   220   312   347
+   206   126   163
+  0.65  0.55  0.50
+  0.44  0.44  0.30
+  0.28  0.39  0.31
+  0.35  0.33  0.25
+ 18.93 21.58 23.79
+ 20.41 21.38 22.76
+ 21.55 22.24 22.78
+ 22.46 23.04 23.86

tests/test_plots.py

@@ -13,10 +13,12 @@
     plot_srf_cumulative_moment,
     plot_srf_distribution,
     plot_srf_moment,
+    plot_stoch,
 )
 
 PLOT_IMAGE_DIRECTORY = Path("wiki/images")
 SRF_FFP = Path(__file__).parent / "srfs" / "rupture_1.srf"
+STOCH_FFP = Path(__file__).parent / "stoch"
 MULTI_SUMMARY_SRF_FFP = Path(__file__).parent / "srfs" / "nevis.srf"
 REALISATION_FFP = Path(__file__).parent / "srfs" / "realisation.json"
 
@@ -56,7 +58,6 @@ def test_plot_functions(
 
 
 def test_plot_mw_contributions(tmp_path: Path):
-
     original = PLOT_IMAGE_DIRECTORY / "example_mw_contributions.png"
     generated = tmp_path / "output.png"
 
@@ -100,3 +101,20 @@ def test_plot_slip_rise_rake_segment(tmp_path: Path):
 
     diff = diffimg.diff(original, output_image_path)
     assert diff <= 0.05
+
+
+def test_plot_stoch(tmp_path: Path):
+    output_image_path = tmp_path / "output.png"
+    original = PLOT_IMAGE_DIRECTORY / "stoch_example.png"
+
+    plot_stoch.plot_stoch_to_file(
+        STOCH_FFP / "realisation.stoch",
+        output_image_path,
+        width=60,
+        height=20,
+        dpi=300,
+        title="Stoch file",
+    )
+
+    diff = diffimg.diff(original, output_image_path)
+    assert diff <= 0.05

visualisation/sources/plot_stoch.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""Plot multi-segment rupture stoch file with slip."""
+
+from pathlib import Path
+from typing import Annotated, Optional
+
+import numpy as np
+import typer
+from matplotlib import pyplot as plt
+
+from qcore import cli
+from source_modelling import stoch
+from visualisation import utils
+
+app = typer.Typer()
+
+
+def plot_stoch(
+    stoch_data: stoch.StochFile,
+    width: float = 10,
+    height: float = 10,
+) -> tuple[plt.Figure, list[list[plt.Axes]] | plt.Axes]:
+    """Plot multi-segment rupture with slip.
+
+    The plot is a heatmap of the slip distribution of the stoch file.
+    The heatmap is labelled with the length along the x-axis and the
+    width along the y-axis (both in kilometres). The heatmap is
+    coloured with a reverse hot colourmap (like the SRF segment
+    plots). The text labels are coloured white for high values and
+    black for low values. See the examples in the wiki.
+
+    Parameters
+    ----------
+    stoch_data : StochFile
+        Stoch file to plot.
+    width : float, optional
+        Width of plot (in cm). Default is 10.
+    height : float, optional
+        Height of plot (in cm). Default is 10.
+
+    Returns
+    -------
+    plt.Figure
+        Figure of the plot.
+    list[list[plt.Axes]] | plt.Axes
+        Axes of the plot. Will be a list if there are multiple segments.
+
+    Examples
+    --------
+    >>> stoch_data = stoch.StochFile("tests/stochs/rupture_1.stoch")
+    >>> fig, axes = plot_stoch(stoch_data, width=10, height=10)
+    >>> # The above code would plot the slip distribution of the stoch file 'rupture_1.stoch'.
+    >>> # The plot will have a width of 10 cm and a height of 10 cm.
+    """
+    cm = 1 / 2.54
+    rows = int(np.sqrt(len(stoch_data.data)))
+    cols = int(np.ceil(len(stoch_data.data) / rows))
+    fig, axes = plt.subplots(rows, cols, figsize=(width * cm, height * cm))
+    for i, (ax, plane_data, slip) in enumerate(
+        zip(axes.ravel(), stoch_data.data, stoch_data.slip)
+    ):
+        dx = plane_data.header.dx
+        dy = plane_data.header.dy
+        length = dx * slip.shape[1]
+        width = dy * slip.shape[0]
+        description = utils.format_description(slip, compact=True, units="cm")
+        ax.set_title(f"Segment {i + 1}\n{description}")
+
+        # Plot slip array as a heatmap labelled with length along the x-axis and width along the y-axis.
+        ax.set_ylim(width, 0)
+        ax.imshow(
+            slip[::-1],
+            cmap="hot_r",
+            extent=[0, length, 0, width],
+        )
+        ax.set_xlabel("Length (km)")
+        ax.set_ylabel("Width (km)")
+
+        for j, k in np.ndindex(slip.shape):
+            # Add text labels to the heatmap, use white text for high values, and black text for low values.
+            colour = "black" if slip[j, k] < np.median(slip) else "white"
+            ax.text(
+                k * dx + dx / 2,
+                (j * dy + dy / 2),
+                f"{int(slip[j,  k])}",
+                ha="center",
+                va="center",
+                color=colour,
+            )
+
+    # empty the unused axes
+    for ax in axes.ravel()[len(stoch_data.data) :]:
+        ax.axis("off")
+
+    return fig, axes
+
+
+@cli.from_docstring(app)
+def plot_stoch_to_file(
+    stoch_ffp: Annotated[Path, typer.Argument(exists=True, dir_okay=False)],
+    output_ffp: Annotated[Path, typer.Argument(dir_okay=False)],
+    width: Annotated[float, typer.Option()] = 10,
+    height: Annotated[float, typer.Option()] = 10,
+    dpi: Annotated[int, typer.Option()] = 300,
+    title: Annotated[Optional[str], typer.Option()] = None,
+) -> None:
+    """Plot multi-segment rupture with slip.
+
+    Parameters
+    ----------
+    stoch_ffp : Path
+        Path to stoch file to plot.
+    output_ffp : Path
+        Output plot image.
+    width : float
+        Width of plot (in cm).
+    height : float
+        Width of plot (in cm).
+    dpi : int
+        Plot output DPI (higher is better).
+    title : Optional[str]
+        Plot title to use.
+
+
+    Examples
+    --------
+    >>> plot_stoch_to_file(
+    ...     stoch_ffp="tests/stochs/rupture_1.stoch",
+    ...     output_ffp="slip_plot.png",
+    ...     dpi=300,
+    ...     title="Rupture Slip Distribution",
+    ...     latitude_pad=0.5,
+    ...     longitude_pad=0.5,
+    ...     annotations=True,
+    ...     width=15,
+    ...     show_inset=True,
+    ... )
+    >>> # The above code would plot the slip distribution of the stoch file 'rupture_1.stoch'
+    >>> # and save it as 'slip_plot.png'.
+    >>> # The plot will have a DPI of 300.
+    >>> # The plot will have the title "Rupture Slip Distribution".
+    >>> # The plot will have jump points marked from the realisation file 'realisation.json'.
+    >>> # The plot will have a latitude and longitude padding of 0.5 degrees.
+    >>> # The plot will have annotations of slip times and an inset map.
+    """
+    stoch_data = stoch.StochFile(stoch_ffp)
+    fig, axes = plot_stoch(stoch_data, width, height)
+    if title:
+        fig.suptitle(title)
+    fig.tight_layout()
+    fig.savefig(output_ffp, dpi=dpi)
+
+
+if __name__ == "__main__":
+    app()

wiki/Sources.md

@@ -63,6 +63,16 @@ See the help text to find more formatting options. You will probably be interest
 ### With an Inset
 It often helps to provide an overview map that helps the reader know where you rupture is occurring relative to the whole country. To see one, pass the `--show-inset` flag to `plot-srf`.
 
+# Plotting Stoch Files
+Slip models are downsampled at a low resolution and provided to the high frequency simulation code. These stochastic input files (`.stoch` files) contain most of the same information as the SRF files. To plot a stoch file, run
+
+``` bash
+$ plot-stoch STOCH_FILE OUTPUT_FILE --title 'Stoch file' --width 40 --height 20
+```
+Adjust the title, height and width to best suit your plots. Each segment of the stoch file is plotted, with the slip statistics and slip in each patch.
+
+![](images/stoch_example.png)
+
 # How Do I Plot a Moment Rate Function?
 
 The tool for this job is `plot-srf-moment`. To plot the SRF moment for a given SRF file type