diff --git a/pytomoatt/checkerboard.py b/pytomoatt/checkerboard.py index 589c577..a65db5c 100644 --- a/pytomoatt/checkerboard.py +++ b/pytomoatt/checkerboard.py @@ -22,7 +22,10 @@ def __init__(self, model_fname:str, para_fname='input_params.yml') -> None: self.vel = f['vel'][:] self.eta = f['eta'][:] self.xi = f['xi'][:] - self.zeta = f['zeta'][:] + if 'zeta' in f: # some model may not have zeta + self.zeta = f['zeta'][:] + else: + self.zeta = np.zeros_like(self.vel) self._init_axis() def _init_axis(self): diff --git a/pytomoatt/io/crustmodel.py b/pytomoatt/io/crustmodel.py index f7fa937..15bdc66 100644 --- a/pytomoatt/io/crustmodel.py +++ b/pytomoatt/io/crustmodel.py @@ -2,6 +2,7 @@ from os.path import dirname, abspath, join from ..utils.common import init_axis from ..setuplog import SetupLog +from ..utils.rotate import rtp_rotation_reverse import pickle import sys from tqdm import tqdm @@ -40,7 +41,7 @@ def __init__(self, fname=join(dirname(dirname(abspath(__file__))), 'data', 'crus self.points_dict = pickle.load(f) self.log = SetupLog() - def griddata(self, min_max_dep, min_max_lat, min_max_lon, n_rtp, type='vp'): + def griddata(self, min_max_dep, min_max_lat, min_max_lon, n_rtp, type='vp', rotate=None): """Linearly interpolate velocity into regular grids :param min_max_dep: min and max depth, ``[min_dep, max_dep]`` @@ -63,25 +64,40 @@ def griddata(self, min_max_dep, min_max_lat, min_max_lon, n_rtp, type='vp'): else: self.log.Modellog.error(f"Velocity type {type} not supported in CRUST1.0 model") sys.exit(1) + self.dd, self.tt, self.pp, _, _, _, = init_axis( min_max_dep, min_max_lat, min_max_lon, n_rtp ) + tt_2d, pp_2d = np.meshgrid(self.tt, self.pp, indexing='ij') + + # rotate reversely, from computational grid to physical grid + if rotate is not None: + try: + central_lat, central_lon, rotation_angle = rotate + except (TypeError, ValueError): + self.log.Modellog.error("rotate must be a 3-item sequence: [central_lat, central_lon, rotation_angle]") + sys.exit(1) + tt_2d, pp_2d = rtp_rotation_reverse(tt_2d, pp_2d, central_lat, central_lon, rotation_angle) + # Grid data self.log.Modellog.info('Grid data, please wait for a few minutes') vel = np.zeros(n_rtp) with tqdm(total=self.n_rtp[1] * self.n_rtp[2], desc='Gridding') as pbar: for ilat in range(self.n_rtp[1]): - new_lat = self.tt[ilat] - idx_lat_left, ratio_lat = degree_to_idx_and_ratio(new_lat) - idx_lat_right = idx_lat_left + 1 - if idx_lat_left == -1: - idx_lat_left = 0 - idx_lat_right = 1 - for ilon in range(self.n_rtp[2]): pbar.update(1) - new_lon = self.pp[ilon] + + # latitude index and ratio + new_lat = tt_2d[ilat, ilon] + idx_lat_left, ratio_lat = degree_to_idx_and_ratio(new_lat) + idx_lat_right = idx_lat_left + 1 + if idx_lat_left == -1: + idx_lat_left = 0 + idx_lat_right = 1 + + # longitude index and ratio + new_lon = pp_2d[ilat, ilon] idx_lon_left, ratio_lon = degree_to_idx_and_ratio(new_lon) idx_lon_right = idx_lon_left + 1 if idx_lon_left == -1: # between -179.5 and +179.5 diff --git a/pytomoatt/model.py b/pytomoatt/model.py index 14c5102..ef1adc8 100644 --- a/pytomoatt/model.py +++ b/pytomoatt/model.py @@ -102,18 +102,20 @@ def to_xarray(self): ) return dataset - def grid_data_crust1(self, type='vp'): + def grid_data_crust1(self, type='vp', rotate=None): """Grid data from CRUST1.0 model :param type: Specify velocity type of ``vp`` or ``vs``, defaults to 'vp' :type type: str, optional + :param rotate: Rotation parameters [central_lat, central_lon, rotation_angle] in degrees, defaults to None """ cm = CrustModel() self.vel = cm.griddata( self.min_max_dep, self.min_max_lat, self.min_max_lon, - self.n_rtp, type=type + self.n_rtp, type=type, + rotate=rotate ) def grid_data_ascii(self, model_fname:str, **kwargs):