import numpy as np
import matplotlib.pyplot as plt
from typing import Optional,List,Union
from ADFWI.survey import Survey
from ADFWI.utils import gpu2cpu,tensor2numpy
from ADFWI.view import plot_waveform2D,plot_waveform_wiggle,plot_waveform_trace
[docs]class SeismicData():
"""
A class to handle and process seismic data collected during a survey.
Parameters
----------
survey : Survey
The Survey object containing information about sources and receivers.
src_num : int
Number of sources in the survey.
rcv_num : int
Number of receivers in the survey.
src_loc : np.ndarray
Locations of the sources.
rcv_loc : np.ndarray
Locations of the receivers.
src_type : str
Type of sources (e.g., 'point', 'line').
rcv_type : str
Type of receivers (e.g., 'pressure').
nt : int
Number of time samples.
dt : float
Time step size.
t : np.ndarray
Array of time steps.
data : dict
Seismic data (pressure, velocity, etc.).
data_masks : dict
Masks for the seismic data.
"""
def __init__(self,survey:Survey):
"""
Initialize SeismicData object with survey information.
"""
self.survey = survey
# get the survey information
self.src_num = survey.source.num
self.rcv_num = survey.receiver.num
self.src_loc = survey.source.get_loc()
self.rcv_loc = survey.receiver.get_loc()
self.src_type = survey.source.get_type()
self.rcv_type = survey.receiver.get_type()
self.nt = survey.receiver.nt
self.dt = survey.receiver.dt
self.t = np.arange(self.nt)*self.dt
# data
self.data = None
self.data_masks = None
def __repr__(self):
""" Print the survey information
"""
info = f"Seismic Data:\n"
info += f" Source number : {self.src_num}\n"
info += f" Receiver number : {self.rcv_num}\n"
info += f" Time samples : {self.nt} samples at {self.dt * 1000:.2f} ms\n"
return info
[docs] def record_data(self, data: dict):
""" Add the shot gather data to the class
"""
for key,value in data.items():
value = tensor2numpy(gpu2cpu(value)).copy()
data[key] = value
self.data = data
[docs] def save(self,path:str):
"""
Save the seismic data to a .npz file.
"""
data_save = { 'data' : self.data,
'src_loc' : self.src_loc,
'rcv_loc' : self.rcv_loc,
'src_num' : self.src_num,
'rcv_num' : self.rcv_num,
'rcv_type' : self.rcv_type,
'src_type' : self.src_type,
't' : self.t,
'nt' : self.nt,
'dt' : self.dt
}
np.savez(path, **data_save)
[docs] def load(self, path: str):
"""
Load the shot gather data
"""
data = np.load(path, allow_pickle=True)
# load the data
self.data = data['data'].item()
self.src_loc = data['src_loc']
self.rcv_loc = data['rcv_loc']
self.src_num = data['src_num']
self.rcv_num = data['rcv_num']
self.rcv_type = data['rcv_type']
self.src_type = data['src_type']
self.t = data['t']
self.nt = data['nt']
self.dt = data['dt']
return
[docs] def normalize_and_mask(self,array):
"""
Normalize the seismic data array by dividing by its maximum value along the time axis, and create a mask for invalid data (time sum == 0).
"""
time_sum = np.sum(np.abs(array), axis=1, keepdims=True)
mask = time_sum == 0
max_val = np.max(np.abs(array), axis=1, keepdims=True)
max_val = np.where(mask, 1, max_val)
array = array / max_val
return array
[docs] def parse_elastic_data(self,normalize=False):
"""
Parse elastic data (stress and velocity components) from the seismic data.
"""
txx = self.data["txx"]
tzz = self.data["tzz"]
txz = self.data["txz"]
vx = self.data["vx"]
vz = self.data["vz"]
pressure = -(txx + tzz)
if normalize:
pressure = self.normalize_and_mask(pressure)
txz = self.normalize_and_mask(txz)
vx = self.normalize_and_mask(vx)
vz = self.normalize_and_mask(vz)
return pressure,txz,vx,vz
[docs] def parse_acoustic_data(self,normalize=False):
"""
Parse acoustic data (pressure and velocity components) from the seismic data.
"""
pressure = self.data["p"]
u = self.data["u"]
w = self.data["w"]
if normalize:
pressure = self.normalize_and_mask(pressure)
u = self.normalize_and_mask(u)
w = self.normalize_and_mask(w)
return pressure,u,w