Implement vector fitting to replace external vectfit package

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

openmc/data/multipole.py

@@ -9,13 +9,14 @@
 from scipy.signal import find_peaks
 
 import openmc.checkvalue as cv
+
 from ..exceptions import DataError
 from ..mixin import EqualityMixin
 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 +175,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 +191,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 +248,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 +265,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 +385,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
@@ -493,7 +490,7 @@ def vectfit_nuclide(endf_file, njoy_error=5e-4, vf_pieces=None,
     piece_width = (sqrt(E_max) - sqrt(E_min)) / vf_pieces
 
     alpha = nuc_ce.atomic_weight_ratio/(K_BOLTZMANN*TEMPERATURE_LIMIT)
-
+    
     poles, residues = [], []
     # VF piece by piece
     for i_piece in range(vf_pieces):
@@ -571,10 +568,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 +638,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 +652,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)