Merge pull request #65 from SteSeg/uq_w_schema

UQ-TMC workflow - new

Author: SteSeg

.gitignore

@@ -59,3 +59,7 @@ docs/source/_build/
 
 # Visual Studio Code configuration files
 .vscode/
+
+# Virtualenv
+.venv/
+venv/

TODO.md

@@ -3,6 +3,13 @@
 
 
 ## Benchmarks API
+- UQ-TMC: Handle the NJOY dependency
+- UQ-TMC: separate the preparation of run from the execution of the tmc loops (xml files outside the loop, tmc-manager like this: https://github.com/openmc-dev/openmc/pull/3508)
+- UQ-TMC: implement list of models in tmc engine
+- UQ-TMC: the first realization should be the original xs "mean" (baseline)
+- UQ-TMC: allow more options to user (e.g. library to perturb)
+- UQ-TMC: allow a list of nuclides to perturb
+- UQ-TMC: allow list of nuclides and perturbations from specifications file
 - Add capability of merging multiple surface results for oktavian (should happen only with surface tallies)
 - Add expected tally results shape in `schema` and `specifications` to deal with the point here above
 - Try to manage results as numpy arrays instead of pandas dataframes in `OpenmcBenchmark.postprocess()`

pyproject.toml

@@ -25,7 +25,8 @@ dependencies = [
     "xarray",
     "tables",
     "cad-to-dagmc",
-    "pydagmc @ git+https://github.com/svalinn/pydagmc.git"
+    "pydagmc @ git+https://github.com/svalinn/pydagmc.git",
+    "sandy",
 ]
 
 [tool.setuptools.package-data]

src/openmc_fusion_benchmarks/__init__.py

@@ -1,4 +1,5 @@
 from openmc_fusion_benchmarks.benchmark import *
 from openmc_fusion_benchmarks.validate import *
+import openmc_fusion_benchmarks.uq
 
 __version__ = "0.1.0"

src/openmc_fusion_benchmarks/benchmark.py

@@ -6,6 +6,8 @@
 import xarray as xr
 import h5py
 from .validate import validate_benchmark
+from .utils import _openmc_to_ofb, _save_result
+from .uq.tmc_engine import tmc_engine
 
 import openmc
 import pydagmc
@@ -59,7 +61,7 @@ def _build_model(self):
         pass
 
     @abstractmethod
-    def postprocess(self):
+    def _postprocess(self):
         """Post-process the model after running."""
         pass
 
@@ -68,6 +70,11 @@ def run(self):
         """Run the benchmark simulation."""
         pass
 
+    @abstractmethod
+    def _uncertainty_quantification(self):
+        """Perform uncertainty quantification for the benchmark."""
+        pass
+
     def _read_metadata(self):
         """Read metadata from the benchmark specification."""
         metadata = self._benchmark_spec['metadata']
@@ -137,8 +144,7 @@ def _build_materials(self):
 
         materials = openmc.Materials()
         for m in material_data:
-            mat = openmc.Material(name=m['name'])
-            mat.id = m['id']
+            mat = openmc.Material(material_id=m['id'], name=m['name'])
             mat.set_density(m['density']['units'], m['density']['value'])
 
             # Ensure fraction type is valid
@@ -395,77 +401,55 @@ def _build_model(self):
         )
         return model
 
-    def postprocess(self, statepoint: str = 'statepoint.100.h5', mesh: str = 'mesh.h5m'):
+    def _postprocess(self, statepoint: openmc.StatePoint, mesh: str = 'mesh.h5m'):
         """Post-process the model after running."""
         # Retrieve tallies data from specifications
         tallies_data = self._benchmark_spec['tallies']
