From 435b921df97b646784bea433a274cac192f45c25 Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Tue, 18 Jun 2024 19:09:31 -0500
Subject: [PATCH 1/9] parallelized coordinate transform

---
 parastell/nwl_utils.py | 202 ++++++++++++++++++++++++-----------------
 1 file changed, 120 insertions(+), 82 deletions(-)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 2eba5fff..487a486b 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -4,6 +4,8 @@
 import openmc
 import pystell.read_vmec as read_vmec
 import matplotlib.pyplot as plt
+import concurrent.futures
+import os
 
 
 def extract_ss(ss_file):
@@ -17,8 +19,8 @@ def extract_ss(ss_file):
             Returned only if source mesh is generated.
     """
     strengths = []
-    
-    file_obj = open(ss_file, 'r')
+
+    file_obj = open(ss_file, "r")
     data = file_obj.readlines()
     for line in data:
         strengths.append(float(line))
@@ -37,48 +39,43 @@ def nwl_transport(dagmc_geom, source_mesh, tor_ext, ss_file, num_parts):
         num_parts (int): number of source particles to simulate.
     """
     tor_ext = np.deg2rad(tor_ext)
-    
+
     strengths = extract_ss(ss_file)
 
     # Initialize OpenMC model
     model = openmc.model.Model()
 
-    dag_univ = openmc.DAGMCUniverse(dagmc_geom, auto_geom_ids = False)
+    dag_univ = openmc.DAGMCUniverse(dagmc_geom, auto_geom_ids=False)
 
     # Define problem boundaries
-    vac_surf = openmc.Sphere(
-        r = 10000, surface_id = 9999, boundary_type = 'vacuum'
-    )
-    per_init = openmc.YPlane(
-        boundary_type = 'periodic',
-        surface_id = 9991
-    )
+    vac_surf = openmc.Sphere(r=10000, surface_id=9999, boundary_type="vacuum")
+    per_init = openmc.YPlane(boundary_type="periodic", surface_id=9991)
     per_fin = openmc.Plane(
-        a = np.sin(tor_ext),
-        b = -np.cos(tor_ext),
-        c = 0,
-        d = 0,
-        boundary_type = 'periodic',
-        surface_id = 9990
+        a=np.sin(tor_ext),
+        b=-np.cos(tor_ext),
+        c=0,
+        d=0,
+        boundary_type="periodic",
+        surface_id=9990,
     )
 
     # Define first period of geometry
-    region  = -vac_surf & +per_init & +per_fin
-    period = openmc.Cell(cell_id = 9996, region = region, fill = dag_univ)
+    region = -vac_surf & +per_init & +per_fin
+    period = openmc.Cell(cell_id=9996, region=region, fill=dag_univ)
     geometry = openmc.Geometry([period])
     model.geometry = geometry
 
     # Define run settings
     settings = openmc.Settings()
-    settings.run_mode = 'fixed source'
+    settings.run_mode = "fixed source"
     settings.particles = num_parts
     settings.batches = 1
 
     # Define neutron source
-    mesh = openmc.UnstructuredMesh(source_mesh, 'moab')
+    mesh = openmc.UnstructuredMesh(source_mesh, "moab")
     src = openmc.IndependentSource()
     src.space = openmc.stats.MeshSpatial(
-        mesh, strengths = strengths, volume_normalized = False
+        mesh, strengths=strengths, volume_normalized=False
     )
     src.angle = openmc.stats.Isotropic()
     src.energy = openmc.stats.Discrete([14.1e6], [1.0])
@@ -86,8 +83,8 @@ def nwl_transport(dagmc_geom, source_mesh, tor_ext, ss_file, num_parts):
 
     # Track surface crossings
     settings.surf_source_write = {
-        'surface_ids': [1],
-        'max_particles': num_parts
+        "surface_ids": [1],
+        "max_particles": num_parts,
     }
 
     model.settings = settings
@@ -112,8 +109,8 @@ def min_problem(theta, vmec, wall_s, phi, pt):
     # Compute first wall point
     fw_pt = np.array(vmec.vmec2xyz(wall_s, theta, phi))
     m2cm = 100
-    fw_pt = fw_pt*m2cm
-    
+    fw_pt = fw_pt * m2cm
+
     diff = np.linalg.norm(pt - fw_pt)
 
     return diff
