diff --git a/stardis/radiation_field/radiation_field_solvers/base.py b/stardis/radiation_field/radiation_field_solvers/base.py index 981c7d03..3f83319e 100644 --- a/stardis/radiation_field/radiation_field_solvers/base.py +++ b/stardis/radiation_field/radiation_field_solvers/base.py @@ -120,8 +120,8 @@ def single_theta_trace_parallel( mean_alphas = np.exp((np.log(alphas[1:]) + np.log(alphas[:-1])) * 0.5) taus = np.zeros_like(mean_alphas, dtype=np.float64) - for gap_index in numba.prange(taus.shape[0]): - for nu_index in range(taus.shape[1]): + for nu_index in numba.prange(taus.shape[1]): + for gap_index in range(taus.shape[0]): taus[gap_index, nu_index] = ( mean_alphas[gap_index, nu_index] * ray_dist_to_next_depth_point[gap_index] @@ -133,7 +133,7 @@ def single_theta_trace_parallel( I_nu_theta = np.zeros((no_of_depth_gaps + 1, len(tracing_nus))) I_nu_theta[0] = source[ 0 - ] # the innermost depth point is approximated as a blackbody + ]*0 # the innermost depth point is approximated as a blackbody w0, w1, w2 = calc_weights_parallel(taus) @@ -329,16 +329,14 @@ def raytrace( end_theta = (np.pi / 2) - (dtheta / 2) thetas = np.linspace(start_theta, end_theta, no_of_thetas) weights = 2 * np.pi * dtheta * np.sin(thetas) * np.cos(thetas) - - if True: + spherical=True + if spherical: ray_distances = calculate_spherical_ray(thetas, stellar_model.geometry.r) - print(ray_distances) - # print(ray_distances.shape) + # print(ray_distances) else: ray_distances = stellar_model.geometry.dist_to_next_depth_point.reshape( -1, 1 ) / np.cos(thetas) - # print(ray_distances.shape) ###TODO: Thetas should probably be held by the model? Then can be passed in from there. if n_threads == 1: # Single threaded @@ -374,28 +372,16 @@ def raytrace( def calculate_spherical_ray(thetas, depth_points_radii): ###NOTE: This will need to be revisited to handle some rays more carefully if they don't go through the star ray_distance_through_layer_by_impact_parameter = np.zeros( - (len(depth_points_radii) - 1, len(thetas)) + (len(depth_points_radii)-1, len(thetas)) ) for theta_index, theta in enumerate(thetas): b = depth_points_radii[-1] * np.sin(theta) # impact parameter of the ray - - deepest_ray_layer = np.argmin(np.abs(depth_points_radii - b)) - for i in range(deepest_ray_layer + 1, len(depth_points_radii)): - ray_z_coordinate_grid = np.sqrt(depth_points_radii**2 - b**2) - ray_distance = np.diff(ray_z_coordinate_grid) - ray_distance_through_layer_by_impact_parameter[i - 1, theta_index] = ( - ray_distance[i - 1] - ) - - # ray_depth_selection_mask = ( - # b < depth_points_radii[:-1] - # ) # mask for the depth points that the ray will pass through. - # ray_z_coordinate_grid = np.sqrt(depth_points_radii**2 - b**2) - - # ray_distance = np.diff(ray_z_coordinate_grid) - # ray_distance_through_layer_by_impact_parameter[ - # ray_depth_selection_mask, theta_index - # ] = ray_distance[ray_depth_selection_mask] - + deepest_ray_layer_index = np.argmin(np.abs(depth_points_radii - b)) + ray_z_coordinate_grid = np.sqrt(depth_points_radii**2 - b**2) + ray_distance = np.diff(ray_z_coordinate_grid) + # ray_distance_through_layer_by_impact_parameter[~np.isnan(ray_distance), theta_index] = ray_distance[~np.isnan(ray_distance)] + valid_layers = np.arange(deepest_ray_layer_index, len(depth_points_radii)-1) + ray_distance_through_layer_by_impact_parameter[valid_layers, theta_index] = ray_distance[valid_layers].value + return ray_distance_through_layer_by_impact_parameter