diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 38a49b60219..4109fb4f1c8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,7 +31,6 @@ jobs: dagmc: [n] libmesh: [n] event: [n] - vectfit: [n] include: - python-version: "3.12" @@ -56,14 +55,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 }} @@ -71,7 +65,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 @@ -131,11 +124,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: | diff --git a/openmc/data/__init__.py b/openmc/data/__init__.py index c2d35565a8a..a45d026a030 100644 --- a/openmc/data/__init__.py +++ b/openmc/data/__init__.py @@ -33,5 +33,6 @@ from .multipole import * from .grid import * from .function import * +from .vectfit import * from .effective_dose.dose import dose_coefficients diff --git a/openmc/data/multipole.py b/openmc/data/multipole.py index dd14e0d1945..6523cf6ab7e 100644 --- a/openmc/data/multipole.py +++ b/openmc/data/multipole.py @@ -15,7 +15,7 @@ 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) diff --git a/openmc/data/vectfit.py b/openmc/data/vectfit.py new file mode 100644 index 00000000000..b6f4ce81bf5 --- /dev/null +++ b/openmc/data/vectfit.py @@ -0,0 +1,801 @@ +""" +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 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 + ) + + 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 + + column_scale = np.linalg.norm(lhs_matrix, axis=0) + column_scale[column_scale == 0] = 1.0 + lhs_matrix /= column_scale + + solution, *_ = lstsq(lhs_matrix, rhs_vector) + solution *= column_scale + 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 + + column_scale = np.linalg.norm(lhs_matrix, axis=0) + column_scale[column_scale == 0] = 1.0 + lhs_matrix /= column_scale + coeffs, *_ = lstsq(lhs_matrix, rhs_vector) + coeffs *= column_scale + + 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)) + + scale_column = np.linalg.norm(lhs_matrix, axis=0) + scale_column[scale_column == 0] = 1.0 + lhs_matrix /= scale_column + x = lstsq(lhs_matrix, rhs_vector)[0] + x *= scale_column + + 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 diff --git a/tests/unit_tests/test_data_multipole.py b/tests/unit_tests/test_data_multipole.py index bc7136a97c8..083a2b1c6d8 100644 --- a/tests/unit_tests/test_data_multipole.py +++ b/tests/unit_tests/test_data_multipole.py @@ -49,7 +49,6 @@ def test_export_to_hdf5(tmpdir, u235): def test_from_endf(): - pytest.importorskip('vectfit') endf_data = os.environ['OPENMC_ENDF_DATA'] endf_file = os.path.join(endf_data, 'neutrons', 'n-001_H_001.endf') assert openmc.data.WindowedMultipole.from_endf( @@ -57,7 +56,6 @@ def test_from_endf(): def test_from_endf_search(): - pytest.importorskip('vectfit') endf_data = os.environ['OPENMC_ENDF_DATA'] endf_file = os.path.join(endf_data, 'neutrons', 'n-095_Am_244.endf') assert openmc.data.WindowedMultipole.from_endf( diff --git a/tests/unit_tests/test_vectfit.py b/tests/unit_tests/test_vectfit.py new file mode 100644 index 00000000000..c5426cbb856 --- /dev/null +++ b/tests/unit_tests/test_vectfit.py @@ -0,0 +1,218 @@ +""" +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 = 5000 + 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 = 20000 + N = 1000 + 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(poles, expected_poles, rtol=1e-7) + assert np.allclose(residues, expected_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(poles, expected_poles, rtol=1e-5) + assert np.allclose(residues, expected_residues, 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(10): + poles, residues, _, fit, _ = vectfit(f, s, poles, weight) + assert np.allclose(np.sort(poles), np.sort(ref_poles)) + assert np.allclose(np.sort(residues), np.sort(ref_residues)) + 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) diff --git a/tools/ci/gha-install-vectfit.sh b/tools/ci/gha-install-vectfit.sh deleted file mode 100755 index bd38e1ea8cd..00000000000 --- a/tools/ci/gha-install-vectfit.sh +++ /dev/null @@ -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 diff --git a/tools/ci/gha-install.sh b/tools/ci/gha-install.sh index 74c3947f183..223c864b31b 100755 --- a/tools/ci/gha-install.sh +++ b/tools/ci/gha-install.sh @@ -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