diff --git a/image.png b/image.png new file mode 100644 index 0000000..14160fc Binary files /dev/null and b/image.png differ diff --git a/rainbow.py b/rainbow.py index 973776c..7806115 100644 --- a/rainbow.py +++ b/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") diff --git a/test.png b/test.png new file mode 100644 index 0000000..02ab425 Binary files /dev/null and b/test.png differ