Merge pull request #1257 from makeclean/dlopen_source

Dynamically loaded custom source function

docs/source/io_formats/settings.rst

@@ -142,7 +142,7 @@ materials in the problem and is specified using a :ref:`mesh_element`.
 ----------------------------
 
 Determines whether to use event-based parallelism instead of the default
-history-based parallelism. 
+history-based parallelism.
 
   *Default*: false
 
@@ -459,6 +459,15 @@ attributes/sub-elements:
 
     *Default*: None
 
+  :library:
+    If this attribute is given, it indicates that the source is to be
+    instantiated from an externally compiled source function. This source can be
+    as complex as is required to define the source for your problem. The only
+    requirement is that there is a function called ``sample_source()``. More
+    documentation on how to build sources can be found in :ref:`custom_source`.
+
+    *Default*: None
+
   :space:
     An element specifying the spatial distribution of source sites. This element
     has the following attributes:
@@ -591,6 +600,67 @@ attributes/sub-elements:
 
     *Default*: false
 
+.. _custom_source:
+
+Custom Sources
+++++++++++++++
+
+It is often the case that one may wish to simulate a complex source
+distribution, which may include physics not present within OpenMC or to be phase
+space complex. It is possible to define a complex source with an externally
+defined source function that is loaded at runtime. A simple example source is
+shown below.
+
+.. code-block:: c++
+
+   #include "openmc/random_lcg.h"
+   #include "openmc/source.h"
+   #include "openmc/particle.h"
+
+   // you must have external C linkage here
+   extern "C" openmc::Particle::Bank sample_source(uint64_t* seed) {
+     openmc::Particle::Bank particle;
+     // weight
+     particle.particle = openmc::Particle::Type::neutron;
+     particle.wgt = 1.0;
+     // position
+     double angle = 2.0 * M_PI * openmc::prn(seed);
+     double radius = 3.0;
+     particle.r.x = radius * std::cos(angle);
+     particle.r.y = radius * std::sin(angle);
+     particle.r.z = 0.0;
+     // angle
+     particle.u = {1.0, 0.0, 0.0};
+     particle.E = 14.08e6;
+     particle.delayed_group = 0;
+     return particle;
+  }
+
+The above source, creates 14.08 MeV neutrons, with an istropic direction
+vector but distributed in a ring with a 3 cm radius. This routine is
+not particularly complex, but should serve as an example upon which to build
+more complicated sources.
+
+  .. note:: The function signature must be declared to be extern "C".
+
+  .. note:: You should only use the openmc::prn() random number generator
+
+In order to build your external source, you will need to link it against the
+OpenMC shared library. This can be done by writing a CMakeLists.txt file:
+
+.. code-block:: cmake
+
+   cmake_minimum_required(VERSION 3.3 FATAL_ERROR)
+   project(openmc_sources CXX)
+   add_library(source SHARED source_ring.cpp)
+   find_package(OpenMC REQUIRED HINTS <path to openmc>)
+   target_link_libraries(source OpenMC::libopenmc)
+
+After running ``cmake`` and ``make``, you will have a libsource.so (or .dylib)
+file in your build directory. Setting the :attr:`openmc.Source.library`
+attribute to the path of this shared library will indicate that it should be
+used for sampling source particles at runtime.
+
 .. _univariate:
 
 Univariate Probability Distributions

examples/xml/custom_source/CMakeLists.txt

@@ -0,0 +1,8 @@
+cmake_minimum_required(VERSION 3.3 FATAL_ERROR)
+project(openmc_sources CXX)
+add_library(source SHARED source_ring.cpp)
+find_package(OpenMC REQUIRED)
+if (OpenMC_FOUND)
+  message(STATUS "Found OpenMC: ${OpenMC_DIR}")
+endif()
+target_link_libraries(source OpenMC::libopenmc)

examples/xml/custom_source/geometry.xml

