From b1267f05448d80bce3069779f03de68749c44297 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Fri, 15 Jul 2022 13:23:07 -0500
Subject: [PATCH] Allow openmc_sample_external_source to produce multiple
 samples

---
 docs/source/pythonapi/capi.rst |  1 +
 include/openmc/capi.h          |  1 +
 openmc/lib/core.py             | 33 ++++++++++++++++++++-------------
 src/source.cpp                 | 13 +++++++++++--
 4 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/docs/source/pythonapi/capi.rst b/docs/source/pythonapi/capi.rst
index 03be25abf..882e3c71b 100644
--- a/docs/source/pythonapi/capi.rst
+++ b/docs/source/pythonapi/capi.rst
@@ -29,6 +29,7 @@ Functions
    reset
    run
    run_in_memory
+   sample_external_source
    simulation_init
    simulation_finalize
    source_bank
diff --git a/include/openmc/capi.h b/include/openmc/capi.h
index 0929a11f7..69d3ff1f6 100644
--- a/include/openmc/capi.h
+++ b/include/openmc/capi.h
@@ -121,6 +121,7 @@ int openmc_regular_mesh_set_params(int32_t index, int n, const double* ll,
 int openmc_reset();
 int openmc_reset_timers();
 int openmc_run();
+int openmc_sample_external_source(size_t n, uint64_t* seed, void* sites);
 void openmc_set_seed(int64_t new_seed);
 int openmc_set_n_batches(
   int32_t n_batches, bool set_max_batches, bool add_statepoint_batch);
diff --git a/openmc/lib/core.py b/openmc/lib/core.py
index 3f4a301cd..3a437a295 100644
--- a/openmc/lib/core.py
+++ b/openmc/lib/core.py
@@ -1,7 +1,7 @@
 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_uint64)
+                    c_uint64, c_size_t)
 import sys
 import os
 from random import getrandbits
@@ -98,7 +98,7 @@ _dll.openmc_global_bounding_box.argtypes = [POINTER(c_double),
                                             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 = [POINTER(c_uint64), POINTER(_SourceSite)]
+_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
 
@@ -421,35 +421,42 @@ def run(output=True):
         _dll.openmc_run()
 
 
-def sample_external_source(prn_seed=None):
+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
         PRNG seed; if None, one will be generated randomly
 
     Returns
     -------
-    openmc.SourceParticle
-        Sampled source particle
+    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
-    site = _SourceSite()
-    _dll.openmc_sample_external_source(c_uint64(prn_seed), site)
+    sites_array = (_SourceSite * n_samples)()
+    _dll.openmc_sample_external_source(c_size_t(n_samples), c_uint64(prn_seed), sites_array)
 
-    # Convert to 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)
-    )
+    # 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():
diff --git a/src/source.cpp b/src/source.cpp
index 8b634ec45..3716720df 100644
--- a/src/source.cpp
+++ b/src/source.cpp
@@ -403,9 +403,18 @@ void free_memory_source()
 // C API
 //==============================================================================
 
-extern "C" int openmc_sample_external_source(uint64_t* seed, SourceSite* site)
+extern "C" int openmc_sample_external_source(
+  size_t n, uint64_t* seed, void* sites)
 {
-  *site = sample_external_source(seed);
+  if (!sites || !seed) {
+    set_errmsg("Received null pointer.");
+    return OPENMC_E_INVALID_ARGUMENT;
+  }
+
+  auto sites_array = static_cast<SourceSite*>(sites);
+  for (size_t i = 0; i < n; ++i) {
+    sites_array[i] = sample_external_source(seed);
+  }
   return 0;
 }