Repository: openmc-dev/openmc

Issue:
optional writing of source for fixed_source simulations
A recent PR #1518 implemented a new bank model which is a nice improvement for fixed source simulations. Thanks @alex-lyons 

The new bank model no longer exports source information like particle x,y,z, E and u,v,w to the statepoint file. This is great as we often run simulation with fixed_source mode with many particles and the automatic writing of this data to the statepoint file was generating large files and taking a long time.

However it would be nice to be able to optinally add source information in some cases so we can do things like plot the source locations, distributions etc

So I have been looking at the documentation the and it appears we have a `sourcepoint` setting called `write` in` openmc.settings ` that should control if we write the source to the h5 file or not. 

https://docs.openmc.org/en/latest/pythonapi/generated/openmc.Settings.html

Then there is an additional check in the simulations.cpp that prohibit writing of source points for fixed source.

```
if (contains(settings::sourcepoint_batch, simulation::current_batch)
        && settings::source_write && !settings::source_separate) {
      bool b = (settings::run_mode == RunMode::EIGENVALUE);
```

https://github.com/openmc-dev/openmc/blob/571ed3b6ab449e555fc9c8452106425d26b19bd3/src/simulation.cpp#L367

Perhaps we should just use the `openmc.settings sourcepoint write` to determine if source points should be writing and not involve the simulation type.

https://github.com/Shimwell/openmc/blob/8b8392f94326a965de47d3ba6fb010a9bbb8f922/src/simulation.cpp#L366-L367

Unfortunately this would not be enough to write the source points to the statepoint file and some additional code is probably needed in the bank model but I get a bit lost and was wondering if anyone has any suggestions or perhaps this is not an issue for anyone else? 


Discussion:
Just in case anyone else is looking for a solution to this there are a few work arounds, here is an example of getting data points by pretending to perform an eiganvalue simulation

```


import xml.etree.ElementTree as ET

import h5py
import openmc
import plotly.graph_objects as go

sett = openmc.Settings()
sett.run_mode = "eigenvalue" # this will fail but it will write the inital_source.h5 file first
sett.particles = 10
sett.batches = 3
sett.inactive = 0
sett.write_initial_source = True

# Create a DT point source
source = openmc.Source()
source.space = openmc.stats.Point((0, 0, 0))
source.angle = openmc.stats.Isotropic()
source.energy = openmc.stats.Muir(kt=10e3) # might want to change the 10KeV plasma temperature
sett.source = source

sett.export_to_xml()

# this just adds write_initial_source == True to the settings.xml
tree = ET.parse("settings.xml")
root = tree.getroot()
elem = ET.SubElement(root, "write_initial_source")
elem.text = "true"
tree.write("settings.xml")

# no materials needed
openmc.Materials([]).export_to_xml()

# simple geometry
outer_surface = openmc.Sphere(r=100, boundary_type='vacuum')
cell = openmc.Cell(region=-outer_surface)
universe = openmc.Universe(cells=[cell])
geom = openmc.Geometry(universe)
geom.export_to_xml()

# this will crash hence the try except loop, but it writes the inital source.h5
try:
    openmc.run()
except:
    pass

# Read the intial_source.h5 and pull the coordinates from the dataset
filename = "initial_source.h5"
x = []
y = []
z = []
dir_x = []
dir_y = []
dir_z = []
e = []
with h5py.File(filename, "r") as f:
    dataset = f["source_bank"]
    for data in dataset:
        x += [data[0][0]]
        y += [data[0][1]]
        z += [data[0][2]]
        e += [data[2]]
        dir_x += [data[1][0]]
        dir_y += [data[1][1]]
        dir_z += [data[1][2]]

print('neutron x birth positions', x)
print('neutron y birth positions', y)
print('neutron z birth positions', z)
print('neutron dir_x birth positions', dir_x)
print('neutron dir_y birth positions', dir_y)
print('neutron dir_z birth positions', dir_z)
print('neutron birth energy', e)

```

