G4-DESCSS/utils/draw.py
2022-05-12 20:39:38 +08:00

91 lines
2.6 KiB
Python

import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.basemap import Basemap
def data_read(file):
with open(file, 'r') as f:
data = f.readlines()
row_header = int(data[0].split(',')[1].strip())
row_variable = int(data[0].split(',')[5].strip())
row_column = int(data[0].split(',')[6].strip())
var = data[row_header - row_variable:row_header]
data = [tuple(s.replace('\n', '').strip() for s in row.split(',')) for row in data[row_header:-1]]
dtype = []
for row in var:
row = [s.replace("'", '').strip() for s in row.split(',')]
num = int(row[2])
type_s = 'float' if num > 0 else 'string'
for k in range(num):
name_s = row[0] + ('_{:d}'.format(k + 1) if num > 1 else '')
dtype.append((name_s, type_s))
assert len(dtype) == row_column
data = np.array(data, dtype=dtype)
return dtype, data
def orbit():
_, data = data_read('assets/orbit.csv')
lat, lon = data['Latitude'], data['Longitude'] - 180
m = Basemap()
m.drawcoastlines()
m.fillcontinents(color='white', lake_color='lightskyblue')
m.drawmapboundary(fill_color='skyblue')
m.scatter(lon, lat, s=1)
plt.savefig('docs/orbit.png', bbox_inches='tight')
def GCR(M=1, ax=None):
_, data = data_read('assets/CREME86-M{:d}.csv'.format(M))
E = data['Energy']
proton_i, proton_d = data['IFlux_1'], data['DFlux_1']
alpha_i, alpha_d = data['IFlux_2'], data['DFlux_2']
if ax:
ax1 = ax
else:
_, ax1 = plt.subplots(1, 1)
ax1.plot(E, proton_i, label=r'$p$')
ax1.plot(E, alpha_i, label=r'$\alpha$')
ax1.set_ylim([1, 5 * 1e2])
ax1.set_yscale('log')
ax1.set_ylabel(r'$Integrated\ Flux\ (\mathrm{m^{-2}sr^{-1}s^{-1}})$')
ax1.legend(loc="upper left")
ax2 = ax1.twinx()
ax2.plot(E, proton_d, linestyle=':', label=r'$p$')
ax2.plot(E, alpha_d, linestyle=':', label=r'$\alpha$')
ax2.set_ylim([1e-5, 5 * 1e-2])
ax2.set_yscale('log')
ax2.set_ylabel(r'$Differential\ Flux\ (\mathrm{m^{-2}sr^{-1}s^{-1}MeV^{-1}})$')
ax2.legend(loc="upper right")
ax1.set_xlabel(r'$Energy\ (\mathrm{MeV})$')
ax1.set_xlim(1, 1e5)
if ax is None:
plt.title('Average spectra - M = {:d}'.format(M))
plt.xscale('log')
plt.savefig('docs/spectra-M{:d}.png'.format(M), bbox_inches='tight')
else:
ax.set_title('Average spectra - M = {:d}'.format(M))
ax.set_xscale('log')
fig = plt.figure(figsize=(20, 10), dpi=150)
GCR(1, fig.add_subplot(1, 2, 1))
GCR(3, fig.add_subplot(1, 2, 2))
fig.savefig('docs/spectra.png', bbox_inches='tight')
GCR(1)
GCR(3)