linearization.py 12.7 KB
Newer Older
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/>.
#
Philipp Arras's avatar
Philipp Arras committed
14
# Copyright(C) 2013-2020 Max-Planck-Society
15
16
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
Martin Reinecke's avatar
Martin Reinecke committed
17
18
19
20

import numpy as np

from .sugar import makeOp
Martin Reinecke's avatar
misc    
Martin Reinecke committed
21
from .operators.operator import Operator
Martin Reinecke's avatar
Martin Reinecke committed
22
23


Martin Reinecke's avatar
misc    
Martin Reinecke committed
24
class Linearization(Operator):
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
25
    """Let `A` be an operator and `x` a field. `Linearization` stores the value
26
27
28
29
30
    of the operator application (i.e. `A(x)`), the local Jacobian
    (i.e. `dA(x)/dx`) and, optionally, the local metric.

    Parameters
    ----------
Philipp Arras's avatar
Philipp Arras committed
31
32
    val : Field or MultiField
        The value of the operator application.
33
    jac : LinearOperator
Philipp Arras's avatar
Philipp Arras committed
34
35
36
37
38
39
        The Jacobian.
    metric : LinearOperator or None
        The metric. Default: None.
    want_metric : bool
        If True, the metric will be computed for other Linearizations derived
        from this one. Default: False.
40
    """
41
    def __init__(self, val, jac, metric=None, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
42
43
        self._val = val
        self._jac = jac
Martin Reinecke's avatar
Martin Reinecke committed
44
45
        if self._val.domain != self._jac.target:
            raise ValueError("domain mismatch")
46
        self._want_metric = want_metric
Martin Reinecke's avatar
Martin Reinecke committed
47
48
        self._metric = metric

49
    def new(self, val, jac, metric=None):
50
51
52
53
54
        """Create a new Linearization, taking the `want_metric` property from
           this one.

        Parameters
        ----------
Philipp Arras's avatar
Philipp Arras committed
55
        val : Field or MultiField
56
57
58
            the value of the operator application
        jac : LinearOperator
            the Jacobian
Philipp Arras's avatar
Philipp Arras committed
59
60
        metric : LinearOperator or None
            The metric. Default: None.
61
        """
62
63
        return Linearization(val, jac, metric, self._want_metric)

Philipp Arras's avatar
Philipp Arras committed
64
    def trivial_jac(self):
Martin Reinecke's avatar
Martin Reinecke committed
65
        return self.make_var(self._val, self._want_metric)
Philipp Arras's avatar
Philipp Arras committed
66

Philipp Arras's avatar
Philipp Arras committed
67
68
69
70
71
72
73
    def prepend_jac(self, jac):
        metric = None
        if self._metric is not None:
            from .operators.sandwich_operator import SandwichOperator
            metric = None if self._metric is None else SandwichOperator.make(jac, self._metric)
        return self.new(self._val, self._jac @ jac, metric)

Martin Reinecke's avatar
Martin Reinecke committed
74
75
    @property
    def domain(self):
Philipp Arras's avatar
Philipp Arras committed
76
        """DomainTuple or MultiDomain : the Jacobian's domain"""
Martin Reinecke's avatar
Martin Reinecke committed
77
78
79
80
        return self._jac.domain

    @property
    def target(self):
Philipp Arras's avatar
Philipp Arras committed
81
        """DomainTuple or MultiDomain : the Jacobian's target (i.e. the value's domain)"""
Martin Reinecke's avatar
Martin Reinecke committed
82
83
84
85
        return self._jac.target

    @property
    def val(self):
Philipp Arras's avatar
Philipp Arras committed
86
        """Field or MultiField : the value"""
Martin Reinecke's avatar
Martin Reinecke committed
87
88
89
90
        return self._val

    @property
    def jac(self):
91
        """LinearOperator : the Jacobian"""
Martin Reinecke's avatar
Martin Reinecke committed
92
93
        return self._jac

Martin Reinecke's avatar
Martin Reinecke committed
94
95
    @property
    def gradient(self):
Philipp Arras's avatar
Philipp Arras committed
96
        """Field or MultiField : the gradient
97
98
99
100
101

        Notes
        -----
        Only available if target is a scalar
        """
Martin Reinecke's avatar
Martin Reinecke committed
102
        from .field import Field
Martin Reinecke's avatar
Martin Reinecke committed
103
        return self._jac.adjoint_times(Field.scalar(1.))
Martin Reinecke's avatar
Martin Reinecke committed
104

105
106
    @property
    def want_metric(self):
Martin Reinecke's avatar
Martin Reinecke committed
107
        """bool : True iff the metric was requested in the constructor"""
108
109
        return self._want_metric

Martin Reinecke's avatar
Martin Reinecke committed
110
111
    @property
    def metric(self):
112
113
114
115
116
117
        """LinearOperator : the metric

        Notes
        -----
        Only available if target is a scalar
        """
Martin Reinecke's avatar
Martin Reinecke committed
118
119
        return self._metric

Martin Reinecke's avatar
Martin Reinecke committed
120
    def __getitem__(self, name):
Martin Reinecke's avatar
Martin Reinecke committed
121
        return self.new(self._val[name], self._jac.ducktape_left(name))
Martin Reinecke's avatar
Martin Reinecke committed
122

Martin Reinecke's avatar
Martin Reinecke committed
123
    def __neg__(self):
124
125
        return self.new(-self._val, -self._jac,
                        None if self._metric is None else -self._metric)
Martin Reinecke's avatar
Martin Reinecke committed
126

Martin Reinecke's avatar
Martin Reinecke committed
127
    def conjugate(self):
128
        return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
129
130
131
132
133
            self._val.conjugate(), self._jac.conjugate(),
            None if self._metric is None else self._metric.conjugate())

    @property
    def real(self):
