Support 2D projections

parametric_plasma_source/enums.hpp

@@ -0,0 +1,10 @@
+namespace plasma_source {
+
+    // Defines a selection of valid basis values to sample from.
+    enum valid_basis {
+        XYZ, // 3-D basis
+        RY, // 2-D basis in radial plane along y-axis
+        RZ, // 2-D basis in radial plane along z-axis
+    };
+
+}

parametric_plasma_source/source_sampling.cpp

@@ -3,6 +3,7 @@
 #include "openmc/source.h"
 #include "openmc/particle.h"
 #include "plasma_source.hpp"
+#include "enums.hpp"
 
 
 // Spherical tokamak SOURCE 
@@ -24,6 +25,7 @@ const double shafranov_shift = 0.0; //metres
 const std::string name = "parametric_plasma_source";
 const int number_of_bins  = 100;
 const int plasma_type = 1; // 1 is default; //0 = L mode anything else H/A mode
+const plasma_source::valid_basis basis = plasma_source::valid_basis::XYZ; // XYZ for 3D, RY or RZ for 2D
 
 
 
@@ -70,9 +72,28 @@ extern "C" openmc::Particle::Bank sample_source(uint64_t* seed) {
     source.SampleSource(randoms,particle.r.x,particle.r.y,particle.r.z,
                         u,v,w,E); 
 
+    // Convert m to cm
     particle.r.x *= 100.;
     particle.r.y *= 100.;
-    particle.r.z *= 100.;    
+    particle.r.z *= 100.;
+
+    if(basis == plasma_source::valid_basis::XYZ) {
+        // Use values as-is
+    }
+    else if(basis == plasma_source::valid_basis::RY) {
+        particle.r.x = std::sqrt(std::pow(particle.r.x, 2) + std::pow(particle.r.y, 2));
+        particle.r.y = particle.r.z;
+        particle.r.z = 0.;
+    }
+    else if(basis == plasma_source::valid_basis::RZ) {
+        particle.r.x = std::sqrt(std::pow(particle.r.x, 2) + std::pow(particle.r.y, 2));
+        particle.r.y = 0.;
+        particle.r.z = particle.r.z;
+    }
+    else {
+        throw std::runtime_error("Parametric plasma source: incorrect basis provided, "
+                                 "please use XYZ, RY, or RZ.");
+    }
 
     // particle.E = 14.08e6;
     particle.E = E*1e6; // convert from MeV -> eV