Commit dfe6fb20 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

Merge branch 'NIFTy_5' into operator_work_pa

parents fa052fee e8105d79
...@@ -32,8 +32,9 @@ from .operators.domain_distributor import DomainDistributor ...@@ -32,8 +32,9 @@ from .operators.domain_distributor import DomainDistributor
from .operators.endomorphic_operator import EndomorphicOperator from .operators.endomorphic_operator import EndomorphicOperator
from .operators.exp_transform import ExpTransform from .operators.exp_transform import ExpTransform
from .operators.fft_operator import FFTOperator from .operators.fft_operator import FFTOperator
from .operators.fft_smoothing_operator import FFTSmoothingOperator
from .operators.field_zero_padder import FieldZeroPadder from .operators.field_zero_padder import FieldZeroPadder
from .operators.hartley_operator import HartleyOperator
from .operators.harmonic_smoothing_operator import HarmonicSmoothingOperator
from .operators.geometry_remover import GeometryRemover from .operators.geometry_remover import GeometryRemover
from .operators.harmonic_transform_operator import HarmonicTransformOperator from .operators.harmonic_transform_operator import HarmonicTransformOperator
from .operators.inversion_enabler import InversionEnabler from .operators.inversion_enabler import InversionEnabler
......
...@@ -25,7 +25,7 @@ from ..models.local_nonlinearity import PointwiseExponential ...@@ -25,7 +25,7 @@ from ..models.local_nonlinearity import PointwiseExponential
from ..models.variable import Variable from ..models.variable import Variable
from ..multi.multi_field import MultiField from ..multi.multi_field import MultiField
from ..operators.domain_distributor import DomainDistributor from ..operators.domain_distributor import DomainDistributor
from ..operators.fft_operator import FFTOperator from ..operators.hartley_operator import HartleyOperator
from ..operators.harmonic_transform_operator import HarmonicTransformOperator from ..operators.harmonic_transform_operator import HarmonicTransformOperator
from ..operators.power_distributor import PowerDistributor from ..operators.power_distributor import PowerDistributor
...@@ -41,7 +41,7 @@ def make_correlated_field(s_space, amplitude_model): ...@@ -41,7 +41,7 @@ def make_correlated_field(s_space, amplitude_model):
amplitude_model : model for correlation structure amplitude_model : model for correlation structure
''' '''
h_space = s_space.get_default_codomain() h_space = s_space.get_default_codomain()
ht = FFTOperator(h_space, s_space) ht = HartleyOperator(h_space, s_space)
p_space = amplitude_model.value.domain[0] p_space = amplitude_model.value.domain[0]
power_distributor = PowerDistributor(h_space, p_space) power_distributor = PowerDistributor(h_space, p_space)
......
...@@ -25,26 +25,16 @@ from ..compat import * ...@@ -25,26 +25,16 @@ from ..compat import *
from ..domain_tuple import DomainTuple from ..domain_tuple import DomainTuple
from ..field import Field from ..field import Field
from .linear_operator import LinearOperator from .linear_operator import LinearOperator
from .. import utilities
# MR FIXME: this needs to be rewritten in a generic fashion
class DomainDistributor(LinearOperator): class DomainDistributor(LinearOperator):
def __init__(self, target, axis): def __init__(self, target, spaces):
if dobj.ntask > 1:
raise NotImplementedError('UpProj class does not support MPI.')
assert len(target) == 2
assert axis in [0, 1]
if axis == 0:
domain = target[1]
self._size = target[0].size
else:
domain = target[0]
self._size = target[1].size
self._axis = axis
self._domain = DomainTuple.make(domain)
self._target = DomainTuple.make(target) self._target = DomainTuple.make(target)
self._spaces = utilities.parse_spaces(spaces, len(self._target))
self._domain = [tgt for i, tgt in enumerate(self._target)
if i in self._spaces]
self._domain = DomainTuple.make(self._domain)
@property @property
def domain(self): def domain(self):
...@@ -57,23 +47,16 @@ class DomainDistributor(LinearOperator): ...@@ -57,23 +47,16 @@ class DomainDistributor(LinearOperator):
def apply(self, x, mode): def apply(self, x, mode):
self._check_input(x, mode) self._check_input(x, mode)
if mode == self.TIMES: if mode == self.TIMES:
x = x.local_data ldat = x.local_data if 0 in self._spaces else x.to_global_data()
otherDirection = np.ones(self._size) shp = []
if self._axis == 0: for i, tgt in enumerate(self._target):
res = np.outer(otherDirection, x) tmp = tgt.shape if i > 0 else tgt.local_shape
else: shp += tmp if i in self._spaces else(1,)*len(tgt.shape)
res = np.outer(x, otherDirection) ldat = np.broadcast_to(ldat.reshape(shp), self._target.local_shape)
res = res.reshape(dobj.local_shape(self.target.shape)) return Field.from_local_data(self._target, ldat)
return Field.from_local_data(self.target, res)
else: else:
if self._axis == 0: return x.sum([s for s in range(len(x.domain))
x = x.local_data.reshape(self._size, -1) if s not in self._spaces])
res = np.sum(x, axis=0)
else:
x = x.local_data.reshape(-1, self._size)
res = np.sum(x, axis=1)
res = res.reshape(dobj.local_shape(self.domain.shape))
return Field.from_local_data(self.domain, res)
@property @property
def capability(self): def capability(self):
......
...@@ -27,12 +27,13 @@ from ..domains.power_space import PowerSpace ...@@ -27,12 +27,13 @@ from ..domains.power_space import PowerSpace
from ..domains.rg_space import RGSpace from ..domains.rg_space import RGSpace
from ..field import Field from ..field import Field
from .linear_operator import LinearOperator from .linear_operator import LinearOperator
from .. import utilities
class ExpTransform(LinearOperator): class ExpTransform(LinearOperator):
def __init__(self, target, dof, space=0): def __init__(self, target, dof, space=0):
self._target = DomainTuple.make(target) self._target = DomainTuple.make(target)
self._space = int(space) self._space = utilities.infer_space(self._target, space)
tgt = self._target[self._space] tgt = self._target[self._space]
if not ((isinstance(tgt, RGSpace) and tgt.harmonic) or if not ((isinstance(tgt, RGSpace) and tgt.harmonic) or
isinstance(tgt, PowerSpace)): isinstance(tgt, PowerSpace)):
......
...@@ -43,6 +43,13 @@ class FFTOperator(LinearOperator): ...@@ -43,6 +43,13 @@ class FFTOperator(LinearOperator):
The index of the subdomain on which the operator should act The index of the subdomain on which the operator should act
If None, it is set to 0 if `domain` contains exactly one space. If None, it is set to 0 if `domain` contains exactly one space.
`domain[space]` must be an RGSpace. `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): def __init__(self, domain, target=None, space=None):
...@@ -67,41 +74,35 @@ class FFTOperator(LinearOperator): ...@@ -67,41 +74,35 @@ class FFTOperator(LinearOperator):
utilities.fft_prep() utilities.fft_prep()
def apply(self, x, mode): def apply(self, x, mode):
from pyfftw.interfaces.numpy_fft import fftn, ifftn
self._check_input(x, mode) self._check_input(x, mode)
if np.issubdtype(x.dtype, np.complexfloating): ncells = x.domain[self._space].size
return (self._apply_cartesian(x.real, mode) + if x.domain[self._space].harmonic: # harmonic -> position
1j*self._apply_cartesian(x.imag, mode)) func = fftn
fct = 1.
else: else:
return self._apply_cartesian(x, mode) func = ifftn
fct = ncells
def _apply_cartesian(self, x, mode):
axes = x.domain.axes[self._space] axes = x.domain.axes[self._space]
tdom = self._tgt(mode) tdom = self._tgt(mode)
oldax = dobj.distaxis(x.val) oldax = dobj.distaxis(x.val)
if oldax not in axes: # straightforward, no redistribution needed if oldax not in axes: # straightforward, no redistribution needed
ldat = x.local_data ldat = x.local_data
ldat = utilities.hartley(ldat, axes=axes) ldat = func(ldat, axes=axes)
tmp = dobj.from_local_data(x.val.shape, ldat, distaxis=oldax) tmp = dobj.from_local_data(x.val.shape, ldat, distaxis=oldax)
elif len(axes) < len(x.shape) or len(axes) == 1: elif len(axes) < len(x.shape) or len(axes) == 1:
# we can use one Hartley pass in between the redistributions # we can use one FFT pass in between the redistributions
tmp = dobj.redistribute(x.val, nodist=axes) tmp = dobj.redistribute(x.val, nodist=axes)
newax = dobj.distaxis(tmp) newax = dobj.distaxis(tmp)
ldat = dobj.local_data(tmp) ldat = dobj.local_data(tmp)
ldat = utilities.hartley(ldat, axes=axes) ldat = func(ldat, axes=axes)
tmp = dobj.from_local_data(tmp.shape, ldat, distaxis=newax) tmp = dobj.from_local_data(tmp.shape, ldat, distaxis=newax)
tmp = dobj.redistribute(tmp, dist=oldax) tmp = dobj.redistribute(tmp, dist=oldax)
else: # two separate, full FFTs needed else: # two separate 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) rem_axes = tuple(i for i in axes if i != oldax)
tmp = x.val tmp = x.val
ldat = dobj.local_data(tmp) ldat = dobj.local_data(tmp)
ldat = utilities.my_fftn_r2c(ldat, axes=rem_axes) ldat = func(ldat, axes=rem_axes)
if oldax != 0: if oldax != 0:
raise ValueError("bad distribution") raise ValueError("bad distribution")
ldat2 = ldat.reshape((ldat.shape[0], ldat2 = ldat.reshape((ldat.shape[0],
...@@ -110,17 +111,16 @@ class FFTOperator(LinearOperator): ...@@ -110,17 +111,16 @@ class FFTOperator(LinearOperator):
tmp = dobj.from_local_data(shp2d, ldat2, distaxis=0) tmp = dobj.from_local_data(shp2d, ldat2, distaxis=0)
tmp = dobj.transpose(tmp) tmp = dobj.transpose(tmp)
ldat2 = dobj.local_data(tmp) ldat2 = dobj.local_data(tmp)
ldat2 = utilities.my_fftn(ldat2, axes=(1,)) ldat2 = func(ldat2, axes=(1,))
ldat2 = ldat2.real+ldat2.imag
tmp = dobj.from_local_data(tmp.shape, ldat2, distaxis=0) tmp = dobj.from_local_data(tmp.shape, ldat2, distaxis=0)
tmp = dobj.transpose(tmp) tmp = dobj.transpose(tmp)
ldat2 = dobj.local_data(tmp).reshape(ldat.shape) ldat2 = dobj.local_data(tmp).reshape(ldat.shape)
tmp = dobj.from_local_data(x.val.shape, ldat2, distaxis=0) tmp = dobj.from_local_data(x.val.shape, ldat2, distaxis=0)
Tval = Field(tdom, tmp) Tval = Field(tdom, tmp)
if mode & (LinearOperator.TIMES | LinearOperator.ADJOINT_TIMES): if mode & (LinearOperator.TIMES | LinearOperator.ADJOINT_TIMES):
fct = self._domain[self._space].scalar_dvol fct *= self._domain[self._space].scalar_dvol
else: else:
fct = self._target[self._space].scalar_dvol fct *= self._target[self._space].scalar_dvol
return Tval if fct == 1 else Tval*fct return Tval if fct == 1 else Tval*fct
@property @property
......
...@@ -8,13 +8,14 @@ from ..domain_tuple import DomainTuple ...@@ -8,13 +8,14 @@ from ..domain_tuple import DomainTuple
from ..domains.rg_space import RGSpace from ..domains.rg_space import RGSpace
from ..field import Field from ..field import Field
from .linear_operator import LinearOperator from .linear_operator import LinearOperator
from .. import utilities
class FieldZeroPadder(LinearOperator): class FieldZeroPadder(LinearOperator):
def __init__(self, domain, factor, space=0): def __init__(self, domain, factor, space=0):
super(FieldZeroPadder, self).__init__() super(FieldZeroPadder, self).__init__()
self._domain = DomainTuple.make(domain) self._domain = DomainTuple.make(domain)
self._space = int(space) self._space = utilities.infer_space(self._domain, space)
dom = self._domain[self._space] dom = self._domain[self._space]
if not isinstance(dom, RGSpace): if not isinstance(dom, RGSpace):
raise TypeError("RGSpace required") raise TypeError("RGSpace required")
...@@ -52,11 +53,11 @@ class FieldZeroPadder(LinearOperator): ...@@ -52,11 +53,11 @@ class FieldZeroPadder(LinearOperator):
curax = dobj.distaxis(x) curax = dobj.distaxis(x)
if mode == self.ADJOINT_TIMES: if mode == self.ADJOINT_TIMES:
newarr = np.empty(dobj.local_shape(shp_out), dtype=x.dtype) newarr = np.empty(dobj.local_shape(shp_out, curax), dtype=x.dtype)
newarr[()] = dobj.local_data(x)[(slice(None),)*ax + newarr[()] = dobj.local_data(x)[(slice(None),)*ax +
(slice(0, shp_out[ax]),)] (slice(0, shp_out[ax]),)]
else: else:
newarr = np.zeros(dobj.local_shape(shp_out), dtype=x.dtype) newarr = np.zeros(dobj.local_shape(shp_out, curax), dtype=x.dtype)
newarr[(slice(None),)*ax + newarr[(slice(None),)*ax +
(slice(0, shp_in[ax]),)] = dobj.local_data(x) (slice(0, shp_in[ax]),)] = dobj.local_data(x)
newarr = dobj.from_local_data(shp_out, newarr, distaxis=curax) newarr = dobj.from_local_data(shp_out, newarr, distaxis=curax)
......
...@@ -22,11 +22,11 @@ from ..compat import * ...@@ -22,11 +22,11 @@ from ..compat import *
from ..domain_tuple import DomainTuple from ..domain_tuple import DomainTuple
from ..utilities import infer_space from ..utilities import infer_space
from .diagonal_operator import DiagonalOperator from .diagonal_operator import DiagonalOperator
from .fft_operator import FFTOperator from .hartley_operator import HartleyOperator
from .scaling_operator import ScalingOperator from .scaling_operator import ScalingOperator
def FFTSmoothingOperator(domain, sigma, space=None): def HarmonicSmoothingOperator(domain, sigma, space=None):
""" This function returns an operator that carries out a smoothing with """ This function returns an operator that carries out a smoothing with
a Gaussian kernel of width `sigma` on the part of `domain` given by a Gaussian kernel of width `sigma` on the part of `domain` given by
`space`. `space`.
...@@ -59,12 +59,12 @@ def FFTSmoothingOperator(domain, sigma, space=None): ...@@ -59,12 +59,12 @@ def FFTSmoothingOperator(domain, sigma, space=None):
space = infer_space(domain, space) space = infer_space(domain, space)
if domain[space].harmonic: if domain[space].harmonic:
raise TypeError("domain must not be harmonic") raise TypeError("domain must not be harmonic")
FFT = FFTOperator(domain, space=space) Hartley = HartleyOperator(domain, space=space)
codomain = FFT.domain[space].get_default_codomain() codomain = Hartley.domain[space].get_default_codomain()
kernel = codomain.get_k_length_array() kernel = codomain.get_k_length_array()
smoother = codomain.get_fft_smoothing_kernel_function(sigma) smoother = codomain.get_fft_smoothing_kernel_function(sigma)
kernel = smoother(kernel) kernel = smoother(kernel)
ddom = list(domain) ddom = list(domain)
ddom[space] = codomain ddom[space] = codomain
diag = DiagonalOperator(kernel, ddom, space) diag = DiagonalOperator(kernel, ddom, space)
return FFT.inverse*diag*FFT return Hartley.inverse*diag*Hartley
...@@ -22,7 +22,7 @@ from .. import utilities ...@@ -22,7 +22,7 @@ from .. import utilities
from ..compat import * from ..compat import *
from ..domain_tuple import DomainTuple from ..domain_tuple import DomainTuple
from ..domains.rg_space import RGSpace from ..domains.rg_space import RGSpace
from .fft_operator import FFTOperator from .hartley_operator import HartleyOperator
from .linear_operator import LinearOperator from .linear_operator import LinearOperator
from .sht_operator import SHTOperator from .sht_operator import SHTOperator
...@@ -65,7 +65,7 @@ class HarmonicTransformOperator(LinearOperator): ...@@ -65,7 +65,7 @@ class HarmonicTransformOperator(LinearOperator):
raise TypeError( raise TypeError(
"HarmonicTransformOperator only works on a harmonic space") "HarmonicTransformOperator only works on a harmonic space")
if isinstance(hspc, RGSpace): if isinstance(hspc, RGSpace):
self._op = FFTOperator(domain, target, space) self._op = HartleyOperator(domain, target, space)
else: else:
self._op = SHTOperator(domain, target, space) self._op = SHTOperator(domain, target, space)
......
# 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-2018 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.
from __future__ import absolute_import, division, print_function
import numpy as np
from .. import dobj, utilities
from ..compat import *
from ..domain_tuple import DomainTuple
from ..domains.rg_space import RGSpace
from ..field import Field
from .linear_operator import LinearOperator
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
of the result field, respectivey.
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-symmetrc kernel, the full FFT must be used instead.
"""
def __init__(self, domain, target=None, space=None):
super(HartleyOperator, self).__init__()
# Initialize domain and target
self._domain = DomainTuple.make(domain)
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)
utilities.fft_prep()
def apply(self, x, mode):
self._check_input(x, mode)
if np.issubdtype(x.dtype, np.complexfloating):
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
ldat = utilities.hartley(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 Hartley pass in between the redistributions
tmp = dobj.redistribute(x.val, nodist=axes)
newax = dobj.distaxis(tmp)
ldat = dobj.local_data(tmp)
ldat = utilities.hartley(ldat, axes=axes)
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)
ldat = utilities.my_fftn_r2c(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 = utilities.my_fftn(ldat2, axes=(1,))
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
@property
def domain(self):
return self._domain
@property
def target(self):
return self._target
@property
def capability(self):
return self._all_ops
...@@ -22,7 +22,7 @@ from .. import dobj ...@@ -22,7 +22,7 @@ from .. import dobj
from ..compat import * from ..compat import *
from ..domain_tuple import DomainTuple from ..domain_tuple import DomainTuple
from ..field import Field from ..field import Field
from ..utilities import hartley from ..utilities import hartley, infer_space
from .linear_operator import LinearOperator from .linear_operator import LinearOperator
...@@ -47,7 +47,7 @@ class QHTOperator(LinearOperator): ...@@ -47,7 +47,7 @@ class QHTOperator(LinearOperator):
""" """
def __init__(self, domain, target, space=0): def __init__(self, domain, target, space=0):
self._domain = DomainTuple.make(domain) self._domain = DomainTuple.make(domain)
self._space = int(space) self._space = infer_space(self._domain, space)
from ..domains.log_rg_space import LogRGSpace from ..domains.log_rg_space import LogRGSpace
if not isinstance(self._domain[self._space], LogRGSpace): if not isinstance(self._domain[self._space], LogRGSpace):
...@@ -57,7 +57,7 @@ class QHTOperator(LinearOperator): ...@@ -57,7 +57,7 @@ class QHTOperator(LinearOperator):
if not self._domain[self._space].harmonic: if not self._domain[self._space].harmonic:
raise TypeError( raise TypeError(
"HarmonicTransformOperator only works on a harmonic space") "QHTOperator only works on a harmonic space")
if target.harmonic: if target.harmonic:
raise TypeError("Target is not a codomain of domain") raise TypeError("Target is not a codomain of domain")
......
...@@ -24,15 +24,16 @@ from ..domain_tuple import DomainTuple ...@@ -24,15 +24,16 @@ from ..domain_tuple import DomainTuple
from ..domains.log_rg_space import LogRGSpace from ..domains.log_rg_space import LogRGSpace
from ..field import Field from ..field import Field
from .endomorphic_operator import EndomorphicOperator from .endomorphic_operator import EndomorphicOperator
from .. import utilities
class SymmetrizingOperator(EndomorphicOperator):