import numpy as np from abc import ABC, abstractmethod class Ray: def __init__(self, origin, direction): self.origin = np.array(origin) self.direction = np.array(direction) / np.linalg.norm(direction) class Shapes(ABC): def __init__(self,refractive_index) -> None: self.refractive_index = refractive_index @abstractmethod def find_intersection(self, ray:Ray): pass def get_normal(self, point): pass class Scheme: def __init__(self, shapes, rays) -> None: self.shapes = shapes self.rays = rays self.raytrace() def raytrace(self): for ray in self.rays: t, intersection_point, first_shape = None, None, None for shape in self.shapes: t1,intersection_point1 = shape.find_intersection(ray) if t is None or ((t1 is not None) and (t1 < t)): t, intersection_point, first_shape = t1, intersection_point1, shape if t is None: continue else: self.reflection_and_refraction(ray, first_shape, intersection_point) def reflection_and_refraction(ray:Ray, shape:Shapes, intersection_point): pass class Disk(Shapes): def __init__(self, center, radius, refractive_index): super().__init__(refractive_index) self.center = np.array(center) self.radius = radius def find_intersection(self, ray:Ray): oc = ray.origin - self.center a = np.dot(ray.direction, ray.direction) b = 2.0 * np.dot(oc, ray.direction) c = np.dot(oc, oc) - self.radius * self.radius discriminant = b*b - 4*a*c if discriminant < 0: return None, None # 没有交点 else: dis2 = np.dot(oc,oc) epsilon = 1e-12 if dis2 > 1 - epsilon and dis2 < 1 + epsilon: normal = self.get_normal(ray.origin) if np.dot(normal, ray.direction) < 0: t = (-b + np.sqrt(discriminant)) / (2.0 * a) else: return None, None elif dis2 <= 1 - epsilon: t = (-b + np.sqrt(discriminant)) / (2.0 * a) else: t = (-b - np.sqrt(discriminant)) / (2.0 * a) if t < 0: return None, None intersection_point = ray.origin + t * ray.direction intersection_point = self.center + self.radius * (intersection_point - self.center)/np.linalg.norm(intersection_point - self.center) return t,intersection_point def get_normal(self, point): normal = (point - self.center)/self.radius normal = normal / np.linalg.norm(normal) return normal class half_plane(Shapes): def __init__(self, point, normal, refractive_index) -> None: super().__init__(refractive_index) self.point = point self.normal = normal/np.linalg.norm(normal) def find_intersection(self, ray:Ray): dot = np.dot(ray.direction, self.normal) if dot == 0: return None,None else: t = np.dot(self.point - ray.origin, self.normal) / dot return t, ray.origin + t*ray.direction def get_normal(self, point): return self.normal