linearization.py 12.8 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 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/>.
#
# Copyright(C) 2013-2019 Max-Planck-Society
#
# 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
Martin Reinecke committed
21
from . import utilities
Martin Reinecke's avatar
misc    
Martin Reinecke committed
22
from .operators.operator import Operator
Martin Reinecke's avatar
Martin Reinecke committed
23
24


Martin Reinecke's avatar
misc    
Martin Reinecke committed
25
class Linearization(Operator):
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
26
    """Let `A` be an operator and `x` a field. `Linearization` stores the value
27
28
29
30
31
    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
32
33
    val : Field or MultiField
        The value of the operator application.
34
    jac : LinearOperator
Philipp Arras's avatar
Philipp Arras committed
35
36
37
38
39
40
        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.
41
    """
42
    def __init__(self, val, jac, metric=None, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
43
44
        self._val = val
        self._jac = jac
Martin Reinecke's avatar
Martin Reinecke committed
45
46
        if self._val.domain != self._jac.target:
            raise ValueError("domain mismatch")
47
        self._want_metric = want_metric
Martin Reinecke's avatar
Martin Reinecke committed
48
49
        self._metric = metric

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

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

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

Philipp Arras's avatar
Philipp Arras committed
68
69
70
71
72
73
74
    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
75
76
    @property
    def domain(self):
Philipp Arras's avatar
Philipp Arras committed
77
        """DomainTuple or MultiDomain : the Jacobian's domain"""
Martin Reinecke's avatar
Martin Reinecke committed
78
79
80
81
        return self._jac.domain

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

194
    def outer(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
195
196
197
198
199
200
201
202
203
204
205
206
        """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
        """
207
        if np.isscalar(other):
Martin Reinecke's avatar
Martin Reinecke committed
208
            return self.__mul__(other)
Martin Reinecke's avatar
Martin Reinecke committed
209
210
        from .operators.outer_product_operator import OuterProduct
        if other.jac is None:
Sebastian Hutschenreuter's avatar
Sebastian Hutschenreuter committed
211
212
            return self.new(OuterProduct(self._val, other.domain)(other),
                            OuterProduct(self._jac(self._val), other.domain))
Martin Reinecke's avatar
Martin Reinecke committed
213
214
215
216
        return self.new(
            OuterProduct(self._val, other.target)(other._val),
            OuterProduct(self._jac(self._val), other.target)._myadd(
                OuterProduct(self._val, other.target)(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
        """
274
        from .operators.contraction_operator import ContractionOperator
Martin Reinecke's avatar
Martin Reinecke committed
275
276
277
        return self.new(
            self._val.integrate(spaces),
            ContractionOperator(self._jac.target, spaces, 1)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
278

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

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

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

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

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

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

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