diff --git a/parastell/nwl_utils.py b/parastell/nwl_utils.py index 4fa8ab11..6823b38e 100644 --- a/parastell/nwl_utils.py +++ b/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 \ No newline at end of file