From 7b8a2dae260679739a12a0456c689f54b7790d5d Mon Sep 17 00:00:00 2001 From: Daniel Mevec Date: Sat, 11 Feb 2023 13:04:58 +0100 Subject: [PATCH] refactor to OOP --- torus.py | 506 +++++++++++++++++++++++++++++-------------------------- 1 file changed, 269 insertions(+), 237 deletions(-) diff --git a/torus.py b/torus.py index 653d8bd..ab343e0 100644 --- a/torus.py +++ b/torus.py @@ -1,272 +1,304 @@ import numpy as np from numpy.polynomial.polynomial import Polynomial import matplotlib.pyplot as plt +import matplotlib.colors as colors from matplotlib.widgets import Slider, Button -# The parametrized function to be plotted -def mantle(rfrac, n=1000): - theta = np.linspace(-np.pi, np.pi, n) - rmaj = 1 - rmin = rmaj*rfrac - u = theta * rmin - func = (np.pi)*(rmaj - rmin*np.cos(theta)) - data = np.array([func, u]) - data2 = np.array([-func, u[::-1]]) - - data = np.append(data, data2, axis=1) - data = np.append(data, data[:, 0:1], axis=1) - return data - - -def crossection(rfrac, n=1000): - theta = np.linspace(0, 2*np.pi, n) - rmaj = 1 - rmin = rmaj * rfrac - x = rmaj + rmin*np.cos(theta) - y = rmin*np.sin(theta) - 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) - return data - - -def sunpath_side(rfrac, n=1000): - theta = np.linspace(0, 2*np.pi, n) - rmaj = 1 - x = rmaj + rmaj*np.cos(theta) - y = rmaj*np.sin(theta) - return np.array([x, y]) - - -def sunpath_top(rfrac, n=1000): - theta = np.linspace(0, np.pi, n) - rmaj = 1 - x = rmaj + rmaj*np.cos(theta) - y = np.zeros_like(theta) - return np.array([x, y]) - - -def sun_side(rfrac, sunpos): - rmaj = 1 - x = rmaj - rmaj*np.cos(sunpos) - y = rmaj*np.sin(sunpos) - return [x, y] - - -def sun_top(rfrac, sunpos): - rmaj = 1 - x = rmaj - rmaj*np.cos(sunpos) - y = 0 - return [x, y] - - -def sun_map(rfrac, sunpos): - rmaj = 1 - rmin = rmaj*rfrac - y = sunpos * rmin - x = 0 - return [x, y] - - -def topview(rfrac, n=1000): - theta = np.linspace(0, np.pi, n) - rmaj = 1 - rmin = rmaj*rfrac - x1 = (rmaj + rmin) * np.cos(theta) - y1 = -1 * (rmaj + rmin) * np.sin(theta) - data1 = np.array([x1, y1]) - x2 = (rmaj - rmin) * np.cos(theta) - y2 = -1 * (rmaj - rmin) * 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 contour_map(rfrac, sunpos, n=100): - rmaj = 1 - rmin = rmaj*rfrac - phi = np.linspace(-np.pi, np.pi, n) - theta = np.linspace(-np.pi, np.pi, int(n*rfrac)) - - y = np.array([[w]*n for w in rmin*theta]) - func = (np.pi)*(rmaj - rmin*np.cos(theta)) - x = np.array([np.linspace(-f, f, n) for f in func]) - # x, y = np.meshgrid(np.linspace(-func[0], func[0], n), rmin*theta) - - sun = [rmaj - rmaj*np.cos(sunpos), 0, rmaj*np.sin(sunpos)] - - def raytraced(ph, th): - def r(ph, th): - rx = (rmaj - rmin*np.cos(th)) * np.cos(ph) - ry = (rmaj - rmin*np.cos(th)) * np.sin(ph) - rz = rmin * np.sin(th) - return rx, ry, rz - - def n(ph, th): - rx, ry, rz = r(ph, th) - cx = rmaj * np.cos(ph) - cy = rmaj * np.sin(ph) - cz = 0 - retx, rety, retz = rx-cx, ry-cy, rz-cz - return -1*np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz]) - - def b(ph, th): - rx, ry, rz = r(ph, th) - sx, sy, sz = sun - retx, rety, retz = rx-sx, ry-sy, rz-sz - return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz]) - o = sun - d = b(ph, th) - s = r(ph, th) - - pol = Polynomial(coef(rmaj, rmin, d, o)) - roots = uproot(pol.roots()) - - if len(roots) == 0: - return 0 - poi = o + roots[0]*d - return (1 if np.allclose(s, poi) else -1) - - z = np.array([[raytraced(ph, th) for ph in phi] for th in theta]) - # z = np.array([[np.dot(b(ph, th), n(ph, th)) for ph in phi] for th in theta]) - - return x, y, z - - -def coef(rmaj, rmin, d, o): - 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 uproot(arr): mask = np.logical_and(arr > 0, np.isreal(arr)) masked = arr[mask] return np.real(sorted(masked)) -r_fraction = 0.5 -sun_pos = np.pi/2 +class TorusWorld: -# Create the figure and the line that we will manipulate -fig, (ax_side, ax_top, ax_map) = plt.subplots(3, 1, gridspec_kw={'height_ratios': [1, 1, 1]}) + @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) -circles_top, = ax_top.plot(*topview(r_fraction), 'k') -path_top, = ax_top.plot(*sunpath_top(r_fraction), 'k:') -pos_top, = ax_top.plot(*sun_top(r_fraction, sun_pos), marker='o', color='r', markersize=10,) -ax_top.set_xlim(-2.05, 2.05) -ax_top.set_ylim(-2.05, None) -ax_top.set_aspect('equal') + 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] -ax_top.axis('off') + def __init__(self, rfrac): + self.r_maj = 1 + self.r_min = rfrac + self.sun = np.array([self.r_maj, 0, self.r_maj]) + self.sun_r = np.array([np.pi/2, 0]) + self.surface_map = None -circles_side, = ax_side.plot(*crossection(r_fraction), 'k') -path_side, = ax_side.plot(*sunpath_side(r_fraction), 'k:') -pos_side, = ax_side.plot(*sun_side(r_fraction, sun_pos), marker='o', color='r', markersize=10,) -ax_side.set_xlim(-2.05, 2.05) -ax_side.set_aspect('equal') -# ax_side.set_ylim(-r_min, None) -ax_side.axis('off') + def update(self, rfrac=None): + rfrac = rfrac if rfrac else self.r_min + self.r_min = rfrac + self.put_sun(*self.sun_r) -map_border, = ax_map.plot(*mantle(r_fraction), 'k') -pos_map, = ax_map.plot(*sun_map(r_fraction, sun_pos), marker='o', color='r', markersize=10,) -dawnline_map = [ax_map.contourf(*contour_map(r_fraction, sun_pos), [-1, 0, 1], cmap='YlOrBr_r')] -ax_map.set_aspect('equal') -ax_map.set_ylabel('Longitude') -ax_map.set_xlabel('Lattitude') -ax_map.axis('off') -# cbar = fig.colorbar(dawnline_map[0]) + def put_sun(self, phi, theta): + 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), + ]) -# adjust the main plot to make room for the sliders -fig.subplots_adjust(left=0.25, bottom=0.25) + 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 -# Make a horizontal slider to control the frequency. -axsun = fig.add_axes([0.25, 0.1, 0.65, 0.03]) -slider_sun = Slider( - ax=axsun, - label='Angle of Sun', - valmin=-np.pi, - valmax=np.pi, - valinit=sun_pos, -) + 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]) -# Make a vertically oriented slider to control the amplitude -axrad = fig.add_axes([0.1, 0.25, 0.0225, 0.63]) -slider_rf = Slider( - ax=axrad, - label="Fraction of Radii (r/R)", - valmin=0, - valmax=1, - valinit=r_fraction, - orientation="vertical" -) + 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 + return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz]) + + 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): + # return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta)) + return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta)) * self.is_illuminated(phi, theta) -# The function to be called anytime a slider's value changes -def update_torus(val): - def update(ax, line, func): - u, m = func(slider_rf.val) - line.set_ydata(m) - line.set_xdata(u) +class NoInamge(): + def __init__(self, rfrac_init=0.5, sun_init=np.pi/2): + fig, (ax_side, ax_top, ax_map) = plt.subplots(3, 1, gridspec_kw={'height_ratios': [2, 2, 1]}) + self.fig = fig + self.ax = dict( + map=ax_map, + top=ax_top, + side=ax_side, + ) + 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) + self.init_top_view() + self.init_side_view() + self.init_map_view() + + def init_map_view(self): + self.lines['map_border'], = self.ax['map'].plot(*self._mantle_map(), 'k') + self.lines['pos_map'], = self.ax['map'].plot(*self._sunpos_map(), marker='o', color='r', markersize=10,) + self.lines['dawnline_map'] = [ + self.ax['map'].contourf(*self._contour_map(), self.levels, **self.contour_kwargs), + ] + self.ax['map'].set_aspect('equal') + self.ax['map'].set_ylabel('Longitude') + self.ax['map'].set_xlabel('Lattitude') + self.ax['map'].axis('off') + + def init_top_view(self): + self.lines['circles_top'], = self.ax['top'].plot(*self._top_section(), 'k') + self.lines['path_top'], = self.ax['top'].plot(*self._sunpath_top(), 'k:') + self.lines['pos_top'], = self.ax['top'].plot(*self._sunpos_top(), **self.sun_kwargs) + self.ax['top'].set_xlim(-2.05, 2.05) + self.ax['top'].set_ylim(-2.05, None) + self.ax['top'].set_aspect('equal') + + self.ax['top'].axis('off') + + def init_side_view(self): + self.lines['circles_side'], = self.ax['side'].plot(*self._crossection(), 'k') + self.lines['path_side'], = self.ax['side'].plot(*self._sunpath_side(), 'k:') + self.lines['pos_side'], = self.ax['side'].plot(*self._sunpos_side(), **self.sun_kwargs) + self.ax['side'].set_xlim(-2.05, 2.05) + self.ax['side'].set_aspect('equal') + # self.ax['side'].set_ylim(-self.torus.r_min, None) + self.ax['side'].axis('off') + + def _mantle_map(self, n=1000): + 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) + return data + + def _contour_map(self, n=100): + theta = np.linspace(-np.pi, np.pi, n) + phi = np.linspace(-np.pi, np.pi, int(n*self.torus.r_min)) + + y = np.array([[w]*n for w in self.torus.r_min*phi]) + func = (np.pi)*(self.torus.r_maj - self.torus.r_min*np.cos(phi)) + x = np.array([np.linspace(-f, f, n) for f in func]) + # x, y = np.meshgrid(np.linspace(-func[0], func[0], n), rmin*theta) + + z = np.array([[self.torus.illumination(ph, th) for th in theta] for ph in phi]) + + return x, y, z + + def _sunpos_map(self): + 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 + return x, y + + def _top_section(self, n=1000): + phi = np.linspace(0, np.pi, n) + x1 = (self.torus.r_maj + self.torus.r_min) * np.cos(phi) + y1 = -1 * (self.torus.r_maj + self.torus.r_min) * np.sin(phi) + data1 = np.array([x1, y1]) + x2 = (self.torus.r_maj - self.torus.r_min) * np.cos(phi) + y2 = -1 * (self.torus.r_maj - self.torus.r_min) * np.sin(phi) + 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) + x = self.torus.r_maj + self.torus.r_maj*np.cos(phi) + y = np.zeros_like(phi) + return np.array([x, y]) + + def _sunpos_top(self): + phi, theta = self.torus.sun_r + x = self.torus.r_maj + self.torus.r_maj*np.cos(phi) + y = 0 + 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) + return data + + def _sunpath_side(self, n=1000): + phi = np.linspace(0, 2*np.pi, n) + x = self.torus.r_maj + self.torus.r_maj*np.cos(phi) + y = self.torus.r_maj*np.sin(phi) + return np.array([x, y]) + + def _sunpos_side(self): + phi, theta = self.torus.sun_r + x = self.torus.r_maj - self.torus.r_maj*np.cos(phi) + y = self.torus.r_maj*np.sin(phi) + return x, y + + @staticmethod + def redraw_plot(ax, line, func): + x, y = func() + line.set_xdata(x) + line.set_ydata(y) ax.relim() ax.autoscale_view() - update(ax_map, map_border, mantle) - update(ax_side, circles_side, crossection) - update(ax_side, path_side, sunpath_side) - update(ax_top, circles_top, topview) - update(ax_top, path_top, sunpath_top) + @staticmethod + def redraw_contourf(ax, 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] = ax.contourf(*func(), levels, **contour_kwargs) - for coll in dawnline_map[0].collections: - coll.remove() - dawnline_map[0] = ax_map.contourf(*contour_map(slider_rf.val, slider_sun.val), [-1, 0, 1], cmap='YlOrBr_r') + def update_torus(self, rfrac): + self.torus.update(rfrac) + self.redraw_plot(self.ax['map'], self.lines['map_border'], self._mantle_map) + self.redraw_plot(self.ax['side'], self.lines['circles_side'], self._crossection) + self.redraw_plot(self.ax['side'], self.lines['path_side'], self._sunpath_side) + self.redraw_plot(self.ax['top'], self.lines['circles_top'], self._top_section) + self.redraw_plot(self.ax['top'], self.lines['path_top'], self._sunpath_top) + self.redraw_contourf(self.ax['map'], self.lines['dawnline_map'], + self._contour_map, self.levels, self.contour_kwargs) + self.fig.canvas.draw_idle() - fig.canvas.draw_idle() + def update_sun(self, phi, theta): + self.torus.put_sun(phi, theta) + self.redraw_plot(self.ax['map'], self.lines['pos_map'], self._sunpos_map) + self.redraw_plot(self.ax['side'], self.lines['pos_side'], self._sunpos_side) + self.redraw_plot(self.ax['top'], self.lines['pos_top'], self._sunpos_top) + self.redraw_contourf(self.ax['map'], self.lines['dawnline_map'], + self._contour_map, self.levels, self.contour_kwargs) + self.fig.canvas.draw_idle() -# The function to be called anytime a slider's value changes -def update_sun(val): - def update(ax, line, func): - u, m = func(slider_rf.val, slider_sun.val) - line.set_ydata(m) - line.set_xdata(u) +class ImageStatic(NoInamge): - update(ax_map, pos_map, sun_map) - update(ax_side, pos_side, sun_side) - update(ax_top, pos_top, sun_top) - - for coll in dawnline_map[0].collections: - coll.remove() - dawnline_map[0] = ax_map.contourf(*contour_map(slider_rf.val, slider_sun.val), [-1, 0, 1], cmap='YlOrBr_r') - - fig.canvas.draw_idle() + def __init__(self, rfrac_init=0.5, sun_init=np.pi/2): + super().__init__(rfrac_init, sun_init) + plt.show() -# register the update function with each slider -slider_sun.on_changed(update_sun) -slider_rf.on_changed(update_torus) +class ImageInteractive(NoInamge): -# Create a `matplotlib.widgets.Button` to reset the sliders to initial values. -resetax = fig.add_axes([0.8, 0.025, 0.1, 0.04]) -button = Button(resetax, 'Reset', hovercolor='0.975') + 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) + plt.show() + + def init_interactivity(self, rfrac_init, sun_init): + self.fig.subplots_adjust(left=0.25, bottom=0.25) + self.ax['slider_sun'] = self.fig.add_axes([0.25, 0.1, 0.65, 0.03]) + self.ax['slider_rf'] = self.fig.add_axes([0.1, 0.25, 0.0225, 0.63]) + self.ax['button_reset'] = self.fig.add_axes([0.8, 0.025, 0.1, 0.04]) + self.sliders = dict( + sun_phi=Slider( + ax=self.ax['slider_sun'], + label='Angle of Sun', + valmin=-np.pi, + valmax=np.pi, + valinit=sun_init, + ), + rfrac=Slider( + ax=self.ax['slider_rf'], + label="Fraction of Radii (r/R)", + valmin=0, + valmax=1, + valinit=rfrac_init, + orientation="vertical" + ) + + ) + self.sliders['sun_phi'].on_changed(self._slider_update_sun) + self.sliders['rfrac'].on_changed(self._slider_update_torus) + button = Button(self.ax['button_reset'], 'Reset', hovercolor='0.975') + button.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.sliders['sun_phi'].reset() + self.sliders['rfrac'].reset() -def reset(event): - slider_sun.reset() - slider_rf.reset() - - -button.on_clicked(reset) - -plt.show() +if __name__ == '__main__': + # ImageStatic() + ImageInteractive()