-        # Read openmc statepoint file
-        sp = openmc.StatePoint(statepoint)
-        # Read mesh file
-        mesh = pydagmc.DAGModel(mesh)
-
-        # Cycle tallies in specifications
-        for spec_t in tallies_data:
-            # Get corresponding tally from statepoint
-            df = sp.get_tally(name=spec_t['name']).get_pandas_dataframe()
-
-            # Preparing tally dataframe
-            df = df.drop(columns=['surface', 'cell', 'particle', 'nuclide',
-                                  'score', 'energyfunction'], errors='ignore')
-            # Cyle tally filters
-            norm = 1
-            for f in spec_t['filters']:
-                if f['type'] == 'cell':
-                    # Get cell volumes for normalization
-                    norm = [mesh.volumes_by_id[v].area for v in f['values']]
-                elif f['type'] == 'surface':
-                    # Get surface areas for normalization
-                    norm = [mesh.surfaces_by_id[v].area for v in f['values']]
-                elif f['type'] == 'material':
-                    raise NotImplementedError(
-                        'Material filter not implemented in postprocess yet.')
-
-                # Normalize the tally data
-                df['mean'] = df['mean'] / norm
-                df['std. dev.'] = df['std. dev.'] / norm
-
-            # Convert to xarray and add dimensions
-            t = xr.DataArray(
-                df.values[np.newaxis, :, :],  # shape: (1, r, c)
-                dims=["case", "row", "column"],
-                coords={
-                    "case": ["0"],
-                    "column": df.columns,
-                    "row": np.arange(df.shape[0])
-                },
-                name=spec_t['name']
-            )
-
-            # Save the tally data to a netCDF file
-            t.to_netcdf(f"benchmark_results.h5",
-                        group=f"{spec_t['name']}", mode='a')
-
-        # Add some metadata attributes
-        with h5py.File(f"benchmark_results.h5", "a") as f:
-            f.attrs['benchmark_name'] = self.name
-            # f.attrs['benchmark_version'] = self._benchmark_spec['metadata'].get(
-            #     'version', 'N/A')
-            # f.attrs['description'] = self._benchmark_spec['metadata'].get(
-            #     'description', 'N/A')
-            # f.attrs['title'] = self._benchmark_spec['metadata'].get(
-            #     'title', 'N/A')
-            # f.attrs['literature'] = self._benchmark_spec['metadata'].get(
-            #     'references', [])
-            f.attrs['code'] = f"openmc {sp.version[0]}.{sp.version[1]}.{sp.version[2]}"
+
+        _openmc_to_ofb(
+            spec_tallies=tallies_data,
+            statepoint=statepoint,
+            mesh=mesh
+        )
+
+        return
+
+    def _uncertainty_quantification(self, *args, **kwargs):
+        """Perform uncertainty quantification for the benchmark."""
+        uq_data = self._benchmark_spec['uncertainty_quantification']
+        tallies_data = self._benchmark_spec['tallies']
+
+        mesh = 'mesh.h5m'
+
+        # Run a TMC for every nuclide present in specifications
+        for n, r in zip(uq_data[0]['nuclides'], uq_data[0]['realizations']):
+            tmc_engine(model=self.model, realizations=r,
+                       lib_name='endfb_80', nuclide=n[0], reaction=None,
+                       perturb_xs=True, _is_benchmark=True, _spec_tallies=tallies_data,
+                       _mesh=mesh, *args, **kwargs)
 
         return
 
-    def run(self, *args, **kwargs):
+    def run(self, uq: bool = False, *args, **kwargs):
         """Run the benchmark simulation."""
-        # Run the OpenMC model
-        statepoint = self.model.run(*args, **kwargs)
 
-        # Post-process the results
-        self.postprocess(statepoint=statepoint)
+        # Check if benchmark_results.h5 already exists and delete it
+        path = Path("benchmark_results.h5")
+        if path.exists():
+            path.unlink()
+            warnings.warn(
+                f"Existing file '{path}' was found and deleted.", UserWarning)
+
+        # Run the model
+        if uq:
+            # Perform uncertainty quantification
+            self._uncertainty_quantification(*args, **kwargs)
+        else:
+            # Run the OpenMC model
+            # self._run_model(*args, **kwargs)
+            # Run the OpenMC model
+            statepoint = self.model.run(*args, **kwargs)
+            # Post-process the results
+            self._postprocess(statepoint=statepoint)
 
         return

src/openmc_fusion_benchmarks/benchmarks/benchmark_schema.yaml

@@ -43,7 +43,9 @@ components:
         irradiation_schedule:
           $ref: '#/components/schemas/IrradiationSchedule'
         uncertainty_quantification:
-          $ref: '#/components/schemas/UncertaintyQuantification'
+          type: array
+          items:
+            $ref: '#/components/schemas/UncertaintyQuantification'
 
     Metadata:
       type: object
@@ -312,21 +314,37 @@ components:
     UncertaintyQuantification:
       type: object
       properties:
-        # nuclear_data:
-        #   type: boolean
-        # materials:
-        #   type: object
-        #   properties:
-        #     density:
-        #       type: boolean
-        #     composition:
-        #       type: boolean
-        # source:
-        #   type: object
-        #   properties:
-        #     geometry:
-        #       type: boolean
-        #     angle:
-        #       type: boolean
-        #     energy:
-        #       type: boolean
+        realizations:
+          type: array
+          items:
+            type: integer
+      required: [realizations]
+      allOf:
+        - oneOf:
+          - type: object
+            properties:
+              type:
+                type: string
+                const: nuclear_data
+              nuclides:
+                type: array
+                items:
+                  type: array
+                  items:
+                    type: string
+              reactions:
+                type: array
+                items:
+                  type: string
+            required: [type, nuclides, reactions]
+
+          # - type: object
+          #   properties:
+          #     type:
+          #       type: string
+          #       const: material
+          #     materials:
+          #       type: array
+          #       items:
+          #         type: string
+          #   required: [type, materials]

