linearization.py 15 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 .field import Field
Martin Reinecke's avatar
Martin Reinecke committed
21
from .multi_field import MultiField
Martin Reinecke's avatar
Martin Reinecke committed
22 23 24 25
from .sugar import makeOp


class Linearization(object):
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)

Martin Reinecke's avatar
Martin Reinecke committed
65 66
    @property
    def domain(self):
Philipp Arras's avatar
Philipp Arras committed
67
        """DomainTuple or MultiDomain : the Jacobian's domain"""
Martin Reinecke's avatar
Martin Reinecke committed
68 69 70 71
        return self._jac.domain

    @property
    def target(self):
Philipp Arras's avatar
Philipp Arras committed
72
        """DomainTuple or MultiDomain : the Jacobian's target (i.e. the value's domain)"""
Martin Reinecke's avatar
Martin Reinecke committed
73 74 75 76
        return self._jac.target

    @property
    def val(self):
Philipp Arras's avatar
Philipp Arras committed
77
        """Field or MultiField : the value"""
Martin Reinecke's avatar
Martin Reinecke committed
78 79 80 81
        return self._val

    @property
    def jac(self):
82
        """LinearOperator : the Jacobian"""
Martin Reinecke's avatar
Martin Reinecke committed
83 84
        return self._jac

Martin Reinecke's avatar
Martin Reinecke committed
85 86
    @property
    def gradient(self):
Philipp Arras's avatar
Philipp Arras committed
87
        """Field or MultiField : the gradient
88 89 90 91 92

        Notes
        -----
        Only available if target is a scalar
        """
Martin Reinecke's avatar
Martin Reinecke committed
93
        return self._jac.adjoint_times(Field.scalar(1.))
Martin Reinecke's avatar
Martin Reinecke committed
94

95 96
    @property
    def want_metric(self):
Martin Reinecke's avatar
Martin Reinecke committed
97
        """bool : True iff the metric was requested in the constructor"""
98 99
        return self._want_metric

Martin Reinecke's avatar
Martin Reinecke committed
100 101
    @property
    def metric(self):
102 103 104 105 106 107
        """LinearOperator : the metric

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

Martin Reinecke's avatar
Martin Reinecke committed
110
    def __getitem__(self, name):
Martin Reinecke's avatar
Martin Reinecke committed
111
        from .operators.simple_linear_operators import ducktape
Martin Reinecke's avatar
Martin Reinecke committed
112
        return self.new(self._val[name], self._jac.ducktape_left(name))
Martin Reinecke's avatar
Martin Reinecke committed
113

Martin Reinecke's avatar
Martin Reinecke committed
114
    def __neg__(self):
115 116
        return self.new(-self._val, -self._jac,
                        None if self._metric is None else -self._metric)
Martin Reinecke's avatar
Martin Reinecke committed
117

Martin Reinecke's avatar
Martin Reinecke committed
118
    def conjugate(self):
119
        return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
120 121 122 123 124
            self._val.conjugate(), self._jac.conjugate(),
            None if self._metric is None else self._metric.conjugate())

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

Martin Reinecke's avatar
Martin Reinecke committed
127
    def _myadd(self, other, neg):
Martin Reinecke's avatar
Martin Reinecke committed
128 129 130
        if isinstance(other, Linearization):
            met = None
            if self._metric is not None and other._metric is not None:
Martin Reinecke's avatar
Martin Reinecke committed
131
                met = self._metric._myadd(other._metric, neg)
132
            return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
133 134
                self._val.flexible_addsub(other._val, neg),
                self._jac._myadd(other._jac, neg), met)
Martin Reinecke's avatar
Martin Reinecke committed
135
        if isinstance(other, (int, float, complex, Field, MultiField)):
Martin Reinecke's avatar
Martin Reinecke committed
136
            if neg:
137
                return self.new(self._val-other, self._jac, self._metric)
Martin Reinecke's avatar
Martin Reinecke committed
138
            else:
139
                return self.new(self._val+other, self._jac, self._metric)
Martin Reinecke's avatar
Martin Reinecke committed
140 141 142

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

    def __radd__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
145
        return self._myadd(other, False)
Martin Reinecke's avatar
Martin Reinecke committed
146 147

    def __sub__(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
148
        return self._myadd(other, True)
Martin Reinecke's avatar
Martin Reinecke committed
149 150 151 152

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

153 154
    def __truediv__(self, other):
        if isinstance(other, Linearization):
Philipp Frank's avatar
Philipp Frank committed
155
            return self.__mul__(other.one_over())
156 157 158
        return self.__mul__(1./other)

    def __rtruediv__(self, other):
Philipp Frank's avatar
Philipp Frank committed
159
        return self.one_over().__mul__(other)
160

Martin Reinecke's avatar
Martin Reinecke committed
161 162 163
    def __pow__(self, power):
        if not np.isscalar(power):
            return NotImplemented
Martin Reinecke's avatar
Martin Reinecke committed
164 165
        return self.new(self._val**power,
                        makeOp(self._val**(power-1)).scale(power)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
166

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

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

189
    def outer(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
190 191 192 193 194 195 196 197 198 199 200 201
        """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
        """
202 203 204
        from .operators.outer_product_operator import OuterProduct
        if isinstance(other, Linearization):
            return self.new(
Sebastian Hutschenreuter's avatar
Sebastian Hutschenreuter committed
205 206 207
                OuterProduct(self._val, other.target)(other._val),
                OuterProduct(self._jac(self._val), other.target)._myadd(
                    OuterProduct(self._val, other.target)(other._jac), False))
208
        if np.isscalar(other):
Martin Reinecke's avatar
Martin Reinecke committed
209
            return self.__mul__(other)
210
        if isinstance(other, (Field, MultiField)):
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))
213

Martin Reinecke's avatar
Martin Reinecke committed
214
    def vdot(self, other):
Martin Reinecke's avatar
Martin Reinecke committed
215 216 217 218 219 220 221 222 223 224 225 226
        """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
