diff --git a/requirements.txt b/requirements.txt index fdfb8cd..a70d53b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,7 +10,7 @@ qcore @ git+https://github.com/ucgmsim/qcore.git geopandas pandas shapely -scipy +scipy >= 1.15.0 pooch pytest pytest-cov diff --git a/visualisation/plot_ts.py b/visualisation/plot_ts.py index daea48e..447592a 100644 --- a/visualisation/plot_ts.py +++ b/visualisation/plot_ts.py @@ -1,6 +1,7 @@ """Create simulation video of surface ground motion levels.""" import functools +import multiprocessing as mp import os import shutil import subprocess @@ -10,23 +11,23 @@ import cartopy.crs as ccrs import cartopy.feature as cfeature +import cartopy.io.img_tiles as cimgt import matplotlib -import shapely matplotlib.use("Agg") -import multiprocessing as mp - -import cartopy.io.img_tiles as cimgt import matplotlib.colors as mcolors import matplotlib.pyplot as plt import numpy as np +import shapely import tqdm import typer +from matplotlib.animation import FFMpegWriter, FuncAnimation from qcore import cli, coordinates from qcore.xyts import XYTSFile -from workflow.realisations import SourceConfig +from source_modelling import srf +from workflow.realisations import DomainParameters, SourceConfig app = typer.Typer() @@ -280,14 +281,17 @@ def zoom_extents( return new_x_min, new_x_max, new_y_min, new_y_max -def xyts_waveform_coordinates(xyts_file: XYTSFile) -> np.ndarray: +def waveform_coordinates(nztm_corners: np.ndarray, nx: int, ny: int) -> np.ndarray: """Compute gridpoint coordinates for XYTS waveform. Parameters ---------- - xyts_file : XYTSFile - The xyts file containing gridded data. - + nztm_corners : np.ndarray + The corners of the waveform grid. + nx : int + The number of x-points in the output grid. + ny : int + The number of y-points in the output grid. Returns ------- @@ -295,12 +299,7 @@ def xyts_waveform_coordinates(xyts_file: XYTSFile) -> np.ndarray: A numpy array of shape (2 x ny x nx) containing the x and y coordinates of gridpoints in the NZTM coordinate system. """ - corners_geo = np.array(xyts_file.corners()) - nztm_corners = coordinates.wgs_depth_to_nztm(corners_geo[:, ::-1]) - - norm_xi, norm_eta = np.meshgrid( - np.linspace(0, 1, xyts_file.nx), np.linspace(0, 1, xyts_file.ny) - ) + norm_xi, norm_eta = np.meshgrid(np.linspace(0, 1, nx), np.linspace(0, 1, ny)) origin = nztm_corners[0] # Bottom-left corner (x0, y0) in NZTM vec_x = nztm_corners[1] - origin # Vector along xi axis (bottom edge) in NZTM vec_y = nztm_corners[3] - origin # Vector along eta axis (left edge) in NZTM @@ -330,7 +329,7 @@ def render_single_frame( title: str | None, width: float, height: float, - dpi: float, + dpi: int, ) -> str: """Render a single frame of the animation. @@ -368,7 +367,7 @@ def render_single_frame( The width of the figure in cm. height : float The height of the figure in cm. - dpi : float + dpi : int The DPI for the figure. Returns @@ -468,7 +467,7 @@ def render_single_frame( return frame_filename -@cli.from_docstring(app) +@cli.from_docstring(app, name="xyts") def animate_low_frequency_mpl_nztm( realisation_ffp: Annotated[Path, typer.Argument(exists=True, dir_okay=False)], xyts_ffp: Annotated[Path, typer.Argument(exists=True, dir_okay=False)], @@ -483,8 +482,8 @@ def animate_low_frequency_mpl_nztm( frame_count: Annotated[int | None, typer.Option()] = None, width: Annotated[float, typer.Option()] = 30.0, height: Annotated[float, typer.Option()] = 30.0, - dpi: Annotated[float, typer.Option()] = 150.0, - fps: Annotated[float, typer.Option()] = 15.0, + dpi: Annotated[int, typer.Option()] = 150, + fps: Annotated[int, typer.Option()] = 15, title: Annotated[str | None, typer.Option()] = None, zoom: Annotated[float, typer.Option()] = 1, simple_map: Annotated[bool, typer.Option()] = False, @@ -516,10 +515,10 @@ def animate_low_frequency_mpl_nztm( The width of the figure in cm, by default 30. height : float, optional The height of the figure in cm, by default 30. - dpi : float, optional + dpi : int, optional The DPI for the figure, by default 150.0. - fps : float, optional - The frames per second for the animation, by default 15.0. + fps : int, optional + The frames per second for the animation, by default 15. title : str | None, optional The title for the animation, by default None (no title). zoom : float, optional @@ -538,9 +537,6 @@ def animate_low_frequency_mpl_nztm( ) raise typer.Exit(code=1) - if dpi % 2: - dpi += 1 - source_config = SourceConfig.read_from_realisation(realisation_ffp) xyts_file = XYTSFile(xyts_ffp) @@ -562,7 +558,7 @@ def animate_low_frequency_mpl_nztm( ) frame_count = frame_count or xyts_file.nt - xr, yr = xyts_waveform_coordinates(xyts_file) + xr, yr = waveform_coordinates(nztm_corners, xyts_file.nx, xyts_file.ny) with tempfile.TemporaryDirectory() as temp_dir: render_frame = functools.partial( @@ -623,3 +619,248 @@ def animate_low_frequency_mpl_nztm( ] subprocess.run(ffmpeg_cmd, check=True) + + +def non_zero_data_points( + x: np.ndarray, y: np.ndarray, z: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Get the non-zero data points in the 3D array. + + Parameters + ---------- + x : np.ndarray + The x coordinates of the data points. + y : np.ndarray + The y coordinates of the data points. + z : np.ndarray + The z coordinates of the data points. + + Returns + ------- + tuple[np.ndarray, np.ndarray, np.ndarray] + The non-zero data points in the 3D array. + """ + mask = z > 0 + return x[mask], y[mask], z[mask] + + +@cli.from_docstring(app, name="srf") +def animate_srf_slip_times( + realisation_ffp: Annotated[Path, typer.Argument(exists=True, dir_okay=False)], + srf_ffp: Annotated[Path, typer.Argument(exists=True, dir_okay=False)], + output_mp4: Annotated[ + Path, typer.Argument(writable=True, dir_okay=False, resolve_path=True) + ], + max_slip: Annotated[float, typer.Option()] = 10.0, + padding: Annotated[float, typer.Option()] = 5.0, + cmap: Annotated[str, typer.Option()] = "hot", + scale: Annotated[str, typer.Option()] = "10m", + frame_count: Annotated[int | None, typer.Option()] = None, + width: Annotated[float, typer.Option()] = 30.0, + height: Annotated[float, typer.Option()] = 30.0, + dpi: Annotated[int, typer.Option()] = 150.0, + fps: Annotated[int, typer.Option()] = 15, + title: Annotated[str | None, typer.Option()] = None, + zoom: Annotated[float, typer.Option()] = 1, + simple_map: Annotated[bool, typer.Option()] = False, + map_quality: Annotated[int, typer.Option()] = 4, + frame_dt: Annotated[int, typer.Option(min=0)] = 20, +) -> None: + """Render SRF slip times as a 2D video. + + Parameters + ---------- + realisation_ffp : Path + The input realisation file. + srf_ffp : Path + The input srf file containing the simulation data. + output_mp4 : Path + The output file path for the generated animation. + max_slip : float, optional + The slip (not ground motion) for color scaling, by default 10.0 cm. + padding : float, optional + The padding in km for the map extent, by default 5.0. + cmap : str, optional + The colormap to use for the animation, by default "hot". + scale : str, optional + The scale for cartopy features, by default "10m". + frame_count : int | None, optional + The number of frames to display in the animation, by default None (uses all frames). + width : float, optional + The width of the figure in cm, by default 30. + height : float, optional + The height of the figure in cm, by default 30. + dpi : int, optional + The DPI for the figure, by default 150.0. + fps : int, optional + The frames per second for the animation, by default 15.0. + title : str | None, optional + The title for the animation, by default None (no title). + zoom : float, optional + Zoom factor for the map, by default 1.0, on a log-scale. Zoom + centres on centre of source geometry. + simple_map : bool, optional + If True, disable OpenStreetMap background and use a simple map. + map_quality : int, optional + The quality of the map, by default 4. Has no effect if using a + simple map. Lower values have lower quality but render faster. + frame_dt : int, optional + The number of timeslices per dt-step, default is 20. + """ + ffmpeg = shutil.which("ffmpeg") + if not ffmpeg: + print( + "You must have ffmpeg installed. See https://ffmpeg.org/download.html.", + ) + raise typer.Exit(code=1) + + if dpi % 2: + dpi += 1 + + source_config = SourceConfig.read_from_realisation(realisation_ffp) + domain_config = DomainParameters.read_from_realisation(realisation_ffp) + srf_file = srf.read_srf(srf_ffp) + + nztm_corners = coordinates.wgs_depth_to_nztm(domain_config.domain.corners)[:, ::-1] + slip = srf_file.slipt1_array.tocsc() + map_extent_nztm = map_extents(nztm_corners, padding) + + if zoom != 1: + centre = shapely.centroid( + shapely.union_all( + [fault.geometry for fault in source_config.source_geometries.values()] + ) + ) + map_extent_nztm = zoom_extents( + map_extent_nztm, + (centre.y, centre.x), + zoom, + ) + + frame_count = frame_count or srf_file.nt + + # Create figure and initial setup + cm = 1 / 2.54 + fig = plt.figure(figsize=(width * cm, height * cm)) + ax = fig.add_subplot(1, 1, 1, projection=NZTM_CRS) + ax.set_extent(map_extent_nztm, crs=NZTM_CRS) + + # Add time text + + time_text = ax.text( + 0.98, + 0.02, + "Time: 0s", + transform=ax.transAxes, + fontsize=12, + color="black", + fontweight="bold", + ha="right", + va="bottom", + bbox={"boxstyle": "round", "facecolor": "white", "alpha": 0.8}, + ) + + if simple_map: + plot_cartographic_features(ax, scale) + plot_towns(ax, map_extent_nztm) + else: + request = cimgt.OSM(cache=True) + request._MAX_THREADS = ( + 1 # Limit to one thread because it is in a multiprocess pool. + ) + ax.add_image( + request, + 10, + interpolation="spline36", + regrid_shape=map_quality * 1000, + zorder=0, + ) + + ax.add_geometries( + [shapely.Polygon(nztm_corners)], + facecolor="none", + edgecolor="black", + linestyle="--", + zorder=1, + crs=NZTM_CRS, + ) + + ax.add_geometries( + [ + shapely.transform(fault.geometry, lambda coords: coords[:, ::-1]) + for fault in sorted( + source_config.source_geometries.values(), + key=lambda fault: -fault.centroid[-1], + ) + ], + facecolor="red", + edgecolor="black", + zorder=2, + crs=NZTM_CRS, + ) + + if title: + fig.suptitle(title, fontsize=16) + coords = coordinates.wgs_depth_to_nztm(srf_file.points[["lat", "lon"]].values)[ + :, ::-1 + ] + x, y = coords[:, 0], coords[:, 1] + init_x, init_y, init_z = non_zero_data_points(x, y, slip[:, 0].todense()) + scat = ax.scatter( + init_x, + init_y, + c=init_z, + cmap=cmap, + vmin=0, + vmax=max_slip, + transform=NZTM_CRS, + zorder=100, + ) + fig.colorbar( + scat, + ax=ax, + orientation="vertical", + pad=0.02, + aspect=30, + shrink=0.8, + label="Slip (cm)", + ) + + def initial_frame() -> None: # numpydoc ignore=GL08 + time_text.set_text("Time: 0s") + return [scat, time_text] + + # Setup the animation function + def render_single_frame( + frame_index: int, + ) -> list: # numpydoc ignore=GL08 + # Create a new figure for this frame + slip_index = frame_index * frame_dt + slip_end = min(slip_index + frame_dt, srf_file.nt) + interval_slip_mean = slip[:, list(range(slip_index, slip_end))].mean(axis=1) + # Add the actual data for this frame + cur_x, cur_y, z = non_zero_data_points( + x, + y, + interval_slip_mean, + ) + scat.set_offsets(np.c_[cur_x, cur_y]) + scat.set_array(z) + time_text.set_text(f"Time: {slip_index * srf_file.dt:.2f} s") + return [scat, time_text] + + # Create the animation + anim = FuncAnimation( + fig, + render_single_frame, + init_func=initial_frame, + frames=tqdm.trange( + frame_count // frame_dt, desc="Rendering frames", unit="frame" + ), + blit=True, + ) + + # Save the animation + writer = FFMpegWriter(fps=fps) + anim.save(output_mp4, writer=writer) + plt.close(fig)