light_cone_operator.py 4.94 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
        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:
66
                res[i] = np.sum(self._derivatives[i]*x.real)
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 Frank's avatar
Philipp Frank committed
99
    '''Constructs a Light cone from a set of lightspeed parameters.
Martin Reinecke's avatar
Martin Reinecke committed
100

Philipp Frank's avatar
Philipp Frank committed
101
    The resulting cone is defined as follows
Martin Reinecke's avatar
Martin Reinecke committed
102

Philipp Frank's avatar
Philipp Frank committed
103
104
    .. math::
        \\exp \\left(- \\frac{1}{2} \\Re \\left( \\Delta \\right)^2 \\right)
Martin Reinecke's avatar
Martin Reinecke committed
105

Philipp Frank's avatar
Philipp Frank committed
106
    with
Martin Reinecke's avatar
Martin Reinecke committed
107

Philipp Frank's avatar
Philipp Frank committed
108
109
110
    .. math::
        \\Delta = \\sqrt{- \\left(t^2 - \\frac{x^\\dagger C^{-1} x}
        {\\sigma_x^2} \\right)}
Martin Reinecke's avatar
Martin Reinecke committed
111

Philipp Frank's avatar
Philipp Frank committed
112
113
114
115
116
    where t and x are the coordinates of the target space. Note that axis zero
    of the space is interpreted as the time axis. C denotes the input
    paramters of the operator and parametrizes the shape of the cone.
    sigx is the width of the asymptotic Gaussian in x necessary for
    discretization.
Martin Reinecke's avatar
Martin Reinecke committed
117

Philipp Frank's avatar
Philipp Frank committed
118
119
120
121
122
123
124
125
126
127
    Parameters
    ----------
    domain : Domain, tuple of Domain or DomainTuple
        The domain of the input parameters of the light cone, the values of the
        lightspeed tensor.
    target : Domain, tuple of Domain or DomainTuple
        Output space on which the lightcone should be defined. The zeroth axis
        of this space is interpreted as the time axis.
    sigx : float
        Width of the Gaussian for the discretized representation of the cone.
Philipp Arras's avatar
Philipp Arras committed
128
    '''
129
    def __init__(self, domain, target, sigx):
Philipp Frank's avatar
Philipp Frank committed
130
131
        self._domain = DomainTuple.make(domain)
        self._target = DomainTuple.make(target)
132
133
134
        self._sigx = sigx

    def apply(self, x):
Philipp Frank's avatar
cleanup  
Philipp Frank committed
135
136
        islin = isinstance(x, Linearization)
        val = x.val.to_global_data() if islin else x.to_global_data()
Philipp Arras's avatar
Philipp Arras committed
137
        a, derivs = _cone_arrays(val, self.target, self._sigx, islin)
Philipp Frank's avatar
cleanup  
Philipp Frank committed
138
139
140
        res = Field.from_global_data(self.target, a)
        if not islin:
            return res
Philipp Arras's avatar
Philipp Arras committed
141
        jac = _LightConeDerivative(x.jac.target, self.target, derivs)(x.jac)
Philipp Frank's avatar
cleanup  
Philipp Frank committed
142
        return Linearization(res, jac, want_metric=x.want_metric)