83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
import numpy as np
|
|
from raytrace_2D import Disk
|
|
import csv
|
|
|
|
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.005
|
|
N = int(2*r /dx - 1)
|
|
min_intensity = 0.001
|
|
max_ray = 10000
|
|
|
|
stack = []
|
|
result = []
|
|
points = []
|
|
|
|
def init():
|
|
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))
|
|
|
|
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
|
|
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)
|
|
|
|
init()
|
|
# stack.append(Ray([0.8, 2*r], [0,-1], 1))
|
|
trace(disk, stack, max_ray)
|
|
|
|
with open("result.csv", 'w', newline='') as f:
|
|
writer = csv.writer(f)
|
|
writer.writerows(result)
|
|
|
|
with open("points.csv", 'w', newline='') as f:
|
|
writer = csv.writer(f)
|
|
writer.writerows(points)
|