Merge pull request #33 from fusion-energy/develop

 pep8 automatic formatting and relaxes matplotlib version requirement

.github/workflows/black.yml

@@ -0,0 +1,32 @@
+name: black
+
+on:
+  push:
+    paths:
+      - '**.py'
+
+defaults:
+  run:
+    shell: bash
+
+jobs:
+  black:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v2
+      with:
+        ref: ${{ github.head_ref }}
+    - name: Setup Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.x
+    - name: Install black
+      run: |
+        python -m pip install --upgrade pip
+        pip install black
+    - name: Run black
+      run: |
+        black .
+    - uses: stefanzweifel/git-auto-commit-action@v4
+      with:
+        commit_message: "[skip ci] Apply formatting changes"

examples/point_source_example.py

@@ -1,12 +1,7 @@
-
 import openmc
 from openmc_plasma_source import FusionPointSource
 
-my_source = FusionPointSource(
-        coordinate = (0,0,0),
-        temperature = 20000.,
-        fuel='DT'
-)
+my_source = FusionPointSource(coordinate=(0, 0, 0), temperature=20000.0, fuel="DT")
 
 # Create a single material
 iron = openmc.Material()

examples/ring_source_example.py

@@ -1,11 +1,10 @@
-
 import openmc
 from openmc_plasma_source import FusionRingSource
 
 my_source = FusionRingSource(
-    radius = 700,
-    angles =(0., 6.28318530718),  # 360deg source
-    temperature = 20000.,
+    radius=700,
+    angles=(0.0, 6.28318530718),  # 360deg source
+    temperature=20000.0,
 )
 
 # Create a single material

examples/tokamak_source_example.py

@@ -19,9 +19,8 @@
     mode="H",
     shafranov_factor=0.44789,
     triangularity=0.270,
-    ion_temperature_beta=6
-    )
-
+    ion_temperature_beta=6,
+)
 
 
 # Create a single material

openmc_plasma_source/plotting/__init__.py

@@ -1 +1 @@
-from .plot_tokamak_source import plot_tokamak_source_3D, scatter_tokamak_source
+from .plot_tokamak_source import plot_tokamak_source_3D, scatter_tokamak_source

openmc_plasma_source/plotting/plot_tokamak_source.py

@@ -33,7 +33,9 @@ def scatter_tokamak_source(source, quantity=None, **kwargs):
     return plt.scatter(source.RZ[0], source.RZ[1], c=colours, **kwargs)
 
 
-def plot_tokamak_source_3D(source, quantity=None, angles=[0, 1/2*np.pi], colorbar="viridis", **kwargs):
+def plot_tokamak_source_3D(
+    source, quantity=None, angles=[0, 1 / 2 * np.pi], colorbar="viridis", **kwargs
+):
     """Creates a 3D plot of the tokamak source.
     See https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.plot.html#matplotlib.pyplot.plot
     for more arguments.
@@ -68,7 +70,7 @@ def plot_tokamak_source_3D(source, quantity=None, angles=[0, 1/2*np.pi], colorba
     theta = np.linspace(*angles, 100)
     for i in range(source.sample_size):
         if values is not None:
-            colour = colorbar(values[i]/max(values))
+            colour = colorbar(values[i] / max(values))
         else:
             colour = None
         x = source.RZ[0][i] * np.sin(theta)

openmc_plasma_source/point_source.py

@@ -1,4 +1,3 @@
-
 from typing import Tuple
 
 import openmc
@@ -10,24 +9,27 @@ class FusionPointSource(openmc.Source):
     after initialization if required. Default isotropic point source at the
     origin with a Muir energy distribution.
     """
+
     def __init__(
         self,
         coordinate: Tuple[float, float, float] = (0, 0, 0),
-        temperature: float = 20000.,
-        fuel: str = 'DT'
+        temperature: float = 20000.0,
+        fuel: str = "DT",
     ):
 
         super().__init__()
 
         # performed after the super init as these are Source attributes
         self.space = openmc.stats.Point(coordinate)
         self.angle = openmc.stats.Isotropic()
-        if fuel == 'DT':
-            mean_energy = 14080000.  # mean energy in eV
+        if fuel == "DT":
+            mean_energy = 14080000.0  # mean energy in eV
             mass_of_reactants = 5  # mass of the reactants (D + T) AMU
-        elif fuel == 'DD':
-            mean_energy = 2450000.  # mean energy in eV
+        elif fuel == "DD":
+            mean_energy = 2450000.0  # mean energy in eV
             mass_of_reactants = 4  # mass of the reactants (D + D) AMU
         else:
             raise ValueError(f'fuel must be either "DT" or "DD", not {fuel}')
-        self.energy = openmc.stats.Muir(e0=mean_energy , m_rat=mass_of_reactants , kt=temperature)
+        self.energy = openmc.stats.Muir(
+            e0=mean_energy, m_rat=mass_of_reactants, kt=temperature
+        )

openmc_plasma_source/ring_source.py

@@ -1,4 +1,3 @@
-
 import openmc
 import numpy as np
 
