<jupyter_start><jupyter_text>Import libraries<jupyter_code>import os
import yaml
import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import matplotlib
matplotlib.rcParams.update({'font.size': 8})

from obspy.core import Stream, Trace, UTCDateTime, Stats
from obspy.io.sac import SACTrace<jupyter_output><empty_output><jupyter_text>Event<jupyter_code># read event location
with open('./input/inparam.source.yaml', 'r') as file:
    source_yaml = yaml.load(file, Loader=yaml.FullLoader)
loc_leaf = source_yaml['list_of_sources'][0]['VIRGINIA_201108231751A']['location']
event_latlon = loc_leaf['latitude_longitude']
event_depth = loc_leaf['depth']<jupyter_output><empty_output><jupyter_text>GSN  Station info and map<jupyter_code># read station info
info_GSN = np.loadtxt('./input/GSN.txt', dtype=str, skiprows=3)

####################################
# draw a map of event and stations #
####################################

plt.figure(dpi=150)
# draw map
map = Basemap(projection='cyl', resolution='l', lon_0=0)
map.drawcoastlines(linewidth=0.25)
map.fillcontinents(color='ivory',lake_color='lightblue')
map.drawmapboundary(fill_color='lightblue', linewidth=0)
# draw event
map.scatter(event_latlon[1], event_latlon[0], latlon=True, 
            s=150, c='r', marker='*', lw=0, zorder=100)
# draw stations
map.scatter(info_GSN[:, 3].astype(float), info_GSN[:, 2].astype(float), latlon=True, 
            s=30, c='b', marker=7, lw=0, zorder=10)
plt.show()<jupyter_output><empty_output><jupyter_text>Read and plot seismograms<jupyter_code># specify a station key (network.name)
station_key = 'IU.ANMO'

# read time and displacement
gsn_dir = './output/stations/global_seismic_network_GSN'
time = np.loadtxt(gsn_dir + '/data_time.ascii')
disp = np.loadtxt(gsn_dir + '/%s.ascii' % station_key)

# plot
fig, ax = plt.subplots(3, sharex=True, dpi=150)
for ich, ch in enumerate('RTZ'):
    # change unit to mm
    ax[ich].plot(time, disp[:, ich] * 1e6, lw=1)
    ax[ich].text(.95, .9, 'channel = ' + ch, transform = ax[ich].transAxes, ha='right', va='top')
ax[1].set_ylabel('Amplitude (mm)')
ax[0].set_xlim(time[0], time[-1])
plt.xlabel('Time after source origin (s)')
plt.show()<jupyter_output><empty_output><jupyter_text>Processing using obspy<jupyter_code># trace header
stats = Stats()
stats.starttime = UTCDateTime(time[0])
stats.delta = UTCDateTime(time[1] - time[0])
stats.npts = len(time)

# stream
stream = Stream()
for ich, ch in enumerate('RTZ'):
    stats.channel = ch   
    stream.append(Trace(disp[:, ich], header=stats))

# process (filter, resample, slice, ...)
stream.filter('lowpass', freq=1/50)
stream.resample(1.)
stream = stream.slice(UTCDateTime(0.), UTCDateTime(1800.))

# print & plot
print(stream)
stream.plot()<jupyter_output><empty_output><jupyter_text>Save to SAC after down-sampling<jupyter_code># create dir
os.makedirs(gsn_dir + '/sac', exist_ok=True)

# sac header
sac_header = {}
sac_header['evla'] = event_latlon[0]
sac_header['evlo'] = event_latlon[1]
sac_header['evdp'] = float(event_depth) / 1e3

# loop over stations
print('Saving to SAC...')
for ist, st in enumerate(info_GSN):
    print('%d / %d' % (ist + 1, len(info_GSN)), end='\r')
    # sac header
    sac_header['kstnm'] = st[0]
    sac_header['knetwk'] = st[1]
    sac_header['stla'] = float(st[2])
    sac_header['stlo'] = float(st[3])
    sac_header['stdp'] = float(st[5])
    # read data
    disp = np.loadtxt(gsn_dir + '/%s.%s.ascii' % (st[1], st[0]))
    # loop over channels
    for ich, ch in enumerate('RTZ'):
        # sac header
        sac_header['kcmpnm'] = ch
        # add sac header to trace header
        stats.sac = sac_header
        # create and process trace
        tr = Trace(data=disp[:, ich], header=stats)
        tr.resample(1.)
        tr = tr.slice(UTCDateTime(0.), UTCDateTime(1800.))
        # create sac from trace
        sac = SACTrace.from_obspy_trace(tr)
        sac.write(gsn_dir + '/sac/%s.%s.%s.sac' % (st[1], st[0], ch))
print('Done with %d stations.' % len(info_GSN))<jupyter_output><empty_output><jupyter_text>USArray  Station info<jupyter_code># read station locations
info_US_TA = np.loadtxt('./input/US_TA.txt', dtype=str, skiprows=3)

