simpler switch statements

examples/point_source_example.py

@@ -11,7 +11,7 @@
 geometry = openmc.Geometry([cell])
 
 # define the source
-my_source = FusionPointSource(coordinate=(0, 0, 0), temperature=20000.0, fuel="DT")
+my_source = FusionPointSource(coordinate=(0, 0, 0), temperature=20000.0, fuel={"D":0.5, "T":0.5})
 
 # Tell OpenMC we're going to use our custom source
 settings = openmc.Settings()

examples/ring_source_example.py

@@ -16,6 +16,7 @@
     radius=700,
     angles=(0.0, 2 * math.pi),  # 360deg source
     temperature=20000.0,
+    fuel={"D":0.5,"T":0.5}
 )
 
 settings = openmc.Settings()

examples/tokamak_source_example.py

@@ -29,6 +29,7 @@
     shafranov_factor=0.44789,
     triangularity=0.270,
     ion_temperature_beta=6,
+    fuel={"D":0.5,"T":0.5}
 )
 
 # Tell OpenMC we're going to use our custom source

src/openmc_plasma_source/__init__.py

@@ -14,3 +14,4 @@
 from .tokamak_source import TokamakSource
 from .ring_source import FusionRingSource
 from .point_source import FusionPointSource
+from .fuel_types import get_neutron_energy_distribution

src/openmc_plasma_source/fuel_types.py

@@ -5,7 +5,7 @@
 
 def get_neutron_energy_distribution(
     ion_temperature: float,
-    fuel: dict = {"D": 0.5, "T": 0.5},
+    fuel: dict,
 ) -> openmc.stats.Discrete:
     """Finds the energy distribution
 
@@ -33,7 +33,7 @@ def get_neutron_energy_distribution(
     if sum_fuel_isotopes < 0.0:
         raise ValueError(f"isotope must sum to be above 0. Not {sum_fuel_isotopes}")
 
-    for k, v in fuel.dict:
+    for k, v in fuel.items():
         if k not in ["D", "T"]:
             raise ValueError(
                 f'Fuel dictionary keys must be either "D" or "T" not "{k}".'
@@ -43,54 +43,49 @@ def get_neutron_energy_distribution(
         if v > 1:
             raise ValueError(f'Fuel dictionary values must be below 1 not "{k}".')
 
-    # Set atomic fraction of D and T in scattering medium and source
-    if "D" in fuel.keys():
-        nst.frac_D_default = fuel["D"]
+    if ['D'] == sorted(set(fuel.keys())):
         max_energy_mev = 5
-    else:
-        nst.frac_D_default = 0
-
-    if "T" in fuel.keys():
-        nst.frac_T_default = fuel["T"]
+    elif ['T'] == sorted(set(fuel.keys())):
         max_energy_mev = 12
-    else:
-        nst.frac_T_default = 0
-
-    if "T" in fuel.keys() and "D" in fuel.keys():
+    elif ['D', 'T'] == sorted(set(fuel.keys())):
         max_energy_mev = 20
 
+    print(max_energy_mev,'MeV')
+
     # 1.0 neutron yield, all reactions scaled by this value
-    num_of_vals = 500
+    num_of_vals = 50
     # single grid for DT, DD and TT grid
     E_pspec = np.linspace(0, max_energy_mev, num_of_vals)  # accepts MeV units
 
     dNdE_DT_DD_TT = np.zeros(num_of_vals)
-    if "D" in fuel.keys() and "T" in fuel.keys():
+    if ['D', 'T'] == sorted(set(fuel.keys())):
         DTmean, DTvar = nst.DTprimspecmoments(ion_temperature)
         DDmean, DDvar = nst.DDprimspecmoments(ion_temperature)
 
         Y_DT = 1.0
-        Y_DD = nst.yield_from_dt_yield_ratio("dd", Y_DT, ion_temperature)
-        Y_TT = nst.yield_from_dt_yield_ratio("tt", Y_DT, ion_temperature)
+        Y_DD = nst.yield_from_dt_yield_ratio("dd", Y_DT, ion_temperature, fuel["D"], fuel["T"])
+        Y_TT = nst.yield_from_dt_yield_ratio("tt", Y_DT, ion_temperature, fuel["D"], fuel["T"])
 
         dNdE_DT = Y_DT * nst.Qb(E_pspec, DTmean, DTvar)  # Brysk shape i.e. Gaussian
         dNdE_DD = Y_DD * nst.Qb(E_pspec, DDmean, DDvar)  # Brysk shape i.e. Gaussian
         dNdE_TT = Y_TT * nst.dNdE_TT(E_pspec, ion_temperature)
         dNdE_DT_DD_TT = dNdE_DT + dNdE_DD + dNdE_TT
-
-    if "D" in fuel.keys() and "T" not in fuel.keys():
+        print('dt')
+        return openmc.stats.Discrete(E_pspec * 1e6, dNdE_DT_DD_TT)
+    elif ['D'] == sorted(set(fuel.keys())):
         DTmean, DTvar = nst.DTprimspecmoments(ion_temperature)
         DDmean, DDvar = nst.DDprimspecmoments(ion_temperature)
 
+        print('d')
         Y_DD = 1.0
 
         dNdE_DD = Y_DD * nst.Qb(E_pspec, DDmean, DDvar)  # Brysk shape i.e. Gaussian
-        dNdE_DT_DD_TT = dNdE_DD
+        return openmc.stats.Discrete(E_pspec * 1e6, dNdE_DD)
 
-    if "D" not in fuel.keys() and "T" in fuel.keys():
+    elif ['T'] == sorted(set(fuel.keys())):
         Y_TT = 1.0
 
+        print('t')
         dNdE_TT = Y_TT * nst.dNdE_TT(E_pspec, ion_temperature)
-        dNdE_DT_DD_TT = dNdE_TT
+        return openmc.stats.Discrete(E_pspec * 1e6, dNdE_TT)
 
-    return openmc.stats.Discrete(E_pspec * 1e6, dNdE_DT_DD_TT)

src/openmc_plasma_source/point_source.py

@@ -26,8 +26,7 @@ def __init__(
         # Set local attributes
         self.coordinate = coordinate
         self.temperature = temperature
-        self.fuel_type = fuel
-        self.fuel = fuel_types[self.fuel_type]
+        self.fuel = fuel
 
         # Call init for openmc.Source
         super().__init__()
@@ -64,14 +63,3 @@ def temperature(self, value):
             self._temperature = value
         else:
             raise ValueError("Temperature must be strictly positive float.")
-
-    @property
-    def fuel_type(self):
-        return self._fuel_type
-
-    @fuel_type.setter
-    def fuel_type(self, value):
-        if value in fuel_types:
-            self._fuel_type = value
-        else:
-            raise KeyError("Invalid fuel type")