-
Notifications
You must be signed in to change notification settings - Fork 0
/
utilities.py
198 lines (171 loc) · 7.34 KB
/
utilities.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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
import numpy as np
from os import chdir, mkdir, path
from matplotlib import pyplot as pp
from Chandra.Time import DateTime
from os import mkdir, chdir
def mkdir_cd(dir):
# Make a directory (if doesn't already exist) and cd to it.
if not path.exists(dir):
mkdir(dir)
chdir(dir)
def smaller_axes(ax1):
# Replace existing axes with one only 90% width and 75% height,
# allowing more room for ticks and titles
ax_pos = ax1.get_position().get_points()
ax_width = ax_pos[1,0] - ax_pos[0,0]
ax_height = ax_pos[1,1] - ax_pos[0,1]
pp.delaxes(ax1)
ax2 = pp.axes([ax_pos[0,0] + .10 * ax_width, ax_pos[0,1] + .20 * ax_height,
.90 * ax_width, .75 * ax_height])
return ax2
def reshape_array(array, num_rows, order='F'):
# Return an array based on the values of the input array, but reshaped
# with a user-specified number of rows. Excess data will be trucated.
#
# Order = 'F' (column-major, FORTRAN-style)
# = 'C' (row-major, C-style)
if np.mod(len(array),num_rows) > 0:
out = array[:-np.mod(len(array), num_rows)].reshape((num_rows, np.floor(len(array) / int(num_rows))), order=order)
else:
out = array.reshape((num_rows, len(stdev) / int(num_rows)), order=order)
return out
def count_by_month(times):
"""Counts the number of times an event occurs each month.
Inputs a set of DateTime times (in any format).
Outputs an array of DateTime times spanning the timeframe with one entry per month
and the count of events that occur within that month
"""
len_mo = 365.25/12*24*3600
#Create an array that spans the timeframe with one date in the middle-ish of every month
x_start = DateTime(min(DateTime(times).secs)).iso[:7] + '-15 12:00:00.00'
x_stop = DateTime(max(DateTime(times).secs)).iso[:7] + '-15 12:00:00.00'
x_times = np.arange(DateTime(x_start).secs, DateTime(x_stop).secs, len_mo)
#Collect the year and month for each date
t_month = [DateTime(t).iso[:7] for t in times]
x_month = [DateTime(x).iso[:7] for x in x_times]
num_months = len(x_times)
month_counts = np.zeros(num_months)
for i in np.arange(num_months):
match = [t == x_month[i] for t in t_month]
month_counts[i] = np.sum(match)
month_times = x_times
return month_times, month_counts
def sum_by_month(times, vals):
"""Sums events by month.
Inputs a set of DateTime times (in any format).
Outputs an array of DateTime times spanning the timeframe with one entry per month
and the sum of events that occur within that month
"""
len_mo = 365.25/12*24*3600
#Create an array that spans the timeframe with one date in the middle-ish of every month
x_start = DateTime(min(DateTime(times).secs)).iso[:7] + '-15 12:00:00.00'
x_stop = DateTime(max(DateTime(times).secs)).iso[:7] + '-15 12:00:00.00'
x_times = np.arange(DateTime(x_start).secs, DateTime(x_stop).secs, len_mo)
#Collect the year and month for each date
t_month = [DateTime(t).iso[:7] for t in times]
x_month = [DateTime(x).iso[:7] for x in x_times]
num_months = len(x_times)
month_sums = np.zeros(num_months)
for i in np.arange(num_months):
match = np.array([t == x_month[i] for t in t_month])
month_sums[i] = np.sum(vals[match])
month_times = x_times
return month_times, month_sums
def mean_by_month(times, vals):
"""Sums events by month.
Inputs a set of DateTime times (in any format).
Outputs an array of DateTime times spanning the timeframe with one entry per month
and the sum of events that occur within that month
"""
len_mo = 365.25/12*24*3600
#Create an array that spans the timeframe with one date in the middle-ish of every month
x_start = DateTime(min(DateTime(times).secs)).iso[:7] + '-15 12:00:00.00'
x_stop = DateTime(max(DateTime(times).secs)).iso[:7] + '-15 12:00:00.00'
x_times = np.arange(DateTime(x_start).secs, DateTime(x_stop).secs, len_mo)
#Collect the year and month for each date
t_month = [DateTime(t).iso[:7] for t in times]
x_month = [DateTime(x).iso[:7] for x in x_times]
num_months = len(x_times)
month_means = np.zeros(num_months)
for i in np.arange(num_months):
match = np.array([t == x_month[i] for t in t_month])
month_means[i] = np.mean(vals[match])
month_times = x_times
return month_times, month_means
def heat_map(x, y, bins=(20,20), colorbar=True, **kwargs):
"""Plots a "heat map", i.e. a scatter plot highlighting density
Based on an exmple from
http://www.physics.ucdavis.edu/~dwittman/Matplotlib-examples/
Inputs:
x = x vals (as from a scatter plot)
y = y vals (as from a scatter plot)
bins = int or tuple bins to feed into histogram2d
colorbar = whether or not to display colorbar (default=True)
x_lim = x limits of plot
y_lim = y_limits of plot
"""
if kwargs.has_key('x_lim'):
x_lim = kwargs.pop('x_lim')
else:
x_lim = [np.min(x), np.max(x)]
if kwargs.has_key('y_lim'):
y_lim = kwargs.pop('y_lim')
else:
y_lim = [min(y), max(y)]
hist,xedges,yedges = np.histogram2d(x,y,bins=bins,range=[x_lim,y_lim])
extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
pp.imshow(hist.T,extent=extent,interpolation='nearest',origin='lower', aspect='auto')
if colorbar==True:
pp.colorbar()
pp.show()
def same_limits(subs):
"""Sets the x-limits and y-limits for a given set of subplots
equal to the collective mins and maxes
Inputs:
subs = subplots (tuple of 3-digit subplot numbers)
e.g. same_limits((411, 412, 413, 414))
"""
all_limits = np.zeros((len(subs), 4))
for sub, i in zip(subs, range(len(subs))):
pp.subplot(sub)
all_limits[i,:] = pp.axis()
best_limits = (min(all_limits[:,0]), max(all_limits[:,1]),
min(all_limits[:,2]), max(all_limits[:,3]))
for i in subs:
pp.subplot(i)
pp.axis(best_limits)
def append_to_array(a, pos=-1, val=0):
"""Appends a zero (or user-defined value) to a given one-dimensional array,
either at the end (pos=-1) or beginning (pos=0).
e.g. append_to_array(arange(5),pos=-1)
returns array([0, 1, 2, 3, 4, 0])
"""
val_a = np.array([val])
if pos==0:
out = np.concatenate([val_a, a])
elif pos==-1:
out = np.concatenate([a, val_a])
return out
def find_closest(a, b):
"""This function returns an array of length a with the indices of
array b that are closest to the values of array a.
"""
a = np.atleast_1d(np.array(a))
b = np.atleast_1d(np.array(b))
out = [np.argmin(abs(b - a1)) for a1 in a]
return out
def find_last_before(a, b):
"""This function returns an array of length a with the indices of
array b that are closest without going over the values of array a.
(Bob Barker rules. Assumes b is sorted by magnitude.)
"""
out = np.searchsorted(b,a,side='left') - 1
out[out==-1] = 0
return out
def find_first_after(a, b):
"""This function returns an array of length a with the indices of
array b that are closest without being less than the values of array a.
(Opposite of Bob Barker rules. Assumes b is sorted by magnitude.)
"""
out = np.searchsorted(b,a,side='right')
return out