some functions

rainbow.py

@@ -1,6 +1,9 @@
 import numpy as np
 from raytrace_2D import Disk
 import csv
+import colorspace
+from rich.progress import track
+from PIL import Image
 
 class Ray:
 	def __init__(self, origin, direction, intensity):
@@ -12,7 +15,7 @@ center = [0,0]
 r = 1
 n = 1.3
 disk = Disk(center, r, n)
-dx = 0.005
+dx = 0.01
 N = int(2*r /dx - 1)
 min_intensity = 0.001
 max_ray = 10000
@@ -22,9 +25,10 @@ result = []
 points = []
 
 def init():
+	stack = []
 	for x in np.linspace(-r + dx, r-dx, N):
 		stack.append(Ray([x, 2*r], [0,-1], 10*np.abs(x)))
-	print(np.linspace(-r + dx, r-dx, N))
+	return stack
 
 def reflection_and_refraction(ray:Ray, intersection_point, normal, n):
 	# print("reflection/refraction at the point:", intersection_point)
@@ -49,6 +53,7 @@ def reflection_and_refraction(ray:Ray, intersection_point, normal, n):
 
 def trace(disk:Disk, stack:list, max_ray:int):
 	ray_count = 0
+	result = []
 	while stack and ray_count < max_ray:
 		ray = stack.pop()
 		ray_count += 1
@@ -67,7 +72,8 @@ def trace(disk:Disk, stack:list, max_ray:int):
 				points.append(np.concatenate((ray.origin, ray.origin+ray.direction,[ray.intensity])))
 				if direction[1] > 0:
 					result.append([np.arccos(direction[1]), ray.intensity*direction[1]])
-	print("ray trace finish, total ray count:",ray_count)
+	# print("ray trace finish, total ray count:",ray_count)
+	return result
 
 def water_density(t):
 	data_t = np.linspace(0,101,102)
@@ -76,7 +82,7 @@ def water_density(t):
 
 def water_refraction_index(t, wavelength):
 	rhobar = 0.999974950 * (1 - (t - 3.983035)*(t - 3.983035)*(t + 301.797)/(522528.9*(t+69.34881)))
-	print(rhobar)
+	# print(rhobar)
 	a0 = 0.244257733
 	a1 = 0.00974634476
 	a2 = -0.00373234996
@@ -96,4 +102,68 @@ def water_refraction_index(t, wavelength):
 	n = np.sqrt((2*C+1)/(1-C))
 	return n
 
+def sun_spectral(wavelength):
+	return 1e16/(np.power(wavelength,5)*(np.exp(6.62607015e6*2.99792458/(wavelength*1.380649*(5250+273.15)))-1))
 
+def rainbow(n_theta, temp):
+	angles, d_theta = np.linspace(0,np.pi/2,n_theta, retstep=True)
+	mid_angles = angles[:-1] + d_theta
+
+	angle_spectral = np.zeros((n_theta, len(colorspace.data_wavelength)))
+	angle_XYZ = np.zeros((n_theta, 3))
+	angle_sRGB = np.zeros((n_theta, 3))
+
+	num_wavelength = len(colorspace.data_wavelength)
+	for wavelength_index in track(range(num_wavelength), description="simulating..."):
+		wavelength = colorspace.data_wavelength[wavelength_index]
+		n = water_refraction_index(temp, wavelength)
+		disk = Disk(center, r, n)
+		stack = []
+		for x in np.linspace(-r + dx, r-dx, N):
+			stack.append(Ray([x, 2*r], [0,-1], 10*np.abs(x)))
+		result = np.array(trace(disk, stack, max_ray))
+		result[:,1] *= sun_spectral(wavelength)
+		angle_indexes = np.floor(result[:,0]/d_theta)
+		for i,angle_index in enumerate(angle_indexes):
+			angle_spectral[int(angle_index), wavelength_index] += result[i,1]
+
+	for i,spectral in enumerate(angle_spectral):
+		angle_XYZ[i] = colorspace.spectral2XYZ(spectral)
+		angle_sRGB[i] = colorspace.XYZ2RGB(angle_XYZ[i])
+
+	maxRGB = np.max(angle_sRGB)
+	angle_sRGB/= maxRGB
+	colorspace.vectorized_gamma_correct(angle_sRGB)
+	np.clip(angle_sRGB, 0, 1)
+	return angles,angle_sRGB
+
+def bin_find(x, list):
+	l,r = 0, len(list)-1
+	if x < list[l]:
+		return 0
+	if x >= list[r]:
+		return r
+	while l <= r:
+		mid = l + (r-l)//2
+		if list[mid] <= x < list[mid + 1]:
+			return mid
+		elif list[mid] < x:
+			l = mid + 1
+		else:
+			r = mid - 1
+	return -1
+
+def take_picture(angles, angle_sRGB, w, h, distance, filename="image.png"):
+	image_array = np.zeros((w, h, 3), dtype=np.uint8)
+	radius = distance * np.tan(angles)
+	for i in range(w):
+		for j in range(h):
+			r = np.linalg.norm(np.array([i,j]) - [w/2, 0])
+			index = bin_find(r, radius)
+			image_array[i,j] = angle_sRGB[index]*255
+	image = Image.fromarray(image_array, "RGB")
+	image.save(filename)
+
+if __name__=="__main__":
+	angles,sRGB=rainbow(10, 10)
+	take_picture(angles,sRGB, 1920, 1080, 500, "image.png")