harmonic_operators.py 17.3 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke 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/>.
#
14
# Copyright(C) 2013-2019 Max-Planck-Society
Martin Reinecke's avatar
Martin Reinecke committed
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

20
from .. import dobj, utilities, fft
Martin Reinecke's avatar
Martin Reinecke committed
21 22
from ..domain_tuple import DomainTuple
from ..domains.gl_space import GLSpace
Philipp Arras's avatar
Philipp Arras committed
23 24
from ..domains.lm_space import LMSpace
from ..domains.rg_space import RGSpace
Martin Reinecke's avatar
Martin Reinecke committed
25 26
from ..field import Field
from .diagonal_operator import DiagonalOperator
Philipp Arras's avatar
Philipp Arras committed
27
from .linear_operator import LinearOperator
Martin Reinecke's avatar
Martin Reinecke committed
28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57
from .scaling_operator import ScalingOperator


class FFTOperator(LinearOperator):
    """Transforms between a pair of position and harmonic RGSpaces.

    Parameters
    ----------
    domain: Domain, tuple of Domain or DomainTuple
        The domain of the data that is input by "times" and output by
        "adjoint_times".
    target: Domain, optional
        The target (sub-)domain of the transform operation.
        If omitted, a domain will be chosen automatically.
    space: int, optional
        The index of the subdomain on which the operator should act
        If None, it is set to 0 if `domain` contains exactly one space.
        `domain[space]` must be an RGSpace.

    Notes
    -----
    This operator performs full FFTs, which implies that its output field will
    always have complex type, regardless of the type of the input field.
    If a real field is desired after a forward/backward transform couple, it
    must be manually cast to real.
    """

    def __init__(self, domain, target=None, space=None):
        # Initialize domain and target
        self._domain = DomainTuple.make(domain)
Martin Reinecke's avatar
Martin Reinecke committed
58
        self._capability = self._all_ops
Martin Reinecke's avatar
Martin Reinecke committed
59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
        self._space = utilities.infer_space(self._domain, space)

        adom = self._domain[self._space]
        if not isinstance(adom, RGSpace):
            raise TypeError("FFTOperator only works on RGSpaces")
        if target is None:
            target = adom.get_default_codomain()

        self._target = [dom for dom in self._domain]
        self._target[self._space] = target
        self._target = DomainTuple.make(self._target)
        adom.check_codomain(target)
        target.check_codomain(adom)

    def apply(self, x, mode):
        self._check_input(x, mode)
        ncells = x.domain[self._space].size
        if x.domain[self._space].harmonic:  # harmonic -> position
Martin Reinecke's avatar
Martin Reinecke committed
77
            func = fft.fftn
Martin Reinecke's avatar
Martin Reinecke committed
78 79
            fct = 1.
        else:
Martin Reinecke's avatar
Martin Reinecke committed
80
            func = fft.ifftn
