mirror of
https://github.com/openmc-dev/openmc.git
synced 2026-07-21 14:35:27 -04:00
360ec24b41
Co-authored-by: azim_givron <a.givron@naarea.fr> Co-authored-by: Paul Romano <paul.k.romano@gmail.com> Co-authored-by: GuySten <guyste@post.bgu.ac.il> Co-authored-by: GuySten <62616591+GuySten@users.noreply.github.com>
224 lines
8.1 KiB
Python
224 lines
8.1 KiB
Python
"""
|
|
Initially from Jingang Liang: https://github.com/mit-crpg/vectfit.git
|
|
"""
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from openmc.data.vectfit import evaluate, vectfit
|
|
|
|
|
|
@pytest.fixture
|
|
def ref_poles():
|
|
"""Reference poles for real-pole test."""
|
|
return np.array(
|
|
[
|
|
9.709261771920490e02 + 0.0j,
|
|
-1.120960794075339e03 + 0.0j,
|
|
1.923889557426567e00 + 7.543700246109742e01j,
|
|
1.923889557426567e00 - 7.543700246109742e01j,
|
|
1.159741300380281e02 + 3.595650922556496e-02j,
|
|
1.159741300380281e02 - 3.595650922556496e-02j,
|
|
1.546932165729394e02 + 8.728391144940301e-02j,
|
|
1.546932165729394e02 - 8.728391144940301e-02j,
|
|
2.280349190818197e02 + 2.814037559718684e-01j,
|
|
2.280349190818197e02 - 2.814037559718684e-01j,
|
|
2.313004772627853e02 + 3.004628477692201e-01j,
|
|
2.313004772627853e02 - 3.004628477692201e-01j,
|
|
2.787470098364861e02 + 3.414179169920170e-01j,
|
|
2.787470098364861e02 - 3.414179169920170e-01j,
|
|
3.570711338764254e02 + 4.485587371149193e-01j,
|
|
3.570711338764254e02 - 4.485587371149193e-01j,
|
|
4.701059001346060e02 + 6.598089307174224e-01j,
|
|
4.701059001346060e02 - 6.598089307174224e-01j,
|
|
7.275819506342254e02 + 1.189678974845038e03j,
|
|
7.275819506342254e02 - 1.189678974845038e03j,
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def ref_residues():
|
|
"""Reference residues for real-pole test."""
|
|
return np.array(
|
|
[
|
|
[
|
|
-3.269879776751686e07 + 0.0j,
|
|
1.131087935798761e09 + 0.0j,
|
|
1.634151281869857e04 + 2.251103589277891e05j,
|
|
1.634151281869857e04 - 2.251103589277891e05j,
|
|
3.281792303833561e03 - 1.756079516325274e04j,
|
|
3.281792303833561e03 + 1.756079516325274e04j,
|
|
1.110800880243503e04 - 4.324813594540043e04j,
|
|
1.110800880243503e04 + 4.324813594540043e04j,
|
|
8.812700704117636e04 - 2.256520243571103e05j,
|
|
8.812700704117636e04 + 2.256520243571103e05j,
|
|
5.842090495551535e04 - 1.442159380741478e05j,
|
|
5.842090495551535e04 + 1.442159380741478e05j,
|
|
1.339410514130921e05 - 2.640767909713812e05j,
|
|
1.339410514130921e05 + 2.640767909713812e05j,
|
|
2.211245633333130e05 - 3.222447758311512e05j,
|
|
2.211245633333130e05 + 3.222447758311512e05j,
|
|
4.124430059785149e05 - 4.076023108323907e05j,
|
|
4.124430059785149e05 + 4.076023108323907e05j,
|
|
1.607378314999252e09 - 1.401163320110452e08j,
|
|
1.607378314999252e09 + 1.401163320110452e08j,
|
|
]
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def vector_test_data():
|
|
"""Simple 2-signal test with known poles and residues."""
|
|
Ns = 101
|
|
s = np.linspace(3.0, 7.0, Ns)
|
|
poles = [5.0 + 0.1j, 5.0 - 0.1j]
|
|
residues = [[0.5 - 11.0j, 0.5 + 11.0j], [1.5 - 20.0j, 1.5 + 20.0j]]
|
|
f = np.zeros((2, Ns))
|
|
for i in range(2):
|
|
f[i, :] = np.real(
|
|
residues[i][0] / (s - poles[0]) + residues[i][1] / (s - poles[1])
|
|
)
|
|
weight = 1.0 / f
|
|
init_poles = [3.5 + 0.035j, 3.5 - 0.035j]
|
|
return s, poles, residues, f, weight, init_poles
|
|
|
|
|
|
@pytest.fixture
|
|
def poly_test_data():
|
|
"""Test data with rational function plus polynomial terms."""
|
|
Ns = 201
|
|
s = np.linspace(0.0, 5.0, Ns)
|
|
poles = [-20.0 + 30.0j, -20.0 - 30.0j]
|
|
residues = [[5.0 + 10.0j, 5.0 - 10.0j]]
|
|
polys = [[1.0, 2.0, 0.3]]
|
|
f = evaluate(s, poles, residues, polys)
|
|
weight = 1.0 / f
|
|
init_poles = [2.5 + 0.025j, 2.5 - 0.025j]
|
|
return s, poles, residues, polys, f, weight, init_poles
|
|
|
|
|
|
@pytest.fixture
|
|
def real_poles_data(ref_poles, ref_residues):
|
|
"""Large-scale signal using complex and real poles."""
|
|
Ns = 2000
|
|
s = np.linspace(1.0e-2, 5.0e3, Ns)
|
|
f = np.zeros((1, Ns))
|
|
for p, r in zip(ref_poles, ref_residues[0]):
|
|
f[0] += (r / (s - p)).real
|
|
weight = 1.0 / f
|
|
poles = np.linspace(1.1e-2, 4.8e3, 10)
|
|
poles = poles + poles * 0.01j
|
|
poles = np.sort(np.append(poles, np.conj(poles)))
|
|
return s, f, weight, poles
|
|
|
|
|
|
@pytest.fixture
|
|
def large_test_data():
|
|
"""Stress test data with thousands of poles and samples."""
|
|
Ns = 3000
|
|
N = 200
|
|
s = np.linspace(1.0e-2, 5.0e3, Ns)
|
|
poles = np.linspace(1.1e-2, 4.8e3, N // 2) + 0.01j * np.linspace(
|
|
1.1e-2, 4.8e3, N // 2
|
|
)
|
|
poles = np.sort(np.append(poles, np.conj(poles)))
|
|
residues = np.linspace(1e2, 1e6, N // 2) + 0.5j * np.linspace(1e2, 1e6, N // 2)
|
|
residues = np.sort(np.append(residues, np.conj(residues))).reshape((1, N))
|
|
f = np.zeros((1, Ns))
|
|
for p, r in zip(poles, residues[0]):
|
|
f[0] += (r / (s - p)).real
|
|
weight = 1.0 / f
|
|
init_poles = np.linspace(1.2e-2, 4.7e3, N // 2) + 0.01j * np.linspace(
|
|
1.2e-2, 4.7e3, N // 2
|
|
)
|
|
init_poles = np.sort(np.append(init_poles, np.conj(init_poles)))
|
|
return s, f, weight, init_poles
|
|
|
|
|
|
@pytest.fixture
|
|
def eval_test_data():
|
|
"""Reference data for evaluating rational + polynomial models."""
|
|
Ns = 101
|
|
s = np.linspace(-5.0, 5.0, Ns)
|
|
poles = [-2.0 + 30.0j, -2.0 - 30.0j]
|
|
residues = [5.0 + 10.0j, 5.0 - 10.0j]
|
|
polys = [1.0, 2.0, 0.3]
|
|
return s, poles, residues, polys
|
|
|
|
|
|
def test_vector(vector_test_data):
|
|
"""Test vectfit with vector samples and simple poles.
|
|
It is expected to get exact results with one iteration.
|
|
"""
|
|
s, expected_poles, expected_residues, f, weight, init_poles = vector_test_data
|
|
poles, residues, _, fit, _ = vectfit(f, s, init_poles, weight)
|
|
assert np.allclose(
|
|
np.sort_complex(poles), np.sort_complex(expected_poles), rtol=1e-7
|
|
)
|
|
assert np.allclose(f, evaluate(s, poles, residues), rtol=1e-7)
|
|
assert np.allclose(f, fit, rtol=1e-5)
|
|
|
|
|
|
def test_poly(poly_test_data):
|
|
"""Test vectfit with polynomials."""
|
|
s, expected_poles, expected_residues, expected_polys, f, weight, init_poles = (
|
|
poly_test_data
|
|
)
|
|
poles, residues, cf, fit, _ = vectfit(f, s, init_poles, weight, n_polys=3)
|
|
poles, residues, cf, fit, _ = vectfit(f, s, poles, weight, n_polys=3)
|
|
assert np.allclose(
|
|
np.sort_complex(poles), np.sort_complex(expected_poles), rtol=1e-5
|
|
)
|
|
assert np.allclose(f, evaluate(s, poles, residues, cf), rtol=1e-5)
|
|
assert np.allclose(cf, expected_polys, rtol=1e-5)
|
|
assert np.allclose(f, fit, rtol=1e-4)
|
|
|
|
|
|
def test_real_poles(real_poles_data, ref_poles, ref_residues):
|
|
"""Test vectfit with more poles including real poles"""
|
|
s, f, weight, poles = real_poles_data
|
|
for _ in range(6):
|
|
poles, residues, _, fit, _ = vectfit(f, s, poles, weight)
|
|
assert np.allclose(
|
|
np.sort_complex(poles), np.sort_complex(ref_poles), rtol=1e-5, atol=1e-8
|
|
)
|
|
assert np.allclose(f, evaluate(s, poles, residues), rtol=1e-4)
|
|
assert np.allclose(f, fit, rtol=1e-3)
|
|
|
|
|
|
def test_large(large_test_data):
|
|
"""Test vectfit with a large set of poles and samples"""
|
|
s, f, weight, init_poles = large_test_data
|
|
poles_fit, residues_fit, _, f_fit, _ = vectfit(f, s, init_poles, weight)
|
|
assert np.allclose(f, f_fit, rtol=1e-3)
|
|
|
|
|
|
def test_evaluate(eval_test_data):
|
|
"""Test evaluate function"""
|
|
s, poles, residues, polys = eval_test_data
|
|
|
|
# Single signal, no polynomial
|
|
f_ref = np.real(residues[0] / (s - poles[0]) + residues[1] / (s - poles[1]))
|
|
f = evaluate(s, poles, residues)
|
|
assert np.allclose(f[0], f_ref)
|
|
|
|
# Single signal, with polynomial
|
|
for n, c in enumerate(polys):
|
|
f_ref += c * np.power(s, n)
|
|
f = evaluate(s, poles, residues, polys)
|
|
assert np.allclose(f[0], f_ref)
|
|
|
|
# Multi-signal, multi-residue, multi-poly
|
|
poles = [5.0 + 0.1j, 5.0 - 0.1j]
|
|
residues = [[0.5 - 11.0j, 0.5 + 11.0j], [1.5 - 20.0j, 1.5 + 20.0j]]
|
|
polys = [[1.0, 2.0, 0.3], [4.0, -2.0, -10.0]]
|
|
f_ref = np.zeros((2, len(s)))
|
|
for i in range(2):
|
|
f_ref[i, :] = np.real(
|
|
residues[i][0] / (s - poles[0]) + residues[i][1] / (s - poles[1])
|
|
)
|
|
for n, c in enumerate(polys[i]):
|
|
f_ref[i, :] += c * np.power(s, n)
|
|
f = evaluate(s, poles, residues, polys)
|
|
assert np.allclose(f, f_ref)
|