Just a thought -- one of the simplest approaches here might be to expose the `sample_external_source` C++ function to Python via `openmc.lib`. This would give a power user the ability to sample the source as they wish completely separate from a simulation.
I am keen understand how to contribute to the openmc.lib so I had a little go and am wondering if I am on the right track.

Would core.py need these extra lines
https://github.com/shimwell/openmc-1/blob/f7b968663c098848e5f5da0f25ac50b61f3597d6/openmc/lib/core.py#L93-L95

Also would core.py need a function link this
https://github.com/shimwell/openmc-1/blob/f7b968663c098848e5f5da0f25ac50b61f3597d6/openmc/lib/core.py#L102-L111


We'll need these changes, but there's a bit more to it. We'll need to expose the function through the CAPI in `capi.h` and add a wrapper function `openmc_sample_external_source` that matches the C-style pattern we have in place for errorcodes and such.
Is [this line](https://github.com/shimwell/openmc-1/blob/254884576ea65705138d7d6742e53ed719ab2908/include/openmc/capi.h#L124) all that is needed to expose the function through the CAPI in ```capi.h```

Resolving patch:
diff --git a/openmc/lib/core.py b/openmc/lib/core.py
--- a/openmc/lib/core.py
+++ b/openmc/lib/core.py
@@ -1,8 +1,10 @@
 from contextlib import contextmanager
 from ctypes import (c_bool, c_int, c_int32, c_int64, c_double, c_char_p,
-                    c_char, POINTER, Structure, c_void_p, create_string_buffer)
+                    c_char, POINTER, Structure, c_void_p, create_string_buffer,
+                    c_uint64, c_size_t)
 import sys
 import os
+from random import getrandbits
 
 import numpy as np
 from numpy.ctypeslib import as_array
@@ -10,6 +12,7 @@
 from . import _dll
 from .error import _error_handler
 import openmc.lib
+import openmc
 
 
 class _SourceSite(Structure):
@@ -95,6 +98,9 @@ class _SourceSite(Structure):
                                             POINTER(c_double)]
 _dll.openmc_global_bounding_box.restype = c_int
 _dll.openmc_global_bounding_box.errcheck = _error_handler
+_dll.openmc_sample_external_source.argtypes = [c_size_t, POINTER(c_uint64), POINTER(_SourceSite)]
+_dll.openmc_sample_external_source.restype = c_int
+_dll.openmc_sample_external_source.errcheck = _error_handler
 
 
 def global_bounding_box():
@@ -415,6 +421,45 @@ def run(output=True):
         _dll.openmc_run()
 
 
+def sample_external_source(n_samples=1, prn_seed=None):
+    """Sample external source
+
+    .. versionadded:: 0.13.1
+
+    Parameters
+    ----------
+    n_samples : int
+        Number of samples
+    prn_seed : int
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
+
+    Returns
+    -------
+    list of openmc.SourceParticle
+        List of samples source particles
+
+    """
+    if n_samples <= 0:
+        raise ValueError("Number of samples must be positive")
+    if prn_seed is None:
+        prn_seed = getrandbits(63)
+
+    # Call into C API to sample source
+    sites_array = (_SourceSite * n_samples)()
+    _dll.openmc_sample_external_source(c_size_t(n_samples), c_uint64(prn_seed), sites_array)
+
+    # Convert to list of SourceParticle and return
+    return [
+        openmc.SourceParticle(
+            r=site.r, u=site.u, E=site.E, time=site.time, wgt=site.wgt,
+            delayed_group=site.delayed_group, surf_id=site.surf_id,
+            particle=openmc.ParticleType(site.particle)
+        )
+        for site in sites_array
+    ]
+
+
 def simulation_init():
     """Initialize simulation"""
     _dll.openmc_simulation_init()
diff --git a/openmc/lib/math.py b/openmc/lib/math.py
--- a/openmc/lib/math.py
+++ b/openmc/lib/math.py
@@ -1,12 +1,11 @@
 from ctypes import c_int, c_double, POINTER, c_uint64
+from random import getrandbits
 
 import numpy as np
 from numpy.ctypeslib import ndpointer
 
 from . import _dll
 
