/ concept-collection / benchcompress
Sign in
concept-collection / benchcompress
87 lines · 2.4 KBBlameHistoryRaw
1import numpy as np
2import segyio
3import os
4import requests
7SOURCE_FILE = "seismic/__init__.py"
10def _load_long_description():
11 current_dir = os.path.dirname(os.path.abspath(__file__))
12 md_path = os.path.join(current_dir, "seismic.md")
13 with open(md_path, "r", encoding="utf-8") as f:
14 return f.read()
17LONG_DESCRIPTION = _load_long_description()
19tags = ["real", "seismic", "continuous", "timeseries", "1d"]
20tags_float = tags + ["float"]
21tags_integer = tags + ["integer"]
24def _load_seismic_data() -> np.ndarray:
25 """Load seismic data from the SEG-Y file.
27 Returns:
28 Array containing the loaded seismic data
29 """
30 file_path = "04A+04B.segy"
31 if not os.path.exists(file_path):
32 # Download the SEG-Y file
33 url = "https://zenodo.org/records/8152964/files/04A+04B.segy?download=1"
34 response = requests.get(url)
35 with open(file_path, "wb") as f:
36 f.write(response.content)
37 print(f"Downloaded {file_path}")
38 else:
39 print(f"{file_path} already exists locally.")
41 # Open the SEG-Y file
42 with segyio.open(file_path, "r", ignore_geometry=True) as f:
43 # Read the seismic data
44 data = f.trace.raw[:]
46 # Consider only the first 3000 traces, because the others have a bunch of zeros
47 X = data[:3000]
48 # The first part of each trace is zeros
49 first_nonzero_indices = []
50 for j in range(X.shape[0]):
51 inds = np.where(X[j] != 0)[0]
52 first_nonzero_indices.append(inds[0] if len(inds) > 0 else -1)
53 # plt.figure(figsize=(10, 5))
54 # plt.hist(first_nonzero_indices, bins=20)
55 # print(np.max(first_nonzero_indices)) # 1607
56 X = X[:, 1700:]
58 return X.ravel()
61def _load_quantized_seismic_data():
62 X = _load_seismic_data()
63 step = 10000
64 X = np.round(X / step).astype(np.int32)
65 return X
68datasets = [
69 {
70 "name": "seismic-04A-04B",
71 "version": "1",
72 "description": "Seismic data from Roger Revelle voyage RR1508.",
73 "create": lambda: _load_seismic_data(),
74 "tags": tags_float,
75 "source_file": SOURCE_FILE,
76 "long_description": LONG_DESCRIPTION,
77 },
78 {
79 "name": "seismic-04A-04B-quantized",
80 "version": "1",
81 "description": "Seismic data from Roger Revelle voyage RR1508, quantized.",
82 "create": lambda: _load_quantized_seismic_data(),
83 "tags": tags_integer,
84 "source_file": SOURCE_FILE,
85 "long_description": LONG_DESCRIPTION,
86 },
moveopenescclose