Skip to content
GitLab
Menu
Projects
Groups
Snippets
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Neel Shah
NIFTy
Commits
a63ef1fc
Commit
a63ef1fc
authored
Jan 28, 2018
by
Philipp Arras
Browse files
Merge branch 'NIFTy_4' into energy_tests
parents
e0da17b8
7fe10569
Changes
14
Hide whitespace changes
Inline
Side-by-side
.gitlab-ci.yml
View file @
a63ef1fc
...
...
@@ -17,10 +17,10 @@ test_min:
stage
:
test
script
:
-
pip install --user .
-
nosetests -
x
--with-coverage --cover-package=nifty4 --cover-branches
-
OMP_NUM_THREADS=1 mpiexec --allow-run-as-root -n 4 nosetests -
x --with-coverage --cover-package=nifty4 --cover-branches
-
nosetests -
q
--with-coverage --cover-package=nifty4 --cover-branches
-
OMP_NUM_THREADS=1 mpiexec --allow-run-as-root -n 4 nosetests -
q
-
pip3 install --user .
-
nosetests3 -
x
-
OMP_NUM_THREADS=1 mpiexec --allow-run-as-root -n 4 nosetests3 -
x
-
nosetests3 -
q
-
OMP_NUM_THREADS=1 mpiexec --allow-run-as-root -n 4 nosetests3 -
q
-
>
coverage report | grep TOTAL | awk '{ print "TOTAL: "$6; }'
coverage report | grep TOTAL | awk '{ print "TOTAL: "$6; }'
|| true
demos/critical_filtering.py
View file @
a63ef1fc
...
...
@@ -12,7 +12,7 @@ if __name__ == "__main__":
# Define harmonic transformation and associated harmonic space
h_space
=
s_space
.
get_default_codomain
()
fft
=
ift
.
FFT
Operator
(
h_space
,
s_space
)
fft
=
ift
.
HarmonicTransform
Operator
(
h_space
,
s_space
)
# Set up power space
p_space
=
ift
.
PowerSpace
(
h_space
,
...
...
demos/wiener_filter_via_curvature.py
View file @
a63ef1fc
...
...
@@ -2,6 +2,7 @@ import numpy as np
import
nifty4
as
ift
import
numericalunits
as
nu
if
__name__
==
"__main__"
:
# In MPI mode, the random seed for numericalunits must be set by hand
#nu.reset_units(43)
...
...
@@ -38,7 +39,7 @@ if __name__ == "__main__":
signal_space
=
ift
.
RGSpace
(
shape
,
distances
=
L
/
N_pixels
)
harmonic_space
=
signal_space
.
get_default_codomain
()
fft
=
ift
.
FFT
Operator
(
harmonic_space
,
target
=
signal_space
)
fft
=
ift
.
HarmonicTransform
Operator
(
harmonic_space
,
target
=
signal_space
)
power_space
=
ift
.
PowerSpace
(
harmonic_space
)
# Creating the mock data
...
...
@@ -48,34 +49,27 @@ if __name__ == "__main__":
mock_power
=
ift
.
PS_field
(
power_space
,
power_spectrum
)
mock_harmonic
=
ift
.
power_synthesize
(
mock_power
,
real_signal
=
True
)
print
mock_harmonic
.
val
[
0
]
/
nu
.
K
/
(
nu
.
m
**
dimensionality
)
mock_signal
=
fft
(
mock_harmonic
)
print
"msig"
,
mock_signal
.
val
[
0
:
10
]
/
nu
.
K
sensitivity
=
(
1.
/
nu
.
m
)
**
dimensionality
/
nu
.
K
R
=
ift
.
ResponseOperator
(
signal_space
,
sigma
=
(
0.
*
response_sigma
,),
R
=
ift
.
ResponseOperator
(
signal_space
,
sigma
=
(
response_sigma
,),
sensitivity
=
(
sensitivity
,))
data_domain
=
R
.
target
[
0
]
R_harmonic
=
R
*
fft
noise_amplitude
=
1.
/
signal_to_noise
*
field_sigma
*
sensitivity
*
((
L
/
N_pixels
)
**
dimensionality
)
print
noise_amplitude
print
"noise amplitude:"
,
noise_amplitude
N
=
ift
.
DiagonalOperator
(
ift
.
Field
.
full
(
data_domain
,
noise_amplitude
**
2
))
noise
=
ift
.
Field
.
from_random
(
domain
=
data_domain
,
random_type
=
'normal'
,
std
=
noise_amplitude
,
mean
=
0
)
data
=
R
(
mock_signal
)
print
data
.
val
[
5
:
10
]
data
+=
noise
print
data
.
val
[
5
:
10
]
data
=
R
(
mock_signal
)
+
noise
# Wiener filter
j
=
R_harmonic
.
adjoint_times
(
N
.
inverse_times
(
data
))
print
"xx"
,
j
.
val
[
0
]
*
nu
.
K
*
(
nu
.
m
**
dimensionality
)
exit
()
ctrl
=
ift
.
GradientNormController
(
verbose
=
True
,
tol_abs_gradnorm
=
1e-
40
/
(
nu
.
K
*
(
nu
.
m
**
dimensionality
)))
verbose
=
True
,
tol_abs_gradnorm
=
1e-
5
/
(
nu
.
K
*
(
nu
.
m
**
dimensionality
)))
inverter
=
ift
.
ConjugateGradient
(
controller
=
ctrl
)
wiener_curvature
=
ift
.
library
.
WienerFilterCurvature
(
S
=
S
,
N
=
N
,
R
=
R_harmonic
,
inverter
=
inverter
)
...
...
@@ -85,8 +79,10 @@ if __name__ == "__main__":
sspace2
=
ift
.
RGSpace
(
shape
,
distances
=
L
/
N_pixels
/
nu
.
m
)
ift
.
plot
(
ift
.
Field
(
sspace2
,
mock_signal
.
val
)
/
nu
.
K
,
name
=
"mock_signal.png"
)
data
=
ift
.
dobj
.
to_global_data
(
data
.
val
).
reshape
(
sspace2
.
shape
)
data
=
ift
.
Field
(
sspace2
,
val
=
ift
.
dobj
.
from_global_data
(
data
))
ift
.
plot
(
ift
.
Field
(
sspace2
,
val
=
data
),
name
=
"data.png"
)
ift
.
plot
(
ift
.
Field
(
sspace2
,
m_s
.
val
)
/
nu
.
K
,
name
=
"map.png"
)
ift
.
plot
(
ift
.
Field
(
sspace2
,
mock_signal
.
val
)
/
nu
.
K
,
title
=
"mock_signal.png"
)
#data = ift.dobj.to_global_data(data.val).reshape(sspace2.shape)
#data = ift.Field(sspace2, val=ift.dobj.from_global_data(data))
ift
.
plot
(
ift
.
Field
(
sspace2
,
val
=
R
.
adjoint_times
(
data
).
val
),
title
=
"data.png"
)
print
"msig"
,
np
.
min
(
mock_signal
.
val
)
/
nu
.
K
,
np
.
max
(
mock_signal
.
val
)
/
nu
.
K
print
"map"
,
np
.
min
(
m_s
.
val
)
/
nu
.
K
,
np
.
max
(
m_s
.
val
)
/
nu
.
K
ift
.
plot
(
ift
.
Field
(
sspace2
,
m_s
.
val
)
/
nu
.
K
,
title
=
"map.png"
)
demos/wiener_filter_via_hamiltonian.py
View file @
a63ef1fc
...
...
@@ -11,7 +11,7 @@ if __name__ == "__main__":
# Define associated harmonic space and harmonic transformation
h_space
=
s_space
.
get_default_codomain
()
fft
=
ift
.
FFT
Operator
(
h_space
,
s_space
)
fft
=
ift
.
HarmonicTransform
Operator
(
h_space
,
s_space
)
# Setting up power space
p_space
=
ift
.
PowerSpace
(
h_space
)
...
...
docs/source/start.rst
View file @
a63ef1fc
...
...
@@ -116,7 +116,6 @@ typical examples for this category are the :py:class:`ScalingOperator`, which si
Nifty4 allows simple and intuitive construction of combined operators.
As an example, if :math:`A`, :math:`B` and :math:`C` are of type :py:class:`LinearOperator` and :math:`f_1` and :math:`f_2` are fields, writing::
X = A*B.inverse*A.adjoint + C
f2 = X(f1)
...
...
nifty4/__init__.py
View file @
a63ef1fc
...
...
@@ -19,6 +19,7 @@ from .operators.linear_operator import LinearOperator
from
.operators.endomorphic_operator
import
EndomorphicOperator
from
.operators.scaling_operator
import
ScalingOperator
from
.operators.diagonal_operator
import
DiagonalOperator
from
.operators.harmonic_transform_operator
import
HarmonicTransformOperator
from
.operators.fft_operator
import
FFTOperator
from
.operators.fft_smoothing_operator
import
FFTSmoothingOperator
from
.operators.direct_smoothing_operator
import
DirectSmoothingOperator
...
...
nifty4/minimization/scipy_minimizer.py
View file @
a63ef1fc
...
...
@@ -45,7 +45,7 @@ class ScipyMinimizer(Minimizer):
self
.
_need_hessp
=
need_hessp
def
__call__
(
self
,
energy
):
class
_MinimizationDone
:
class
_MinimizationDone
(
BaseException
)
:
pass
class
_MinHelper
(
object
):
...
...
nifty4/minimization/scipy_minimizer.py.bak
0 → 100644
View file @
a63ef1fc
# 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-2017 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 division
from .minimizer import Minimizer
from ..field import Field
from .. import dobj
class ScipyMinimizer(Minimizer):
"""Scipy-based minimizer
Parameters
----------
controller : IterationController
Object that decides when to terminate the minimization.
"""
def __init__(self, controller, method="trust-ncg"):
super(ScipyMinimizer, self).__init__()
if not dobj.is_numpy():
raise NotImplementedError
self._controller = controller
self._method = method
def __call__(self, energy):
class _MinimizationDone:
pass
class _MinHelper(object):
def __init__(self, controller, energy):
self._controller = controller
self._energy = energy
self._domain = energy.position.domain
def _update(self, x):
pos = Field(self._domain, x.reshape(self._domain.shape))
if (pos.val != self._energy.position.val).any():
self._energy = self._energy.at(pos)
status = self._controller.check(self._energy)
if status != self._controller.CONTINUE:
raise _MinimizationDone
def fun(self, x):
self._update(x)
return self._energy.value
def jac(self, x):
self._update(x)
return self._energy.gradient.val.reshape(-1)
def hessp(self, x, p):
self._update(x)
vec = Field(self._domain, p.reshape(self._domain.shape))
res = self._energy.curvature(vec)
return res.val.reshape(-1)
import scipy.optimize as opt
status = self._controller.start(energy)
if status != self._controller.CONTINUE:
return energy, status
hlp = _MinHelper(self._controller, energy)
options = {'disp': False,
'xtol': 1e-15,
'eps': 1.4901161193847656e-08,
'return_all': False,
'maxiter': None}
options = {'disp': False,
'ftol': 1e-15,
'gtol': 1e-15,
'eps': 1.4901161193847656e-08}
try:
opt.minimize(hlp.fun, energy.position.val.reshape(-1),
method=self._method, jac=hlp.jac,
hessp=hlp.hessp,
options=options)
except _MinimizationDone:
energy = hlp._energy
status = self._controller.check(energy)
return energy, status
return hlp._energy, self._controller.ERROR
nifty4/operators/fft_operator.py
View file @
a63ef1fc
...
...
@@ -23,42 +23,22 @@ from .linear_operator import LinearOperator
from
..
import
dobj
from
..
import
utilities
from
..field
import
Field
from
..spaces.gl_space
import
GLSpace
class
FFTOperator
(
LinearOperator
):
"""Transforms between a pair of position and harmonic domains.
Built-in domain pairs are
- a harmonic and a non-harmonic RGSpace (with matching distances)
- a HPSpace and a LMSpace
- a GLSpace and a LMSpace
Within a domain pair, both orderings are possible.
For RGSpaces, the operator provides the full set of operations.
For the sphere-related domains, it only supports the transform from
harmonic to position space and its adjoint; if the operator domain is
harmonic, this will be times() and adjoint_times(), otherwise
inverse_times() and adjoint_inverse_times()
"""Transforms between a pair of position and harmonic RGSpaces.
Parameters
----------
domain: Space
or single-element tuple of Spaces
domain: Space
, tuple of Spaces or DomainObject
The domain of the data that is input by "times" and output by
"adjoint_times".
target: Space
The target space of the transform operation.
If omitted, a space will be chosen automatically.
space: the index of the space on which the operator should act
If None, it is set to 0 if domain contains exactly one space
target: Space or single-element tuple of Spaces (optional)
The domain of the data that is output by "times" and input by
"adjoint_times".
If omitted, a co-domain will be chosen automatically.
Whenever "domain" is an RGSpace, the codomain (and its parameters) are
uniquely determined.
For GLSpace, HPSpace, and LMSpace, a sensible (but not unique)
co-domain is chosen that should work satisfactorily in most situations,
but for full control, the user should explicitly specify a codomain.
If None, it is set to 0 if domain contains exactly one space.
domain[space] must be an RGSpace.
"""
def
__init__
(
self
,
domain
,
target
=
None
,
space
=
None
):
...
...
@@ -69,6 +49,8 @@ class FFTOperator(LinearOperator):
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
()
...
...
@@ -78,47 +60,18 @@ class FFTOperator(LinearOperator):
adom
.
check_codomain
(
target
)
target
.
check_codomain
(
adom
)
if
isinstance
(
adom
,
RGSpace
):
self
.
_applyfunc
=
self
.
_apply_cartesian
self
.
_capability
=
self
.
_all_ops
import
pyfftw
pyfftw
.
interfaces
.
cache
.
enable
()
else
:
from
pyHealpix
import
sharpjob_d
self
.
_applyfunc
=
self
.
_apply_spherical
hspc
=
adom
if
adom
.
harmonic
else
target
pspc
=
target
if
adom
.
harmonic
else
adom
self
.
lmax
=
hspc
.
lmax
self
.
mmax
=
hspc
.
mmax
self
.
sjob
=
sharpjob_d
()
self
.
sjob
.
set_triangular_alm_info
(
self
.
lmax
,
self
.
mmax
)
if
isinstance
(
pspc
,
GLSpace
):
self
.
sjob
.
set_Gauss_geometry
(
pspc
.
nlat
,
pspc
.
nlon
)
else
:
self
.
sjob
.
set_Healpix_geometry
(
pspc
.
nside
)
if
adom
.
harmonic
:
self
.
_capability
=
self
.
TIMES
|
self
.
ADJOINT_TIMES
else
:
self
.
_capability
=
(
self
.
INVERSE_TIMES
|
self
.
INVERSE_ADJOINT_TIMES
)
import
pyfftw
pyfftw
.
interfaces
.
cache
.
enable
()
def
apply
(
self
,
x
,
mode
):
self
.
_check_input
(
x
,
mode
)
if
np
.
issubdtype
(
x
.
dtype
,
np
.
complexfloating
):
return
(
self
.
_apply
func
(
x
.
real
,
mode
)
+
1j
*
self
.
_apply
func
(
x
.
imag
,
mode
))
return
(
self
.
_apply
_cartesian
(
x
.
real
,
mode
)
+
1j
*
self
.
_apply
_cartesian
(
x
.
imag
,
mode
))
else
:
return
self
.
_apply
func
(
x
,
mode
)
return
self
.
_apply
_cartesian
(
x
,
mode
)
def
_apply_cartesian
(
self
,
x
,
mode
):
"""
RG -> RG transform method.
Parameters
----------
x : Field
The field to be transformed
"""
from
pyfftw.interfaces.numpy_fft
import
fftn
axes
=
x
.
domain
.
axes
[
self
.
_space
]
tdom
=
self
.
_target
if
x
.
domain
==
self
.
_domain
else
self
.
_domain
...
...
@@ -171,47 +124,6 @@ class FFTOperator(LinearOperator):
return
Tval
def
_slice_p2h
(
self
,
inp
):
rr
=
self
.
sjob
.
alm2map_adjoint
(
inp
)
assert
len
(
rr
)
==
((
self
.
mmax
+
1
)
*
(
self
.
mmax
+
2
))
//
2
+
\
(
self
.
mmax
+
1
)
*
(
self
.
lmax
-
self
.
mmax
)
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
)
assert
len
(
res
)
==
((
self
.
mmax
+
1
)
*
(
self
.
mmax
+
2
))
//
2
+
\
(
self
.
mmax
+
1
)
*
(
self
.
lmax
-
self
.
mmax
)
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
]
tval
=
x
.
val
if
dobj
.
distaxis
(
tval
)
==
axis
:
tval
=
dobj
.
redistribute
(
tval
,
nodist
=
(
axis
,))
distaxis
=
dobj
.
distaxis
(
tval
)
p2h
=
not
x
.
domain
[
self
.
_space
].
harmonic
tdom
=
self
.
_target
if
x
.
domain
==
self
.
_domain
else
self
.
_domain
func
=
self
.
_slice_p2h
if
p2h
else
self
.
_slice_h2p
idat
=
dobj
.
local_data
(
tval
)
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
)
if
distaxis
!=
dobj
.
distaxis
(
x
.
val
):
odat
=
dobj
.
redistribute
(
odat
,
dist
=
dobj
.
distaxis
(
x
.
val
))
return
Field
(
tdom
,
odat
)
@
property
def
domain
(
self
):
return
self
.
_domain
...
...
@@ -222,4 +134,4 @@ class FFTOperator(LinearOperator):
@
property
def
capability
(
self
):
return
self
.
_
capability
return
self
.
_
all_ops
nifty4/operators/harmonic_transform_operator.py
0 → 100644
View file @
a63ef1fc
# 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-2017 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.
import
numpy
as
np
from
..domain_tuple
import
DomainTuple
from
..spaces.rg_space
import
RGSpace
from
.linear_operator
import
LinearOperator
from
..
import
dobj
from
..
import
utilities
from
..field
import
Field
from
..spaces.gl_space
import
GLSpace
class
HarmonicTransformOperator
(
LinearOperator
):
"""Transforms between a harmonic domain and a position space counterpart.
Built-in domain pairs are
- a harmonic and a non-harmonic RGSpace (with matching distances)
- an LMSpace and a LMSpace
- an LMSpace and a GLSpace
The supported operations are times() and adjoint_times().
Parameters
----------
domain: Space, tuple of Spaces or DomainObject
The domain of the data that is input by "times" and output by
"adjoint_times".
target: Space (optional)
The target space of the transform operation.
If omitted, a space will be chosen automatically.
Whenever the input space of the transform is an RGSpace, the codomain
(and its parameters) are uniquely determined.
For LMSpace, a GLSpace of sufficient resolution is chosen.
space: the index of the space on which the operator should act
If None, it is set to 0 if domain contains exactly one space.
domain[space] must be a harmonic space.
"""
def
__init__
(
self
,
domain
,
target
=
None
,
space
=
None
):
super
(
HarmonicTransformOperator
,
self
).
__init__
()
# Initialize domain and target
self
.
_domain
=
DomainTuple
.
make
(
domain
)
self
.
_space
=
utilities
.
infer_space
(
self
.
_domain
,
space
)
hspc
=
self
.
_domain
[
self
.
_space
]
if
not
hspc
.
harmonic
:
raise
TypeError
(
"HarmonicTransformOperator only works on a harmonic space"
)
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
)
if
isinstance
(
hspc
,
RGSpace
):
self
.
_applyfunc
=
self
.
_apply_cartesian
import
pyfftw
pyfftw
.
interfaces
.
cache
.
enable
()
else
:
from
pyHealpix
import
sharpjob_d
self
.
_applyfunc
=
self
.
_apply_spherical
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
)
def
apply
(
self
,
x
,
mode
):
self
.
_check_input
(
x
,
mode
)
if
np
.
issubdtype
(
x
.
dtype
,
np
.
complexfloating
):
return
(
self
.
_applyfunc
(
x
.
real
,
mode
)
+
1j
*
self
.
_applyfunc
(
x
.
imag
,
mode
))
else
:
return
self
.
_applyfunc
(
x
,
mode
)
def
_apply_cartesian
(
self
,
x
,
mode
):
from
pyfftw.interfaces.numpy_fft
import
fftn
axes
=
x
.
domain
.
axes
[
self
.
_space
]
tdom
=
self
.
_target
if
x
.
domain
==
self
.
_domain
else
self
.
_domain
oldax
=
dobj
.
distaxis
(
x
.
val
)
if
oldax
not
in
axes
:
# straightforward, no redistribution needed
ldat
=
dobj
.
local_data
(
x
.
val
)
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
=
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
)
fct
=
self
.
_domain
[
self
.
_space
].
scalar_dvol
()
if
fct
!=
1
:
Tval
*=
fct
return
Tval
def
_slice_p2h
(
self
,
inp
):
rr
=
self
.
sjob
.
alm2map_adjoint
(
inp
)
assert
len
(
rr
)
==
((
self
.
mmax
+
1
)
*
(
self
.
mmax
+
2
))
//
2
+
\
(
self
.
mmax
+
1
)
*
(
self
.
lmax
-
self
.
mmax
)
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
)
assert
len
(
res
)
==
((
self
.
mmax
+
1
)
*
(
self
.
mmax
+
2
))
//
2
+
\
(
self
.
mmax
+
1
)
*
(
self
.
lmax
-
self
.
mmax
)
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
]
tval
=
x
.
val
if
dobj
.
distaxis
(
tval
)
==
axis
:
tval
=
dobj
.
redistribute
(
tval
,
nodist
=
(
axis
,))
distaxis
=
dobj
.
distaxis
(
tval
)
p2h
=
not
x
.
domain
[
self
.
_space
].
harmonic
tdom
=
self
.
_target
if
x
.
domain
==
self
.
_domain
else
self
.
_domain
func
=
self
.
_slice_p2h
if
p2h
else
self
.
_slice_h2p
idat
=
dobj
.
local_data
(
tval
)
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
)
if
distaxis
!=
dobj
.
distaxis
(
x
.
val
):
odat
=
dobj
.
redistribute
(
odat
,
dist
=
dobj
.
distaxis
(
x
.
val
))
return
Field
(
tdom
,
odat
)
@
property
def
domain
(
self
):
return
self
.
_domain