Merge remote-tracking branch 'svalinn/main' into oo_version

parastell/nwl_utils.py

@@ -210,17 +210,48 @@ def plot(nwl_mat, phi_pts, theta_pts, num_levels):
     fig, ax = plt.subplots()
     CF = ax.contourf(phi_pts, theta_pts, nwl_mat, levels = levels)
     cbar = plt.colorbar(CF)
-    cbar.ax.set_ylabel('NWL (MW)')
+    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):
+    """
+    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 
+            points in the order given should result in a polygon
+
+    Returns:
+        area (float): approximation of area
+    """
+    # triangle 1
+    v1 = corners[3] - corners[0]
+    v2 = corners[2] - corners[0]
+
+    v3 = np.cross(v1,v2)
+
+    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)
+
+    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
+    source_file, ss_file, plas_eq, tor_ext, pol_ext, wall_s, num_phi = 100,
+    num_theta = 101, num_levels = 10, num_crossings = None
     ):
-    """Computes and plots NWL.
+    """Computes and plots NWL. Assumes toroidal extent is less than 360 degrees
 
     Arguments:
         source_file (str): path to OpenMC surface source file.
@@ -232,11 +263,22 @@ def nwl_plot(
         num_phi (int): number of toroidal angle bins (defaults to 101).
         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
+
+    Returns:
+        nwl_mat (numpy array): array used to create the NWL plot
+        phi_pts (numpy array): phi axis of NWL plot
+        theta_pts (numpy array): theta axis of NWL plot
+        area_array (numpy array): area array used to normalize nwl_mat
     """
     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)
@@ -257,6 +299,12 @@ def nwl_plot(
         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
+
     # 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)
@@ -275,4 +323,30 @@ def nwl_plot(
 
     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))
+
+    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]
+
+    # construct area array
+    area_array = np.zeros((num_phi, num_theta))
+
+    for phi_index in range(num_phi):
+        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])
+            area = area_from_corners(corners)
+            area_array[phi_index,theta_index] = area
+
+    nwl_mat = nwl_mat/area_array
     plot(nwl_mat, phi_pts, theta_pts, num_levels)
+
+    return nwl_mat, phi_pts, theta_pts, area_array