@@ -0,0 +1,15 @@
+<?xml version="1.0"?>
+<geometry>
+
+  <!-- Definition of Cells -->
+  <cell id="1" universe="0" fill="37" region="-2" />
+  <cell id="100" universe="37" material="40" region="-1" />
+  <cell id="101" universe="37" material="41" region="1" />
+  <cell id="2" universe="0" material="41" region="2 -3" />
+
+  <!-- Defition of Surfaces -->
+  <surface id="1" type="z-cylinder" coeffs="0 0 7" />
+  <surface id="2" type="z-cylinder" coeffs="0 0 9" />
+  <surface id="3" type="z-cylinder" coeffs="0 0 11" boundary="vacuum" />
+
+</geometry>

examples/xml/custom_source/materials.xml

@@ -0,0 +1,16 @@
+<?xml version="1.0"?>
+<materials>
+
+  <material id="40">
+    <density value="4.5" units="g/cc" />
+    <nuclide name="U235" ao="1.0" />
+  </material>
+
+  <material id="41">
+    <density value="1.0" units="g/cc" />
+    <nuclide name="H1" ao="2.0" />
+    <nuclide name="O16" ao="1.0" />
+    <sab name="c_H_in_H2O"/>
+  </material>
+
+</materials>

examples/xml/custom_source/settings.xml

@@ -0,0 +1,14 @@
+<?xml version="1.0"?>
+<settings>
+
+  <run_mode>fixed source</run_mode>
+  <batches>10</batches>
+  <inactive>0</inactive>
+  <particles>100000</particles>
+
+  <!-- Starting source -->
+  <source>
+    <library>build/libsource.so</library>
+  </source>
+
+</settings>

examples/xml/custom_source/source_ring.cpp

@@ -0,0 +1,26 @@
+#include <cmath> // for M_PI
+
+#include "openmc/random_lcg.h"
+#include "openmc/source.h"
+#include "openmc/particle.h"
+
+// you must have external C linkage here otherwise
+// dlopen will not find the file
+extern "C" openmc::Particle::Bank sample_source(uint64_t* seed)
+{
+  openmc::Particle::Bank particle;
+  // wgt
+  particle.particle = openmc::Particle::Type::neutron;
+  particle.wgt = 1.0;
+  // position
+  double angle = 2. * M_PI * openmc::prn(seed);
+  double radius = 3.0;
+  particle.r.x = radius * std::cos(angle);
+  particle.r.y = radius * std::sin(angle);
+  particle.r.z = 0.0;
+  // angle
+  particle.u = {1.0, 0.0, 0.0};
+  particle.E = 14.08e6;
+  particle.delayed_group = 0;
+  return particle;
+}

examples/xml/custom_source/tallies.xml

@@ -0,0 +1,17 @@
+<?xml version="1.0"?>
+<tallies>
+
+  <filter id="1" type="cell">
+    <bins>100</bins>
+  </filter>
+
+  <filter id="2" type="energy">
+    <bins>0 20.0e6</bins>
+  </filter>
+
+  <tally id="3">
+    <filters>1 2 </filters>
+    <scores>flux</scores>
+  </tally>
+
+</tallies>

include/openmc/settings.h

@@ -60,6 +60,7 @@ extern std::string path_input;            //!< directory where main .xml files r
 extern std::string path_output;           //!< directory where output files are written
 extern std::string path_particle_restart; //!< path to a particle restart file
 extern std::string path_source;
+extern std::string path_source_library;   //!< path to the source shared object
 extern std::string path_sourcepoint;      //!< path to a source file
 extern "C" std::string path_statepoint;   //!< path to a statepoint file
 

include/openmc/source.h

@@ -68,6 +68,9 @@ Particle::Bank sample_external_source(uint64_t* seed);
 //! Fill source bank at end of generation for fixed source simulations
 void fill_source_bank_fixedsource();
 
+//! Fill source bank at the end of a generation for dlopen based source simulation
+void fill_source_bank_custom_source();
+
 void free_memory_source();
 
 } // namespace openmc

openmc/source.py

@@ -8,20 +8,22 @@ from openmc.stats.multivariate import UnitSphere, Spatial
 import openmc.checkvalue as cv
 
 