134
        return self.new(self._val.real, self._jac.real)
Martin Reinecke's avatar
Martin Reinecke committed
135

Martin Reinecke's avatar
Martin Reinecke committed
136
    def _myadd(self, other, neg):
Martin Reinecke's avatar
Martin Reinecke committed
137
138
139
140
141
142
143
144
145
        if np.isscalar(other) or other.jac is None:
            return self.new(self._val-other if neg else self._val+other,
                            self._jac, self._metric)
        met = None
        if self._metric is not None and other._metric is not None:
            met = self._metric._myadd(other._metric, neg)
        return self.new(
            self.val.flexible_addsub(other.val, neg),
            self.jac._myadd(other.jac, neg), met)
Martin Reinecke's avatar
Martin Reinecke committed
146
147
148

    def __add__(self, other):
        return self._myadd(other, False)
Martin Reinecke's avatar
Martin Reinecke committed
149
150

    def __radd__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
151
        return self._myadd(other, False)
Martin Reinecke's avatar
Martin Reinecke committed
152
153

    def __sub__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
154
        return self._myadd(other, True)
Martin Reinecke's avatar
Martin Reinecke committed
155
156
157
158

    def __rsub__(self, other):
        return (-self).__add__(other)

159
    def __truediv__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
160
161
162
        if np.isscalar(other):
            return self.__mul__(1/other)
        return self.__mul__(other.ptw("reciprocal"))
163
164

    def __rtruediv__(self, other):
165
        return self.ptw("reciprocal").__mul__(other)
166

Martin Reinecke's avatar
Martin Reinecke committed
167
    def __pow__(self, power):
Martin Reinecke's avatar
Martin Reinecke committed
168
        if not (np.isscalar(power) or power.jac is None):
Martin Reinecke's avatar
Martin Reinecke committed
169
            return NotImplemented
Martin Reinecke's avatar
misc    
Martin Reinecke committed
170
        return self.ptw("power", power)
Martin Reinecke's avatar
Martin Reinecke committed
171

Martin Reinecke's avatar
Martin Reinecke committed
172
    def __mul__(self, other):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
173
174
175
176
        if np.isscalar(other):
            if other == 1:
                return self
            met = None if self._metric is None else self._metric.scale(other)
177
            return self.new(self._val*other, self._jac.scale(other), met)
Martin Reinecke's avatar
Martin Reinecke committed
178
179
        from .sugar import makeOp
        if other.jac is None:
Martin Reinecke's avatar
Martin Reinecke committed
180
181
            if self.target != other.domain:
                raise ValueError("domain mismatch")
182
            return self.new(self._val*other, makeOp(other)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
183
184
185
186
187
188
        if self.target != other.target:
            raise ValueError("domain mismatch")
        return self.new(
            self.val*other.val,
            (makeOp(other.val)(self.jac))._myadd(
             makeOp(self.val)(other.jac), False))
Martin Reinecke's avatar
Martin Reinecke committed
189
190

    def __rmul__(self, other):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
191
        return self.__mul__(other)
Martin Reinecke's avatar
Martin Reinecke committed
192

193
    def outer(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
194
195
196
197
198
199
200
201
202
203
204
205
        """Computes the outer product of this Linearization with a Field or
        another Linearization

        Parameters
        ----------
        other : Field or MultiField or Linearization

        Returns
        -------
        Linearization
            the outer product of self and other
        """
206
        if np.isscalar(other):
Martin Reinecke's avatar
Martin Reinecke committed
207
            return self.__mul__(other)
Martin Reinecke's avatar
Martin Reinecke committed
208
209
        from .operators.outer_product_operator import OuterProduct
        if other.jac is None:
Martin Reinecke's avatar
Martin Reinecke committed
210
211
212
            return self.new(OuterProduct(other.domain, self._val)(other),
                            OuterProduct(other.domain, self._jac(self._val)))
        tmp_op = OuterProduct(other.target, self._val)
Martin Reinecke's avatar
Martin Reinecke committed
213
        return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
214
215
216
            tmp_op(other._val),
            OuterProduct(other.target, self._jac(self._val))._myadd(
                tmp_op(other._jac), False))
217

Martin Reinecke's avatar
Martin Reinecke committed
218
    def vdot(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
219
220
221
222
223
224
225
226
227
228
229
230
        """Computes the inner product of this Linearization with a Field or
        another Linearization

        Parameters
        ----------
        other : Field or MultiField or Linearization

        Returns
        -------
        Linearization
            the inner product of self and other
        """
Martin Reinecke's avatar
Martin Reinecke committed
231
        from .operators.simple_linear_operators import VdotOperator
Martin Reinecke's avatar
Martin Reinecke committed
232
        if other.jac is None:
233
            return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
234
                self._val.vdot(other),
Martin Reinecke's avatar
Martin Reinecke committed
235
                VdotOperator(other)(self._jac))
236
        return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
237
            self._val.vdot(other._val),
Martin Reinecke's avatar
Martin Reinecke committed
238
239
            VdotOperator(self._val)(other._jac) +
            VdotOperator(other._val)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
240

241
    def sum(self, spaces=None):
Martin Reinecke's avatar
Martin Reinecke committed
242
243
244
245
246
247
248
249
250
251
252
253
254
        """Computes the (partial) sum over self

        Parameters
        ----------
        spaces : None, int or list of int
            - if None, sum over the entire domain
            - else sum over the specified subspaces

        Returns
        -------
        Linearization
            the (partial) sum
        """
255
        from .operators.contraction_operator import ContractionOperator
Martin Reinecke's avatar
Martin Reinecke committed
256
257
258
        return self.new(
            self._val.sum(spaces),
            ContractionOperator(self._jac.target, spaces)(self._jac))
259
260

    def integrate(self, spaces=None):
Martin Reinecke's avatar
Martin Reinecke committed
261
262
263
264
265
266
267
268
269
270
271
272
273
        """Computes the (partial) integral over self

        Parameters
        ----------
        spaces : None, int or list of int
            - if None, integrate over the entire domain
            - else integrate over the specified subspaces

        Returns
        -------
        Linearization
            the (partial) integral
        """
Philipp Arras's avatar
Philipp Arras committed
274
275
        from .operators.contraction_operator import IntegrationOperator
        return IntegrationOperator(self._target, spaces)(self)
Martin Reinecke's avatar
Martin Reinecke committed
276

Martin Reinecke's avatar
misc    
Martin Reinecke committed
277
278
    def ptw(self, op, *args, **kwargs):
        t1, t2 = self._val.ptw_with_deriv(op, *args, **kwargs)
Martin Reinecke's avatar
Martin Reinecke committed
279
        return self.new(t1, makeOp(t2)(self._jac))
Philipp Arras's avatar
Philipp Arras committed
280

Martin Reinecke's avatar
Martin Reinecke committed
281
    def add_metric(self, metric):
282
        return self.new(self._val, self._jac, metric)
Martin Reinecke's avatar
Martin Reinecke committed
283

Martin Reinecke's avatar
Martin Reinecke committed
284
285
286
    def with_want_metric(self):
        return Linearization(self._val, self._jac, self._metric, True)

Martin Reinecke's avatar
Martin Reinecke committed
287
    @staticmethod
288
    def make_var(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
        """Converts a Field to a Linearization, with a unity Jacobian

        Parameters
        ----------
        field : Field or Multifield
            the field to be converted
        want_metric : bool
            If True, the metric will be computed for other Linearizations
            derived from this one. Default: False.

        Returns
        -------
        Linearization
            the requested Linearization
        """
Martin Reinecke's avatar
Martin Reinecke committed
304
        from .operators.scaling_operator import ScalingOperator
305
        return Linearization(field, ScalingOperator(field.domain, 1.),
306
                             want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
307
308

    @staticmethod
309
    def make_const(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
        """Converts a Field to a Linearization, with a zero Jacobian

        Parameters
        ----------
        field : Field or Multifield
            the field to be converted
        want_metric : bool
            If True, the metric will be computed for other Linearizations
            derived from this one. Default: False.

        Returns
        -------
        Linearization
            the requested Linearization

        Notes
        -----
        The Jacobian is square and contains only zeroes.
        """
Martin Reinecke's avatar
Martin Reinecke committed
329
        from .operators.simple_linear_operators import NullOperator
330
331
        return Linearization(field, NullOperator(field.domain, field.domain),
                             want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
332

Martin Reinecke's avatar
Martin Reinecke committed
333
334
    @staticmethod
    def make_const_empty_input(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
        """Converts a Field to a Linearization, with a zero Jacobian

        Parameters
        ----------
        field : Field or Multifield
            the field to be converted
        want_metric : bool
            If True, the metric will be computed for other Linearizations
            derived from this one. Default: False.

        Returns
        -------
        Linearization
            the requested Linearization

        Notes
        -----
        The Jacobian has an empty input domain, i.e. its matrix representation
        has 0 columns.
        """
Martin Reinecke's avatar
Martin Reinecke committed
355
356
        from .operators.simple_linear_operators import NullOperator
        from .multi_domain import MultiDomain
Martin Reinecke's avatar
Martin Reinecke committed
357
358
359
        return Linearization(
            field, NullOperator(MultiDomain.make({}), field.domain),
            want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
360

Martin Reinecke's avatar
Martin Reinecke committed
361
362
    @staticmethod
    def make_partial_var(field, constants, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
        """Converts a MultiField to a Linearization, with a Jacobian that is
        unity for some MultiField components and a zero matrix for others.

        Parameters
        ----------
        field : Multifield
            the field to be converted
        constants : list of string
            the MultiField components for which the Jacobian should be
            a zero matrix.
        want_metric : bool
            If True, the metric will be computed for other Linearizations
            derived from this one. Default: False.

        Returns
        -------
        Linearization
            the requested Linearization

        Notes
        -----
        The Jacobian is square.
        """
Martin Reinecke's avatar
Martin Reinecke committed
386
        from .operators.scaling_operator import ScalingOperator
Philipp Arras's avatar
Typo    
Philipp Arras committed
387
        from .operators.block_diagonal_operator import BlockDiagonalOperator
Martin Reinecke's avatar
Martin Reinecke committed
388
389
390
        if len(constants) == 0:
            return Linearization.make_var(field, want_metric)
        else:
391
            ops = {key: ScalingOperator(dom, 0. if key in constants else 1.)
392
393
                   for key, dom in field.domain.items()}
            bdop = BlockDiagonalOperator(field.domain, ops)
Martin Reinecke's avatar
Martin Reinecke committed
394
            return Linearization(field, bdop, want_metric=want_metric)