pylot/pylot/tomography/fmtomo_tools/quantify_ray_crossing.py
2025-04-10 13:58:01 +02:00

270 lines
10 KiB
Python

#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
Estimate ray crossing at each depth using rays (as npy objects) precalculated from FMTOMO file rays.dat
(e.g. script: plot_rays_on_plane)
'''
import os
import glob
import json
import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import RegularGridInterpolator
from pylot.tomography.fmtomo_tools.fmtomo_grid_utils import read_vgrid_regular, read_vgrid
from pylot.tomography.map_utils import angle_marker, make_map
pjoin = os.path.join
def main(fdir, it=1, station_coords_file='/data/AdriaArray_Data/various/station_coords.json'):
fnin_vgrid = pjoin(fdir, f'it_{it}/vgrids.in')
fnout_vtk = pjoin(fdir, f'vgrids_it{it}_w_ray_crossings.vtk')
fnout_npy = pjoin(fdir, f'vgrids_hf_it{it}' + '_crossing_bins_{}-{}_mincount_{}_STEP_{}.npy')
# binsize (degrees), min_count for each direction to raise quality by 1 increment
binsize_az = 90.
binsize_incl = 60.
min_count = 5
plot = False
write = True
draw_hists = True
draw_simple_hists = True
step = 2 # take nth sample of vgrid for calculating cells
# for stations plotting
with open(station_coords_file, 'r') as infile:
stations = json.load(infile)
lonsStations, latsStations = zip(*[(sta['longitude'], sta['latitude']) for sta in stations.values()])
# some constants etc.
R = 6371.
bins_azim = np.arange(0, 360 + binsize_az, binsize_az)
bins_incl = np.arange(0., 60. + binsize_incl, binsize_incl)
nbins_azim = len(bins_azim) - 1
nbins_incl = len(bins_incl) - 1
# for plotting:
bins = bins_azim
binsize = binsize_az
nbins = nbins_azim
# get vgrid and estimate dlat/dlon
vgrid_reg = read_vgrid_regular(fnin_vgrid)
lons, lats, rads = [array[::step] for array in vgrid_reg[0]]
vps, covs, sms, pdvs = [array[::step, ::step, ::step] for array in vgrid_reg[1:]]
dlat = lats[1] - lats[0]
dlon = lons[1] - lons[0]
# grid = init_dict()
# grid['vps'] = list(vps.ravel())
# grid['covs'] = list(covs.ravel())
# grid['sms'] = list(sms.ravel())
# grid['pdvs'] = list(pdvs.ravel())
# grid['res'] = []
#
# for rad in rads:
# for lat in lats:
# for lon in lons:
# depth = R - rad
# grid['lons'].append(lon)
# grid['lats'].append(lat)
# grid['depths'].append(depth)
# x, y, z = pol2cart(lat, lon, R - depth)
# grid['xs'].append(x)
# grid['ys'].append(y)
# grid['zs'].append(z)
# just for plotting at certain depths (not for write!!!)
depths = np.arange(0., 800., 100)
if not write:
rads = R - depths
resolutions = []
lons_export = []
lats_export = []
rads_export = []
print(rads)
for rad in rads:
# prepare dict containing horizontal angle for key tuple (ilat, ilon)
vgrid_azim_ids = {}
vgrid_incl_ids = {}
print('Working on radius', rad)
# iterate over rays
for fnin_rays in glob.glob1(fpath_events, '*.npz'):
rays = np.load(os.path.join(fpath_events, fnin_rays), allow_pickle=True)
# raypoints = np.zeros((len(rays), 3))
#for index in range(len(rays)):
for ray in rays.values():
#ray = rays[index]
# get index of closest depth here from ray (TAKE CARE OF ALIASING! Using nearest value)
ind_min = np.abs(ray[:, 0] - rad).argmin()
# in case ind_min is at upper boundary (surface)
if ind_min == len(ray[:, 0]) - 1:
ind_min -= 1
# get diffs to following ray index (km, deg, deg)!
ray_diff = np.diff(ray[ind_min:ind_min + 2], axis=0)[0]
lat = ray[ind_min, 1]
lat_diff_km = ray_diff[1] * (np.pi * R) / 180.
lon_diff_km = ray_diff[2] * (np.pi * np.cos(np.deg2rad(lat)) * R) / 180.
r_diff_km = ray_diff[0]
dlateral = np.sqrt(lat_diff_km ** 2 + lon_diff_km ** 2)
# calculate horizontal angle
azim = np.rad2deg(np.arctan2(lon_diff_km, lat_diff_km))
incl = np.rad2deg(np.arctan(dlateral / r_diff_km))
# correct angle from -180-180 to 0-360 and also change azim to bazim
bazim = azim + 180.
# angles[index] = angle
# raypoints[index] = ray[ind_min]
lat, lon = ray[ind_min, 1:]
lati = np.where((lats <= lat + dlat / 2.) & (lats > lat - dlat / 2.))[0][0]
loni = np.where((lons <= lon + dlon / 2.) & (lons > lon - dlon / 2.))[0][0]
key = (lati, loni)
if not key in vgrid_azim_ids.keys():
vgrid_azim_ids[key] = []
vgrid_incl_ids[key] = []
vgrid_azim_ids[key].append(bazim)
vgrid_incl_ids[key].append(incl)
# vgrid_ids[index] = np.array([lati, loni])
# plt.scatter(lons[loni], lats[lati], c='r', marker='o', facecolor='none', alpha=0.5)
# sc = plt.scatter(raypoints[:,2], raypoints[:,1], c=angles[:])
vgrid_angles_quality = {}
vgrid_azim_hist = {}
for inds in vgrid_azim_ids.keys():
azims = vgrid_azim_ids[inds]
incls = vgrid_incl_ids[inds]
lati, loni = inds
hist_az, _ = np.histogram(azims, bins=bins_azim)
hist_incl, _ = np.histogram(incls, bins=bins_incl)
hist2d, _, _ = np.histogram2d(azims, incls, bins=[bins_azim, bins_incl])
# hist_az, hist_inc = hist2d
hist = hist2d.ravel()
quality = len(hist[hist > min_count - 1]) / (nbins_azim * nbins_incl)
# quality = len(hist_az[hist_az > min_count - 1]) / nbins_azim
# quality = len(hist_incl[hist_incl > min_count - 1]) / nbins_incl
vgrid_angles_quality[inds] = quality
vgrid_azim_hist[inds] = hist_az
# LATS = []
# LONS = []
qualities = []
hists = []
for ilat, lat in enumerate(lats):
for ilon, lon in enumerate(lons):
# LATS.append(lat)
# LONS.append(lon)
key = (ilat, ilon)
quality = vgrid_angles_quality.get(key)
hist = vgrid_azim_hist.get(key)
# print(key, quality)
if not quality:
quality = 0.
# hist for plotting only!
if hist is None:
hist = np.zeros(nbins)
qualities.append(quality)
hists.append(hist)
lons_export.append(lon)
lats_export.append(lat)
rads_export.append(rad)
resolutions += qualities
if plot:
# TODO: still old basemap code
raise NotImplementedError('Still using old basemap code')
fig = plt.figure()
ax = fig.add_subplot(111)
bmap = make_map(ax, resolution='l')
LONS, LATS = np.meshgrid(lons, lats)
sc = bmap.pcolormesh(LONS - dlon / 2., LATS - dlat / 2.,
np.array(qualities).reshape(LATS.shape), latlon=True, zorder=1.5,
shading='nearest') # , s=np.array(qualities)*100)
# sc = plt.contourf(LONS, LATS, np.array(qualities).reshape(LATS.shape), levels=21)
if draw_hists:
for index, bin in enumerate(bins[:-1]):
marker = angle_marker(bin, bin + binsize)
s = np.array(hists)[:, index]
if draw_simple_hists:
s[s < min_count] = 0
s[s >= min_count] = 150
else:
s *= 10.
sc_angle = bmap.scatter(LONS, LATS, s=s, marker=marker, edgecolors='0.75', alpha=1., linewidths=0.6,
latlon=True, zorder=1.5, )
sc_angle.set_facecolor('none')
# sc_angle = plt.scatter(LONS, LATS, s=1, marker='.', edgecolors='1.', alpha=1., linewidths=1.)
bmap.scatter(lonsStations, latsStations, c='k', s=0.5, latlon=True, zorder=1.5, ) # , alpha=0.5)
plt.title('Azimuthal coverage at depth of {}km, {} bins'.format((R - rad), nbins_azim * nbins_incl))
# plt.xlim([0, 22])
# plt.ylim([40, 53])
# plt.gca().set_aspect('equal')
cbar = plt.colorbar(sc)
cbar.set_label('Azimuthal coverage')
plt.show()
if write:
rginter_res = RegularGridInterpolator((rads, lats, lons),
np.array(resolutions).reshape((len(rads), len(lats), len(lons))),
bounds_error=False, fill_value=0.)
grid = read_vgrid(fnin_vgrid)[0]
grid['res'] = []
for lon, lat, depth in zip(grid['lons'], grid['lats'], grid['depths']):
grid['res'].append(rginter_res((R - depth, lat, lon)))
a = np.array([lons_export, lats_export, rads_export, resolutions])
np.save(fnout_npy.format(nbins_azim, nbins_incl, min_count, step), a)
write_vtk(grid, fnout_vtk,
write_data=['vps', 'res', 'covs', 'frechs']) # , clon=11., clat=46., dlon=12., dlat=6., sort=True)
def rays_to_npy(infile, fnout_npy, n_points=10):
rays = {}
i = 0
src_id_old = 1
while True:
i += 1
l1 = infile.readline()
if l1 == '': break
rec_id, src_id, ray_id = [int(item) for item in l1.split()[:3]]
if not src_id in rays.keys():
rays[src_id] = []
l2 = infile.readline()
n = int(l2.split()[0])
ray = np.zeros((n, 3))
for index in range(n):
r, lat, lon = [float(item) for item in infile.readline().split()]
ray[index] = np.array([r, np.rad2deg(lat), np.rad2deg(lon)])
rays[src_id].append(ray[::n_points])
dirname = os.path.split(fnout_npy)[0]
if not os.path.isdir(dirname):
os.mkdir(dirname)
for src_id, ray in rays.items():
np.savez(fnout_npy.format(src_id), *ray)
if __name__ == '__main__':
fmtomodir = '/data/AdriaArray_Data/fmtomo_adriaarray/alpadege/crust_incl_TESAURO_sm30_dm1/'
it = 12
fpath_events = pjoin(fmtomodir, f'rays_npy_it{it}/')
infile = open(pjoin(fmtomodir, f'it_{it}/rays.dat'), 'r')
fnout_npy = pjoin(fpath_events, f'it_{it}_rays_event_' + '{}.npz')
rays_to_npy(infile, fnout_npy)
main(fmtomodir, it=it)