From 6a2402d43c179c8948c6a032008e03bd8efe8c0d Mon Sep 17 00:00:00 2001 From: Christopher Date: Fri, 6 Feb 2026 16:26:50 -0500 Subject: [PATCH] coincident source --- include/openmc/source.h | 55 ++++++++- openmc/source.py | 191 ++++++++++++++++++++++++++++++++ src/particle.cpp | 14 +-- src/simulation.cpp | 8 +- src/source.cpp | 180 ++++++++++++++++++++++++++++-- tests/unit_tests/test_source.py | 124 +++++++++++++++++++++ 6 files changed, 548 insertions(+), 24 deletions(-) diff --git a/include/openmc/source.h b/include/openmc/source.h index 2f32aa2a058..066421ad856 100644 --- a/include/openmc/source.h +++ b/include/openmc/source.h @@ -77,6 +77,23 @@ class Source { //! \return Sampled site virtual SourceSite sample(uint64_t* seed) const = 0; + //! Sample one or more source sites (without applying constraints) + // + //! For most sources this returns a single site. CorrelatedSource overrides + //! this to return multiple correlated sites. + //! \param[inout] seed Pseudorandom seed pointer + //! \return Vector of sampled sites + virtual vector sample_sites(uint64_t* seed) const + { + return {sample(seed)}; + } + + //! Sample one or more source sites and apply constraints + // + //! \param[inout] seed Pseudorandom seed pointer + //! \return Vector of sampled sites + vector sample_sites_with_constraints(uint64_t* seed) const; + static unique_ptr create(pugi::xml_node node); protected: @@ -250,6 +267,36 @@ class MeshSource : public Source { vector> sources_; //!< Source distributions }; +//============================================================================== +//! Source that emits multiple correlated particles per source event +//! +//! All particles share the same position and time, but each has its own +//! particle type, energy, and angular distributions. Used for simulating +//! correlated emissions such as Co-60 summation peaks. +//============================================================================== + +class CorrelatedSource : public Source { +public: + // Constructors + explicit CorrelatedSource(pugi::xml_node node); + + //! Sample a single source site (returns first correlated particle) + //! \param[inout] seed Pseudorandom seed pointer + //! \return Sampled site + SourceSite sample(uint64_t* seed) const override; + + //! Sample all correlated source sites + //! \param[inout] seed Pseudorandom seed pointer + //! \return Vector of sampled sites sharing the same position and time + vector sample_sites(uint64_t* seed) const override; + +private: + UPtrSpace space_; //!< Shared spatial distribution + UPtrDist time_; //!< Shared time distribution + vector> sources_; //!< Sub-source distributions + vector probabilities_; //!< Emission probability per sub-source +}; + //============================================================================== // Functions //============================================================================== @@ -257,11 +304,11 @@ class MeshSource : public Source { //! Initialize source bank from file/distribution extern "C" void initialize_source(); -//! Sample a site from all external source distributions in proportion to their -//! source strength +//! Sample one or more sites from all external source distributions in +//! proportion to their source strength //! \param[inout] seed Pseudorandom seed pointer -//! \return Sampled source site -SourceSite sample_external_source(uint64_t* seed); +//! \return Vector of sampled source sites (multiple for correlated sources) +vector sample_external_source(uint64_t* seed); void free_memory_source(); diff --git a/openmc/source.py b/openmc/source.py index a11cfd6c8d8..4fc24a46eb2 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -203,6 +203,8 @@ def from_xml_element(cls, elem: ET.Element, meshes=None) -> SourceBase: return FileSource.from_xml_element(elem) elif source_type == 'mesh': return MeshSource.from_xml_element(elem, meshes) + elif source_type == 'correlated': + return CorrelatedSource.from_xml_element(elem, meshes) else: raise ValueError( f'Source type {source_type} is not recognized') @@ -658,6 +660,195 @@ def from_xml_element(cls, elem: ET.Element, meshes) -> openmc.MeshSource: return cls(mesh, sources, constraints=constraints) +class CorrelatedSource(SourceBase): + """A source that emits multiple correlated particles per source event. + + All particles share the same spatial position and time, but each has + independent particle type, energy, and angular distributions. This is + useful for simulating correlated emissions such as Co-60 summation peaks. + + .. versionadded:: 0.15.1 + + Parameters + ---------- + space : openmc.stats.Spatial, optional + Shared spatial distribution for all emitted particles + time : openmc.stats.Univariate, optional + Shared time distribution for all emitted particles + sources : list of openmc.IndependentSource + Sub-sources defining particle type, energy, and angle for each + correlated particle. Must contain at least 2 sources. + strength : float + Strength of the source + probabilities : list of float, optional + Emission probability for each sub-source, values in (0, 1]. + Default is 1.0 (always emit) for every sub-source. + constraints : dict, optional + Constraints on sampled source particles. + + Attributes + ---------- + space : openmc.stats.Spatial or None + Shared spatial distribution + time : openmc.stats.Univariate or None + Shared time distribution + sources : list of openmc.IndependentSource + Sub-source distributions + probabilities : list of float + Emission probability per sub-source + strength : float + Strength of the source + type : str + Indicator of source type: 'correlated' + + """ + + def __init__( + self, + space: openmc.stats.Spatial | None = None, + time: openmc.stats.Univariate | None = None, + sources: list[IndependentSource] | None = None, + strength: float = 1.0, + probabilities: list[float] | None = None, + constraints: dict[str, Any] | None = None + ): + super().__init__(strength=strength, constraints=constraints) + + self._space = None + self._time = None + self._sources = [] + self._probabilities = None + + if space is not None: + self.space = space + if time is not None: + self.time = time + if sources is not None: + self.sources = sources + if probabilities is not None: + self.probabilities = probabilities + + @property + def type(self) -> str: + return 'correlated' + + @property + def space(self): + return self._space + + @space.setter + def space(self, space): + cv.check_type('spatial distribution', space, Spatial) + self._space = space + + @property + def time(self): + return self._time + + @time.setter + def time(self, time): + cv.check_type('time distribution', time, Univariate) + self._time = time + + @property + def sources(self): + return self._sources + + @sources.setter + def sources(self, sources): + cv.check_type('sub-sources', sources, list) + for s in sources: + cv.check_type('sub-source', s, IndependentSource) + if len(sources) < 2: + raise ValueError('A correlated source must have at least 2 ' + 'sub-sources.') + self._sources = list(sources) + + @property + def probabilities(self): + return self._probabilities + + @probabilities.setter + def probabilities(self, probabilities): + cv.check_type('probabilities', probabilities, list) + for p in probabilities: + cv.check_type('probability', p, Real) + if p <= 0.0 or p > 1.0: + raise ValueError( + f'Probability {p} is not in the range (0, 1].') + if self._sources and len(probabilities) != len(self._sources): + raise ValueError( + f'Length of probabilities ({len(probabilities)}) must match ' + f'number of sub-sources ({len(self._sources)}).') + self._probabilities = list(probabilities) + + def populate_xml_element(self, element): + """Add necessary source information to an XML element + + Returns + ------- + element : lxml.etree._Element + XML element containing source data + + """ + if self.space is not None: + element.append(self.space.to_xml_element()) + if self.time is not None: + element.append(self.time.to_xml_element('time')) + probs = self.probabilities + for i, src in enumerate(self.sources): + sub_elem = src.to_xml_element() + if probs is not None and probs[i] != 1.0: + sub_elem.set("probability", str(probs[i])) + element.append(sub_elem) + + @classmethod + def from_xml_element(cls, elem: ET.Element, meshes=None) -> CorrelatedSource: + """Generate correlated source from an XML element + + Parameters + ---------- + elem : lxml.etree._Element + XML element + meshes : dict, optional + Dictionary with mesh IDs as keys and openmc.MeshBase instances as + values + + Returns + ------- + openmc.CorrelatedSource + Source generated from XML element + + """ + constraints = cls._get_constraints(elem) + source = cls(constraints=constraints) + + strength = get_text(elem, 'strength') + if strength is not None: + source.strength = float(strength) + + space = elem.find('space') + if space is not None: + source.space = Spatial.from_xml_element(space, meshes) + + time = elem.find('time') + if time is not None: + source.time = Univariate.from_xml_element(time) + + sub_sources = [] + probabilities = [] + for e in elem.iterchildren('source'): + sub_sources.append(IndependentSource.from_xml_element(e)) + prob_str = e.get('probability') + probabilities.append(float(prob_str) if prob_str is not None else 1.0) + if sub_sources: + source.sources = sub_sources + if any(p != 1.0 for p in probabilities): + source.probabilities = probabilities + + return source + + def Source(*args, **kwargs): """ A function for backward compatibility of sources. Will be removed in the diff --git a/src/particle.cpp b/src/particle.cpp index a035e76b917..de09fe7abc4 100644 --- a/src/particle.cpp +++ b/src/particle.cpp @@ -853,13 +853,13 @@ void Particle::write_restart() const settings::n_particles + simulation::work_index[mpi::rank] + i; uint64_t seed = init_seed(id, STREAM_SOURCE); - // re-sample source site - auto site = sample_external_source(&seed); - write_dataset(file_id, "weight", site.wgt); - write_dataset(file_id, "energy", site.E); - write_dataset(file_id, "xyz", site.r); - write_dataset(file_id, "uvw", site.u); - write_dataset(file_id, "time", site.time); + // re-sample source site (use primary site only) + auto sites = sample_external_source(&seed); + write_dataset(file_id, "weight", sites[0].wgt); + write_dataset(file_id, "energy", sites[0].E); + write_dataset(file_id, "xyz", sites[0].r); + write_dataset(file_id, "uvw", sites[0].u); + write_dataset(file_id, "time", sites[0].time); } // Close file diff --git a/src/simulation.cpp b/src/simulation.cpp index 18e40a8bc73..a957ec801a1 100644 --- a/src/simulation.cpp +++ b/src/simulation.cpp @@ -583,8 +583,12 @@ void initialize_history(Particle& p, int64_t index_source) simulation::work_index[mpi::rank] + index_source; uint64_t seed = init_seed(id, STREAM_SOURCE); // sample from external source distribution or custom library then set - auto site = sample_external_source(&seed); - p.from_source(&site); + auto sites = sample_external_source(&seed); + p.from_source(&sites[0]); + // For correlated sources, add extra particles to secondary bank + for (size_t i = 1; i < sites.size(); ++i) { + p.secondary_bank().push_back(sites[i]); + } } p.current_work() = index_source; diff --git a/src/source.cpp b/src/source.cpp index 8bd9c789352..c527271ddf8 100644 --- a/src/source.cpp +++ b/src/source.cpp @@ -95,6 +95,8 @@ unique_ptr Source::create(pugi::xml_node node) return make_unique(node); } else if (source_type == "mesh") { return make_unique(node); + } else if (source_type == "correlated") { + return make_unique(node); } else { fatal_error(fmt::format("Invalid source type '{}' found.", source_type)); } @@ -227,6 +229,49 @@ SourceSite Source::sample_with_constraints(uint64_t* seed) const return site; } +vector Source::sample_sites_with_constraints(uint64_t* seed) const +{ + bool accepted = false; + static int64_t n_reject = 0; + static int64_t n_accept = 0; + vector sites; + + while (!accepted) { + // Sample all sites from the source + sites = this->sample_sites(seed); + + if (constraints_applied()) { + accepted = true; + } else { + // Check whether all sampled sites satisfy constraints + accepted = true; + for (const auto& site : sites) { + if (!satisfies_spatial_constraints(site.r) || + !satisfies_energy_constraints(site.E) || + !satisfies_time_constraints(site.time)) { + accepted = false; + break; + } + } + + if (!accepted) { + ++n_reject; + check_rejection_fraction(n_reject, n_accept); + + if (rejection_strategy_ == RejectionStrategy::KILL) { + accepted = true; + for (auto& site : sites) { + site.wgt = 0.0; + } + } + } + } + } + + ++n_accept; + return sites; +} + bool Source::satisfies_energy_constraints(double E) const { return E > energy_bounds_.first && E < energy_bounds_.second; @@ -623,6 +668,111 @@ SourceSite MeshSource::sample(uint64_t* seed) const return source(element)->sample_with_constraints(seed); } +//============================================================================== +// CorrelatedSource implementation +//============================================================================== + +CorrelatedSource::CorrelatedSource(pugi::xml_node node) : Source(node) +{ + // Correlated sources are only valid for fixed-source simulations + if (settings::run_mode == RunMode::EIGENVALUE) { + fatal_error("Correlated sources cannot be used in eigenvalue mode."); + } + + // Read shared spatial distribution + if (check_for_node(node, "space")) { + space_ = SpatialDistribution::create(node.child("space")); + } else { + space_ = UPtrSpace {new SpatialPoint()}; + } + + // Read shared time distribution + if (check_for_node(node, "time")) { + pugi::xml_node node_dist = node.child("time"); + time_ = distribution_from_xml(node_dist); + } else { + double T[] {0.0}; + double p[] {1.0}; + time_ = UPtrDist {new Discrete {T, p, 1}}; + } + + // Read child elements as sub-sources + for (auto source_node : node.children("source")) { + // Read emission probability (default 1.0) + double prob = 1.0; + if (source_node.attribute("probability")) { + prob = std::stod(source_node.attribute("probability").value()); + if (prob <= 0.0 || prob > 1.0) { + fatal_error("Sub-source probability must be in (0, 1]."); + } + } + + auto src = Source::create(source_node); + if (auto ptr = dynamic_cast(src.get())) { + src.release(); + sources_.emplace_back(ptr); + } else { + fatal_error( + "Sub-sources of a correlated source must be IndependentSource."); + } + probabilities_.push_back(prob); + } + + // Validate at least 2 sub-sources + if (sources_.size() < 2) { + fatal_error("A correlated source must have at least 2 sub-sources."); + } +} + +SourceSite CorrelatedSource::sample(uint64_t* seed) const +{ + auto sites = sample_sites(seed); + return sites[0]; +} + +vector CorrelatedSource::sample_sites(uint64_t* seed) const +{ + // Sample shared position once + auto [r, r_wgt] = space_->sample(seed); + + // Sample shared time once + auto [time, time_wgt] = time_->sample(seed); + + // Build a site for each sub-source, applying emission probabilities + vector sites; + sites.reserve(sources_.size()); + + for (size_t i = 0; i < sources_.size(); ++i) { + // Roll against emission probability for this sub-source + if (prn(seed) >= probabilities_[i]) + continue; + + // Sample from the sub-source to get particle type, energy, angle + SourceSite site = sources_[i]->sample(seed); + + // Override position and time with shared values + site.r = r; + site.time = time; + + // Apply shared spatial and time weights + site.wgt *= (r_wgt * time_wgt); + + sites.push_back(site); + } + + // If all rolls failed, push first sub-source with zero weight so the + // history can still initialize + if (sites.empty()) { + SourceSite site = sources_[0]->sample(seed); + site.r = r; + site.time = time; + site.wgt = 0.0; + sites.push_back(site); + } + + return sites; +} + //============================================================================== // Non-member functions //============================================================================== @@ -639,8 +789,9 @@ void initialize_source() simulation::work_index[mpi::rank] + i + 1; uint64_t seed = init_seed(id, STREAM_SOURCE); - // sample external source distribution - simulation::source_bank[i] = sample_external_source(&seed); + // sample external source distribution (store only primary site) + auto sites = sample_external_source(&seed); + simulation::source_bank[i] = sites[0]; } // Write out initial source @@ -653,7 +804,7 @@ void initialize_source() } } -SourceSite sample_external_source(uint64_t* seed) +vector sample_external_source(uint64_t* seed) { // Sample from among multiple source distributions int i = 0; @@ -666,26 +817,32 @@ SourceSite sample_external_source(uint64_t* seed) } } - // Sample source site from i-th source distribution - SourceSite site {model::external_sources[i]->sample_with_constraints(seed)}; + // Sample source sites from i-th source distribution + vector sites { + model::external_sources[i]->sample_sites_with_constraints(seed)}; // For uniform source sampling, multiply the weight by the ratio of the actual // probability of sampling source i to the biased probability of sampling // source i, which is (strength_i / total_strength) / (1 / n) if (n_sources > 1 && settings::uniform_source_sampling) { double total_strength = model::external_sources_probability.integral(); - site.wgt *= + double wgt_factor = model::external_sources[i]->strength() * n_sources / total_strength; + for (auto& site : sites) { + site.wgt *= wgt_factor; + } } // If running in MG, convert site.E to group if (!settings::run_CE) { - site.E = lower_bound_index(data::mg.rev_energy_bins_.begin(), - data::mg.rev_energy_bins_.end(), site.E); - site.E = data::mg.num_energy_groups_ - site.E - 1.; + for (auto& site : sites) { + site.E = lower_bound_index(data::mg.rev_energy_bins_.begin(), + data::mg.rev_energy_bins_.end(), site.E); + site.E = data::mg.num_energy_groups_ - site.E - 1.; + } } - return site; + return sites; } void free_memory_source() @@ -712,7 +869,8 @@ extern "C" int openmc_sample_external_source( auto sites_array = static_cast(sites); for (size_t i = 0; i < n; ++i) { - sites_array[i] = sample_external_source(seed); + auto sampled = sample_external_source(seed); + sites_array[i] = sampled[0]; } return 0; } diff --git a/tests/unit_tests/test_source.py b/tests/unit_tests/test_source.py index bb8a1b78528..ff305a4a218 100644 --- a/tests/unit_tests/test_source.py +++ b/tests/unit_tests/test_source.py @@ -125,6 +125,130 @@ def test_source_dlopen(): assert 'library' in elem.attrib +def test_correlated_source_xml_roundtrip(): + # Create a correlated source (e.g., Co-60 two-gamma emission) + gamma1 = openmc.IndependentSource( + energy=openmc.stats.Discrete([1.17e6], [1.0]), + angle=openmc.stats.Isotropic(), + particle='photon' + ) + gamma2 = openmc.IndependentSource( + energy=openmc.stats.Discrete([1.33e6], [1.0]), + angle=openmc.stats.Isotropic(), + particle='photon' + ) + space = openmc.stats.Point((1.0, 2.0, 3.0)) + time = openmc.stats.Discrete([0.0], [1.0]) + src = openmc.CorrelatedSource( + space=space, + time=time, + sources=[gamma1, gamma2], + strength=5.0 + ) + + # Write to XML element + elem = src.to_xml_element() + + # Read back from XML element + new_src = openmc.SourceBase.from_xml_element(elem) + assert isinstance(new_src, openmc.CorrelatedSource) + assert new_src.type == 'correlated' + assert new_src.strength == approx(5.0) + + # Check spatial distribution preserved + assert isinstance(new_src.space, openmc.stats.Point) + np.testing.assert_allclose(new_src.space.xyz, [1.0, 2.0, 3.0]) + + # Check time distribution preserved + assert isinstance(new_src.time, openmc.stats.Discrete) + np.testing.assert_allclose(new_src.time.x, [0.0]) + np.testing.assert_allclose(new_src.time.p, [1.0]) + + # Check sub-sources preserved + assert len(new_src.sources) == 2 + for i, (orig, new) in enumerate(zip([gamma1, gamma2], new_src.sources)): + assert isinstance(new, openmc.IndependentSource) + assert new.particle == orig.particle + assert isinstance(new.energy, openmc.stats.Discrete) + np.testing.assert_allclose(new.energy.x, orig.energy.x) + np.testing.assert_allclose(new.energy.p, orig.energy.p) + + +def test_correlated_source_probabilities_roundtrip(): + # Create a correlated source with emission probabilities + gamma1 = openmc.IndependentSource( + energy=openmc.stats.Discrete([1.17e6], [1.0]), + angle=openmc.stats.Isotropic(), + particle='photon' + ) + gamma2 = openmc.IndependentSource( + energy=openmc.stats.Discrete([1.33e6], [1.0]), + angle=openmc.stats.Isotropic(), + particle='photon' + ) + src = openmc.CorrelatedSource( + space=openmc.stats.Point((0, 0, 0)), + sources=[gamma1, gamma2], + probabilities=[1.0, 0.5] + ) + + # Write to XML and read back + elem = src.to_xml_element() + new_src = openmc.SourceBase.from_xml_element(elem) + assert isinstance(new_src, openmc.CorrelatedSource) + assert new_src.probabilities == [1.0, 0.5] + assert len(new_src.sources) == 2 + + # Verify the probability attribute is only written when != 1.0 + sub_elems = list(elem.iterchildren('source')) + assert sub_elems[0].get('probability') is None # 1.0 omitted + assert sub_elems[1].get('probability') == '0.5' + + +def test_correlated_source_probabilities_default(): + # Without probabilities, they should be None + gamma1 = openmc.IndependentSource(particle='photon') + gamma2 = openmc.IndependentSource(particle='photon') + src = openmc.CorrelatedSource( + space=openmc.stats.Point(), + sources=[gamma1, gamma2] + ) + assert src.probabilities is None + + # XML roundtrip should also give None + elem = src.to_xml_element() + new_src = openmc.SourceBase.from_xml_element(elem) + assert new_src.probabilities is None + + +def test_correlated_source_validation(): + # Must have at least 2 sub-sources + gamma1 = openmc.IndependentSource(particle='photon') + with pytest.raises(ValueError, match='at least 2'): + openmc.CorrelatedSource(sources=[gamma1]) + + +def test_correlated_source_probabilities_validation(): + gamma1 = openmc.IndependentSource(particle='photon') + gamma2 = openmc.IndependentSource(particle='photon') + src = openmc.CorrelatedSource( + space=openmc.stats.Point(), + sources=[gamma1, gamma2] + ) + + # Wrong length + with pytest.raises(ValueError, match='Length of probabilities'): + src.probabilities = [0.5] + + # Out of range (0) + with pytest.raises(ValueError, match='not in the range'): + src.probabilities = [0.0, 0.5] + + # Out of range (> 1) + with pytest.raises(ValueError, match='not in the range'): + src.probabilities = [1.5, 0.5] + + def test_source_xml_roundtrip(): # Create a source and write to an XML element space = openmc.stats.Box([-5., -5., -5.], [5., 5., 5.])