import numpy as np
from scipy.ndimage import gaussian_filter
import h5py
from .para import ATTPara
from .io.crustmodel import CrustModel
from .io.asciimodel import ASCIIModel
from .attarray import Dataset
from .utils.common import init_axis
import copy
[docs]
class ATTModel():
"""Create initial model from external models
"""
def __init__(self, para_fname='input_params.yml') -> None:
"""
:param para_fname: Path to parameter file, defaults to 'input_params.yml'
:type para_fname: str, optional
"""
self.para_fname = para_fname
self.d_rtp = np.zeros(3)
self.read_param()
self.eta = np.zeros(self.n_rtp)
self.xi = np.zeros(self.n_rtp)
self.zeta = np.zeros(self.n_rtp)
self.vel = np.zeros(self.n_rtp)
[docs]
def read_param(self):
"""Read ``n_rtp``, ``min_max_dep``, ``min_max_lat`` and ``min_max_lon`` from ``para_fname``
"""
para = ATTPara(self.para_fname)
self.n_rtp = para.input_params['domain']['n_rtp']
self.min_max_dep = para.input_params['domain']['min_max_dep']
self.min_max_lat = para.input_params['domain']['min_max_lat']
self.min_max_lon = para.input_params['domain']['min_max_lon']
self.depths, self.latitudes, self.longitudes, \
self.d_rtp[0], self.d_rtp[1], self.d_rtp[2] = init_axis(
self.min_max_dep, self.min_max_lat, self.min_max_lon, self.n_rtp
)
self.radius = 6371. - self.depths
[docs]
@classmethod
def read(cls, model_fname: str, para_fname='input_params.yml'):
"""Read an exists model
:param model_fname: Path to the exists model
:type model_fname: str
:param para_fname: Path to parameter file, defaults to 'input_params.yml'
:type para_fname: str, optional
"""
mod = cls(para_fname)
f = h5py.File(model_fname)
mod.vel = f['vel'][:]
mod.xi = f['xi'][:]
mod.eta = f['eta'][:]
# mod.zeta = f['zeta'][:]
mod._check_axis()
if not ((mod.xi==0).all() and (mod.eta==0).all()):
mod.to_ani()
f.close()
return mod
[docs]
def copy(self):
"""Copy the model
"""
return copy.deepcopy(self)
def _check_axis(self):
if self.vel.shape != tuple(self.n_rtp):
raise ValueError('conflicting size of data and n_rtp in parameters')
[docs]
def to_ani(self):
"""Convert to anisotropic strength (epsilon) and azimuth (phi)
"""
self.epsilon = np.sqrt(self.eta**2+self.xi**2)
# self.phi = np.zeros_like(self.epsilon)
self.phi = np.rad2deg(0.5*np.arctan2(self.eta, self.xi))
# idx = np.where(self.xi <= 0)
# self.phi[idx] = 90 + 0.5*atand(self.eta[idx]/self.xi[idx])
# idx = np.where((self.xi > 0) & (self.eta <= 0))
# self.phi[idx] = 180 + 0.5*atand(self.eta[idx]/self.xi[idx])
# idx = np.where((self.xi > 0) & (self.eta > 0))
# self.phi[idx] = 0.5*atand(self.eta[idx]/self.xi[idx])
[docs]
def to_xarray(self):
"""Convert to xarray
:return: attarray dataset
:rtype: pytomoatt.attarray.Dataset
"""
data_dict = {}
data_dict['vel'] = (["r", "t", "p"], self.vel)
data_dict['xi'] = (["r", "t", "p"], self.xi)
data_dict['eta'] = (["r", "t", "p"], self.eta)
# data_dict['zeta'] = (["r", "t", "p"], self.zeta)
if hasattr(self, 'epsilon') and hasattr(self, 'phi'):
data_dict['epsilon'] = (["r", "t", "p"], self.epsilon)
data_dict['phi'] = (["r", "t", "p"], self.phi)
if hasattr(self, 'dlnv'):
data_dict['dlnv'] = (["r", "t", "p"], self.dlnv)
dataset = Dataset(
data_dict,
coords={
'dep': (['r'], self.depths),
'rad': (['r'], self.radius),
'lat': (['t'], self.latitudes),
'lon': (['p'], self.longitudes),
}
)
return dataset
[docs]
def grid_data_crust1(self, type='vp'):
"""Grid data from CRUST1.0 model
:param type: Specify velocity type of ``vp`` or ``vs``, defaults to 'vp'
:type type: str, optional
"""
cm = CrustModel()
self.vel = cm.griddata(
self.min_max_dep,
self.min_max_lat,
self.min_max_lon,
self.n_rtp, type=type
)
[docs]
def grid_data_ascii(self, model_fname:str, **kwargs):
"""Grid data from custom model file in ASCII format
:param model_fname: Path to model file
:type model_fname: str
:param usecols: Columns order by longitude, latitude, depth and velocity, defaults to [0, 1, 2, 3]
:type usecols: list or tuple
"""
am = ASCIIModel(model_fname)
self.vel = am.read_ascii(**kwargs)
self.vel = am.griddata(
self.min_max_dep,
self.min_max_lat,
self.min_max_lon,
self.n_rtp,
)
[docs]
def smooth(self, sigma=5.0):
"""Gaussian smooth the 3D velocity model
:param sigma: Standard division of gaussian kernel in km, defaults to 10
:type sigma: scalar or sequence of scalars , optional
"""
if isinstance(sigma, (int, float)):
sigma_all = np.ones(3)*sigma/self.d_rtp/2/np.pi
elif len(sigma) == 3:
sigma_all = np.array(sigma)/self.d_rtp/2/np.pi
sigma_all[0:2] /= 111.19
self.vel = gaussian_filter(self.vel, sigma)
[docs]
def calc_dv_avg(self):
"""calculate anomalies relative to average velocity at each depth
"""
self.dlnv = np.zeros_like(self.vel)
for i, _ in enumerate(self.depths):
avg = np.mean(self.vel[i, :, :])
self.dlnv[i, :, :] = 100 * (self.vel[i, :, :] - avg)/avg
[docs]
def calc_dv(self, ref_mod_fname: str):
"""calculate anomalies relative to another model
:param ref_mod_fname: Path to reference model
:type ref_mod_fname: str
"""
with h5py.File(ref_mod_fname) as f:
ref_vel = f['vel'][:]
if self.vel.shape != ref_vel.shape:
raise ValueError('reference model should be in same size as input model')
self.dlnv = 100*(self.vel - ref_vel)/ref_vel
[docs]
def write(self, out_fname=None):
"""Write to h5 file with TomoATT format.
:param fname: file name of output model, defaults to 'model_crust1.0.h5'
:type fname: str, optional
"""
if out_fname is None:
out_fname = 'Sub_CRUST1.0_{}_{:d}_{:d}_{:d}.h5'.format(self.type, *self.n_rtp)
with h5py.File(out_fname, 'w') as f:
f.create_dataset('eta', data=self.eta)
f.create_dataset('xi', data=self.xi)
f.create_dataset('zeta', data=self.zeta)
f.create_dataset('vel', data=self.vel)