finish!

rainbow.py

@@ -2,8 +2,9 @@ import numpy as np
 from raytrace_2D import Disk
 import csv
 import colorspace
-from rich.progress import track
+from rich.progress import Progress, track
 from PIL import Image
+from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
 
 class Ray:
 	def __init__(self, origin, direction, intensity):
@@ -13,8 +14,8 @@ class Ray:
 
 center = [0,0]
 r = 1
-n = 1.3
-disk = Disk(center, r, n)
+# n = 1.3
+# disk = Disk(center, r, n)
 dx = 0.00001
 N = int(2*r /dx - 1)
 min_intensity = 0.00001
@@ -105,44 +106,57 @@ def water_refraction_index(t, wavelength):
 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)
+def modified_trace(wavelength, temp, center, r, dx, N, n_theta, d_theta, min_intensity, max_ray, max_angle):
+	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)))
+	ray_count = 0
+	XYZ = colorspace.wavelength2XYZ(wavelength)*sun_spectral(wavelength)
+	own_angle_XYZ = np.zeros((n_theta, 3))
+	while stack and ray_count < max_ray:
+		ray = stack.pop()
+		ray_count += 1
+		if ray is None:
+			continue
+		if isinstance(ray, Ray):
+			if ray.intensity < min_intensity:
+				continue
+			direction = ray.direction
+			t,intersection_point = disk.find_intersection(ray)
+			if intersection_point is not None:
+				points.append(np.concatenate((ray.origin, intersection_point,[ray.intensity])))
+				normal = disk.get_normal(intersection_point)
+				stack.extend(reflection_and_refraction(ray, intersection_point, normal, disk.refractive_index))
+			else:
+				points.append(np.concatenate((ray.origin, ray.origin+ray.direction,[ray.intensity])))
+				if direction[1] > 0:
+					angle = np.arccos(direction[1])
+					if angle < max_angle:
+						own_angle_XYZ[int(angle/d_theta)] += XYZ*ray.intensity*direction[1]
+	return own_angle_XYZ
+
+def rainbow(n_theta, max_theta, temp):
+	angles, d_theta = np.linspace(0,max_theta,n_theta, retstep=True)
 
 	angle_XYZ = np.zeros((n_theta, 3))
 	angle_sRGB = np.zeros((n_theta, 3))
 
 	num_wavelength = len(colorspace.data_wavelength)
 
-	def modified_trace(wavelength_index):
-		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)))
-		ray_count = 0
-		XYZ = colorspace.wavelength2XYZ(wavelength)*sun_spectral(wavelength)
-		while stack and ray_count < max_ray:
-			ray = stack.pop()
-			ray_count += 1
-			if ray is None:
-				continue
-			if isinstance(ray, Ray):
-				if ray.intensity < min_intensity:
-					continue
-				direction = ray.direction
-				t,intersection_point = disk.find_intersection(ray)
-				if intersection_point is not None:
-					points.append(np.concatenate((ray.origin, intersection_point,[ray.intensity])))
-					normal = disk.get_normal(intersection_point)
-					stack.extend(reflection_and_refraction(ray, intersection_point, normal, disk.refractive_index))
-				else:
-					points.append(np.concatenate((ray.origin, ray.origin+ray.direction,[ray.intensity])))
-					if direction[1] > 0:
-						angle = np.arccos(direction[1])
-						angle_XYZ[int(angle/d_theta)] += XYZ*ray.intensity*direction[1]
+	with Progress() as progress:
+		# 创建一个进度条
+		task = progress.add_task("[green]Simulating...", total=num_wavelength)
 
-	map(modified_trace, range(num_wavelength))
+		# 使用 ThreadPoolExecutor 并行执行
+		with ProcessPoolExecutor() as executor:
+			futures = [executor.submit(modified_trace, wavelength, temp, center, r, dx, N, n_theta, d_theta, min_intensity, max_ray, max_theta) for wavelength in colorspace.data_wavelength]
+
+			for future in as_completed(futures):
+				result = future.result()
+				angle_XYZ += result
+				progress.update(task, advance=1)
 
 	for i,XYZ in enumerate(angle_XYZ):
 		angle_sRGB[i] = colorspace.XYZ2RGB(XYZ)
@@ -153,12 +167,21 @@ def rainbow(n_theta, temp):
 	np.clip(angle_sRGB, 0, 1)
 	return angles,angle_sRGB
 
-def bin_find(x, list):
+def bin_find(x, list, start):
 	l,r = 0, len(list)-1
+
 	if x < list[l]:
 		return 0
 	if x >= list[r]:
 		return r
+
+	if list[start] <= x < list[start + 1]:
+		return start
+	elif list[start] < x:
+		l = start
+	else:
+		r = start
+
 	while l <= r:
 		mid = l + (r-l)//2
 		if list[mid] <= x < list[mid + 1]:
@@ -169,18 +192,39 @@ def bin_find(x, list):
 			r = mid - 1
 	return -1
 
+def draw_column(i, w, h, radius, angle_sRGB):
+		index = 0
+		column = np.zeros((h, 3), dtype=np.int8)
+		for j in range(h):
+			r = np.linalg.norm(np.array([i,j]) - [w/2, 0])
+			index = bin_find(r, radius, index)
+			column[j] = angle_sRGB[index]*255
+		return i,column
+
 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 track(range(w), description="generating picture..."):
-		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
+
+	with Progress() as progress:
+		# 创建一个进度条
+		task2 = progress.add_task("[green]Rendering...", total=w)
+
+		# 使用 ThreadPoolExecutor 并行执行
+		with ProcessPoolExecutor() as executor:
+			futures2 = [executor.submit(draw_column, i, w,h,radius,angle_sRGB) for i in range(w)]
+
+			for future in as_completed(futures2):
+				ind,column = future.result()
+				image_array[ind] = column
+				progress.update(task2, advance=1)
 	image = Image.fromarray(image_array, "RGB")
 	image.save(filename)
 
 if __name__=="__main__":
-	angles,sRGB=rainbow(10000, 10)
+	w = 7680
+	h = 4320
+	dis = 2400
+	max_angle = 1.1*np.arctan(np.sqrt(1+(h*h+w*w/4)/(dis*dis)))
+	angles,sRGB=rainbow(100, max_angle, 10)
 	np.savez_compressed("saved.npz", a=angles, b=sRGB)
 	take_picture(angles,sRGB, 7680, 4320, 2400, "image.png")