@@ -15,13 +14,14 @@ class FusionRingSource(openmc.Source):
         z_placement: Location of the ring source (m). Defaults to 0.
         temperature: the temperature to use in the Muir distribution in eV,
     """
+
     def __init__(
         self,
         radius,
-        angles=(0, 2*np.pi),
+        angles=(0, 2 * np.pi),
         z_placement=0,
-        temperature: float = 20000.,
-        fuel='DT'
+        temperature: float = 20000.0,
+        fuel="DT",
     ):
 
         super().__init__()
@@ -31,18 +31,17 @@ def __init__(
         z_values = openmc.stats.Discrete([z_placement], [1])
         angle = openmc.stats.Uniform(a=angles[0], b=angles[1])
         self.space = openmc.stats.CylindricalIndependent(
-            r=radius,
-            phi=angle,
-            z=z_values,
-            origin=(0.0, 0.0, 0.0)
+            r=radius, phi=angle, z=z_values, origin=(0.0, 0.0, 0.0)
         )
         self.angle = openmc.stats.Isotropic()
-        if fuel == 'DT':
-            mean_energy = 14080000.  # mean energy in eV
+        if fuel == "DT":
+            mean_energy = 14080000.0  # mean energy in eV
             mass_of_reactants = 5  # mass of the reactants (D + T) AMU
-        elif fuel == 'DD':
-            mean_energy = 2450000.  # mean energy in eV
+        elif fuel == "DD":
+            mean_energy = 2450000.0  # mean energy in eV
             mass_of_reactants = 4  # mass of the reactants (D + D) AMU
         else:
             raise ValueError(f'fuel must be either "DT" or "DD", not {fuel}')
-        self.energy = openmc.stats.Muir(e0=mean_energy , m_rat=mass_of_reactants , kt=temperature)
+        self.energy = openmc.stats.Muir(
+            e0=mean_energy, m_rat=mass_of_reactants, kt=temperature
+        )

openmc_plasma_source/tokamak_source.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 
-class TokamakSource():
+class TokamakSource:
     """Plasma neutron source sampling.
     This class greatly relies on models described in [1]
 
@@ -45,6 +45,7 @@ class TokamakSource():
         sample_size (int, optional): number of neutron sources. Defaults
             to 1000.
     """
+
     def __init__(
         self,
         major_radius,
@@ -63,8 +64,8 @@ def __init__(
         ion_temperature_separatrix,
         pedestal_radius,
         shafranov_factor,
-        angles=(0, 2*np.pi),
-        sample_size=1000
+        angles=(0, 2 * np.pi),
+        sample_size=1000,
     ) -> None:
         self.major_radius = major_radius
         self.minor_radius = minor_radius
@@ -103,21 +104,27 @@ def ion_density(self, r):
             float, ndarray: ion density in m-3
         """
         if self.mode == "L":
-            density = self.ion_density_centre * \
-                (1 - (r/self.major_radius)**2)**self.ion_density_peaking_factor
+            density = (
+                self.ion_density_centre
+                * (1 - (r / self.major_radius) ** 2) ** self.ion_density_peaking_factor
+            )
         elif self.mode in ["H", "A"]:
             # TODO: find an alternative to iterating through the array
             if isinstance(r, np.ndarray):
                 density = []
                 for radius in r:
                     if 0 < radius < self.pedestal_radius:
                         prod = self.ion_density_centre - self.ion_density_pedestal
-                        prod *= (1 - (radius/self.pedestal_radius)**2)**self.ion_density_peaking_factor
+                        prod *= (
+                            1 - (radius / self.pedestal_radius) ** 2
+                        ) ** self.ion_density_peaking_factor
 
                         density_loc = self.ion_density_pedestal + prod
                     else:
                         prod = self.ion_density_pedestal - self.ion_density_separatrix
-                        prod *= (self.major_radius - radius)/(self.major_radius - self.pedestal_radius)
+                        prod *= (self.major_radius - radius) / (
+                            self.major_radius - self.pedestal_radius
+                        )
                         density_loc = self.ion_density_separatrix + prod
                     density.append(density_loc)
                 density = np.array(density)
