py-nfi 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
py_nfi-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,3 @@
1
+ Metadata-Version: 2.4
2
+ Name: py-nfi
3
+ Version: 0.1.0
py_nfi-0.1.0/README.md ADDED
@@ -0,0 +1,4 @@
1
+ # py-nfi
2
+
3
+
4
+ ...
@@ -0,0 +1,11 @@
1
+ [build-system]
2
+ requires = ["setuptools>=64"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "py-nfi"
7
+ version = "0.1.0"
8
+ # ...
9
+
10
+ [tool.setuptools.packages.find]
11
+ where = ["src"]
py_nfi-0.1.0/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1 @@
1
+ __version__ = "0.1.0"
@@ -0,0 +1,1221 @@
1
+ #!/usr/bin/env python3
2
+
3
+ """
4
+ beta_compute.py
5
+
6
+ Optimized version of beta_compute.py targeting:
7
+ 1) Memory reduction (~3-5x for 10M+ records)
8
+ 2) Multiprocessing of compute_dlogbeta via shared memory
9
+
10
+ Key changes from original:
11
+ - target_data replaced with contiguous numpy arrays (eliminates dict-of-arrays overhead)
12
+ - metadata_calib stored as raw arrays, not a DataFrame copy
13
+ - group_events / explode_events use index-based mapping instead of list-in-cell DataFrames
14
+ - Uncertainty bookkeeping uses sparse storage (only events that produce results)
15
+ - compute_dlogbeta parallelized with multiprocessing.Pool + shared memory arrays
16
+ - np.isin replaced with pre-built set lookups where beneficial
17
+ - convert_utc_to_posix_ns uses vectorized pd.to_datetime instead of per-row obspy UTC
18
+ - groupby uses observed=True to avoid Cartesian product explosion with categoricals
19
+ - changed dlogbeta_corr to nfi
20
+ - save_fwf() now sorts by edatetime when available, otherwise event_name
21
+ - added 'nearest-neighbor' option for calib_search_method
22
+ - Now, explosion events are dropped before calibration event selection
23
+ - Added 'source-separate' option for mag_corr_method: separate magnitude corrections for each source
24
+ - Changed magnitude correction to extend to lowest magnitude of catalog, instead of dM/2 + Mmin
25
+
26
+ Last Modified:
27
+ 2026-04-21
28
+
29
+ Future improvements:
30
+ - Use nearest neighbors for calibration events
31
+ """
32
+
33
+ import numpy as np
34
+ import pandas as pd
35
+ from tqdm import trange, tqdm
36
+ import time
37
+ import os
38
+ import pickle as pkl
39
+ import inspect
40
+ import multiprocessing as mp
41
+ from multiprocessing import shared_memory
42
+ from functools import partial
43
+
44
+ import matplotlib.pyplot as plt
45
+ from obspy import UTCDateTime as UTC
46
+
47
+
48
+ # ============================================================
49
+ # Module-level worker function (must be picklable for mp.Pool)
50
+ # ============================================================
51
+
52
+ def _dlogbeta_worker(event_indices, shm_names, shm_shapes, shm_dtypes,
53
+ calib_n_records_min, calib_hdist_max, calib_zdist_max,
54
+ calib_time_filter, calib_time_method, calib_time_ndays,
55
+ compute_uncertainty, confidence_level,
56
+ save_full_uncertainty, calib_search_method, calib_nn,
57
+ calib_depth_scale
58
+ ):
59
+ """
60
+ Worker function for parallel dlogbeta computation.
61
+
62
+ Attaches to shared memory arrays, processes a chunk of events,
63
+ returns results as dicts.
64
+ """
65
+ # Attach to shared memory
66
+ shm_handles = {}
67
+ arrays = {}
68
+ for name in shm_names:
69
+ shm = shared_memory.SharedMemory(name=shm_names[name])
70
+ arr = np.ndarray(shm_shapes[name], dtype=shm_dtypes[name], buffer=shm.buf)
71
+ shm_handles[name] = shm
72
+ arrays[name] = arr
73
+
74
+ # Unpack arrays
75
+ # Target event arrays (nevents x ...)
76
+ t_edep = arrays['t_edep']
77
+ t_elat = arrays['t_elat']
78
+ t_elon = arrays['t_elon']
79
+ t_rec_start = arrays['t_rec_start'] # start index into record arrays
80
+ t_rec_count = arrays['t_rec_count'] # number of records for this event
81
+
82
+ # Per-record arrays (nrecords)
83
+ r_cid = arrays['r_cid']
84
+ r_logbeta = arrays['r_logbeta']
85
+
86
+ # Calibration arrays (ncalib_records)
87
+ c_edep = arrays['c_edep']
88
+ c_elat = arrays['c_elat']
89
+ c_elon = arrays['c_elon']
90
+ c_cid = arrays['c_cid']
91
+ c_logbeta = arrays['c_logbeta']
92
+ c_slat = arrays['c_slat']
93
+ c_slon = arrays['c_slon']
94
+
95
+ # NN-specific arrays
96
+ if calib_search_method == 'nearest-neighbor':
97
+ c_event_idx = arrays['c_event_idx'] # record -> unique calib event index
98
+ cu_elat = arrays['cu_elat']
99
+ cu_elon = arrays['cu_elon']
100
+ cu_edep = arrays['cu_edep']
101
+
102
+ # Optional time arrays
103
+ if calib_time_filter:
104
+ t_etime = arrays['t_etime']
105
+ c_etime = arrays['c_etime']
106
+
107
+ # Confidence intervals
108
+ if compute_uncertainty:
109
+ alpha = 1 - confidence_level
110
+ lower_percentile = 100 * (alpha / 2)
111
+ upper_percentile = 100 * (1 - alpha / 2)
112
+
113
+ # Results for this chunk
114
+ results = []
115
+
116
+ for ev_idx in event_indices:
117
+ start = t_rec_start[ev_idx]
118
+ count = t_rec_count[ev_idx]
119
+ if count == 0:
120
+ continue
121
+
122
+ target_edep = t_edep[ev_idx]
123
+ target_elat = t_elat[ev_idx]
124
+ target_elon = t_elon[ev_idx]
125
+ target_cids = r_cid[start:start+count]
126
+ target_logbeta = r_logbeta[start:start+count]
127
+
128
+ # ---- CALIBRATION EVENT SEARCH ----
129
+ # Depth filter
130
+ depth_mask = (c_edep >= (target_edep - calib_zdist_max)) & \
131
+ (c_edep <= (target_edep + calib_zdist_max))
132
+
133
+ # Channel overlap filter
134
+ target_cid_set = set(target_cids)
135
+ cid_mask = np.array([cid in target_cid_set for cid in c_cid[depth_mask]], dtype=bool)
136
+
137
+ # Get indices into original calib arrays that pass depth
138
+ depth_indices = np.where(depth_mask)[0]
139
+ combined_indices = depth_indices[cid_mask]
140
+
141
+ if len(combined_indices) == 0:
142
+ continue
143
+
144
+ # Time filter
145
+ if calib_time_filter and calib_time_method == 'interval':
146
+ target_etime = t_etime[ev_idx]
147
+ time_delta_days = np.abs(c_etime[combined_indices] - target_etime) / 86400.0 / 1E9
148
+ time_mask = time_delta_days <= calib_time_ndays
149
+ combined_indices = combined_indices[time_mask]
150
+ if len(combined_indices) == 0:
151
+ continue
152
+
153
+ # Distance filter
154
+ dists = _haversine_km(
155
+ np.full(len(combined_indices), target_elat),
156
+ np.full(len(combined_indices), target_elon),
157
+ c_elat[combined_indices],
158
+ c_elon[combined_indices]
159
+ )
160
+ dist_mask = dists <= calib_hdist_max
161
+
162
+ if calib_search_method == 'cylinder':
163
+ if dist_mask.sum() < calib_n_records_min:
164
+ continue
165
+ calib_mask_indices = combined_indices[dist_mask]
166
+
167
+ elif calib_search_method == 'nearest-neighbor':
168
+ combined_indices = combined_indices[dist_mask]
169
+ if len(combined_indices) < calib_n_records_min:
170
+ continue
171
+
172
+ # Compute scaled 3D distance per record (via parent event)
173
+ rec_event_idx = c_event_idx[combined_indices]
174
+ ce_elat = cu_elat[rec_event_idx]
175
+ ce_elon = cu_elon[rec_event_idx]
176
+ ce_edep = cu_edep[rec_event_idx]
177
+
178
+ hdists_ev = _haversine_km(
179
+ np.full(len(combined_indices), target_elat),
180
+ np.full(len(combined_indices), target_elon),
181
+ ce_elat, ce_elon
182
+ )
183
+ vdists_ev = np.abs(ce_edep - target_edep) * calib_depth_scale
184
+ rec_scaled_dists = np.sqrt(hdists_ev**2 + vdists_ev**2)
185
+
186
+ calib_mask_indices = combined_indices
187
+
188
+ cc_cid = c_cid[calib_mask_indices]
189
+ cc_logbeta = c_logbeta[calib_mask_indices]
190
+
191
+ if calib_search_method == 'nearest-neighbor':
192
+ cc_rec_dists = rec_scaled_dists # already aligned with calib_mask_indices
193
+
194
+ # ---- MATCH CALIBRATION & TARGET RECORDS ----
195
+
196
+ mask_t = np.isin(target_cids, cc_cid)
197
+ t_cid = target_cids[mask_t]
198
+ t_lb = target_logbeta[mask_t]
199
+
200
+ sorter = np.argsort(t_cid)
201
+ t_cid_sorted = t_cid[sorter]
202
+ t_lb_sorted = t_lb[sorter]
203
+
204
+ v = np.searchsorted(t_cid_sorted, cc_cid)
205
+ differences = t_lb_sorted[v] - cc_logbeta
206
+
207
+ # Group by channel
208
+ unique_channels, inverse_indices = np.unique(cc_cid, return_inverse=True)
209
+ sorted_order = np.argsort(inverse_indices)
210
+ sorted_diffs = differences[sorted_order]
211
+ sorted_inverse = inverse_indices[sorted_order]
212
+
213
+ # If NN mode, also sort the distances for per-channel trimming
214
+ if calib_search_method == 'nearest-neighbor':
215
+ sorted_rec_dists = cc_rec_dists[sorted_order]
216
+
217
+ split_indices = np.r_[
218
+ 0,
219
+ np.where(np.diff(sorted_inverse) != 0)[0] + 1,
220
+ len(sorted_inverse)
221
+ ]
222
+ ngroups = len(split_indices) - 1
223
+
224
+ # Compute per-channel medians and counts
225
+ dlogbeta_j = np.empty(ngroups, dtype=np.float64)
226
+ ncalib_j = np.empty(ngroups, dtype=np.int64)
227
+ for g in range(ngroups):
228
+ s, e = split_indices[g], split_indices[g+1]
229
+ grp = sorted_diffs[s:e]
230
+
231
+ if calib_search_method == 'nearest-neighbor' and len(grp) > calib_nn:
232
+ # Keep only the nearest calib_nn records for this channel
233
+ grp_dists = sorted_rec_dists[s:e]
234
+ nearest = np.argpartition(grp_dists, calib_nn)[:calib_nn]
235
+ grp = grp[nearest]
236
+
237
+ dlogbeta_j[g] = np.median(grp)
238
+ ncalib_j[g] = len(grp)
239
+
240
+ final_mask = ncalib_j >= calib_n_records_min
241
+ if final_mask.sum() < 3:
242
+ continue
243
+
244
+ dlogbeta_j_kept = dlogbeta_j[final_mask]
245
+ ncalib_j_kept = ncalib_j[final_mask]
246
+ dlogbeta = np.median(dlogbeta_j_kept)
247
+
248
+ result_entry = {
249
+ 'ev_idx': ev_idx,
250
+ 'dlogbeta': dlogbeta,
251
+ }
252
+
253
+ if compute_uncertainty:
254
+ # Compute residuals from kept groups
255
+ kept_group_indices = np.where(final_mask)[0]
256
+ residuals = []
257
+ for g in kept_group_indices:
258
+ grp = sorted_diffs[split_indices[g]:split_indices[g+1]]
259
+ residuals.append(grp - np.median(grp))
260
+ all_residuals = np.concatenate(residuals)
261
+ samples = all_residuals + dlogbeta
262
+
263
+ result_entry['dlogbeta_std'] = np.std(samples)
264
+ result_entry['dlogbeta_lower'] = np.percentile(samples, lower_percentile)
265
+ result_entry['dlogbeta_upper'] = np.percentile(samples, upper_percentile)
266
+ result_entry['dlogbeta_median'] = np.median(samples)
267
+
268
+ if save_full_uncertainty:
269
+ # Store per-station details
270
+ cc_slat = c_slat[calib_mask_indices]
271
+ cc_slon = c_slon[calib_mask_indices]
272
+ sorted_cc_slat = cc_slat[sorted_order]
273
+ sorted_cc_slon = cc_slon[sorted_order]
274
+
275
+ rep_slats = np.array([sorted_cc_slat[split_indices[g]] for g in kept_group_indices])
276
+ rep_slons = np.array([sorted_cc_slon[split_indices[g]] for g in kept_group_indices])
277
+
278
+ bearings = _get_bearing(
279
+ np.full(len(rep_slats), target_elat),
280
+ np.full(len(rep_slons), target_elon),
281
+ rep_slats, rep_slons
282
+ )
283
+ cdists = _haversine_km(
284
+ np.full(len(rep_slats), target_elat),
285
+ np.full(len(rep_slons), target_elon),
286
+ rep_slats, rep_slons
287
+ )
288
+
289
+
290
+ result_entry['dlogbeta_j'] = dlogbeta_j_kept.tolist()
291
+ result_entry['ncalib_j'] = ncalib_j_kept.tolist()
292
+ result_entry['bearing'] = bearings.tolist()
293
+ result_entry['deldist'] = cdists.tolist()
294
+ result_entry['slat'] = rep_slats.tolist()
295
+ result_entry['slon'] = rep_slons.tolist()
296
+ result_entry['dlogbeta_j_residual'] = all_residuals.tolist()
297
+
298
+ results.append(result_entry)
299
+
300
+ # Detach shared memory (don't unlink — the parent owns it)
301
+ for shm in shm_handles.values():
302
+ shm.close()
303
+
304
+ return results
305
+
306
+ class BetaEstimator:
307
+ def __init__(self,
308
+ df_records,
309
+ spectra,
310
+ f,
311
+ low_window_desired = (1.0, 5.0),
312
+ high_window_desired = (15.0, 22.0),
313
+ calib_mag_range = (1.4, 1.6),
314
+ calib_hdist_max = 10.0,
315
+ calib_zdist_max = 2.0,
316
+ calib_n_records_min = 10,
317
+ mag_corr_method = 'smoothedspline',
318
+ mag_corr_dM = 0.2,
319
+ quiet = False,
320
+ print_time = True,
321
+ save_dir = None,
322
+ compute_uncertainty = True,
323
+ confidence_level = 0.95,
324
+ calib_time_filter = False,
325
+ calib_time_method = 'interval',
326
+ calib_time_ndays = 180,
327
+ calib_time_n_nearest = 25,
328
+ n_workers = None,
329
+ save_full_uncertainty = False,
330
+ calib_search_method = 'cylinder',
331
+ calib_nn = 10,
332
+ calib_depth_scale = 5,
333
+ ):
334
+ t0 = time.time()
335
+
336
+ if not quiet: print("Initializing BetaEstimator")
337
+ if not quiet: print("--------------------------")
338
+
339
+ # These are parameters that shouldn't change results (assuming
340
+ # df_records and spectra are the same)
341
+ excl = ['self', 'df_records', 'spectra', 'f', 'quiet', 'save_dir',
342
+ 'print_time', 'n_workers']
343
+ sig = inspect.signature(self.__init__)
344
+ self.input_parameter_names = [k for k in sig.parameters if k not in excl]
345
+
346
+ # Store parameters as attributes
347
+ self.df_records = df_records.copy()
348
+ self.spectra = spectra
349
+ self.f = f
350
+ self.df = f[1] - f[0]
351
+ self.low_window_desired = low_window_desired
352
+ self.high_window_desired = high_window_desired
353
+ self.calib_mag_range = calib_mag_range
354
+ self.calib_hdist_max = calib_hdist_max
355
+ self.calib_zdist_max = calib_zdist_max
356
+ self.calib_n_records_min = calib_n_records_min
357
+ self.mag_corr_method = mag_corr_method
358
+ self.mag_corr_dM = mag_corr_dM
359
+ self.quiet = quiet
360
+ self.save_dir = save_dir
361
+ self.compute_uncertainty = compute_uncertainty
362
+ self.confidence_level = confidence_level
363
+ self.print_time = print_time
364
+ self.calib_time_filter = calib_time_filter
365
+ self.calib_time_method = calib_time_method
366
+ self.calib_time_ndays = calib_time_ndays
367
+ self.calib_time_n_nearest = calib_time_n_nearest
368
+ self.save_full_uncertainty = save_full_uncertainty
369
+ self.calib_search_method = calib_search_method
370
+ self.calib_nn = calib_nn
371
+ self.calib_depth_scale = calib_depth_scale
372
+
373
+ # Number of workers for parallel computation
374
+ if n_workers is None:
375
+ self.n_workers = max(1, mp.cpu_count() // 2)
376
+ else:
377
+ self.n_workers = n_workers
378
+
379
+ # Validate input
380
+ self.validate_input()
381
+ if not self.quiet: self.print_metadata_information()
382
+
383
+ # Do some basic processing before computing logbeta
384
+ self.digitize_names()
385
+ self.compute_column_dependencies()
386
+
387
+ # Calculate actual indices and frequency bands for beta computation
388
+ self.compute_frequency_bands()
389
+ if not self.quiet: self.print_frequency_information()
390
+
391
+ # Get indices of calibration events
392
+ self.store_calibration_events()
393
+ if not self.quiet: self.print_calibration_information()
394
+
395
+ self.fprint("STATUS:")
396
+ self.fprint(f"BetaEstimator initialized in {time.time()-t0:.4f} s. "
397
+ f"Using {self.n_workers} workers. Run 'compute()' to continue.")
398
+ self.fprint("--------------------------------------------------------------------")
399
+
400
+ def update_nrec(self):
401
+ self.df_events['nrec'] = self.df_events['_cid'].apply(len)
402
+
403
+ def store_calibration_events(self):
404
+ self.calibration_inds = np.where(np.logical_and(
405
+ self.df_records['emag'].values >= self.calib_mag_range[0],
406
+ self.df_records['emag'].values <= self.calib_mag_range[1]
407
+ ))[0]
408
+ self.metadata_calib = self.df_records.iloc[self.calibration_inds].reset_index(drop=True)
409
+
410
+ if "etype" in self.metadata_calib.columns:
411
+ l0 = len(self.metadata_calib)
412
+ self.metadata_calib = self.metadata_calib[np.logical_and(
413
+ self.metadata_calib['etype']!='qb',
414
+ self.metadata_calib['etype']!='ex'
415
+ )].reset_index(drop=True)
416
+
417
+ print(f"Dropped {l0-len(self.metadata_calib):,} quarry blast "\
418
+ "and explosion calibration events. Remaining: "\
419
+ f"{len(self.metadata_calib):,} calibration events.")
420
+
421
+ def compute(self, recompute=False):
422
+ self.fprint('')
423
+ self.fprint("Computing nFI")
424
+ self.fprint("-----------------")
425
+ params_filepath = os.path.join(self.save_dir, 'params.logbeta')
426
+ data_filepath = os.path.join(self.save_dir, 'data.logbeta')
427
+
428
+ files_exist = os.path.exists(params_filepath) and os.path.exists(data_filepath)
429
+
430
+ if not files_exist:
431
+ self.fprint(f"Results not found in {self.save_dir} \nComputing.")
432
+ do_computation = True
433
+ else:
434
+ if recompute:
435
+ self.fprint("Recomputing & overwriting previous results.")
436
+ do_computation = True
437
+ else:
438
+ t0 = time.time()
439
+ current_params = {key: getattr(self, key) for key in self.input_parameter_names}
440
+ loaded_params = self.load_parameters()
441
+ if current_params == loaded_params:
442
+ self.fprint(f"Parameters unchanged from {params_filepath} \nSkipping computation.")
443
+ do_computation = False
444
+ else:
445
+ self.fprint(f"Parameters changed from {params_filepath} \nRecomputing.")
446
+ do_computation = True
447
+
448
+ if do_computation:
449
+ self.compute_logbeta()
450
+ self.store_calibration_events()
451
+ self.group_events()
452
+ self.compute_dlogbeta()
453
+ self.apply_magnitude_correction(corr_type=self.mag_corr_method)
454
+
455
+ self.store_calibration_events()
456
+ if not self.quiet: self.print_calibration_information()
457
+ self.update_nrec()
458
+ self.save_parameters()
459
+ self.store_edatetimes()
460
+ self.save_fwf()
461
+ self.save_data()
462
+
463
+ return self
464
+ else:
465
+ t0 = time.time()
466
+ self.fprint(f"Loading data & parameters from {self.save_dir}")
467
+ loaded_instance = self.load_data(self.save_dir)
468
+ self.load_parameters()
469
+ self.fprint(f"Done. ({time.time()-t0:.2f}s)")
470
+ return loaded_instance
471
+
472
+ def save_data(self):
473
+ t0 = time.time()
474
+ self.fprint(f"Saving data to {self.save_dir}/data.logbeta")
475
+ with open(f"{self.save_dir}/data.logbeta", 'wb') as fs:
476
+ pkl.dump(self, fs)
477
+ self.tprint(f" |---> ({time.time()-t0:.2f}s)")
478
+
479
+ def store_edatetimes(self):
480
+
481
+ pass
482
+
483
+ def save_fwf(self):
484
+ t0 = time.time()
485
+ self.fprint(f"Saving fixed-width format results to {self.save_dir}/logbeta.txt")
486
+
487
+ col_fmts = {
488
+ 'event_name': '%12s',
489
+ 'edatetime': '%27s',
490
+ 'etype': '%2s',
491
+ 'emag': '%5.2f',
492
+ 'emagtype': '%2s',
493
+ 'elon': '%10.5f',
494
+ 'elat': '%9.5f',
495
+ 'edep': '%7.3f',
496
+ 'nrec': '%5i',
497
+ 'dlogbeta': '%9.5f',
498
+ 'dlogbeta_std': '%8.3e',
499
+ 'dlogbeta_lower':'%8.3e',
500
+ 'dlogbeta_upper':'%8.3e',
501
+ 'nfi': '%9.5f',
502
+ }
503
+
504
+ # Filter to columns that exist in the output dataframe
505
+ available = {c: f for c, f in col_fmts.items() if c in self.df_events.columns}
506
+ cols = list(available.keys())
507
+ fmts = list(available.values())
508
+
509
+ # Sort by time if available, otherwise by name
510
+ sort_col = 'edatetime' if 'edatetime' in cols else 'event_name'
511
+ df_out = self.df_events[cols].sort_values(by=sort_col).reset_index(drop=True)
512
+ self.fprint(f"Output sorted by {sort_col}")
513
+
514
+ df_out['edatetime'] = np.datetime_as_string(df_out['edatetime'].values, unit='ms')
515
+
516
+ np.savetxt(
517
+ f"{self.save_dir}/logbeta.txt",
518
+ df_out.values,
519
+ fmt=" ".join(fmts),
520
+ header=' '.join(cols),
521
+ comments='',
522
+ )
523
+ self.tprint(f" |---> ({time.time()-t0:.2f}s)")
524
+
525
+ @staticmethod
526
+ def load_data(save_dir):
527
+ with open(f"{save_dir}/data.logbeta", 'rb') as fs:
528
+ loaded_instance = pkl.load(fs)
529
+ return loaded_instance
530
+
531
+ def save_parameters(self):
532
+ t0 = time.time()
533
+ self.fprint(f"Saving parameters to {self.save_dir}/params.logbeta")
534
+ params = {key: getattr(self, key) for key in self.input_parameter_names}
535
+ with open(f"{self.save_dir}/params.logbeta", 'wb') as fs:
536
+ pkl.dump(params, fs)
537
+ self.tprint(f" |---> ({time.time()-t0:.2f}s)")
538
+
539
+ def load_parameters(self):
540
+ self.fprint(f"Loading parameters from {self.save_dir}/params.logbeta")
541
+ with open(f"{self.save_dir}/params.logbeta", 'rb') as fs:
542
+ params = pkl.load(fs)
543
+ return params
544
+
545
+ def convert_utc_to_posix_ns(self):
546
+ """Convert edatetime column to POSIX nanoseconds (int64).
547
+
548
+ Handles three input types efficiently:
549
+ 1) obspy UTCDateTime objects → use .ns attribute
550
+ 2) datetime strings → vectorized pd.to_datetime (~100x faster than per-row UTC())
551
+ 3) already numeric → pass through
552
+ """
553
+ col = self.df_records['edatetime']
554
+
555
+ # Check the first non-null value to determine type
556
+ first_val = col.iloc[0]
557
+
558
+ if hasattr(first_val, 'ns'):
559
+ # obspy UTCDateTime objects — use .ns (original behavior)
560
+ self.df_records['etime'] = col.apply(lambda x: x.ns)
561
+ elif isinstance(first_val, (int, np.integer)):
562
+ # Already numeric
563
+ self.df_records['etime'] = col.values.astype(np.int64)
564
+ else:
565
+ # String or datetime-like — use vectorized pandas parsing
566
+ # pd.to_datetime is ~100x faster than per-row obspy UTC()
567
+ dt = pd.to_datetime(col)
568
+ self.df_records['etime'] = dt.astype(np.int64) # nanoseconds since epoch
569
+
570
+ # self.df_records.drop(columns=['edatetime'], inplace=True)
571
+
572
+ def compute_dlogbeta(self):
573
+ """Correct path and station effects simultaneously (dlogbeta).
574
+
575
+ Parallelized version using multiprocessing with shared memory.
576
+ Memory-optimized: uses contiguous numpy arrays instead of
577
+ dict-of-arrays for target data.
578
+ """
579
+ t0 = time.time()
580
+ self.fprint(f"Computing dlogbeta (parallel, {self.n_workers} workers)")
581
+ self.fprint(f"Using {self.calib_search_method} calibration search method")
582
+ self.fprint("----------------------------------------")
583
+
584
+ eids = self.df_events['_eid'].values
585
+ nevents = len(eids)
586
+ nrecords = len(self.df_records)
587
+
588
+ # ---------------------------------------------------------
589
+ # Build contiguous arrays for target events
590
+ # ---------------------------------------------------------
591
+ self.df_records = self.df_records.sort_values('_eid').reset_index(drop=True)
592
+
593
+ # These store event metadata, one entry per event, aligned with eids
594
+ t_edep = np.empty(nevents, dtype=np.float64)
595
+ t_elat = np.empty(nevents, dtype=np.float64)
596
+ t_elon = np.empty(nevents, dtype=np.float64)
597
+ t_rec_start = np.empty(nevents, dtype=np.int64)
598
+ t_rec_count = np.empty(nevents, dtype=np.int64)
599
+ t_event_name = []
600
+
601
+ has_time = self.calib_time_filter
602
+ if has_time:
603
+ t_etime = np.empty(nevents, dtype=np.int64)
604
+
605
+ # Store _eid-sorted logbeta and channel identifiers
606
+ r_cid = self.df_records['_cid'].values.astype(np.int64)
607
+ r_logbeta = self.df_records['logbeta'].values.astype(np.float64)
608
+
609
+ # df_records is sorted by _eid, so going through records in order
610
+ # we will see all records for one event together, then all records
611
+ # for the next event, etc.
612
+
613
+ # Build event offset arrays vectorized instead of 720K get_group() calls.
614
+ # Since df_records is sorted by _eid, we can find boundaries directly.
615
+ eid_col = self.df_records['_eid'].values
616
+ # Find where each eid starts in the sorted records
617
+ change_points = np.r_[0, np.where(np.diff(eid_col) != 0)[0] + 1, len(eid_col)] # tested, works
618
+ eid_at_boundary = eid_col[change_points[:-1]] # tested, works
619
+
620
+ # Build lookup: eid -> (start, count)
621
+ eid_to_offset = {}
622
+ for k in range(len(change_points) - 1):
623
+ eid_to_offset[eid_at_boundary[k]] = (change_points[k], change_points[k+1] - change_points[k])
624
+ # tested, works
625
+
626
+ # Extract per-event scalars using first record of each group
627
+ edep_all = self.df_records['edep'].values
628
+ elat_all = self.df_records['elat'].values
629
+ elon_all = self.df_records['elon'].values
630
+ ename_all = self.df_records['event_name'].values
631
+ if has_time:
632
+ etime_all = self.df_records['etime'].values
633
+
634
+ for i, eid in enumerate(eids):
635
+ start, count = eid_to_offset[eid]
636
+ t_edep[i] = edep_all[start]
637
+ t_elat[i] = elat_all[start]
638
+ t_elon[i] = elon_all[start]
639
+ t_event_name.append(ename_all[start])
640
+ t_rec_start[i] = start
641
+ t_rec_count[i] = count
642
+ if has_time:
643
+ t_etime[i] = etime_all[start]
644
+
645
+ # ---------------------------------------------------------
646
+ # Calibration arrays
647
+ # ---------------------------------------------------------
648
+
649
+ # Pull the calibration record metadata into contiguous numpy arrays
650
+ c_edep = self.metadata_calib['edep'].values.astype(np.float64)
651
+ c_elat = self.metadata_calib['elat'].values.astype(np.float64)
652
+ c_elon = self.metadata_calib['elon'].values.astype(np.float64)
653
+ c_cid = self.metadata_calib['_cid'].values.astype(np.int64)
654
+ c_logbeta_arr = self.metadata_calib['logbeta'].values.astype(np.float64)
655
+ c_slat = self.metadata_calib['slat'].values.astype(np.float64)
656
+ c_slon = self.metadata_calib['slon'].values.astype(np.float64)
657
+
658
+ if has_time:
659
+ c_etime = self.metadata_calib['etime'].values.astype(np.int64)
660
+
661
+ self.fprint(f" Target events: {nevents:,}, Records: {nrecords:,}, "
662
+ f"Calib records: {len(c_edep):,}")
663
+
664
+ # ---------------------------------------------------------
665
+ # NN-specific: map each calibration record to its unique
666
+ # calibration event index
667
+ # ---------------------------------------------------------
668
+ use_nn = (self.calib_search_method == 'nearest-neighbor')
669
+
670
+ if use_nn:
671
+ calib_eids = self.metadata_calib['_eid'].values.astype(np.int64)
672
+ unique_calib_eids, c_event_idx = np.unique(
673
+ calib_eids, return_inverse=True)
674
+ c_event_idx = c_event_idx.astype(np.int64)
675
+ n_unique_calib = len(unique_calib_eids)
676
+
677
+ # First occurrence of each unique event -> its location
678
+ first_occurrence = np.zeros(n_unique_calib, dtype=np.int64)
679
+ seen = set()
680
+ for i, ev_idx_val in enumerate(c_event_idx):
681
+ if ev_idx_val not in seen:
682
+ first_occurrence[ev_idx_val] = i
683
+ seen.add(ev_idx_val)
684
+
685
+ cu_elat = c_elat[first_occurrence]
686
+ cu_elon = c_elon[first_occurrence]
687
+ cu_edep = c_edep[first_occurrence]
688
+ if has_time:
689
+ cu_etime = c_etime[first_occurrence]
690
+
691
+ self.fprint(f" NN mode: {n_unique_calib:,} unique calib events, "
692
+ f"calib_nn={self.calib_nn}, "
693
+ f"depth_scale={self.calib_depth_scale}")
694
+
695
+ # ---------------------------------------------------------
696
+ # Create shared memory blocks
697
+ # ---------------------------------------------------------
698
+ def _create_shm(name, arr):
699
+ shm = shared_memory.SharedMemory(create=True, size=arr.nbytes, name=name)
700
+ shared_arr = np.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf)
701
+ shared_arr[:] = arr[:]
702
+ return shm
703
+
704
+ shm_blocks = []
705
+ shm_names = {}
706
+ shm_shapes = {}
707
+ shm_dtypes = {}
708
+
709
+ arrays_to_share = {
710
+ 't_edep': t_edep, 't_elat': t_elat, 't_elon': t_elon,
711
+ 't_rec_start': t_rec_start, 't_rec_count': t_rec_count,
712
+ 'r_cid': r_cid, 'r_logbeta': r_logbeta,
713
+ 'c_edep': c_edep, 'c_elat': c_elat, 'c_elon': c_elon,
714
+ 'c_cid': c_cid, 'c_logbeta': c_logbeta_arr,
715
+ 'c_slat': c_slat, 'c_slon': c_slon,
716
+ }
717
+ if has_time:
718
+ arrays_to_share['t_etime'] = t_etime
719
+ arrays_to_share['c_etime'] = c_etime
720
+
721
+ if use_nn:
722
+ arrays_to_share['c_event_idx'] = c_event_idx
723
+ arrays_to_share['cu_elat'] = cu_elat
724
+ arrays_to_share['cu_elon'] = cu_elon
725
+ arrays_to_share['cu_edep'] = cu_edep
726
+ if has_time:
727
+ arrays_to_share['cu_etime'] = cu_etime
728
+
729
+ shm_prefix = f"beta_{os.getpid()}_"
730
+
731
+ for key, arr in arrays_to_share.items():
732
+ name = shm_prefix + key
733
+ shm = _create_shm(name, arr)
734
+ shm_blocks.append(shm)
735
+ shm_names[key] = name
736
+ shm_shapes[key] = arr.shape
737
+ shm_dtypes[key] = arr.dtype
738
+
739
+ # ---------------------------------------------------------
740
+ # Split events into many small chunks for fine-grained progress.
741
+ # Small chunks (~100-500 events) give good load balancing and
742
+ # per-event resolution in the progress bar, while keeping IPC
743
+ # overhead manageable (~1ms per chunk × a few thousand chunks = seconds).
744
+ # ---------------------------------------------------------
745
+ chunk_size = max(1, min(200, nevents // (self.n_workers * 8)))
746
+ chunks = []
747
+ for start in range(0, nevents, chunk_size):
748
+ end = min(start + chunk_size, nevents)
749
+ chunks.append(list(range(start, end)))
750
+
751
+ self.fprint(f" Dispatching {len(chunks)} chunks ({chunk_size} events/chunk) "
752
+ f"to {self.n_workers} workers...")
753
+
754
+ worker_fn = partial(
755
+ _dlogbeta_worker,
756
+ shm_names=shm_names,
757
+ shm_shapes=shm_shapes,
758
+ shm_dtypes=shm_dtypes,
759
+ calib_n_records_min=self.calib_n_records_min,
760
+ calib_hdist_max=self.calib_hdist_max,
761
+ calib_zdist_max=self.calib_zdist_max,
762
+ calib_time_filter=self.calib_time_filter,
763
+ calib_time_method=self.calib_time_method,
764
+ calib_time_ndays=self.calib_time_ndays,
765
+ compute_uncertainty=self.compute_uncertainty,
766
+ confidence_level=self.confidence_level,
767
+ save_full_uncertainty=self.save_full_uncertainty,
768
+ calib_search_method=self.calib_search_method,
769
+ calib_nn=self.calib_nn,
770
+ calib_depth_scale=self.calib_depth_scale,
771
+ )
772
+
773
+ # ---------------------------------------------------------
774
+ # Run parallel computation
775
+ # Progress bar tracks events completed (not chunks), so you
776
+ # see fine-grained progress even with multiprocessing.
777
+ # ---------------------------------------------------------
778
+ all_results = []
779
+ events_done = 0
780
+
781
+ pbar = tqdm(total=nevents, desc="Computing dlogbeta", unit="events")
782
+
783
+ if self.n_workers == 1:
784
+ for chunk in chunks:
785
+ chunk_results = worker_fn(chunk)
786
+ all_results.extend(chunk_results)
787
+ events_done += len(chunk)
788
+ pbar.update(len(chunk))
789
+ else:
790
+ with mp.Pool(processes=self.n_workers) as pool:
791
+ for chunk_results in pool.imap_unordered(worker_fn, chunks):
792
+ all_results.extend(chunk_results)
793
+ # Each chunk has chunk_size events (last may be smaller)
794
+ # We track by chunk completion but update by event count
795
+ events_done += chunk_size # approximate; close enough for progress
796
+ pbar.n = min(events_done, nevents)
797
+ pbar.refresh()
798
+
799
+ pbar.n = nevents
800
+ pbar.refresh()
801
+ pbar.close()
802
+
803
+ # ---------------------------------------------------------
804
+ # Collect results back into arrays
805
+ # ---------------------------------------------------------
806
+ dlogbeta_results = np.full(nrecords, np.nan)
807
+
808
+ if self.compute_uncertainty:
809
+ dlogbeta_std_results = np.full(nrecords, np.nan)
810
+ dlogbeta_lower_results = np.full(nrecords, np.nan)
811
+ dlogbeta_upper_results = np.full(nrecords, np.nan)
812
+ dlogbeta_median_results = np.full(nrecords, np.nan)
813
+
814
+ uncertainty_records = {}
815
+
816
+ for res in all_results:
817
+ ev_idx = res['ev_idx']
818
+ start = t_rec_start[ev_idx]
819
+ count = t_rec_count[ev_idx]
820
+ record_indices = np.arange(start, start + count)
821
+
822
+ dlogbeta_results[record_indices] = res['dlogbeta']
823
+
824
+ if self.compute_uncertainty:
825
+ dlogbeta_std_results[record_indices] = res.get('dlogbeta_std', np.nan)
826
+ dlogbeta_lower_results[record_indices] = res.get('dlogbeta_lower', np.nan)
827
+ dlogbeta_upper_results[record_indices] = res.get('dlogbeta_upper', np.nan)
828
+ dlogbeta_median_results[record_indices] = res.get('dlogbeta_median', np.nan)
829
+
830
+ if self.save_full_uncertainty:
831
+ uncertainty_records[ev_idx] = {
832
+ 'event_name': t_event_name[ev_idx],
833
+ 'dlogbeta_j': res.get('dlogbeta_j', []),
834
+ 'dlogbeta_j_residual': res.get('dlogbeta_j_residual', []),
835
+ 'bearing': res.get('bearing', []),
836
+ 'deldist': res.get('deldist', []),
837
+ 'ncalib_j': res.get('ncalib_j', []),
838
+ 'slat': res.get('slat', []),
839
+ 'slon': res.get('slon', []),
840
+ }
841
+
842
+ # ---------------------------------------------------------
843
+ # Clean up shared memory
844
+ # ---------------------------------------------------------
845
+ for shm in shm_blocks:
846
+ shm.close()
847
+ shm.unlink()
848
+
849
+ # ---------------------------------------------------------
850
+ # Assign results back to DataFrame
851
+ # ---------------------------------------------------------
852
+ self.df_records['dlogbeta'] = dlogbeta_results
853
+
854
+ if self.compute_uncertainty:
855
+
856
+ if self.save_full_uncertainty:
857
+ sorted_keys = sorted(uncertainty_records.keys())
858
+ self.df_dlogbeta = pd.DataFrame([uncertainty_records[k] for k in sorted_keys])
859
+
860
+ self.df_records['dlogbeta_std'] = dlogbeta_std_results
861
+ self.df_records['dlogbeta_lower'] = dlogbeta_lower_results
862
+ self.df_records['dlogbeta_upper'] = dlogbeta_upper_results
863
+ self.df_records['dlogbeta_median'] = dlogbeta_median_results
864
+
865
+ # Drop rows with NaN dlogbeta
866
+ n_before = len(self.df_records)
867
+ self.df_records = self.df_records.dropna(subset=['dlogbeta']).reset_index(drop=True)
868
+ n_after = len(self.df_records)
869
+
870
+ elapsed = time.time() - t0
871
+ self.fprint(f" dlogbeta complete: {len(all_results):,}/{nevents:,} events, "
872
+ f"{n_after:,}/{n_before:,} records kept. ({elapsed:.1f}s)")
873
+
874
+
875
+ def apply_magnitude_correction(self, corr_type='smoothedspline'):
876
+ def _smoothed_spline_correction(emag, dlogbeta, mag_edges):
877
+ from scipy.signal import savgol_filter
878
+ from scipy.stats import binned_statistic
879
+
880
+ midpoints = (mag_edges[:-1] + mag_edges[1:])/2
881
+ bin_inds = np.digitize(emag, mag_edges) - 1
882
+
883
+ median_dlbeta = binned_statistic(emag, dlogbeta, statistic='median', bins=mag_edges)[0]
884
+
885
+ naninds = np.isnan(median_dlbeta)
886
+ realmids = midpoints[~naninds]
887
+ median_dlbeta_valid = median_dlbeta[~naninds]
888
+
889
+ filter_window_len = int(np.floor(len(median_dlbeta_valid)/4))
890
+ if filter_window_len % 2 == 0:
891
+ filter_window_len -= 1
892
+ filter_window_len = max(filter_window_len, 3)
893
+
894
+ polyorder = min(3, filter_window_len - 1)
895
+
896
+ median_dlbeta_smooth = savgol_filter(median_dlbeta_valid, filter_window_len, polyorder)
897
+
898
+ corrections = np.interp(emag, realmids, median_dlbeta_smooth)
899
+ correction_mags = realmids
900
+ correction_function = median_dlbeta_smooth
901
+
902
+ nfi = dlogbeta - corrections
903
+
904
+ return nfi, correction_mags, correction_function
905
+
906
+ self.group_events()
907
+ mag_corr_dM = self.mag_corr_dM
908
+
909
+ M_range = np.round((
910
+ np.floor(np.min(self.df_events['emag'].values) / mag_corr_dM) * mag_corr_dM - mag_corr_dM/2,
911
+ np.ceil(np.max(self.df_events['emag'].values) / mag_corr_dM) * mag_corr_dM + mag_corr_dM/2
912
+ ), 4)
913
+
914
+ mag_edges = np.arange(M_range[0], M_range[1]+1*mag_corr_dM, mag_corr_dM)
915
+
916
+ # midpoint_0 = (mag_edges[0] + mag_edges[1])/2
917
+
918
+ # self.df_events = self.df_events[self.df_events['emag']>=midpoint_0].reset_index(drop=True)
919
+
920
+
921
+ if corr_type == 'smoothedspline':
922
+ nfi, correction_mags, correction_function = _smoothed_spline_correction(
923
+ self.df_events['emag'].values,
924
+ self.df_events['dlogbeta'].values,
925
+ mag_edges
926
+ )
927
+ self.df_events['nfi'] = nfi
928
+ self.correction_mags = correction_mags
929
+ self.correction_function = correction_function
930
+ elif corr_type == 'source-separate':
931
+ source = self.df_events['event_name'].str[0]
932
+ self.correction_mags = {}
933
+ self.correction_function = {}
934
+ self.df_events['nfi'] = np.nan
935
+ for s in np.unique(source):
936
+ df_source = self.df_events[source==s]
937
+ nfi, correction_mags, correction_function = _smoothed_spline_correction(
938
+ df_source['emag'].values,
939
+ df_source['dlogbeta'].values,
940
+ mag_edges
941
+ )
942
+ self.df_events.loc[source==s, 'nfi'] = nfi
943
+ self.correction_mags[s] = correction_mags
944
+ self.correction_function[s] = correction_function
945
+
946
+
947
+ self.explode_events()
948
+ self.group_channels()
949
+
950
+ def group_events(self):
951
+ self.fprint("group_events()")
952
+ t0 = time.time()
953
+ self.compute_column_dependencies()
954
+ # observed=True: only group by values actually present in the data.
955
+ # This is critical when event_name/channel_name are categorical —
956
+ # observed=False (the pandas default) would try to create the full
957
+ # Cartesian product of all category levels, which with 720K events
958
+ # × 4K channels overflows memory.
959
+ self.df_events = self.df_records.groupby(
960
+ self.ev_dep, as_index=False, observed=True
961
+ )[self.ch_dep+self.pair_dep].agg(list)
962
+ self.tprint(f" |---> {time.time()-t0:.4f} seconds")
963
+
964
+ def group_channels(self):
965
+ self.fprint("group_channels()")
966
+ self.compute_column_dependencies()
967
+ self.df_channels = self.df_records.groupby(
968
+ self.ch_dep, as_index=False, dropna=False, observed=True
969
+ )[self.ev_dep+self.pair_dep].agg(list)
970
+
971
+ def explode_events(self):
972
+ self.fprint('explode_events()')
973
+ t0 = time.time()
974
+ self.compute_column_dependencies()
975
+ self.df_records = self.df_events.explode(self.ch_dep+self.pair_dep)
976
+ self.tprint(f" |---> {time.time()-t0:.4f} seconds")
977
+
978
+ def explode_channels(self):
979
+ t0 = time.time()
980
+ self.fprint('explode_channels()')
981
+ self.compute_column_dependencies()
982
+ self.df_records = self.df_channels.explode(self.ev_dep+self.pair_dep)
983
+ self.tprint(f" |---> {time.time()-t0:.4f} seconds")
984
+
985
+ def compute_column_dependencies(self):
986
+ ev_dep_columns = [
987
+ 'event_name', 'emag', 'elat', 'elon', 'edep', 'etime', 'edatetime', '_eid', 'event_id',
988
+ 'evid', 'event', 'ex', 'ey', 'nts', 'dlogbeta', 'nfi',
989
+ 'nrec', 'dlogbeta_std', 'dlogbeta_lower', 'dlogbeta_upper',
990
+ 'dlogbeta_median', 'etype', 'emagtype']
991
+ cha_dep_columns = [
992
+ 'channel_name', 'slat', 'slon', 'selev', '_cid', 'sx', 'sy', 'kappa0'
993
+ ]
994
+ pair_dep_columns = [
995
+ 'logbeta', 'deldist', 'a2_max', 's1', 's2', 'stn'
996
+ ]
997
+ columns = self.df_records.columns
998
+ self.ev_dep = []
999
+ self.ch_dep = []
1000
+ self.pair_dep = []
1001
+
1002
+ for col in columns:
1003
+ if col in ev_dep_columns:
1004
+ self.ev_dep.append(col)
1005
+ elif col in cha_dep_columns:
1006
+ self.ch_dep.append(col)
1007
+ elif col in pair_dep_columns:
1008
+ self.pair_dep.append(col)
1009
+ else:
1010
+ raise Warning("Column {} not recognized".format(col))
1011
+
1012
+
1013
+ def validate_input(self):
1014
+ def _to_datetime64(x):
1015
+ if isinstance(x, np.datetime64):
1016
+ return x
1017
+ if isinstance(x, UTC):
1018
+ return np.datetime64(int(x.ns), 'ns')
1019
+ if isinstance(x, str):
1020
+ return np.datetime64(pd.Timestamp(x).to_datetime64())
1021
+ if isinstance(x, pd.Timestamp):
1022
+ return x.to_datetime64()
1023
+ raise TypeError(f"Unhandled edatetime type: {type(x)}")
1024
+
1025
+ t0 = time.time()
1026
+ self.fprint("validate_input()")
1027
+ columns = self.df_records.columns
1028
+ # Rename columns to standardized names. Conventions:
1029
+ # 1) event-dependent quantities begin with e-
1030
+ # 2) station/channel-dependent quantities begin with s-
1031
+ # 3) event_name and channel_name are unique identifiers for
1032
+ # events and channels, respectively
1033
+
1034
+ ### Check column names and verify all required columns are present ###
1035
+ # Check for an event_name column
1036
+ event_name_columns = ['event_id', 'evid', 'event']
1037
+ if 'event_name' not in columns:
1038
+ for col in event_name_columns:
1039
+ if col in columns:
1040
+ self.df_records.rename(
1041
+ columns={col: 'event_name'}, inplace=True)
1042
+ self.fprint("Renamed column {} to event_name".format(col))
1043
+ break
1044
+ if "event_name" not in self.df_records.columns:
1045
+ raise ValueError("No event_name column found. Check input data.")
1046
+
1047
+ channel_name_columns = ['station_id', 'sid', 'station', 'stname',
1048
+ 'st_name', 'stid']
1049
+ if 'channel_name' not in columns:
1050
+ for col in channel_name_columns:
1051
+ if col in columns:
1052
+ self.df_records.rename(
1053
+ columns={col: 'channel_name'}, inplace=True)
1054
+ self.fprint("Renamed column {} to channel_name".format(col))
1055
+ break
1056
+ if "channel_name" not in self.df_records.columns:
1057
+ raise ValueError("No channel_name column found. Check input data.")
1058
+
1059
+ if 'qmag' in columns:
1060
+ self.df_records.rename(columns={'qmag': 'emag'}, inplace=True)
1061
+ if 'qlat' in columns:
1062
+ self.df_records.rename(columns={'qlat': 'elat'}, inplace=True)
1063
+ if 'qlon' in columns:
1064
+ self.df_records.rename(columns={'qlon': 'elon'}, inplace=True)
1065
+ if 'qdep' in columns:
1066
+ self.df_records.rename(columns={'qdep': 'edep'}, inplace=True)
1067
+
1068
+ required_columns = [
1069
+ 'event_name', 'channel_name', 'emag', 'elat', 'elon', 'edep',
1070
+ 'slat', 'slon', 'selev', 'deldist']
1071
+ for col in required_columns:
1072
+ if col not in self.df_records.columns:
1073
+ raise ValueError(f"Missing required column: {col}")
1074
+
1075
+ # Make sure 'edatetime' is proper format
1076
+ if 'edatetime' in columns:
1077
+ self.df_records['edatetime'] = self.df_records['edatetime'].apply(_to_datetime64)
1078
+
1079
+ # Convert edatetime to posix ns (handles strings, UTCDateTime, etc.)
1080
+ if 'edatetime' in columns:
1081
+ self.convert_utc_to_posix_ns()
1082
+
1083
+ if len(self.df_records) != self.spectra.shape[0]:
1084
+ raise ValueError("Metadata and spectra row counts do not match.")
1085
+ if len(self.f) != self.spectra.shape[1]:
1086
+ raise ValueError(
1087
+ f"Spectra width {self.spectra.shape[1]} does not match "
1088
+ f"f array length {len(self.f)}.")
1089
+
1090
+ if (np.median(self.df_records['deldist']) > 100) or \
1091
+ (np.median(self.df_records['deldist']) < 10):
1092
+ raise Warning(
1093
+ f"Is deldist in km? Median is"
1094
+ f" {np.median(self.df_records['deldist'])}")
1095
+
1096
+ assert np.sum(np.isnan(self.spectra)) == 0, \
1097
+ "Spectra contains NaNs."
1098
+
1099
+ self.nrecords_initial = len(self.df_records)
1100
+ self.nevents_initial = len(self.df_records['event_name'].unique())
1101
+ self.nchannels_initial = len(self.df_records['channel_name'].unique())
1102
+
1103
+ if self.save_dir is None: self.save_dir = 'tmp/'
1104
+ os.makedirs(self.save_dir, exist_ok=True)
1105
+
1106
+ self.tprint(f" |---> {time.time()-t0:.4f} seconds")
1107
+
1108
+
1109
+ def digitize_names(self):
1110
+ self.df_records = self.df_records.sort_values(by=['event_name', 'channel_name'])
1111
+ v = self.df_records.index.values
1112
+ self.df_records = self.df_records.reset_index(drop=True)
1113
+
1114
+ self.spectra = self.spectra[v, :]
1115
+
1116
+ self.df_records, self.event_name_dict = _get_id_from_column(self.df_records, 'event_name', '_eid')
1117
+ self.df_records, self.channel_name_dict = _get_id_from_column(self.df_records, 'channel_name', '_cid')
1118
+
1119
+ def compute_logbeta(self):
1120
+ t0 = time.time()
1121
+ self.fprint("compute_logbeta()")
1122
+ low_inds = self.low_window_inds
1123
+ high_inds = self.high_window_inds
1124
+
1125
+ low_band = np.median(np.log10(self.spectra[:, low_inds[0]:low_inds[1]+1]), axis=1)
1126
+ high_band = np.median(np.log10(self.spectra[:, high_inds[0]:high_inds[1]+1]), axis=1)
1127
+ self.df_records['logbeta'] = high_band - low_band
1128
+ del self.spectra
1129
+ self.tprint(f" |---> {time.time()-t0:.4f} seconds")
1130
+
1131
+ def compute_frequency_bands(self):
1132
+ self.low_window_inds = _get_inds_of_values_in_array(self.f, self.low_window_desired)
1133
+ self.high_window_inds = _get_inds_of_values_in_array(self.f, self.high_window_desired)
1134
+
1135
+ self.low_window = [self.f[self.low_window_inds[0]], self.f[self.low_window_inds[1]]]
1136
+ self.high_window = [self.f[self.high_window_inds[0]], self.f[self.high_window_inds[1]]]
1137
+
1138
+ def fprint(self, S):
1139
+ if not self.quiet:
1140
+ print(S)
1141
+
1142
+ def tprint(self, S):
1143
+ if self.print_time:
1144
+ print(S)
1145
+
1146
+ def print_frequency_information(self):
1147
+ f = self.f
1148
+ print("")
1149
+ print("FREQUENCY ARRAY INFORMATION")
1150
+ print("----------------------------")
1151
+ print(f"Frequency array ranges from {f[0]:.2f} to {f[-1]:.2f} Hz with {len(f)} elements (df = {self.df:.3f} Hz). ")
1152
+ print(f"Desired | Actual low-frequency band: {self.low_window_desired[0]:7.3f} -{self.low_window_desired[1]:7.3f} Hz | {self.low_window[0]:7.3f} -{self.low_window[1]:7.3f} Hz")
1153
+ print(f"Desired | Actual high-frequency band: {self.high_window_desired[0]:7.3f} -{self.high_window_desired[1]:7.3f} Hz | {self.high_window[0]:7.3f} -{self.high_window[1]:7.3f} Hz")
1154
+
1155
+ def print_calibration_information(self):
1156
+ ncalib = len(self.metadata_calib['event_name'].unique())
1157
+ ncalibrec = len(self.metadata_calib)
1158
+ nev = len(self.df_records['event_name'].unique())
1159
+ nrec = len(self.df_records)
1160
+ print("")
1161
+ print("CALIBRATION INFORMATION")
1162
+ print("----------------------------")
1163
+ print(f"Calibration range: M {self.calib_mag_range[0]:.2f} to {self.calib_mag_range[1]:.2f}")
1164
+ print(f"Calibration events: {ncalib:,} ({ncalib/nev*100:.2f}%)")
1165
+ print(f"Calibration records: {ncalibrec:,} ({ncalibrec/nrec*100:.2f}%)")
1166
+ print("")
1167
+
1168
+ def print_metadata_information(self):
1169
+ nev = len(self.df_records['event_name'].unique())
1170
+ nst = len(self.df_records['channel_name'].unique())
1171
+ nrec = len(self.df_records)
1172
+ print("")
1173
+ print("METADATA INFORMATION")
1174
+ print("----------------------------")
1175
+ if nev==self.nevents_initial:
1176
+ print(f"Events: {nev:,}")
1177
+ else:
1178
+ print(f"Events: {nev:,} of {self.nevents_initial:,} inital ({nev/self.nevents_initial*100:.2f}%)")
1179
+
1180
+ if nst==self.nchannels_initial:
1181
+ print(f"Channels: {nst:,}")
1182
+ else:
1183
+ print(f"Channels: {nst:,} of {self.nchannels_initial:,} inital ({nst/self.nchannels_initial*100:.2f}%)")
1184
+
1185
+ if nrec==self.nrecords_initial:
1186
+ print(f"Records: {nrec:,}")
1187
+ else:
1188
+ print(f"Records: {nrec:,} of {self.nrecords_initial:,} inital ({nrec/self.nrecords_initial*100:.2f}%)")
1189
+
1190
+
1191
+ def _get_inds_of_values_in_array(x, values):
1192
+ return np.array([np.argmin(np.abs(x - val)) for val in values], dtype=int)
1193
+
1194
+ def _get_id_from_column(df, column, id_name):
1195
+ ids = df[column].unique()
1196
+ id_dict = {ids[i]: i for i in range(len(ids))}
1197
+ df[id_name] = df[column].map(id_dict)
1198
+ return df, id_dict
1199
+
1200
+ def _haversine_km(lon1, lat1, lon2, lat2):
1201
+ """
1202
+ Calculate the great circle distance in kilometers between two points
1203
+ on the earth (specified in decimal degrees)
1204
+ """
1205
+ lon1r, lat1r, lon2r, lat2r = map(np.radians, [lon1, lat1, lon2, lat2])
1206
+ dlonr = lon2r - lon1r
1207
+ dlatr = lat2r - lat1r
1208
+ a = np.sin(dlatr/2)**2 + np.cos(lat1r) * np.cos(lat2r) * np.sin(dlonr/2)**2
1209
+ c = 2 * np.arcsin(np.sqrt(a))
1210
+ r = 6371.0
1211
+ return c * r
1212
+
1213
+ def _get_bearing(lat1, lon1, lat2, lon2):
1214
+ """
1215
+ Calculate the bearing between two points, from point 1 to point 2.
1216
+ """
1217
+ y = np.sin(np.radians(lon2 - lon1)) * np.cos(np.radians(lat2))
1218
+ x = np.cos(np.radians(lat1)) * np.sin(np.radians(lat2)) - np.sin(np.radians(lat1)) * np.cos(np.radians(lat2)) * np.cos(np.radians(lon2 - lon1))
1219
+ bearing = np.degrees(np.arctan2(y, x))
1220
+ bearing = (bearing + 360) % 360
1221
+ return bearing
@@ -0,0 +1,3 @@
1
+ Metadata-Version: 2.4
2
+ Name: py-nfi
3
+ Version: 0.1.0
@@ -0,0 +1,8 @@
1
+ README.md
2
+ pyproject.toml
3
+ src/nfi/__init__.py
4
+ src/nfi/nfi.py
5
+ src/py_nfi.egg-info/PKG-INFO
6
+ src/py_nfi.egg-info/SOURCES.txt
7
+ src/py_nfi.egg-info/dependency_links.txt
8
+ src/py_nfi.egg-info/top_level.txt
@@ -0,0 +1 @@
1
+ nfi