1#!/usr/bin/env python
3"""
4This script prepares a seismic dataset from Lim et al. 2024 for use in the benchcompress project.
5It downloads seismic waveform data from Zenodo, extracts .sac files from a specific folder,
6and concatenates them into a single numpy array. The resulting dataset will be used to evaluate
7various compression algorithms in benchcompress to determine their effectiveness on real-world
8seismic waveform data.
10The original data is from:
11Lim, H., & Zhang, M. (2024). Machine Learning Phase Picker Training Dataset in Shanghai.
12Zenodo. https://doi.org/10.5281/zenodo.10457508
13"""
15import os
16import sys
17from urllib.request import urlretrieve
18import tarfile
19import glob
20from obspy import read
21import numpy as np
23def download_file(url, local_filename):
24 if os.path.exists(local_filename):
25 print(f"File {local_filename} already exists locally")
26 return True
28 print(f"Downloading {url} to {local_filename}...")
29 try:
30 urlretrieve(url, local_filename)
31 print("Download completed successfully")
32 return True
33 except Exception as e:
34 print(f"Error downloading file: {e}")
35 return False
37def extract_data_folder(archive_path):
38 if not os.path.exists(archive_path):
39 print(f"Archive file {archive_path} not found")
40 return False
42 print(f"Extracting /data/01 folder from {archive_path}...")
43 try:
44 with tarfile.open(archive_path, 'r:gz') as tar:
45 # Extract only files in the /data/01 directory
46 members = [m for m in tar.getmembers() if m.name.startswith('data/01/')]
47 for member in members:
48 tar.extract(member)
49 print("Extraction completed successfully")
50 return True
51 except Exception as e:
52 print(f"Error extracting archive: {e}")
53 return False
55def load_and_concatenate_sac_files(directory):
56 # Find all .sac files in the directory
57 sac_files = glob.glob(f"{directory}/*.sac")
58 if not sac_files:
59 print(f"No .sac files found in {directory}")
60 return None
62 print(f"Found {len(sac_files)} .sac files")
64 all_data = []
65 for i, filename in enumerate(sac_files, 1):
66 try:
67 print(f"Loading file {i}/{len(sac_files)}: {filename}")
68 st = read(filename)
69 tr = st[0] # Get the first trace
70 all_data.append(tr.data)
71 print(f" Shape: {tr.data.shape}, dtype: {tr.data.dtype}")
72 except Exception as e:
73 print(f"Error loading {filename}: {e}")
74 continue
76 if not all_data:
77 print("No data was successfully loaded")
78 return None
80 # Concatenate all data arrays
81 concatenated_data = np.concatenate(all_data)
82 print(f"\nConcatenation complete!")
83 print(f"Final array shape: {concatenated_data.shape}")
84 print(f"Final array dtype: {concatenated_data.dtype}")
86 return concatenated_data
88def main():
89 url = "https://zenodo.org/records/10457508/files/data.tar.gz?download=1"
90 local_file = "lim_et_al_2024.zip"
91 sac_directory = "data/01"
92 output_file = "lim_et_al_2024.01.concat.npy"
94 # Download if needed
95 if not download_file(url, local_file):
96 print("Failed to download file")
97 sys.exit(1)
99 # Extract data folder
100 if not extract_data_folder(local_file):
101 print("Failed to extract data folder")
102 sys.exit(1)
104 # Load and concatenate SAC files
105 concatenated_data = load_and_concatenate_sac_files(sac_directory)
106 if concatenated_data is None:
107 print("Failed to process SAC files")
108 sys.exit(1)
110 # Save concatenated data
111 print(f"\nSaving concatenated data to {output_file}...")
112 try:
113 np.save(output_file, concatenated_data)
114 print("Data saved successfully")
115 except Exception as e:
116 print(f"Error saving data: {e}")
117 sys.exit(1)
119if __name__ == "__main__":
120 main()