Skip to content

Commit fa61aa0

Browse files
vitalsongvitalsong
authored
Add fftshift, circshift, sqrt functions. (#74)
1 parent 59dc45c commit fa61aa0

File tree

9 files changed

+201
-67
lines changed

9 files changed

+201
-67
lines changed

‎CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
cmake_minimum_required(VERSION 3.10)
2-
project(dsplib LANGUAGES CXX VERSION 0.54.5)
2+
project(dsplib LANGUAGES CXX VERSION 0.54.6)
33

44
set(CMAKE_CXX_STANDARD 17)
55
set(CMAKE_CXX_STANDARD_REQUIRED ON)

‎include/dsplib/math.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,11 @@ cmplx_t power(cmplx_t x, int n);
161161
arr_real power(const arr_real& x, int n);
162162
arr_cmplx power(const arr_cmplx& x, int n);
163163

164+
//square root (only positive values)
165+
//TODO: add complex result for negative or complex input
166+
real_t sqrt(real_t x);
167+
arr_real sqrt(const arr_real& x);
168+
164169
//array log
165170
arr_real log(const arr_real& arr);
166171
arr_real log2(const arr_real& arr);

‎include/dsplib/utils.h‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -232,4 +232,12 @@ PeakList findpeaks(arr_real data, int npeaks);
232232
// generate linearly spaced vector
233233
arr_real linspace(real_t x1, real_t x2, size_t n);
234234

235+
// shift array circularly
236+
arr_real circshift(const arr_real& x, int k);
237+
arr_cmplx circshift(const arr_cmplx& x, int k);
238+
239+
//shift zero-frequency component to center of spectrum (swaps the left and right halves of X)
240+
arr_real fftshift(const arr_real& x);
241+
arr_cmplx fftshift(const arr_cmplx& x);
242+
235243
} // namespace dsplib

‎lib/fft/factory.cpp‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
#include <dsplib/fft.h>
2+
#include <dsplib/math.h>
3+
#include <dsplib/utils.h>
4+
5+
#include "fft/fact-fft.h"
6+
#include "fft/primes-fft.h"
7+
#include "fft/factory.h"
8+
#include "fft/pow2-fft.h"
9+
#include "fft/real-fft.h"
10+
#include "fft/small-fft.h"
11+
#include "lru-cache.h"
12+
13+
#include <cassert>
14+
#include <memory>
15+
16+
namespace dsplib {
17+
18+
namespace {
19+
20+
constexpr int FFT_CACHE_SIZE = DSPLIB_FFT_CACHE_SIZE;
21+
22+
static_assert(FFT_CACHE_SIZE > 0);
23+
24+
std::shared_ptr<BaseFftPlanC> _get_fft_plan(int n) {
25+
if (isprime(n)) {
26+
return std::make_shared<PrimesFftC>(n);
27+
}
28+
if (ispow2(n)) {
29+
return std::make_shared<Pow2FftPlan>(n);
30+
}
31+
return std::make_shared<FactorFFTPlan>(n);
32+
}
33+
34+
std::shared_ptr<BaseFftPlanR> _get_rfft_plan(int n) {
35+
if (isprime(n)) {
36+
return std::make_shared<PrimesFftR>(n);
37+
}
38+
if (n % 2 == 0) {
39+
return std::make_shared<RealFftPlan>(n);
40+
}
41+
return std::make_shared<FactorFFTPlanR>(n);
42+
}
43+
44+
} // namespace
45+
46+
//-------------------------------------------------------------------------------------------------
47+
std::shared_ptr<BaseFftPlanC> create_fft_plan(int n) {
48+
//dont cache small fft plans
49+
if ((n == 1) || (n == 2) || (n == 4) || (n == 8)) {
50+
return std::make_shared<SmallFftPow2C>(n);
51+
}
52+
53+
//TODO: use weak_ptr cache to prevent duplication
54+
thread_local LRUCache<int, std::shared_ptr<BaseFftPlanC>> cache{FFT_CACHE_SIZE};
55+
if (!cache.exists(n)) {
56+
auto plan = _get_fft_plan(n);
57+
cache.put(n, plan);
58+
return plan;
59+
}
60+
return cache.get(n);
61+
}
62+
63+
std::shared_ptr<BaseFftPlanR> create_rfft_plan(int n) {
64+
if ((n == 1) || (n == 2) || (n == 4) || (n == 8)) {
65+
return std::make_shared<SmallFftPow2R>(n);
66+
}
67+
68+
thread_local LRUCache<int, std::shared_ptr<BaseFftPlanR>> cache{FFT_CACHE_SIZE};
69+
if (!cache.exists(n)) {
70+
auto plan = _get_rfft_plan(n);
71+
cache.put(n, plan);
72+
return plan;
73+
}
74+
return cache.get(n);
75+
}
76+
77+
} // namespace dsplib