Martin Reinecke's avatar
Martin Reinecke committed
81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145
            fct = ncells
        axes = x.domain.axes[self._space]
        tdom = self._tgt(mode)
        oldax = dobj.distaxis(x.val)
        if oldax not in axes:  # straightforward, no redistribution needed
            ldat = x.local_data
            ldat = func(ldat, axes=axes)
            tmp = dobj.from_local_data(x.val.shape, ldat, distaxis=oldax)
        elif len(axes) < len(x.shape) or len(axes) == 1:
            # we can use one FFT pass in between the redistributions
            tmp = dobj.redistribute(x.val, nodist=axes)
            newax = dobj.distaxis(tmp)
            ldat = dobj.local_data(tmp)
            ldat = func(ldat, axes=axes)
            tmp = dobj.from_local_data(tmp.shape, ldat, distaxis=newax)
            tmp = dobj.redistribute(tmp, dist=oldax)
        else:  # two separate FFTs needed
            rem_axes = tuple(i for i in axes if i != oldax)
            tmp = x.val
            ldat = dobj.local_data(tmp)
            ldat = func(ldat, axes=rem_axes)
            if oldax != 0:
                raise ValueError("bad distribution")
            ldat2 = ldat.reshape((ldat.shape[0],
                                  np.prod(ldat.shape[1:])))
            shp2d = (x.val.shape[0], np.prod(x.val.shape[1:]))
            tmp = dobj.from_local_data(shp2d, ldat2, distaxis=0)
            tmp = dobj.transpose(tmp)
            ldat2 = dobj.local_data(tmp)
            ldat2 = func(ldat2, axes=(1,))
            tmp = dobj.from_local_data(tmp.shape, ldat2, distaxis=0)
            tmp = dobj.transpose(tmp)
            ldat2 = dobj.local_data(tmp).reshape(ldat.shape)
            tmp = dobj.from_local_data(x.val.shape, ldat2, distaxis=0)
        Tval = Field(tdom, tmp)
        if mode & (LinearOperator.TIMES | LinearOperator.ADJOINT_TIMES):
            fct *= self._domain[self._space].scalar_dvol
        else:
            fct *= self._target[self._space].scalar_dvol
        return Tval if fct == 1 else Tval*fct


class HartleyOperator(LinearOperator):
    """Transforms between a pair of position and harmonic RGSpaces.

    Parameters
    ----------
    domain: Domain, tuple of Domain or DomainTuple
        The domain of the data that is input by "times" and output by
        "adjoint_times".
    target: Domain, optional
        The target (sub-)domain of the transform operation.
        If omitted, a domain will be chosen automatically.
    space: int, optional
        The index of the subdomain on which the operator should act
        If None, it is set to 0 if `domain` contains exactly one space.
        `domain[space]` must be an RGSpace.

    Notes
    -----
    This operator always produces output fields with the same data type as
    its input. This is achieved by performing so-called Hartley transforms
    (https://en.wikipedia.org/wiki/Discrete_Hartley_transform).
    For complex input fields, the operator will transform the real and
    imaginary parts separately and use the results as real and imaginary parts
146
    of the result field, respectively.
Martin Reinecke's avatar
Martin Reinecke committed
147 148 149 150 151 152 153 154
    In many contexts the Hartley transform is a perfect substitute for the
    Fourier transform, but in some situations (e.g. convolution with a general,
    non-symmetric kernel, the full FFT must be used instead.
    """

    def __init__(self, domain, target=None, space=None):
        # Initialize domain and target
        self._domain = DomainTuple.make(domain)
Martin Reinecke's avatar
Martin Reinecke committed
155
        self._capability = self._all_ops
Martin Reinecke's avatar
Martin Reinecke committed
156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183
        self._space = utilities.infer_space(self._domain, space)

        adom = self._domain[self._space]
        if not isinstance(adom, RGSpace):
            raise TypeError("HartleyOperator only works on RGSpaces")
        if target is None:
            target = adom.get_default_codomain()

        self._target = [dom for dom in self._domain]
        self._target[self._space] = target
        self._target = DomainTuple.make(self._target)
        adom.check_codomain(target)
        target.check_codomain(adom)

    def apply(self, x, mode):
        self._check_input(x, mode)
        if utilities.iscomplextype(x.dtype):
            return (self._apply_cartesian(x.real, mode) +
                    1j*self._apply_cartesian(x.imag, mode))
        else:
            return self._apply_cartesian(x, mode)

    def _apply_cartesian(self, x, mode):
        axes = x.domain.axes[self._space]
        tdom = self._tgt(mode)
        oldax = dobj.distaxis(x.val)
        if oldax not in axes:  # straightforward, no redistribution needed
            ldat = x.local_data
184
            ldat = fft.hartley(ldat, axes=axes)
Martin Reinecke's avatar
Martin Reinecke committed
185 186 187 188 189 190
            tmp = dobj.from_local_data(x.val.shape, ldat, distaxis=oldax)
        elif len(axes) < len(x.shape) or len(axes) == 1:
            # we can use one Hartley pass in between the redistributions
            tmp = dobj.redistribute(x.val, nodist=axes)
            newax = dobj.distaxis(tmp)
            ldat = dobj.local_data(tmp)
191
            ldat = fft.hartley(ldat, axes=axes)
Martin Reinecke's avatar
Martin Reinecke committed
192 193 194 195 196 197 198 199 200 201 202 203 204
            tmp = dobj.from_local_data(tmp.shape, ldat, distaxis=newax)
            tmp = dobj.redistribute(tmp, dist=oldax)
        else:  # two separate, full FFTs needed
            # ideal strategy for the moment would be:
            # - do real-to-complex FFT on all local axes
            # - fill up array
            # - redistribute array
            # - do complex-to-complex FFT on remaining axis
            # - add re+im
            # - redistribute back
            rem_axes = tuple(i for i in axes if i != oldax)
            tmp = x.val
            ldat = dobj.local_data(tmp)
205
            ldat = fft.my_fftn_r2c(ldat, axes=rem_axes)
Martin Reinecke's avatar
Martin Reinecke committed
206 207 208 209 210 211 212 213
            if oldax != 0:
                raise ValueError("bad distribution")
            ldat2 = ldat.reshape((ldat.shape[0],
                                  np.prod(ldat.shape[1:])))
            shp2d = (x.val.shape[0], np.prod(x.val.shape[1:]))
            tmp = dobj.from_local_data(shp2d, ldat2, distaxis=0)
            tmp = dobj.transpose(tmp)
            ldat2 = dobj.local_data(tmp)
214
            ldat2 = fft.my_fftn(ldat2, axes=(1,))
Martin Reinecke's avatar
Martin Reinecke committed
215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257
            ldat2 = ldat2.real+ldat2.imag
            tmp = dobj.from_local_data(tmp.shape, ldat2, distaxis=0)
            tmp = dobj.transpose(tmp)
            ldat2 = dobj.local_data(tmp).reshape(ldat.shape)
            tmp = dobj.from_local_data(x.val.shape, ldat2, distaxis=0)
        Tval = Field(tdom, tmp)
        if mode & (LinearOperator.TIMES | LinearOperator.ADJOINT_TIMES):
            fct = self._domain[self._space].scalar_dvol
        else:
            fct = self._target[self._space].scalar_dvol
        return Tval if fct == 1 else Tval*fct


class SHTOperator(LinearOperator):
    """Transforms between a harmonic domain on the sphere and a position
    domain counterpart.

    Built-in domain pairs are
      - an LMSpace and a HPSpace
      - an LMSpace and a GLSpace

    The supported operations are times() and adjoint_times().

    Parameters
    ----------
    domain : Domain, tuple of Domain or DomainTuple
        The domain of the data that is input by "times" and output by
        "adjoint_times".
    target : Domain, optional
        The target domain of the transform operation.
        If omitted, a domain will be chosen automatically.
        Whenever the input domain of the transform is an RGSpace, the codomain
        (and its parameters) are uniquely determined.
        For LMSpace, a GLSpace of sufficient resolution is chosen.
    space : int, optional
        The index of the domain on which the operator should act
        If None, it is set to 0 if domain contains exactly one subdomain.
        domain[space] must be a LMSpace.
    """

    def __init__(self, domain, target=None, space=None):
        # Initialize domain and target
        self._domain = DomainTuple.make(domain)
Martin Reinecke's avatar
Martin Reinecke committed
258
        self._capability = self.TIMES | self.ADJOINT_TIMES