@@ -134,9 +131,7 @@ def find_coords(vmec, wall_s, phi, pt):
     """
     # Solve for the poloidal angle via minimization
     theta = direct(
-        min_problem,
-        bounds = [(-np.pi, np.pi)],
-        args = (vmec, wall_s, phi, pt)
+        min_problem, bounds=[(-np.pi, np.pi)], args=(vmec, wall_s, phi, pt)
     )
     # Extract angle
     theta = theta.x[0]
@@ -144,10 +139,19 @@ def find_coords(vmec, wall_s, phi, pt):
     return theta
 
 
-def flux_coords(vmec, wall_s, coords):
+def find_coord(data):
+    print("pid ", os.getpid())
+    thetas = []
+    vmec = read_vmec.VMECData(data[0])
+    for t in data[2]:
+        thetas.append(find_coords(vmec, data[1], t[0], t[1]))
+    return thetas
+
+
+def flux_coords(plas_eq, wall_s, coords):
     """Computes flux-coordinate toroidal and poloidal angles corresponding to
     specified Cartesian coordinates.
-    
+
     Arguments:
         vmec (object): plasma equilibrium object.
         wall_s (float): closed flux surface label extrapolation at wall.
@@ -159,15 +163,41 @@ def flux_coords(vmec, wall_s, coords):
         theta_coords (array of float): poloidal angles of surface crossings
             (rad).
     """
-    phi_coords = np.arctan2(coords[:,1], coords[:,0])
-    theta_coords = []
-    
-    for pt, phi in zip(coords, phi_coords):
-        theta = find_coords(vmec, wall_s, phi, pt)
-        theta_coords.append(theta)
+
+    phi_coords = np.arctan2(coords[:, 1], coords[:, 0])
+
+    num_chunks = 5
+
+    chunk_size = len(phi_coords) // num_chunks
+
+    chunks = []
+    for i in range(num_chunks):
+        subchunk = list(
+            zip(
+                phi_coords[i * chunk_size : (i + 1) * chunk_size],
+                coords[i * chunk_size : (i + 1) * chunk_size],
+            )
+        )
+        chunks.append((plas_eq, wall_s, subchunk))
+
+    with concurrent.futures.ProcessPoolExecutor(
+        max_workers=num_chunks
+    ) as executor:
+        theta_coords = np.array(
+            list(executor.map(find_coord, chunks))
+        ).flatten()
 
     return phi_coords, theta_coords
 
+    # phi_coords = np.arctan2(coords[:, 1], coords[:, 0])
+    # theta_coords = []
+
+    # for pt, phi in zip(coords, phi_coords):
+    #     theta = find_coords(vmec, wall_s, phi, pt)
+    #     theta_coords.append(theta)
+
+    # return phi_coords, theta_coords
+
 
 def extract_coords(source_file):
     """Extracts Cartesian coordinates of particle surface crossings given an
@@ -175,20 +205,20 @@ def extract_coords(source_file):
 
     Arguments:
         source_file (str): path to OpenMC surface source file.
