Commit 6ea15ac4 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

Merge branch 'FFT' into 'NIFTy_5'

"Real" FFT Operator

See merge request ift/nifty-dev!51
parents 57006538 01f3ae91
......@@ -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)
# FIXME Remove tau and phi stuff from here. Should not be necessary
......
......@@ -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
......
......@@ -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
......@@ -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)
......
......@@ -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]))
def test_composed_fft(self, index, dtype):
[np.float64, np.float32, np.complex64, np.complex128],
[ift.HartleyOperator, ift.FFTOperator]))
def test_composed_fft(self, index, dtype, op):
tol = _get_rtol(dtype)
a = [a1, a2, a3] = [ift.RGSpace((32,)), ift.RGSpace((4, 4)),
ift.RGSpace((5, 6))]
fft = ift.FFTOperator(domain=a, space=index)
fft = op(domain=a, space=index)
inp = ift.Field.from_random(domain=(a1, a2, a3), random_type='normal',
std=7, mean=3, dtype=dtype)
......@@ -96,15 +99,16 @@ class FFTOperatorTests(unittest.TestCase):
@expand(product([ift.RGSpace(128, distances=3.76, harmonic=True),
ift.RGSpace((15, 27), distances=(.7, .33), harmonic=True),
ift.RGSpace(73, distances=0.5643)],
[np.float64, np.float32, np.complex64, np.complex128]))
def test_normalisation(self, space, tp):
[np.float64, np.float32, np.complex64, np.complex128],
[ift.HartleyOperator, ift.FFTOperator]))
def test_normalisation(self, space, tp, op):
tol = 10 * _get_rtol(tp)
cospace = space.get_default_codomain()
fft = ift.FFTOperator(space, cospace)
fft = op(space, cospace)
inp = ift.Field.from_random(domain=space, random_type='normal',
std=1, mean=2, dtype=tp)
out = fft.times(inp)
fft2 = ift.FFTOperator(cospace, space)
fft2 = op(cospace, space)
out2 = fft2.inverse_times(inp)
zero_idx = tuple([0]*len(space.shape))
assert_allclose(inp.to_global_data()[zero_idx], out.integrate(),
......
......@@ -37,13 +37,13 @@ class SmoothingOperator_Tests(unittest.TestCase):
@expand(product(spaces, [0., .5, 5.]))
def test_property(self, space, sigma):
op = ift.FFTSmoothingOperator(space, sigma=sigma)
op = ift.HarmonicSmoothingOperator(space, sigma=sigma)
if op.domain[0] != space:
raise TypeError
@expand(product(spaces, [0., .5, 5.]))
def test_adjoint_times(self, space, sigma):
op = ift.FFTSmoothingOperator(space, sigma=sigma)
op = ift.HarmonicSmoothingOperator(space, sigma=sigma)
rand1 = ift.Field.from_random('normal', domain=space)
rand2 = ift.Field.from_random('normal', domain=space)
tt1 = rand1.vdot(op.times(rand2))
......@@ -52,7 +52,7 @@ class SmoothingOperator_Tests(unittest.TestCase):
@expand(product(spaces, [0., .5, 5.]))
def test_times(self, space, sigma):
op = ift.FFTSmoothingOperator(space, sigma=sigma)
op = ift.HarmonicSmoothingOperator(space, sigma=sigma)
fld = np.zeros(space.shape, dtype=np.float64)
fld[0] = 1.
rand1 = ift.Field.from_global_data(space, fld)
......@@ -64,7 +64,7 @@ class SmoothingOperator_Tests(unittest.TestCase):
def test_smooth_regular1(self, sz, d, sigma, tp):
tol = _get_rtol(tp)
sp = ift.RGSpace(sz, distances=d)
smo = ift.FFTSmoothingOperator(sp, sigma=sigma)
smo = ift.HarmonicSmoothingOperator(sp, sigma=sigma)
inp = ift.Field.from_random(domain=sp, random_type='normal', std=1,
mean=4, dtype=tp)
out = smo(inp)
......@@ -75,7 +75,7 @@ class SmoothingOperator_Tests(unittest.TestCase):
def test_smooth_regular2(self, sz1, sz2, d1, d2, sigma, tp):
tol = _get_rtol(tp)
sp = ift.RGSpace([sz1, sz2], distances=[d1, d2])
smo = ift.FFTSmoothingOperator(sp, sigma=sigma)
smo = ift.HarmonicSmoothingOperator(sp, sigma=sigma)
inp = ift.Field.from_random(domain=sp, random_type='normal', std=1,
mean=4, dtype=tp)
out = smo(inp)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment