From 64f9e1aacc78c541a74b9bf5ff4105473bc35fb7 Mon Sep 17 00:00:00 2001
From: liangjg <liangjg2008@gmail.com>
Date: Wed, 2 Jan 2019 09:28:41 -0500
Subject: [PATCH] separate vectfit and windowing

---
 openmc/data/multipole.py | 403 +++++++++++++++++++++++----------------
 1 file changed, 241 insertions(+), 162 deletions(-)

diff --git a/openmc/data/multipole.py b/openmc/data/multipole.py
index edf205e8d9..46b60ec2c1 100644
--- a/openmc/data/multipole.py
+++ b/openmc/data/multipole.py
@@ -136,46 +136,46 @@ def _broaden_wmp_polynomials(E, dopp, n):
     return factors
 
 def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
-                n_iter_vf=30, log=False, path_plot=None):
+                n_vf_iter=30, log=False, path_out=None, **kwargs):
     r"""Generate multipole data from point-wise cross sections.
 
     Parameters
     ----------
     energy : np.ndarray
-        Energy array.
+        Energy array
     cs_xs : np.ndarray
-        Point-wise cross sections to be fitted.
+        Point-wise cross sections to be fitted
     mts : Iterable of Integral
-        Reaction list.
+        Reaction list
     rtol : Real, optional
-        Relative error tolerance.
+        Relative error tolerance
     atol : Real, optional
-        Absolute error tolerance.
+        Absolute error tolerance
     orders : Iterable of Integral, optional
-        A list of orders (number of poles) to be searched.
-    n_iter_vf : Integral, optional
-        Number of maximum VF iterations.
+        A list of orders (number of poles) to be searched
+    n_vf_iter : Integral, optional
+        Number of maximum VF iterations
     log : Bool, optional
-        Whether to print log.
-    path_plot : str, optional
-        Path to save the figures.
+        Whether to print log
+    path_out : str, optional
+        Path to save the figures
+    **kwargs
+        Additional keyword arguments
 
     Returns
     -------
     Tuple
-        (poles, residues).
+        (poles, residues)
 
     """
 
-    MIN_CROSS_SECTION = 1e-7
-    N_FINER = 10
-
     ne = energy.size
     nmt = len(mts)
     if ce_xs.shape != (nmt, ne):
         raise ValueError('Inconsistent cross section data.')
 
     # construct test data: interpolate xs with finer grids
+    N_FINER = 10
     ne_test = (ne-1)*N_FINER + 1
     test_energy = np.interp(np.arange(ne_test),
                             np.arange(ne_test, step=N_FINER), energy)
@@ -187,12 +187,16 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
     if log:
         print("Energy: {:.3e} to {:.3e} eV ({} points)".format(
               energy[0], energy[-1], ne))
+
     # inputs
     f = ce_xs * energy # sigma*E
     s = np.sqrt(energy) # sqrt(E)
     test_s = np.sqrt(test_energy)
     weight = 1.0/f
-    # very small cross sections can lead to huge weights which harm accuracy
+
+    # very small cross sections can lead to huge weights, which will harm the 
+    # fitting accuracy
+    MIN_CROSS_SECTION = 1e-7
     for i in range(nmt):
         if np.all(ce_xs[i]<=MIN_CROSS_SECTION):
             weight[i] = 1.0
@@ -200,9 +204,9 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
             weight[i, ce_xs[i]<=MIN_CROSS_SECTION] = \
                max(weight[i, ce_xs[i]>MIN_CROSS_SECTION])
 
-    # order search
     peaks, _ = find_peaks(ce_xs[0]+ce_xs[1])
     n_peaks = peaks.size
+    # order search
     if orders is not None:
         # make sure orders are even integers
         orders = list(set([int(i/2)*2 for i in orders if i>=2]))
@@ -228,9 +232,9 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
 
         found_better = False
         # fitting iteration
-        for i_vf in range(n_iter_vf):
+        for i_vf in range(n_vf_iter):
             if log:
-                print("VF iteration {}/{}".format(i_vf+1, n_iter_vf))
+                print("VF iteration {}/{}".format(i_vf+1, n_vf_iter))
 
             # call vf
             poles, residues, cf, f_fit, rms = m.vectfit(f, s, poles, weight)
@@ -266,7 +270,7 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
                 ratio = np.sum((relerr<rtol) | (abserr<atol)) / relerr.size
                 ratio2 = np.sum((relerr<10*rtol) | (abserr<atol)) / relerr.size
 
-            quality = 100*(ratio+ratio2-min(0.1*maxre, 1)-5e-4*new_poles.size)
+            quality = ratio + ratio2 - min(0.1*maxre, 1) - 0.001*new_poles.size
 
             if np.any(test_xs < -atol):
                 quality = -np.inf
@@ -332,12 +336,12 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
     if log:
         print("Final number of poles: {}".format(mp_poles.size))
 
-    if path_plot:
+    if path_out:
         import matplotlib
         matplotlib.use("agg")
         import matplotlib.pyplot as plt
-        if not os.path.exists(path_plot):
-            os.makedirs(path_plot)
+        if not os.path.exists(path_out):
+            os.makedirs(path_out)
         for i, mt in enumerate(mts):
             fig, ax1 = plt.subplots()
             lns1 = ax1.loglog(test_energy, test_xs_ref[i], 'g', label="ACE xs")
@@ -356,14 +360,188 @@ def _vectfit_xs(energy, ce_xs, mts, rtol=1e-3, atol=1e-5, orders=None,
 
             plt.title("MT {} vectfitted with {} poles".format(mt, mp_poles.size))
             fig.tight_layout()
-            figfile = os.path.join(path_plot, "{}_vf.png".format(mt))
-            plt.savefig(figfile)
+            fig_file = os.path.join(path_out, "{:.0f}-{:.0f}_MT{}.png".format(
+                                    energy[0], energy[-1], mt))
+            plt.savefig(fig_file)
             plt.close()
             if log:
-                print("Plot figure: {}".format(figfile))
+                print("Saved figure: {}".format(fig_file))
 
     return (mp_poles, mp_residues)
 
+def _vectfit_nuclide(endf_file, njoy_error=5e-4, vf_error=1e-3, vf_pieces=None,
+                     log=False, path_out=None, mp_filename=None, **kwargs):
+    """Convert point-wise cross section to multipole data via Vector Fitting.
+
+    Parameters
+    ----------
+    endf_file : str
+        Path to ENDF evaluation
+    njoy_error : float, optional
+        Fractional error tolerance for processing point-wise data with NJOY
+    vf_error : float, optional
+        Fractional error tolerance for data fitting
+    vf_pieces : integer, optional
+        Number of equal-in-mementum spaced energy pieces for data fitting
+    log : bool, optional
+        Whether to display log
+    path_out : str, optional
+        Path to write out data files
+    mp_filename : str, optional
+        File name to write out multipole data
+    **kwargs
+        Keyword arguments passed to :func:`openmc.data.multipole._vectfit_xs`
+
+    Returns
+    -------
+    mp_data
+        Dictionary with multipole data of the nuclide
+
+    """
+
+    # ======================================================================
+    # PREPARE POINT-WISE XS
+
+    # make 0K ACE data using njoy
+    if log:
+        print("Running NJOY to get 0K point-wise data...")
+
+    nuc_ce = IncidentNeutron.from_njoy(endf_file, temperatures=[0.0],
+             error=njoy_error, broadr=False, heatr=False, purr=False)
+
+    if log:
+        print("Parsing cross sections within resolved resonance range...")
+
+    # Determine upper energy: the lower of RRR upper bound and first threshold
+    endf_res = IncidentNeutron.from_endf(endf_file).resonances
+    if hasattr(endf_res, 'resolved') and \
+       hasattr(endf_res.resolved, 'energy_max') and \
+       type(endf_res.resolved) is not ResonanceRange:
+        E_max = endf_res.resolved.energy_max
+    elif hasattr(endf_res, 'unresolved') and \
+         hasattr(endf_res.unresolved, 'energy_min'):
+        E_max = endf_res.unresolved.energy_min
+    else:
+        E_max = nuc_ce.energy['0K'][-1]
+    E_max_idx = np.searchsorted(nuc_ce.energy['0K'], E_max, side='right') - 1
+    for mt in nuc_ce.reactions:
+        if hasattr(nuc_ce.reactions[mt].xs['0K'], '_threshold_idx'):
+            threshold_idx = nuc_ce.reactions[mt].xs['0K']._threshold_idx
+            if 0 < threshold_idx < E_max_idx:
+                E_max_idx = threshold_idx
+
+    # parse energy and cross sections
+    energy = nuc_ce.energy['0K'][:E_max_idx+1]
+    E_min, E_max = energy[0], energy[-1]
+    n_points = energy.size
+    total_xs = nuc_ce[1].xs['0K'](energy)
+    if 2 in nuc_ce:
+        elastic_xs = nuc_ce[2].xs['0K'](energy)
+    else:
+        elastic_xs = np.zeros_like(total_xs)
+    if 27 in nuc_ce:
+        absorption_xs = nuc_ce[27].xs['0K'](energy)
+    else:
+        if 101 in nuc_ce:
+            absorption_xs = nuc_ce[101].xs['0K'](energy)
+            if 18 in nuc_ce:
+                absorption_xs += nuc_ce[18].xs['0K'](energy)
+        else:
+            absorption_xs = np.zeros_like(total_xs)
+    fissionable = False
+    if 18 in nuc_ce:
+        fission_xs = nuc_ce[18].xs['0K'](energy)
+        fissionable = True
+
+    # make vectors
+    if fissionable:
+        ce_xs = np.vstack((elastic_xs, absorption_xs, fission_xs))
+        mts = [2, 27, 18]
+    else:
+        ce_xs = np.vstack((elastic_xs, absorption_xs))
+        mts = [2, 27]
+
+    if log:
+        print("  MTs: {})".format(mts))
+        print("  Energy range: {:.3e} to {:.3e} eV ({} points)".format(
+              E_min, E_max, n_points))
+
+    # ======================================================================
+    # PERFORM VECTOR FITTING
+
+    if vf_pieces is None:
+        # divide into pieces for complex nuclides
+        peaks, _ = find_peaks(total_xs)
+        n_peaks = peaks.size
+        if n_peaks > 300 or n_points > 50000 or n_peaks * n_points > 100*30000:
+            vf_pieces = max(5, n_peaks // 80,  n_points // 5000)
+        else:
+            vf_pieces = 1
+    piece_width = (sqrt(E_max) - sqrt(E_min)) / vf_pieces
+
+    alpha = nuc_ce.atomic_weight_ratio/(K_BOLTZMANN*TEMPERATURE_LIMIT)
+
+    poles, residues = None, None
+    piece_idxs = [0]
+    # VF piece by piece
+    for i_piece in range(vf_pieces):
+        if log:
+            print("Processing piece {}/{}...".format(i_piece+1, vf_pieces))
+        # start E of this piece
+        e_bound = (sqrt(E_min) + piece_width*(i_piece-0.5))**2
+        if i_piece == 0 or sqrt(alpha*e_bound) < 4.0:
+            e_start = E_min
+            e_start_idx = 0
+        else:
+            e_start = max(E_min, (sqrt(alpha*e_bound)-4.0)**2/alpha)
+            e_start_idx = np.searchsorted(energy, e_start, side='right') - 1
+        # end E of this piece
+        e_bound = (sqrt(E_min) + piece_width*(i_piece+1))**2
+        e_end = min(E_max, (sqrt(alpha*e_bound) + 4.0)**2/alpha)
+        e_end_idx = np.searchsorted(energy, e_end, side='left') + 1
+        e_idx = range(e_start_idx, min(e_end_idx+1, n_points))
+
+        p, r = _vectfit_xs(energy[e_idx], ce_xs[:, e_idx], mts, rtol=vf_error,
+                       log=log, path_out=path_out, **kwargs)
+
+        if poles is None:
+            poles, residues = p, r
+        else:
+            poles = np.hstack(poles, p)
+            residues = np.hstack(residues, r)
+
+        piece_idxs.append(piece_idxs[-1] + p.size)
+
+    # gather multipole data into a dictionary
+    mp_data = {"name": nuc_ce.name,
+               "AWR": nuc_ce.atomic_weight_ratio,
+               "E_min": E_min,
+               "E_max": E_max,
+               "poles": poles,
+               "residues": residues,
+               "piece_idxs": piece_idxs}
+
+    # dump multipole data to files
+    if path_out:
+        import pickle
+        if not os.path.exists(path_out):
+            os.makedirs(path_out)
+        if not mp_filename:
+            mp_filename = "{}_mp.pickle".format(nuc_ce.name)
+        mp_filename = os.path.join(path_out, mp_filename)
+        with open(mp_filename, 'wb') as f:
+            pickle.dump(mp_data, f)
+        if log:
+            print("Dumped multipole data to file: {}".format(mp_filename))
+
+    return mp_data
+
+def _windowing(mp_data):
+    """
+
+    """
+    pass
+
 class WindowedMultipole(EqualityMixin):
     """Resonant cross sections represented in the windowed multipole format.
 
@@ -627,22 +805,43 @@ class WindowedMultipole(EqualityMixin):
         return out
 
     @classmethod
-    def from_vectfit(cls, endf_file, error=1e-3, njoy_error=5e-4, log=False,
-                     path=None, **kwargs):
-        """Generate windowed multipole neutron data via Vector Fitting.
+    def from_endf(cls, endf_file, **kwargs):
+        """Generate windowed multipole neutron data from an ENDF evaluation.
 
         Parameters
         ----------
         endf_file : str
             Path to ENDF evaluation
-        error : float, optional
-            Fractional error tolerance for data fitting
-        njoy_error : float, optional
-            Fractional error tolerance for processing point-wise data with NJOY
+        **kwargs
+            Keyword arguments passed to :func:`openmc.data.multipole._vectfit_nuclide`
+            and :func:`openmc.data.WindowedMultipole.from_multipole`
+
+        Returns
+        -------
+        openmc.data.WindowedMultipole
+            Resonant cross sections represented in the windowed multipole
+            format.
+
+        """
+
+        # generate multipole data from EDNF
+        mp_data = _vectfit_nuclide(endf_file, **kwargs)
+
+        # windowing
+        return cls.from_multipole(mp_data, **kwargs)
+
+    @classmethod
+    def from_multipole(cls, mp_data, **kwargs):
+        """Generate windowed multipole neutron data from multipole data.
+
+        Parameters
+        ----------
+        mp_data : dictionary or str
+            Dictionary or Path to the multipole data
         log : bool, optional
             Whether to display log
-        path : str, optional
-            Path to save the figures.
+        path_out : str, optional
+            Path to write out data files
         **kwargs
             Keyword arguments passed to :func:`openmc.data.multipole._vectfit_xs`
 
@@ -654,134 +853,14 @@ class WindowedMultipole(EqualityMixin):
 
         """
 
-        # ======================================================================
-        # PREPARE POINT-WISE XS
+        if isinstance(mp_data, str):
+            # load multipole data from file
+            import pickle
+            with open(mp_data, 'rb') as f:
+                mp_data = pickle.load(f)
 
-        # make 0K ACE data using njoy
-        if log:
-            print("Running NJOY to get 0K point-wise data...")
-        nuc_ce = IncidentNeutron.from_njoy(endf_file, temperatures=[0.0],
-                        error=njoy_error, broadr=False, heatr=False, purr=False)
+        # windowing
 
-        if log:
-            print("Parsing cross section within resolved resonance range...")
-        # RRR bound
-        endf_res = IncidentNeutron.from_endf(endf_file).resonances
-
-        rrr_bound_energy = nuc_ce.energy['0K'][-1]
-        try:
-            rrr = endf_res.resolved
-        except:
-            rrr = None
-        # if resolved resonance parameters exist
-        if rrr is not None and hasattr(rrr, 'energy_max') and \
-             type(rrr) is not ResonanceRange:
-            rrr_bound_energy = rrr.energy_max
-        else:
-            try:
-                # set rrr bound as lower bound of unresolved
-                rrr_bound_energy = endf_res.unresolved.energy_min
-            except:
-                pass
-        rrr_bound_idx = np.searchsorted(nuc_ce.energy['0K'], rrr_bound_energy,
-                                        side='right') - 1
-
-        # first threshold
-        first_threshold_idx = float("inf")
-        first_threshold_mt = None
-        for mt in nuc_ce.reactions:
-            if hasattr(nuc_ce.reactions[mt].xs['0K'], '_threshold_idx'):
-                threshold_idx = nuc_ce.reactions[mt].xs['0K']._threshold_idx
-                if 0 < threshold_idx < first_threshold_idx:
-                    first_threshold_idx = threshold_idx
-                    first_threshold_mt = mt
-
-        # lower of RRR bound and first threshold
-        e_max_idx = min(rrr_bound_idx, first_threshold_idx)
-
-        if log:
-            print("  RRR idx: {}, first threshold idx: {}".format(rrr_bound_idx,
-                  first_threshold_idx))
-
-        # parse energy and summed cross sections
-        energy = nuc_ce.energy['0K'][:e_max_idx+1]
-        E_min, E_max = energy[0], energy[-1]
-        n_points = energy.size
-
-        total_xs = nuc_ce[1].xs['0K'](energy)
-        if 2 in nuc_ce:
-            elastic_xs = nuc_ce[2].xs['0K'](energy)
-        else:
-            elastic_xs = np.zeros_like(total_xs)
-        if 27 in nuc_ce:
-            absorption_xs = nuc_ce[27].xs['0K'](energy)
-        else:
-            if 101 in nuc_ce:
-                absorption_xs = nuc_ce[101].xs['0K'](energy)
-                if 18 in nuc_ce:
-                    absorption_xs += nuc_ce[18].xs['0K'](energy)
-            else:
-                absorption_xs = np.zeros_like(total_xs)
-        fissionable = False
-        if 18 in nuc_ce:
-            fission_xs = nuc_ce[18].xs['0K'](energy)
-            fissionable = True
-
-        # make vectors
-        if fissionable:
-            ce_xs = np.vstack((elastic_xs, absorption_xs, fission_xs))
-            mts = [2, 27, 18]
-        else:
-            ce_xs = np.vstack((elastic_xs, absorption_xs))
-            mts = [2, 27]
-
-        if log:
-            print("  MTs: {})".format(mts))
-            print("  Energy range: {:.3e} to {:.3e} eV ({} points)".format(
-                  E_min, E_max, n_points))
-
-        # ======================================================================
-        # PERFORM VECTOR FITTING
-
-        alpha = nuc_ce.atomic_weight_ratio/(K_BOLTZMANN*TEMPERATURE_LIMIT)
-
-        # divide into pieces for complex nuclides
-        peaks, _ = find_peaks(total_xs)
-        n_peaks = peaks.size
-        if n_peaks > 300 or n_points > 50000 or n_peaks * n_points > 100*30000:
-            n_pieces = max(5, n_peaks // 80,  n_points // 5000)
-        else:
-            n_pieces = 1
-        piece_width = (sqrt(E_max) - sqrt(E_min)) / n_pieces
-
-        # VF piece by piece
-        for i_piece in range(n_pieces):
-            if log:
-                print("Processing piece {}/{}...".format(i_piece+1, n_pieces))
-            # start E of this piece
-            e_bound = (sqrt(E_min) + piece_width*(i_piece-0.5))**2
-            if i_piece == 0 or sqrt(alpha*e_bound) < 4.0:
-                e_start = E_min
-                e_start_idx = 0
-            else:
-                e_start = max(E_min, (sqrt(alpha*e_bound)-4.0)**2/alpha)
-                e_start_idx = np.searchsorted(energy, e_start, side='right') - 1
-            # end E of this piece
-            e_bound = (sqrt(E_min) + piece_width*(i_piece+1))**2
-            e_end = min(E_max, (sqrt(alpha*e_bound) + 4.0)**2/alpha)
-            e_end_idx = np.searchsorted(energy, e_end, side='left') + 1
-            e_idx = range(e_start_idx, min(e_end_idx+1, n_points))
-
-            # fitting xs
-            if path:
-                path_plot = os.path.join(path, "{}".format(i_piece+1))
-            else:
-                path_plot = None
-            p, r = _vectfit_xs(energy[e_idx], ce_xs[:, e_idx], mts, rtol=error,
-                               log=log, path_plot=path_plot, **kwargs)
-
-        # ======================================================================
-        # WINDOWING
 
     def _evaluate(self, E, T):
         """Compute scattering, absorption, and fission cross sections.