src/openmc_fusion_benchmarks/benchmarks/oktavian_al/specifications.yaml

@@ -191,4 +191,11 @@ tallies:
                 4700000.,  4800000.,  4900000.,  5000000.,  5500000.,  6000000.,
                 6500000.,  7000000.,  7500000.,  8000000.,  8500000.,  9000000.,
                 9500000., 10000000., 10500000., 11000000.]
-        units: particle/src/cm**2
+        units: particle/src/cm**2
+
+uncertainty_quantification:
+  - type: nuclear_data
+    description: "Uncertainty in nuclear data"
+    nuclides: [[Al27]]
+    reactions: [total]
+    realizations: [500]

src/openmc_fusion_benchmarks/uq/__init__.py

@@ -0,0 +1,2 @@
+from .uq_utils import *
+from .tmc_engine import *

src/openmc_fusion_benchmarks/uq/tmc_engine.py

@@ -0,0 +1,96 @@
+from .uq_utils import perturb_xs_xml, get_nuclide_gnds, perturb_to_hdf5
+from pathlib import Path
+import openmc
+import numpy as np
+import xarray as xr
+
+from ..utils import _openmc_to_ofb, _save_result
+
+
+def tmc_engine(model: openmc.Model, realizations: int, lib_name: str, nuclide,
+               reaction: int = None, perturb_xs: bool = True,
+               _is_benchmark: bool = False, _mesh=None, _spec_tallies=None,
+               *args, **kwargs):
+    """Runs a TMC simulation on a given OpenMC model object. 
+    With perturb_xs=True it is possible to perturb the cross sections of a 
+    specific nuclide and reaction from a given nuclear data library 
+    automatically before the starting of the actual TMC simulation.
+    The results of the TMC simulation are stored in a .h5 file as OpenMC
+    tallies in Xarray DataArray format.
+
+    Parameters
+    ----------
+    model : openmc.Model
+        OpenMC model object to run TMC simulations on
+    realizations : int
+        Number of samples to run in the TMC simulation 
+        (i.e. number of times the xs is perturbed)
+    lib_name : str
+        Name of the nuclear data library to perturb the xs from 
+        (e.g. 'ENDF/B-VIII.0')
+    nuclide : str or int
+        Identifier of the nuclide for which the cross section is perturbed
+        (GNDS, ZAID or ZAM)
+    reaction : int, optional
+        MT value for the specific reaction to perturb, by default None
+    perturb_xs : bool, optional
+        Flag for the automatic generation of perturbed .h5 xs files right
+        before running the TMC simulation. Set to False if the use has already
+        a set of perturbed xs in .h5 format
+        to point to, by default True
+    """
+
+    # convert nuclide to gnds name
+    nuclide = get_nuclide_gnds(nuclide)
+    xs_file = f'cross_sections_mod.xml'
+
+    # runs sandy and generates perturbed xs only if perturb_xs is True
+    if perturb_xs:
+        perturb_to_hdf5(realizations, lib_name, nuclide, reaction, nprocesses=1,
+                        error=.001)
+
+    for n in np.arange(realizations):
+
+        directory = f"{nuclide}_{lib_name}"
+        xs_h5_file = f"{directory}/{nuclide}_{n}_{lib_name}.h5"
+        perturb_xs_xml(xs_file, xs_h5_file, nuclide)
+        openmc.config['cross_sections'] = xs_file
+
+        # run simulation
+        statepoint = model.run(*args, **kwargs)
+
+        # Postprocess results in case of standalone TMC run
+        sp = openmc.StatePoint(statepoint)
+
+        # realization_label = f'{nuclide}_{n}_{lib_name}'
+        realization_label = f'{nuclide}_{n}'
+        if _is_benchmark:
+            # Postprocess results in case of OFB benchmark model
+            _openmc_to_ofb(statepoint=sp,
+                           realization_label=realization_label,
+                           mesh=_mesh,
+                           spec_tallies=_spec_tallies,
+                           *args, **kwargs)
+        else:
+            # open tally and push to netcdf ofb style
+            for t in sp.tallies:
+                tally = sp.get_tally(id=t)
+                df = tally.get_pandas_dataframe()
+
+                # Convert to xarray and add dimensions
+                t = xr.DataArray(
+                    df.values[np.newaxis, :, :],  # shape: (1, r, c)
+                    dims=["realization", "row", "column"],
+                    coords={
+                        "realization": [realization_label],
+                        "column": df.columns,
+                        "row": np.arange(df.shape[0]),
+                    },
+                    name=t.name
+                )
+
+                _save_result(new_result=t, filename="tmc_results.h5",
+                             group=t.name, realization_label=f'{nuclide}_{n}_{lib_name}')
+
+        Path('summary.h5').unlink(missing_ok=True)
+        Path(statepoint).unlink(missing_ok=True)