def test_ffsn_ref(self):
T = [1, 1]
T_c = [0, 0]
N_FS = [3, 3]
N_s = [4, 3]
for mod in AVAILABLE_MOD:
# Sample the kernel and do the transform.
sample_points, _ = ffsn_sample(T=T,
N_FS=N_FS,
T_c=T_c,
N_s=N_s,
mod=mod)
diric_samples = dirichlet_2D(sample_points=sample_points,
T=T,
T_c=T_c,
N_FS=N_FS)
diric_FS = _ffsn(x=diric_samples, T=T, N_FS=N_FS, T_c=T_c)
# Compare with theoretical result.
diric_FS_exact = mod.outer(
dirichlet_fs(N_FS[0], T[0], T_c[0], mod=mod),
dirichlet_fs(N_FS[1], T[1], T_c[1], mod=mod),
)
assert mod.allclose(diric_FS[:N_FS[0], :N_FS[1]], diric_FS_exact)
# Inverse transform.
diric_samples_recov = _iffsn(x_FS=diric_FS,
T=T,
T_c=T_c,
N_FS=N_FS)
# Compare with original samples.
assert mod.allclose(diric_samples, diric_samples_recov)
def test_ffsn_sample(self):
N_s = [4, 3]
for mod in AVAILABLE_MOD:
sample_points, idx = ffsn_sample(T=[1, 1],
N_FS=[3, 3],
T_c=[0, 0],
N_s=N_s,
mod=mod)
# check sample points
assert sample_points[0].shape == (N_s[0], 1)
assert sample_points[1].shape == (1, N_s[1])
mod.testing.assert_array_equal(
sample_points[0][:, 0],
mod.array([0.125, 0.375, -0.375, -0.125]))
mod.testing.assert_array_equal(sample_points[1][0, :],
mod.array([0, 1 / 3, -1 / 3]))
# check index values
assert idx[0].shape == (N_s[0], 1)
assert idx[1].shape == (1, N_s[1])
mod.testing.assert_array_equal(idx[0][:, 0],
mod.array([2, 3, 0, 1]))
mod.testing.assert_array_equal(idx[1][0, :], mod.array([1, 2, 0]))
assert all([get_array_module(s) == mod for s in sample_points])
assert all([get_array_module(i) == mod for i in idx])
def test_ffsn_axes(self):
T_x = T_y = 1
T_cx = T_cy = 0
N_FSx = N_FSy = 3
N_sx, N_sy = 4, 3
for mod in AVAILABLE_MOD:
# Sample the kernel.
sample_points, _ = ffsn_sample(T=[T_x, T_y],
N_FS=[N_FSx, N_FSy],
T_c=[T_cx, T_cy],
N_s=[N_sx, N_sy],
mod=mod)
diric_samples = dirichlet_2D(sample_points=sample_points,
T=[T_x, T_y],
T_c=[T_cx, T_cy],
N_FS=[N_FSx, N_FSy])
# Add new dimension.
diric_samples = diric_samples[:, mod.newaxis, :]
axes = (0, 2)
# Perform transform.
diric_FS = ffsn(x=diric_samples,
T=[T_x, T_y],
T_c=[T_cx, T_cy],
N_FS=[N_FSx, N_FSy],
axes=axes)
# Compare with theoretical result.
diric_FS_exact = mod.outer(dirichlet_fs(N_FSx, T_x, T_cx, mod=mod),
dirichlet_fs(N_FSy, T_y, T_cy, mod=mod))
assert mod.allclose(diric_FS[:N_FSx, 0, :N_FSy], diric_FS_exact)
# Inverse transform.
diric_samples_recov = iffsn(x_FS=diric_FS,
T=[T_x, T_y],
T_c=[T_cx, T_cy],
N_FS=[N_FSx, N_FSy],
axes=axes)
# Compare with original samples.
assert mod.allclose(diric_samples, diric_samples_recov)
def test_ffsn_sample_shape(self):
for mod in AVAILABLE_MOD:
D = 5
T = mod.ones(D)
N_FS = mod.arange(D) * 2 + 3
T_c = mod.zeros(D)
N_s = N_FS
sample_points, idx = ffsn_sample(T=T,
N_FS=N_FS,
T_c=T_c,
N_s=N_s,
mod=mod)
# check shape
for d in range(D):
sh = [1] * D
sh[d] = N_s[d]
assert list(sample_points[d].shape) == sh
assert list(idx[d].shape) == sh
assert all([get_array_module(s) == mod for s in sample_points])
assert all([get_array_module(i) == mod for i in idx])
def profile_ffsn(n_trials):
print(f"\nCOMPARING FFSN APPROACHES WITH {n_trials} TRIALS")
T = [1, 1]
T_c = [0, 0]
N_FS_vals = [11, 31, 101, 301, 1001, 3001, 10001]
n_std = 1
func = {"ffsn_fftn": ffsn, "ffsn_ref": _ffsn}
proc_time = dict()
proc_time_std = dict()
for _N_FS in N_FS_vals:
print("\nN_FS : {}".format(_N_FS))
N_FS = [_N_FS, _N_FS]
proc_time[_N_FS] = dict()
proc_time_std[_N_FS] = dict()
# Loop through modules
for mod in AVAILABLE_MOD:
backend = get_module_name(mod)
# fastest FFT length depends on module
_N_s = next_fast_len(_N_FS, mod=mod)
N_s = [_N_s, _N_s]
print("-- module : {}, Length-{} FFT".format(backend, N_s))
for _f in func:
# Sample the kernel and do the transform.
sample_points, _ = ffsn_sample(T=T,
N_FS=N_FS,
T_c=T_c,
N_s=N_s,
mod=mod)
diric_samples = dirichlet_2D(sample_points=sample_points,
T=T,
T_c=T_c,
N_FS=N_FS)
_key = "{}_{}".format(_f, backend)
timings = []
for _ in range(n_trials):
start_time = time.time()
func[_f](x=diric_samples, T=T, N_FS=N_FS, T_c=T_c)
timings.append(time.time() - start_time)
proc_time[_N_FS][_key] = np.mean(timings)
proc_time_std[_N_FS][_key] = np.std(timings)
print("{} : {} seconds".format(_f, proc_time[_N_FS][_key]))
# plot results
fig, ax = plt.subplots()
util.comparison_plot(proc_time, proc_time_std, n_std, ax)
ax.set_xlabel("Number of FS coefficients")
fig.tight_layout()
fname = pathlib.Path(__file__).resolve().parent / "profile_ffsn_2d.png"
fig.savefig(fname, dpi=300)