-from random import getrandbits
-
 
 _dll.t_percentile.restype = c_double
 _dll.t_percentile.argtypes = [c_double, c_int]
@@ -200,7 +199,8 @@ def rotate_angle(uvw0, mu, phi, prn_seed=None):
     phi : float
         Azimuthal angle; if None, one will be sampled uniformly
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -232,7 +232,8 @@ def maxwell_spectrum(T, prn_seed=None):
     T : float
         Spectrum parameter
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -240,10 +241,10 @@ def maxwell_spectrum(T, prn_seed=None):
         Sampled outgoing energy
 
     """
-    
+
     if prn_seed is None:
         prn_seed = getrandbits(63)
-	
+
     return _dll.maxwell_spectrum(T, c_uint64(prn_seed))
 
 
@@ -257,7 +258,8 @@ def watt_spectrum(a, b, prn_seed=None):
     b : float
         Spectrum parameter b
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -265,7 +267,7 @@ def watt_spectrum(a, b, prn_seed=None):
         Sampled outgoing energy
 
     """
-    
+
     if prn_seed is None:
         prn_seed = getrandbits(63)
 
@@ -282,7 +284,8 @@ def normal_variate(mean_value, std_dev, prn_seed=None):
     std_dev : float
         Standard deviation of the normal distribution
     prn_seed : int
-        PRNG seed; if None, one will be generated randomly
+        Pseudorandom number generator (PRNG) seed; if None, one will be
+        generated randomly.
 
     Returns
     -------
@@ -290,7 +293,7 @@ def normal_variate(mean_value, std_dev, prn_seed=None):
         Sampled outgoing normally distributed value
 
     """
-    
+
     if prn_seed is None:
         prn_seed = getrandbits(63)
 


Test patch:
diff --git a/tests/unit_tests/test_lib.py b/tests/unit_tests/test_lib.py
--- a/tests/unit_tests/test_lib.py
+++ b/tests/unit_tests/test_lib.py
@@ -807,3 +807,46 @@ def test_cell_rotation(pincell_model_w_univ, mpi_intracomm):
     cell.rotation = (180., 0., 0.)
     assert cell.rotation == pytest.approx([180., 0., 0.])
     openmc.lib.finalize()
+
+
+def test_sample_external_source(run_in_tmpdir, mpi_intracomm):
+    # Define a simple model and export
+    mat = openmc.Material()
+    mat.add_nuclide('U235', 1.0e-2)
+    sph = openmc.Sphere(r=100.0, boundary_type='vacuum')
+    cell = openmc.Cell(fill=mat, region=-sph)
+    model = openmc.Model()
+    model.geometry = openmc.Geometry([cell])
+    model.settings.source = openmc.Source(
+        space=openmc.stats.Box([-5., -5., -5.], [5., 5., 5.]),
+        angle=openmc.stats.Monodirectional((0., 0., 1.)),
+        energy=openmc.stats.Discrete([1.0e5], [1.0])
+    )
+    model.settings.particles = 1000
+    model.settings.batches = 10
+    model.export_to_xml()
+
+    # Sample some particles and make sure they match specified source
+    openmc.lib.init()
+    particles = openmc.lib.sample_external_source(10, prn_seed=3)
+    assert len(particles) == 10
+    for p in particles:
+        assert -5. < p.r[0] < 5.
+        assert -5. < p.r[1] < 5.
+        assert -5. < p.r[2] < 5.
+        assert p.u[0] == 0.0
+        assert p.u[1] == 0.0
+        assert p.u[2] == 1.0
+        assert p.E == 1.0e5
+
+    # Using the same seed should produce the same particles
+    other_particles = openmc.lib.sample_external_source(10, prn_seed=3)
+    assert len(other_particles) == 10
+    for p1, p2 in zip(particles, other_particles):
+        assert p1.r == p2.r
+        assert p1.u == p2.u
+        assert p1.E == p2.E
+        assert p1.time == p2.time
+        assert p1.wgt == p2.wgt
+
+    openmc.lib.finalize()