‎lib/fft/fft.cpp‎

Lines changed: 0 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -2,78 +2,12 @@
22
#include <dsplib/math.h>
33
#include <dsplib/utils.h>
44

5-
#include "fft/fact-fft.h"
6-
#include "fft/primes-fft.h"
75
#include "fft/factory.h"
8-
#include "fft/pow2-fft.h"
9-
#include "fft/real-fft.h"
10-
#include "fft/small-fft.h"
11-
#include "lru-cache.h"
126

137
#include <cassert>
14-
#include <memory>
158

169
namespace dsplib {
1710

18-
namespace {
19-
20-
constexpr int FFT_CACHE_SIZE = DSPLIB_FFT_CACHE_SIZE;
21-
22-
static_assert(FFT_CACHE_SIZE > 0);
23-
24-
std::shared_ptr<BaseFftPlanC> _get_fft_plan(int n) {
25-
if (isprime(n)) {
26-
return std::make_shared<PrimesFftC>(n);
27-
}
28-
if (ispow2(n)) {
29-
return std::make_shared<Pow2FftPlan>(n);
30-
}
31-
return std::make_shared<FactorFFTPlan>(n);
32-
}
33-
34-
std::shared_ptr<BaseFftPlanR> _get_rfft_plan(int n) {
35-
if (isprime(n)) {
36-
return std::make_shared<PrimesFftR>(n);
37-
}
38-
if (n % 2 == 0) {
39-
return std::make_shared<RealFftPlan>(n);
40-
}
41-
return std::make_shared<FactorFFTPlanR>(n);
42-
}
43-
44-
} // namespace
45-
46-
//-------------------------------------------------------------------------------------------------
47-
std::shared_ptr<BaseFftPlanC> create_fft_plan(int n) {
48-
//dont cache small fft plans
49-
if ((n == 1) || (n == 2) || (n == 4) || (n == 8)) {
50-
return std::make_shared<SmallFftPow2C>(n);
51-
}
52-
53-
//TODO: use weak_ptr cache to prevent duplication
54-
thread_local LRUCache<int, std::shared_ptr<BaseFftPlanC>> cache{FFT_CACHE_SIZE};
55-
if (!cache.exists(n)) {
56-
auto plan = _get_fft_plan(n);
57-
cache.put(n, plan);
58-
return plan;
59-
}
60-
return cache.get(n);
61-
}
62-
63-
std::shared_ptr<BaseFftPlanR> create_rfft_plan(int n) {
64-
if ((n == 1) || (n == 2) || (n == 4) || (n == 8)) {
65-
return std::make_shared<SmallFftPow2R>(n);
66-
}
67-
68-
thread_local LRUCache<int, std::shared_ptr<BaseFftPlanR>> cache{FFT_CACHE_SIZE};
69-
if (!cache.exists(n)) {
70-
auto plan = _get_rfft_plan(n);
71-
cache.put(n, plan);
72-
return plan;
73-
}
74-
return cache.get(n);
75-
}
76-
7711
//-------------------------------------------------------------------------------------------------
7812
FftPlan::FftPlan(int n)
7913
: _d{create_fft_plan(n)} {

‎lib/math.cpp‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -497,6 +497,19 @@ arr_cmplx power(const arr_cmplx& x, int n) {
497497
return _power(x, n);
498498
}
499499

500+
//-------------------------------------------------------------------------------------------------
501+
real_t sqrt(real_t x) {
502+
return std::sqrt(x);
503+
}
504+
505+
arr_real sqrt(const arr_real& x) {
506+
arr_real y(x.size());
507+
for (int i = 0; i < x.size(); ++i) {
508+
y[i] = std::sqrt(x[i]);
509+
}
510+
return y;
511+
}
512+
500513
//-------------------------------------------------------------------------------------------------
501514
arr_real angle(const arr_cmplx& arr) {
502515
arr_real r(arr.size());

‎lib/utils.cpp‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,4 +224,47 @@ arr_real linspace(real_t x1, real_t x2, size_t n) {
224224
return out;
225225
}
226226

227+
template<typename T>
228+
static base_array<T> _circshift(const base_array<T>& x, int k) {
229+
const int n = x.size();
230+
if (n == 1 || k == 0) {
231+
return x;
232+
}
233+
base_array<T> y(n);
234+
for (size_t i = 0; i < n; ++i) {
235+
const size_t p = (i + k + n) % n;
236+
y[i] = x[p];
237+
}
238+
return y;
239+
}
240+
241+
arr_real circshift(const arr_real& x, int k) {
242+
return _circshift(x, -k);
243+
}
244+
245+
arr_cmplx circshift(const arr_cmplx& x, int k) {
246+
return _circshift(x, -k);
247+
}
248+
249+
template<typename T>
250+
static base_array<T> _fftshift(const base_array<T>& x) {
251+
const int n = x.size();
252+
if (n == 1) {
253+
return x;
254+
}
255+
const int n2 = n / 2;
256+
dsplib::base_array<T> y(x);
257+
y.slice(0, n2) = x.slice(n - n2, n);
258+
y.slice(n2, n) = x.slice(0, n - n2);
259+
return y;
260+
}
261+
262+
arr_real fftshift(const arr_real& x) {
263+
return _fftshift(x);
264+
}
265+
266+
arr_cmplx fftshift(const arr_cmplx& x) {
267+
return _fftshift(x);
268+
}
269+
227270
} // namespace dsplib

‎tests/math_test.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -519,3 +519,9 @@ TEST(MathTest, Log2) {
519519
arr_real r = {0, 1, 1.58496250072116, 2, 2.32192809488736};
520520
ASSERT_EQ_ARR_REAL(y, r);
521521
}
522+
523+
TEST(MathTest, Sqrt) {
524+
arr_real x = {0, 1, 2, 3, 4, 5};
525+
arr_real y = dsplib::sqrt(abs2(x));
526+
ASSERT_EQ_ARR_REAL(y, x);
527+
}

‎tests/utils_test.cpp‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,4 +245,52 @@ TEST(Utils, Linspace) {
245245
const arr_real x = linspace(-5, 5, 2);
246246
ASSERT_EQ_ARR_REAL(x, arr_real{-5, 5}, 1e-3);
247247
}
248+
}
249+
250+
//-------------------------------------------------------------------------------------------------
251+
TEST(Utils, Circshift) {
252+
{
253+
arr_real x1 = {1};
254+
arr_real x2 = circshift(x1, 1);
255+
ASSERT_EQ_ARR_REAL(x2, arr_real{1});
256+
}
257+
{
258+
arr_real x1 = {1, 2, 3, 4, 5};
259+
arr_real x2 = circshift(x1, 0);
260+
ASSERT_EQ_ARR_REAL(x2, arr_real{1, 2, 3, 4, 5});
261+
}
262+
{
263+
arr_real x1 = {1, 2, 3, 4, 5};
264+
arr_real x2 = circshift(x1, 1);
265+
ASSERT_EQ_ARR_REAL(x2, arr_real{5, 1, 2, 3, 4});
266+
}
267+
{
268+
arr_real x1 = {1, 2, 3, 4, 5};
269+
arr_real x2 = circshift(x1, -1);
270+
ASSERT_EQ_ARR_REAL(x2, arr_real{2, 3, 4, 5, 1});
271+
}
272+
}
273+
274+
//-------------------------------------------------------------------------------------------------
275+
TEST(Utils, FftShift) {
276+
{
277+
arr_real x1 = {1, 2, 3, 4, 5};
278+
arr_real x2 = fftshift(x1);
279+
ASSERT_EQ_ARR_REAL(x2, arr_real{4, 5, 1, 2, 3});
280+
}
281+
{
282+
arr_real x1 = {1, 2, 3, 4, 5, 6};
283+
arr_real x2 = fftshift(x1);
284+
ASSERT_EQ_ARR_REAL(x2, arr_real{4, 5, 6, 1, 2, 3});
285+
}
286+
{
287+
arr_real x1 = {1, 2};
288+
arr_real x2 = fftshift(x1);
289+
ASSERT_EQ_ARR_REAL(x2, arr_real{2, 1});
290+
}
291+
{
292+
arr_real x1 = {1};
293+
arr_real x2 = fftshift(x1);
294+
ASSERT_EQ_ARR_REAL(x2, arr_real{1});
295+
}
248296
}

0 commit comments

Comments
 (0)