reconstruction.py 6.09 KB
Newer Older
Philipp Arras's avatar
Philipp Arras committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2019-2020 Max-Planck-Society
Philipp Frank's avatar
Philipp Frank committed
15
16
# Author: Philipp Arras, Philipp Frank, Philipp Haim, Reimar Leike,
# Jakob Knollmueller
Philipp Arras's avatar
Philipp Arras committed
17
18

import sys
Philipp Frank's avatar
Philipp Frank committed
19
20
import gc
from mpi4py import MPI
Philipp Arras's avatar
Philipp Arras committed
21
22
from functools import reduce
from operator import add
Philipp Frank's avatar
Philipp Frank committed
23
from time import time
Philipp Arras's avatar
Philipp Arras committed
24
25

import nifty6 as ift
Philipp Arras's avatar
Philipp Arras committed
26
import src as vlbi
Philipp Frank's avatar
Philipp Frank committed
27
from config import comm, nranks, rank, master
Philipp Frank's avatar
Philipp Frank committed
28
from config import doms, eps, min_timestamps_per_bin, nthreads
Philipp Arras's avatar
Philipp Arras committed
29
from config import sky_movie_mf as sky
Philipp Frank's avatar
Philipp Frank committed
30
from config import startt, dt, npixt
Philipp Arras's avatar
Philipp Arras committed
31

Philipp Frank's avatar
Philipp Frank committed
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def stat_plotting(pos, KL):
    if master:
        sc = ift.StatCalculator()
        sc_mean = ift.StatCalculator()
        sc_spectral = ift.StatCalculator()
        skysamples = {'hi': [], 'lo': []}

    # CAUTION: the loop and the if-clause cannot be interchanged!
    for ss in KL.samples:
        if master:
            samp = sky(pos+ss)
            skysamples['hi'].append(samp['hi'])
            skysamples['lo'].append(samp['lo'])
            sc.add(samp)
            sc_mean.add(0.5*(samp['hi']+samp['lo']).val)
            sc_spectral.add(2 * (samp['lo']-samp['hi']).val / (samp['lo']+samp['hi']).val )

49
50
def optimization_heuristic(ii, likelihoods):
    lh_full, lh_full_amp, lh_full_ph, lh_cut, lh_cut_amp, lh_cut_ph = likelihoods 
