P_M.py 4.18 KB
Newer Older
jansen31's avatar
jansen31 committed
1
2
from numpy import where, log, sum, errstate

3
4
from proxtoolbox.proxoperators.proxoperator import ProxOperator, magproj

Matthijs's avatar
Matthijs committed
5
__all__ = ['P_M', 'P_M_masked', "Approx_P_M", 'Approx_P_M_masked']
6

7
8
9
10
11
12
13
14
15
16

class P_M(ProxOperator):
    """
    Projection onto Fourier magnitude constraints
    """

    def __init__(self, experiment):
        """
        Initialization
        """
jansen31's avatar
jansen31 committed
17
        super().__init__(experiment)
18
        self.data = experiment.data
19
20
        self.prop = experiment.propagator(experiment)
        self.invprop = experiment.inverse_propagator(experiment)
21

22
    def eval(self, u, prox_idx=None):
23
24
25
26
27
28
29
30
31
        """
        Applies the proxoperator P_M
        
        Parameters
        ----------
        u : array_like - function in the physical domain to be projected
        
        Returns
        -------
32
33
        array_like - [p_M,phat_M] ,where p_M = the projection IN THE PHYSICAL (time) DOMAIN
            and phat_M = projection IN THE FOURIER DOMAIN
34
35
36
        """
        m = self.data
        a = self.prop.eval(u)
37
        b = magproj(m, a)
38
39
40
41
42
        return self.invprop.eval(b)


class P_M_masked(P_M):
    """
Matthijs's avatar
Matthijs committed
43
    Projection onto Fourier magnitude constraints, leaving alone points masked by 'fmask'
44
45
46
47
48
49
50
51
    """

    def __init__(self, experiment):
        """
        Initialization

        Parameters
        ----------
52
        experiment: experiment class
53
54
        """
        super(P_M_masked, self).__init__(experiment)
jansen31's avatar
jansen31 committed
55
        self.mask = experiment.fmask  # True means not updated
56

57
    def eval(self, u, prox_idx=None):
58
59
60
61
62
63
64
65
66
67
68
        """
        Applies the proxoperator P_M_masked

        Parameters
        ----------
        u : array_like - function in the physical domain to be projected

        Returns
        -------
        array_like - p_M: the projection IN THE PHYSICAL (time) DOMAIN
        """
jansen31's avatar
jansen31 committed
69
        # assert u.shape == self.mask.shape, 'Non-matching iterate array!'
70
        fourier_space_iterate = self.prop.eval(u)
71
        constrained = magproj(self.data, fourier_space_iterate.copy())
72
        update = where(self.mask, fourier_space_iterate, constrained)
jansen31's avatar
jansen31 committed
73
        return self.invprop.eval(update)
Matthijs's avatar
Matthijs committed
74
75
76
77
78
79
80
81
82
83


class Approx_P_M(P_M):
    def __init__(self, experiment):
        super(Approx_P_M, self).__init__(experiment)
        self.data_sq = self.data ** 2
        self.data_zeros = where(self.data == 0)
        self.data_ball = experiment.data_ball
        self.TOL2 = experiment.TOL2

84
    def eval(self, u, prox_idx=None):
Matthijs's avatar
Matthijs committed
85
86
87
        u_hat = self.prop.eval(u)
        u_hat_sq = abs(u_hat)**2

jansen31's avatar
jansen31 committed
88
89
        with errstate(divide='ignore', invalid='ignore'):
            tmp = u_hat_sq / self.data_sq
Matthijs's avatar
Matthijs committed
90
91
92
93
94
95
96
97
98
        tmp[self.data_zeros] = 1
        u_hat_sq[self.data_zeros] = 0
        I_u_hat = tmp == 0
        tmp[I_u_hat] = 1
        tmp = log(tmp)
        h_u_hat = sum(u_hat_sq * tmp + self.data_sq - u_hat_sq).real
        # Now see that the propagated field is within the ball around the data (if any).
        # If not, the projection is calculated, otherwise we do nothing.
        if h_u_hat >= self.data_ball + self.TOL2:
99
            b = magproj(self.data, u_hat)
Matthijs's avatar
Matthijs committed
100
            return self.invprop.eval(b)
Matthijs's avatar
Matthijs committed
101
102
103
104
105
106
107
108
109
110
111
112
        else:
            return u


class Approx_P_M_masked(P_M_masked):
    def __init__(self, experiment):
        super(Approx_P_M_masked, self).__init__(experiment)
        self.data_sq = self.data ** 2
        self.data_zeros = where(self.data == 0)
        self.data_ball = experiment.data_ball
        self.TOL2 = experiment.TOL2

113
    def eval(self, u, prox_idx=None):
Matthijs's avatar
Matthijs committed
114
115
116
117
118
119
120
121
122
123
124
125
        u_hat = self.prop.eval(u)
        u_hat_sq = abs(u_hat) ** 2
        tmp = u_hat_sq / self.data_sq
        tmp[self.data_zeros] = 1
        u_hat_sq[self.data_zeros] = 0
        I_u_hat = tmp == 0
        tmp[I_u_hat] = 1
        tmp = log(tmp)
        h_u_hat = sum(u_hat_sq * tmp + self.data_sq - u_hat_sq).real
        # Now see that the propagated field is within the ball around the data (if any).
        # If not, the projection is calculated, otherwise we do nothing.
        if h_u_hat >= self.data_ball + self.TOL2:
126
            constrained = magproj(self.data, u_hat.copy())  # Apply constraint
Matthijs's avatar
Matthijs committed
127
128
            update = where(self.mask, u_hat, constrained)  # Masking operation
            return self.invprop(update)  # Propagate back
Matthijs's avatar
Matthijs committed
129
130
        else:
            return u