Martin Reinecke's avatar
Martin Reinecke committed
259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282
        self._space = utilities.infer_space(self._domain, space)

        hspc = self._domain[self._space]
        if not isinstance(hspc, LMSpace):
            raise TypeError("SHTOperator only works on a LMSpace domain")
        if target is None:
            target = hspc.get_default_codomain()

        self._target = [dom for dom in self._domain]
        self._target[self._space] = target
        self._target = DomainTuple.make(self._target)
        hspc.check_codomain(target)
        target.check_codomain(hspc)

        from pyHealpix import sharpjob_d
        self.lmax = hspc.lmax
        self.mmax = hspc.mmax
        self.sjob = sharpjob_d()
        self.sjob.set_triangular_alm_info(self.lmax, self.mmax)
        if isinstance(target, GLSpace):
            self.sjob.set_Gauss_geometry(target.nlat, target.nlon)
        else:
            self.sjob.set_Healpix_geometry(target.nside)

Martin Reinecke's avatar
Martin Reinecke committed
283 284
    def __reduce__(self):
        return (_unpickleSHTOperator,
Martin Reinecke's avatar
Martin Reinecke committed
285
                (self._domain, self._target[self._space], self._space))
Martin Reinecke's avatar
Martin Reinecke committed
286

Martin Reinecke's avatar
Martin Reinecke committed
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319
    def apply(self, x, mode):
        self._check_input(x, mode)
        if utilities.iscomplextype(x.dtype):
            return (self._apply_spherical(x.real, mode) +
                    1j*self._apply_spherical(x.imag, mode))
        else:
            return self._apply_spherical(x, mode)

    def _slice_p2h(self, inp):
        rr = self.sjob.alm2map_adjoint(inp)
        if len(rr) != ((self.mmax+1)*(self.mmax+2))//2 + \
                      (self.mmax+1)*(self.lmax-self.mmax):
            raise ValueError("array length mismatch")
        res = np.empty(2*len(rr)-self.lmax-1, dtype=rr[0].real.dtype)
        res[0:self.lmax+1] = rr[0:self.lmax+1].real
        res[self.lmax+1::2] = np.sqrt(2)*rr[self.lmax+1:].real
        res[self.lmax+2::2] = np.sqrt(2)*rr[self.lmax+1:].imag
        return res/np.sqrt(np.pi*4)

    def _slice_h2p(self, inp):
        res = np.empty((len(inp)+self.lmax+1)//2, dtype=(inp[0]*1j).dtype)
        if len(res) != ((self.mmax+1)*(self.mmax+2))//2 + \
                       (self.mmax+1)*(self.lmax-self.mmax):
            raise ValueError("array length mismatch")
        res[0:self.lmax+1] = inp[0:self.lmax+1]
        res[self.lmax+1:] = np.sqrt(0.5)*(inp[self.lmax+1::2] +
                                          1j*inp[self.lmax+2::2])
        res = self.sjob.alm2map(res)
        return res/np.sqrt(np.pi*4)

    def _apply_spherical(self, x, mode):
        axes = x.domain.axes[self._space]
        axis = axes[0]
320 321 322
        v = x.val
        v, idat = dobj.ensure_not_distributed(v, (axis,))
        distaxis = dobj.distaxis(v)
Martin Reinecke's avatar
Martin Reinecke committed
323 324 325 326 327 328 329 330 331

        p2h = not x.domain[self._space].harmonic
        tdom = self._tgt(mode)
        func = self._slice_p2h if p2h else self._slice_h2p
        odat = np.empty(dobj.local_shape(tdom.shape, distaxis=distaxis),
                        dtype=x.dtype)
        for slice in utilities.get_slice_list(idat.shape, axes):
            odat[slice] = func(idat[slice])
        odat = dobj.from_local_data(tdom.shape, odat, distaxis)
332
        return Field(tdom, dobj.ensure_default_distributed(odat))
Martin Reinecke's avatar
Martin Reinecke committed
333 334


Martin Reinecke's avatar
Martin Reinecke committed
335 336 337 338
def _unpickleSHTOperator(*args):
    return SHTOperator(*args)


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 359 360 361 362 363
class HarmonicTransformOperator(LinearOperator):
    """Transforms between a harmonic domain and a position domain counterpart.

    Built-in domain pairs are
      - a harmonic and a non-harmonic RGSpace (with matching distances)
      - an LMSpace and a HPSpace
      - an LMSpace and a GLSpace

    The supported operations are times() and adjoint_times().

    Parameters
    ----------
    domain : Domain, tuple of Domain or DomainTuple
        The domain of the data that is input by "times" and output by
        "adjoint_times".
    target : Domain, optional
        The target domain of the transform operation.
        If omitted, a domain will be chosen automatically.
        Whenever the input domain of the transform is an RGSpace, the codomain
        (and its parameters) are uniquely determined.
        For LMSpace, a GLSpace of sufficient resolution is chosen.
    space : int, optional
        The index of the domain on which the operator should act
        If None, it is set to 0 if domain contains exactly one subdomain.
        domain[space] must be a harmonic domain.
364 365 366

    Notes
    -----
Martin Reinecke's avatar
Martin Reinecke committed
367 368 369 370
    HarmonicTransformOperator uses a Hartley transformation to transform
    between harmonic and non-harmonic RGSpaces. This has the advantage that all
    field values are real in either space. If you require a true Fourier
    transform you should use FFTOperator instead.
Martin Reinecke's avatar
Martin Reinecke committed
371 372 373 374 375 376 377 378 379 380 381 382 383 384
    """

    def __init__(self, domain, target=None, space=None):
        domain = DomainTuple.make(domain)
        space = utilities.infer_space(domain, space)

        hspc = domain[space]
        if not hspc.harmonic:
            raise TypeError(
                "HarmonicTransformOperator only works on a harmonic space")
        if isinstance(hspc, RGSpace):
            self._op = HartleyOperator(domain, target, space)
        else:
            self._op = SHTOperator(domain, target, space)
Martin Reinecke's avatar
Martin Reinecke committed
385 386
        self._domain = self._op.domain
        self._target = self._op.target
Martin Reinecke's avatar
Martin Reinecke committed
387
        self._capability = self.TIMES | self.ADJOINT_TIMES
Martin Reinecke's avatar
Martin Reinecke committed
388 389 390 391 392 393 394

    def apply(self, x, mode):
        self._check_input(x, mode)
        return self._op.apply(x, mode)


def HarmonicSmoothingOperator(domain, sigma, space=None):
395
    """Returns an operator that carries out a smoothing with a Gaussian kernel
396
    of width `sigma` on the part of `domain` given by `space`.
Martin Reinecke's avatar
Martin Reinecke committed
397 398 399 400 401 402 403 404 405 406 407 408 409 410 411

    Parameters
    ----------
    domain : Domain, tuple of Domain, or DomainTuple
       The total domain of the operator's input and output fields
    sigma : float>=0
       The sigma of the Gaussian used for smoothing. It has the same units as
       the RGSpace the operator is working on.
       If `sigma==0`, an identity operator will be returned.
    space : int, optional
       The index of the sub-domain on which the smoothing is performed.
       Can be omitted if `domain` only has one sub-domain.

    Notes
    -----
412 413
    The sub-domain on which the smoothing is carried out *must* be a
    non-harmonic `RGSpace`.
Martin Reinecke's avatar
Martin Reinecke committed
414 415 416 417
    """

    sigma = float(sigma)
    if sigma < 0.:
418
        raise ValueError("sigma must be non-negative")
Martin Reinecke's avatar
Martin Reinecke committed
419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434
    if sigma == 0.:
        return ScalingOperator(1., domain)

    domain = DomainTuple.make(domain)
    space = utilities.infer_space(domain, space)
    if domain[space].harmonic:
        raise TypeError("domain must not be harmonic")
    Hartley = HartleyOperator(domain, space=space)
    codomain = Hartley.domain[space].get_default_codomain()
    kernel = codomain.get_k_length_array()
    smoother = codomain.get_fft_smoothing_kernel_function(sigma)
    kernel = smoother(kernel)
    ddom = list(domain)
    ddom[space] = codomain
    diag = DiagonalOperator(kernel, ddom, space)
    return Hartley.inverse(diag(Hartley))