UQ pipeline: TMCManager, TMCStatePoint and TMCTally

[SteSeg]

In this PR a new version of the TMC-UQ workflow is proposed, specifically for openmc Model objects.

Workflow

The workflow relies on the new TMCManager, the TMCStatePoint and the TMCTally classes.

The TMCManager class takes in an openmc.Model, a list of callables called perturbations, a number of realizations and a numpy random number generator object. The integer realizations specifies the number of realizations per perturbation to run in the total monte carlo run. Perturbations are external functions, that have to be written by the user, that perturb the model, and have to have the following structure:

def perturb_something(model=mymodel, rng=myrng):

    # Actual perturbing function
    def perturbation:
        # perturb something in the model according to user defined distribution and rng
        ...
        return model

    return perturbation

Once the TMCManager is instantiated , it is possile to run the tmc simulation with TMCManager.run() which takes similar arguments of openmc.run(). During the run the manager will build a folder structure as follows (a tmc/ folder with perturbation_{p} subfolders which in turn have realization_{r} subfolders. In every realization_{r} subfolder there are all the information about the specific openmc run associated with that specific perturbation and realization (i.e. model.xml, statepoint etc.)
The folder structure is saved in the tmc_manifest.jsonl file in the main working directory.
In the main working directory there will be also the tmc_statepoint file. It collects all the openmc tallies in xarray Dataset adding a dimention to the classic openmc.Tally object in order to stack all the perturbation and realization results of the same tally.

Here the folder structure:

cwd/
├── tmc_statepoint.r.h5
├── tmc_manifest.jsonl
└── tmc/
    ├── perturbation_0/
    │   ├── realization_0
    │   │   ├── model.xml
    │   │   ├── statepoint.h5
    │   │   └── summary.h5
    │   ├── realization_1/
    │   │   ├── model.xml
    │   │   └── ...
    │   ├── ...
    │   └── realization_r/
    ├── perturbation_1/
    │   └── ...
    ├── perturbation_3/
    ├── ...
    └── perturbation_p/  

The TMCStatePoint class is responsible for instantiating the tmc_statepoint file. The inspection (class methods and properties) is similar to the openmc.Statepoint object.

Single tallies in the tmc_statepoint are extracted via TMCStatepoint(...).get_tally(...) (similarly to openmc) as TMCTally objects. They can be treated similar to an openmc.Tally object but they have the perturbation times realization extra dimension. So, it is possible to get tally.name, TMCTally.scores, TMCTally.id info but the TMCTally.mean will provide the mean of the tmc run (i.e. the mean of the single means) and the TMCTally.std_dev will provide the standard deviation of the tmc run (i.e. the std_dev of the means). To get the statistics of the single openmc runs there is the TMCTally.realization_means and TMC.realization_std_devs that will provide such statistics embedded in a r times p (where r is the number of realizations and p is the number of perturbations) long array.
To get the most numerical information regarding the tmc run for a specific tally, there is the TMCTally.data that can be inspected and used.

Example usage

Here an example usage with an openmc model that has a water material and gets its density perturbed.

Build openmc model

import openmc_fusion_benchmark as ofb
import openmc 

# Build openmc model
model = openmc.Model()

water = openmc.Material(name='Water')
water.set_density('g/cm3', 1.0)
water.add_element('H', 0.6667)
water.add_element('O', 0.3333)

model.materials = openmc.Materials([water])
...

Define perturbing function(s)

# Define perturbing function (perturb water density):
def perturb_water_density(model, rng=None):
    # rng is a numpy random generator object 
    # if None, create a new one
    if rng is None:
        rng = np.random.default_rng()
    
    def perturbation(model, rng=rng):

        for mat in model.materials:
            if mat.name == 'Water':
                factor = 0.9 + 0.2 * rng.random()  # random factor between 0.9 and 1.1
                current_density = mat.density
                new_density = current_density * factor
                mat.set_density('g/cm3', new_density)

        return model
    
    return perturbation

Instantiate and run TMC

# Instantiate tmc manager
tmc_manager = ofb.uq.TMCManager(base_model=model, perturbations=[perturb_water_density], realizations=100, rng=np.random.default_rng(42))
# Run tmc
tmc_manager.run()

Extract and inspect TMC results

# Extract results
tsp = ofb.uq.TMCStatepoint('tmc_statepoint.100.h5')
# Print tallies available
print(tsp.tallies)

# Extract one tally
tmc_tally = tsp.get_tally(name='tally_name')

# Print tally metadata - including info on the realizations
print(tmc_tally.name, tmc_tally.filters, tmc_tally.nuclides, tmc_tally.realizations)

# Print TMC statistics (mean of means, std_dev of means)
print(tmc_tally.mean, tmc_tally.std_dev)

# Print all openmc single run statistics
print(tmc_tally.realization_means, tmc_tally.realization_std_devs)

Next steps

A possible next step is to allow the perturbation functions to have perturbation_id and name options and transfer such information in the tmc_statepoint file as metadata in order to be able to slice the TMCTally per perturbation (currently the tmc manager just adds one dimension r times p long

Spefications and ofb benchmarks currently not upgraded to tmc manager

Regarding running uq on the benchmarks of this repository, currently the specifications and the uq methods stick with the older tmc_engine function. We will update the pipeline to the TMCManager once it gets finalized and merged.

[coveralls]

Pull Request Test Coverage Report for Build 20931348909

Details

• 0 of 0 changed or added relevant lines in 0 files are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage decreased (-17.2%) to 62.189%

---

Totals
Change from base Build 20582698747:	-17.2%
Covered Lines:	500
Relevant Lines:	804

---

💛 - Coveralls