wiener_filter_harmonic.py 3.5 KB
Newer Older
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
from nifty import *
from mpi4py import MPI
import plotly.offline as py
import plotly.graph_objs as go


comm = MPI.COMM_WORLD
rank = comm.rank


def plot_maps(x, name):

    trace = [None]*len(x)

    keys = x.keys()
    field = x[keys[0]]
    domain = field.domain[0]
    shape = len(domain.shape)
    max_n = domain.shape[0]*domain.distances[0]
    step = domain.distances[0]
    x_axis = np.arange(0, max_n, step)

    if shape == 1:
        for ii in xrange(len(x)):
            trace[ii] = go.Scatter(x= x_axis, y=x[keys[ii]].val.get_full_data(), name=keys[ii])
        fig = go.Figure(data=trace)

        py.plot(fig, filename=name)

    elif shape == 2:
        for ii in xrange(len(x)):
            py.plot([go.Heatmap(z=x[keys[ii]].val.get_full_data())], filename=keys[ii])
    else:
        raise TypeError("Only 1D and 2D field plots are supported")

def plot_power(x, name):

    layout = go.Layout(
        xaxis=dict(
            type='log',
            autorange=True
        ),
        yaxis=dict(
            type='log',
            autorange=True
        )
    )

    trace = [None]*len(x)

    keys = x.keys()
    field = x[keys[0]]
    domain = field.domain[0]
    x_axis = domain.kindex

    for ii in xrange(len(x)):
        trace[ii] = go.Scatter(x= x_axis, y=x[keys[ii]].val.get_full_data(), name=keys[ii])

    fig = go.Figure(data=trace, layout=layout)
    py.plot(fig, filename=name)

np.random.seed(42)



if __name__ == "__main__":

    distribution_strategy = 'not'

    # setting spaces
    npix = np.array([500])  # number of pixels
    total_volume = 1.  # total length

    # setting signal parameters
    lambda_s = .05  # signal correlation length
    sigma_s = 10.  # signal variance


    #setting response operator parameters
    length_convolution = .025
    exposure = 1.

    # calculating parameters
    k_0 = 4. / (2 * np.pi * lambda_s)
    a_s = sigma_s ** 2. * lambda_s * total_volume

    # creation of spaces
88
89
90
91
92
93
    # x1 = RGSpace([npix,npix], distances=total_volume / npix,
    #              zerocenter=False)
    # k1 = RGRGTransformation.get_codomain(x1)

    x1 = HPSpace(64)
    k1 = HPLMTransformation.get_codomain(x1)
94
    p1 = PowerSpace(harmonic_partner=k1, logarithmic=False)
95
96
97
98
99
100
101

    # creating Power Operator with given spectrum
    spec = (lambda k: a_s / (1 + (k / k_0) ** 2) ** 2)
    p_field = Field(p1, val=spec)
    S_op = create_power_operator(k1, spec)

    # creating FFT-Operator and Response-Operator with Gaussian convolution
102
    # adjust dtype_target probperly
103
104
    Fft_op = FFTOperator(domain=x1, target=k1,
                        domain_dtype=np.float64,
105
                        target_dtype=np.float64)
106
107
108
109
    R_op = ResponseOperator(x1, sigma=[length_convolution],
                            exposure=[exposure])

    # drawing a random field
110
    sk = p_field.power_synthesize(real_power=True, mean=0.)
111
    s = Fft_op.adjoint_times(sk)
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127

    signal_to_noise = 1
    N_op = DiagonalOperator(R_op.target, diagonal=s.var()/signal_to_noise, bare=True)
    n = Field.from_random(domain=R_op.target,
                          random_type='normal',
                          std=s.std()/np.sqrt(signal_to_noise),
                          mean=0)

    d = R_op(s) + n

    # Wiener filter
    j = Fft_op.times(R_op.adjoint_times(N_op.inverse_times(d)))
    D = HarmonicPropagatorOperator(S=S_op, N=N_op, R=R_op)

    mk = D(j)

128
    m = Fft_op.adjoint_times(mk)
129

130
131
132
133
134
135
136
    # z={}
    # z["signal"] = s
    # z["reconstructed_map"] = m
    # z["data"] = d
    # z["lambda"] = R_op(s)
    #
    # plot_maps(z, "Wiener_filter.html")
137
138