Philipp Frank's avatar
Philipp Frank committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
    N_samples = 10 * (1 + ii // 8)
    N_iterations = 4 * (4 + ii // 4) if ii<50 else 20

    eps = 0.1
    clvl = 3 if ii < 20 else 2

    ic_newton = ift.AbsDeltaEnergyController(eps, iteration_limit=N_iterations,
                                             name=f'it_{ii}' if master else None,
                                             convergence_level=clvl)

    if ii < 50:
        minimizer = ift.VL_BFGS(ic_newton)
    else:
        minimizer = ift.NewtonCG(ic_newton)

    if ii < 30:
        if ii % 2 == 0 or ii < 10:
68
            lh = lh_cut
Philipp Frank's avatar
Philipp Frank committed
69
        elif ii % 4 == 1:
70
            lh = lh_cut_amp
Philipp Frank's avatar
Philipp Frank committed
71
        else:
72
            lh = lh_cut_ph
Philipp Frank's avatar
Philipp Frank committed
73
74
75
76
    else:
        if ii % 2 == 0 or ii > 50:
            lh = lh_full
        elif ii % 4 == 1:
77
            lh = lh_full_amp
Philipp Frank's avatar
Philipp Frank committed
78
        else:
79
            lh = lh_full_ph
Philipp Frank's avatar
Philipp Frank committed
80
81
    return minimizer, N_samples, N_iterations, lh

Philipp Frank's avatar
Philipp Frank committed
82
def build_likelihood(rawdd, startt, npixt, dt, mode):
83
    lh = []
Philipp Frank's avatar
Philipp Frank committed
84
    active_inds = []
Philipp Arras's avatar
Philipp Arras committed
85
    for freq in vlbi.data.FREQS:
86
        rawd = rawdd[freq]
Philipp Frank's avatar
Philipp Frank committed
87
88
89
        args = {'tmin': startt, 'tmax': npixt*dt, 'delta_t': dt}

        for ii, dd in enumerate(vlbi.time_binning(rawd, **args)):
Philipp Arras's avatar
Philipp Arras committed
90
91
92
93
            if len(dd) == 0:
                continue
            ind = str(ii) + "_" + freq
            active_inds.append(ind)
Philipp Frank's avatar
Philipp Frank committed
94
95
96

            nfft = vlbi.RadioResponse(doms, dd['uv'], eps, nthreads)
            vis = ift.makeField(nfft.target, dd['vis'])
97
98
99
100
101
102
103
104
            llh = []
            if mode in ['amp', 'full']:
                vis2closampl, evalsampl = vlbi.Visibilities2ClosureAmplitudes(dd)
                llh.append(ift.GaussianEnergy(mean=vis2closampl(vis)) @ vis2closampl)
            if mode in ['ph', 'full']:
                vis2closph, evalsph, _  = vlbi.Visibilities2ClosurePhases(dd)
                llh.append(ift.GaussianEnergy(mean=vis2closph(vis)) @ vis2closph)
            llh_op = reduce(add, llh) @ nfft.ducktape(ind)
Philipp Frank's avatar
Philipp Frank committed
105
106
107
            if mode == 'full' and freq == vlbi.data.FREQS[0]:
                ift.extra.check_jacobian_consistency(llh_op, ift.from_random(llh_op.domain),
                                                     tol=1e-5, ntries=10)
108
            lh.append(llh_op)
Philipp Frank's avatar
Philipp Frank committed
109

110
    conv = vlbi.DomainTuple2MultiField(sky.target, active_inds)
Philipp Frank's avatar
Philipp Frank committed
111
    return reduce(add, lh) @ conv
Philipp Frank's avatar
Philipp Frank committed
112

113
114
115
116
def setup():
    if len(sys.argv) != 3:
        raise RuntimeError
    _, pre_data, fname_input = sys.argv
Philipp Frank's avatar
Philipp Frank committed
117

118
119
120
121
122
    rawd = {}
    for freq in vlbi.data.FREQS:
        rawd[freq] = vlbi.combined_data(f'data/{pre_data}', [freq], min_timestamps_per_bin) if master else None
    if nranks > 1:
        rawd = comm.bcast(rawd)
Philipp Frank's avatar
Philipp Frank committed
123

124
    pre_output = pre_data
Philipp Frank's avatar
Philipp Frank committed
125

Philipp Frank's avatar
Philipp Frank committed
126
127
128
    lh_full = build_likelihood(rawd, startt, npixt, dt, 'full')
    lh_full_ph = build_likelihood(rawd, startt, npixt, dt, 'ph')
    lh_full_amp = build_likelihood(rawd, startt, npixt, dt, 'amp')
Philipp Frank's avatar
Philipp Frank committed
129

Philipp Frank's avatar
Philipp Frank committed
130
131
132
    lh_cut = build_likelihood(rawd, startt, npixt//2, dt, 'full')
    lh_cut_ph = build_likelihood(rawd, startt, npixt//2, dt, 'ph')
    lh_cut_amp = build_likelihood(rawd, startt, npixt//2, dt, 'amp')
Philipp Frank's avatar
Philipp Frank committed
133
134
135
136
137
138
139
140
141
142
143

    pos = vlbi.load_hdf5(fname_input, sky.domain) if master else None
    if nranks > 1:
        pos = comm.bcast(pos)

    ic = ift.AbsDeltaEnergyController(0.5, iteration_limit=200, name=f'Sampling(task {rank})', convergence_level=3)

    if master:
        t0 = time()
        (lh_full @ sky)(pos)
        print(f'Likelihood call: {1000*(time()-t0):.0f} ms')
144
    return pos, sky, ic, pre_output, (lh_full, lh_full_amp, lh_full_ph, lh_cut, lh_cut_amp, lh_cut_ph)
Philipp Frank's avatar
Philipp Frank committed
145
146
147
148


# Encapsulate everything in functions, to avoid as many (unintended) global variables as possible
def main():
Philipp Frank's avatar
Philipp Frank committed
149
    pos, sky, ic, pre_output, likelihoods = setup()
Philipp Frank's avatar
Philipp Frank committed
150
151
152

    for ii in range(60):
        gc.collect()
153
        minimizer, N_samples, N_iterations, lh = optimization_heuristic(ii, likelihoods)
Philipp Frank's avatar
Philipp Frank committed
154
155
156

        if master:
            print(f'Iter: {ii}, N_samples: {N_samples}, N_iter: {N_iterations}')
Philipp Frank's avatar
Philipp Frank committed
157

Philipp Frank's avatar
Philipp Frank committed
158
159
160
161
162
163
        ll = lh @ sky
        H = ift.StandardHamiltonian(ll, ic)
        KL = ift.MetricGaussianKL(pos, H, N_samples, comm=comm, mirror_samples=True)
        del ll, H
        gc.collect()

Philipp Arras's avatar
Philipp Arras committed
164
165
        KL, _ = minimizer(KL)
        pos = KL.position
Philipp Frank's avatar
Philipp Frank committed
166
167
168
169
170
171
172
173
174
175

        vlbi.save_state(sky, KL.position, f'{pre_output}_', ii, samples=KL.samples, master=master)
        del KL
        gc.collect()


if __name__ == '__main__':
    with open("initial_random_state.txt", "rb") as f:
        ift.random.setState(f.read())
    main()