initial commit

.gitignore

@@ -0,0 +1 @@
+*.pyc

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()