wyj/rainbow
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

performance

Browse Source
This commit is contained in:
Yingjie Wang 2023-11-10 09:53:12 -05:00
parent 1d5e61c30b
commit a98e40d16b
1 changed files with 27 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
rainbow.py
View File

@ -108,27 +108,44 @@ def sun_spectral(wavelength):
def rainbow(n_theta, temp):
angles, d_theta = np.linspace(0,np.pi/2,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..."):
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)))
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]
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]
for i,spectral in enumerate(angle_spectral):
angle_XYZ[i] = colorspace.spectral2XYZ(spectral)
angle_sRGB[i] = colorspace.XYZ2RGB(angle_XYZ[i])
map(modified_trace, range(num_wavelength))
for i,XYZ in enumerate(angle_XYZ):
angle_sRGB[i] = colorspace.XYZ2RGB(XYZ)
maxRGB = np.max(angle_sRGB)
angle_sRGB/= maxRGB