From 0b0c8624c58f2a7b118bc3b2a4c75bebca8de8d5 Mon Sep 17 00:00:00 2001 From: GaugeAndGravity Date: Fri, 10 Nov 2023 02:01:47 -0500 Subject: [PATCH] some functions --- .gitignore | 1 + colorspace.py | 18 ++++++++++-- rainbow.py | 78 ++++++++++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 90 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 5d709e8..b5860c2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ *.csv __pycache__/ +*.png diff --git a/colorspace.py b/colorspace.py index 88dd64b..ed5c1c1 100644 --- a/colorspace.py +++ b/colorspace.py @@ -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]) diff --git a/rainbow.py b/rainbow.py index 4143a10..d2697e1 100644 --- a/rainbow.py +++ b/rainbow.py @@ -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")