init

rainbow.py

@@ -0,0 +1,78 @@
+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.01
+N = int(2*r /dx - 1)
+min_intensity = 0.001
+
+stack = []
+result = []
+points = []
+
+def init():
+	for x in np.linspace(-r + dx, r-dx, N):
+		stack.append(Ray([x, 2*r], [0,-1], 1))
+	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):
+	while stack:
+		ray = stack.pop()
+		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]])
+
+# init()
+stack.append(Ray([0.8, 2*r], [0,-1], 1))
+trace(disk, stack)
+
+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)