@@ -134,21 +141,35 @@ def ion_temperature(self, r):
             float, ndarray: ion temperature (keV)
         """
         if self.mode == "L":
-            temperature = self.ion_temperature_centre * \
-                (1 - (r/self.major_radius)**2)**self.ion_temperature_peaking_factor
+            temperature = (
+                self.ion_temperature_centre
+                * (1 - (r / self.major_radius) ** 2)
+                ** self.ion_temperature_peaking_factor
+            )
         elif self.mode in ["H", "A"]:
             # TODO: find an alternative to iterating through the array
             if isinstance(r, np.ndarray):
                 temperature = []
                 for radius in r:
                     if 0 < radius < self.pedestal_radius:
-                        prod = self.ion_temperature_centre - self.ion_temperature_pedestal
-                        prod *= (1 - (radius/self.pedestal_radius)**self.ion_temperature_beta)**self.ion_temperature_peaking_factor
+                        prod = (
+                            self.ion_temperature_centre - self.ion_temperature_pedestal
+                        )
+                        prod *= (
+                            1
+                            - (radius / self.pedestal_radius)
+                            ** self.ion_temperature_beta
+                        ) ** self.ion_temperature_peaking_factor
 
                         temperature_loc = self.ion_temperature_pedestal + prod
                     else:
-                        prod = self.ion_temperature_pedestal - self.ion_temperature_separatrix
-                        prod *= (self.major_radius - radius)/(self.major_radius - self.pedestal_radius)
+                        prod = (
+                            self.ion_temperature_pedestal
+                            - self.ion_temperature_separatrix
+                        )
+                        prod *= (self.major_radius - radius) / (
+                            self.major_radius - self.pedestal_radius
+                        )
                         temperature_loc = self.ion_temperature_separatrix + prod
                     temperature.append(temperature_loc)
                 temperature = np.array(temperature)
@@ -165,29 +186,32 @@ def convert_a_alpha_to_R_Z(self, a, alpha):
         Returns:
             ((float, ndarray), (float, ndarray)): (R, Z) coordinates
         """
-        shafranov_shift = self.shafranov_factor*(1.0-(a/self.minor_radius)**2)
-        R = self.major_radius + \
-            a*np.cos(alpha + (self.triangularity*np.sin(alpha))) + \
-            shafranov_shift
-        Z = self.elongation*a*np.sin(alpha)
+        shafranov_shift = self.shafranov_factor * (1.0 - (a / self.minor_radius) ** 2)
+        R = (
+            self.major_radius
+            + a * np.cos(alpha + (self.triangularity * np.sin(alpha)))
+            + shafranov_shift
+        )
+        Z = self.elongation * a * np.sin(alpha)
         return (R, Z)
 
     def sample_sources(self):
         """Samples self.sample_size neutrons and creates attributes .densities
-            (ion density), .temperatures (ion temperature), .strengths
-            (neutron source density) and .RZ (coordinates)
+        (ion density), .temperatures (ion temperature), .strengths
+        (neutron source density) and .RZ (coordinates)
         """
         # create a sample of (a, alpha) coordinates
-        a = np.random.random(self.sample_size)*self.minor_radius
-        alpha = np.random.random(self.sample_size)*2*np.pi
+        a = np.random.random(self.sample_size) * self.minor_radius
+        alpha = np.random.random(self.sample_size) * 2 * np.pi
 
         # compute densities, temperatures, neutron source densities and
         # convert coordinates
         self.densities = self.ion_density(a)
         self.temperatures = self.ion_temperature(a)
         self.neutron_source_density = neutron_source_density(
-            self.densities, self.temperatures)
-        self.strengths = self.neutron_source_density/sum(self.neutron_source_density)
+            self.densities, self.temperatures
+        )
+        self.strengths = self.neutron_source_density / sum(self.neutron_source_density)
         self.RZ = self.convert_a_alpha_to_R_Z(a, alpha)
 
     def make_openmc_sources(self):
@@ -218,11 +242,13 @@ def make_openmc_sources(self):
 
             # create a ring source
             my_source.space = openmc.stats.CylindricalIndependent(
-                r=radius, phi=angle, z=z_values, origin=(0.0, 0.0, 0.0))
+                r=radius, phi=angle, z=z_values, origin=(0.0, 0.0, 0.0)
+            )
 
             my_source.angle = openmc.stats.Isotropic()
             my_source.energy = openmc.stats.Muir(
-                e0=14080000.0, m_rat=5.0, kt=self.temperatures[i])
+                e0=14080000.0, m_rat=5.0, kt=self.temperatures[i]
+            )
 
             # the strength of the source (its probability) is given by
             # self.strengths
@@ -244,7 +270,7 @@ def neutron_source_density(ion_density, ion_temperature):
     Returns:
         float, ndarray: Neutron source density (neutron/s/m3)
     """
-    return ion_density**2*DT_xs(ion_temperature)
+    return ion_density ** 2 * DT_xs(ion_temperature)
 
 
 def DT_xs(ion_temperature):
@@ -264,12 +290,12 @@ def DT_xs(ion_temperature):
         2.5773408e-3,
         6.1880463e-5,
         6.6024089e-2,
-        8.1215505e-3
-        ]
-    prod = ion_temperature*(c[2]+ion_temperature*(c[3]-c[4]*ion_temperature))
-    prod *= 1/(1.0+ion_temperature*(c[5]+c[6]*ion_temperature))
+        8.1215505e-3,
+    ]
+    prod = ion_temperature * (c[2] + ion_temperature * (c[3] - c[4] * ion_temperature))
+    prod *= 1 / (1.0 + ion_temperature * (c[5] + c[6] * ion_temperature))
     U = 1 - prod
 
-    val = c[0]/(U**(5/6)*ion_temperature**(2/3))
-    val *= np.exp(-c[1]*(U/ion_temperature)**(1/3))
+    val = c[0] / (U ** (5 / 6) * ion_temperature ** (2 / 3))
+    val *= np.exp(-c[1] * (U / ion_temperature) ** (1 / 3))
     return val