72 lines
1.5 KiB
Python
72 lines
1.5 KiB
Python
#!/usr/bin/env python3
|
|
|
|
from math import sqrt, pi
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
class Ellipse:
|
|
def __init__(self, a, b):
|
|
self.a = float(a)
|
|
self.b = float(b)
|
|
|
|
def normal(self, vec):
|
|
x, y = vec
|
|
return norm([self.b**2 * x, self.a**2 * y])
|
|
|
|
def quadratic(a, b, c):
|
|
D = b**2 - 4*a*c
|
|
if D < 0:
|
|
return False
|
|
r1 = (-b + sqrt(D)) / (2*a)
|
|
r2 = (-b - sqrt(D)) / (2*a)
|
|
return r1, r2
|
|
|
|
def norm(v):
|
|
l = sqrt(sum(x**2 for x in v))
|
|
return np.array([x/l for x in v])
|
|
|
|
def intersections(ellipse, point, dir):
|
|
m = dir[1] / dir[0]
|
|
yint = point[1] - point[0]*m
|
|
|
|
a = ellipse.a**2 * m**2 + ellipse.b**2
|
|
b = 2 * ellipse.a**2 * m * yint
|
|
c = ellipse.a**2 * (yint**2 - ellipse.b**2)
|
|
|
|
x1, x2 = quadratic(a, b, c)
|
|
y1, y2 = m*x1 + yint, m*x2 + yint
|
|
|
|
return np.array([x1, y1]), np.array([x2, y2])
|
|
|
|
def reflect(ellipse, A, B):
|
|
point = B
|
|
dir = B - A
|
|
n = ellipse.normal(point)
|
|
ref = 2 * (n.dot(dir) * n - dir) + dir
|
|
ps = intersections(ellipse, point, ref)
|
|
return max(ps, key=lambda v: (B-v).dot(B-v))
|
|
|
|
def plot_ellipse(ellipse):
|
|
t = np.linspace(0, 2*pi, 100)
|
|
plt.plot(ellipse.a*np.cos(t), ellipse.b*np.sin(t))
|
|
|
|
ellipse = Ellipse(5, 10)
|
|
A = np.array([0.0, 10.1])
|
|
B = np.array([1.4, -9.6])
|
|
x = 0.01
|
|
|
|
|
|
plt.axis('equal')
|
|
plot_ellipse(ellipse)
|
|
|
|
count = 0
|
|
|
|
while not (-x <= B[0] <= x and B[1] > 0):
|
|
plt.plot(*zip(A, B))
|
|
A, B = B, reflect(ellipse, A, B)
|
|
count += 1
|
|
|
|
|
|
plt.show()
|
|
print(count)
|