-
Notifications
You must be signed in to change notification settings - Fork 3
/
hymodel.py
156 lines (145 loc) · 5.36 KB
/
hymodel.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
import glob
import pandas as pd
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap, addcyclic, cm
def normalize(data):
data = (data - data.min()) / (data.max() - data.min())
return data
class HyReceptor:
def __init__(
self,
ozone,
traj,
station_name="Davis",
latx=np.arange(-90, 91, 1),
lonx=np.arange(-180, 181, 1),
):
self.ozone = ozone
self.traj = traj[station_name]
self.lat = self.traj.sel(geo="lat").to_pandas()
self.lon = self.traj.sel(geo="lon").to_pandas()
self.latx = latx
self.lonx = lonx
self.dates = self.lat.columns
self.station = station_name
self.location = self.traj.isel(step=0, time=0).sel(geo=["lat", "lon"]).values
def get_density(self, lon, lat, density=True):
density, lonx, latx = np.histogram2d(
lon, lat, [self.lonx, self.latx], density=density
)
return density
def get_weight(self, psc):
wgt = np.ones_like(psc)
avg = np.average(psc)
wgt[psc >= 2 * avg] = 1.0
wgt[(avg <= psc) & (psc < 2 * avg)] = 0.75
wgt[(0.5 * avg <= psc) & (psc < avg)] = 0.5
wgt[(psc < 0.5 * avg)] = 0.25
return wgt
def to_nc(self, data, name="Potential Source Contribution Function"):
pscf = xr.DataArray(
data,
coords=[self.lonx[:-1], self.latx[:-1]],
dims=["Longitude", "Latitude"],
)
pscf.name = name
pscf.attrs["station_name"] = self.station
return pscf
def calculate_pscf(self, thresh=0.5, multi=False, weighted=True):
self.ozone_threshold = self.ozone.quantile(thresh)
sub = self.ozone[self.ozone > self.ozone_threshold]
slat = self.lat[sub.index]
slon = self.lon[sub.index]
self.dates_with_more_than_threshold_ozone = sub.index
num1 = self.get_density(
self.lon.values.flatten(), self.lat.values.flatten(), density=False
)
num2 = self.get_density(
slon.values.flatten(), slat.values.flatten(), density=False
)
self.pscf = num2 / num1
self.rt_pscf = num2
if weighted:
wgt = self.get_weight(self.rt_pscf)
self.pscf *= wgt
self.pscf = self.to_nc(self.pscf, name="PSCF")
if not multi:
return self.pscf
else:
return num2, num1, np.sum(num1)
def calculate_cwt(self, weighted=True):
cwt = np.zeros((self.lonx.shape[0] - 1, self.latx.shape[0] - 1))
tau = np.zeros_like(cwt)
nums = np.zeros_like(cwt)
for column in self.lat.columns:
tlen = float(len(self.lat[column].values))
num = self.get_density(
self.lon[column].values.flatten(),
self.lat[column].values.flatten(),
density=False,
)
ttau = num / tlen
cwt += ttau * self.ozone[column]
tau += ttau
nums += num
self.cwt = cwt / tau
self.rt_cwt = nums
if weighted:
wgt = self.get_weight(self.rt_cwt)
self.cwt *= wgt
self.cwt = self.to_nc(self.cwt, name="CWT")
return self.cwt
def calculate_rtwc(self, normalise=True):
err = 100
bcwt = self.calculate_cwt(weighted=True)
bcwt = np.nan_to_num(bcwt.values)
while err >= 0.001:
cwt = np.zeros((self.lonx.shape[0] - 1, self.latx.shape[0] - 1))
tau = np.zeros_like(cwt)
nums = np.zeros_like(cwt)
for column in self.lat.columns:
ozo = self.ozone[column]
num = self.get_density(
self.lon[column].values.flatten(),
self.lat[column].values.flatten(),
density=False,
)
tlen = float(len(self.lat[column].values))
ttau = num / tlen
tavg = np.sum(bcwt * num / tlen)
cwt += ttau * ozo * bcwt / tavg
tau += ttau
nums += num
cwt = np.nan_to_num(cwt / tau)
err = (cwt - bcwt) / bcwt
err = np.nansum(err)
bcwt = cwt
if normalise:
bcwt = normalize(bcwt)
bcwt[bcwt == 0] = np.nan
self.rtwc = self.to_nc(bcwt, name="RTWC")
return self.rtwc
def plot_map(self, pscf, boundinglat=-20, axes=[]):
if len(axes) == 0:
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
cax = fig.add_axes([0.15, 0.06, 0.7, 0.04])
else:
fig, ax, cax = axes
m = Basemap(
projection="spstere", lon_0=180, boundinglat=boundinglat, round=True, ax=ax
)
data, lonx = addcyclic(pscf.T.values, pscf["Longitude"])
latx, lonx = np.meshgrid(pscf["Latitude"], lonx)
lonx, latx = m(lonx, latx)
if pscf.max().values <= 1:
h = m.contourf(lonx, latx, data.T, levels=np.arange(0, 1.1, 0.1))
else:
levels = np.arange(0, pscf.max(), 2)
h = m.contourf(lonx, latx, data.T, levels=levels)
plt.colorbar(h, cax=cax, orientation="horizontal")
m.drawcoastlines(linewidth=1.5)
m.drawcountries(linewidth=0.55)
ax.set_title(self.station, fontweight="bold")
return fig, ax, m