mirror of
https://github.com/openmc-dev/openmc.git
synced 2026-07-27 13:45:36 -04:00
Implement vector fitting to replace external vectfit package (#3493)
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>
This commit is contained in:
azim-givron
GitHub
parent
44da7022b8
commit
360ec24b41
8 changed files with 1102 additions and 92 deletions
14
.github/workflows/ci.yml
vendored
14
.github/workflows/ci.yml
vendored
|
|
@ -49,7 +49,6 @@ jobs:
|
|||
dagmc: [n]
|
||||
libmesh: [n]
|
||||
event: [n]
|
||||
vectfit: [n]
|
||||
|
||||
include:
|
||||
- python-version: "3.12"
|
||||
|
|
@ -74,14 +73,9 @@ jobs:
|
|||
python-version: "3.11"
|
||||
omp: y
|
||||
mpi: n
|
||||
- vectfit: y
|
||||
python-version: "3.11"
|
||||
omp: n
|
||||
mpi: y
|
||||
name: "Python ${{ matrix.python-version }} (omp=${{ matrix.omp }},
|
||||
mpi=${{ matrix.mpi }}, dagmc=${{ matrix.dagmc }},
|
||||
libmesh=${{ matrix.libmesh }}, event=${{ matrix.event }}
|
||||
vectfit=${{ matrix.vectfit }})"
|
||||
libmesh=${{ matrix.libmesh }}, event=${{ matrix.event }}"
|
||||
|
||||
env:
|
||||
MPI: ${{ matrix.mpi }}
|
||||
|
|
@ -89,7 +83,6 @@ jobs:
|
|||
OMP: ${{ matrix.omp }}
|
||||
DAGMC: ${{ matrix.dagmc }}
|
||||
EVENT: ${{ matrix.event }}
|
||||
VECTFIT: ${{ matrix.vectfit }}
|
||||
LIBMESH: ${{ matrix.libmesh }}
|
||||
NPY_DISABLE_CPU_FEATURES: "AVX512F AVX512_SKX"
|
||||
OPENBLAS_NUM_THREADS: 1
|
||||
|
|
@ -150,11 +143,6 @@ jobs:
|
|||
sudo update-alternatives --set mpirun /usr/bin/mpirun.mpich
|
||||
sudo update-alternatives --set mpi-x86_64-linux-gnu /usr/include/x86_64-linux-gnu/mpich
|
||||
|
||||
- name: Optional apt dependencies for vectfit
|
||||
shell: bash
|
||||
if: ${{ matrix.vectfit == 'y' }}
|
||||
run: sudo apt install -y libblas-dev liblapack-dev
|
||||
|
||||
- name: install
|
||||
shell: bash
|
||||
run: |
|
||||
|
|
|
|||
|
|
@ -33,5 +33,6 @@ from .resonance_covariance import *
|
|||
from .multipole import *
|
||||
from .grid import *
|
||||
from .function import *
|
||||
from .vectfit import *
|
||||
|
||||
from .effective_dose.dose import dose_coefficients
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ from . import WMP_VERSION, WMP_VERSION_MAJOR
|
|||
from .data import K_BOLTZMANN
|
||||
from .neutron import IncidentNeutron
|
||||
from .resonance import ResonanceRange
|
||||
|
||||
from .vectfit import vectfit, evaluate
|
||||
|
||||
# Constants that determine which value to access
|
||||
_MP_EA = 0 # Pole
|
||||
|
|
@ -174,10 +174,6 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
|
|||
(poles, residues)
|
||||
|
||||
"""
|
||||
|
||||
# import vectfit package: https://github.com/liangjg/vectfit
|
||||
import vectfit as vf
|
||||
|
||||
ne = energy.size
|
||||
nmt = len(mts)
|
||||
if ce_xs.shape != (nmt, ne):
|
||||
|
|
@ -194,8 +190,8 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
|
|||
test_xs_ref[i] = np.interp(test_energy, energy, ce_xs[i])
|
||||
|
||||
if log:
|
||||
print(f" energy: {energy[0]:.3e} to {energy[-1]:.3e} eV ({ne} points)")
|
||||
print(f" error tolerance: rtol={rtol}, atol={atol}")
|
||||
print(f"\tenergy: {energy[0]:.3e} to {energy[-1]:.3e} eV ({ne} points)")
|
||||
print(f"\terror tolerance: rtol={rtol}, atol={atol}")
|
||||
|
||||
# transform xs (sigma) and energy (E) to f (sigma*E) and s (sqrt(E)) to be
|
||||
# compatible with the multipole representation
|
||||
|
|
@ -251,7 +247,7 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
|
|||
print(f"VF iteration {i_vf + 1}/{n_vf_iter}")
|
||||
|
||||
# call vf
|
||||
poles, residues, cf, f_fit, rms = vf.vectfit(f, s, poles, weight)
|
||||
poles, residues, *_ = vectfit(f, s, poles, weight)
|
||||
|
||||
# convert real pole to conjugate pairs
|
||||
n_real_poles = 0
|
||||
|
|
@ -268,11 +264,11 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
|
|||
if n_real_poles > 0:
|
||||
if log >= DETAILED_LOGGING:
|
||||
print(f" # real poles: {n_real_poles}")
|
||||
new_poles, residues, cf, f_fit, rms = \
|
||||
vf.vectfit(f, s, new_poles, weight, skip_pole=True)
|
||||
new_poles, residues, *_ = \
|
||||
vectfit(f, s, new_poles, weight, skip_pole_update=True)
|
||||
|
||||
# assess the result on test grid
|
||||
test_xs = vf.evaluate(test_s, new_poles, residues) / test_energy
|
||||
test_xs = evaluate(test_s, new_poles, residues) / test_energy
|
||||
abserr = np.abs(test_xs - test_xs_ref)
|
||||
with np.errstate(invalid='ignore', divide='ignore'):
|
||||
relerr = abserr / test_xs_ref
|
||||
|
|
@ -388,9 +384,9 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
|
|||
|
||||
return (mp_poles, mp_residues)
|
||||
|
||||
|
||||
def vectfit_nuclide(endf_file, njoy_error=5e-4, vf_pieces=None,
|
||||
log=False, path_out=None, mp_filename=None, **kwargs):
|
||||
log=False, path_out=None, mp_filename=None,
|
||||
**kwargs):
|
||||
r"""Generate multipole data for a nuclide from ENDF.
|
||||
|
||||
Parameters
|
||||
|
|
@ -571,10 +567,6 @@ def _windowing(mp_data, n_cf, rtol=1e-3, atol=1e-5, n_win=None, spacing=None,
|
|||
format.
|
||||
|
||||
"""
|
||||
|
||||
# import vectfit package: https://github.com/liangjg/vectfit
|
||||
import vectfit as vf
|
||||
|
||||
# unpack multipole data
|
||||
name = mp_data["name"]
|
||||
awr = mp_data["AWR"]
|
||||
|
|
@ -645,7 +637,7 @@ def _windowing(mp_data, n_cf, rtol=1e-3, atol=1e-5, n_win=None, spacing=None,
|
|||
|
||||
# reference xs from multipole form, note the residue terms in the
|
||||
# multipole and vector fitting representations differ by a 1j
|
||||
xs_ref = vf.evaluate(energy_sqrt, poles, residues*1j) / energy
|
||||
xs_ref = evaluate(energy_sqrt, poles, residues*1j) / energy
|
||||
|
||||
# curve fit matrix
|
||||
matrix = np.vstack([energy**(0.5*i - 1) for i in range(n_cf + 1)]).T
|
||||
|
|
@ -659,7 +651,7 @@ def _windowing(mp_data, n_cf, rtol=1e-3, atol=1e-5, n_win=None, spacing=None,
|
|||
|
||||
# calculate the cross sections contributed by the windowed poles
|
||||
if rp > lp:
|
||||
xs_wp = vf.evaluate(energy_sqrt, poles[lp:rp],
|
||||
xs_wp = evaluate(energy_sqrt, poles[lp:rp],
|
||||
residues[:, lp:rp]*1j) / energy
|
||||
else:
|
||||
xs_wp = np.zeros_like(xs_ref)
|
||||
|
|
@ -1054,7 +1046,15 @@ class WindowedMultipole(EqualityMixin):
|
|||
return cls.from_multipole(mp_data, **wmp_options)
|
||||
|
||||
@classmethod
|
||||
def from_multipole(cls, mp_data, search=None, log=False, **kwargs):
|
||||
def from_multipole(
|
||||
cls,
|
||||
mp_data,
|
||||
search=None,
|
||||
log=False,
|
||||
search_n_win=20,
|
||||
search_cf_orders=None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Generate windowed multipole neutron data from multipole data.
|
||||
|
||||
Parameters
|
||||
|
|
@ -1066,8 +1066,14 @@ class WindowedMultipole(EqualityMixin):
|
|||
Defaults to True if no windowing parameters are specified.
|
||||
log : bool or int, optional
|
||||
Whether to print running logs (use int for verbosity control)
|
||||
search_n_win : int, optional
|
||||
Number of window sizes to consider in the search grid when
|
||||
``search`` is True.
|
||||
search_cf_orders : iterable of int, optional
|
||||
Curve-fit orders to consider in the search grid when ``search`` is
|
||||
True. Defaults to integers from 10 down to 2.
|
||||
**kwargs
|
||||
Keyword arguments passed to :func:`openmc.data.multipole._windowing`
|
||||
Keyword arguments passed to :func:`openmc.data.multipole._windowing`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
|
|
@ -1098,12 +1104,17 @@ class WindowedMultipole(EqualityMixin):
|
|||
# search optimal WMP from a range of window sizes and CF orders
|
||||
if log:
|
||||
print("Start searching ...")
|
||||
if search_cf_orders is None:
|
||||
search_cf_orders = range(10, 1, -1)
|
||||
|
||||
n_poles = sum([p.size for p in mp_data["poles"]])
|
||||
n_win_min = max(5, n_poles // 20)
|
||||
n_win_max = 2000 if n_poles < 2000 else 8000
|
||||
best_wmp = best_metric = None
|
||||
for n_w in np.unique(np.linspace(n_win_min, n_win_max, 20, dtype=int)):
|
||||
for n_cf in range(10, 1, -1):
|
||||
for n_w in np.unique(
|
||||
np.linspace(n_win_min, n_win_max, search_n_win, dtype=int)
|
||||
):
|
||||
for n_cf in search_cf_orders:
|
||||
if log:
|
||||
print(f"Testing N_win={n_w} N_cf={n_cf}")
|
||||
|
||||
|
|
|
|||
811
openmc/data/vectfit.py
Normal file
811
openmc/data/vectfit.py
Normal file
|
|
@ -0,0 +1,811 @@
|
|||
"""
|
||||
Fast Relaxed Vector Fitting function
|
||||
|
||||
Approximate f(s) with a rational function:
|
||||
f(s)=R*(s*I-A)^(-1) + Polynomials*s
|
||||
where f(s) is a vector of elements.
|
||||
|
||||
When f(s) is a vector, all elements become fitted with a common pole set. The
|
||||
identification is done using the pole relocating method known as Vector Fitting
|
||||
[1] with relaxed non-triviality constraint for faster convergence and smaller
|
||||
fitting errors [2], and utilization of matrix structure for fast solution of the
|
||||
pole identifion step [3].
|
||||
|
||||
[1] B. Gustavsen and A. Semlyen, "Rational approximation of frequency
|
||||
domain responses by Vector Fitting", IEEE Trans. Power Delivery, vol. 14,
|
||||
no. 3, pp. 1052-1061, July 1999.
|
||||
[2] B. Gustavsen, "Improving the pole relocating properties of vector
|
||||
fitting", IEEE Trans. Power Delivery, vol. 21, no. 3, pp. 1587-1592, July
|
||||
2006.
|
||||
[3] D. Deschrijver, M. Mrozowski, T. Dhaene, and D. De Zutter,
|
||||
"Macromodeling of Multiport Systems Using a Fast Implementation of the
|
||||
Vector Fitting Method", IEEE Microwave and Wireless Components Letters, vol.
|
||||
18, no. 6, pp. 383-385, June 2008.
|
||||
|
||||
All credit goes to:
|
||||
- Bjorn Gustavsen for his MATLAB implementation.
|
||||
(http://www.sintef.no/Projectweb/VECTFIT/)
|
||||
- Jingang Liang for his C++ implementation.
|
||||
(https://github.com/mit-crpg/vectfit.git)
|
||||
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
from scipy.linalg import eigvals, lstsq, norm, qr
|
||||
|
||||
|
||||
def wlstsq(a, b):
|
||||
"""Apply least-squares solve with column normalization.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This routine rescales columns of `a` to improve conditioning. Columns with
|
||||
zero norm are left unscaled to avoid divide-by-zero warnings.
|
||||
"""
|
||||
col_norm = np.linalg.norm(a, axis=0)
|
||||
scale = np.ones_like(col_norm, dtype=float)
|
||||
nonzero = col_norm > 0.0
|
||||
scale[nonzero] = 1.0 / col_norm[nonzero]
|
||||
scale = np.nan_to_num(scale, nan=1.0, posinf=1.0, neginf=1.0)
|
||||
|
||||
sol = lstsq(a * scale, b)
|
||||
return (sol[0] * scale, sol[1:])
|
||||
|
||||
|
||||
def evaluate(
|
||||
eval_points: np.ndarray,
|
||||
pole_values: np.ndarray,
|
||||
residue_matrix: np.ndarray,
|
||||
poly_coefficients: np.ndarray | None = None,
|
||||
) -> np.ndarray:
|
||||
"""Evaluate the rational function approximation:
|
||||
f(s) ≈ sum(residue / (s - pole)) + sum(poly_coefficients * s^j)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
eval_points : np.ndarray
|
||||
1D array of real scalar frequency values (s).
|
||||
pole_values : np.ndarray
|
||||
1D array of complex poles.
|
||||
residue_matrix : np.ndarray
|
||||
2D or 1D array of complex residues (shape: [num_vectors, num_poles] or [num_poles]).
|
||||
poly_coefficients : np.ndarray, optional
|
||||
2D or 1D array of real polynomial coefficients (shape: [num_vectors, num_polys]).
|
||||
|
||||
Returns
|
||||
-------
|
||||
np.ndarray
|
||||
2D array of evaluated real function values (shape: [num_vectors, num_samples]).
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If input arrays have incompatible shapes.
|
||||
"""
|
||||
eval_points = np.asarray(eval_points)
|
||||
pole_values = np.asarray(pole_values)
|
||||
residue_matrix = np.asarray(residue_matrix)
|
||||
|
||||
if eval_points.ndim != 1:
|
||||
raise ValueError("eval_points must be a 1D array")
|
||||
if pole_values.ndim != 1:
|
||||
raise ValueError("pole_values must be a 1D array")
|
||||
|
||||
if residue_matrix.ndim == 1:
|
||||
residue_matrix = residue_matrix.reshape((1, -1))
|
||||
num_vectors, _ = residue_matrix.shape
|
||||
num_samples = len(eval_points)
|
||||
|
||||
if poly_coefficients is not None and isinstance(poly_coefficients, list):
|
||||
poly_coefficients = np.array(poly_coefficients)
|
||||
if poly_coefficients is None or poly_coefficients.size == 0:
|
||||
poly_coefficients = np.zeros((num_vectors, 0))
|
||||
else:
|
||||
poly_coefficients = np.asarray(poly_coefficients)
|
||||
if poly_coefficients.ndim == 1:
|
||||
poly_coefficients = poly_coefficients.reshape((1, -1))
|
||||
elif poly_coefficients.shape[0] != num_vectors:
|
||||
raise ValueError("Mismatch in residues and poly_coefficients shapes")
|
||||
|
||||
num_coeffs = poly_coefficients.shape[1]
|
||||
result = np.zeros((num_vectors, num_samples))
|
||||
|
||||
# term: sum over poles of (residues / (eval_points - poles))
|
||||
denominator = (
|
||||
eval_points[np.newaxis, :] - pole_values[:, np.newaxis]
|
||||
) # shape: (num_poles, num_eval)
|
||||
pole_terms = residue_matrix @ (1.0 / denominator) # shape: (num_vectors, num_eval)
|
||||
result = np.real(pole_terms)
|
||||
|
||||
# polynomial part: sum over poly_idx of (coeff * eval_points**poly_idx)
|
||||
if num_coeffs > 0:
|
||||
powers = (
|
||||
eval_points[np.newaxis, :] ** np.arange(num_coeffs)[:, np.newaxis]
|
||||
) # shape: (num_coeffs, num_eval)
|
||||
result += poly_coefficients @ powers # shape: (num_vectors, num_eval)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def vectfit(
|
||||
response_matrix: np.ndarray,
|
||||
eval_points: np.ndarray,
|
||||
initial_poles: np.ndarray,
|
||||
weights: np.ndarray,
|
||||
n_polys: int = 0,
|
||||
skip_pole_update: bool = False,
|
||||
skip_residue_update: bool = False,
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]:
|
||||
"""Perform vector fitting using the Fast Relaxed Vector Fitting algorithm.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
response_matrix : np.ndarray
|
||||
Complex matrix of frequency responses (shape: [num_vectors, num_samples]).
|
||||
eval_points : np.ndarray
|
||||
Real frequency samples (s), shape (num_samples,).
|
||||
initial_poles : np.ndarray
|
||||
Initial guess for poles (complex), shape (num_poles,).
|
||||
weights : np.ndarray
|
||||
Weighting matrix for fitting (same shape as response_matrix).
|
||||
n_polys : int, optional
|
||||
Number of real polynomial terms to include.
|
||||
skip_pole_update : bool, optional
|
||||
Whether to skip pole relocation step.
|
||||
skip_residue_update : bool, optional
|
||||
Whether to skip residue fitting step.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, float]
|
||||
- Updated poles (np.ndarray)
|
||||
- Residues (np.ndarray)
|
||||
- Polynomial coefficients (np.ndarray)
|
||||
- Fitted response matrix (np.ndarray)
|
||||
- Root-mean-square error (float)
|
||||
"""
|
||||
tol_low = 1e-18
|
||||
tol_high = 1e18
|
||||
response_matrix = np.asarray(response_matrix)
|
||||
eval_points = np.asarray(eval_points)
|
||||
initial_poles = np.asarray(initial_poles)
|
||||
weights = np.asarray(weights)
|
||||
|
||||
num_vectors, num_samples = response_matrix.shape
|
||||
num_poles = len(initial_poles)
|
||||
|
||||
if n_polys < 0 or n_polys > 11:
|
||||
raise ValueError("n_polys must be in [0, 11]")
|
||||
|
||||
residue_matrix = np.zeros((num_vectors, num_poles), dtype=np.complex128)
|
||||
poly_coefficients = np.zeros((num_vectors, n_polys))
|
||||
fit_result = np.zeros_like(response_matrix)
|
||||
rms_error = 0.0
|
||||
|
||||
if num_poles == 0 and n_polys == 0:
|
||||
rms_error = norm(response_matrix) / np.sqrt(num_vectors * num_samples)
|
||||
return initial_poles, residue_matrix, poly_coefficients, fit_result, rms_error
|
||||
|
||||
if not skip_pole_update and num_poles > 0:
|
||||
updated_poles = identify_poles(
|
||||
num_poles,
|
||||
num_samples,
|
||||
n_polys,
|
||||
initial_poles,
|
||||
eval_points,
|
||||
tol_high,
|
||||
weights,
|
||||
response_matrix,
|
||||
num_vectors,
|
||||
tol_low,
|
||||
)
|
||||
else:
|
||||
updated_poles = initial_poles
|
||||
|
||||
if not skip_residue_update:
|
||||
fit_result, rms_error = identify_residues(
|
||||
num_poles,
|
||||
updated_poles,
|
||||
num_samples,
|
||||
n_polys,
|
||||
eval_points,
|
||||
num_vectors,
|
||||
weights,
|
||||
response_matrix,
|
||||
poly_coefficients,
|
||||
residue_matrix,
|
||||
)
|
||||
|
||||
return updated_poles, residue_matrix, poly_coefficients, fit_result, rms_error
|
||||
|
||||
|
||||
def compute_dk_matrix(
|
||||
dk_matrix: np.ndarray,
|
||||
eval_points: np.ndarray,
|
||||
poles: np.ndarray,
|
||||
conj_index: np.ndarray,
|
||||
num_poles: int,
|
||||
num_polys: int,
|
||||
tol_high: float = None,
|
||||
):
|
||||
"""Compute the dk_matrix used in windowed multipole evaluations.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dk_matrix : ndarray of shape (len(eval_points), M)
|
||||
The full matrix used in least-squares fitting or evaluation.
|
||||
eval_points : ndarray of shape (N,)
|
||||
Energy points at which to evaluate.
|
||||
poles : ndarray of shape (num_poles,)
|
||||
Complex poles used in the resonance model.
|
||||
conj_index : ndarray of shape (num_poles,)
|
||||
Index array indicating pole conjugacy behavior: 0 (normal), 1 (add conjugate), 2 (imaginary part).
|
||||
num_poles : int
|
||||
Number of complex poles.
|
||||
num_polys : int
|
||||
Number of polynomial terms (including constant term).
|
||||
tol_high : float
|
||||
Replacement value for infinities.
|
||||
"""
|
||||
# Broadcast shapes
|
||||
eval_points_col = eval_points[:, np.newaxis]
|
||||
poles_row = poles[np.newaxis, :]
|
||||
|
||||
# Compute base terms
|
||||
term1 = 1.0 / (eval_points_col - poles_row)
|
||||
term2 = 1.0 / (eval_points_col - np.conj(poles_row))
|
||||
term3 = 1j / (eval_points_col - np.conj(poles_row)) - 1j / (
|
||||
eval_points_col - poles_row
|
||||
)
|
||||
|
||||
# Masks for different conjugacy types
|
||||
mask0 = conj_index == 0
|
||||
mask1 = conj_index == 1
|
||||
mask2 = conj_index == 2
|
||||
|
||||
# Fill dk_matrix with pole terms
|
||||
dk_matrix[:, :num_poles][:, mask0] = term1[:, mask0]
|
||||
dk_matrix[:, :num_poles][:, mask1] = term1[:, mask1] + term2[:, mask1]
|
||||
dk_matrix[:, :num_poles][:, mask2] = term3[:, mask2]
|
||||
|
||||
# Replace infinities with high tolerance value
|
||||
if tol_high is not None:
|
||||
inf_mask = np.isinf(dk_matrix)
|
||||
dk_matrix[inf_mask] = tol_high + 0j
|
||||
|
||||
# Add polynomial basis (Chebyshev-like, just powers here)
|
||||
powers = np.arange(num_polys)
|
||||
dk_matrix[:, num_poles : num_poles + num_polys] = eval_points_col**powers + 0j
|
||||
return dk_matrix
|
||||
|
||||
|
||||
def row_block_matrix(
|
||||
dk_matrix: np.ndarray,
|
||||
weights: np.ndarray,
|
||||
response_matrix: np.ndarray,
|
||||
vec_idx: int,
|
||||
num_poles: int,
|
||||
num_polys: int,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Construct a single matrix row block for the given vector index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dk_matrix : ndarray of shape (num_samples, num_poles + num_polys)
|
||||
Basis function evaluations at each sample point.
|
||||
weights : ndarray of shape (num_vectors, num_samples)
|
||||
Sample weights for each vector.
|
||||
response_matrix : ndarray of shape (num_vectors, num_samples)
|
||||
Response values at each sample point.
|
||||
vec_idx : int
|
||||
Index of the vector to construct the A1 block for.
|
||||
num_poles : int
|
||||
Number of poles used in the model.
|
||||
num_polys : int
|
||||
Number of polynomial basis terms.
|
||||
|
||||
Returns
|
||||
-------
|
||||
A : ndarray of shape (num_samples, num_poles + num_polys + num_poles + 1)
|
||||
Weighted and assembled matrix block for the current vector.
|
||||
"""
|
||||
num_samples = dk_matrix.shape[0]
|
||||
A = np.zeros(
|
||||
(num_samples, num_poles + num_polys + num_poles + 1), dtype=np.complex128
|
||||
)
|
||||
|
||||
# Weighted basis terms
|
||||
A[:, : num_poles + num_polys] = (
|
||||
weights[vec_idx][:, np.newaxis] * dk_matrix[:, : num_poles + num_polys]
|
||||
)
|
||||
|
||||
# Weighted response terms (includes poles + 1)
|
||||
A[:, num_poles + num_polys : num_poles + num_polys + num_poles + 1] = (
|
||||
-weights[vec_idx][:, np.newaxis]
|
||||
* dk_matrix[:, : num_poles + 1]
|
||||
* response_matrix[vec_idx][:, np.newaxis]
|
||||
)
|
||||
|
||||
return A
|
||||
|
||||
|
||||
def process_constrained_block(
|
||||
vec_idx: int,
|
||||
dk_matrix: np.ndarray,
|
||||
weights: np.ndarray,
|
||||
response_matrix: np.ndarray,
|
||||
num_samples: int,
|
||||
num_poles: int,
|
||||
num_polys: int,
|
||||
scale_factor: float,
|
||||
num_vectors: int,
|
||||
) -> Tuple[int, np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Construct a constrained least-squares system block for the given vector index.
|
||||
|
||||
This function computes the A matrix using weighted evaluations of the basis functions
|
||||
and response terms. It appends a constraint row to enforce physical properties
|
||||
(e.g., normalization) **only for the final vector index**. The full matrix A is
|
||||
decomposed via QR, and the resulting triangular block is returned.
|
||||
|
||||
This routine is intended for use in the main vector fitting loop when the denominator
|
||||
is well-conditioned but requires an additional constraint row for physical consistency.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vec_idx : int
|
||||
Index of the vector to process.
|
||||
dk_matrix : ndarray of shape (num_samples, num_poles + num_polys)
|
||||
Evaluated basis functions at sample points.
|
||||
weights : ndarray of shape (num_vectors, num_samples)
|
||||
Weight matrix per vector.
|
||||
response_matrix : ndarray of shape (num_vectors, num_samples)
|
||||
Response function values for each vector.
|
||||
num_samples : int
|
||||
Number of sample points.
|
||||
num_poles : int
|
||||
Number of poles in the model.
|
||||
num_polys : int
|
||||
Number of polynomial terms in the model.
|
||||
scale_factor : float
|
||||
Scaling factor applied to the final constraint row.
|
||||
num_vectors : int
|
||||
Total number of vectors to process.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vec_idx : int
|
||||
Index of the processed vector.
|
||||
lhs_block : ndarray of shape (num_poles + 1, num_poles + 1)
|
||||
Triangular matrix block from QR decomposition.
|
||||
rhs_block : ndarray of shape (num_poles + 1,) or None
|
||||
Right-hand side vector block (only returned for final vec_idx), else None.
|
||||
"""
|
||||
A1 = row_block_matrix(
|
||||
dk_matrix, weights, response_matrix, vec_idx, num_poles, num_polys
|
||||
)
|
||||
A = np.zeros((2 * num_samples + 1, num_poles + num_polys + num_poles + 1))
|
||||
A[:num_samples] = A1.real
|
||||
A[num_samples : 2 * num_samples] = A1.imag
|
||||
# Handle final row only if vec_idx is last
|
||||
if vec_idx == num_vectors - 1:
|
||||
A[
|
||||
2 * num_samples,
|
||||
num_poles + num_polys : num_poles + num_polys + num_poles + 1,
|
||||
] = scale_factor * np.real(dk_matrix[:, : num_poles + 1].sum(axis=0))
|
||||
|
||||
Q, R = qr(A, mode="economic")
|
||||
|
||||
lhs_block = R[
|
||||
num_poles + num_polys : num_poles + num_polys + num_poles + 1,
|
||||
num_poles + num_polys : num_poles + num_polys + num_poles + 1,
|
||||
]
|
||||
|
||||
if vec_idx == num_vectors - 1:
|
||||
rhs_block = (
|
||||
num_samples
|
||||
* scale_factor
|
||||
* Q[-1, num_poles + num_polys : num_poles + num_polys + num_poles + 1]
|
||||
)
|
||||
else:
|
||||
rhs_block = np.zeros_like(
|
||||
Q[-1, num_poles + num_polys : num_poles + num_polys + num_poles + 1]
|
||||
)
|
||||
|
||||
return vec_idx, lhs_block, rhs_block
|
||||
|
||||
|
||||
def process_unconstrained_block(
|
||||
vec_idx: int,
|
||||
dk_matrix: np.ndarray,
|
||||
weights: np.ndarray,
|
||||
response_matrix: np.ndarray,
|
||||
denom: float,
|
||||
num_poles: int,
|
||||
num_polys: int,
|
||||
) -> Tuple[int, np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Construct an unconstrained least-squares system block for the given vector index.
|
||||
|
||||
This function is used when the fitting denominator becomes ill-conditioned
|
||||
(too small or too large), and the original constrained system is replaced by
|
||||
an alternative regularized least-squares problem. The A matrix is built by stacking
|
||||
the real and imaginary parts of the basis evaluations, and the RHS vector b is
|
||||
scaled by `denom`.
|
||||
|
||||
A standard QR decomposition is used to extract the square block of the system,
|
||||
which can be solved independently from the constrained system.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vec_idx : int
|
||||
Index of the vector to process.
|
||||
dk_matrix : ndarray of shape (num_samples, num_poles + num_polys)
|
||||
Evaluated basis functions at sample points.
|
||||
weights : ndarray of shape (num_vectors, num_samples)
|
||||
Weight matrix per vector.
|
||||
response_matrix : ndarray of shape (num_vectors, num_samples)
|
||||
Response function values for each vector.
|
||||
denom : float
|
||||
Scaling factor applied to the right-hand side vector b.
|
||||
num_poles : int
|
||||
Number of poles in the model.
|
||||
num_polys : int
|
||||
Number of polynomial terms in the model.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vec_idx : int
|
||||
Index of the processed vector.
|
||||
lhs_block : ndarray of shape (num_poles, num_poles)
|
||||
Triangular matrix block from QR decomposition.
|
||||
rhs_block : ndarray of shape (num_poles,)
|
||||
Right-hand side vector block for this vector.
|
||||
"""
|
||||
A1 = row_block_matrix(
|
||||
dk_matrix, weights, response_matrix, vec_idx, num_poles, num_polys
|
||||
)
|
||||
A = np.vstack((A1.real, A1.imag))
|
||||
|
||||
b1 = denom * weights[vec_idx] * response_matrix[vec_idx]
|
||||
b = np.concatenate((b1.real, b1.imag))
|
||||
|
||||
Q, R = qr(A, mode="economic")
|
||||
|
||||
lhs_block = R[
|
||||
num_poles + num_polys : num_poles + num_polys + num_poles,
|
||||
num_poles + num_polys : num_poles + num_polys + num_poles,
|
||||
]
|
||||
rhs_block = Q[:, num_poles + num_polys : num_poles + num_polys + num_poles].T @ b
|
||||
return vec_idx, lhs_block, rhs_block
|
||||
|
||||
|
||||
def identify_poles(
|
||||
num_poles: int,
|
||||
num_samples: int,
|
||||
num_polys: int,
|
||||
poles: np.ndarray,
|
||||
eval_points: np.ndarray,
|
||||
tol_high: float,
|
||||
weights: np.ndarray,
|
||||
response_matrix: np.ndarray,
|
||||
num_vectors: int,
|
||||
tol_low: float,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Internal routine to update poles via relaxed vector fitting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
num_poles : int
|
||||
Number of poles.
|
||||
num_samples : int
|
||||
Number of frequency samples.
|
||||
num_polys : int
|
||||
Number of polynomial terms.
|
||||
poles : np.ndarray
|
||||
Initial poles (complex), shape (num_poles,).
|
||||
eval_points : np.ndarray
|
||||
Real frequency values, shape (num_samples,).
|
||||
tol_high : float
|
||||
Upper tolerance threshold for denominator.
|
||||
weights : np.ndarray
|
||||
Weighting matrix, shape (num_vectors, num_samples).
|
||||
response_matrix : np.ndarray
|
||||
Complex frequency responses, shape (num_vectors, num_samples).
|
||||
num_vectors : int
|
||||
Number of response vectors.
|
||||
tol_low : float
|
||||
Lower tolerance threshold for denominator.
|
||||
|
||||
Returns
|
||||
-------
|
||||
np.ndarray
|
||||
Updated poles as eigenvalues (shape: [num_poles]).
|
||||
"""
|
||||
conj_index = label_conjugate_poles(poles)
|
||||
dk_matrix = np.zeros(
|
||||
(num_samples, num_poles + max(num_polys, 1)), dtype=np.complex128
|
||||
)
|
||||
compute_dk_matrix(
|
||||
dk_matrix, eval_points, poles, conj_index, num_poles, num_polys, tol_high
|
||||
)
|
||||
# For relaxed vector fitting, include the constant term in the sigma(s)
|
||||
# function even when no polynomial terms are requested. This ensures the
|
||||
# constrained system is well-posed when n_polys == 0.
|
||||
if num_polys == 0:
|
||||
dk_matrix[:, num_poles] = 1.0 + 0j
|
||||
|
||||
scale_factor = (
|
||||
np.sqrt(
|
||||
sum(norm(weights[m] * response_matrix[m]) ** 2 for m in range(num_vectors))
|
||||
)
|
||||
/ num_samples
|
||||
)
|
||||
lhs_matrix = np.zeros((num_vectors * (num_poles + 1), num_poles + 1))
|
||||
rhs_vector = np.zeros(num_vectors * (num_poles + 1))
|
||||
|
||||
for vec_idx in range(num_vectors):
|
||||
vec_idx, lhs_block, rhs_block = process_constrained_block(
|
||||
vec_idx,
|
||||
dk_matrix,
|
||||
weights,
|
||||
response_matrix,
|
||||
num_samples,
|
||||
num_poles,
|
||||
num_polys,
|
||||
scale_factor,
|
||||
num_vectors,
|
||||
)
|
||||
i0 = vec_idx * (num_poles + 1)
|
||||
i1 = (vec_idx + 1) * (num_poles + 1)
|
||||
lhs_matrix[i0:i1] = lhs_block
|
||||
rhs_vector[i0:i1] = rhs_block
|
||||
|
||||
solution, *_ = wlstsq(lhs_matrix, rhs_vector)
|
||||
coeffs = solution[:-1]
|
||||
denom = solution[-1]
|
||||
|
||||
if abs(denom) < tol_low or abs(denom) > tol_high:
|
||||
lhs_matrix = np.zeros((num_vectors * num_poles, num_poles))
|
||||
rhs_vector = np.zeros(num_vectors * num_poles)
|
||||
# Adjust denom
|
||||
if denom == 0.0:
|
||||
denom = 1.0
|
||||
elif abs(denom) < tol_low:
|
||||
denom = np.sign(denom) * tol_low
|
||||
elif abs(denom) > tol_high:
|
||||
denom = np.sign(denom) * tol_high
|
||||
|
||||
# Allocate output
|
||||
lhs_matrix = np.zeros((num_vectors * num_poles, num_poles))
|
||||
rhs_vector = np.zeros(num_vectors * num_poles)
|
||||
|
||||
for vec_idx in range(num_vectors):
|
||||
vec_idx, lhs_block, rhs_block = process_unconstrained_block(
|
||||
vec_idx,
|
||||
dk_matrix,
|
||||
weights,
|
||||
response_matrix,
|
||||
denom,
|
||||
num_poles,
|
||||
num_polys,
|
||||
)
|
||||
i0 = vec_idx * num_poles
|
||||
i1 = (vec_idx + 1) * num_poles
|
||||
lhs_matrix[i0:i1] = lhs_block
|
||||
rhs_vector[i0:i1] = rhs_block
|
||||
|
||||
coeffs, *_ = wlstsq(lhs_matrix, rhs_vector)
|
||||
|
||||
lambda_matrix = np.zeros((num_poles, num_poles))
|
||||
scale_vector = np.ones((num_poles, 1))
|
||||
|
||||
# Mask for real poles (conj_index == 0)
|
||||
mask_real = conj_index == 0
|
||||
real_indices = np.where(mask_real)[0]
|
||||
lambda_matrix[real_indices, real_indices] = np.real(poles[real_indices])
|
||||
|
||||
# Mask for start of complex conjugate pairs (conj_index == 1)
|
||||
mask_cplx_start = conj_index == 1
|
||||
cplx_indices = np.where(mask_cplx_start)[0]
|
||||
|
||||
# Extract real and imaginary parts of complex conjugate poles
|
||||
real_parts = np.real(poles[cplx_indices])
|
||||
imag_parts = np.imag(poles[cplx_indices])
|
||||
|
||||
# Diagonal assignments
|
||||
lambda_matrix[cplx_indices, cplx_indices] = real_parts
|
||||
lambda_matrix[cplx_indices + 1, cplx_indices + 1] = real_parts
|
||||
|
||||
# Off-diagonal assignments
|
||||
lambda_matrix[cplx_indices, cplx_indices + 1] = imag_parts
|
||||
lambda_matrix[cplx_indices + 1, cplx_indices] = -imag_parts
|
||||
|
||||
# Scaling vector adjustments
|
||||
scale_vector[cplx_indices, 0] = 2.0
|
||||
scale_vector[cplx_indices + 1, 0] = 0.0
|
||||
|
||||
residue_matrix = lambda_matrix - np.outer(scale_vector.squeeze(), coeffs) / denom
|
||||
return eigvals(residue_matrix)
|
||||
|
||||
|
||||
def solve_vector_block(
|
||||
vec_idx: int,
|
||||
dk_matrix: np.ndarray,
|
||||
weights: np.ndarray,
|
||||
response_matrix: np.ndarray,
|
||||
num_poles: int,
|
||||
num_polys: int,
|
||||
) -> Tuple[int, np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Solve the least-squares system for a single vector index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vec_idx : int
|
||||
Index of the vector to solve.
|
||||
dk_matrix : ndarray
|
||||
Basis function evaluations of shape (num_samples, num_poles + num_polys).
|
||||
weights : ndarray
|
||||
Weight array of shape (num_vectors, num_samples).
|
||||
response_matrix : ndarray
|
||||
Response array of shape (num_vectors, num_samples).
|
||||
num_poles : int
|
||||
Number of poles.
|
||||
num_polys : int
|
||||
Number of polynomial coefficients.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vec_idx : int
|
||||
The index of the solved vector.
|
||||
residues : ndarray
|
||||
Solution vector for the residues (length = num_poles).
|
||||
poly_coeffs : ndarray or None
|
||||
Solution vector for polynomial coefficients (length = num_polys), or None if num_polys == 0.
|
||||
"""
|
||||
A = dk_matrix * weights[vec_idx][:, np.newaxis]
|
||||
b = weights[vec_idx] * response_matrix[vec_idx]
|
||||
|
||||
lhs_matrix = np.vstack((A.real, A.imag))
|
||||
rhs_vector = np.concatenate((b.real, b.imag))
|
||||
|
||||
x = wlstsq(lhs_matrix, rhs_vector)[0]
|
||||
|
||||
residues = x[:num_poles]
|
||||
poly_coeffs = x[num_poles : num_poles + num_polys] if num_polys > 0 else None
|
||||
|
||||
return vec_idx, residues, poly_coeffs
|
||||
|
||||
|
||||
def identify_residues(
|
||||
num_poles: int,
|
||||
poles: np.ndarray,
|
||||
num_samples: int,
|
||||
num_polys: int,
|
||||
eval_points: np.ndarray,
|
||||
num_vectors: int,
|
||||
weights: np.ndarray,
|
||||
response_matrix: np.ndarray,
|
||||
poly_coefficients: np.ndarray,
|
||||
residue_matrix: np.ndarray,
|
||||
) -> Tuple[np.ndarray, float]:
|
||||
"""
|
||||
Internal routine to compute residues and polynomial coefficients.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
num_poles : int
|
||||
Number of poles.
|
||||
poles : np.ndarray
|
||||
Current poles (complex), shape (num_poles,).
|
||||
num_samples : int
|
||||
Number of frequency samples.
|
||||
num_polys : int
|
||||
Number of polynomial terms.
|
||||
eval_points : np.ndarray
|
||||
Real frequency values, shape (num_samples,).
|
||||
num_vectors : int
|
||||
Number of response vectors.
|
||||
weights : np.ndarray
|
||||
Weighting matrix, shape (num_vectors, num_samples).
|
||||
response_matrix : np.ndarray
|
||||
Complex frequency responses, shape (num_vectors, num_samples).
|
||||
poly_coefficients : np.ndarray
|
||||
Array to store output polynomial coefficients (in-place).
|
||||
residue_matrix : np.ndarray
|
||||
Array to store output residues (in-place).
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tuple[np.ndarray, float]
|
||||
- Fitted response matrix (np.ndarray)
|
||||
- Root-mean-square fitting error (float)
|
||||
"""
|
||||
conj_index = label_conjugate_poles(poles)
|
||||
dk_matrix = np.zeros((num_samples, num_poles + num_polys), dtype=np.complex128)
|
||||
|
||||
compute_dk_matrix(dk_matrix, eval_points, poles, conj_index, num_poles, num_polys)
|
||||
|
||||
real_residues = np.zeros((num_vectors, num_poles), dtype=np.float64)
|
||||
|
||||
for vec_idx in range(num_vectors):
|
||||
vec_idx, residues, poly_coeffs = solve_vector_block(
|
||||
vec_idx,
|
||||
dk_matrix,
|
||||
weights,
|
||||
response_matrix,
|
||||
num_poles,
|
||||
num_polys,
|
||||
)
|
||||
real_residues[vec_idx] = residues
|
||||
if poly_coeffs is not None:
|
||||
poly_coefficients[vec_idx] = poly_coeffs
|
||||
|
||||
# Mask for real poles
|
||||
mask_real = conj_index == 0
|
||||
real_indices = np.where(mask_real)[0]
|
||||
residue_matrix[:, real_indices] = real_residues[:, real_indices]
|
||||
|
||||
# Mask for first of complex conjugate pairs
|
||||
mask_cplx_start = conj_index == 1
|
||||
cplx_indices = np.where(mask_cplx_start)[0]
|
||||
|
||||
# Compute complex residues using vectorized operations
|
||||
residue_matrix[:, cplx_indices] = (
|
||||
real_residues[:, cplx_indices] + 1j * real_residues[:, cplx_indices + 1]
|
||||
)
|
||||
residue_matrix[:, cplx_indices + 1] = (
|
||||
real_residues[:, cplx_indices] - 1j * real_residues[:, cplx_indices + 1]
|
||||
)
|
||||
|
||||
fit_result = evaluate(eval_points, poles, residue_matrix, poly_coefficients)
|
||||
rms_error = norm(fit_result - response_matrix) / np.sqrt(num_vectors * num_samples)
|
||||
return fit_result, rms_error
|
||||
|
||||
|
||||
def label_conjugate_poles(poles: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Ensure complex poles appear in conjugate pairs and label them accordingly.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
poles : np.ndarray
|
||||
1D array of complex poles.
|
||||
|
||||
Returns
|
||||
-------
|
||||
np.ndarray
|
||||
Array of integers indicating pole type:
|
||||
- 0: real pole
|
||||
- 1: first in a complex-conjugate pair
|
||||
- 2: second in a complex-conjugate pair
|
||||
|
||||
Raises
|
||||
------
|
||||
ValueError
|
||||
If any complex pole does not have a valid conjugate pair.
|
||||
"""
|
||||
num_poles = len(poles)
|
||||
conj_index = np.zeros(num_poles, dtype=int)
|
||||
|
||||
# Identify complex poles (nonzero imaginary part)
|
||||
is_complex = np.imag(poles) != 0.0
|
||||
|
||||
# Find conjugate pairs: poles[i+1] ≈ conj(poles[i])
|
||||
is_pair_start = is_complex[:-1] & np.isclose(np.conj(poles[:-1]), poles[1:])
|
||||
|
||||
# Mark valid conjugate pair entries
|
||||
conj_index[:-1][is_pair_start] = 1 # mark i with 1
|
||||
conj_index[1:][is_pair_start] = 2 # mark i+1 with 2
|
||||
|
||||
# Now validate: all complex poles must be part of valid conjugate pairs
|
||||
unmatched_complex = is_complex & (conj_index == 0)
|
||||
if np.any(unmatched_complex):
|
||||
raise ValueError("Complex poles must appear in conjugate pairs")
|
||||
|
||||
return conj_index
|
||||
|
|
@ -49,14 +49,40 @@ def test_export_to_hdf5(tmpdir, u235):
|
|||
|
||||
|
||||
def test_from_endf(endf_data):
|
||||
pytest.importorskip('vectfit')
|
||||
endf_file = os.path.join(endf_data, 'neutrons', 'n-001_H_001.endf')
|
||||
assert openmc.data.WindowedMultipole.from_endf(
|
||||
endf_file, log=True, wmp_options={"n_win": 400, "n_cf": 3})
|
||||
endf_file,
|
||||
log=True,
|
||||
# Keep the test lightweight
|
||||
vf_options={
|
||||
"njoy_error": 5e-3,
|
||||
"vf_pieces": 1,
|
||||
"rtol": 5e-2,
|
||||
"atol": 1e-3,
|
||||
"orders": [8, 12],
|
||||
"n_vf_iter": 6,
|
||||
},
|
||||
wmp_options={"n_win": 50, "n_cf": 3, "rtol": 5e-2, "atol": 1e-3},
|
||||
)
|
||||
|
||||
|
||||
def test_from_endf_search(endf_data):
|
||||
pytest.importorskip('vectfit')
|
||||
endf_file = os.path.join(endf_data, 'neutrons', 'n-095_Am_244.endf')
|
||||
endf_file = os.path.join(endf_data, 'neutrons', 'n-001_H_001.endf')
|
||||
assert openmc.data.WindowedMultipole.from_endf(
|
||||
endf_file, log=True, wmp_options={"search": True, 'rtol':1e-2})
|
||||
endf_file,
|
||||
log=True,
|
||||
vf_options={
|
||||
"njoy_error": 5e-3,
|
||||
"vf_pieces": 1,
|
||||
"rtol": 5e-2,
|
||||
"atol": 1e-3,
|
||||
"orders": [8, 12],
|
||||
"n_vf_iter": 6,
|
||||
},
|
||||
wmp_options={
|
||||
"search": True,
|
||||
"rtol": 5e-2,
|
||||
"search_n_win": 3,
|
||||
"search_cf_orders": [5, 3],
|
||||
},
|
||||
)
|
||||
|
|
|
|||
224
tests/unit_tests/test_vectfit.py
Normal file
224
tests/unit_tests/test_vectfit.py
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
"""
|
||||
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)
|
||||
|
|
@ -1,46 +0,0 @@
|
|||
#!/bin/bash
|
||||
set -ex
|
||||
|
||||
PYBIND_BRANCH='master'
|
||||
PYBIND_REPO='https://github.com/pybind/pybind11'
|
||||
|
||||
XTL_BRANCH='0.6.13'
|
||||
XTL_REPO='https://github.com/xtensor-stack/xtl'
|
||||
|
||||
XTENSOR_BRANCH='0.21.3'
|
||||
XTENSOR_REPO='https://github.com/xtensor-stack/xtensor'
|
||||
|
||||
XTENSOR_PYTHON_BRANCH='0.24.1'
|
||||
XTENSOR_PYTHON_REPO='https://github.com/xtensor-stack/xtensor-python'
|
||||
|
||||
XTENSOR_BLAS_BRANCH='0.17.1'
|
||||
XTENSOR_BLAS_REPO='https://github.com/xtensor-stack/xtensor-blas'
|
||||
|
||||
cd $HOME
|
||||
git clone -b $PYBIND_BRANCH $PYBIND_REPO
|
||||
cd pybind11 && mkdir build && cd build && cmake .. && sudo make install
|
||||
pip install $HOME/pybind11
|
||||
|
||||
cd $HOME
|
||||
git clone -b $XTL_BRANCH $XTL_REPO
|
||||
cd xtl && mkdir build && cd build && cmake .. && sudo make install
|
||||
|
||||
cd $HOME
|
||||
git clone -b $XTENSOR_BRANCH $XTENSOR_REPO
|
||||
cd xtensor && mkdir build && cd build && cmake .. && sudo make install
|
||||
|
||||
cd $HOME
|
||||
git clone -b $XTENSOR_PYTHON_BRANCH $XTENSOR_PYTHON_REPO
|
||||
cd xtensor-python && mkdir build && cd build && cmake .. && sudo make install
|
||||
|
||||
cd $HOME
|
||||
git clone -b $XTENSOR_BLAS_BRANCH $XTENSOR_BLAS_REPO
|
||||
cd xtensor-blas && mkdir build && cd build && cmake .. && sudo make install
|
||||
|
||||
# Install wheel (remove when vectfit supports installation with build isolation)
|
||||
pip install wheel
|
||||
|
||||
# Install vectfit
|
||||
cd $HOME
|
||||
git clone https://github.com/liangjg/vectfit.git
|
||||
pip install --no-build-isolation ./vectfit
|
||||
|
|
@ -18,11 +18,6 @@ fi
|
|||
pip install 'ncrystal>=4.1.0'
|
||||
nctool --test
|
||||
|
||||
# Install vectfit for WMP generation if needed
|
||||
if [[ $VECTFIT = 'y' ]]; then
|
||||
./tools/ci/gha-install-vectfit.sh
|
||||
fi
|
||||
|
||||
# Install libMesh if needed
|
||||
if [[ $LIBMESH = 'y' ]]; then
|
||||
./tools/ci/gha-install-libmesh.sh
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue