From afcce3b98b4c12fa70466031841b087053872040 Mon Sep 17 00:00:00 2001 From: Daniel Mevec Date: Sat, 11 Feb 2023 13:03:45 +0100 Subject: [PATCH] initial commit --- .gitignore | 1 + torus.py | 272 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 273 insertions(+) create mode 100644 .gitignore create mode 100644 torus.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0d20b64 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +*.pyc diff --git a/torus.py b/torus.py new file mode 100644 index 0000000..653d8bd --- /dev/null +++ b/torus.py @@ -0,0 +1,272 @@ +import numpy as np +from numpy.polynomial.polynomial import Polynomial +import matplotlib.pyplot as plt +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 + +# 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]}) + +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') + +ax_top.axis('off') + +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') + +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]) + +# adjust the main plot to make room for the sliders +fig.subplots_adjust(left=0.25, bottom=0.25) + +# 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, +) + +# 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" +) + + +# 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) + 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) + + 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() + + +# 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) + + 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() + + +# register the update function with each slider +slider_sun.on_changed(update_sun) +slider_rf.on_changed(update_torus) + +# 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 reset(event): + slider_sun.reset() + slider_rf.reset() + + +button.on_clicked(reset) + +plt.show()