465 lines
18 KiB
Python
465 lines
18 KiB
Python
import numpy as np
|
|
from numpy.polynomial.polynomial import Polynomial
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.animation as animation
|
|
from matplotlib.widgets import Slider, Button
|
|
|
|
RAYTRACING = False
|
|
DISTANCE = False
|
|
|
|
|
|
def uproot(arr):
|
|
mask = np.logical_and(arr > 0, np.isreal(arr))
|
|
masked = arr[mask]
|
|
return np.real(sorted(masked))
|
|
|
|
|
|
class TorusWorld:
|
|
|
|
@staticmethod
|
|
def coef(rmaj, rmin, d, o):
|
|
"""
|
|
curtesy of: http://blog.marcinchwedczuk.pl/ray-tracing-torus
|
|
"""
|
|
k1 = np.inner(d, d)
|
|
k2 = np.inner(o, d)
|
|
k3 = np.inner(o, o) - (rmin**2 + rmaj**2)
|
|
|
|
c4 = k1**2
|
|
c3 = 4*k1*k2
|
|
c2 = 2*k1*k3 + 4*k2**2 + 4*(rmaj*d[2])**2
|
|
c1 = 4*k3*k2 + 8*(rmaj**2)*o[2]*d[2]
|
|
c0 = k3**2 - 4*(rmaj**2)*(rmin**2 - o[2]**2)
|
|
return [c0, c1, c2, c3, c4]
|
|
|
|
def __init__(self, rfrac, sun=np.pi/2):
|
|
self.r_maj = 1
|
|
self.r_min = rfrac
|
|
self.sun = None
|
|
self.sun_r = None
|
|
self.put_sun(sun, 0)
|
|
self.surface_map = None
|
|
|
|
def update(self, rfrac=None):
|
|
rfrac = rfrac if rfrac else self.r_min
|
|
self.r_min = rfrac
|
|
self.put_sun(*self.sun_r)
|
|
|
|
def put_sun(self, phi_in, theta_in):
|
|
phi, theta = (np.array([phi_in, theta_in]) + np.pi) % (2 * np.pi) - np.pi
|
|
self.sun_r = np.array([phi, theta])
|
|
self.sun = np.array([
|
|
(self.r_maj - self.r_maj*np.cos(phi))*np.cos(theta),
|
|
(self.r_maj - self.r_maj*np.cos(phi))*np.sin(theta),
|
|
self.r_maj * np.sin(phi),
|
|
])
|
|
|
|
def surface_point(self, phi, theta):
|
|
rx = (self.r_maj - self.r_min*np.cos(phi)) * np.cos(theta)
|
|
ry = (self.r_maj - self.r_min*np.cos(phi)) * np.sin(theta)
|
|
rz = self.r_min * np.sin(phi)
|
|
return rx, ry, rz
|
|
|
|
def normal_vector(self, phi, theta):
|
|
rx, ry, rz = self.surface_point(phi, theta)
|
|
cx = self.r_maj * np.cos(theta)
|
|
cy = self.r_maj * np.sin(theta)
|
|
cz = 0
|
|
retx, rety, retz = rx-cx, ry-cy, rz-cz
|
|
return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz])
|
|
|
|
def ray_vector(self, phi, theta):
|
|
rx, ry, rz = self.surface_point(phi, theta)
|
|
sx, sy, sz = self.sun
|
|
retx, rety, retz = rx-sx, ry-sy, rz-sz
|
|
length = np.linalg.norm([retx, rety, retz])
|
|
norm = length*length if DISTANCE else length
|
|
return np.array([retx, rety, retz])/norm
|
|
|
|
def is_illuminated(self, phi, theta):
|
|
ray = self.ray_vector(phi, theta)
|
|
pol = Polynomial(self.coef(self.r_maj, self.r_min, ray, self.sun))
|
|
roots = uproot(pol.roots())
|
|
if len(roots) == 0:
|
|
return False
|
|
return np.allclose(self.surface_point(phi, theta), self.sun+roots[0]*ray)
|
|
|
|
def illumination(self, phi, theta):
|
|
rt = self.is_illuminated(phi, theta) if RAYTRACING else 1
|
|
return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta)) * rt
|
|
|
|
|
|
class Image():
|
|
|
|
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
|
|
self.fig, self.ax = plt.subplots(figsize=(10, 16))
|
|
|
|
self.lines = dict()
|
|
|
|
self.sun_kwargs = dict(marker='o', color='r', markersize=10,)
|
|
self.contour_kwargs = dict(cmap=plt.colormaps['hot'], vmin=0, vmax=1)
|
|
self.levels = 25 # [-1, 0, 1]
|
|
|
|
self.torus = TorusWorld(rfrac_init, sun_init)
|
|
self.illumination = None
|
|
self.res = 200
|
|
self.rectangular_map = True
|
|
|
|
self.update_illumination()
|
|
|
|
self.ax.set_aspect('equal')
|
|
self.ax.axis('off')
|
|
|
|
self.ax.set_xlim(-2.05, 2.05)
|
|
self.init_top_view()
|
|
self.init_side_view()
|
|
self.init_map_view()
|
|
|
|
def save(self, path="./torus.png", dpi=72):
|
|
self.fig.tight_layout()
|
|
|
|
plt.savefig(path, dpi=dpi)
|
|
|
|
def update_illumination(self):
|
|
# TODO: refactor TorusWorld.illumination() do use vectoization!!
|
|
phi = np.linspace(-np.pi, np.pi, self.res)
|
|
theta = np.linspace(-np.pi, np.pi, self.res)
|
|
self.illumination = np.array([[self.torus.illumination(ph, th) for th in theta] for ph in phi])
|
|
|
|
def init_map_view(self):
|
|
self.lines['map_border'], = self.ax.plot(*self._mantle_map(), 'k')
|
|
self.lines['pos_map'], = self.ax.plot(*self._sunpos_map(), marker='o', color='r', markersize=10,)
|
|
self.lines['dawnline_map'] = [
|
|
self.ax.contourf(*self._contour_map(), self.levels, **self.contour_kwargs),
|
|
]
|
|
|
|
def init_top_view(self):
|
|
self.lines['circles_top'], = self.ax.plot(*self._top_section(), 'k')
|
|
self.lines['path_top'], = self.ax.plot(*self._sunpath_top(), 'k:')
|
|
self.lines['pos_top'], = self.ax.plot(*self._sunpos_top(), **self.sun_kwargs)
|
|
self.lines['dawnline_top'] = [
|
|
self.ax.contourf(*self._contour_top(), self.levels, **self.contour_kwargs),
|
|
]
|
|
|
|
def init_side_view(self):
|
|
self.lines['circles_side'], = self.ax.plot(*self._crossection(), 'k')
|
|
self.lines['path_side'], = self.ax.plot(*self._sunpath_side(), 'k:')
|
|
self.lines['pos_side'], = self.ax.plot(*self._sunpos_side(), **self.sun_kwargs)
|
|
self.lines['dawnline_side'] = [
|
|
self.ax.contourf(*self._contour_side(), self.levels, **self.contour_kwargs),
|
|
]
|
|
|
|
def _offset_map(self):
|
|
return -self.torus.r_maj*1.3 - (1+np.pi*self._scale_map())*self.torus.r_min
|
|
|
|
def _offset_side(self):
|
|
return self.torus.r_maj*1.2
|
|
|
|
def _scale_map(self):
|
|
if self.rectangular_map:
|
|
return (self.torus.r_maj+self.torus.r_min)/(self.torus.r_maj*np.pi)
|
|
else:
|
|
return 1/np.pi
|
|
|
|
def _mantle_map(self, n=1000):
|
|
if self.rectangular_map:
|
|
a = self.torus.r_maj*np.pi
|
|
b = self.torus.r_min*np.pi
|
|
data = np.transpose([[a, b], [-a, b], [-a, -b], [a, -b], [a, b]])
|
|
else:
|
|
phi = np.linspace(-np.pi, np.pi, n)
|
|
u = phi * self.torus.r_min
|
|
width = np.pi*(self.torus.r_maj - self.torus.r_min*np.cos(phi))
|
|
data = np.array([width, u])
|
|
data2 = np.array([-width, u[::-1]])
|
|
|
|
data = np.append(data, data2, axis=1)
|
|
data = np.append(data, data[:, 0:1], axis=1)
|
|
data *= self._scale_map()
|
|
data[1, :] += self._offset_map()
|
|
return data
|
|
|
|
def _contour_map(self):
|
|
if self.rectangular_map:
|
|
a = self.torus.r_maj*np.pi
|
|
b = self.torus.r_min*np.pi
|
|
phi = np.linspace(-b, b, self.res)
|
|
theta = np.linspace(-a, a, self.res)
|
|
x, y = np.meshgrid(theta, phi)
|
|
else:
|
|
phi = np.linspace(-np.pi, np.pi, self.res)
|
|
width = (np.pi)*(self.torus.r_maj - self.torus.r_min*np.cos(phi))
|
|
x = np.array([np.linspace(-w, w, self.res) for w in width])
|
|
y = np.array([[row]*self.res for row in self.torus.r_min*phi])
|
|
x, y = np.meshgrid(np.linspace(-width[0], width[0], self.res), self.torus.r_min*phi)
|
|
z = self.illumination
|
|
x *= self._scale_map()
|
|
y *= self._scale_map()
|
|
return x, y + self._offset_map(), z
|
|
|
|
def _contour_top(self):
|
|
n = self.res
|
|
phi = np.linspace(0, np.pi, int(np.ceil(n/2)))
|
|
theta = np.linspace(0, -np.pi, int(np.ceil(n/2)))
|
|
phx, thy = np.meshgrid(phi, theta)
|
|
x, y, _ = self.torus.surface_point(np.transpose(phx), np.transpose(thy))
|
|
z = self.illumination[int(np.floor(n/2)):, int(np.floor(n/2)):]
|
|
return x, y, z
|
|
|
|
def _contour_side(self):
|
|
n = self.res
|
|
phi = np.linspace(-np.pi/2, np.pi/2, int(np.ceil(n/2)))
|
|
theta = np.linspace(np.pi, 0, int(np.ceil(n/2)))
|
|
phx, thy = np.meshgrid(phi, theta)
|
|
x, _, y = self.torus.surface_point(np.transpose(phx), np.transpose(thy))
|
|
z = self.illumination[int(np.floor(n/4)):int(np.ceil(n*3/4)), :int(np.floor(n/2))]
|
|
return x, y + self._offset_side(), z
|
|
|
|
def _sunpos_map(self):
|
|
if self.rectangular_map:
|
|
y, x = self.torus.sun_r * np.array([self.torus.r_min, self.torus.r_maj])
|
|
else:
|
|
phi, theta = self.torus.sun_r
|
|
x = (self.torus.r_maj - self.torus.r_min*np.cos(phi)) * theta
|
|
y = self.torus.r_min * phi
|
|
x *= self._scale_map()
|
|
y *= self._scale_map()
|
|
return x, y + self._offset_map()
|
|
|
|
def _top_section(self, n=1000):
|
|
theta = np.linspace(0, np.pi, n)
|
|
x1 = (self.torus.r_maj + self.torus.r_min) * np.cos(theta)
|
|
y1 = -1 * (self.torus.r_maj + self.torus.r_min) * np.sin(theta)
|
|
data1 = np.array([x1, y1])
|
|
x2 = (self.torus.r_maj - self.torus.r_min) * np.cos(theta)
|
|
y2 = -1 * (self.torus.r_maj - self.torus.r_min) * np.sin(theta)
|
|
data2 = np.array([x2, y2])
|
|
|
|
data = np.append(data1, [[np.nan,], [np.nan,]], axis=1)
|
|
data = np.append(data, data2, axis=1)
|
|
return data
|
|
|
|
def _sunpath_top(self, n=1000):
|
|
phi = np.linspace(0, np.pi, n)
|
|
_, theta = self.torus.sun_r
|
|
x = (self.torus.r_maj + self.torus.r_maj*np.cos(phi)) * np.cos(theta)
|
|
y = (self.torus.r_maj + self.torus.r_maj*np.cos(phi)) * np.sin(theta)
|
|
return np.array([x, y])
|
|
|
|
def _sunpos_top(self):
|
|
x, y, _ = self.torus.sun
|
|
return x, y
|
|
|
|
def _crossection(self, n=1000):
|
|
phi = np.linspace(0, 2*np.pi, n)
|
|
x = self.torus.r_maj + self.torus.r_min*np.cos(phi)
|
|
y = self.torus.r_min*np.sin(phi)
|
|
data = np.array([x, y])
|
|
data2 = np.array([-x, y])
|
|
data = np.append(data, [[np.nan,], [np.nan,]], axis=1)
|
|
data = np.append(data, data2, axis=1)
|
|
data[1, :] += self._offset_side()
|
|
return data
|
|
|
|
def _sunpath_side(self, n=1000):
|
|
phi = np.linspace(0, 2*np.pi, n)
|
|
_, theta = self.torus.sun_r
|
|
x = (self.torus.r_maj + self.torus.r_maj*np.cos(phi)) * np.cos(theta)
|
|
y = self.torus.r_maj*np.sin(phi)
|
|
return np.array([x, y+self._offset_side()])
|
|
|
|
def _sunpos_side(self):
|
|
x, _, z = self.torus.sun
|
|
return x, z+self._offset_side()
|
|
|
|
def redraw_plot(self, line, func):
|
|
line.set_data(*func())
|
|
self.ax.relim()
|
|
self.ax.autoscale_view()
|
|
|
|
def redraw_contourf(self, container, func, levels=None, contour_kwargs=None):
|
|
levels = levels if levels else [-1, 0, 1]
|
|
contour_kwargs = contour_kwargs if contour_kwargs else {}
|
|
for coll in container[0].collections:
|
|
coll.remove()
|
|
container[0] = self.ax.contourf(*func(), levels, **contour_kwargs)
|
|
|
|
def update_torus(self, rfrac):
|
|
self.torus.update(rfrac)
|
|
self.update_illumination()
|
|
self.redraw_plot(self.lines['map_border'], self._mantle_map)
|
|
self.redraw_plot(self.lines['pos_map'], self._sunpos_map)
|
|
self.redraw_plot(self.lines['circles_side'], self._crossection)
|
|
self.redraw_plot(self.lines['path_side'], self._sunpath_side)
|
|
self.redraw_plot(self.lines['circles_top'], self._top_section)
|
|
self.redraw_plot(self.lines['path_top'], self._sunpath_top)
|
|
self.redraw_contourf(self.lines['dawnline_map'],
|
|
self._contour_map, self.levels, self.contour_kwargs)
|
|
self.redraw_contourf(self.lines['dawnline_top'],
|
|
self._contour_top, self.levels, self.contour_kwargs)
|
|
self.redraw_contourf(self.lines['dawnline_side'],
|
|
self._contour_side, self.levels, self.contour_kwargs)
|
|
self.fig.canvas.draw_idle()
|
|
|
|
def update_sun(self, phi, theta):
|
|
self.torus.put_sun(phi, theta)
|
|
self.update_illumination()
|
|
self.redraw_plot(self.lines['pos_map'], self._sunpos_map)
|
|
self.redraw_plot(self.lines['pos_side'], self._sunpos_side)
|
|
self.redraw_plot(self.lines['pos_top'], self._sunpos_top)
|
|
self.redraw_contourf(self.lines['dawnline_map'],
|
|
self._contour_map, self.levels, self.contour_kwargs)
|
|
self.redraw_contourf(self.lines['dawnline_top'],
|
|
self._contour_top, self.levels, self.contour_kwargs)
|
|
self.redraw_contourf(self.lines['dawnline_side'],
|
|
self._contour_side, self.levels, self.contour_kwargs)
|
|
self.fig.canvas.draw_idle()
|
|
|
|
|
|
class InteractiveImage(Image):
|
|
|
|
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
|
|
super().__init__(rfrac_init, sun_init)
|
|
self.init_interactivity(rfrac_init, sun_init)
|
|
|
|
def init_interactivity(self, rfrac_init, sun_init):
|
|
self.fig.subplots_adjust(left=0.25, bottom=0.25)
|
|
ax1 = self.fig.add_axes([0.25, 0.1, 0.65, 0.03])
|
|
ax2 = self.fig.add_axes([0.1, 0.25, 0.0225, 0.63])
|
|
ax3 = self.fig.add_axes([0.8, 0.025, 0.1, 0.04])
|
|
self.interactions = dict(
|
|
slider_sun=Slider(
|
|
ax=ax1,
|
|
label='Angle of Sun',
|
|
valmin=-np.pi,
|
|
valmax=np.pi,
|
|
valinit=sun_init,
|
|
),
|
|
slider_rfrac=Slider(
|
|
ax=ax2,
|
|
label="Fraction of Radii (r/R)",
|
|
valmin=0,
|
|
valmax=1,
|
|
valinit=rfrac_init,
|
|
orientation="vertical"
|
|
),
|
|
button_reset=Button(ax3, 'Reset', hovercolor='0.975'),
|
|
)
|
|
self.interactions['slider_sun'].on_changed(self._slider_update_sun)
|
|
self.interactions['slider_rfrac'].on_changed(self._slider_update_torus)
|
|
self.interactions['button_reset'].on_clicked(self._reset)
|
|
|
|
|
|
def _slider_update_torus(self, val):
|
|
self.update_torus(val)
|
|
|
|
def _slider_update_sun(self, val):
|
|
self.update_sun(val, 0)
|
|
|
|
def _reset(self, event):
|
|
self.interactions['slider_sun'].reset()
|
|
self.interactions['slider_rfrac'].reset()
|
|
|
|
|
|
class AnimatedImage(Image):
|
|
|
|
def __init__(self, rfrac_init=0.5, sun_init=-np.pi, runtime_sec=5, fps=30):
|
|
super().__init__(rfrac_init, sun_init)
|
|
|
|
self.fig.tight_layout()
|
|
|
|
self.fps = fps
|
|
self.frame_num = self.fps * runtime_sec
|
|
self.sun_init = sun_init
|
|
self.ani = animation.FuncAnimation(self.fig, self.animate, frames=self.frame_num)
|
|
|
|
def animate(self, frame_i):
|
|
phi = (frame_i*2*np.pi/self.frame_num) + self.sun_init
|
|
self.update_sun(phi, 0)
|
|
return self.lines.values()
|
|
|
|
def save(self, path="./torus.mp4"):
|
|
if path.endswith('.gif'):
|
|
writer = animation.PillowWriter(fps=self.fps)
|
|
else:
|
|
writer = animation.FFMpegWriter(fps=self.fps)
|
|
self.ani.save(path, writer)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import argparse
|
|
import pathlib
|
|
|
|
def parse_args():
|
|
|
|
def create_static(args):
|
|
img = Image(args.r_minor, args.sun)
|
|
|
|
if args.output:
|
|
img.save(args.output)
|
|
else:
|
|
plt.show()
|
|
|
|
def create_interactive(args):
|
|
InteractiveImage(args.r_minor, args.sun)
|
|
plt.show()
|
|
|
|
def create_animation(args):
|
|
img = AnimatedImage(args.r_minor, args.sun, args.runtime, args.fps)
|
|
|
|
if args.output:
|
|
img.save(args.output)
|
|
else:
|
|
plt.show()
|
|
|
|
class Range(object):
|
|
def __init__(self, start, end):
|
|
self.start = start
|
|
self.end = end
|
|
|
|
def __eq__(self, other):
|
|
return self.start <= other <= self.end
|
|
|
|
def __contains__(self, item):
|
|
return self.__eq__(item)
|
|
|
|
def __iter__(self):
|
|
yield self
|
|
|
|
def __str__(self):
|
|
return '[{0},{1}]'.format(self.start, self.end)
|
|
|
|
p = argparse.ArgumentParser()
|
|
sp = p.add_subparsers()
|
|
p_static = sp.add_parser('static', aliases=['s'], help='create a static image')
|
|
p_interactive = sp.add_parser('interactive', aliases=['i'], help='create an interactive image')
|
|
p_animated = sp.add_parser('animated', aliases=['a'], help='create a animated image')
|
|
|
|
p_static.add_argument('r_minor', help='Ratio of minor radius to major radius',
|
|
nargs='?', type=float, choices=Range(0., 1.), metavar='RMIN', default=0.5)
|
|
p_static.add_argument('sun', help="Position of the sun as an angle measured from the torus' origin",
|
|
nargs='?', type=float, metavar='PHI', default=np.pi/2)
|
|
p_static.add_argument('--output', '-o', help='output to a file instead of displaying directly',
|
|
type=pathlib.Path, metavar='FILE')
|
|
p_static.set_defaults(func=create_static)
|
|
|
|
p_interactive.add_argument('r_minor', help='Ratio of minor radius to major radius',
|
|
nargs='?', type=float, choices=Range(0., 1.), metavar='RMIN', default=0.5)
|
|
p_interactive.add_argument('sun', help="Position of the sun as an angle measured from the torus' origin",
|
|
nargs='?', type=float, metavar='PHI', default=np.pi/2)
|
|
p_interactive.set_defaults(func=create_interactive)
|
|
|
|
p_animated.add_argument('r_minor', help='Ratio of minor radius to major radius',
|
|
nargs='?', type=float, choices=Range(0., 1.), metavar='RMIN', default=0.5)
|
|
p_animated.add_argument('sun', help="Position of the sun as an angle measured from the torus' origin",
|
|
nargs='?', type=float, metavar='PHI', default=np.pi/2)
|
|
p_animated.add_argument('--output', '-o', help='output to a file instead of displaying directly',
|
|
metavar='FILE')
|
|
p_animated.add_argument('--fps', '-f', help='framerate of animation', metavar='FPS', type=int, default=20)
|
|
p_animated.add_argument('--runtime', '-r', help='runtime of animation', metavar='SEC', type=int, default=5)
|
|
p_animated.set_defaults(func=create_animation)
|
|
|
|
return (p.parse_args())
|
|
|
|
args = parse_args()
|
|
args.func(args)
|