Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Martin Reinecke
pypocketfft
Commits
4e178a24
Commit
4e178a24
authored
Apr 17, 2019
by
Martin Reinecke
Browse files
more
parent
9dae4cdd
Changes
2
Hide whitespace changes
Inline
Side-by-side
pocketfft.cc
View file @
4e178a24
...
...
@@ -2036,6 +2036,50 @@ template<typename T> void pocketfft_general_c(int ndim, const size_t *shape,
}
}
template
<
typename
T
>
void
pocketfft_general_hartley
(
int
ndim
,
const
size_t
*
shape
,
const
int64_t
*
stride_in
,
const
int64_t
*
stride_out
,
int
nax
,
const
size_t
*
axes
,
const
T
*
data_in
,
T
*
data_out
,
T
fct
)
{
// allocate temporary 1D array storage, if necessary
size_t
tmpsize
=
0
;
for
(
int
iax
=
0
;
iax
<
nax
;
++
iax
)
if
(
shape
[
axes
[
iax
]]
>
tmpsize
)
tmpsize
=
shape
[
axes
[
iax
]];
arr
<
T
>
tdata
(
tmpsize
);
multiarr
a_in
(
ndim
,
shape
,
stride_in
),
a_out
(
ndim
,
shape
,
stride_out
);
unique_ptr
<
pocketfft_r
<
T
>>
plan
;
for
(
int
iax
=
0
;
iax
<
nax
;
++
iax
)
{
int
axis
=
axes
[
iax
];
multi_iter
it_in
(
a_in
,
axis
),
it_out
(
a_out
,
axis
);
if
((
!
plan
)
||
(
it_in
.
length
()
!=
plan
->
length
()))
plan
.
reset
(
new
pocketfft_r
<
T
>
(
it_in
.
length
()));
while
(
!
it_in
.
done
())
{
for
(
size_t
i
=
0
;
i
<
it_in
.
length
();
++
i
)
tdata
[
i
]
=
data_in
[
it_in
.
offset
()
+
i
*
it_in
.
stride
()];
plan
->
forward
((
T
*
)
tdata
.
data
(),
fct
);
// Hartley order
data_out
[
it_out
.
offset
()]
=
tdata
[
0
];
size_t
i
=
1
,
i1
=
1
,
i2
=
it_out
.
length
()
-
1
;
for
(
i
=
1
;
i
<
it_out
.
length
()
-
1
;
i
+=
2
,
++
i1
,
--
i2
)
{
data_out
[
it_out
.
offset
()
+
i1
*
it_out
.
stride
()]
=
tdata
[
i
]
+
tdata
[
i
+
1
];
data_out
[
it_out
.
offset
()
+
i2
*
it_out
.
stride
()]
=
tdata
[
i
]
-
tdata
[
i
+
1
];
}
if
(
i
<
it_out
.
length
())
data_out
[
it_out
.
offset
()
+
i1
*
it_out
.
stride
()]
=
tdata
[
i
];
it_in
.
advance
();
it_out
.
advance
();
}
// after the first dimension, take data from output array
a_in
=
a_out
;
data_in
=
data_out
;
// factor has been applied, use 1 for remaining axes
fct
=
T
(
1
);
}
}
template
<
typename
T
>
void
pocketfft_general_r
(
int
ndim
,
const
size_t
*
shape
,
const
int64_t
*
stride_in
,
const
int64_t
*
stride_out
,
size_t
axis
,
bool
forward
,
const
T
*
data_in
,
T
*
data_out
,
T
fct
)
...
...
@@ -2145,10 +2189,13 @@ auto f64 = py::dtype("float64");
void
check_args
(
const
py
::
array
&
in
,
vector
<
size_t
>
&
axes
)
{
if
(
axes
.
size
()
==
0
)
axes
.
resize
(
in
.
ndim
());
for
(
int
i
=
0
;
i
<
in
.
ndim
();
++
i
)
axes
[
i
]
=
i
;
if
(
axes
.
size
()
==
0
)
throw
runtime_error
(
"no FFT axis specified"
);
if
((
axes
.
size
()
==
1
)
&&
axes
[
0
]
==
10000
)
// indicator for "axes not given"
{
axes
.
resize
(
in
.
ndim
());
for
(
int
i
=
0
;
i
<
in
.
ndim
();
++
i
)
axes
[
i
]
=
i
;
}
if
(
axes
.
size
()
>
size_t
(
in
.
ndim
()))
throw
runtime_error
(
"bad axes argument"
);
for
(
size_t
i
=
0
;
i
<
axes
.
size
();
++
i
)
...
...
@@ -2195,8 +2242,9 @@ py::array fftn(const py::array &a, const vector<size_t> &axes)
py
::
array
ifftn
(
const
py
::
array
&
a
,
const
vector
<
size_t
>
&
axes
)
{
return
execute
(
a
,
axes
,
false
);
}
py
::
array
rfftn
(
const
py
::
array
&
in
,
int
ax
i
s
)
py
::
array
rfftn
(
const
py
::
array
&
in
,
vector
<
size_t
>
ax
e
s
)
{
check_args
(
in
,
axes
);
py
::
array
res
;
vector
<
size_t
>
dims_in
(
in
.
ndim
()),
dims_out
(
in
.
ndim
());
vector
<
int64_t
>
s_i
(
in
.
ndim
()),
s_o
(
in
.
ndim
());
...
...
@@ -2205,7 +2253,7 @@ py::array rfftn(const py::array &in, int axis)
dims_in
[
i
]
=
in
.
shape
(
i
);
dims_out
[
i
]
=
in
.
shape
(
i
);
}
dims_out
[
ax
is
]
=
(
dims_out
[
ax
is
]
>>
1
)
+
1
;
dims_out
[
ax
es
.
back
()
]
=
(
dims_out
[
ax
es
.
back
()
]
>>
1
)
+
1
;
if
(
in
.
dtype
().
is
(
f64
))
res
=
py
::
array_t
<
complex
<
double
>>
(
dims_out
);
else
if
(
in
.
dtype
().
is
(
f32
))
...
...
@@ -2221,35 +2269,95 @@ py::array rfftn(const py::array &in, int axis)
throw
runtime_error
(
"weird strides"
);
}
if
(
in
.
dtype
().
is
(
f64
))
{
pocketfft_general_r2c
<
double
>
(
in
.
ndim
(),
dims_in
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
axis
,
(
const
double
*
)
in
.
data
(),
axes
.
back
(),
(
const
double
*
)
in
.
data
(),
(
cmplx
<
double
>
*
)
res
.
mutable_data
(),
1.
);
pocketfft_general_c
<
double
>
(
res
.
ndim
(),
dims_out
.
data
(),
s_o
.
data
(),
s_o
.
data
(),
axes
.
size
()
-
1
,
axes
.
data
(),
true
,
(
const
cmplx
<
double
>
*
)
res
.
data
(),
(
cmplx
<
double
>
*
)
res
.
mutable_data
(),
1.
);
}
else
{
pocketfft_general_r2c
<
float
>
(
in
.
ndim
(),
dims_in
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
axis
,
(
const
float
*
)
in
.
data
(),
axes
.
back
(),
(
const
float
*
)
in
.
data
(),
(
cmplx
<
float
>
*
)
res
.
mutable_data
(),
1.
);
pocketfft_general_c
<
float
>
(
res
.
ndim
(),
dims_out
.
data
(),
s_o
.
data
(),
s_o
.
data
(),
axes
.
size
()
-
1
,
axes
.
data
(),
true
,
(
const
cmplx
<
float
>
*
)
res
.
data
(),
(
cmplx
<
float
>
*
)
res
.
mutable_data
(),
1.
);
}
return
res
;
}
py
::
array
irfftn
(
const
py
::
array
&
in
,
int
ax
i
s
,
int
o
size
)
py
::
array
irfftn
(
const
py
::
array
&
in
,
vector
<
size_t
>
ax
e
s
,
size_t
last
size
)
{
py
::
array
res
;
if
(
osize
/
2
+
1
!=
in
.
shape
(
axis
))
throw
runtime_error
(
"bad output size"
);
vector
<
size_t
>
dims_in
(
in
.
ndim
()),
dims_out
(
in
.
ndim
());
vector
<
int64_t
>
s_i
(
in
.
ndim
()),
s_o
(
in
.
ndim
());
for
(
int
i
=
0
;
i
<
in
.
ndim
();
++
i
)
check_args
(
in
,
axes
);
py
::
array
inter
;
if
(
axes
.
size
()
>
1
)
// need to do complex transforms
{
dims_in
[
i
]
=
in
.
shape
(
i
);
dims_out
[
i
]
=
in
.
shape
(
i
);
vector
<
size_t
>
axes2
(
axes
.
size
()
-
1
);
for
(
size_t
i
=
0
;
i
<
axes2
.
size
();
++
i
)
axes2
[
i
]
=
axes
[
i
];
inter
=
ifftn
(
in
,
axes2
);
}
dims_out
[
axis
]
=
osize
;
if
(
in
.
dtype
().
is
(
c128
))
else
inter
=
in
;
size_t
axis
=
axes
.
back
();
if
(
lastsize
==
0
)
lastsize
=
2
*
(
inter
.
shape
(
axis
)
/
2
);
if
(
lastsize
/
2
+
1
!=
inter
.
shape
(
axis
))
throw
runtime_error
(
"bad lastsize"
);
vector
<
size_t
>
dims_in
(
inter
.
ndim
()),
dims_out
(
inter
.
ndim
());
vector
<
int64_t
>
s_i
(
inter
.
ndim
()),
s_o
(
inter
.
ndim
());
for
(
int
i
=
0
;
i
<
inter
.
ndim
();
++
i
)
{
dims_in
[
i
]
=
inter
.
shape
(
i
);
dims_out
[
i
]
=
inter
.
shape
(
i
);
}
dims_out
[
axis
]
=
lastsize
;
py
::
array
res
;
if
(
inter
.
dtype
().
is
(
c128
))
res
=
py
::
array_t
<
double
>
(
dims_out
);
else
if
(
in
.
dtype
().
is
(
c64
))
else
if
(
in
ter
.
dtype
().
is
(
c64
))
res
=
py
::
array_t
<
float
>
(
dims_out
);
else
throw
runtime_error
(
"unsupported data type"
);
for
(
int
i
=
0
;
i
<
inter
.
ndim
();
++
i
)
{
s_i
[
i
]
=
inter
.
strides
(
i
)
/
inter
.
itemsize
();
if
(
s_i
[
i
]
*
inter
.
itemsize
()
!=
inter
.
strides
(
i
))
throw
runtime_error
(
"weird strides"
);
s_o
[
i
]
=
res
.
strides
(
i
)
/
res
.
itemsize
();
if
(
s_o
[
i
]
*
res
.
itemsize
()
!=
res
.
strides
(
i
))
throw
runtime_error
(
"weird strides"
);
}
if
(
inter
.
dtype
().
is
(
c128
))
pocketfft_general_c2r
<
double
>
(
inter
.
ndim
(),
dims_out
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
axis
,
(
const
cmplx
<
double
>
*
)
inter
.
data
(),
(
double
*
)
res
.
mutable_data
(),
1.
);
else
pocketfft_general_c2r
<
float
>
(
inter
.
ndim
(),
dims_out
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
axis
,
(
const
cmplx
<
float
>
*
)
inter
.
data
(),
(
float
*
)
res
.
mutable_data
(),
1.
);
return
res
;
}
py
::
array
hartley
(
const
py
::
array
&
in
,
vector
<
size_t
>
axes
)
{
check_args
(
in
,
axes
);
py
::
array
res
;
vector
<
size_t
>
dims
(
in
.
ndim
());
vector
<
int64_t
>
s_i
(
in
.
ndim
()),
s_o
(
in
.
ndim
());
for
(
int
i
=
0
;
i
<
in
.
ndim
();
++
i
)
dims
[
i
]
=
in
.
shape
(
i
);
if
(
in
.
dtype
().
is
(
f64
))
res
=
py
::
array_t
<
double
>
(
dims
);
else
if
(
in
.
dtype
().
is
(
f32
))
res
=
py
::
array_t
<
float
>
(
dims
);
else
throw
runtime_error
(
"unsupported data type"
);
for
(
int
i
=
0
;
i
<
in
.
ndim
();
++
i
)
{
s_i
[
i
]
=
in
.
strides
(
i
)
/
in
.
itemsize
();
...
...
@@ -2259,15 +2367,15 @@ py::array irfftn(const py::array &in, int axis, int osize)
if
(
s_o
[
i
]
*
res
.
itemsize
()
!=
res
.
strides
(
i
))
throw
runtime_error
(
"weird strides"
);
}
if
(
in
.
dtype
().
is
(
c128
))
pocketfft_general_
c2r
<
double
>
(
in
.
ndim
(),
dims
_out
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
ax
is
,
(
const
cmplx
<
double
>
*
)
in
.
data
(),
if
(
in
.
dtype
().
is
(
f64
))
pocketfft_general_
hartley
<
double
>
(
in
.
ndim
(),
dims
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
axes
.
size
(),
ax
es
.
data
()
,
(
const
double
*
)
in
.
data
(),
(
double
*
)
res
.
mutable_data
(),
1.
);
else
pocketfft_general_
c2r
<
float
>
(
in
.
ndim
(),
dims
_out
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
ax
is
,
(
const
cmplx
<
float
>
*
)
in
.
data
(),
pocketfft_general_
hartley
<
float
>
(
in
.
ndim
(),
dims
.
data
(),
s_i
.
data
(),
s_o
.
data
(),
axes
.
size
(),
ax
es
.
data
()
,
(
const
float
*
)
in
.
data
(),
(
float
*
)
res
.
mutable_data
(),
1.
);
return
res
;
}
...
...
@@ -2283,9 +2391,10 @@ PYBIND11_MODULE(pypocketfft, m)
{
using
namespace
pybind11
::
literals
;
m
.
def
(
"fftn"
,
&
fftn
,
"a"
_a
,
"axes"
_a
=
vector
<
size_t
>
());
m
.
def
(
"ifftn"
,
&
ifftn
,
"a"
_a
,
"axes"
_a
=
vector
<
size_t
>
());
m
.
def
(
"rfftn"
,
&
rfftn
);
m
.
def
(
"irfftn"
,
&
irfftn
);
m
.
def
(
"fftn"
,
&
fftn
,
"a"
_a
,
"axes"
_a
=
vector
<
size_t
>
(
1
,
10000
));
m
.
def
(
"ifftn"
,
&
ifftn
,
"a"
_a
,
"axes"
_a
=
vector
<
size_t
>
(
1
,
10000
));
m
.
def
(
"rfftn"
,
&
rfftn
,
"a"
_a
,
"axes"
_a
=
vector
<
size_t
>
(
1
,
10000
));
m
.
def
(
"irfftn"
,
&
irfftn
,
"a"
_a
,
"axes"
_a
=
vector
<
size_t
>
(
1
,
10000
),
"lastsize"
_a
=
0
);
m
.
def
(
"hartley"
,
&
hartley
,
"a"
_a
,
"axes"
_a
=
vector
<
size_t
>
(
1
,
10000
));
m
.
def
(
"xtest"
,
&
xtest
,
"a"
_a
,
"s"
_a
=
vector
<
size_t
>
(),
"axes"
_a
=
vector
<
size_t
>
(),
"norm"
_a
=
""
);
}
test.py
View file @
4e178a24
...
...
@@ -6,7 +6,10 @@ from numpy.testing import assert_allclose
pmp
=
pytest
.
mark
.
parametrize
shapes
=
((
10
,),
(
127
,),
(
128
,
128
),
(
128
,
129
),
(
1
,
129
),
(
129
,
1
),
(
32
,
17
,
39
))
shapes1D
=
((
10
,),
(
127
,))
shapes2D
=
((
128
,
128
),
(
128
,
129
),
(
1
,
129
),
(
129
,
1
))
shapes3D
=
((
32
,
17
,
39
),)
shapes
=
shapes1D
+
shapes2D
+
shapes3D
@
pmp
(
"shp"
,
shapes
)
def
test_fftn
(
shp
):
...
...
@@ -16,6 +19,32 @@ def test_fftn(shp):
a
=
a
.
astype
(
np
.
complex64
)
assert_allclose
(
pyfftw
.
interfaces
.
numpy_fft
.
fftn
(
a
),
pypocketfft
.
fftn
(
a
),
atol
=
6e-7
*
vmax
,
rtol
=
0
)
@
pmp
(
"shp"
,
shapes2D
)
@
pmp
(
"axes"
,
((
0
,),(
1
,),(
0
,
1
),(
1
,
0
)))
def
test_fftn2D
(
shp
,
axes
):
a
=
np
.
random
.
rand
(
*
shp
)
-
0.5
+
1j
*
np
.
random
.
rand
(
*
shp
)
-
0.5j
vmax
=
np
.
max
(
np
.
abs
(
pyfftw
.
interfaces
.
numpy_fft
.
fftn
(
a
,
axes
=
axes
)))
assert_allclose
(
pyfftw
.
interfaces
.
numpy_fft
.
fftn
(
a
,
axes
=
axes
),
pypocketfft
.
fftn
(
a
,
axes
=
axes
),
atol
=
2e-15
*
vmax
,
rtol
=
0
)
a
=
a
.
astype
(
np
.
complex64
)
assert_allclose
(
pyfftw
.
interfaces
.
numpy_fft
.
fftn
(
a
,
axes
=
axes
),
pypocketfft
.
fftn
(
a
,
axes
=
axes
),
atol
=
6e-7
*
vmax
,
rtol
=
0
)
@
pmp
(
"shp"
,
shapes
)
def
test_rfftn
(
shp
):
a
=
np
.
random
.
rand
(
*
shp
)
-
0.5
vmax
=
np
.
max
(
np
.
abs
(
pyfftw
.
interfaces
.
numpy_fft
.
fftn
(
a
)))
assert_allclose
(
pyfftw
.
interfaces
.
numpy_fft
.
rfftn
(
a
),
pypocketfft
.
rfftn
(
a
),
atol
=
2e-15
*
vmax
,
rtol
=
0
)
a
=
a
.
astype
(
np
.
float32
)
assert_allclose
(
pyfftw
.
interfaces
.
numpy_fft
.
rfftn
(
a
),
pypocketfft
.
rfftn
(
a
),
atol
=
6e-7
*
vmax
,
rtol
=
0
)
@
pmp
(
"shp"
,
shapes2D
)
@
pmp
(
"axes"
,
((
0
,),(
1
,),(
0
,
1
),(
1
,
0
)))
def
test_rfftn2D
(
shp
,
axes
):
a
=
np
.
random
.
rand
(
*
shp
)
-
0.5
vmax
=
np
.
max
(
np
.
abs
(
pyfftw
.
interfaces
.
numpy_fft
.
rfftn
(
a
,
axes
=
axes
)))
assert_allclose
(
pyfftw
.
interfaces
.
numpy_fft
.
rfftn
(
a
,
axes
=
axes
),
pypocketfft
.
rfftn
(
a
,
axes
=
axes
),
atol
=
2e-15
*
vmax
,
rtol
=
0
)
a
=
a
.
astype
(
np
.
float32
)
assert_allclose
(
pyfftw
.
interfaces
.
numpy_fft
.
rfftn
(
a
,
axes
=
axes
),
pypocketfft
.
rfftn
(
a
,
axes
=
axes
),
atol
=
6e-7
*
vmax
,
rtol
=
0
)
@
pmp
(
"shp"
,
shapes
)
def
test_identity
(
shp
):
a
=
np
.
random
.
rand
(
*
shp
)
-
0.5
+
1j
*
np
.
random
.
rand
(
*
shp
)
-
0.5j
...
...
@@ -25,7 +54,7 @@ def test_identity(shp):
assert_allclose
(
pypocketfft
.
ifftn
(
pypocketfft
.
fftn
(
a
))
/
a
.
size
,
a
,
atol
=
6e-7
*
vmax
,
rtol
=
0
)
@
pmp
(
"shp"
,
shapes
)
def
test_identity_r
(
shp
):
def
x
test_identity_r
(
shp
):
a
=
np
.
random
.
rand
(
*
shp
)
-
0.5
b
=
a
.
astype
(
np
.
float32
)
vmax
=
np
.
max
(
np
.
abs
(
a
))
...
...
@@ -33,3 +62,19 @@ def test_identity_r(shp):
n
=
a
.
shape
[
ax
]
assert_allclose
(
pypocketfft
.
irfftn
(
pypocketfft
.
rfftn
(
a
,
ax
),
ax
)
/
n
,
a
,
atol
=
2e-15
*
vmax
,
rtol
=
0
)
assert_allclose
(
pypocketfft
.
irfftn
(
pypocketfft
.
rfftn
(
b
,
ax
),
ax
)
/
n
,
b
,
atol
=
6e-7
*
vmax
,
rtol
=
0
)
@
pmp
(
"shp"
,
shapes3D
)
def
xtest_hartley
(
shp
):
a
=
np
.
random
.
rand
(
*
shp
)
-
0.5
v1
=
pypocketfft
.
hartley
(
a
)
v2
=
pypocketfft
.
fftn
(
a
.
astype
(
np
.
complex128
))
vmax
=
np
.
max
(
np
.
abs
(
v1
))
v2
=
v2
.
real
+
v2
.
imag
assert_allclose
(
v1
,
v2
,
atol
=
2e-15
*
vmax
,
rtol
=
0
)
@
pmp
(
"shp"
,
shapes3D
)
def
test_hartley_identity
(
shp
):
a
=
np
.
random
.
rand
(
*
shp
)
-
0.5
v1
=
pypocketfft
.
hartley
(
pypocketfft
.
hartley
(
a
))
/
a
.
size
vmax
=
np
.
max
(
np
.
abs
(
a
))
assert_allclose
(
a
,
v1
,
atol
=
2e-15
*
vmax
,
rtol
=
0
)
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment