light_cone_operator.py 4.74 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
from ..operators.linear_operator import LinearOperator
from ..operators.operator import Operator


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


Philipp Arras's avatar
Cleanup    
Philipp Arras committed
32
def _make_coords(domain, absolute=False):
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    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
49
class _LightConeDerivative(LinearOperator):
50
    def __init__(self, domain, target, derivatives):
Philipp Arras's avatar
Philipp Arras committed
51
        super(_LightConeDerivative, self).__init__()
52
53
54
55
56
57
58
        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)
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
59
        x = x.val
60
61
62
63
64
        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:
65
                res[i] = np.sum(self._derivatives[i]*x.real)
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
66
        return Field(self._tgt(mode), res)
67

Philipp Arras's avatar
Cleanup    
Philipp Arras committed
68

Philipp Arras's avatar
Philipp Arras committed
69
def _cone_arrays(c, domain, sigx, want_gradient):
Philipp Arras's avatar
Cleanup    
Philipp Arras committed
70
    x = _make_coords(domain)
71
    a = np.zeros(domain.shape, dtype=np.complex128)
72
    if want_gradient:
73
        derivs = np.zeros((c.size,) + domain.shape, dtype=np.complex128)
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    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
96

97
class LightConeOperator(Operator):
Philipp Frank's avatar
Philipp Frank committed
98
    '''Constructs a Light cone from a set of lightspeed parameters.
Martin Reinecke's avatar
Martin Reinecke committed
99

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

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

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

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

Philipp Frank's avatar
Philipp Frank committed
111
112
113
114
115
    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
116

Philipp Frank's avatar
Philipp Frank committed
117
118
119
120
121
122
123
124
125
126
    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
127
    '''
128
    def __init__(self, domain, target, sigx):
Philipp Frank's avatar
Philipp Frank committed
129
130
        self._domain = DomainTuple.make(domain)
        self._target = DomainTuple.make(target)
131
132
        self._sigx = sigx

Philipp Arras's avatar
Philipp Arras committed
133
    def apply(self, x):
Martin Reinecke's avatar
more    
Martin Reinecke committed
134
        lin = x.jac is not None
Philipp Arras's avatar
Fixes    
Philipp Arras committed
135
        a, derivs = _cone_arrays(x.val.val if lin else x.val, self.target, self._sigx, lin)
Martin Reinecke's avatar
stage 3    
Martin Reinecke committed
136
        res = Field(self.target, a)
Philipp Arras's avatar
Philipp Arras committed
137
        if not lin:
Philipp Frank's avatar
cleanup    
Philipp Frank committed
138
            return res
Philipp Arras's avatar
Philipp Arras committed
139
        jac = _LightConeDerivative(self._domain, self._target, derivs)
Philipp Arras's avatar
Philipp Arras committed
140
        return x.new(res, jac)