170 lines
5.9 KiB
Python
170 lines
5.9 KiB
Python
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):
|
|
self.origin = np.array(origin)
|
|
self.direction = np.array(direction) / np.linalg.norm(direction)
|
|
self.intensity = intensity
|
|
|
|
center = [0,0]
|
|
r = 1
|
|
n = 1.3
|
|
disk = Disk(center, r, n)
|
|
dx = 0.00001
|
|
N = int(2*r /dx - 1)
|
|
min_intensity = 0.00001
|
|
max_ray = 1000000
|
|
|
|
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 reflection_and_refraction(ray:Ray, intersection_point, normal, n):
|
|
# print("reflection/refraction at the point:", intersection_point)
|
|
k_n = np.dot(ray.direction, normal)
|
|
if k_n > 0:
|
|
n1,n2 = n,1
|
|
else:
|
|
n1, n2 = 1,n
|
|
kp = ray.direction - k_n*normal
|
|
k2p = n1*kp/n2
|
|
if np.dot(k2p, k2p) >= 1:
|
|
ray1 = Ray(intersection_point, ray.direction-2*k_n*normal, ray.intensity)
|
|
return ray1,None
|
|
else:
|
|
k2n = np.sqrt(1 - np.dot(k2p, k2p))*np.sign(k_n)
|
|
Rs = (n1*k_n - n2*k2n)*(n1*k_n - n2*k2n)/((n1*k_n + n2*k2n)*(n1*k_n + n2*k2n))
|
|
Rp = (n1*k2n - n2*k_n)*(n1*k2n - n2*k_n)/((n1*k2n + n2*k_n)*(n1*k2n + n2*k_n))
|
|
R = (Rs+Rp)/2
|
|
ray1 = Ray(intersection_point, ray.direction-2*k_n*normal, ray.intensity*R)
|
|
ray2 = Ray(intersection_point, k2n*normal+k2p, ray.intensity*(1-R))
|
|
return ray1,ray2
|
|
|
|
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
|
|
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:
|
|
result.append([np.arccos(direction[1]), ray.intensity*direction[1]])
|
|
# print("ray trace finish, total ray count:",ray_count)
|
|
return result
|
|
|
|
def water_density(t):
|
|
data_t = np.linspace(0,101,102)
|
|
data_rho = np.array([999.79,999.84,999.88,999.92,999.97,999.92,999.88,999.84,999.79,999.72,999.65,999.55,999.44,999.32,999.19,999.05,998.90,998.74,998.56,998.36,998.16,997.96,997.74,997.50,997.25,996.99,996.74,996.48,996.20,995.90,995.59,995.29,995.97,994.65,994.33,994.98,993.64,993.28,992.93,992.55,992.17,991.78,991.39,990.99,990.57,990.16,989.75,989.33,988.90,988.45,988.04,987.55,987.06,986.58,986.19,985.70,985.22,984.73,984.67,983.18,983.13,982.70,982.22,981.64,981.06,980.48,979.91,979.33,978.85,978.28,977.70,977.13,976.56,975.99,975.41,974.84,974.27,973.70,973.14,972.47,971.81,971.25,970.59,969.93,969.27,968.61,967.96,967.30,966.65,965.99,965.34,964.69,964.04,963.39,962.64,961.90,961.26,960.52,959.78,959.04,958.40,957.67])
|
|
return np.interp(t, data_t, data_rho)
|
|
|
|
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)
|
|
a0 = 0.244257733
|
|
a1 = 0.00974634476
|
|
a2 = -0.00373234996
|
|
a3 = 0.000268678472
|
|
a4 = 0.0015892057
|
|
a5 = 0.00245934259
|
|
a6 = 0.90070492
|
|
a7 = -0.0166626219
|
|
Tstar = 273.15
|
|
Tbar = 1 + t/Tstar
|
|
lambdastar = 589
|
|
lambdabar = wavelength/lambdastar
|
|
lambda_IR = 5.432937
|
|
lambda_UV = 0.229202
|
|
C = a0 + a1*rhobar + a2*Tbar + a3*lambdabar*lambdabar*Tbar + a4/(lambdabar*lambdabar) + a5/((lambdabar+lambda_UV)*(lambdabar-lambda_UV)) + a6/((lambdabar + lambda_IR)*(lambdabar - lambda_IR)) + a7*rhobar*rhobar
|
|
C *= rhobar
|
|
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)
|
|
|
|
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 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
|
|
image = Image.fromarray(image_array, "RGB")
|
|
image.save(filename)
|
|
|
|
if __name__=="__main__":
|
|
angles,sRGB=rainbow(10000, 10)
|
|
np.savez_compressed("saved.npz", a=angles, b=sRGB)
|
|
take_picture(angles,sRGB, 3840, 2160, 1200, "image.png")
|