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
from .operators.endomorphic_operator import EndomorphicOperator
from .operators.exp_transform import ExpTransform
from .operators.fft_operator import FFTOperator
from .operators.fft_smoothing_operator import FFTSmoothingOperator
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.harmonic_transform_operator import HarmonicTransformOperator
from .operators.inversion_enabler import InversionEnabler
......
......@@ -25,7 +25,7 @@ from ..models.local_nonlinearity import PointwiseExponential
from ..models.variable import Variable
from ..multi.multi_field import MultiField
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.power_distributor import PowerDistributor
......@@ -41,7 +41,7 @@ def make_correlated_field(s_space, amplitude_model):
amplitude_model : model for correlation structure
'''
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]
power_distributor = PowerDistributor(h_space, p_space)
......
......@@ -25,26 +25,16 @@ from ..compat import *
from ..domain_tuple import DomainTuple
from ..field import Field
from .linear_operator import LinearOperator
from .. import utilities
# MR FIXME: this needs to be rewritten in a generic fashion
class DomainDistributor(LinearOperator):
def __init__(self, target, axis):
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)
def __init__(self, target, spaces):
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
def domain(self):
......@@ -57,23 +47,16 @@ class DomainDistributor(LinearOperator):
def apply(self, x, mode):
self._check_input(x, mode)
if mode == self.TIMES:
x = x.local_data
otherDirection = np.ones(self._size)
if self._axis == 0:
res = np.outer(otherDirection, x)
else:
res = np.outer(x, otherDirection)
res = res.reshape(dobj.local_shape(self.target.shape))
return Field.from_local_data(self.target, res)
ldat = x.local_data if 0 in self._spaces else x.to_global_data()
shp = []
for i, tgt in enumerate(self._target):
tmp = tgt.shape if i > 0 else tgt.local_shape
shp += tmp if i in self._spaces else(1,)*len(tgt.shape)
ldat = np.broadcast_to(ldat.reshape(shp), self._target.local_shape)
return Field.from_local_data(self._target, ldat)
else:
if self._axis == 0:
x = x.local_data.reshape(self._size, -1)
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)
return x.sum([s for s in range(len(x.domain))
if s not in self._spaces])
@property
def capability(self):
......
......@@ -27,12 +27,13 @@ from ..domains.power_space import PowerSpace
from ..domains.rg_space import RGSpace
from ..field import Field
from .linear_operator import LinearOperator
from .. import utilities
class ExpTransform(LinearOperator):
def __init__(self, target, dof, space=0):
self._target = DomainTuple.make(target)
self._space = int(space)
self._space = utilities.infer_space(self._target, space)
tgt = self._target[self._space]
if not ((isinstance(tgt, RGSpace) and tgt.harmonic) or
isinstance(tgt, PowerSpace)):
......
......@@ -43,6 +43,13 @@ class FFTOperator(LinearOperator):
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):
......@@ -67,41 +74,35 @@ class FFTOperator(LinearOperator):
utilities.fft_prep()
def apply(self, x, mode):
from pyfftw.interfaces.numpy_fft import fftn, ifftn
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))
ncells = x.domain[self._space].size
if x.domain[self._space].harmonic: # harmonic -> position
func = fftn
fct = 1.
else:
return self._apply_cartesian(x, mode)
def _apply_cartesian(self, x, mode):
func = ifftn
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 = utilities.hartley(ldat, axes=axes)
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 Hartley pass in between the redistributions
# 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 = utilities.hartley(ldat, axes=axes)
ldat = func(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
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 = utilities.my_fftn_r2c(ldat, axes=rem_axes)
ldat = func(ldat, axes=rem_axes)
if oldax != 0:
raise ValueError("bad distribution")
ldat2 = ldat.reshape((ldat.shape[0],
......@@ -110,17 +111,16 @@ class FFTOperator(LinearOperator):
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
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
fct *= self._domain[self._space].scalar_dvol
else:
fct = self._target[self._space].scalar_dvol
fct *= self._target[self._space].scalar_dvol
return Tval if fct == 1 else Tval*fct
@property
......
......@@ -8,13 +8,14 @@ from ..domain_tuple import DomainTuple
from ..domains.rg_space import RGSpace
from ..field import Field
from .linear_operator import LinearOperator
from .. import utilities
class FieldZeroPadder(LinearOperator):
def __init__(self, domain, factor, space=0):
super(FieldZeroPadder, self).__init__()
self._domain = DomainTuple.make(domain)
self._space = int(space)
self._space = utilities.infer_space(self._domain, space)
dom = self._domain[self._space]
if not isinstance(dom, RGSpace):
raise TypeError("RGSpace required")
......@@ -52,11 +53,11 @@ class FieldZeroPadder(LinearOperator):
curax = dobj.distaxis(x)
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 +
(slice(0, shp_out[ax]),)]
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 +
(slice(0, shp_in[ax]),)] = dobj.local_data(x)
newarr = dobj.from_local_data(shp_out, newarr, distaxis=curax)
......
......@@ -22,11 +22,11 @@ from ..compat import *
from ..domain_tuple import DomainTuple
from ..utilities import infer_space
from .diagonal_operator import DiagonalOperator
from .fft_operator import FFTOperator
from .hartley_operator import HartleyOperator
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
a Gaussian kernel of width `sigma` on the part of `domain` given by
`space`.
......@@ -59,12 +59,12 @@ def FFTSmoothingOperator(domain, sigma, space=None):
space = infer_space(domain, space)
if domain[space].harmonic:
raise TypeError("domain must not be harmonic")
FFT = FFTOperator(domain, space=space)
codomain = FFT.domain[space].get_default_codomain()
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 FFT.inverse*diag*FFT
return Hartley.inverse*diag*Hartley
......@@ -22,7 +22,7 @@ from .. import utilities
from ..compat import *
from ..domain_tuple import DomainTuple
from ..domains.rg_space import RGSpace
from .fft_operator import FFTOperator
from .hartley_operator import HartleyOperator
from .linear_operator import LinearOperator
from .sht_operator import SHTOperator
......@@ -65,7 +65,7 @@ class HarmonicTransformOperator(LinearOperator):
raise TypeError(
"HarmonicTransformOperator only works on a harmonic space")
if isinstance(hspc, RGSpace):
self._op = FFTOperator(domain, target, space)
self._op = HartleyOperator(domain, target, space)
else:
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
from ..compat import *
from ..domain_tuple import DomainTuple
from ..field import Field
from ..utilities import hartley
from ..utilities import hartley, infer_space
from .linear_operator import LinearOperator
......@@ -47,7 +47,7 @@ class QHTOperator(LinearOperator):
"""
def __init__(self, domain, target, space=0):
self._domain = DomainTuple.make(domain)
self._space = int(space)
self._space = infer_space(self._domain, space)
from ..domains.log_rg_space import LogRGSpace
if not isinstance(self._domain[self._space], LogRGSpace):
......@@ -57,7 +57,7 @@ class QHTOperator(LinearOperator):
if not self._domain[self._space].harmonic:
raise TypeError(
"HarmonicTransformOperator only works on a harmonic space")
"QHTOperator only works on a harmonic space")
if target.harmonic:
raise TypeError("Target is not a codomain of domain")
......
......@@ -24,15 +24,16 @@ from ..domain_tuple import DomainTuple
from ..domains.log_rg_space import LogRGSpace
from ..field import Field
from .endomorphic_operator import EndomorphicOperator
from .. import utilities
class SymmetrizingOperator(EndomorphicOperator):
def __init__(self, domain, space=0):
self._domain = DomainTuple.make(domain)
self._space = int(space)
self._space = utilities.infer_space(self._domain, space)
dom = self._domain[self._space]
if not (isinstance(dom, LogRGSpace) and not dom.harmonic):
raise TypeError
raise TypeError("nonharmonic LogRGSpace needed")
@property
def domain(self):
......
......@@ -56,6 +56,14 @@ class Consistency_Tests(unittest.TestCase):
op = ift.FFTOperator(sp.get_default_codomain())
ift.extra.consistency_check(op, dtype, dtype)
@expand(product(_h_RG_spaces+_p_RG_spaces,
[np.float64, np.complex128]))
def testHartley(self, sp, dtype):
op = ift.HartleyOperator(sp)
ift.extra.consistency_check(op, dtype, dtype)
op = ift.HartleyOperator(sp.get_default_codomain())
ift.extra.consistency_check(op, dtype, dtype)
@expand(product(_h_spaces, [np.float64, np.complex128]))
def testHarmonic(self, sp, dtype):
op = ift.HarmonicTransformOperator(sp)
......@@ -93,3 +101,50 @@ class Consistency_Tests(unittest.TestCase):
def testGeometryRemover(self, sp, dtype):
op = ift.GeometryRemover(sp)
ift.extra.consistency_check(op, dtype, dtype)
@expand(product([0, 1, 2, 3, (0, 1), (0, 2), (0, 1, 2), (0, 2, 3), (1, 3)],
[np.float64, np.complex128]))
def testDomainDistributor(self, spaces, dtype):
dom = (ift.RGSpace(10), ift.UnstructuredDomain(13), ift.GLSpace(5),
ift.HPSpace(4))
op = ift.DomainDistributor(dom, spaces)
ift.extra.consistency_check(op, dtype, dtype)
@expand(product([0, 2], [np.float64, np.complex128]))
def testSymmetrizingOperator(self, space, dtype):
dom = (ift.LogRGSpace(10, [2.], [1.]), ift.UnstructuredDomain(13),
ift.LogRGSpace((5, 27), [1., 2.7], [0., 4.]), ift.HPSpace(4))
op = ift.SymmetrizingOperator(dom, space)
ift.extra.consistency_check(op, dtype, dtype)
@expand(product([0, 2], [2, 2.7], [np.float64, np.complex128]))
def testZeroPadder(self, space, factor, dtype):
dom = (ift.RGSpace(10), ift.UnstructuredDomain(13), ift.RGSpace(7),
ift.HPSpace(4))
op = ift.FieldZeroPadder(dom, factor, space)
ift.extra.consistency_check(op, dtype, dtype)
@expand(product([(ift.RGSpace(10, harmonic=True), 4, 0),
(ift.RGSpace((24, 31), distances=(0.4, 2.34),
harmonic=True), (4, 3), 0),
((ift.HPSpace(4), ift.RGSpace(27, distances=0.3,
harmonic=True)), (10,), 1),
(ift.PowerSpace(ift.RGSpace(10, distances=0.3,
harmonic=True)), 6, 0)],
[np.float64, np.complex128]))
def testExpTransform(self, args, dtype):
op = ift.ExpTransform(args[0], args[1], args[2])
ift.extra.consistency_check(op, dtype, dtype)
@expand(product([(ift.LogRGSpace([10, 17], [2., 3.], [1., 0.]), 0),
((ift.LogRGSpace(10, [2.], [1.]),
ift.UnstructuredDomain(13)), 0),
((ift.UnstructuredDomain(13),
ift.LogRGSpace(17, [3.], [.7])), 1)],
[np.float64]))
def testQHTOperator(self, args, dtype):
dom = ift.DomainTuple.make(args[0])
tgt = list(dom)
tgt[args[1]] = tgt[args[1]].get_default_codomain()
op = ift.QHTOperator(tgt, dom[args[1]], args[1])
ift.extra.consistency_check(op, dtype, dtype)
......@@ -36,14 +36,15 @@ def _get_rtol(tp):
class FFTOperatorTests(unittest.TestCase):
@expand(product([16, ], [0.1, 1, 3.7],
[np.float64, np.float32, np.complex64, np.complex128]))
def test_fft1D(self, dim1, d, itp):
[np.float64, np.float32, np.complex64, np.complex128],
[ift.HartleyOperator, ift.FFTOperator]))
def test_fft1D(self, dim1, d, itp, op):
tol = _get_rtol(itp)
a = ift.RGSpace(dim1, distances=d)
b = ift.RGSpace(dim1, distances=1./(dim1*d), harmonic=True)
np.random.seed(16)
fft = ift.FFTOperator(domain=a, target=b)
fft = op(domain=a, target=b)
inp = ift.Field.from_random(domain=a, random_type='normal',
std=7, mean=3, dtype=itp)
out = fft.inverse_times(fft.times(inp))
......@@ -59,14 +60,15 @@ class FFTOperatorTests(unittest.TestCase):
@expand(product([12, 15], [9, 12], [0.1, 1, 3.7],
[0.4, 1, 2.7],
[np.float64, np.float32, np.complex64, np.complex128]))
def test_fft2D(self, dim1, dim2, d1, d2, itp):
[np.float64, np.float32, np.complex64, np.complex128],
[ift.HartleyOperator, ift.FFTOperator]))
def test_fft2D(self, dim1, dim2, d1, d2, itp, op):
tol = _get_rtol(itp)
a = ift.RGSpace([dim1, dim2], distances=[d1, d2])
b = ift.RGSpace([dim1, dim2],
distances=[1./(dim1*d1), 1./(dim2*d2)], harmonic=True)
fft = ift.FFTOperator(domain=a, target=b)
fft = op(domain=a, target=b)
inp = ift.Field.from_random(domain=a, random_type='normal',
std=7, mean=3, dtype=itp)
out = fft.inverse_times(fft.times(inp))
......@@ -81,12 +83,13 @@ class FFTOperatorTests(unittest.TestCase):
assert_allclose(inp.local_data, out.local_data, rtol=tol, atol=tol)
@expand(product([0, 1, 2],
[np.float64, np.float32, np.complex64, np.complex128]))