-    
+
     Returns:
         coords (array of array of float): Cartesian coordinates of all particle
             surface crossings.
     """
     # Load source file
-    file = h5py.File(source_file, 'r')
+    file = h5py.File(source_file, "r")
     # Extract source information
-    dataset = file['source_bank']['r']
+    dataset = file["source_bank"]["r"]
     # Construct matrix of particle crossing coordinates
     coords = np.empty((len(dataset), 3))
-    coords[:,0] = dataset['x']
-    coords[:,1] = dataset['y']
-    coords[:,2] = dataset['z']
+    coords[:, 0] = dataset["x"]
+    coords[:, 1] = dataset["y"]
+    coords[:, 2] = dataset["z"]
 
     return coords
 
@@ -206,14 +236,14 @@ def plot(nwl_mat, phi_pts, theta_pts, num_levels):
     phi_pts = np.rad2deg(phi_pts)
     theta_pts = np.rad2deg(theta_pts)
 
-    levels = np.linspace(np.min(nwl_mat), np.max(nwl_mat), num = num_levels)
+    levels = np.linspace(np.min(nwl_mat), np.max(nwl_mat), num=num_levels)
     fig, ax = plt.subplots()
-    CF = ax.contourf(phi_pts, theta_pts, nwl_mat.T, levels = levels)
+    CF = ax.contourf(phi_pts, theta_pts, nwl_mat.T, levels=levels)
     cbar = plt.colorbar(CF)
-    cbar.ax.set_ylabel('NWL (MW/m2)')
-    plt.xlabel('Toroidal Angle (degrees)')
-    plt.ylabel('Poloidal Angle (degrees)')
-    fig.savefig('NWL.png')
+    cbar.ax.set_ylabel("NWL (MW/m2)")
+    plt.xlabel("Toroidal Angle (degrees)")
+    plt.ylabel("Poloidal Angle (degrees)")
+    fig.savefig("NWL.png")
 
 
 def area_from_corners(corners):
@@ -221,7 +251,7 @@ def area_from_corners(corners):
     Calculate an approximation of the area defined by 4 xyz points
 
     Arguments:
-        corners (4x3 numpy array): list of 4 (x,y,z) points. Connecting the 
+        corners (4x3 numpy array): list of 4 (x,y,z) points. Connecting the
             points in the order given should result in a polygon
 
     Returns:
@@ -231,26 +261,35 @@ def area_from_corners(corners):
     v1 = corners[3] - corners[0]
     v2 = corners[2] - corners[0]
 
-    v3 = np.cross(v1,v2)
+    v3 = np.cross(v1, v2)
 
-    area1 = np.sqrt(np.sum(np.square(v3)))/2
+    area1 = np.sqrt(np.sum(np.square(v3))) / 2
 
     # triangle 2
     v1 = corners[1] - corners[0]
     v2 = corners[2] - corners[0]
 
-    v3 = np.cross(v1,v2)
+    v3 = np.cross(v1, v2)
 
-    area2 = np.sqrt(np.sum(np.square(v3)))/2
+    area2 = np.sqrt(np.sum(np.square(v3))) / 2
 
     area = area1 + area2
 
     return area
 
+
 def nwl_plot(
-    source_file, ss_file, plas_eq, tor_ext, pol_ext, wall_s, num_phi = 101,
-    num_theta = 101, num_levels = 10, num_crossings = None
-    ):
+    source_file,
+    ss_file,
+    plas_eq,
+    tor_ext,
+    pol_ext,
+    wall_s,
+    num_phi=101,
+    num_theta=101,
+    num_levels=10,
+    num_crossings=None,
+):
     """Computes and plots NWL. Assumes toroidal extent is less than 360 degrees
 
     Arguments:
