P_incoherent.py 1.9 KB
Newer Older
jansen31's avatar
jansen31 committed
1
2
3
4
5
6
7
"""
Written by Matthijs Jansen, October 2020.

Contains two proxoperators which apply to a set of iterates, similar to incoherent modes of a light field
or (in my application) degenerate molecular orbitals.
"""

8
from numpy import sum, where
9
10
11
12
13
14
15
16
17
18
from numpy.core._multiarray_umath import sqrt, array

from proxtoolbox.proxoperators import P_M, P_Sparsity_real

__all__ = ['P_M_incoherent', 'P_Sparsity_real_incoherent']


class P_M_incoherent(P_M):
    """
    Apply the Fourier-domain modulus constraint to a set of incoherent fields:
jansen31's avatar
jansen31 committed
19
20
21
    scale the amplitudes such that the combined intensity matches the data.

    Data scaling is done while avoiding division by zero errors similar to how it is done in proxoperator.magproj
22
23
24
25
    """
    def __init__(self, experiment):
        super(P_M_incoherent, self).__init__(experiment)
        self.intensity = self.data ** 2
jansen31's avatar
jansen31 committed
26
        self.eps = 1e-10
27
28
29
30
31
32
33
34
35
36
37

    def eval(self, u, prox_idx=None):
        amplitudes = self.prop.eval(u)  # Propagate to measurement domain
        pred = sum(abs(amplitudes)**2, axis=0)  # Predicted total intensity
        divisor = sqrt(pred + self.eps**2)  #
        scaling = 1 - (pred + 2*self.eps**2) / divisor**3 * (pred / divisor - self.data)
        out = array([scaling * u_i for u_i in amplitudes])
        return self.invprop.eval(out)  # Propagate back


class P_Sparsity_real_incoherent(P_Sparsity_real):
38
39
40
    """
    Apply a sparsity constraint to the combined intensity of a set of incoherent iterates
    """
41
42
43
44
    def __init__(self, experiment):
        super(P_Sparsity_real_incoherent, self).__init__(experiment)

    def eval(self, u, prox_idx=None):
45
46
47
48
49
        u_real = u.real
        density = sum(abs(u_real)**2, axis=0)
        sparse = super(P_Sparsity_real_incoherent, self).eval(density)
        out = array([where(sparse != 0, u_i, 0) for u_i in u_real])
        # out = array([super(P_Sparsity_real_incoherent, self).eval(u_i) for u_i in u])
50
        return out