# dict: station key -> [lat, lon]
nstation = len(info_US_TA)
stlatlon_dict = {}
for ist in np.arange(nstation):
    key = info_US_TA[ist, 1] + '.' + info_US_TA[ist, 0]
    stlatlon_dict[key] = np.array([float(info_US_TA[ist, 2]), float(info_US_TA[ist, 3])])<jupyter_output><empty_output><jupyter_text>Rank-to-station map<jupyter_code># data dir
us_ta_dir = './output/stations/USArray_transportable'

# read rank-station info
rank_station_info = np.loadtxt(us_ta_dir + '/rank_station.info', dtype=str, skiprows=1)

# dict: mpi-rank -> [station keys]
rank_station_dict = {}

# (lat, lon) of stations re-ordered by data
stlatlon_in_data_order = []

for item in rank_station_info:
    rank = item[0]
    stkey = item[1]
    # initialize with an empty array if rank does not exists in rank_station_dict
    if rank not in rank_station_dict.keys():
        rank_station_dict[rank] = []
    # append the station
    rank_station_dict[rank].append(stkey)
    stlatlon_in_data_order.append(stlatlon_dict[stkey])
    
# convert to numpy array
stlatlon_in_data_order = np.array(stlatlon_in_data_order)

# read time
time = np.loadtxt(us_ta_dir + '/data_time.ascii')
ntime = len(time)<jupyter_output><empty_output><jupyter_text>Animations on array<jupyter_code># choose a channel to animate
# U3   -- vertical displacement
# E_I1 -- trace of strain
# R3   -- vertical rotation
channel = 'R3'

# colormap norm in animation
if channel == 'U3':
    norm = 1e-6
elif channel == 'E_I1':
    norm = 1e-11
elif channel == 'R3':
    norm = 1e-11
else:
    assert False, "Invalid channel."
    
# allocate data
data = np.ndarray((ntime, nstation))

# loop over mpi-ranks to read data
for rank in rank_station_dict.keys():
    data_on_rank = np.loadtxt('%s/dir_rank%s/%s.ascii' % (us_ta_dir, rank, channel))
    data[:, 0:len(rank_station_dict[rank])] = data_on_rank
#############################
###### plot a snapshot ######
#############################

# specify a time step (0~384)
tstep = 100

# plot the snapshot
plt.figure(dpi=150)
plt.gca().axis('off')
plt.scatter(stlatlon_in_data_order[:, 1], stlatlon_in_data_order[:, 0], s=1, 
            c=data[tstep, :], vmin=-norm, vmax=norm, cmap='coolwarm')
plt.text(0, 0, 'Time = %.1f s' % (time[tstep]), transform = plt.gca().transAxes)
plt.colorbar(orientation='vertical', shrink=.5, label=channel)
plt.gca().set_aspect(1.3)
plt.show()
############################
###### make animation ######
############################

# create dir
os.makedirs(us_ta_dir + '/animation', exist_ok=True)

# create all snapshots
print('Making snapshots...')
for tstep in np.arange(len(time)):
    print('%d / %d' % (tstep + 1, len(time)), end='\r')
    plt.figure(dpi=150)
    plt.gca().axis('off')
    plt.scatter(stlatlon_in_data_order[:, 1], stlatlon_in_data_order[:, 0], s=1, 
                c=data[tstep, :], vmin=-norm, vmax=norm, cmap='coolwarm')
    plt.text(0, 1, 'Time = %.1f s' % (time[tstep]), transform = plt.gca().transAxes)
    plt.colorbar(orientation='vertical', shrink=.5, label=channel)
    plt.gca().set_aspect(1.3)
    plt.savefig(us_ta_dir + '/animation/%s.%04d.png' % (channel, tstep))
    plt.close()
    
# use ffmepg to combine snapshots to animation
print('Creating vedio using ffmpeg...')
os.system("ffmpeg -y -i %s/animation/%s.%%04d.png %s/animation/%s.mp4" % 
          (us_ta_dir, channel, us_ta_dir, channel))

# remove snapshots
os.system('rm ' + us_ta_dir + '/animation/%s.*.png' % (channel,))
print('Done.')
# play animation
from IPython.display import Video
Video("%s/animation/%s.mp4" % (us_ta_dir, channel))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18212090163934427, 'fraction_numerical': 0.019082991803278687, 'mean_word_length': 4.801634472511144, 'pattern_counts': {'":': 0, '<': 37, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '112510', 'license': 'MIT', 'n_tokens_mistral': 3034, 'n_tokens_neox': 2923, 'n_words': 595, 'path': 'examples/01_spherical_Earth_PREM_50s/post_processing.ipynb', 'repo_name': 'nicklinyi/AxiSEM-3D', 'size': '7808'}