-class Source(object):
+class Source:
     """Distribution of phase space coordinates for source sites.
 
     Parameters
     ----------
-    space : openmc.stats.Spatial, optional
+    space : openmc.stats.Spatial
         Spatial distribution of source sites
-    angle : openmc.stats.UnitSphere, optional
+    angle : openmc.stats.UnitSphere
         Angular distribution of source sites
-    energy : openmc.stats.Univariate, optional
+    energy : openmc.stats.Univariate
         Energy distribution of source sites
-    filename : str, optional
+    filename : str
         Source file from which sites should be sampled
-    strength : Real
+    library : str
+        Path to a custom source library
+    strength : float
         Strength of the source
     particle : {'neutron', 'photon'}
         Source particle type
@@ -36,7 +38,9 @@ class Source(object):
         Energy distribution of source sites
     file : str or None
         Source file from which sites should be sampled
-    strength : Real
+    library : str or None
+        Path to a custom source library
+    strength : float
         Strength of the source
     particle : {'neutron', 'photon'}
         Source particle type
@@ -44,11 +48,12 @@ class Source(object):
     """
 
     def __init__(self, space=None, angle=None, energy=None, filename=None,
-                 strength=1.0, particle='neutron'):
+                 library=None, strength=1.0, particle='neutron'):
         self._space = None
         self._angle = None
         self._energy = None
         self._file = None
+        self._library = None
 
         if space is not None:
             self.space = space
@@ -58,6 +63,8 @@ class Source(object):
             self.energy = energy
         if filename is not None:
             self.file = filename
+        if library is not None:
+            self.library = library
         self.strength = strength
         self.particle = particle
 
@@ -65,6 +72,10 @@ class Source(object):
     def file(self):
         return self._file
 
+    @property
+    def library(self):
+        return self._library
+
     @property
     def space(self):
         return self._space
@@ -90,6 +101,11 @@ class Source(object):
         cv.check_type('source file', filename, str)
         self._file = filename
 
+    @library.setter
+    def library(self, library_name):
+        cv.check_type('library', library_name, str)
+        self._library = library_name
+
     @space.setter
     def space(self, space):
         cv.check_type('spatial distribution', space, Spatial)
@@ -131,6 +147,8 @@ class Source(object):
             element.set("particle", self.particle)
         if self.file is not None:
             element.set("file", self.file)
+        if self.library is not None:
+            element.set("library", self.library)
         if self.space is not None:
             element.append(self.space.to_xml_element())
         if self.angle is not None:
@@ -168,6 +186,10 @@ class Source(object):
         if filename is not None:
             source.file = filename
 
+        library = get_text(elem, 'library')
+        if library is not None:
+            source.library = library
+
         space = elem.find('space')
         if space is not None:
             source.space = Spatial.from_xml_element(space)

src/settings.cpp

@@ -74,6 +74,7 @@ std::string path_input;
 std::string path_output;
 std::string path_particle_restart;
 std::string path_source;
+std::string path_source_library;
 std::string path_sourcepoint;
 std::string path_statepoint;
 

src/source.cpp

@@ -1,7 +1,15 @@
 #include "openmc/source.h"
 
+#if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
+#define HAS_DYNAMIC_LINKING
+#endif
+
 #include <algorithm> // for move
 
+#ifdef HAS_DYNAMIC_LINKING
+#include <dlfcn.h> // for dlopen, dlsym, dlclose, dlerror
+#endif
+
 #include <fmt/core.h>
 #include "xtensor/xadapt.hpp"
 
@@ -71,7 +79,12 @@ SourceDistribution::SourceDistribution(pugi::xml_node node)
       fatal_error(fmt::format("Source file '{}' does not exist.",
         settings::path_source));
     }
