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

some functions

Browse Source
This commit is contained in:
Yingjie Wang 2023-11-10 02:01:47 -05:00
parent 0d92049889
commit 0b0c8624c5
3 changed files with 90 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -1,2 +1,3 @@
*.csv
__pycache__/
*.png

18
colorspace.py
View File

@ -1,6 +1,8 @@
import numpy as np
XYZ2RGB_matrix = np.array([[0.49000,0.17697, 0],[0.31000, 0.81240, 0.01000],[0.20000, 0.01063, 0.99000]])
XYZ2RGB_matrix = np.array([[3.2404542, -1.5371385, -0.4985314],
[-0.9692660, 1.8760108, 0.0415560],
[ 0.0556434, -0.2040259, 1.0572252]])
RGB2XYZ_matrix = np.linalg.inv(XYZ2RGB_matrix)
_CIEXYZ_1931_table = [
@ -482,12 +484,21 @@ data_x = np.array(_CIEXYZ_1931_table)[:,1]
data_y = np.array(_CIEXYZ_1931_table)[:,2]
data_z = np.array(_CIEXYZ_1931_table)[:,3]
def XYZ2RGB(XYZ:list)->list:
def XYZ2RGB(XYZ):
return np.dot(XYZ2RGB_matrix, XYZ)
def RGB2XYZ(RGB:list)->list:
def RGB2XYZ(RGB):
return np.dot(RGB2XYZ_matrix, RGB)
def gamma_correct(c):
if c <= 0.0031308:
return 12.92 * c
else:
return 1.055 * (c ** (1 / 2.4)) - 0.055
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):
X = np.interp(wavelength, data_wavelength, data_x, 0, 0)
Y = np.interp(wavelength, data_wavelength, data_y, 0, 0)
@ -498,3 +509,4 @@ def spectral2XYZ(spectral):
X = np.dot(spectral, data_x)
Y = np.dot(spectral, data_y)
Z = np.dot(spectral, data_z)
return np.array([X,Y,Z])

78
rainbow.py
View File

@ -1,6 +1,9 @@
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):
@ -12,7 +15,7 @@ center = [0,0]
r = 1
n = 1.3
disk = Disk(center, r, n)
dx = 0.005
dx = 0.01
N = int(2*r /dx - 1)
min_intensity = 0.001
max_ray = 10000
@ -22,9 +25,10 @@ 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)))
print(np.linspace(-r + dx, r-dx, N))
return stack
def reflection_and_refraction(ray:Ray, intersection_point, normal, n):
# print("reflection/refraction at the point:", intersection_point)
@ -49,6 +53,7 @@ def reflection_and_refraction(ray:Ray, intersection_point, normal, n):
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
@ -67,7 +72,8 @@ def trace(disk:Disk, stack:list, max_ray:int):
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)
# print("ray trace finish, total ray count:",ray_count)
return result
def water_density(t):
data_t = np.linspace(0,101,102)
@ -76,7 +82,7 @@ def water_density(t):
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)
# print(rhobar)
a0 = 0.244257733
a1 = 0.00974634476
a2 = -0.00373234996
@ -96,4 +102,68 @@ def water_refraction_index(t, wavelength):
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)
mid_angles = angles[:-1] + d_theta
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 range(w):
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(10, 10)
take_picture(angles,sRGB, 1920, 1080, 500, "image.png")