light_cone_operator.py 3.84 KB
Newer Older
Philipp Arras's avatar
Philipp Arras committed
1
2
3
4
5
6
7
8
9
10
11
12
13
# 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/>.
#
Martin Reinecke's avatar
Martin Reinecke committed
14
# Copyright(C) 2013-2019 Max-Planck-Society
Philipp Arras's avatar
Philipp Arras committed
15
#
Martin Reinecke's avatar
Martin Reinecke committed
16
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
17

Philipp Arras's avatar
Philipp Arras committed
18
19
import numpy as np

20
from ..domain_tuple import DomainTuple
Philipp Arras's avatar
Philipp Arras committed
21
from ..field import Field
22
23
24
25
26
from ..linearization import Linearization
from ..operators.linear_operator import LinearOperator
from ..operators.operator import Operator


Philipp Arras's avatar
Fixup    
Philipp Arras committed
27
28
29
30
31
32
def _field_from_function(domain, func, absolute=False):
    domain = DomainTuple.make(domain)
    k_array = _make_coords(domain, absolute=absolute)
    return Field.from_global_data(domain, func(k_array))


Philipp Arras's avatar
Cleanup    
Philipp Arras committed
33
def _make_coords(domain, absolute=False):
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
    domain = DomainTuple.make(domain)
    dim = len(domain.shape)
    dist = domain[0].distances
    shape = domain.shape
    k_array = np.zeros((dim,) + shape)
    for i in range(dim):
        ks = np.minimum(shape[i] - np.arange(shape[i]), np.arange(
            shape[i]))*dist[i]
        if not absolute:
            ks[int(shape[i]/2) + 1:] *= -1
        fst_dims = (1,)*i
        lst_dims = (1,)*(dim - i - 1)
        k_array[i] += ks.reshape(fst_dims + (shape[i],) + lst_dims)
    return k_array


Philipp Arras's avatar
Philipp Arras committed
50
class _LightConeDerivative(LinearOperator):
51
    def __init__(self, domain, target, derivatives):
Philipp Arras's avatar
Philipp Arras committed
52
        super(_LightConeDerivative, self).__init__()
53
54
55
56
57
58
59
        self._domain = domain
        self._target = target
        self._derivatives = derivatives
        self._capability = self.TIMES | self.ADJOINT_TIMES

    def apply(self, x, mode):
        self._check_input(x, mode)
Philipp Frank's avatar
Philipp Frank committed
60
        x = x.to_global_data()
61
62
63
64
65
66
        res = np.zeros(self._tgt(mode).shape, dtype=self._derivatives.dtype)
        for i in range(self.domain.shape[0]):
            if mode == self.TIMES:
                res += self._derivatives[i]*x[i]
            else:
                res[i] = np.sum(self._derivatives[i]*x)
67
        return Field.from_global_data(self._tgt(mode), res)
68

Philipp Arras's avatar
Cleanup    
Philipp Arras committed
69

Philipp Arras's avatar
Philipp Arras committed
70
def _cone_arrays(c, domain, sigx, want_gradient):
Philipp Arras's avatar
Cleanup    
Philipp Arras committed
71
    x = _make_coords(domain)
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
    a = np.zeros(domain.shape, dtype=np.complex)
    if want_gradient:
        derivs = np.zeros((c.size,) + domain.shape, dtype=np.complex)
    else:
        derivs = None
    a -= (x[0]/(sigx*domain[0].distances[0]))**2
    for i in range(c.size):
        res = (x[i + 1]/(sigx*domain[0].distances[i + 1]))**2
        a += c[i]*res
        if want_gradient:
            derivs[i] = res
    a = np.sqrt(a)
    if want_gradient:
        derivs *= -0.5
        for i in range(c.size):
            derivs[i][a == 0] = 0.
            derivs[i][a != 0] /= a[a != 0]
    a = a.real
    if want_gradient:
        derivs *= a
    a = np.exp(-0.5*a**2)
    if want_gradient:
        derivs = a*derivs.real
    return a, derivs

Philipp Arras's avatar
Cleanup    
Philipp Arras committed
97

98
class LightConeOperator(Operator):
Philipp Arras's avatar
Philipp Arras committed
99
100
101
    '''
    FIXME
    '''
102
103
104
105
106
107
    def __init__(self, domain, target, sigx):
        self._domain = domain
        self._target = target
        self._sigx = sigx

    def apply(self, x):
Philipp Frank's avatar
cleanup    
Philipp Frank committed
108
109
        islin = isinstance(x, Linearization)
        val = x.val.to_global_data() if islin else x.to_global_data()
Philipp Arras's avatar
Philipp Arras committed
110
        a, derivs = _cone_arrays(val, self.target, self._sigx, islin)
Philipp Frank's avatar
cleanup    
Philipp Frank committed
111
112
113
        res = Field.from_global_data(self.target, a)
        if not islin:
            return res
Philipp Arras's avatar
Philipp Arras committed
114
        jac = _LightConeDerivative(x.jac.target, self.target, derivs)(x.jac)
Philipp Frank's avatar
cleanup    
Philipp Frank committed
115
        return Linearization(res, jac, want_metric=x.want_metric)