-
+  } else if (check_for_node(node, "library")) {
+    settings::path_source_library = get_node_value(node, "library", false, true);
+    if (!file_exists(settings::path_source_library)) {
+      fatal_error(fmt::format("Source library '{}' does not exist.",
+        settings::path_source_library));
+    }
   } else {
 
     // Spatial distribution for external source
@@ -237,7 +250,7 @@ void initialize_source()
 {
   write_message("Initializing source particles...", 5);
 
-  if (settings::path_source != "") {
+  if (!settings::path_source.empty()) {
     // Read the source from a binary file instead of sampling from some
     // assumed source distribution
 
@@ -261,6 +274,27 @@ void initialize_source()
 
     // Close file
     file_close(file_id);
+  } else if (!settings::path_source_library.empty()) {
+
+  #ifdef HAS_DYNAMIC_LINKING
+    // Get the source from a library object
+    write_message(fmt::format("Sampling from library source {}...",
+      settings::path_source), 6);
+
+    // Open the library
+    auto source_library = dlopen(settings::path_source_library.c_str(), RTLD_LAZY);
+    if (!source_library) {
+      fatal_error("Couldn't open source library " + settings::path_source_library);
+    }
+
+    // reset errors
+    dlerror();
+
+    fill_source_bank_custom_source();
+  #else
+    fatal_error("Custom source libraries have not yet been implemented for "
+      "non-POSIX systems");
+  #endif
 
   } else {
     // Generation source sites from specified distribution in user input
@@ -321,10 +355,50 @@ void free_memory_source()
   model::external_sources.clear();
 }
 
+// fill the source bank from the external source
+void fill_source_bank_custom_source()
+{
+#ifdef HAS_DYNAMIC_LINKING
+  // Open the library
+  auto source_library = dlopen(settings::path_source_library.c_str(), RTLD_LAZY);
+  if (!source_library) {
+    fatal_error("Couldn't open source library " + settings::path_source_library);
+  }
+
+  // reset errors
+  dlerror();
+
+  // get the function from the library
+  using sample_t = Particle::Bank (*)(uint64_t* seed);
+  auto sample_source = reinterpret_cast<sample_t>(dlsym(source_library, "sample_source"));
+
+  // check for any dlsym errors
+  auto dlsym_error = dlerror();
+  if (dlsym_error) {
+    dlclose(source_library);
+    fatal_error(fmt::format("Couldn't open the sample_source symbol: {}", dlsym_error));
+  }
+
+  // Generation source sites from specified distribution in the
+  // library source
+  for (int64_t i = 0; i < simulation::work_per_rank; ++i) {
+    // initialize random number seed
+    int64_t id = (simulation::total_gen + overall_generation()) *
+      settings::n_particles + simulation::work_index[mpi::rank] + i + 1;
+     uint64_t seed = init_seed(id, STREAM_SOURCE);
+
+    // sample external source distribution
+    simulation::source_bank[i] = sample_source(&seed);
+  }
+
+  // release the library
+  dlclose(source_library);
+#endif
+}
+
 void fill_source_bank_fixedsource()
 {
-  if (settings::path_source.empty()) {
-    #pragma omp parallel for
+  if (settings::path_source.empty() && settings::path_source_library.empty()) {
     for (int64_t i = 0; i < simulation::work_per_rank; ++i) {
       // initialize random number seed
       int64_t id = (simulation::total_gen + overall_generation()) *
@@ -334,6 +408,8 @@ void fill_source_bank_fixedsource()
       // sample external source distribution
       simulation::source_bank[i] = sample_external_source(&seed);
     }
+  } else if (settings::path_source.empty() && !settings::path_source_library.empty()) {
+     fill_source_bank_custom_source();
   }
 }
 

tests/regression_tests/source_dlopen/inputs_true.dat

@@ -0,0 +1,23 @@
+<?xml version='1.0' encoding='utf-8'?>
+<geometry>
+  <cell id="1" material="1" region="-1" universe="1" />
+  <surface boundary="vacuum" coeffs="0.0 0.0 0.0 100" id="1" type="sphere" />
+</geometry>
+<?xml version='1.0' encoding='utf-8'?>
+<materials>
+  <material id="1" name="natural_lead">
+    <density units="g/cm3" value="11.34" />
+    <nuclide ao="0.014" name="Pb204" />
+    <nuclide ao="0.241" name="Pb206" />
+    <nuclide ao="0.221" name="Pb207" />
+    <nuclide ao="0.524" name="Pb208" />
+  </material>
+</materials>
+<?xml version='1.0' encoding='utf-8'?>
+<settings>
+  <run_mode>fixed source</run_mode>
+  <particles>1000</particles>
+  <batches>10</batches>
+  <inactive>0</inactive>
+  <source library="build/libsource.so" strength="1.0" />
+</settings>

tests/regression_tests/source_dlopen/source_sampling.cpp

@@ -0,0 +1,23 @@
+#include <iostream>
+#include "openmc/random_lcg.h"
+#include "openmc/source.h"
+#include "openmc/particle.h"
+
+// you must have external C linkage here otherwise 
+// dlopen will not find the file
+extern "C" openmc::Particle::Bank sample_source(uint64_t *seed) {
+    openmc::Particle::Bank particle;
+    // wgt
+    particle.particle = openmc::Particle::Type::neutron;
+    particle.wgt = 1.0;
+    // position
+    
+    particle.r.x = 0.;
+    particle.r.y = 0.;
+    particle.r.z = 0.;
+    // angle
+    particle.u = {1.0, 0.0, 0.0};
+    particle.E = 14.08e6;
+    particle.delayed_group = 0;
+    return particle;    
+}

tests/regression_tests/source_dlopen/test.py

@@ -0,0 +1,73 @@
+from pathlib import Path
+import os
+import shutil
+import subprocess
+import textwrap
+
+import openmc
+import pytest
+
+from tests.testing_harness import PyAPITestHarness
+
+
+@pytest.fixture
+def compile_source(request):
+    """Compile the external source"""
+
+    # Get build directory and write CMakeLists.txt file
+    openmc_dir = Path(str(request.config.rootdir)) / 'build'
+    with open('CMakeLists.txt', 'w') as f:
+        f.write(textwrap.dedent("""
+            cmake_minimum_required(VERSION 3.3 FATAL_ERROR)
+            project(openmc_sources CXX)
+            add_library(source SHARED source_sampling.cpp)
+            find_package(OpenMC REQUIRED HINTS {})
+            target_link_libraries(source OpenMC::libopenmc)
+            """.format(openmc_dir)))
+
+    # Create temporary build directory and change to there
+    local_builddir = Path('build')
+    local_builddir.mkdir(exist_ok=True)
+    os.chdir(str(local_builddir))
+
+    # Run cmake/make to build the shared libary
+    subprocess.run(['cmake', os.path.pardir], check=True)
+    subprocess.run(['make'], check=True)
+    os.chdir(os.path.pardir)
+
+    yield
+
+    # Remove local build directory when test is complete
+    shutil.rmtree('build')
+
+
+@pytest.fixture
+def model():
+    model = openmc.model.Model()
+    natural_lead = openmc.Material(name="natural_lead")
+    natural_lead.add_element('Pb', 1.0)
+    natural_lead.set_density('g/cm3', 11.34)
+    model.materials.append(natural_lead)
+
+    # geometry
+    surface_sph1 = openmc.Sphere(r=100, boundary_type='vacuum')
+    cell_1 = openmc.Cell(fill=natural_lead, region=-surface_sph1)
+    model.geometry = openmc.Geometry([cell_1])
+
+    # settings
+    model.settings.batches = 10
+    model.settings.inactive = 0
+    model.settings.particles = 1000
+    model.settings.run_mode = 'fixed source'
+
+    # custom source from shared library
+    source = openmc.Source()
+    source.library = 'build/libsource.so'
+    model.settings.source = source
+
+    return model
+
+
+def test_dlopen_source(compile_source, model):
+    harness = PyAPITestHarness('statepoint.10.h5', model)
+    harness.main()

tests/unit_tests/test_source.py

@@ -34,3 +34,11 @@ def test_source_file():
     elem = src.to_xml_element()
     assert 'strength' in elem.attrib
     assert 'file' in elem.attrib
+
+def test_source_dlopen():
+    library = './libsource.so'
+    src = openmc.Source(library=library)
+    assert src.library == library
+
+    elem = src.to_xml_element()
+    assert 'library' in elem.attrib