update: save ram by trace ray one by one

README.md

@@ -0,0 +1,28 @@
+# 几何光学模拟及彩虹模拟
+
+写了几何光学正向光追。考虑了阳光在水珠中的折射反射，统计角度分布从而模拟彩虹
+
+## 使用的数据
+
+- 水的密度：$$\rho(t) = 999.974950 \frac{1 - (t - 3.983035)^2 (t + 301.797)}{522528.9(t+69.34881)}$$
+其中 $t$ 是摄氏温度
+- 水的折射率：
+  $$ n = \sqrt{\frac{2C + 1}{1- C}}, \\ C = \bar{\rho} \left( a_0 + a_1 \bar{\rho} + a_2 \bar{T} + a_{3}{\bar{\lambda }}^{2}{\bar{T}}+{\frac {a_{4}}{{\bar{\lambda }}^{2}}}+{\frac {a_{5}}{{\bar{\lambda }}^{2}-{\bar{\lambda }}_{\mathit {UV}}^{2}}}+{\frac {a_{6}}{{\bar{\lambda }}^{2}-{\bar{\lambda }}_{\mathit {IR}}^{2}}}+a_{7}{\bar{\rho }}^{2} \right)$$
+  其中：$\bar{T} = T/T^*$, $\bar{\rho} = \rho/\rho^*$, $\bar{\lambda} = \lambda/\lambda^*$ 是约化量，$a_{0} = 0.244257733$, $a_{1} = 0.00974634476$, $a_{2} = −0.00373234996$, $a_{3} = 0.000268678472$, $a_4 = 0.0015892057$, $a_{5} = 0.00245934259$, $a_{6} = 0.90070492$, $a_{7} = −0.0166626219$, $T^{*} = 273.15 \ \mathrm{K}$, $\rho^{*} = 1000\ \mathrm{kg/m^3}$, $\lambda^{*} = 589\ \mathrm{nm}$, $\bar\lambda_{\text{IR}} = 5.432937$, $\bar\lambda_{\text{UV}} = 0.229202$。
+- 单色光引起的色觉：见 CIE 1931，详细数据在 `colorspace.py` 中的 `_CIEXYZ_1931_table`。
+- 阳光设为 $5250\ {}\degree\rm C$ 的黑体辐射。这与大气上层吻合较好，但与大气底层相比，忽略了水分子的大量吸收峰和氧分子、二氧化碳分子等的吸收峰。
+
+## 模拟过程
+
+- 给定温度
+- 对每隔 $1\ \rm{nm}$ 的单色光：
+  - 计算折射率，进行正向光追。假设光只与单个水珠相遇。入射光的瞄准距离 $d_i = r \sqrt{u_i}$，其中 $u$ 在 $[0,1)$ 中均匀分布；入射光按照黑体辐射设置
+  - 统计背向出射的光强 - 角度分布
+  - 转化为 XYZ - 角度分布
+- 求和，得到整个频谱的 XYZ - 角度分布
+- 转化为 sRGB - 角度分布，作图
+
+### TODO：
+- 考虑散射效应：
+  - 这将导致背景不是黑色，而是天空蓝
+  - 这将导致水珠反射回的光强被削弱

colorspace.py

@@ -499,11 +499,11 @@ def gamma_correct(c):
 def vectorized_gamma_correct(array):
 	return np.where(array <= 0.0031308, 12.92*array, 1.055*np.power(array, 1/2.4)-0.055)
 
-def wavelength2XYZ(wavelength, intensity):
+def wavelength2XYZ(wavelength):
 	X = np.interp(wavelength, data_wavelength, data_x, 0, 0)
 	Y = np.interp(wavelength, data_wavelength, data_y, 0, 0)
 	Z = np.interp(wavelength, data_wavelength, data_z, 0, 0)
-	return intensity*np.array([X,Y,Z])
+	return np.array([X,Y,Z])
 
 def spectral2XYZ(spectral):
 	X = np.dot(spectral, data_x)

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,22 +14,21 @@ class Ray:
 
 center = [0,0]
 r = 1
-n = 1.3
-disk = Disk(center, r, n)
-dx = 0.00001
-N = int(2*r /dx - 1)
+# n = 1.3
+# disk = Disk(center, r, n)
+N = 100000
 min_intensity = 0.00001
-max_ray = 1000000
+max_ray = 100
 
 stack = []
 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)))
-	return stack
+# 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)))
+# 	return stack
 
 def reflection_and_refraction(ray:Ray, intersection_point, normal, n):
 	# print("reflection/refraction at the point:", intersection_point)
@@ -105,43 +105,82 @@ 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, N, n_theta, d_theta, min_intensity, max_ray, max_angle):
+	n = water_refraction_index(temp, wavelength)
+	disk = Disk(center, r, n)
+	stack = []
+	XYZ = colorspace.wavelength2XYZ(wavelength)*sun_spectral(wavelength)
+	own_angle_XYZ = np.zeros((n_theta, 3))
+	for u in np.linspace(0, 1, N, endpoint=False):
+		stack=[Ray([r*np.sqrt(u), 2*r], [0,-1], 1)]
+		ray_count = 0
+		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_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])
+	with Progress() as progress:
+		# 创建一个进度条
+		task = progress.add_task("[green]Simulating...", total=num_wavelength)
+
+		# 使用 ThreadPoolExecutor 并行执行
+		with ProcessPoolExecutor() as executor:
+			futures = [executor.submit(modified_trace, wavelength, temp, center, r, 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)
 
 	maxRGB = np.max(angle_sRGB)
 	angle_sRGB/= maxRGB
+	angle_sRGB = np.clip(angle_sRGB, 0, 1)
 	colorspace.vectorized_gamma_correct(angle_sRGB)
-	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]:
@@ -152,18 +191,42 @@ 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)
+			ratio = (r - radius[index])/(radius[index+1] - radius[index])
+			column[j] = ((1-ratio)*angle_sRGB[index]+ratio*angle_sRGB[index+1])*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 *2
+	h = 4320 *2
+	dis = 2400 *2
+	max_angle = 1.1*np.arctan(np.sqrt(1+(h*h+w*w/4)/(dis*dis)))
+	angles,sRGB=rainbow(10000, max_angle, 10)
 	np.savez_compressed("saved.npz", a=angles, b=sRGB)
-	take_picture(angles,sRGB, 7680, 4320, 2400, "image.png")
+	# loaded = np.load("saved.npz")
+	# angles, sRGB = loaded['a'], loaded['b']
+	take_picture(angles,sRGB, w, h, dis, "image16k.png")