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

import numpy as np

Martin Reinecke's avatar
misc    
Martin Reinecke committed
20
from .operators.operator import Operator
Philipp Arras's avatar
Philipp Arras committed
21
from .sugar import makeOp
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
    def prepend_jac(self, jac):
Lukas Platz's avatar
Lukas Platz committed
68
69
70
71
        if self._metric is None:
            return self.new(self._val, self._jac @ jac)
        from .operators.sandwich_operator import SandwichOperator
        metric = SandwichOperator.make(jac, self._metric)
Philipp Arras's avatar
Philipp Arras committed
72
73
        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

Philipp Arras's avatar
Philipp Arras committed
136
137
138
139
    @property
    def imag(self):
        return self.new(self._val.imag, self._jac.imag)

Martin Reinecke's avatar
Martin Reinecke committed
140
    def _myadd(self, other, neg):
Martin Reinecke's avatar
Martin Reinecke committed
141
142
143
144
145
146
147
148
149
        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
150
151
152

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

    def __radd__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
155
        return self._myadd(other, False)
Martin Reinecke's avatar
Martin Reinecke committed
156
157

    def __sub__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
158
        return self._myadd(other, True)
Martin Reinecke's avatar
Martin Reinecke committed
159
160
161
162

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

163
    def __truediv__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
164
165
166
        if np.isscalar(other):
            return self.__mul__(1/other)
        return self.__mul__(other.ptw("reciprocal"))
167
168

    def __rtruediv__(self, other):
169
        return self.ptw("reciprocal").__mul__(other)
170

Martin Reinecke's avatar
Martin Reinecke committed
171
    def __pow__(self, power):
Martin Reinecke's avatar
Martin Reinecke committed
172
        if not (np.isscalar(power) or power.jac is None):
Martin Reinecke's avatar
Martin Reinecke committed
173
            return NotImplemented
Martin Reinecke's avatar
misc    
Martin Reinecke committed
174
        return self.ptw("power", power)
Martin Reinecke's avatar
Martin Reinecke committed
175

Martin Reinecke's avatar
Martin Reinecke committed
176
    def __mul__(self, other):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
177
178
179
180
        if np.isscalar(other):
            if other == 1:
                return self
            met = None if self._metric is None else self._metric.scale(other)
181
            return self.new(self._val*other, self._jac.scale(other), met)
Martin Reinecke's avatar
Martin Reinecke committed
182
183
        from .sugar import makeOp
        if other.jac is None:
Martin Reinecke's avatar
Martin Reinecke committed
184
185
            if self.target != other.domain:
                raise ValueError("domain mismatch")
186
            return self.new(self._val*other, makeOp(other)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
187
188
189
190
191
192
        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
193
194

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

197
    def outer(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
198
199
200
201
202
203
204
205
206
207
208
209
        """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
        """
210
        if np.isscalar(other):
Martin Reinecke's avatar
Martin Reinecke committed
211
            return self.__mul__(other)
Martin Reinecke's avatar
Martin Reinecke committed
212
213
        from .operators.outer_product_operator import OuterProduct
        if other.jac is None:
Martin Reinecke's avatar
Martin Reinecke committed
214
215
216
            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
217
        return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
218
219
220
            tmp_op(other._val),
            OuterProduct(other.target, self._jac(self._val))._myadd(
                tmp_op(other._jac), False))
221

Martin Reinecke's avatar
Martin Reinecke committed
222
    def vdot(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
223
224
225
226
227
228
229
230
231
232
233
234
        """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
235
        from .operators.simple_linear_operators import VdotOperator
Martin Reinecke's avatar
Martin Reinecke committed
236
        if other.jac is None:
237
            return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
238
                self._val.vdot(other),
Martin Reinecke's avatar
Martin Reinecke committed
239
                VdotOperator(other)(self._jac))
240
        return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
241
            self._val.vdot(other._val),
Martin Reinecke's avatar
Martin Reinecke committed
242
243
            VdotOperator(self._val)(other._jac) +
            VdotOperator(other._val)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
244

245
    def sum(self, spaces=None):
Martin Reinecke's avatar
Martin Reinecke committed
246
247
248
249
250
251
252
253
254
255
256
257
258
        """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
        """
259
        from .operators.contraction_operator import ContractionOperator
Martin Reinecke's avatar
Martin Reinecke committed
260
261
262
        return self.new(
            self._val.sum(spaces),
            ContractionOperator(self._jac.target, spaces)(self._jac))
263
264

    def integrate(self, spaces=None):
Martin Reinecke's avatar
Martin Reinecke committed
265
266
267
268
269
270
271
272
273
274
275
276
277
        """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
278
279
        from .operators.contraction_operator import IntegrationOperator
        return IntegrationOperator(self._target, spaces)(self)
Martin Reinecke's avatar
Martin Reinecke committed
280

Martin Reinecke's avatar
misc    
Martin Reinecke committed
281
282
    def ptw(self, op, *args, **kwargs):
        t1, t2 = self._val.ptw_with_deriv(op, *args, **kwargs)
Martin Reinecke's avatar
Martin Reinecke committed
283
        return self.new(t1, makeOp(t2)(self._jac))
Philipp Arras's avatar
Philipp Arras committed
284

Martin Reinecke's avatar
Martin Reinecke committed
285
    def add_metric(self, metric):
286
        return self.new(self._val, self._jac, metric)
Martin Reinecke's avatar
Martin Reinecke committed
287

Martin Reinecke's avatar
Martin Reinecke committed
288
289
290
    def with_want_metric(self):
        return Linearization(self._val, self._jac, self._metric, True)

Martin Reinecke's avatar
Martin Reinecke committed
291
    @staticmethod
292
    def make_var(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
        """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
308
        from .operators.scaling_operator import ScalingOperator
309
        return Linearization(field, ScalingOperator(field.domain, 1.),
310
                             want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
311
312

    @staticmethod
313
    def make_const(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
        """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
333
        from .operators.simple_linear_operators import NullOperator
334
335
        return Linearization(field, NullOperator(field.domain, field.domain),
                             want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
336

Martin Reinecke's avatar
Martin Reinecke committed
337
338
    @staticmethod
    def make_const_empty_input(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
        """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
359
        from .multi_domain import MultiDomain
Philipp Arras's avatar
Philipp Arras committed
360
        from .operators.simple_linear_operators import NullOperator
Martin Reinecke's avatar
Martin Reinecke committed
361
362
363
        return Linearization(
            field, NullOperator(MultiDomain.make({}), field.domain),
            want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
364

Martin Reinecke's avatar
Martin Reinecke committed
365
366
    @staticmethod
    def make_partial_var(field, constants, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
        """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.
        """
Philipp Arras's avatar
Typo    
Philipp Arras committed
390
        from .operators.block_diagonal_operator import BlockDiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
391
        from .operators.scaling_operator import ScalingOperator
Martin Reinecke's avatar
Martin Reinecke committed
392
393
394
        if len(constants) == 0:
            return Linearization.make_var(field, want_metric)
        else:
395
            ops = {key: ScalingOperator(dom, 0. if key in constants else 1.)
396
397
                   for key, dom in field.domain.items()}
            bdop = BlockDiagonalOperator(field.domain, ops)
Martin Reinecke's avatar
Martin Reinecke committed
398
            return Linearization(field, bdop, want_metric=want_metric)