@@ -274,40 +313,39 @@ def nwl_plot(
     """
     tor_ext = np.deg2rad(tor_ext)
     pol_ext = np.deg2rad(pol_ext)
-    
+
     coords = extract_coords(source_file)
 
     if num_crossings is not None:
         coords = coords[0:num_crossings]
-    
-    # Load plasma equilibrium data
+
     vmec = read_vmec.VMECData(plas_eq)
 
-    phi_coords, theta_coords = flux_coords(vmec, wall_s, coords)
+    phi_coords, theta_coords = flux_coords(plas_eq, wall_s, coords)
 
     # Define minimum and maximum bin edges for each dimension
-    phi_min = 0 - tor_ext/num_phi/2
-    phi_max = tor_ext + tor_ext/num_phi/2
-    theta_min = -pol_ext/2 - pol_ext/num_theta/2
-    theta_max = pol_ext/2 + pol_ext/num_theta/2
-    
+    phi_min = 0 - tor_ext / num_phi / 2
+    phi_max = tor_ext + tor_ext / num_phi / 2
+    theta_min = -pol_ext / 2 - pol_ext / num_theta / 2
+    theta_max = pol_ext / 2 + pol_ext / num_theta / 2
+
     # Bin particle crossings
     count_mat, phi_bins, theta_bins = np.histogram2d(
         phi_coords,
         theta_coords,
-        bins = [num_phi, num_theta],
-        range = [[phi_min, phi_max], [theta_min, theta_max]]
+        bins=[num_phi, num_theta],
+        range=[[phi_min, phi_max], [theta_min, theta_max]],
     )
 
     # adjust endpoints to eliminate overlap
     phi_bins[0] = 0
     phi_bins[-1] = tor_ext
-    theta_bins[0] = -pol_ext/2
-    theta_bins[-1] = pol_ext/2
+    theta_bins[0] = -pol_ext / 2
+    theta_bins[-1] = pol_ext / 2
 
     # Compute centroids of bin dimensions
-    phi_pts = np.linspace(0, tor_ext, num = num_phi)
-    theta_pts = np.linspace(-pol_ext/2, pol_ext/2, num = num_theta)
+    phi_pts = np.linspace(0, tor_ext, num=num_phi)
+    theta_pts = np.linspace(-pol_ext / 2, pol_ext / 2, num=num_theta)
 
     # Define fusion neutron energy (eV)
     n_energy = 14.1e6
@@ -321,16 +359,16 @@ def nwl_plot(
     # Define number of source particles
     num_parts = len(coords)
 
-    nwl_mat = count_mat*n_energy*eV2J*SS*J2MJ/num_parts
+    nwl_mat = count_mat * n_energy * eV2J * SS * J2MJ / num_parts
 
     # construct array of bin boundaries
-    bin_arr = np.zeros((num_phi+1, num_theta+1, 3))
+    bin_arr = np.zeros((num_phi + 1, num_theta + 1, 3))
 
     for phi_bin, phi in enumerate(phi_bins):
         for theta_bin, theta in enumerate(theta_bins):
 
-            x,y,z = vmec.vmec2xyz(wall_s,theta,phi)
-            bin_arr[phi_bin, theta_bin, :] = [x,y,z]
+            x, y, z = vmec.vmec2xyz(wall_s, theta, phi)
+            bin_arr[phi_bin, theta_bin, :] = [x, y, z]
 
     # construct area array
     area_array = np.zeros((num_phi, num_theta))
@@ -339,14 +377,14 @@ def nwl_plot(
         for theta_index in range(num_theta):
             # each bin has 4 (x,y,z) corners
             corner1 = bin_arr[phi_index, theta_index]
-            corner2 = bin_arr[phi_index, theta_index+1]
-            corner3 = bin_arr[phi_index+1, theta_index+1]
-            corner4 = bin_arr[phi_index+1, theta_index]
-            corners = np.array([corner1,corner2,corner3,corner4])
+            corner2 = bin_arr[phi_index, theta_index + 1]
+            corner3 = bin_arr[phi_index + 1, theta_index + 1]
+            corner4 = bin_arr[phi_index + 1, theta_index]
+            corners = np.array([corner1, corner2, corner3, corner4])
             area = area_from_corners(corners)
-            area_array[phi_index,theta_index] = area
+            area_array[phi_index, theta_index] = area
 
-    nwl_mat = nwl_mat/area_array
+    nwl_mat = nwl_mat / area_array
     plot(nwl_mat, phi_pts, theta_pts, num_levels)
 
     return nwl_mat, phi_pts, theta_pts, area_array

From 39b556b05666465dfac3ddbb0779a2d035b31bb2 Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Tue, 18 Jun 2024 22:20:32 -0500
Subject: [PATCH 2/9] split up computation of large surface_sources due to
 memory

---
 parastell/nwl_utils.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 487a486b..5cf3f495 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -140,7 +140,6 @@ def find_coords(vmec, wall_s, phi, pt):
 
 
 def find_coord(data):
-    print("pid ", os.getpid())
     thetas = []
     vmec = read_vmec.VMECData(data[0])
     for t in data[2]:
@@ -166,7 +165,7 @@ def flux_coords(plas_eq, wall_s, coords):
 
     phi_coords = np.arctan2(coords[:, 1], coords[:, 0])
 
-    num_chunks = 5
+    num_chunks = 10
 
     chunk_size = len(phi_coords) // num_chunks
 
@@ -187,7 +186,7 @@ def flux_coords(plas_eq, wall_s, coords):
             list(executor.map(find_coord, chunks))
         ).flatten()
 
-    return phi_coords, theta_coords
+    return phi_coords.tolist(), theta_coords.tolist()
 
     # phi_coords = np.arctan2(coords[:, 1], coords[:, 0])
     # theta_coords = []
@@ -289,6 +288,7 @@ def nwl_plot(
     num_theta=101,
     num_levels=10,
     num_crossings=None,
+    step_size = 1_000_000
 ):
     """Computes and plots NWL. Assumes toroidal extent is less than 360 degrees
 
@@ -321,7 +321,16 @@ def nwl_plot(
 
     vmec = read_vmec.VMECData(plas_eq)
 
-    phi_coords, theta_coords = flux_coords(plas_eq, wall_s, coords)
+    phi_coords = []
+    theta_coords = []
+
+    iterations = len(coords)//step_size
+
+    for i in range(iterations):
+        print(i, ' million')
+        phi_coord_subset, theta_coord_subset = flux_coords(plas_eq, wall_s, coords[i*step_size:(i+1)*step_size])
+        phi_coords += phi_coord_subset
+        theta_coords += theta_coord_subset
 
     # Define minimum and maximum bin edges for each dimension
     phi_min = 0 - tor_ext / num_phi / 2

From d0fd4b6ec7c678c8566a654821d2a0efb5cba972 Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Fri, 21 Jun 2024 08:27:44 -0500
Subject: [PATCH 3/9] update step sizes for using 10mil particles

---
 parastell/nwl_utils.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 5cf3f495..30768fc0 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -288,7 +288,7 @@ def nwl_plot(
     num_theta=101,
     num_levels=10,
     num_crossings=None,
-    step_size = 1_000_000
+    step_size=1_000_000,
 ):
     """Computes and plots NWL. Assumes toroidal extent is less than 360 degrees
 
@@ -324,11 +324,13 @@ def nwl_plot(
     phi_coords = []
     theta_coords = []
 
-    iterations = len(coords)//step_size
+    num_subsets = len(coords) // step_size
 
-    for i in range(iterations):
-        print(i, ' million')
-        phi_coord_subset, theta_coord_subset = flux_coords(plas_eq, wall_s, coords[i*step_size:(i+1)*step_size])
+    for i in range(num_subsets):
+        print(i, " million")
+        phi_coord_subset, theta_coord_subset = flux_coords(
+            plas_eq, wall_s, coords[i * step_size : (i + 1) * step_size]
+        )
         phi_coords += phi_coord_subset
         theta_coords += theta_coord_subset
 

From d21bd04378053227eef398851700d2070be5961e Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Fri, 21 Jun 2024 09:32:56 -0500
Subject: [PATCH 4/9] added args for number of threads

---
 parastell/nwl_utils.py | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 30768fc0..f820cd70 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -147,7 +147,7 @@ def find_coord(data):
     return thetas
 
 
-def flux_coords(plas_eq, wall_s, coords):
+def flux_coords(plas_eq, wall_s, coords, num_threads):
     """Computes flux-coordinate toroidal and poloidal angles corresponding to
     specified Cartesian coordinates.
 
@@ -165,22 +165,20 @@ def flux_coords(plas_eq, wall_s, coords):
 
     phi_coords = np.arctan2(coords[:, 1], coords[:, 0])
 
-    num_chunks = 10
-
-    chunk_size = len(phi_coords) // num_chunks
+    chunk_size = len(phi_coords) // num_threads
 
     chunks = []
-    for i in range(num_chunks):
-        subchunk = list(
+    for i in range(num_threads):
+        chunk = list(
             zip(
                 phi_coords[i * chunk_size : (i + 1) * chunk_size],
                 coords[i * chunk_size : (i + 1) * chunk_size],
             )
         )
-        chunks.append((plas_eq, wall_s, subchunk))
+        chunks.append((plas_eq, wall_s, chunk))
 
     with concurrent.futures.ProcessPoolExecutor(
-        max_workers=num_chunks
+        max_workers=num_threads
     ) as executor:
         theta_coords = np.array(
             list(executor.map(find_coord, chunks))
@@ -288,7 +286,8 @@ def nwl_plot(
     num_theta=101,
     num_levels=10,
     num_crossings=None,
-    step_size=1_000_000,
+    step_size=None,
+    num_threads=2,
 ):
     """Computes and plots NWL. Assumes toroidal extent is less than 360 degrees
 
@@ -324,12 +323,17 @@ def nwl_plot(
     phi_coords = []
     theta_coords = []
 
-    num_subsets = len(coords) // step_size
+    # split up the work to avoid memory issues
+    iterations = 1
+    if step_size is not None:
+        iterations = len(coords) // step_size
 
-    for i in range(num_subsets):
-        print(i, " million")
+    for i in range(iterations):
         phi_coord_subset, theta_coord_subset = flux_coords(
-            plas_eq, wall_s, coords[i * step_size : (i + 1) * step_size]
+            plas_eq,
+            wall_s,
+            coords[i * step_size : (i + 1) * step_size],
+            num_threads,
         )
         phi_coords += phi_coord_subset
         theta_coords += theta_coord_subset

From 8b66fa8b79458166d5a523c7b8f53759e079b3ca Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Fri, 21 Jun 2024 15:01:58 -0500
Subject: [PATCH 5/9] consolidate helper and doer functions

---
 parastell/nwl_utils.py | 49 ++++++++++++++++--------------------------
 1 file changed, 19 insertions(+), 30 deletions(-)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index f820cd70..4135a4b9 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -116,34 +116,32 @@ def min_problem(theta, vmec, wall_s, phi, pt):
     return diff
 
 
-def find_coords(vmec, wall_s, phi, pt):
+def find_coords(data):
     """Solves for poloidal angle of plasma equilibrium corresponding to
-    specified Cartesian coordinates.
+    specified Cartesian coordinates. Takes a single arg so it works nicely with
+    ProcessPoolExecutor
 
     Arguments:
-        vmec (object): plasma equilibrium object.
-        wall_s (float): closed flux surface label extrapolation at wall.
-        phi (float): toroidal angle (rad).
-        pt (array of float): Cartesian coordinates of interest (cm).
+        data (tuple of (str, float, list of tuple of float)): First element is
+            the path to the plasma equilibrium file, second is the wall_s value,
+            3rd is the list of phi, xyz coordinate pairs to solve for theta at.
 
     Returns:
-        theta (float): poloidal angle (rad).
+        thetas (list of float): poloidal angles (rad) corresponding to phi xyz
+            coordinate pairs.
     """
-    # Solve for the poloidal angle via minimization
-    theta = direct(
-        min_problem, bounds=[(-np.pi, np.pi)], args=(vmec, wall_s, phi, pt)
-    )
-    # Extract angle
-    theta = theta.x[0]
-
-    return theta
-
-
-def find_coord(data):
     thetas = []
     vmec = read_vmec.VMECData(data[0])
-    for t in data[2]:
-        thetas.append(find_coords(vmec, data[1], t[0], t[1]))
+    wall_s = data[1]
+    phi_xyz_coords = data[2]
+
+    for coords in phi_xyz_coords:
+        theta = direct(
+            min_problem,
+            bounds=[(-np.pi, np.pi)],
+            args=(vmec, wall_s, coords[0], coords[1]),
+        )
+        thetas.append(theta.x[0])
     return thetas
 
 
@@ -181,20 +179,11 @@ def flux_coords(plas_eq, wall_s, coords, num_threads):
         max_workers=num_threads
     ) as executor:
         theta_coords = np.array(
-            list(executor.map(find_coord, chunks))
+            list(executor.map(find_coords, chunks))
         ).flatten()
 
     return phi_coords.tolist(), theta_coords.tolist()
 
-    # phi_coords = np.arctan2(coords[:, 1], coords[:, 0])
-    # theta_coords = []
-
-    # for pt, phi in zip(coords, phi_coords):
-    #     theta = find_coords(vmec, wall_s, phi, pt)
-    #     theta_coords.append(theta)
-
-    # return phi_coords, theta_coords
-
 
 def extract_coords(source_file):
     """Extracts Cartesian coordinates of particle surface crossings given an

From 94d85b353fffd3bcbec9a0f0e00d1c77f9f81319 Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Wed, 26 Jun 2024 18:33:31 -0500
Subject: [PATCH 6/9] works with thread/particle counts that dont divide evenly

---
 Examples/nwl_plot_example.py | 14 ++++++++------
 parastell/nwl_utils.py       | 37 +++++++++++++++++++++---------------
 2 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/Examples/nwl_plot_example.py b/Examples/nwl_plot_example.py
index 648ea593..b621fa70 100644
--- a/Examples/nwl_plot_example.py
+++ b/Examples/nwl_plot_example.py
@@ -2,19 +2,21 @@
 
 
 # Define simulation parameters
-dagmc_geom = 'nwl_geom.h5m'
-source_mesh = 'source_mesh.h5m'
+dagmc_geom = "nwl_geom.h5m"
+source_mesh = "source_mesh.h5m"
 tor_ext = 90.0
-ss_file = 'source_strengths.txt'
+ss_file = "source_strengths.txt"
 num_parts = 100000
 
 nwl.nwl_transport(dagmc_geom, source_mesh, tor_ext, ss_file, num_parts)
 
 # Define first wall geometry and plotting parameters
-source_file = 'surface_source.h5'
-plas_eq = 'wout_vmec.nc'
+source_file = "surface_source.h5"
+plas_eq = "wout_vmec.nc"
 tor_ext = 90.0
 pol_ext = 360.0
 wall_s = 1.08
 
-nwl.nwl_plot(source_file, ss_file, plas_eq, tor_ext, pol_ext, wall_s)
+nwl.nwl_plot(
+    source_file, ss_file, plas_eq, tor_ext, pol_ext, wall_s, num_threads=2
+)
diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 4135a4b9..1b34eb1e 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -5,7 +5,7 @@
 import pystell.read_vmec as read_vmec
 import matplotlib.pyplot as plt
 import concurrent.futures
-import os
+import math
 
 
 def extract_ss(ss_file):
@@ -162,9 +162,7 @@ def flux_coords(plas_eq, wall_s, coords, num_threads):
     """
 
     phi_coords = np.arctan2(coords[:, 1], coords[:, 0])
-
-    chunk_size = len(phi_coords) // num_threads
-
+    chunk_size = math.ceil(len(phi_coords) / num_threads)
     chunks = []
     for i in range(num_threads):
         chunk = list(
@@ -178,11 +176,11 @@ def flux_coords(plas_eq, wall_s, coords, num_threads):
     with concurrent.futures.ProcessPoolExecutor(
         max_workers=num_threads
     ) as executor:
-        theta_coords = np.array(
-            list(executor.map(find_coords, chunks))
-        ).flatten()
-
-    return phi_coords.tolist(), theta_coords.tolist()
+        theta_coords = list(executor.map(find_coords, chunks))
+        theta_coords = [
+            theta_coord for chunk in theta_coords for theta_coord in chunk
+        ]
+    return phi_coords.tolist(), theta_coords
 
 
 def extract_coords(source_file):
@@ -275,8 +273,8 @@ def nwl_plot(
     num_theta=101,
     num_levels=10,
     num_crossings=None,
-    step_size=None,
-    num_threads=2,
+    chunk_size=None,
+    num_threads=None,
 ):
     """Computes and plots NWL. Assumes toroidal extent is less than 360 degrees
 
@@ -291,7 +289,12 @@ def nwl_plot(
         num_theta (int): number of poloidal angle bins (defaults to 101).
         num_levels (int): number of contour regions (defaults to 10).
         num_crossings (int): number of crossings to use from the surface source.
-            If none, then all crossings will be used
+            If None, then all crossings will be used
+        chunk_size (int): number of crossings to calculate at once, to help
+            with potential memory limits. If None all crossings will be done
+            at once
+        num_threads (int): number of threads to use for NWL calculations, if
+            None, all threads will be use.
 
     Returns:
         nwl_mat (numpy array): array used to create the NWL plot
@@ -314,19 +317,23 @@ def nwl_plot(
 
     # split up the work to avoid memory issues
     iterations = 1
-    if step_size is not None:
-        iterations = len(coords) // step_size
+
+    if chunk_size is not None:
+        iterations = math.ceil(len(coords) / chunk_size)
 
     for i in range(iterations):
         phi_coord_subset, theta_coord_subset = flux_coords(
             plas_eq,
             wall_s,
-            coords[i * step_size : (i + 1) * step_size],
+            coords[i * chunk_size : (i + 1) * chunk_size],
             num_threads,
         )
         phi_coords += phi_coord_subset
         theta_coords += theta_coord_subset
 
+    print(len(phi_coords))
+    print(len(theta_coords))
+
     # Define minimum and maximum bin edges for each dimension
     phi_min = 0 - tor_ext / num_phi / 2
     phi_max = tor_ext + tor_ext / num_phi / 2

From 6be9698e42369df1f33006d4125af9e72d0f869d Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Wed, 26 Jun 2024 18:42:50 -0500
Subject: [PATCH 7/9] detect number of threads available as default

---
 parastell/nwl_utils.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 1b34eb1e..7fa6cda3 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -5,6 +5,7 @@
 import pystell.read_vmec as read_vmec
 import matplotlib.pyplot as plt
 import concurrent.futures
+from multiprocessing import cpu_count
 import math
 
 
@@ -310,6 +311,9 @@ def nwl_plot(
     if num_crossings is not None:
         coords = coords[0:num_crossings]
 
+    if num_threads is None:
+        num_threads = cpu_count()
+
     vmec = read_vmec.VMECData(plas_eq)
 
     phi_coords = []

From fa8d128b56ea5ab609ac1949b138641611e91747 Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Wed, 3 Jul 2024 09:54:32 -0500
Subject: [PATCH 8/9] handle when chunk_size is none correctly

---
 parastell/nwl_utils.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 7fa6cda3..75921580 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -324,6 +324,8 @@ def nwl_plot(
 
     if chunk_size is not None:
         iterations = math.ceil(len(coords) / chunk_size)
+    else:
+        chunk_size = len(coords)
 
     for i in range(iterations):
         phi_coord_subset, theta_coord_subset = flux_coords(
@@ -335,9 +337,6 @@ def nwl_plot(
         phi_coords += phi_coord_subset
         theta_coords += theta_coord_subset
 
-    print(len(phi_coords))
-    print(len(theta_coords))
-
     # Define minimum and maximum bin edges for each dimension
     phi_min = 0 - tor_ext / num_phi / 2
     phi_max = tor_ext + tor_ext / num_phi / 2

From 37c4ce9b20ffab14ecfc7b9f0d85838a485a9afe Mon Sep 17 00:00:00 2001
From: Edgar <ejpflug@gmail.com>
Date: Tue, 16 Jul 2024 08:59:30 -0500
Subject: [PATCH 9/9] changes per pr comments

---
 parastell/nwl_utils.py | 26 +++++++++++---------------
 1 file changed, 11 insertions(+), 15 deletions(-)

diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py
index 75921580..4a3c9735 100644
--- a/parastell/nwl_utils.py
+++ b/parastell/nwl_utils.py
@@ -177,9 +177,11 @@ def flux_coords(plas_eq, wall_s, coords, num_threads):
     with concurrent.futures.ProcessPoolExecutor(
         max_workers=num_threads
     ) as executor:
-        theta_coords = list(executor.map(find_coords, chunks))
+        theta_coord_chunks = list(executor.map(find_coords, chunks))
         theta_coords = [
-            theta_coord for chunk in theta_coords for theta_coord in chunk
+            theta_coord
+            for theta_coord_chunk in theta_coord_chunks
+            for theta_coord in theta_coord_chunk
         ]
     return phi_coords.tolist(), theta_coords
 
@@ -275,7 +277,7 @@ def nwl_plot(
     num_levels=10,
     num_crossings=None,
     chunk_size=None,
-    num_threads=None,
+    num_threads=1,
 ):
     """Computes and plots NWL. Assumes toroidal extent is less than 360 degrees
 
@@ -294,8 +296,8 @@ def nwl_plot(
         chunk_size (int): number of crossings to calculate at once, to help
             with potential memory limits. If None all crossings will be done
             at once
-        num_threads (int): number of threads to use for NWL calculations, if
-            None, all threads will be use.
+        num_threads (int): number of threads to use for NWL calculations,
+            defaults to 1.
 
     Returns:
         nwl_mat (numpy array): array used to create the NWL plot
@@ -311,23 +313,17 @@ def nwl_plot(
     if num_crossings is not None:
         coords = coords[0:num_crossings]
 
-    if num_threads is None:
-        num_threads = cpu_count()
-
     vmec = read_vmec.VMECData(plas_eq)
 
     phi_coords = []
     theta_coords = []
 
-    # split up the work to avoid memory issues
-    iterations = 1
-
-    if chunk_size is not None:
-        iterations = math.ceil(len(coords) / chunk_size)
-    else:
+    if chunk_size is None:
         chunk_size = len(coords)
 
-    for i in range(iterations):
+    chunks = math.ceil(len(coords) / chunk_size)
+
+    for i in range(chunks):
         phi_coord_subset, theta_coord_subset = flux_coords(
             plas_eq,
             wall_s,