P_M.py 4.2 KB
Newer Older
1
from proxtoolbox.proxoperators.proxoperator import ProxOperator, magproj
Matthijs's avatar
Matthijs committed
2
from numpy import where, log, exp, sum
3

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

6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

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

    def __init__(self, experiment):
        """
        Initialization
        
        Parameters
        ----------
        config : dict - Dictionary containing the problem configuration. It must contain the following mapping:
        'M' : array_like - non-negative real FOURIER DOMAIN CONSTRAINT
        """
        self.data = experiment.data
22
23
        self.prop = experiment.propagator(experiment)
        self.invprop = experiment.inverse_propagator(experiment)
24

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


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

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

        Parameters
        ----------
55
        experiment: experiment class
56
57
        """
        super(P_M_masked, self).__init__(experiment)
58
        self.mask = experiment.fmask == 0
59

60
    def eval(self, u, prox_idx=None):
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
        """
        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
        """
        fourier_space_iterate = self.prop.eval(u)
        constrained = magproj(fourier_space_iterate.copy(), self.data)
        update = where(self.mask, fourier_space_iterate, constrained)
        return self.invprop(update)
Matthijs's avatar
Matthijs committed
76
77
78
79
80
81
82
83
84
85


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

86
    def eval(self, u, prox_idx=None):
Matthijs's avatar
Matthijs committed
87
88
89
90
91
92
93
94
95
96
97
98
99
        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:
Matthijs's avatar
Matthijs committed
100
101
            b = magproj(u_hat, self.data)
            return self.invprop.eval(b)
Matthijs's avatar
Matthijs committed
102
103
104
105
106
107
108
109
110
111
112
113
        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

114
    def eval(self, u, prox_idx=None):
Matthijs's avatar
Matthijs committed
115
116
117
118
119
120
121
122
123
124
125
126
        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:
Matthijs's avatar
Matthijs committed
127
128
129
            constrained = magproj(u_hat.copy(), self.data)  # Apply constraint
            update = where(self.mask, u_hat, constrained)  # Masking operation
            return self.invprop(update)  # Propagate back
Matthijs's avatar
Matthijs committed
130
131
        else:
            return u