-
Notifications
You must be signed in to change notification settings - Fork 18
/
io.jl
149 lines (127 loc) · 3.76 KB
/
io.jl
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
export save_fft, load_fft, load_corr, save_corr
"""
save_fft(F::FFTData, OUT::String)
Save FFTData `F` to JLD2.
"""
function save_fft(F::FFTData, FFTOUT::String)
# check if FFTDIR exists
FFTOUT = expanduser(FFTOUT)
if isdir(FFTOUT) == false
mkpath(FFTOUT)
end
# check if F is on the GPU
if isa(F.fft,AbstractGPUArray)
F = F |> cpu
end
# create JLD2 file and save FFT
net,sta,loc,chan = split(F.name,'.')
filename = joinpath(FFTOUT,"$(F.name).jld2")
file = jldopen(filename, "a+")
if !(chan in keys(file))
group = JLD2.Group(file, chan)
group[F.id] = F
else
file[chan][F.id] = F
end
close(file)
return filename
end
"""
load_fft(filename,chan,day=day)
Loads FFTData for channel `chan` from JLD2 file `filename`.
If day is specified, loads data from day `day`, else
loads data from all days of `chan`.
"""
function load_fft(filename::String,chan::String;day::Union{String,Missing}=missing)
file = jldopen(filename,"r")
if ismissing(day)
days = keys(file[chan])
num_days = length(days)
# get sizes of all corrs
sizes = Array{Int}(undef,num_days,2)
for ii = 1:num_days
sizes[ii,:]=collect(size(file[chan][days[ii]].fft))
end
# create empty arrays and get indicies
M = sizes[1]
N = sum(sizes[:,2])
ffts = Array{eltype(file[chan][days[1]].fft)}(undef,M,N)
t = Array{eltype(file[chan][days[1]].t)}(undef,N)
starts = cumsum(vcat(1 ,sizes[:,2]))[1:end-1]
ends = cumsum(sizes[:,2])
for ii = 1:num_days
ffts[:,starts[ii]:ends[ii]] = file[chan][days[ii]].fft
t[starts[ii]:ends[ii]] = file[chan][days[ii]].t
end
# add data to FFTData
F = file[chan][days[1]]
F.t = t
F.fft = ffts
else
F = file[chan][day]
end
close(file)
return F
end
"""
save_corr(C::CorrData, OUT::String)
Save CorrData `C` to JLD2.
"""
function save_corr(C::CorrData, CORROUT::String)
# check if CORRDIR exists
CORROUT = expanduser(CORROUT)
if isdir(CORROUT) == false
mkpath(CORROUT)
end
# check if C is on the GPU
if isa(C.corr,AbstractGPUArray)
C = C |> cpu
end
# create JLD2 file and save correlation
filename = joinpath(CORROUT,"$(C.name).jld2")
file = jldopen(filename, "a+")
if !(C.comp in keys(file))
group = JLD2.Group(file, C.comp)
group[C.id] = C
else
file[C.comp][C.id] = C
end
close(file)
return filename
end
"""
load_corr(filename,chan,day=day)
Loads CorrData for channel `chan` on day `day` from JLD2 file `filename`.
"""
function load_corr(filename::String,chan::String;
day::Union{String,Missing}=missing)
file = jldopen(filename,"r")
if ismissing(day) # load all files
days = keys(file[chan])
num_days = length(days)
# get sizes of all corrs
sizes = Array{Int}(undef,num_days,2)
for ii = 1:num_days
sizes[ii,:]=collect(size(file[chan][days[ii]].corr))
end
# create empty arrays and get indicies
M = sizes[1]
N = sum(sizes[:,2])
corrs = Array{eltype(file[chan][days[1]].corr)}(undef,M,N)
t = Array{eltype(file[chan][days[1]].t)}(undef,N)
starts = cumsum(vcat(1 ,sizes[:,2]))[1:end-1]
ends = cumsum(sizes[:,2])
for ii = 1:num_days
corrs[:,starts[ii]:ends[ii]] = file[chan][days[ii]].corr
t[starts[ii]:ends[ii]] = file[chan][days[ii]].t
end
# add data to CorrData
C = file[chan][days[1]]
C.t = t
C.corr = corrs
else # return single day
C = file[chan][day]
end
close(file)
return C
end