OpenMC/tests/unit_tests/test_source_biasing.py
Jack Fletcher 0486e433d2
Source biasing capabilities (#3460)
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
2026-01-12 09:51:12 -06:00

276 lines
9.5 KiB
Python

"""Tests for source biasing using C++ sampling routines via openmc.lib
This test module validates that the C++ distribution sampling implementations
correctly handle both unbiased and biased sampling when used in source
definitions. Each test:
1. Creates a minimal model with a source using a specific energy distribution
2. Uses model.sample_external_source() to generate samples via openmc.lib
3. Extracts energies from the returned particle list
4. Validates that:
- Unbiased sampling produces the expected mean
- Biased sampling with importance weighting produces the expected mean
- Weights are correctly applied (non-unity for biased case)
These tests complement the Python-level tests in test_stats.py by exercising
the full C++ sampling codepath that is used during actual simulations.
"""
import numpy as np
import pytest
import openmc
from tests.unit_tests import assert_sample_mean
@pytest.fixture
def model():
"""Create a minimal model for source sampling tests."""
sphere = openmc.Sphere(r=100.0, boundary_type='vacuum')
cell = openmc.Cell(region=-sphere)
geometry = openmc.Geometry([cell])
settings = openmc.Settings(particles=100, batches=1)
space = openmc.stats.Point()
angle = openmc.stats.Monodirectional((1.0, 0.0, 0.0))
settings.source = openmc.IndependentSource(space=space, angle=angle)
return openmc.Model(geometry=geometry, settings=settings)
@pytest.mark.flaky(reruns=1)
def test_discrete(run_in_tmpdir, model):
"""Test Discrete distribution sampling via C++ routines."""
vals = np.array([1.0, 2.0, 3.0])
probs = np.array([0.1, 0.7, 0.2])
exp_mean = (vals * probs).sum()
# Create source with discrete energy distribution
model.settings.source[0].energy = energy_dist = openmc.stats.Discrete(vals, probs)
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, exp_mean)
# Sample from biased distribution
energy_dist.bias = np.array([0.2, 0.1, 0.7])
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, exp_mean)
assert np.any(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_uniform(run_in_tmpdir, model):
"""Test Uniform distribution sampling via C++ routines."""
a, b = 5.0, 10.0
exp_mean = 0.5 * (a + b)
# Create source with uniform energy distribution
model.settings.source[0].energy = energy_dist = openmc.stats.Uniform(a, b)
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, exp_mean)
# Sample from biased distribution
energy_dist.bias = openmc.stats.PowerLaw(a, b, 2)
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, exp_mean)
assert np.any(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_powerlaw(run_in_tmpdir, model):
"""Test PowerLaw distribution sampling via C++ routines."""
a, b, n = 1.0, 20.0, 2.0
# Determine mean of distribution
exp_mean = (n+1)*(b**(n+2) - a**(n+2))/((n+2)*(b**(n+1) - a**(n+1)))
# Create source with powerlaw energy distribution
model.settings.source[0].energy = energy_dist = openmc.stats.PowerLaw(a, b, n)
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, exp_mean)
# Sample from biased distribution
energy_dist.bias = openmc.stats.Uniform(a, b)
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, exp_mean)
assert np.any(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_maxwell(run_in_tmpdir, model):
"""Test Maxwell distribution sampling via C++ routines."""
theta = 1.2895e6
exp_mean = 3/2 * theta
# Create source with Maxwell energy distribution
model.settings.source[0].energy = energy_dist = openmc.stats.Maxwell(theta)
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, exp_mean)
# Sample from biased distribution
energy_dist.bias = openmc.stats.Maxwell(theta * 1.1)
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, exp_mean)
assert np.any(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_watt(run_in_tmpdir, model):
"""Test Watt distribution sampling via C++ routines."""
a, b = 0.965e6, 2.29e-6
exp_mean = 3/2 * a + a**2 * b / 4
# Create source with Watt energy distribution
model.settings.source[0].energy = energy_dist = openmc.stats.Watt(a, b)
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, exp_mean)
# Sample from biased distribution
energy_dist.bias = openmc.stats.Watt(a*1.05, b)
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, exp_mean)
assert np.any(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_tabular(run_in_tmpdir, model):
"""Test Tabular distribution sampling via C++ routines."""
# Test linear-linear sampling
x = np.array([0.0, 5.0, 7.0, 10.0])
p = np.array([10.0, 20.0, 5.0, 6.0])
# Create tabular distribution and normalize to get expected mean
model.settings.source[0].energy = energy_dist = openmc.stats.Tabular(x, p, 'linear-linear')
energy_dist.normalize()
exp_mean = energy_dist.mean()
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, exp_mean)
# Sample from biased distribution
energy_dist.bias = openmc.stats.Uniform(x[0], x[-1])
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, exp_mean)
assert np.any(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_mixture(run_in_tmpdir, model):
"""Test Mixture distribution sampling via C++ routines."""
d1 = openmc.stats.Uniform(0, 5)
d2 = openmc.stats.Uniform(3, 7)
p = [0.5, 0.5]
# Create mixture energy distribution
model.settings.source[0].energy = energy_dist = openmc.stats.Mixture(p, [d1, d2])
exp_mean = (2.5 + 5.0) / 2
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, exp_mean)
# Sample using biased sub-distribution
energy_dist.distribution[0].bias = openmc.stats.PowerLaw(0, 5, 2)
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, exp_mean)
assert np.any(weights != 1.0)
@pytest.mark.flaky(reruns=1)
def test_normal(run_in_tmpdir, model):
"""Test Normal distribution sampling via C++ routines."""
mean_val = 25.0
std_dev = 2.0
# Create source with normal energy distribution
model.settings.source[0].energy = energy_dist = openmc.stats.Normal(mean_val, std_dev)
# Sample using C++ routines and extract energies
n_samples = 10_000
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
# Check unbiased mean
assert_sample_mean(energies, mean_val)
# Sample from biased distribution
energy_dist.bias = openmc.stats.Normal(mean_val * 1.1, std_dev)
particles = model.sample_external_source(n_samples)
energies = np.array([p.E for p in particles])
weights = np.array([p.wgt for p in particles])
# Check biased weighted mean
weighted_energies = energies * weights
assert_sample_mean(weighted_energies, mean_val)
assert np.any(weights != 1.0)