Local adjoint source for Random Ray

[j-fletcher]

Description

Adds ability to specify fixed, localized adjoint sources for the random ray solver. This capability will ultimately be useful for the addition of automated CADIS weight-windowing and source biasing using the adjoint random ray solver, and in this implementation, enables FW-CADIS weight windows to be generated for local variance reduction.

There are two possible methods of constructing local adjoint sources. The first is by the straightforward specification of a geometric region of interest and energy discretization desired for the adjoint source; the initial forward calculation proceeds as previously to set $Q^{\dagger} = 1/\phi$, but only within flat source regions in the region of interest. Elsewhere, the adjoint source simply remains 0. This approach, although still requiring an initial forward calculation, ensures that variance reduction is performed in all regions of phase space counted as part of the response (e.g., each individual detector and each energy group therein) uniformly. The user can designate a target openmc.Tallies object containing each tally on which variance reduction should be performed, then the weight windower will use geometric and energy constraints inferred from CellFilter, CellInstanceFilter, DistribcellFilter, UniverseFilter, MaterialFilter, and EnergyFilter instances on each tally to populate adjoint sources in flat source regions, as described in the method above.

The second is by user specification of the adjoint source term directly, using the distribution objects already available for constructing forward openmc.SourceBase instances, for example when information is available about a certain detector's response function. This approach can potentially result in unequal variance across different detector responses if multiple such sources are used to generate weight window and source biasing parameters in CADIS, but removes the need to perform an initial forward solve when pursuing local variance reduction.

This PR includes both of these capabilities, with the original implementation of the forward-weighted method by @jtramm. It includes an additional model in openmc.examples, based on the existing random_ray_three_region_cube with the addition of two "detector" volumes, which can be used to demonstrate both options as shown below.

Specifying Forward-Weighted Local Adjoint Sources for FW-CADIS

First, we'll import the test geometry and define a weight window mesh over the entire geometry.

import openmc
from openmc.examples import random_ray_three_region_cube_with_detectors

model = random_ray_three_region_cube_with_detectors()

ww_mesh = openmc.RegularMesh()
n = 7
width = 35.0
ww_mesh.dimension = (n, n, n)
ww_mesh.lower_left = (0.0, 0.0, 0.0)
ww_mesh.upper_right = (width, width, width)

Next, although there are several tallies included in the model, we'll pick out the tallies which cover our phase space regions of interest (detectors 1 and 2). We'll then add these to an FW-CADIS WeightWindowGenerator on the model and generate our weight windows for local variance reduction:

target_tallies = openmc.Tallies()
for tally in list(model.tallies):
    if tally.name in {"Detector 1 Tally", "Detector 2 Tally"}:
        target_tallies.append(tally)

wwg = openmc.WeightWindowGenerator(
    method="fw_cadis", 
    targets=target_tallies,
    mesh=ww_mesh, 
    max_realizations=model.settings.batches
)
model.settings.weight_window_generators = wwg
model.settings.random_ray['volume_estimator'] = 'naive'

model.run()

The resulting weight window lower bounds, depicted below, show a decreasing trend in the direction from the forward source (in the "bottom left" corner; $(X, Y, Z) = [0.0, 0.0, 0.0]\times [5.0, 5.0, 5.0]$) towards Detector 1 (centered at (2.5, 32.5, 2.5)) and especially towards Detector 2 (centered further from the source, at (32.5, 32.5, 32.5)), as desired for simultaneous local variance reduction on both tallies.

Specifying Localized Sources for Non-Forward-Weighted Adjoint Simulations

Outside of FW-CADIS weight window generation, this PR also adds the ability to specify localized sources for pure adjoint simulations. If supplied by the user, this removes the requirement to run an initial forward solve before performing adjoint transport, but in general the forward-weighted method is still recommended for applications like weight window generation with multiple detectors because the uniform reduction of variance in tally regions of interest is guaranteed.

At present, these sources are subject to the same constraints as all forward fixed sources for random ray: they must be isotropic in angle, use a discrete (i.e. multigroup) energy distribution, and the spatial distribution must either be a point, or constrained by Cell, Universe or Material. In this demonstration, we use a Box distribution identical to that specified on model.settings.random_ray['ray_source'], but constrain it to the flat source regions which make up Detector 1. As in the previous example, the problem is run with a single energy group.

model = random_ray_three_region_cube_with_detectors()

detector1_cells = model.geometry.get_cells_by_name("detector 1")

strengths = [1.0]
midpoints = [100.0]
energy_distribution = openmc.stats.Discrete(x=midpoints, p=strengths)

adj_space = openmc.stats.Box([0.0, 0.0, 0.0], [35.0, 35.0, 35.0], only_fissionable=False)
adj_source = openmc.IndependentSource(energy=energy_distribution, space=adj_space, constraints={'domains': detector1_cells})

model.settings.random_ray['adjoint'] = True
model.settings.random_ray['adjoint_source'] = adj_source
model.settings.random_ray['volume_estimator'] = 'naive'

model.run()

As shown below, the addition of the 'adjoint_source' entry has allowed us to map the adjoint flux with respect to Detector 1 throughout the geometry without running an initial forward calculation.

Fixes #3710

[jtramm]

PR is looking fantastic -- nice work on this! Just a few minor questions/comments to consider.

[jtramm]

