')
-
- quality_old_record_index = quality_old.set_index("record_id")
- quality_new_record_index = quality_new.set_index("record_id")
+ img_base64_old = plot_usable_period_records(full_old, quality_old, "Old NZGMDB")
+ html_parts.append("
")
+ html_parts.append(f'
')
+ html_parts.append(f'
')
+ html_parts.append("
")
+ else:
+ html_parts.append('
')
+ html_parts.append(f'
')
+ html_parts.append("
")
- # IM Compare
- shared_record_ids = np.intersect1d(
- quality_old_record_index.index.values.astype(str),
- quality_new_record_index.index.values.astype(str),
- )
+ if compare_version_directory:
+ # Add IM Compare
+ html_parts.append("
IM Comparison
")
+ html_parts.append('
')
- plot_ims = [
- "PGV",
- "PGA",
- "pSA_0.01",
- "pSA_0.1",
- "pSA_0.5",
- "pSA_1.0",
- "pSA_3.0",
- "pSA_10.0",
- ]
+ quality_old_record_index = quality_old.set_index("record_id")
+ quality_new_record_index = quality_new.set_index("record_id")
- fig, axs = get_fig_axes(len(plot_ims), 2, -1, ind_figsize=(8, 6))
+ # IM Compare
+ shared_record_ids = np.intersect1d(
+ quality_old_record_index.index.values.astype(str),
+ quality_new_record_index.index.values.astype(str),
+ )
- for i, (cur_im, cur_ax) in enumerate(zip(plot_ims, axs)):
- cur_old = quality_old_record_index.loc[shared_record_ids, cur_im]
- cur_new = quality_new_record_index.loc[shared_record_ids, cur_im]
+ plot_ims = [
+ "PGV",
+ "PGA",
+ "pSA_0.01",
+ "pSA_0.1",
+ "pSA_0.5",
+ "pSA_1.0",
+ "pSA_3.0",
+ "pSA_10.0",
+ ]
- cur_max = max(cur_old.max(), cur_new.max())
+ fig, axs = get_fig_axes(len(plot_ims), 2, -1, ind_figsize=(8, 6))
- cur_ax.scatter(cur_old, cur_new, s=1)
- cur_ax.set_xlabel("Old NZGMDB")
- cur_ax.set_ylabel("New NZGMDB")
- cur_ax.set_title(cur_im)
- cur_ax.plot(
- [0, cur_max], [0, cur_max], color="black", linestyle="--", linewidth=0.5
- )
- cur_ax.set_xlim(0, cur_max)
- cur_ax.set_ylim(0, cur_max)
+ for i, (cur_im, cur_ax) in enumerate(zip(plot_ims, axs)):
+ cur_old = quality_old_record_index.loc[shared_record_ids, cur_im]
+ cur_new = quality_new_record_index.loc[shared_record_ids, cur_im]
- # If the cur_im is PGA or includes pSA, set the x and y scales to log
- if cur_im == "PGA" or cur_im.startswith("pSA_"):
- cur_ax.set_xscale("log")
- cur_ax.set_yscale("log")
- cur_ax.set_xlim(0.001, cur_max)
- cur_ax.set_ylim(0.001, cur_max)
+ cur_max = max(cur_old.max(), cur_new.max())
- cur_ax.grid(which="both", linewidth=0.5, alpha=0.5, linestyle="--")
+ cur_ax.scatter(cur_old, cur_new, s=1)
+ cur_ax.set_xlabel("Old NZGMDB")
+ cur_ax.set_ylabel("New NZGMDB")
+ cur_ax.set_title(cur_im)
+ cur_ax.plot(
+ [0, cur_max], [0, cur_max], color="black", linestyle="--", linewidth=0.5
+ )
+ cur_ax.set_xlim(0, cur_max)
+ cur_ax.set_ylim(0, cur_max)
- # Convert to base64
- buf = BytesIO()
- fig.savefig(buf, format="png", bbox_inches="tight")
- plt.close(fig)
- buf.seek(0)
- img_base64 = base64.b64encode(buf.read()).decode("utf-8")
- # Embed in HTML
- html_parts.append(f'
')
- html_parts.append("
")
+ # If the cur_im is PGA or includes pSA, set the x and y scales to log
+ if cur_im == "PGA" or cur_im.startswith("pSA_"):
+ cur_ax.set_xscale("log")
+ cur_ax.set_yscale("log")
+ cur_ax.set_xlim(0.001, cur_max)
+ cur_ax.set_ylim(0.001, cur_max)
- html_parts.append(
- """
-
GMM Input Comparison
-
- """
- )
- # GMM Input comparison
- input_cols = [
- "ev_lat",
- "ev_lon",
- "ev_depth",
- "mag",
- "strike",
- "dip",
- "rake",
- "z_tor",
- "r_jb",
- "r_rup",
- "r_x",
- "Vs30",
- "Z1.0",
- "Z2.5",
- ]
- # Generate figures and embed as base64
- for i, col in enumerate(input_cols):
- cur_x_data = quality_old_record_index.loc[shared_record_ids, col].values
- cur_y_data = quality_new_record_index.loc[shared_record_ids, col].values
- nan_mask = np.isnan(cur_x_data) | np.isnan(cur_y_data)
- cur_x_data = cur_x_data[~nan_mask]
- cur_y_data = cur_y_data[~nan_mask]
- lims = (
- np.quantile(cur_x_data, 0.01),
- np.quantile(cur_x_data, 0.99),
- )
- fig, ax = plt.subplots(figsize=(8, 6), dpi=300)
- ax.scatter(cur_x_data, cur_y_data, alpha=0.5, s=1)
- ax.plot(lims, lims, color="k")
- ax.set_xlabel("Old NZGMDB")
- ax.set_ylabel("New NZGMDB")
- ax.grid(linewidth=0.5, alpha=0.5, linestyle="--")
- ax.set_xlim(lims)
- ax.set_ylim(lims)
- ax.set_aspect("equal")
- ax.set_title(f"{col} - N: {len(cur_x_data)}")
- fig.tight_layout()
+ cur_ax.grid(which="both", linewidth=0.5, alpha=0.5, linestyle="--")
# Convert to base64
buf = BytesIO()
@@ -2019,6 +2093,62 @@ def generate_report(
img_base64 = base64.b64encode(buf.read()).decode("utf-8")
# Embed in HTML
html_parts.append(f'
')
+ html_parts.append("
")
+
+ html_parts.append(
+ """
+
GMM Input Comparison
+
+ """
+ )
+ # GMM Input comparison
+ input_cols = [
+ "ev_lat",
+ "ev_lon",
+ "ev_depth",
+ "mag",
+ "strike",
+ "dip",
+ "rake",
+ "z_tor",
+ "r_jb",
+ "r_rup",
+ "r_x",
+ "Vs30",
+ "Z1.0",
+ "Z2.5",
+ ]
+ # Generate figures and embed as base64
+ for i, col in enumerate(input_cols):
+ cur_x_data = quality_old_record_index.loc[shared_record_ids, col].values
+ cur_y_data = quality_new_record_index.loc[shared_record_ids, col].values
+ nan_mask = np.isnan(cur_x_data) | np.isnan(cur_y_data)
+ cur_x_data = cur_x_data[~nan_mask]
+ cur_y_data = cur_y_data[~nan_mask]
+ lims = (
+ np.quantile(cur_x_data, 0.01),
+ np.quantile(cur_x_data, 0.99),
+ )
+ fig, ax = plt.subplots(figsize=(8, 6), dpi=300)
+ ax.scatter(cur_x_data, cur_y_data, alpha=0.5, s=1)
+ ax.plot(lims, lims, color="k")
+ ax.set_xlabel("Old NZGMDB")
+ ax.set_ylabel("New NZGMDB")
+ ax.grid(linewidth=0.5, alpha=0.5, linestyle="--")
+ ax.set_xlim(lims)
+ ax.set_ylim(lims)
+ ax.set_aspect("equal")
+ ax.set_title(f"{col} - N: {len(cur_x_data)}")
+ fig.tight_layout()
+
+ # Convert to base64
+ buf = BytesIO()
+ fig.savefig(buf, format="png", bbox_inches="tight")
+ plt.close(fig)
+ buf.seek(0)
+ img_base64 = base64.b64encode(buf.read()).decode("utf-8")
+ # Embed in HTML
+ html_parts.append(f'
')
# End of Section
html_parts.append("
")
@@ -2026,3 +2156,7 @@ def generate_report(
# Save report
with open(output_file, "w") as f:
f.write("".join(html_parts))
+
+
+if __name__ == "__main__":
+ app()
diff --git a/nzgmdb/scripts/real_time_eq_runs.py b/nzgmdb/scripts/real_time_eq_runs.py
index c7607a82..f81d7510 100644
--- a/nzgmdb/scripts/real_time_eq_runs.py
+++ b/nzgmdb/scripts/real_time_eq_runs.py
@@ -129,7 +129,7 @@ def reply_to_message_on_slack(thread_ts: str, reply_message: str):
def download_earthquake_data(
- start_date: datetime, end_date: datetime, mag_filter: float
+ start_date: datetime.datetime, end_date: datetime.datetime, mag_filter: float
) -> pd.DataFrame:
"""
Download the earthquake data from the GeoNet API
@@ -237,9 +237,7 @@ def run_event(
conda_sh: Annotated[Path, typer.Argument(exists=True, file_okay=True)],
gmc_activate: Annotated[str, typer.Argument()],
gmc_predict_activate: Annotated[str, typer.Argument()],
- ko_matrix_path: Annotated[
- Path, typer.Argument(exists=True, file_okay=False)
- ],
+ ko_matrix_path: Annotated[Path, typer.Argument(exists=True, file_okay=False)],
add_seismic_now: Annotated[bool, typer.Option(is_flag=True)] = False,
machine: Annotated[
cfg.MachineName,
@@ -510,7 +508,9 @@ def poll_earthquake_data(
init_start_date = None
while True:
# Get the last 10 minutes worth of data and check if there are any new events
- end_date = datetime.datetime.utcnow() - datetime.timedelta(minutes=1)
+ end_date = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(
+ minutes=1
+ )
# If an event was just executed, ensures we capture any events that may have been missed during the execution
start_date = (
end_date - datetime.timedelta(minutes=10)
diff --git a/nzgmdb/scripts/run_gmc.py b/nzgmdb/scripts/run_gmc.py
index 8b7b39f3..781c3de8 100644
--- a/nzgmdb/scripts/run_gmc.py
+++ b/nzgmdb/scripts/run_gmc.py
@@ -5,7 +5,7 @@
import functools
import multiprocessing
from pathlib import Path
-from typing import Annotated
+from typing import Annotated, Optional
import numpy as np
import pandas as pd
@@ -73,6 +73,7 @@ def process_batch(
gmc_predict_activate: str,
phase_arrival_table_ffp: Path,
prob_series_ffp: Path,
+ xml_dir: Path,
):
"""
Process a single subfolder: extract features and run predictions.
@@ -99,6 +100,8 @@ def process_batch(
The full file path to the phase arrival table
prob_series_ffp : Path
The full file path to the prob_series hdf5 file.
+ xml_dir : Path
+ The directory containing the stationxml files.
Raises
------
@@ -128,7 +131,7 @@ def process_batch(
)
else:
# Activate gmc environment and extract features for the subfolder
- features_command = f"python {gmc_scripts_path}/extract_features.py {gmc_dir} {waveform_dir} mseed --ko_matrices_dir {ko_matrices_dir} --record_list_ffp {batch_txt} --phase_arrival_table {phase_arrival_table_ffp} --prob_series {prob_series_ffp}"
+ features_command = f"python {gmc_scripts_path}/extract_features.py {gmc_dir} {waveform_dir} mseed --ko_matrices_dir {ko_matrices_dir} --record_list_ffp {batch_txt} --phase_arrival_table {phase_arrival_table_ffp} --prob_series {prob_series_ffp} --xml_dir {xml_dir}"
shell_commands.run_command(
features_command, conda_sh, gmc_activate, log_file_path_features
)
@@ -195,22 +198,26 @@ def run_gmc_processing(
int,
typer.Option(),
] = 1,
+ gmc_n_batches: Annotated[
+ Optional[int],
+ typer.Option(),
+ ] = None,
waveform_dir: Annotated[
- Path,
+ Optional[Path],
typer.Option(
exists=True,
file_okay=False,
),
] = None,
output_dir: Annotated[
- Path,
+ Optional[Path],
typer.Option(
exists=True,
file_okay=False,
),
] = None,
bypass_records_ffp: Annotated[
- Path,
+ Optional[Path],
typer.Option(),
] = None,
):
@@ -236,6 +243,9 @@ def run_gmc_processing(
Command to activate the GMC predict environment to run predictions.
n_procs : int, optional
Number of processes to use for multiprocessing (default is 1).
+ gmc_n_batches : int, optional
+ Number of batches to split the mseed files into for processing.
+ If None, it will be set to the number of processes (default is None).
waveform_dir : Path, optional
Directory containing all waveform files.
output_dir : Path, optional
@@ -270,11 +280,13 @@ def run_gmc_processing(
# Get all the mseed files
mseed_files = list(waveform_dir.rglob("*.mseed"))
- # Ensure that n_procs is equal too or less than the number of mseed files
+ # Ensure that n_procs and n_batches is equal too or less than the number of mseed files
n_procs = min(n_procs, len(mseed_files))
+ n_batches = gmc_n_batches or n_procs
+ n_batches = min(n_batches, len(mseed_files))
- # Split them into even batches based on number of mseeds and n_procs
- mseed_batches = np.array_split(mseed_files, n_procs)
+ # Split them into even batches based on number of mseeds and n_batches
+ mseed_batches = np.array_split(mseed_files, n_batches)
# Create a partial function with common arguments pre-filled
process_partial = functools.partial(
@@ -287,17 +299,26 @@ def run_gmc_processing(
gmc_predict_activate=gmc_predict_activate,
phase_arrival_table_ffp=phase_arrival_table_ffp,
prob_series_ffp=prob_series_ffp,
+ xml_dir=file_structure.get_stationxml_dir(main_dir),
)
- # Use multiprocessing with starmap and the partial function
- with multiprocessing.Pool(n_procs) as p:
- p.starmap(
- process_partial,
- [
- (batch, (gmc_dir / f"batch_{idx}"))
- for idx, batch in enumerate(mseed_batches)
- ],
- )
+ pending_jobs = []
+ for idx, batch in enumerate(mseed_batches):
+ batch_dir = gmc_dir / f"batch_{idx}"
+ predictions_output = batch_dir / file_structure.FlatfileNames.GMC_PREDICTIONS
+
+ if predictions_output.exists():
+ print(f"Skipping Batch {idx} (found existing gmc_predictions.csv)")
+ continue
+
+ pending_jobs.append((list(batch), batch_dir))
+
+ if not pending_jobs:
+ print("All batches already have gmc_predictions.csv; nothing to run.")
+ else:
+ # Use multiprocessing with starmap and the partial function (pending only)
+ with multiprocessing.Pool(n_procs) as p:
+ p.starmap(process_partial, pending_jobs)
# For each subfolder combine the gmc_predictions.csv into a single file
dfs = []
diff --git a/nzgmdb/scripts/run_nzgmdb.py b/nzgmdb/scripts/run_nzgmdb.py
index 8a5d2c46..31a7a6ff 100644
--- a/nzgmdb/scripts/run_nzgmdb.py
+++ b/nzgmdb/scripts/run_nzgmdb.py
@@ -4,7 +4,7 @@
from datetime import datetime
from pathlib import Path
-from typing import Annotated
+from typing import Annotated, Optional
import typer
@@ -29,16 +29,24 @@ def fetch_geonet_data(
n_procs: Annotated[int, typer.Option()] = 1,
batch_size: Annotated[int, typer.Option()] = 500,
only_event_ids: Annotated[
- list[str], typer.Option(callback=lambda x: [] if x is None else x[0].split(","))
+ Optional[list[str]],
+ typer.Option(callback=lambda x: [] if x is None else x[0].split(",")),
] = None,
only_sites: Annotated[
- list[str], typer.Option(callback=lambda x: [] if x is None else x[0].split(","))
+ Optional[list[str]],
+ typer.Option(callback=lambda x: [] if x is None else x[0].split(",")),
] = None,
only_record_ids_ffp: Annotated[
- Path, typer.Option(exists=True, dir_okay=False)
+ Optional[Path], typer.Option(exists=True, dir_okay=False)
] = None,
real_time: Annotated[bool, typer.Option()] = False,
mp_sites: Annotated[bool, typer.Option()] = False,
+ add_tmp_arrays: Annotated[
+ bool,
+ typer.Option(
+ is_flag=True,
+ ),
+ ] = False,
):
"""
Fetch earthquake data from Geonet and generate the earthquake source and station magnitude tables.
@@ -65,6 +73,8 @@ def fetch_geonet_data(
If True, the function will run in real-time mode by using a different client (default is False).
mp_sites : bool, optional
If True, the function will use multiprocessing over sites instead of events (default is False).
+ add_tmp_arrays : bool, optional
+ If True, temporary arrays will be added to the database run (default is False).
"""
geonet.parse_geonet_information(
main_dir,
@@ -77,6 +87,7 @@ def fetch_geonet_data(
only_record_ids_ffp,
real_time,
mp_sites,
+ add_tmp_arrays,
)
@@ -98,7 +109,7 @@ def extract_waveforms(
],
n_procs: Annotated[int, typer.Option()] = 1,
only_record_ids_ffp: Annotated[
- Path,
+ Optional[Path],
typer.Option(
exists=True,
dir_okay=False,
@@ -108,6 +119,13 @@ def extract_waveforms(
int,
typer.Option(),
] = 1000,
+ tmp_array_dir: Annotated[
+ Optional[Path],
+ typer.Option(
+ exists=True,
+ file_okay=False,
+ ),
+ ] = None,
):
"""
Extract waveforms using the station extraction table and save them as MiniSEED files.
@@ -125,9 +143,16 @@ def extract_waveforms(
The full file path to a set of record IDs to only run for. If provided, only these records will be processed.
batch_size : int, optional
The batch size for how many extracted waveforms to process before checkpointing (default is 1000).
+ tmp_array_dir : Path, optional
+ The directory the saved temporary array data is to be used for waveform extraction.
"""
waveform_extraction.extract_waveforms(
- main_dir, station_extraction_table_ffp, n_procs, only_record_ids_ffp, batch_size
+ main_dir,
+ station_extraction_table_ffp,
+ n_procs,
+ only_record_ids_ffp,
+ batch_size,
+ tmp_array_dir,
)
@@ -188,6 +213,10 @@ def make_phase_arrival_table(
typer.Argument(),
],
n_procs: Annotated[int, typer.Option()] = 1,
+ n_batches: Annotated[
+ Optional[int],
+ typer.Option(),
+ ] = None,
bypass_records_ffp: Annotated[
Path,
typer.Option(
@@ -214,6 +243,9 @@ def make_phase_arrival_table(
The command to activate the environment for running PhaseNet.
n_procs : int, optional
Number of processes to use (default is 1).
+ n_batches : int, optional
+ The number of batches to split the mseed files into for processing.
+ If not provided, it will be determined automatically based on the number of mseed files and n_procs.
bypass_records_ffp : Path, optional
The full file path to the bypass records file for custom P-wave index values.
This allows you to specify custom P-wave index values for records that may not be
@@ -225,7 +257,9 @@ def make_phase_arrival_table(
conda_sh,
env_activate_command,
n_procs,
+ n_batches,
bypass_records_ffp,
+ xml_dir=file_structure.get_stationxml_dir(main_dir),
)
@@ -493,9 +527,17 @@ def run_im_calculation(
),
] = False,
intensity_measures: Annotated[
- list[IM],
+ str | None,
typer.Option(
- callback=lambda x: [IM(i) for i in x[0].split(",")],
+ callback=lambda x: (
+ None
+ if x is None or (isinstance(x, (list, tuple)) and not x)
+ else [
+ IM(i.strip())
+ for i in (x[0] if isinstance(x, (list, tuple)) else x).split(",")
+ if i.strip()
+ ]
+ ),
),
] = None,
):
@@ -542,6 +584,12 @@ def generate_site_table_basin(
file_okay=False,
),
],
+ add_tmp_arrays: Annotated[
+ bool,
+ typer.Option(
+ is_flag=True,
+ ),
+ ] = False,
):
"""
Generate the site table basin flatfile.
@@ -554,15 +602,20 @@ def generate_site_table_basin(
The main directory of the NZGMDB results (Highest level directory).
nzcvm_data_ffp : Path
The full file path to the nzcvm_data repository that stores the basin information.
+ add_tmp_arrays : bool, optional
+ If True, temporary arrays will be added to the site table (default is False).
"""
main_dir.mkdir(parents=True, exist_ok=True)
# Generate the site basin flatfile
flatfile_dir = file_structure.get_flatfile_dir(main_dir)
flatfile_dir.mkdir(parents=True, exist_ok=True)
- site_df = sites.create_site_table_response()
+ site_df, station_df = sites.create_site_table_response(add_tmp_arrays)
site_df = sites.add_site_basins(site_df, nzcvm_data_ffp)
+ station_df.to_csv(
+ flatfile_dir / file_structure.PreFlatfileNames.STATION_TABLE, index=False
+ )
site_df.to_csv(
flatfile_dir / file_structure.PreFlatfileNames.SITE_TABLE, index=False
)
@@ -782,17 +835,13 @@ def run_full_nzgmdb(
file_okay=False,
),
],
- gmc_procs: Annotated[
- int,
- typer.Option(),
- ] = 1,
n_procs: Annotated[int, typer.Option()] = 1,
checkpoint: Annotated[
bool,
typer.Option(),
] = False,
only_event_ids: Annotated[
- list[str],
+ Optional[list[str]],
typer.Option(
callback=lambda x: [] if x is None else x[0].split(","),
),
@@ -813,11 +862,19 @@ def run_full_nzgmdb(
geonet_batch_size: Annotated[
int,
typer.Option(),
- ] = 500,
+ ] = 100,
snr_batch_size: Annotated[
int,
typer.Option(),
] = 5000,
+ phase_arrival_n_batches: Annotated[
+ int,
+ typer.Option(),
+ ] = None,
+ gmc_n_batches: Annotated[
+ int,
+ typer.Option(),
+ ] = None,
real_time: Annotated[
bool,
typer.Option(),
@@ -837,6 +894,17 @@ def run_full_nzgmdb(
dir_okay=False,
),
] = None,
+ add_tmp_arrays: Annotated[
+ bool,
+ typer.Option(),
+ ] = False,
+ tmp_array_data_dir: Annotated[
+ Path,
+ typer.Option(
+ exists=True,
+ file_okay=False,
+ ),
+ ] = None,
machine: Annotated[
cfg.MachineName,
typer.Option(
@@ -850,7 +918,9 @@ def run_full_nzgmdb(
This function orchestrates the full pipeline of NZGMDB, executing all necessary steps sequentially.
Steps Included:
+ - Generate site table with basin information
- Fetch Geonet data
+ - Waveform extraction
- Merge tectonic domains
- Generate phase arrival table
- Calculate SNR
@@ -883,8 +953,6 @@ def run_full_nzgmdb(
Command to activate gmc_predict environment to run the predictions.
ko_matrix_path : Path
Path to the ko matrix directory
- gmc_procs : int, optional
- Number of processes to use for GMC (default is 1).
n_procs : int, optional
The number of processes to use (default is 1).
checkpoint : bool, optional
@@ -899,6 +967,12 @@ def run_full_nzgmdb(
The batch size for Geonet data retrieval (default is 500).
snr_batch_size : int, optional
The batch size for the SNR calculation (default is 5000).
+ phase_arrival_n_batches : int, optional
+ The number of batches to split the phase arrival calculation into
+ (default is None, which will automatically determine the number of batches based on the number of records and n_procs).
+ gmc_n_batches : int, optional
+ The number of batches to split the GMC prediction step into
+ (default is None, which will automatically determine the number of batches based on the number of records and gmc_procs).
real_time : bool, optional
If True, the function will run in real-time mode using a different client (default is False).
upload : bool, optional
@@ -907,12 +981,22 @@ def run_full_nzgmdb(
If True, the function will create a quality database (default is False).
bypass_records_ffp : Path, optional
The full file path to the bypass records file, if applicable.
+ add_tmp_arrays : bool, optional
+ If True, temporary arrays will be added to the database run (default is False).
+ tmp_array_data_dir : Path, optional
+ The directory containing temporary array data, required if add_tmp_arrays is True.
machine : cfg.MachineName, optional
The machine name to use for process configuration (default is None).
"""
main_dir.mkdir(parents=True, exist_ok=True)
config = cfg.Config()
+ # Check that if add_tmp_arrays is True, tmp_array_data_dir is provided
+ if add_tmp_arrays and tmp_array_data_dir is None:
+ raise ValueError(
+ "tmp_array_data_dir must be provided if add_tmp_arrays is True."
+ )
+
# Generate the site basin flatfile
flatfile_dir = file_structure.get_flatfile_dir(main_dir)
flatfile_dir.mkdir(parents=True, exist_ok=True)
@@ -921,7 +1005,7 @@ def run_full_nzgmdb(
and (flatfile_dir / file_structure.PreFlatfileNames.SITE_TABLE).exists()
):
print("Generating site table basin flatfile")
- generate_site_table_basin(main_dir, nzcvm_data_ffp)
+ generate_site_table_basin(main_dir, nzcvm_data_ffp, add_tmp_arrays)
# Fetch the Geonet data
if not (
@@ -947,6 +1031,7 @@ def run_full_nzgmdb(
only_sites,
only_record_ids_ffp,
real_time,
+ add_tmp_arrays=add_tmp_arrays,
)
# Extract Waveforms
@@ -970,6 +1055,7 @@ def run_full_nzgmdb(
extract_n_procs,
only_record_ids_ffp,
batch_size=geonet_batch_size,
+ tmp_array_dir=tmp_array_data_dir,
)
# Merge the tectonic domains
@@ -1018,7 +1104,9 @@ def run_full_nzgmdb(
conda_sh,
gmc_activate,
phase_n_procs,
+ phase_arrival_n_batches,
bypass_records_ffp,
+ xml_dir=file_structure.get_stationxml_dir(main_dir),
)
# Generate SNR
@@ -1072,7 +1160,7 @@ def run_full_nzgmdb(
):
print("Running GMC")
gmc_n_procs = (
- gmc_procs
+ n_procs
if machine is None
else config.get_n_procs(machine, cfg.WorkflowStep.GMC)
)
@@ -1084,6 +1172,7 @@ def run_full_nzgmdb(
gmc_activate,
gmc_predict_activate,
gmc_n_procs,
+ gmc_n_batches,
bypass_records_ffp=bypass_records_ffp,
)
diff --git a/nzgmdb/scripts/upload_to_dropbox.py b/nzgmdb/scripts/upload_to_dropbox.py
index c1aac9ad..424937df 100644
--- a/nzgmdb/scripts/upload_to_dropbox.py
+++ b/nzgmdb/scripts/upload_to_dropbox.py
@@ -443,7 +443,7 @@ def download_dropbox_archive(
if not ignore_xml:
xml_dir.mkdir(exist_ok=True)
zips_to_download.append(
- f"{dropbox_version_dir}/{xml_zip}", zip_dir / xml_zip, xml_dir
+ (f"{dropbox_version_dir}/{xml_zip}", zip_dir / xml_zip, xml_dir)
)
# Ensure there is something to download
diff --git a/nzgmdb/temp_arrays/backup_to_dropbox.py b/nzgmdb/temp_arrays/backup_to_dropbox.py
new file mode 100755
index 00000000..c6864e87
--- /dev/null
+++ b/nzgmdb/temp_arrays/backup_to_dropbox.py
@@ -0,0 +1,342 @@
+import csv
+import subprocess
+from collections.abc import Iterator
+from pathlib import Path
+from typing import TypedDict
+
+import typer
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+MANIFEST_HEADER = [
+ "type",
+ "net",
+ "name",
+ "local_path",
+ "zip_name",
+ "status",
+ "bytes",
+]
+
+
+class ManifestRow(TypedDict):
+ """A single row in the Dropbox backup manifest CSV."""
+
+ type: str
+ net: str
+ name: str
+ local_path: str
+ zip_name: str
+ status: str
+ bytes: int
+
+
+def zip_directory(src_dir: Path, out_dir: Path) -> Path:
+ """Create a zip file for a directory.
+
+ This uses the system `zip` executable (fast + reliable) and skips work if the
+ expected zip already exists.
+
+ Parameters
+ ----------
+ src_dir : Path
+ Source directory to zip.
+ out_dir : Path
+ Output directory where the zip file will be written.
+
+ Returns
+ -------
+ Path
+ Path to the created (or existing) zip file.
+ """
+ out_dir.mkdir(parents=True, exist_ok=True)
+ zip_path = out_dir / f"{src_dir.name}.zip"
+
+ if zip_path.exists():
+ return zip_path
+
+ subprocess.check_call(
+ ["zip", "-r", "-1", str(zip_path), src_dir.name],
+ cwd=src_dir.parent,
+ )
+
+ return zip_path
+
+
+def upload_and_verify(local_file: Path, dropbox_dir: str) -> bool:
+ """Upload a file to Dropbox using rclone and verify by file size.
+
+ Parameters
+ ----------
+ local_file : Path
+ File to upload.
+ dropbox_dir : str
+ Rclone remote path (directory) to upload into.
+
+ Returns
+ -------
+ bool
+ True if the remote file exists and its size matches the local file.
+ """
+ subprocess.check_call(
+ ["rclone", "copy", str(local_file), dropbox_dir],
+ )
+
+ local_size = local_file.stat().st_size
+
+ out = (
+ subprocess.check_output(
+ [
+ "rclone",
+ "lsf",
+ "--format=s",
+ f"{dropbox_dir}/{local_file.name}",
+ ]
+ )
+ .decode()
+ .strip()
+ )
+
+ return bool(out) and int(out) == local_size
+
+
+def load_manifest(path: Path) -> dict[str, ManifestRow]:
+ """Load an upload manifest keyed by ``zip_name``.
+
+ Parameters
+ ----------
+ path : Path
+ Path to the manifest CSV.
+
+ Returns
+ -------
+ dict[str, ManifestRow]
+ Manifest rows keyed by ``zip_name``. Returns an empty dict if the file
+ does not exist.
+ """
+ rows: dict[str, ManifestRow] = {}
+
+ if not path.exists():
+ return rows
+
+ with path.open(newline="") as f:
+ reader: csv.DictReader[str] = csv.DictReader(f)
+ for row in reader:
+ # DictReader gives us strings; coerce the known int column.
+ bytes_value = int(row.get("bytes") or 0)
+ zip_name = row.get("zip_name")
+ if not zip_name:
+ # Keep behavior conservative: skip malformed rows.
+ continue
+
+ rows[zip_name] = ManifestRow(
+ type=row.get("type", ""),
+ net=row.get("net", ""),
+ name=row.get("name", ""),
+ local_path=row.get("local_path", ""),
+ zip_name=zip_name,
+ status=row.get("status", ""),
+ bytes=bytes_value,
+ )
+
+ return rows
+
+
+def append_manifest_row(path: Path, row: ManifestRow) -> None:
+ """Append a single row to the manifest CSV, creating it if needed.
+
+ Parameters
+ ----------
+ path : Path
+ Path to the manifest CSV.
+ row : ManifestRow
+ The manifest row to append.
+
+ Returns
+ -------
+ None
+ """
+ new_file = not path.exists()
+
+ with path.open("a", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=MANIFEST_HEADER)
+ if new_file:
+ writer.writeheader()
+ writer.writerow(row)
+
+
+def update_manifest_status(path: Path, zip_name: str, status: str, size: int) -> None:
+ """Update an existing manifest row with a new status and size.
+
+ Parameters
+ ----------
+ path : Path
+ Path to the manifest CSV.
+ zip_name : str
+ Zip filename key for the entry to update.
+ status : str
+ New status string (e.g. "DONE", "FAILED").
+ size : int
+ Size of the zip file in bytes.
+ """
+ rows = load_manifest(path)
+ rows[zip_name]["status"] = status
+ rows[zip_name]["bytes"] = size
+
+ with path.open("w", newline="") as f:
+ writer = csv.DictWriter(f, fieldnames=MANIFEST_HEADER)
+ writer.writeheader()
+ writer.writerows(rows.values())
+
+
+def discover_stationxml(stationxml_dir: Path) -> Iterator[ManifestRow]:
+ """Discover the StationXML directory entry.
+
+ Parameters
+ ----------
+ stationxml_dir : Path
+ Path to the directory containing StationXML files.
+
+ Returns
+ -------
+ Iterator[ManifestRow]
+ An iterator yielding a single manifest row representing the StationXML
+ directory.
+ """
+ yield ManifestRow(
+ type="stationxml",
+ net="",
+ name=stationxml_dir.name,
+ local_path=str(stationxml_dir),
+ zip_name=f"{stationxml_dir.name}.zip",
+ status="PENDING",
+ bytes=0,
+ )
+
+
+def discover_waveforms(waveforms_root: Path) -> Iterator[ManifestRow]:
+ """Discover waveform leaf directories to back up.
+
+ The expected directory structure is ``waveforms/
//`` where each
+ ``leaf`` directory is zipped independently.
+
+ Parameters
+ ----------
+ waveforms_root : Path
+ Root directory containing per-network waveform directories.
+
+ Returns
+ -------
+ Iterator[ManifestRow]
+ Iterator of manifest rows for each leaf directory.
+ """
+ for net_dir in sorted(waveforms_root.iterdir()):
+ if not net_dir.is_dir():
+ continue
+
+ for leaf in sorted(net_dir.iterdir()):
+ if not leaf.is_dir():
+ continue
+
+ yield ManifestRow(
+ type="waveforms",
+ net=net_dir.name,
+ name=leaf.name,
+ local_path=str(leaf),
+ zip_name=f"{leaf.name}.zip",
+ status="PENDING",
+ bytes=0,
+ )
+
+
+def process_entry(
+ entry: ManifestRow, tmp_zip_dir: Path, manifest_path: Path, dropbox_path: str
+) -> None:
+ """Zip, upload, verify, and update manifest for a single entry.
+
+ Parameters
+ ----------
+ entry : ManifestRow
+ Manifest entry describing what to back up.
+ tmp_zip_dir : Path
+ Directory where zips are created temporarily.
+ manifest_path : Path
+ Path to the manifest CSV.
+ dropbox_path : str
+ Base rclone Dropbox remote path.
+ """
+ src = Path(entry["local_path"])
+ zip_path = zip_directory(src, tmp_zip_dir)
+
+ if entry["type"] == "stationxml":
+ dropbox_target = f"{dropbox_path}/stationxml"
+ else:
+ dropbox_target = f"{dropbox_path}/waveforms/{entry['net']}"
+
+ try:
+ ok = upload_and_verify(zip_path, dropbox_target)
+ except (OSError, subprocess.CalledProcessError) as e:
+ print(f"ERROR uploading {zip_path.name}: {e}")
+ update_manifest_status(manifest_path, entry["zip_name"], "FAILED", 0)
+ return
+
+ if ok:
+ size = zip_path.stat().st_size
+ update_manifest_status(manifest_path, entry["zip_name"], "DONE", size)
+ zip_path.unlink()
+ print(f"DONE {zip_path.name}")
+ else:
+ update_manifest_status(manifest_path, entry["zip_name"], "FAILED", 0)
+ print(f"FAILED {zip_path.name}")
+
+
+@app.command()
+def run(
+ data_root: Path = typer.Argument(
+ ..., help="Root directory containing waveforms/ and stationxml/"
+ ),
+ dropbox_path: str = typer.Argument(
+ ...,
+ help="Rclone Dropbox path to upload to.",
+ ),
+) -> None:
+ """Resume-safe Dropbox backup with manifest tracking.
+
+ Parameters
+ ----------
+ data_root : Path
+ Root directory containing ``waveforms/`` and ``stationxml/``.
+ dropbox_path : str
+ Base rclone Dropbox remote path to upload to.
+ """
+ manifest = data_root / "dropbox_manifest.csv"
+ manifest_rows = load_manifest(manifest)
+
+ # Discover + register new entries
+ stationxml = data_root / "stationxml"
+ waveforms = data_root / "waveforms"
+
+ tmp_zip_dir = data_root / "tmp_zips"
+ tmp_zip_dir.mkdir(parents=True, exist_ok=True)
+
+ for entry in discover_stationxml(stationxml):
+ if entry["zip_name"] not in manifest_rows:
+ append_manifest_row(manifest, entry)
+
+ for entry in discover_waveforms(waveforms):
+ if entry["zip_name"] not in manifest_rows:
+ append_manifest_row(manifest, entry)
+
+ # Reload after discovery
+ manifest_rows = load_manifest(manifest)
+
+ pending = [row for row in manifest_rows.values() if row["status"] != "DONE"]
+
+ print(f"Pending uploads: {len(pending)}")
+
+ for entry in pending:
+ process_entry(entry, tmp_zip_dir, manifest, dropbox_path)
+
+
+if __name__ == "__main__":
+ app()
diff --git a/nzgmdb/temp_arrays/check_data.py b/nzgmdb/temp_arrays/check_data.py
new file mode 100644
index 00000000..da58933e
--- /dev/null
+++ b/nzgmdb/temp_arrays/check_data.py
@@ -0,0 +1,293 @@
+from pathlib import Path
+
+import pandas as pd
+import typer
+from obspy import UTCDateTime
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+# Month length in seconds (30 days)
+MONTH_SECONDS = 30 * 24 * 3600
+
+
+def _format_end_for_filename(end_dt: UTCDateTime | str) -> str:
+ """Format an end datetime into the MiniSEED filename timestamp form.
+
+ Parameters
+ ----------
+ end_dt : UTCDateTime | str
+ An ObsPy ``UTCDateTime`` or any string parseable by ``UTCDateTime``.
+
+ Returns
+ -------
+ str
+ Timestamp string formatted like ``YYYYMMDDTHHMMSSZ``.
+
+ Raises
+ ------
+ ValueError
+ If ``end_dt`` cannot be parsed by ``UTCDateTime``.
+ """
+ if not isinstance(end_dt, UTCDateTime):
+ end_dt = UTCDateTime(end_dt)
+ return end_dt.strftime("%Y%m%dT%H%M%SZ")
+
+
+def _row_mseed_dir(output_dir: Path, mseed_dirname: str, row: pd.Series) -> Path:
+ """Build the expected output MiniSEED directory for a CSV row.
+
+ Parameters
+ ----------
+ output_dir : Path
+ Root output directory containing the MiniSEED subdirectory.
+ mseed_dirname : str
+ Name of the MiniSEED subdirectory (typically ``"waveforms"``).
+ row : pandas.Series
+ A row from the input CSV containing at least ``net``, ``sta``, ``loc``,
+ and ``chan``.
+
+ Returns
+ -------
+ Path
+ Path like ``///___``.
+ """
+ net = str(row["net"]).strip()
+ sta = str(row["sta"]).strip()
+ loc_field = str(row["loc"])
+ chan_prefix = str(row["chan"]).strip()
+
+ record_sub = f"{net}_{sta}_{chan_prefix}_{loc_field}"
+ return output_dir / mseed_dirname / net / record_sub
+
+
+def is_row_done(row: pd.Series, *, output_dir: Path, mseed_dirname: str) -> bool:
+ """Check whether a CSV row has fully completed waveform download.
+
+ A row is considered **done** if the expected MiniSEED output directory
+ contains at least one ``.mseed`` file whose final ``__END`` timestamp equals
+ the row's ``end_date``.
+
+ Notes
+ -----
+ This check is filesystem-driven and intentionally conservative:
+
+ - Only files whose channel code ends with ``"Z"`` are considered.
+ - Any filesystem error or unexpected filename format causes this function
+ to return ``False`` (so the row can be retried).
+
+ Parameters
+ ----------
+ row : pandas.Series
+ A row of the input CSV.
+ output_dir : Path
+ Root output directory.
+ mseed_dirname : str
+ Name of the MiniSEED subdirectory (typically ``"waveforms"``).
+
+ Returns
+ -------
+ bool
+ ``True`` if the row appears fully completed, otherwise ``False``.
+ """
+ mseed_path = _row_mseed_dir(output_dir, mseed_dirname, row)
+
+ if not mseed_path.is_dir():
+ return False
+
+ target_end = _format_end_for_filename(row["end_date"])
+
+ try:
+ for path in mseed_path.iterdir():
+ if path.suffix != ".mseed":
+ continue
+
+ stem = path.stem
+
+ # Filename parts expected like: NET.STA..CHAN__START__END
+ parts = stem.split(".")
+ if len(parts) <= 3:
+ continue
+
+ chan_check = parts[3].split("__", 1)[0] # remove any suffix after __
+ if not chan_check.endswith("Z"):
+ continue
+
+ if "__" not in stem:
+ continue
+
+ last = stem.rsplit("__", 1)[-1]
+ if last == target_end:
+ return True
+ except OSError:
+ # Any filesystem error -> treat as not done so row will be retried.
+ return False
+
+ return False
+
+
+def is_row_started(row: pd.Series, *, output_dir: Path, mseed_dirname: str) -> bool:
+ """Check whether any waveform data exists for a CSV row.
+
+ Parameters
+ ----------
+ row : pandas.Series
+ A row of the input CSV.
+ output_dir : Path
+ Root output directory.
+ mseed_dirname : str
+ Name of the MiniSEED subdirectory (typically ``"waveforms"``).
+
+ Returns
+ -------
+ bool
+ ``True`` if at least one ``.mseed`` file exists for the row.
+ """
+ mseed_path = _row_mseed_dir(output_dir, mseed_dirname, row)
+
+ if not mseed_path.is_dir():
+ return False
+
+ return any(p.suffix == ".mseed" for p in mseed_path.iterdir())
+
+
+def evaluate_download_completeness(
+ csv_file: Path,
+ output_csv: Path,
+ output_dir: Path,
+ provider: str = "IRIS",
+) -> pd.DataFrame:
+ """Evaluate and summarize download completeness for waveform requests.
+
+ The evaluation is performed at two levels:
+
+ - **completed**: at least one MiniSEED file exists with a final ``__END``
+ timestamp equal to the row ``end_date``.
+ - **started**: at least one MiniSEED file exists for the row.
+
+ Parameters
+ ----------
+ csv_file : Path
+ Input CSV describing waveform requests.
+ output_csv : Path
+ Output CSV to write the evaluation results to.
+ output_dir : Path
+ Root output directory containing the MiniSEED data.
+ provider : str, optional
+ Provider name to filter by, by default "IRIS".
+
+ Returns
+ -------
+ pandas.DataFrame
+ The evaluated dataframe including the ``completed`` and ``started``
+ boolean columns.
+
+ Raises
+ ------
+ ValueError
+ If required columns are missing from the input CSV.
+ """
+ df = pd.read_csv(csv_file, dtype={"loc": str}, keep_default_na=False)
+
+ required_cols = {
+ "provider",
+ "net",
+ "sta",
+ "loc",
+ "chan",
+ "start_date",
+ "end_date",
+ }
+ missing = sorted(required_cols - set(df.columns))
+ if missing:
+ raise ValueError(f"Input CSV missing required columns: {missing}")
+
+ df = df[df["provider"] == provider].reset_index(drop=True)
+
+ print(f"Evaluating {len(df)} rows for completeness...")
+
+ mseed_dirname = "waveforms"
+
+ df["completed"] = df.apply(
+ is_row_done, axis=1, output_dir=output_dir, mseed_dirname=mseed_dirname
+ )
+ df["started"] = df.apply(
+ is_row_started, axis=1, output_dir=output_dir, mseed_dirname=mseed_dirname
+ )
+
+ total = len(df)
+ completed = int(df["completed"].sum())
+ started = int(df["started"].sum())
+
+ partial = int(((df["started"]) & (~df["completed"])).sum())
+ none = int((~df["started"]).sum())
+
+ print("\n===== DOWNLOAD SUMMARY =====")
+ print(f"Total rows : {total}")
+ print(f"Fully completed : {completed}")
+ print(f"Started (any data) : {started}")
+ print(f"Partial (started only): {partial}")
+ print(f"No data at all : {none}")
+
+ cols = ["provider", "net", "sta", "loc", "chan", "start_date", "end_date"]
+
+ if none > 0:
+ print("\n===== ROWS WITH NO DATA =====")
+ print(df.loc[~df["started"], cols].to_string(index=False))
+
+ if partial > 0:
+ print("\n===== PARTIALLY DOWNLOADED ROWS =====")
+ print(df.loc[df["started"] & ~df["completed"], cols].to_string(index=False))
+
+ output_csv.parent.mkdir(parents=True, exist_ok=True)
+ df.to_csv(output_csv, index=False)
+ print(f"\nWrote evaluation CSV to:\n {output_csv}")
+
+ return df
+
+
+@app.command()
+def run(
+ csv_file: Path = typer.Argument(
+ ..., exists=True, dir_okay=False, help="Input waveform request CSV."
+ ),
+ output_csv: Path = typer.Argument(
+ ..., dir_okay=False, help="Output CSV to write evaluation results to."
+ ),
+ output_dir: Path = typer.Argument(
+ ...,
+ exists=False,
+ file_okay=False,
+ help="Root output directory containing the MiniSEED data.",
+ ),
+ provider: str = typer.Option(
+ "IRIS",
+ help="Provider name to filter by.",
+ ),
+) -> None:
+ """Evaluate MiniSEED download completeness for a request table.
+
+ Parameters
+ ----------
+ csv_file : Path
+ Input CSV describing waveform requests.
+ output_csv : Path
+ Output CSV to write the evaluation results.
+ output_dir : Path
+ Root output directory containing the MiniSEED data.
+ provider : str
+ Provider name to filter by.
+
+ Returns
+ -------
+ None
+ """
+ evaluate_download_completeness(
+ csv_file=csv_file,
+ output_csv=output_csv,
+ output_dir=output_dir,
+ provider=provider,
+ )
+
+
+if __name__ == "__main__":
+ app()
diff --git a/nzgmdb/temp_arrays/get_stations.py b/nzgmdb/temp_arrays/get_stations.py
new file mode 100644
index 00000000..d44c093d
--- /dev/null
+++ b/nzgmdb/temp_arrays/get_stations.py
@@ -0,0 +1,381 @@
+"""Fetch station/channel metadata from FDSN providers and write to CSV."""
+
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Annotated
+
+import geopandas as gpd
+import nzgeom.coastlines
+import pandas as pd
+import typer
+from obspy import UTCDateTime
+from obspy.clients.fdsn import Client as FDSN_Client
+from shapely.geometry import Point
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+URL_MAPPINGS = {
+ "AUSPASS": "http://auspass.edu.au",
+ "BGR": "http://eida.bgr.de",
+ "EIDA": "http://eida-federator.ethz.ch",
+ "ETH": "http://eida.ethz.ch",
+ "EMSC": "http://www.seismicportal.eu",
+ "GEONET": "http://service.geonet.org.nz",
+ "GEOFON": "http://geofon.gfz-potsdam.de",
+ "GFZ": "http://geofon.gfz-potsdam.de",
+ "ICGC": "http://ws.icgc.cat",
+ "IESDMC": "http://batsws.earth.sinica.edu.tw",
+ "INGV": "http://webservices.ingv.it",
+ "IPGP": "http://ws.ipgp.fr",
+ "IRIS": "http://service.iris.edu",
+ "IRISPH5": "http://service.iris.edu",
+ "ISC": "http://www.isc.ac.uk",
+ "KNMI": "http://rdsa.knmi.nl",
+ "KOERI": "http://eida.koeri.boun.edu.tr",
+ "LMU": "https://erde.geophysik.uni-muenchen.de",
+ "NCEDC": "https://service.ncedc.org",
+ "NIEP": "http://eida-sc3.infp.ro",
+ "NOA": "http://eida.gein.noa.gr",
+ "ODC": "http://www.orfeus-eu.org",
+ "ORFEUS": "http://www.orfeus-eu.org",
+ "RESIF": "http://ws.resif.fr",
+ "RESIFPH5": "http://ph5ws.resif.fr",
+ "RASPISHAKE": "https://data.raspberryshake.org",
+ "SCEDC": "http://service.scedc.caltech.edu",
+ "TEXNET": "http://rtserve.beg.utexas.edu",
+ "UIB-NORSAR": "http://eida.geo.uib.no",
+ "USGS": "http://earthquake.usgs.gov",
+ "USP": "http://sismo.iag.usp.br",
+}
+
+
+def is_point_inside_nz(lat: float, lon: float, *, nz_coast: gpd.GeoDataFrame) -> bool:
+ """Return True if the given point lies within NZ coastline polygons.
+
+ Parameters
+ ----------
+ lat : float
+ Latitude in degrees.
+ lon : float
+ Longitude in degrees.
+ nz_coast : geopandas.GeoDataFrame
+ Coastline polygons in EPSG:4326.
+
+ Returns
+ -------
+ bool
+ True if the point is contained in any polygon.
+ """
+
+ p = Point(lon, lat)
+
+ # Iterate geometries directly so we don't rely on pandas/geopandas methods here.
+ return any(getattr(g, "contains")(p) for g in nz_coast.geometry)
+
+
+def _parse_time(value: str) -> UTCDateTime:
+ """Parse a time value into ``UTCDateTime``.
+
+ Parameters
+ ----------
+ value : str
+ Input time string parseable by ObsPy. If "now" (case-insensitive), the
+ current time is used.
+
+ Returns
+ -------
+ UTCDateTime
+ Parsed time.
+
+ Raises
+ ------
+ ValueError
+ If the value cannot be parsed.
+ """
+ if value.strip().lower() == "now":
+ return UTCDateTime()
+ return UTCDateTime(value)
+
+
+def _iter_providers(providers: list[str] | None) -> list[str]:
+ """Build the provider list for querying.
+
+ Parameters
+ ----------
+ providers : list[str] | None
+ Optional list of provider names. If None or empty, uses all known
+ providers.
+
+ Returns
+ -------
+ list[str]
+ Provider names.
+ """
+ if not providers:
+ return sorted(URL_MAPPINGS)
+ return providers
+
+
+def collect_station_rows(
+ providers: list[str],
+ starttime: UTCDateTime,
+ endtime: UTCDateTime,
+ minlatitude: float,
+ maxlatitude: float,
+ minlongitude: float,
+ maxlongitude: float,
+ coastline_filter: bool,
+) -> list[list[object]]:
+ """Collect station/channel rows from configured providers.
+
+ Parameters
+ ----------
+ providers : list[str]
+ Provider names.
+ starttime : UTCDateTime
+ Metadata start time.
+ endtime : UTCDateTime
+ Metadata end time.
+ minlatitude, maxlatitude : float
+ Latitude bounding box.
+ minlongitude, maxlongitude : float
+ Longitude bounding box.
+ coastline_filter : bool
+ If True, filter stations to those within NZ coastline polygons.
+
+ Returns
+ -------
+ list[list[object]]
+ Rows matching the output CSV schema.
+ """
+ nz_coast = (
+ nzgeom.coastlines.get_NZ_coastlines().to_crs("EPSG:4326")
+ if coastline_filter
+ else None
+ )
+
+ all_station_info: list[list[object]] = []
+
+ for provider in providers:
+ client = FDSN_Client(base_url=provider)
+ networks = client.get_stations(
+ starttime=starttime,
+ endtime=endtime,
+ minlatitude=minlatitude,
+ maxlatitude=maxlatitude,
+ minlongitude=minlongitude,
+ maxlongitude=maxlongitude,
+ level="network",
+ )
+
+ network_codes = [net.code for net in networks]
+ print("Processing provider:", provider, "with networks:", len(network_codes))
+
+ for net_code in network_codes:
+ inv = client.get_stations(
+ network=net_code,
+ starttime=starttime,
+ endtime=endtime,
+ level="channel",
+ minlatitude=minlatitude,
+ maxlatitude=maxlatitude,
+ minlongitude=minlongitude,
+ maxlongitude=maxlongitude,
+ )
+
+ for network in inv:
+ for station in network:
+ lat = float(station.latitude)
+ lon = float(station.longitude)
+
+ if nz_coast is not None and not is_point_inside_nz(
+ lat, lon, nz_coast=nz_coast
+ ):
+ continue
+
+ current_channels: set[tuple[str, str]] = set()
+ for channel in station:
+ chan_id = (channel.location_code, channel.code[:2])
+ if chan_id in current_channels:
+ continue
+
+ current_channels.add(chan_id)
+ all_station_info.append(
+ [
+ provider,
+ network.code,
+ station.code,
+ lat,
+ lon,
+ getattr(station, "elevation", None),
+ channel.code[:2],
+ channel.location_code,
+ getattr(channel, "start_date", None),
+ getattr(channel, "end_date", None),
+ ]
+ )
+
+ return all_station_info
+
+
+def build_station_dataframe(rows: Iterable[list[object]]) -> pd.DataFrame:
+ """Build the station/channel dataframe.
+
+ Parameters
+ ----------
+ rows : Iterable[list[object]]
+ Station/channel rows.
+
+ Returns
+ -------
+ pandas.DataFrame
+ DataFrame containing station/channel metadata.
+ """
+ station_df = pd.DataFrame(
+ list(rows),
+ columns=[
+ "provider",
+ "net",
+ "sta",
+ "lat",
+ "lon",
+ "elev",
+ "chan",
+ "loc",
+ "start_date",
+ "end_date",
+ ],
+ )
+
+ if station_df.empty:
+ return station_df
+
+ return station_df.drop_duplicates(
+ ["provider", "net", "sta", "chan", "loc"]
+ ).reset_index(drop=True)
+
+
+def write_station_csv(df: pd.DataFrame, out_csv: Path) -> None:
+ """Write station/channel metadata to CSV.
+
+ Parameters
+ ----------
+ df : pandas.DataFrame
+ DataFrame to write.
+ out_csv : Path
+ Output file path.
+
+ Returns
+ -------
+ None
+ This function returns ``None``.
+ """
+ out_csv.parent.mkdir(parents=True, exist_ok=True)
+ df.to_csv(out_csv, index=False)
+
+
+@app.command()
+def get_stations(
+ out_all_csv: Annotated[
+ Path,
+ typer.Argument(
+ dir_okay=False,
+ help="Output CSV for all channels.",
+ ),
+ ],
+ out_filtered_csv: Annotated[
+ Path | None,
+ typer.Option(
+ dir_okay=False,
+ help=(
+ "Optional output CSV for desired channels only. If omitted, the filtered file is not written."
+ ),
+ ),
+ ] = None,
+ provider: Annotated[
+ list[str],
+ typer.Option(
+ "--provider",
+ help="Provider(s) to query. Repeatable. If omitted, queries all known providers.",
+ ),
+ ] = None,
+ starttime: Annotated[
+ str,
+ typer.Option(help="Start time for station metadata (e.g. 2000-01-01)."),
+ ] = "2000-01-01",
+ endtime: Annotated[
+ str,
+ typer.Option(help="End time for station metadata (or 'now')."),
+ ] = "now",
+ min_latitude: Annotated[float, typer.Option(help="Minimum latitude.")] = -49.0,
+ max_latitude: Annotated[float, typer.Option(help="Maximum latitude.")] = -32.0,
+ min_longitude: Annotated[float, typer.Option(help="Minimum longitude.")] = 165.0,
+ max_longitude: Annotated[float, typer.Option(help="Maximum longitude.")] = -176.9,
+ coastline_filter: Annotated[
+ bool,
+ typer.Option(
+ "--coastline-filter/--no-coastline-filter",
+ help="Filter to stations inside NZ coastline polygons.",
+ ),
+ ] = True,
+ desired_channels: Annotated[
+ list[str],
+ typer.Option(
+ "--desired-channel",
+ help="Desired 2-char channel prefixes for the filtered output. Repeatable.",
+ ),
+ ] = None,
+) -> None:
+ """Query station/channel metadata and write CSV outputs.
+
+ Parameters
+ ----------
+ out_all_csv : Path
+ Output CSV for all station/channel rows.
+ out_filtered_csv : Path | None
+ Optional output CSV for filtered station/channel rows.
+ provider : list[str], optional
+ Provider(s) to query. Repeatable.
+ starttime : str
+ Start time for station metadata.
+ endtime : str
+ End time for station metadata.
+ min_latitude : float
+ Minimum latitude.
+ max_latitude : float
+ Maximum latitude.
+ min_longitude : float
+ Minimum longitude.
+ max_longitude : float
+ Maximum longitude.
+ coastline_filter : bool
+ If True, filter to stations in NZ coastline polygons.
+ desired_channels : list[str], optional
+ Desired 2-character channel prefixes for ``out_filtered_csv``.
+ """
+ providers = _iter_providers(provider)
+
+ rows = collect_station_rows(
+ providers=providers,
+ starttime=_parse_time(starttime),
+ endtime=_parse_time(endtime),
+ minlatitude=min_latitude,
+ maxlatitude=max_latitude,
+ minlongitude=min_longitude,
+ maxlongitude=max_longitude,
+ coastline_filter=coastline_filter,
+ )
+
+ station_df = build_station_dataframe(rows)
+ write_station_csv(station_df, out_all_csv)
+
+ if out_filtered_csv is None:
+ return
+
+ desired = desired_channels if desired_channels else ["HH", "BH", "HN", "BN"]
+ filtered_df = station_df[station_df["chan"].isin(desired)].reset_index(drop=True)
+ write_station_csv(filtered_df, out_filtered_csv)
+
+
+if __name__ == "__main__":
+ app()
diff --git a/nzgmdb/temp_arrays/mass_download_data.py b/nzgmdb/temp_arrays/mass_download_data.py
new file mode 100644
index 00000000..7168736b
--- /dev/null
+++ b/nzgmdb/temp_arrays/mass_download_data.py
@@ -0,0 +1,485 @@
+import json
+from pathlib import Path
+from typing import Annotated
+
+import pandas as pd
+import typer
+from obspy import UTCDateTime
+from obspy.clients.fdsn.mass_downloader import (
+ GlobalDomain,
+ MassDownloader,
+ Restrictions,
+)
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+# Month length in seconds (15 days)
+MONTH_SECONDS = 15 * 24 * 3600
+# Minimum chunk size (seconds) when backing off after 413 / manifest-too-large.
+# 3600s = 1 hour.
+MIN_CHUNK_SECONDS: int = 3600
+
+
+def _format_end_for_filename(end_dt: UTCDateTime | str) -> str:
+ """Format an end datetime into the MiniSEED filename timestamp form.
+
+ Parameters
+ ----------
+ end_dt : UTCDateTime | str
+ An ObsPy ``UTCDateTime`` or any string parseable by ``UTCDateTime``.
+
+ Returns
+ -------
+ str
+ Timestamp string formatted like ``YYYYMMDDTHHMMSSZ``.
+ """
+ if not isinstance(end_dt, UTCDateTime):
+ end_dt = UTCDateTime(end_dt)
+ return end_dt.strftime("%Y%m%dT%H%M%SZ")
+
+
+def save_results(
+ results: list[dict[str, object] | object],
+ results_csv: Path,
+) -> None:
+ """Write a single results CSV from per-row download results.
+
+ This function is intentionally defensive:
+
+ - It computes the union of keys across result dictionaries.
+ - It normalizes missing keys to empty strings.
+ - It serializes nested objects (dict/list) as JSON for readability.
+
+ Parameters
+ ----------
+ results : list[dict[str, object] | object]
+ Download results. Most entries are dicts, but non-dicts will be written
+ under a ``value`` column.
+ results_csv : Path
+ Output CSV file.
+ """
+ if not results:
+ results_csv.parent.mkdir(parents=True, exist_ok=True)
+ pd.DataFrame().to_csv(results_csv, index=False)
+ return
+
+ all_keys: set[str] = set()
+ for r in results:
+ if isinstance(r, dict):
+ all_keys.update(r.keys())
+ else:
+ all_keys.add("value")
+
+ preferred = [
+ "idx",
+ "provider",
+ "net",
+ "sta",
+ "status",
+ "error",
+ "attempts",
+ "chunklength",
+ "timestamp",
+ "mseed_path",
+ "xml_path",
+ ]
+ cols = [k for k in preferred if k in all_keys] + sorted(
+ k for k in all_keys if k not in preferred
+ )
+
+ norm_rows: list[dict[str, object]] = []
+ for r in results:
+ if not isinstance(r, dict):
+ row: dict[str, object] = {"value": str(r)}
+ norm_rows.append(row)
+ continue
+
+ row = {}
+ for k in cols:
+ v = r.get(k, "")
+ if isinstance(v, (dict, list)):
+ v = json.dumps(v, ensure_ascii=False)
+ elif v is None:
+ v = ""
+ elif not isinstance(v, (str, int, float, bool)):
+ v = str(v)
+
+ row[k] = v
+ norm_rows.append(row)
+
+ results_csv.parent.mkdir(parents=True, exist_ok=True)
+
+ df = pd.DataFrame(norm_rows, columns=cols)
+ df.to_csv(results_csv, index=False, encoding="utf-8")
+
+
+def is_row_done(
+ row: pd.Series,
+ output_dir: Path,
+ mseed_dirname: str,
+) -> bool:
+ """Check whether a CSV row has fully completed waveform download.
+
+ A row is considered **done** if the expected MiniSEED output directory
+ contains at least one ``.mseed`` file whose final ``__END`` timestamp equals
+ the row's ``end_date``.
+
+ Notes
+ -----
+ This check is filesystem-driven and intentionally conservative:
+
+ - Only files whose channel code ends with ``"Z"`` are considered.
+ - Any filesystem error or unexpected filename format causes this function
+ to return ``False`` (so the row can be retried).
+
+ Parameters
+ ----------
+ row : pandas.Series
+ A row of the input CSV.
+ output_dir : Path
+ Root output directory.
+ mseed_dirname : str
+ Name of the MiniSEED subdirectory (typically ``"waveforms"``).
+
+ Returns
+ -------
+ bool
+ ``True`` if the row appears fully completed, otherwise ``False``.
+ """
+ net = str(row["net"]).strip()
+ sta = str(row["sta"]).strip()
+ loc_field = str(row["loc"]).strip()
+ chan_prefix = str(row["chan"]).strip()
+
+ record_sub = f"{net}_{sta}_{chan_prefix}_{loc_field}"
+ mseed_path = output_dir / mseed_dirname / net / record_sub
+
+ if not mseed_path.is_dir():
+ return False
+
+ target_end = _format_end_for_filename(row["end_date"])
+
+ try:
+ for path in mseed_path.iterdir():
+ if path.suffix != ".mseed":
+ continue
+
+ stem = path.stem
+
+ # Filename parts expected like: NET.STA..CHAN__START__END
+ parts = stem.split(".")
+ if len(parts) <= 3:
+ continue
+
+ chan_check = parts[3].split("__", 1)[0]
+ if not chan_check.endswith("Z"):
+ continue
+
+ if "__" not in stem:
+ continue
+
+ last = stem.rsplit("__", 1)[-1]
+ if last == target_end:
+ return True
+ except OSError:
+ return False
+
+ return False
+
+
+def create_output_dirs(
+ net: str,
+ sta: str,
+ chan_prefix: str,
+ loc: str,
+ output_dir: Path,
+ mseed_dirname: str,
+ stationxml_dirname: str,
+) -> tuple[Path, Path]:
+ """Create output directories for a network/station/channel/location.
+
+ Parameters
+ ----------
+ net : str
+ Network code.
+ sta : str
+ Station code.
+ chan_prefix : str
+ 2-character channel prefix (e.g. ``"HH"``).
+ loc : str
+ Location code used for naming the output directory.
+ output_dir : Path
+ Root output directory.
+ mseed_dirname : str
+ MiniSEED subdirectory name.
+ stationxml_dirname : str
+ StationXML subdirectory name.
+
+ Returns
+ -------
+ tuple[Path, Path]
+ Tuple of ``(mseed_path, xml_path)``.
+ """
+ record_sub = f"{net}_{sta}_{chan_prefix}_{loc}"
+
+ mseed_path = output_dir / mseed_dirname / net / record_sub
+ xml_path = output_dir / stationxml_dirname / net / record_sub
+
+ mseed_path.mkdir(parents=True, exist_ok=True)
+ xml_path.mkdir(parents=True, exist_ok=True)
+
+ return mseed_path, xml_path
+
+
+def download_task(
+ task: tuple[int, str, dict[str, object]],
+ *,
+ output_dir: Path,
+ mseed_dirname: str,
+ stationxml_dirname: str,
+ month_seconds: int,
+ min_chunk_seconds: int,
+) -> dict[str, int | str] | None:
+ """Download waveform and StationXML data for a single CSV row.
+
+ Parameters
+ ----------
+ task : tuple[int, str, dict[str, object]]
+ A 3-tuple of ``(idx, provider, row_dict)``.
+ output_dir : Path
+ Root output directory.
+ mseed_dirname : str
+ MiniSEED subdirectory name.
+ stationxml_dirname : str
+ StationXML subdirectory name.
+ month_seconds : int
+ Default chunk length (seconds) for large windows.
+ min_chunk_seconds : int
+ Minimum chunk length (seconds) when backing off after 413 errors.
+
+ Returns
+ -------
+ dict[str, int | str] | None
+ A result dictionary with keys including ``idx``, ``status``, ``provider``,
+ ``net``, ``sta``, and optionally ``error``.
+
+ Raises
+ ------
+ Exception
+ Re-raises exceptions encountered during download attempts when the
+ request cannot be reduced any further.
+ """
+ idx, provider, row = task
+
+ try:
+ net = str(row["net"]).strip()
+ sta = str(row["sta"]).strip()
+ loc_field = str(row["loc"]).strip()
+ loc = "*" if loc_field == "NA" else loc_field
+ chan_prefix = str(row["chan"]).strip()
+ channel = f"{chan_prefix}?"
+
+ start = UTCDateTime(row["start_date"])
+ end = UTCDateTime(row["end_date"])
+
+ total_window = end - start
+ chunk_base = int(min(total_window, month_seconds))
+
+ max_attempts = 4
+ attempt = 1
+
+ print(
+ f"[{idx}] Provider={provider} Downloading {net}.{sta} {channel} {start} -> {end} chunk={chunk_base}s"
+ )
+
+ mseed_path, xml_path = create_output_dirs(
+ net,
+ sta,
+ chan_prefix,
+ loc_field,
+ output_dir=output_dir,
+ mseed_dirname=mseed_dirname,
+ stationxml_dirname=stationxml_dirname,
+ )
+
+ while attempt <= max_attempts:
+ chunklength = int(min(total_window, chunk_base))
+ if chunklength < 1:
+ chunklength = 1
+
+ domain = GlobalDomain()
+ restrictions = Restrictions(
+ starttime=start,
+ endtime=end,
+ chunklength_in_sec=chunklength,
+ network=net,
+ station=sta,
+ location=loc,
+ channel=channel,
+ reject_channels_with_gaps=False,
+ minimum_length=0.0,
+ minimum_interstation_distance_in_m=0.0,
+ )
+
+ try:
+ mdl = MassDownloader(providers=[provider])
+ mdl.download(
+ domain,
+ restrictions,
+ mseed_storage=str(mseed_path),
+ stationxml_storage=str(xml_path),
+ )
+
+ print(f"[{idx}] Done")
+ return {
+ "idx": idx,
+ "status": "ok",
+ "provider": provider,
+ "net": net,
+ "sta": sta,
+ }
+ except Exception as exc: # noqa: BLE001
+ err_text = f"{exc!r} {exc}"
+ is_manifest_too_large = (
+ "Estimated manifest size" in err_text
+ or "Request Entity Too Large" in err_text
+ or "413" in err_text
+ )
+
+ if is_manifest_too_large:
+ if chunk_base <= min_chunk_seconds:
+ print(
+ f"[{idx}] Server denied request and chunk is already at minimum ({chunk_base}s). Giving up."
+ )
+ raise
+
+ old = chunk_base
+ chunk_base = max(min_chunk_seconds, chunk_base // 2)
+ print(
+ f"[{idx}] Server denied request (413). Reducing chunk base {old}s -> {chunk_base}s and retrying."
+ )
+ attempt += 1
+ continue
+
+ raise
+
+ except Exception as exc: # noqa: BLE001
+ print(
+ f"[{idx}] ERROR provider={provider} net={row.get('net')} sta={row.get('sta')}: {exc}"
+ )
+ return {
+ "idx": idx,
+ "status": "error",
+ "error": str(exc),
+ "provider": provider,
+ }
+
+
+@app.command()
+def run(
+ csv_file: Annotated[
+ Path,
+ typer.Argument(
+ ..., exists=True, dir_okay=False, help="Input station/channel CSV."
+ ),
+ ],
+ output_dir: Annotated[
+ Path,
+ typer.Argument(
+ ...,
+ file_okay=False,
+ help="Root output directory for waveforms and StationXML.",
+ ),
+ ],
+ results_csv: Annotated[
+ Path | None,
+ typer.Option(
+ dir_okay=False,
+ help="Optional output CSV for per-row download results (default: /download_results.csv).",
+ ),
+ ] = None,
+ mseed_dirname: Annotated[
+ str,
+ typer.Option(help="MiniSEED subdirectory name."),
+ ] = "waveforms",
+ stationxml_dirname: Annotated[
+ str,
+ typer.Option(help="StationXML subdirectory name."),
+ ] = "stationxml",
+) -> None:
+ """Download waveforms + StationXML for each row in a station/channel request table.
+
+ Parameters
+ ----------
+ csv_file : Path
+ Input CSV describing requests. Must contain columns:
+ ``net``, ``sta``, ``loc``, ``chan``, ``start_date``, ``end_date``, ``provider``.
+ output_dir : Path
+ Root output directory.
+ results_csv : Path | None, optional
+ Output CSV for results. If omitted, defaults to
+ ``/download_results.csv``.
+ mseed_dirname : str, optional
+ MiniSEED subdirectory name.
+ stationxml_dirname : str, optional
+ StationXML subdirectory name.
+
+ Returns
+ -------
+ None
+ This function returns ``None``.
+
+ Raises
+ ------
+ ValueError
+ If required columns are missing from the input CSV.
+ """
+ df = pd.read_csv(csv_file, dtype={"loc": str}, keep_default_na=False)
+
+ required_cols = {"net", "sta", "loc", "chan", "start_date", "end_date", "provider"}
+ missing = sorted(required_cols - set(df.columns))
+ if missing:
+ raise ValueError(f"CSV must contain: {missing}")
+
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ if results_csv is None:
+ results_csv = output_dir / "download_results.csv"
+
+ tasks: list[tuple[int, str, dict[str, object]]] = []
+ skipped = 0
+ for idx, row in df.reset_index(drop=True).iterrows():
+ if is_row_done(row, output_dir=output_dir, mseed_dirname=mseed_dirname):
+ skipped += 1
+ continue
+
+ tasks.append((int(idx), str(row["provider"]), row.to_dict()))
+
+ print(
+ f"Starting sequential run for {len(tasks)} tasks (skipped {skipped} already done)"
+ )
+
+ results: list[dict[str, object]] = []
+ for task in tasks:
+ results.append(
+ download_task(
+ task,
+ output_dir=output_dir,
+ mseed_dirname=mseed_dirname,
+ stationxml_dirname=stationxml_dirname,
+ month_seconds=MONTH_SECONDS,
+ min_chunk_seconds=MIN_CHUNK_SECONDS,
+ )
+ )
+
+ save_results(results, results_csv=results_csv)
+
+ oks = sum(1 for r in results if r.get("status") == "ok")
+ errs = sum(1 for r in results if r.get("status") == "error")
+ print(
+ f"Completed: {oks} succeeded, {errs} failed (skipped {skipped} previously done)"
+ )
+
+
+if __name__ == "__main__":
+ app()
diff --git a/pyproject.toml b/pyproject.toml
index ffe01990..25eeb5c4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,6 +12,10 @@ readme = "README.md"
requires-python = "==3.12.*"
dynamic = ["version", "dependencies"]
+[project.optional-dependencies]
+test = ["pytest"]
+dev = ["ty"]
+
[tool.setuptools.dynamic]
dependencies = {file = ["requirements.txt"]}
@@ -78,4 +82,12 @@ checks = [
"RT03",
"RT04",
"YD01",
-]
\ No newline at end of file
+]
+
+[tool.ty.rules]
+no-matching-overload = "ignore"
+unresolved-attribute = "ignore"
+not-subscriptable = "ignore"
+possibly-missing-submodule = "ignore"
+invalid-type-arguments = "ignore"
+unresolved-import = "ignore"
diff --git a/requirements.txt b/requirements.txt
index e987d2c4..a61a4ba7 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,6 @@
typer>0.12.3
pandas
-numpy<2.0.0
+numpy>=2.0
obspy
mseedlib
scipy
@@ -13,6 +13,7 @@ gmprocess
h5py
pooch
matplotlib
+rasterio
oq_wrapper>=2025.12.3,
qcore-utils>=2025.12.1,
source_modelling>=2025.12.1,
diff --git a/tests/test_sites.py b/tests/test_sites.py
new file mode 100644
index 00000000..b175eb04
--- /dev/null
+++ b/tests/test_sites.py
@@ -0,0 +1,355 @@
+from collections.abc import Iterator
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Self
+
+import numpy as np
+import pandas as pd
+import pytest
+from rasterio.io import MemoryFile
+from rasterio.transform import from_origin
+
+from nzgmdb.data_retrieval import sites
+
+
+class _DummyFionaCollection:
+ """
+ Minimal context manager that mimics a Fiona Collection.
+
+ Parameters
+ ----------
+ shapes : list
+ Iterable of shape-like records to yield when iterated.
+
+ Returns
+ -------
+ _DummyFionaCollection
+ Context manager instance that can be iterated over.
+ """
+
+ def __init__(self, shapes: list[Any]) -> None: # noqa: D107
+ self._shapes = shapes
+
+ def __enter__(self) -> Self: # noqa: D105
+ return self
+
+ def __exit__( # noqa: D105
+ self,
+ exc_type: type[BaseException] | None,
+ exc: BaseException | None,
+ tb: TracebackType | None,
+ ) -> bool:
+ return False
+
+ def __iter__(self) -> Iterator[Any]: # noqa: D105
+ return iter(self._shapes)
+
+
+def _make_test_geotiff(
+ width: int = 10,
+ height: int = 10,
+ nodata: float = -9999.0,
+) -> MemoryFile:
+ """
+ Create an in-memory single-band GeoTIFF for testing.
+
+ Parameters
+ ----------
+ width : int, default=10
+ Raster width in pixels.
+ height : int, default=10
+ Raster height in pixels.
+ nodata : float, default=-9999.0
+ NoData value written to the dataset (and to one pixel in the raster).
+
+ Returns
+ -------
+ memfile : rasterio.io.MemoryFile
+ In-memory GeoTIFF containing a simple gradient field with a NoData pixel.
+ """
+ data = np.arange(width * height, dtype=np.float32).reshape(height, width)
+ transform = from_origin(0.0, 10.0, 1.0, 1.0)
+ data[0, 0] = nodata
+
+ memfile = MemoryFile()
+ with memfile.open(
+ driver="GTiff",
+ height=height,
+ width=width,
+ count=1,
+ dtype=data.dtype,
+ crs="EPSG:4326",
+ transform=transform,
+ nodata=nodata,
+ ) as ds:
+ ds.write(data, 1)
+
+ return memfile
+
+
+def test_sample_points_from_geotiff_inside_outside_and_nodata() -> None:
+ """
+ Test that GeoTIFF sampling returns finite values for in-bounds points, and
+ returns NaN for out-of-bounds or NoData pixels.
+
+ Returns
+ -------
+ None
+ """
+ memfile = _make_test_geotiff()
+ with memfile.open() as ds:
+ points = np.array(
+ [
+ [9.5, 0.5],
+ [5.5, 5.5],
+ [20.0, 5.0],
+ [-5.0, 5.0],
+ [5.0, 20.0],
+ ],
+ dtype=float,
+ )
+ out = sites.sample_points_from_geotiff(ds.name, points).ravel()
+
+ assert out.shape == (len(points),)
+ assert np.isnan(out[0])
+ assert np.isfinite(out[1])
+ assert np.isnan(out[2])
+ assert np.isnan(out[3])
+ assert np.isnan(out[4])
+
+
+def test_fill_gaps_with_nearest_fills_nans_and_preserves_finite() -> None:
+ """
+ Test that `fill_gaps_with_nearest` fills NaN entries while preserving
+ existing finite values.
+
+ Returns
+ -------
+ None
+ """
+ coords = np.array(
+ [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]],
+ dtype=float,
+ )
+ values = np.array([10.0, 20.0, np.nan, 40.0], dtype=float)
+
+ filled = sites.fill_gaps_with_nearest(coords, values, k=3)
+
+ assert filled.shape == values.shape
+ assert np.isfinite(filled[2])
+ assert filled[0] == 10.0
+ assert filled[1] == 20.0
+ assert filled[3] == 40.0
+
+
+@pytest.mark.parametrize(
+ "points",
+ [
+ np.array(
+ [
+ [-36.8485, 174.7633],
+ [-41.2865, 174.7762],
+ [-43.5321, 172.6362],
+ [-45.0312, 168.6626],
+ ],
+ dtype=float,
+ ),
+ np.array(
+ [[-34.0, 172.5], [-47.5, 166.0], [-41.0, 179.9], [-41.0, 166.5]],
+ dtype=float,
+ ),
+ np.array(
+ [[-30.0, 174.0], [-50.0, 170.0], [-41.0, 160.0], [-41.0, -175.0]],
+ dtype=float,
+ ),
+ np.array(
+ [
+ [48.8566, 2.3522],
+ [51.5074, -0.1278],
+ [52.5200, 13.4050],
+ [41.9028, 12.4964],
+ ],
+ dtype=float,
+ ),
+ np.vstack(
+ [
+ np.array([[-36.8485, 174.7633], [-41.2865, 174.7762]], dtype=float),
+ np.array([[48.8566, 2.3522], [51.5074, -0.1278]], dtype=float),
+ np.array([[-43.5321, 172.6362], [-45.0312, 168.6626]], dtype=float),
+ ]
+ ),
+ ],
+)
+def test_site_updates_vs30_and_z1_fields_only_when_available(
+ monkeypatch: pytest.MonkeyPatch,
+ tmp_path: Path,
+ points: np.ndarray,
+) -> None:
+ """
+ Integration-style test of site table population logic using monkeypatches.
+
+ Verifies expected output columns exist and that Vs30 and Z1.0 / Z2.5 values
+ (and their reference/quality fields) are only set when data are available.
+
+ Parameters
+ ----------
+ monkeypatch : pytest.MonkeyPatch
+ Pytest fixture used to patch external dependencies.
+ tmp_path : pathlib.Path
+ Pytest fixture providing a temporary directory for test artifacts.
+ points : numpy.ndarray, shape (N, 2)
+ Input points as [lat, lon] used to construct a minimal metadata table.
+
+ Returns
+ -------
+ None
+ """
+
+ class _Cfg:
+ def get_value(self, key: str):
+ if key == "channel_codes":
+ return "HHZ"
+ if key == "bbox":
+ return [0.0, -90.0, 360.0, 90.0]
+ if key == "nzcvm_version":
+ return "test"
+ return None
+
+ monkeypatch.setattr(sites.cfg, "Config", _Cfg)
+
+ class _DummyInv(list):
+ pass
+
+ class _DummyClient:
+ def __init__(self, *_args, **_kwargs):
+ pass
+
+ def get_stations(self, *args, **kwargs):
+ return _DummyInv()
+
+ monkeypatch.setattr(sites, "FDSN_Client", _DummyClient)
+
+ monkeypatch.setattr(
+ sites.fiona,
+ "open",
+ lambda *_a, **_k: _DummyFionaCollection([]),
+ )
+
+ # Do NOT monkeypatch NZGMDB_DATA.abspath (read-only property).
+ # Instead, patch fetch() to return a temp file path for the tif and anything else requested.
+ combined_tif = tmp_path / "combined_mvn_wgs84.tif"
+ combined_tif.touch()
+
+ def _fetch(name: str, *args, **kwargs):
+ p = tmp_path / name
+ p.parent.mkdir(parents=True, exist_ok=True)
+ p.touch(exist_ok=True)
+ return str(p)
+
+ monkeypatch.setattr(sites.NZGMDB_DATA, "fetch", _fetch)
+
+ geo = pd.DataFrame(
+ {
+ "Name": [f"S{i}" for i in range(len(points))],
+ "Lat": points[:, 0],
+ "Long": points[:, 1],
+ "Elevation": np.zeros(len(points)),
+ "Vs30_median": [np.nan] * len(points),
+ "Sigmaln_Vs30": [np.nan] * len(points),
+ "T_median": [np.nan] * len(points),
+ "sigmaln_T": [np.nan] * len(points),
+ "Q_T": [None] * len(points),
+ "D_T": [None] * len(points),
+ "T_Ref": [None] * len(points),
+ "Z1.0_median": [np.nan] * len(points),
+ "sigmaln_Z1.0": [np.nan] * len(points),
+ "Z1.0_Ref": [None] * len(points),
+ "Z2.5_median": [np.nan] * len(points),
+ "sigmaln_Z2.5": [np.nan] * len(points),
+ "Z2.5_Ref": [None] * len(points),
+ "NZS1170SiteClass": [None] * len(points),
+ }
+ )
+ monkeypatch.setattr(pd, "read_csv", lambda *_a, **_k: geo)
+
+ def _find_domain_from_shapes(df: pd.DataFrame, _shapes: list) -> pd.DataFrame:
+ out = df.copy()
+ out["domain_no"] = 1
+ return out
+
+ monkeypatch.setattr(
+ sites.tect_domain, "find_domain_from_shapes", _find_domain_from_shapes
+ )
+
+ def _compute_station_thresholds(
+ stations: pd.DataFrame, model_version: str = None
+ ) -> pd.DataFrame:
+ n = len(stations)
+ return pd.DataFrame(
+ {
+ "Z1.0(km)": np.full(n, 0.5),
+ "Z2.5(km)": np.full(n, 2.0),
+ "sigma": np.full(n, 0.25),
+ },
+ index=stations.index,
+ )
+
+ monkeypatch.setattr(
+ sites.threshold, "compute_station_thresholds", _compute_station_thresholds
+ )
+
+ def _sample_points_from_geotiff(
+ file_path: str,
+ latlon_points: np.ndarray,
+ band: int = 1,
+ ) -> np.ndarray:
+ n = len(latlon_points)
+ vals = np.linspace(100.0, 500.0, n).astype(float)
+ if n >= 2:
+ vals[0] = np.nan
+ return vals.reshape(-1, 1)
+
+ monkeypatch.setattr(
+ sites, "sample_points_from_geotiff", _sample_points_from_geotiff
+ )
+
+ def _fill_gaps_with_nearest(
+ coords: np.ndarray,
+ values: np.ndarray,
+ invalid_mask: np.ndarray = None,
+ k: int = 8,
+ ) -> np.ndarray:
+ values = np.asarray(values, dtype=float).copy()
+ values[np.isnan(values)] = 250.0
+ return values
+
+ monkeypatch.setattr(sites, "fill_gaps_with_nearest", _fill_gaps_with_nearest)
+
+ site_df = sites.create_site_table_response()
+
+ for col in [
+ "sta",
+ "lat",
+ "lon",
+ "Vs30",
+ "Q_Vs30",
+ "Vs30_Ref",
+ "Z1.0",
+ "Z2.5",
+ "Z1.0_ref",
+ "Z2.5_ref",
+ "Q_Z1.0",
+ "Q_Z2.5",
+ ]:
+ assert col in site_df.columns
+
+ assert site_df["Z1.0"].notna().any()
+ assert site_df["Z2.5"].notna().any()
+ assert (site_df["Z1.0_ref"].dropna() == "NZCVM (2026)").all()
+ assert (site_df["Z2.5_ref"].dropna() == "NZCVM (2026)").all()
+
+ vs30_non_nan = site_df["Vs30"].notna()
+ assert (site_df.loc[vs30_non_nan, "Vs30_Ref"] == "Foster et al. (2019)").all()
+ assert (site_df.loc[vs30_non_nan, "Q_Vs30"] == "Q3").all()
+ assert site_df.loc[~vs30_non_nan, "Vs30_Ref"].isna().all()
+ assert site_df.loc[~vs30_non_nan, "Q_Vs30"].isna().all()
diff --git a/wiki/Waveform-Extraction.md b/wiki/Waveform-Extraction.md
index 8564bd17..30ad9a3c 100644
--- a/wiki/Waveform-Extraction.md
+++ b/wiki/Waveform-Extraction.md
@@ -21,6 +21,7 @@ python -m nzgmdb.scripts.run_nzgmdb extract-waveforms