227
        from .operators.simple_linear_operators import VdotOperator
Martin Reinecke's avatar
Martin Reinecke committed
228
        if isinstance(other, (Field, MultiField)):
229
            return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
230
                Field.scalar(self._val.vdot(other)),
Martin Reinecke's avatar
Martin Reinecke committed
231
                VdotOperator(other)(self._jac))
232
        return self.new(
Martin Reinecke's avatar
Martin Reinecke committed
233
            Field.scalar(self._val.vdot(other._val)),
Martin Reinecke's avatar
Martin Reinecke committed
234 235
            VdotOperator(self._val)(other._jac) +
            VdotOperator(other._val)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
236

237
    def sum(self, spaces=None):
Martin Reinecke's avatar
Martin Reinecke committed
238 239 240 241 242 243 244 245 246 247 248 249 250
        """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
        """
251
        from .operators.contraction_operator import ContractionOperator
252 253 254
        if spaces is None:
            return self.new(
                Field.scalar(self._val.sum()),
255
                ContractionOperator(self._jac.target, None)(self._jac))
256 257 258
        else:
            return self.new(
                self._val.sum(spaces),
259
                ContractionOperator(self._jac.target, spaces)(self._jac))
260 261

    def integrate(self, spaces=None):
Martin Reinecke's avatar
Martin Reinecke committed
262 263 264 265 266 267 268 269 270 271 272 273 274
        """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
        """
275
        from .operators.contraction_operator import ContractionOperator
276 277 278
        if spaces is None:
            return self.new(
                Field.scalar(self._val.integrate()),
279
                ContractionOperator(self._jac.target, None, 1)(self._jac))
280 281 282
        else:
            return self.new(
                self._val.integrate(spaces),
283
                ContractionOperator(self._jac.target, spaces, 1)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
284 285 286

    def exp(self):
        tmp = self._val.exp()
287
        return self.new(tmp, makeOp(tmp)(self._jac))
Philipp Arras's avatar
Philipp Arras committed
288

Martin Reinecke's avatar
Martin Reinecke committed
289 290
    def clip(self, min=None, max=None):
        tmp = self._val.clip(min, max)
Jakob Knollmueller's avatar
Jakob Knollmueller committed
291
        if (min is None) and (max is None):
Martin Reinecke's avatar
Martin Reinecke committed
292
            return self
Jakob Knollmueller's avatar
Jakob Knollmueller committed
293 294 295 296 297 298
        elif max is None:
            tmp2 = makeOp(1. - (tmp == min))
        elif min is None:
            tmp2 = makeOp(1. - (tmp == max))
        else:
            tmp2 = makeOp(1. - (tmp == min) - (tmp == max))
299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320
        return self.new(tmp, tmp2(self._jac))

    def sin(self):
        tmp = self._val.sin()
        tmp2 = self._val.cos()
        return self.new(tmp, makeOp(tmp2)(self._jac))

    def cos(self):
        tmp = self._val.cos()
        tmp2 = - self._val.sin()
        return self.new(tmp, makeOp(tmp2)(self._jac))

    def tan(self):
        tmp = self._val.tan()
        tmp2 = 1./(self._val.cos()**2)
        return self.new(tmp, makeOp(tmp2)(self._jac))

    def sinc(self):
        tmp = self._val.sinc()
        tmp2 = (self._val.cos()-tmp)/self._val
        return self.new(tmp, makeOp(tmp2)(self._jac))

Martin Reinecke's avatar
Martin Reinecke committed
321 322
    def log(self):
        tmp = self._val.log()
323
        return self.new(tmp, makeOp(1./self._val)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
324

325 326 327 328 329 330 331 332 333 334
    def sinh(self):
        tmp = self._val.sinh()
        tmp2 = self._val.cosh()
        return self.new(tmp, makeOp(tmp2)(self._jac))

    def cosh(self):
        tmp = self._val.cosh()
        tmp2 = self._val.sinh()
        return self.new(tmp, makeOp(tmp2)(self._jac))

Martin Reinecke's avatar
Martin Reinecke committed
335 336
    def tanh(self):
        tmp = self._val.tanh()
337
        return self.new(tmp, makeOp(1.-tmp**2)(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
338

339
    def sigmoid(self):
Martin Reinecke's avatar
Martin Reinecke committed
340 341
        tmp = self._val.tanh()
        tmp2 = 0.5*(1.+tmp)
342
        return self.new(tmp2, makeOp(0.5*(1.-tmp**2))(self._jac))
Martin Reinecke's avatar
Martin Reinecke committed
343

344 345 346 347 348 349 350 351 352 353
    def absolute(self):
        tmp = self._val.absolute()
        tmp2 = self._val.sign()
        return self.new(tmp, makeOp(tmp2)(self._jac))

    def one_over(self):
        tmp = 1./self._val
        tmp2 = - tmp/self._val
        return self.new(tmp, makeOp(tmp2)(self._jac))

Martin Reinecke's avatar
Martin Reinecke committed
354
    def add_metric(self, metric):
355
        return self.new(self._val, self._jac, metric)
Martin Reinecke's avatar
Martin Reinecke committed
356

Martin Reinecke's avatar
Martin Reinecke committed
357 358 359
    def with_want_metric(self):
        return Linearization(self._val, self._jac, self._metric, True)

Martin Reinecke's avatar
Martin Reinecke committed
360
    @staticmethod
361
    def make_var(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
362 363 364 365 366 367 368 369 370 371 372 373 374 375 376
        """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
377
        from .operators.scaling_operator import ScalingOperator
378 379
        return Linearization(field, ScalingOperator(1., field.domain),
                             want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
380 381

    @staticmethod
382
    def make_const(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401
        """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
402
        from .operators.simple_linear_operators import NullOperator
403 404
        return Linearization(field, NullOperator(field.domain, field.domain),
                             want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
405

Martin Reinecke's avatar
Martin Reinecke committed
406 407
    @staticmethod
    def make_const_empty_input(field, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427
        """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
428 429
        from .operators.simple_linear_operators import NullOperator
        from .multi_domain import MultiDomain
Martin Reinecke's avatar
Martin Reinecke committed
430 431 432
        return Linearization(
            field, NullOperator(MultiDomain.make({}), field.domain),
            want_metric=want_metric)
Martin Reinecke's avatar
Martin Reinecke committed
433

Martin Reinecke's avatar
Martin Reinecke committed
434 435
    @staticmethod
    def make_partial_var(field, constants, want_metric=False):
Martin Reinecke's avatar
Martin Reinecke committed
436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458
        """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
459
        from .operators.scaling_operator import ScalingOperator
Philipp Arras's avatar
Typo  
Philipp Arras committed
460
        from .operators.block_diagonal_operator import BlockDiagonalOperator
Martin Reinecke's avatar
Martin Reinecke committed
461 462 463
        if len(constants) == 0:
            return Linearization.make_var(field, want_metric)
        else:
464 465 466
            ops = {key: ScalingOperator(0. if key in constants else 1., dom)
                   for key, dom in field.domain.items()}
            bdop = BlockDiagonalOperator(field.domain, ops)
Martin Reinecke's avatar
Martin Reinecke committed
467
            return Linearization(field, bdop, want_metric=want_metric)