It would be nice if we were able to enforce this restriction (if you're not already doing this?). Probably not feasible for the openmc.Tallies.export_to_xml call, but the openmc.Settings.export_to_xml call seems like it would have the needed info to determine if local variance reduction is active?

[j-fletcher]

The way I'm currently enforcing it is to perform an equivalence check (during Model.export_to_xml/export_to_model_xml) between the tally objects or IDs provided on WeightWindowGenerator.targets, and the tallies currently present in the model. Each of the tallies (or IDs) specified as targets must be equivalent to a tally (or its ID) on Model.tallies. While the WeightWindowGenerator is stored under Model.settings, the tallies are under Model.tallies, hence I wasn't able to put the check under Settings itself.

The nearest comparison that comes to mind is the use of tallies with a CellFilter--when a cell filter is added to a tally, the tally and filter are written to the XML file from Python without verifying that these cells are present in Model.geometry, and instead the check is performed when the filter (with the rest of the tallies subelement) is read off the XML file in C++. However, because the settings (and hence any WeightWindowGenerators) are read from XML before any tallies, I felt it would be cumbersome to perform the check on target tallies with the same timing. In contrast, Model._link_geometry_to_filters() meanwhile is called during Model.export_to_xml/export_to_model_xml, and it ensures that the reference geometry for each DistribcellFilter present in the Model.tallies is the same geometry present under Model.geometry. I felt that Model._assign_fw_cadis_tally_IDs() performed a similar enough task, but let me know if you think it should still be changed!

[jtramm]

I think the conditional check is backwards from the way the text states, as I think we just want to give the warning if we get a non-zero point source that was user-specified. The actual check here only gives the warning if it is a point source at the origin (which is likely the default and not user-specified).

[j-fletcher]

Agreed--the conditional is backwards from what might have been the original intention of this check, judging by the comment above. But I think it would be useful to give the warning if there is both a point source and domain(s) specified, whether or not the point source was assigned a non-default location by the user--so I'd propose removing this if-statement altogether and giving the warning either way.

[jtramm]

I recall adding this restriction in the cadis branch as I didn't yet have the nice mechanism for specifying targets that you do. In that branch, I had originally considered just marking any region that had a tally as a "target", but this became incompatible when a separate mesh filter was applied for visualizing the global flux, which is a super common use case. As such, I added this filter so as to weed out any mesh filters. The above 5 filters seemed to me to indicate that the user probably wanted something "local" rather than a global one, so worked at least for the purposes of an experimental branch.

In your implementation though, I wonder if we still need to impose this restriction on filter types? As you have now added the ability to explicitly set which tallies to use as targets, it seems like any tally (supported by random ray, which is already check at init) marked as a target should be acceptable? A user with a mesh tally for fluxes can simply not add that to the target list to avoid the global VR. Let me know if I'm missing something though!

[j-fletcher]

Great point--since filter compatibility is already checked at init this should be fine to remove.

[jtramm]

Along the same lines of the above comment: if we remove the restriction over which filters are acceptable, then I don't think we'd need this check any longer.

[j-fletcher]

Modified in accordance with previous comment.

[jtramm]

There had been some work done to simplify the logic of this function in develop (e.g., RandomRaySimulation::prepare_adjoint_simulation()). It looks like these housekeeping functions were removed and the openmc_run_random_ray() function has grown a lot in complexity. Was this intentional or just a byproduct from merging?

The logic in the current develop branch was quite simple:

void openmc_run_random_ray()
{
  //////////////////////////////////////////////////////////
  // Run forward simulation
  //////////////////////////////////////////////////////////

  if (openmc::mpi::master) {
    if (openmc::FlatSourceDomain::adjoint_) {
      openmc::FlatSourceDomain::adjoint_ = false;
      openmc::print_adjoint_header();
      openmc::FlatSourceDomain::adjoint_ = true;
    }
  }

  // Initialize OpenMC general data structures
  openmc_simulation_init();

  // Validate that inputs meet requirements for random ray mode
  if (openmc::mpi::master)
    openmc::validate_random_ray_inputs();

  // Initialize Random Ray Simulation Object
  openmc::RandomRaySimulation sim;

  // Initialize fixed sources, if present
  sim.apply_fixed_sources_and_mesh_domains();

  // Run initial random ray simulation
  sim.simulate();

  //////////////////////////////////////////////////////////
  // Run adjoint simulation (if enabled)
  //////////////////////////////////////////////////////////

  if (sim.adjoint_needed_) {
    // Setup for adjoint simulation
    sim.prepare_adjoint_simulation();

    // Run adjoint simulation
    sim.simulate();
  }
}

It would be nice to generally keep the same basic logic flow here as it's fairly high level and easy to read, and avoids a lot of the repeated calls to simulation_int, timer resets, etc, but let me know if you had a motivation for that.

[j-fletcher]

This was indeed an artifact of the merging process--I've now cleaned it up a bit so that openmc_run_random_ray() more closely resembles the above implementation. However I made a couple of changes as a result of the new possibility that a pure adjoint simulation could be run. Because of this, an initial forward calculation might not need to be run, meaning that the "is_first_simulation_" flag was not necessarily useful for all foreseeable cases. I also removed the "print_adjoint_header" function to simplify evaluating the conditions for when "ADJOINT FLUX SOLVE" or "FORWARD FLUX SOLVE" should be printed.

[jtramm]

Were the values in this test expected to change due to this PR?

[jtramm]

Is this change expected?

[jtramm]

Is this change expected?

[jtramm]

As most users are probably more interested in the local vs. global distinction, we might change the MAGIC and FW-CADIS headers to make the difference more explicit. E.g., "Generating Global or Local Weight Windows with FW-CADIS" and "Generating Global Weight Windows with MAGIC"

[j-fletcher]

Sounds good to me!

[jtramm]

"global" instead of "local"

[j-fletcher]

Good catch!