From c1c745320a7c74dd4c1bb8673adfb27567ec5ea1 Mon Sep 17 00:00:00 2001
From: Phoebe Pearce
Date: Thu, 22 Aug 2024 23:31:14 +1000
Subject: [PATCH] more jit, bug fixes
---
examples/analytical_rt_comparison.py | 6 +-
examples/phong_scattering_2.py | 85 ++++++++++---------
.../matrix_formalism/multiply_matrices.py | 6 +-
rayflare/options.py | 2 +-
rayflare/ray_tracing/rt_common.py | 68 +++++++++------
rayflare/ray_tracing/rt_structure.py | 8 +-
tests/test_analytical_rt.py | 9 +-
7 files changed, 104 insertions(+), 80 deletions(-)
diff --git a/examples/analytical_rt_comparison.py b/examples/analytical_rt_comparison.py
index b980620..4caf341 100644
--- a/examples/analytical_rt_comparison.py
+++ b/examples/analytical_rt_comparison.py
@@ -13,8 +13,8 @@
options.wavelength = wl
-options.nx = 50
-options.ny = 50
+options.nx = 100
+options.ny = 100
options.n_rays = 1 * options.nx ** 2
# options.x_limits = [2.5, 7.5]
# options.y_limits = [2.5, 7.5]
@@ -51,7 +51,7 @@
overwrite=True,
)
-angle_in = np.linspace(0, 70, 10)
+angle_in = np.linspace(0, 70, 2)
n_int_R_a = np.zeros((len(angle_in), len(wl)))
n_int_T_a = np.zeros((len(angle_in), len(wl)))
diff --git a/examples/phong_scattering_2.py b/examples/phong_scattering_2.py
index 36b531c..b9d40dc 100644
--- a/examples/phong_scattering_2.py
+++ b/examples/phong_scattering_2.py
@@ -27,7 +27,7 @@
SiN = material("Si3N4")()
# number of x and y points to scan across
-nxy = 30
+nxy = 50
# setting options
options = default_options()
@@ -36,10 +36,12 @@
options.ny = nxy
options.n_rays = 2 * nxy**2
options.depth_spacing = si("0.1um")
-options.parallel = False
+options.parallel = True
options.analytical_ray_tracing = 0
options.I_thresh = 0.005
options.randomize_surface = True
+options.project_name = "phong_scattering"
+options.lookuptable_angles = 200
# analytical: phong and regular pyramids 15 degree rear look the same
front_opening_deg = 50
@@ -58,7 +60,10 @@
args_same = { "materials": [Si],
"widths": [d],
"incidence": Air,
- "transmission": Air}
+ "transmission": Air,
+ "use_TMM": False,
+ "options": options,
+ "overwrite": True}
# set up ray-tracing options
rtstr_planar = rt_structure(
@@ -79,44 +84,44 @@
result_list_phong = []
start = time()
-#
-# for alpha in alpha_values:
-# flat_phong[0].phong_options[0] = alpha
-# rtstr_phong = rt_structure(
-# textures=[triangle_surf, flat_phong],
-# **args_same
-# )
-# result_list_phong.append(rtstr_phong.calculate(options))
-# result_planar = rtstr_planar.calculate(options)
+for alpha in alpha_values:
+ flat_phong[0].phong_options[0] = alpha
+ rtstr_phong = rt_structure(
+ textures=[triangle_surf, flat_phong],
+ **args_same
+ )
+ result_list_phong.append(rtstr_phong.calculate(options))
+
+result_planar = rtstr_planar.calculate(options)
result_flat_tri = rtstr_flat_tri.calculate(options)
-#
-# result_same_tri = rtstr_same_tri.calculate(options)
-#
+
+result_same_tri = rtstr_same_tri.calculate(options)
+
print("total time:", time() - start)
-#
-# R0 = result_planar['R0']
-# A_lambertian = (1-R0)*4*Si.n(options.wavelength)**2*Si.alpha(options.wavelength)*d/(1+4*Si.n(options.wavelength)**2*Si.alpha(options.wavelength)*d)
-#
-#
-# plt.figure()
-#
-# for i1, alpha in enumerate(alpha_values):
-# plt.plot(options.wavelength*1e9, result_list_phong[i1]["A_per_layer"], '-', color=pal[i1],
-# label=r"phong, $\alpha$ = " + str(alpha))
-#
-#
-# plt.plot(options.wavelength * 1e9, result_planar["A_per_layer"], '-.k',
-# label="planar", alpha=0.6)
-# plt.plot(options.wavelength * 1e9, result_flat_tri["A_per_layer"], '-k',
-# label="15 degree pyramid", alpha=0.6)
-# plt.plot(options.wavelength * 1e9, result_same_tri["A_per_layer"], '--k',
-# label="same angle pyramid", alpha=0.6)
-#
-# plt.ylim(0, 1)
-# plt.xlim(np.min(options.wavelength * 1e9), np.max(options.wavelength * 1e9))
-# plt.legend(title='Rear surface:')
-# plt.xlabel("Wavelength (nm)")
-# plt.ylabel("Absorption")
-# plt.show()
+
+R0 = result_planar['R0']
+A_lambertian = (1-R0)*4*Si.n(options.wavelength)**2*Si.alpha(options.wavelength)*d/(1+4*Si.n(options.wavelength)**2*Si.alpha(options.wavelength)*d)
+
+
+plt.figure()
+
+for i1, alpha in enumerate(alpha_values):
+ plt.plot(options.wavelength*1e9, result_list_phong[i1]["A_per_layer"], '-', color=pal[i1],
+ label=r"phong, $\alpha$ = " + str(alpha))
+
+
+plt.plot(options.wavelength * 1e9, result_planar["A_per_layer"], '-.k',
+ label="planar", alpha=0.6)
+plt.plot(options.wavelength * 1e9, result_flat_tri["A_per_layer"], '-k',
+ label="15 degree pyramid", alpha=0.6)
+plt.plot(options.wavelength * 1e9, result_same_tri["A_per_layer"], '--k',
+ label="same angle pyramid", alpha=0.6)
+
+plt.ylim(0, 1)
+plt.xlim(np.min(options.wavelength * 1e9), np.max(options.wavelength * 1e9))
+plt.legend(title='Rear surface:')
+plt.xlabel("Wavelength (nm)")
+plt.ylabel("Absorption")
+plt.show()
diff --git a/rayflare/matrix_formalism/multiply_matrices.py b/rayflare/matrix_formalism/multiply_matrices.py
index 5e73122..4f65ecc 100644
--- a/rayflare/matrix_formalism/multiply_matrices.py
+++ b/rayflare/matrix_formalism/multiply_matrices.py
@@ -163,16 +163,16 @@ def make_D(alphas, thick, thetas):
def dot_wl(mat, vec):
if len(mat.shape) == 3:
- result = einsum('ijk,ik->ij', mat, vec).todense()
+ result = einsum('ijk,ik->ij', mat, COO(vec)).todense()
if len(mat.shape) == 2:
- result = einsum('jk,ik->ij', mat, vec).todense()
+ result = einsum('jk,ik->ij', mat, COO(vec)).todense()
return result
def dot_wl_u2d(mat, vec):
- result = einsum('jk,ik->ij', mat, vec).todense()
+ result = einsum('jk,ik->ij', mat, COO(vec)).todense()
return result
diff --git a/rayflare/options.py b/rayflare/options.py
index e71ce2a..8bcbf6a 100644
--- a/rayflare/options.py
+++ b/rayflare/options.py
@@ -62,6 +62,6 @@ def __init__(self):
self.analytical_threshold = 0.99
# TMM options
- self.lookuptable_angles = 300
+ self.lookuptable_angles = 100
self.coherent = True
self.coherency_list = None
diff --git a/rayflare/ray_tracing/rt_common.py b/rayflare/ray_tracing/rt_common.py
index 05dcecc..5ee8871 100644
--- a/rayflare/ray_tracing/rt_common.py
+++ b/rayflare/ray_tracing/rt_common.py
@@ -10,7 +10,7 @@
from math import atan2
from random import random
from copy import deepcopy
-from numba import jit
+from numba import jit, njit
from scipy.spatial import Delaunay
unit_cell_N = np.array(
@@ -203,6 +203,16 @@ def calculate_tuv(d, tri_size, tri_crossP, r_a, tri_P_0s, tri_P_2s, tri_P_1s):
return t, u, v
+
+@jit(nopython=True, error_model="numpy")
+def normalize(x):
+ return x / norm(x)
+
+@jit(nopython=True)
+def norm(x):
+ s = x[0]**2 + x[1]**2 + x[2]**2
+ return np.sqrt(s)
+
@jit(nopython=True, error_model="numpy")
def calc_intersection_properties(t, which_intersect, r_a, d, tri_N):
t = t[which_intersect]
@@ -213,34 +223,32 @@ def calc_intersection_properties(t, which_intersect, r_a, d, tri_N):
N = tri_N[which_intersect][ind]
- theta = atan(np.linalg.norm(np.cross(N, -d)) / np.dot(N, -d))
+ Nxd = np.cross(N, -d)
+
+ theta = atan(norm(Nxd) / np.dot(N, -d))
return intersn, theta, N
def check_intersect(r_a, d, tri_size, tri_crossP, tri_P_0s, tri_P_2s, tri_P_1s, tri_N):
- # all the stuff which is only surface-dependent (and not dependent on incoming direction) is
- # in the surface object tri.
t, u, v = calculate_tuv(d, tri_size, tri_crossP, r_a, tri_P_0s, tri_P_2s, tri_P_1s)
+ # for some reason numba doesn't like this bit:
which_intersect = (
(u + v <= 1) & (np.all(np.vstack((u, v)) >= -1e-10, axis=0)) & (t > 0)
- ) # get errors if set lower limit exactly to zero.
+ )
+ # get errors if set lower limit exactly to zero, hence 1e-10
if sum(which_intersect) > 0:
-
- # intersn, theta, N = calc_intersection_properties(t, which_intersect, r_a, d, tri_N)
-
return calc_intersection_properties(t, which_intersect, r_a, d, tri_N)
-
else:
return False
def update_ray_d_pol(ray, rnd, R, R_plus_T, Rs, Rp, Ts, Tp, A_per_layer, n0, n1, N, side):
# TODO: is it necessary to normalize here or will it already be normalized?
- if np.abs(np.linalg.norm(ray.d) - 1) > 1e-2:
- raise ValueError(f"Ray direction not normalized {np.linalg.norm(ray.d)}")
+ if np.abs(norm(ray.d) - 1) > 1e-2:
+ raise ValueError(f"Ray direction not normalized {norm(ray.d)}")
if rnd <= R: # REFLECTION
ray.d = np.real(ray.d - 2 * np.dot(ray.d, N) * N)
@@ -251,7 +259,7 @@ def update_ray_d_pol(ray, rnd, R, R_plus_T, Rs, Rp, Ts, Tp, A_per_layer, n0, n1,
# transmission, refraction
# tr_par = (np.real(n0) / np.real(n1)) * (d - np.dot(d, N) * N)
tr_par = (n0 / n1) * (ray.d - np.dot(ray.d, N) * N)
- tr_perp = -sqrt(1 - np.linalg.norm(tr_par) ** 2) * N
+ tr_perp = -sqrt(1 - norm(tr_par) ** 2) * N
side = -side
ray.d = np.real(tr_par + tr_perp)
@@ -262,14 +270,14 @@ def update_ray_d_pol(ray, rnd, R, R_plus_T, Rs, Rp, Ts, Tp, A_per_layer, n0, n1,
# absorption
A = A_per_layer
- ray.d = ray.d / np.linalg.norm(ray.d)
+ ray.d = normalize(ray.d)
ray.pol = ray.pol / np.sum(
ray.pol) # must always sum to 1, these are weights not intensities
return side, A
def decide_RT_Fresnel(ray, n0, n1, theta, N, side, rnd,
- lookuptable=None):
+ lookuptable=None, d_theta=None):
ratio = np.clip(np.real(n1) / np.real(n0), -1, 1)
@@ -286,18 +294,22 @@ def decide_RT_Fresnel(ray, n0, n1, theta, N, side, rnd,
return side, None # never absorbed, A = False
+@jit(nopython=True)
+def get_data(theta, d_theta, R_data, T_data, Alayer_data, pol):
+ angle_ind = round(abs(theta)/d_theta)
+ [Rs, Rp] = R_data[:, angle_ind]*pol
+ [Ts, Tp] = T_data[:, angle_ind]*pol
+ A_per_layer = np.sum(Alayer_data[:, angle_ind].T * pol, 1)
-def decide_RT_TMM(ray, n0, n1, theta, N, side, rnd, lookuptable):
+ return Rs, Rp, Ts, Tp, A_per_layer
- data_sp = lookuptable.loc[dict(side=side)].sel(
- angle=abs(theta), method="nearest",
- ) # why do I have to do this in two steps?
+def decide_RT_TMM(ray, n0, n1, theta, N, side, rnd, lookuptable, d_theta):
- # multiply along the pol dimension by ray.pol:
- [Rs, Rp] = data_sp.R.data*ray.pol
- [Ts, Tp] = data_sp.T.data*ray.pol
- A_per_layer = np.sum(data_sp.Alayer.data.T * ray.pol, 1)
+ data_s = lookuptable.loc[dict(side=side)]
+ Rs, Rp, Ts, Tp, A_per_layer = get_data(theta, d_theta,
+ data_s.R.data, data_s.T.data, data_s.Alayer.data,
+ ray.pol)
R = Rs + Rp
T = Ts + Tp
R_plus_T = R + T
@@ -358,6 +370,7 @@ def single_cell_check(
Ly,
side,
z_cov,
+ d_theta,
n_interactions=0,
wl=None,
Fr_or_TMM=0,
@@ -386,7 +399,7 @@ def single_cell_check(
n_misses += 1
- o_t = np.real(acos(ray.d[2] / np.linalg.norm(ray.d)))
+ o_t = np.real(acos(ray.d[2] / norm(ray.d)))
o_p = np.real(atan2(ray.d[1], ray.d[0]))
if np.sign(d0[2]) == np.sign(ray.d[2]):
@@ -431,7 +444,7 @@ def single_cell_check(
rnd = random()
side, A = decide[Fr_or_TMM](
- ray, n0, n1, theta, N, side, rnd, lookuptable
+ ray, n0, n1, theta, N, side, rnd, lookuptable, d_theta
)
ray.r_a = np.real(
@@ -476,6 +489,7 @@ def single_interface_check(
Ly,
side,
z_cov,
+ d_theta,
n_interactions=0,
wl=None,
Fr_or_TMM=0,
@@ -540,7 +554,7 @@ def single_interface_check(
# assume it is reflected back into layer it came from
ray.d[2] = -ray.d[2]
- o_t = np.real(acos(ray.d[2] / np.linalg.norm(ray.d)))
+ o_t = np.real(acos(ray.d[2] / norm(ray.d)))
o_p = np.real(atan2(ray.d[1], ray.d[0]))
return 0, o_t, o_p, 0, n_interactions, side
@@ -569,7 +583,7 @@ def single_interface_check(
else:
- o_t = np.real(acos(ray.d[2] / np.linalg.norm(ray.d)))
+ o_t = np.real(acos(ray.d[2] / norm(ray.d)))
o_p = np.real(atan2(ray.d[1], ray.d[0]))
if np.sign(d0[2]) == np.sign(ray.d[2]):
@@ -622,7 +636,7 @@ def single_interface_check(
rnd = random()
side, A = decide[Fr_or_TMM](
- ray, n0, n1, theta, N, side, rnd, lookuptable
+ ray, n0, n1, theta, N, side, rnd, lookuptable, d_theta,
)
ray.r_a = np.real(
diff --git a/rayflare/ray_tracing/rt_structure.py b/rayflare/ray_tracing/rt_structure.py
index b132819..b5d33ec 100644
--- a/rayflare/ray_tracing/rt_structure.py
+++ b/rayflare/ray_tracing/rt_structure.py
@@ -313,6 +313,7 @@ def calculate(self, options):
initial_mat,
initial_dir,
periodic,
+ (np.pi/2)/options.lookuptable_angles,
lambertian_approximation,
tmm_args + [wavelengths[i1]],
prop_rays.isel(wl=i1),
@@ -676,6 +677,7 @@ def parallel_inner(
initial_mat,
initial_dir,
periodic,
+ d_theta,
lambertian_approximation=0,
tmm_args=None,
existing_rays=None,
@@ -876,6 +878,7 @@ def parallel_inner(
depth_indices,
I_thresh,
pols[overall_i],
+ d_theta,
randomize,
i_mats[overall_i],
i_dirs[overall_i],
@@ -1173,6 +1176,7 @@ def single_ray_stack(
depth_indices,
I_thresh,
pol,
+ d_theta,
randomize=False,
mat_i=0,
direction=1,
@@ -1184,8 +1188,6 @@ def single_ray_stack(
I_in=1,
):
- if pol is None:
- pol = [0.5, 0.5]
single_surface = {0: single_cell_check, 1: single_interface_check}
# use single_cell_check if not periodic, single_interface_check if is periodic
@@ -1267,11 +1269,11 @@ def single_ray_stack(
surf.Ly,
direction,
surf.zcov,
+ d_theta,
n_interactions,
**tmm_kwargs_list[surf_index]
)
-
if res == 0: # reflection
direction = -direction # changing direction due to reflection
diff --git a/tests/test_analytical_rt.py b/tests/test_analytical_rt.py
index df6bebf..d86ca8e 100644
--- a/tests/test_analytical_rt.py
+++ b/tests/test_analytical_rt.py
@@ -56,6 +56,8 @@ def test_total_RAT_TMM():
from solcore import material
from rayflare.options import default_options
from solcore.structure import Layer
+ import numpy as np
+ from pytest import approx
Si = material("Si")()
Air = material("Air")()
@@ -64,19 +66,20 @@ def test_total_RAT_TMM():
Ge = material("Ge")()
options = default_options()
- options.wavelength = np.linspace(300, 1400, 10)
+ options.wavelength = np.linspace(300, 1400, 5) * 1e-9
options.nx = 10
options.ny = 10
options.n_rays = 1 * options.nx ** 2
options.pol = 's'
options.analytical_ray_tracing = 2
options.project_name = 'test_analytical'
- options.theta_in = 7*np.pi/180
+ options.theta_in = 7 * np.pi / 180
+ options.parallel = False
interface_layers = [Layer(70e-9, MgF2), Layer(500e-9, GaAs)]
pyramids = regular_pyramids(50, True, interface_layers=interface_layers)
- planar = planar_surface(interface_layers=[Layer(200e-9, Ge)])
+ planar = planar_surface()
planar_2 = planar_surface()
rtstr = rt_structure(