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 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, raytrace=True, distance=False): self.r_maj = 1 self.r_min = rfrac self.sun = None self.sun_r = None self.put_sun(sun, 0) self.surface_map = None self._rt = raytrace self._dist = distance 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 self._dist 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 self._rt 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, **t_kwargs): 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, **t_kwargs) 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 run(self, path=None, **kwargs): if path: self.save(path, **kwargs) else: plt.show() 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 SummedImage(Image): 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) day_cycle = np.linspace(-np.pi, np.pi, 24) self.torus.put_sun(day_cycle[0], 0) illumination = np.array([[self.torus.illumination(ph, th) for th in theta] for ph in phi]) for sun_pos in day_cycle[1:]: self.torus.put_sun(sun_pos, 0) illumination += np.array([[self.torus.illumination(ph, th) for th in theta] for ph in phi]) self.illumination = illumination / np.amax(illumination) 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(), 'c:') # 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(), 'c:') # 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), ] class InteractiveImage(Image): def __init__(self, rfrac_init=0.5, sun_init=np.pi/2, **t_kwargs): super().__init__(rfrac_init, sun_init, **t_kwargs) 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, **t_kwargs): super().__init__(rfrac_init, sun_init, **t_kwargs) 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): Image(args.r_minor, args.sun, raytrace=args.raytrace, distance=args.distance).run(args.output) def create_interactive(args): InteractiveImage(args.r_minor, args.sun).run() def create_animation(args): AnimatedImage(args.r_minor, args.sun, args.runtime, args.fps).run() def create_summed(args): SummedImage(args.r_minor).run() 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() p.add_argument('--raytrace', '-r', action='store_true', help='enable raytracing, disable shine-through') p.add_argument('--distance', '-d', action='store_true', help='factor in distance into lightfall intensity') 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_summed = sp.add_parser('daylight', aliases=['d'], help='sum up a day\'s light average') 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) p_summed.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_summed.add_argument('--output', '-o', help='output to a file instead of displaying directly', metavar='FILE') p_summed.add_argument('--hours', '-H', action='store_true', help='sum up hours of daylight instead of intensity') p_summed.set_defaults(func=create_summed) return (p.parse_args()) args = parse_args() args.func(args)