refactor to OOP

This commit is contained in:
Daniel Mevec 2023-02-11 13:04:58 +01:00
parent afcce3b98b
commit 7b8a2dae26

506
torus.py
View file

@ -1,272 +1,304 @@
import numpy as np import numpy as np
from numpy.polynomial.polynomial import Polynomial from numpy.polynomial.polynomial import Polynomial
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.widgets import Slider, Button 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): def uproot(arr):
mask = np.logical_and(arr > 0, np.isreal(arr)) mask = np.logical_and(arr > 0, np.isreal(arr))
masked = arr[mask] masked = arr[mask]
return np.real(sorted(masked)) return np.real(sorted(masked))
r_fraction = 0.5 class TorusWorld:
sun_pos = np.pi/2
# Create the figure and the line that we will manipulate @staticmethod
fig, (ax_side, ax_top, ax_map) = plt.subplots(3, 1, gridspec_kw={'height_ratios': [1, 1, 1]}) 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') c4 = k1**2
path_top, = ax_top.plot(*sunpath_top(r_fraction), 'k:') c3 = 4*k1*k2
pos_top, = ax_top.plot(*sun_top(r_fraction, sun_pos), marker='o', color='r', markersize=10,) c2 = 2*k1*k3 + 4*k2**2 + 4*(rmaj*d[2])**2
ax_top.set_xlim(-2.05, 2.05) c1 = 4*k3*k2 + 8*(rmaj**2)*o[2]*d[2]
ax_top.set_ylim(-2.05, None) c0 = k3**2 - 4*(rmaj**2)*(rmin**2 - o[2]**2)
ax_top.set_aspect('equal') 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') def update(self, rfrac=None):
path_side, = ax_side.plot(*sunpath_side(r_fraction), 'k:') rfrac = rfrac if rfrac else self.r_min
pos_side, = ax_side.plot(*sun_side(r_fraction, sun_pos), marker='o', color='r', markersize=10,) self.r_min = rfrac
ax_side.set_xlim(-2.05, 2.05) self.put_sun(*self.sun_r)
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') def put_sun(self, phi, theta):
pos_map, = ax_map.plot(*sun_map(r_fraction, sun_pos), marker='o', color='r', markersize=10,) self.sun_r = np.array([phi, theta])
dawnline_map = [ax_map.contourf(*contour_map(r_fraction, sun_pos), [-1, 0, 1], cmap='YlOrBr_r')] self.sun = np.array([
ax_map.set_aspect('equal') (self.r_maj - self.r_maj*np.cos(phi))*np.cos(theta),
ax_map.set_ylabel('Longitude') (self.r_maj - self.r_maj*np.cos(phi))*np.sin(theta),
ax_map.set_xlabel('Lattitude') self.r_maj * np.sin(phi),
ax_map.axis('off') ])
# cbar = fig.colorbar(dawnline_map[0])
# adjust the main plot to make room for the sliders def surface_point(self, phi, theta):
fig.subplots_adjust(left=0.25, bottom=0.25) 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. def normal_vector(self, phi, theta):
axsun = fig.add_axes([0.25, 0.1, 0.65, 0.03]) rx, ry, rz = self.surface_point(phi, theta)
slider_sun = Slider( cx = self.r_maj * np.cos(theta)
ax=axsun, cy = self.r_maj * np.sin(theta)
label='Angle of Sun', cz = 0
valmin=-np.pi, retx, rety, retz = rx-cx, ry-cy, rz-cz
valmax=np.pi, return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz])
valinit=sun_pos,
)
# Make a vertically oriented slider to control the amplitude def ray_vector(self, phi, theta):
axrad = fig.add_axes([0.1, 0.25, 0.0225, 0.63]) rx, ry, rz = self.surface_point(phi, theta)
slider_rf = Slider( sx, sy, sz = self.sun
ax=axrad, retx, rety, retz = rx-sx, ry-sy, rz-sz
label="Fraction of Radii (r/R)", return np.array([retx, rety, retz])/np.linalg.norm([retx, rety, retz])
valmin=0,
valmax=1, def is_illuminated(self, phi, theta):
valinit=r_fraction, ray = self.ray_vector(phi, theta)
orientation="vertical" 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 class NoInamge():
def update_torus(val): def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
def update(ax, line, func): fig, (ax_side, ax_top, ax_map) = plt.subplots(3, 1, gridspec_kw={'height_ratios': [2, 2, 1]})
u, m = func(slider_rf.val) self.fig = fig
line.set_ydata(m) self.ax = dict(
line.set_xdata(u) 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.relim()
ax.autoscale_view() ax.autoscale_view()
update(ax_map, map_border, mantle) @staticmethod
update(ax_side, circles_side, crossection) def redraw_contourf(ax, container, func, levels=None, contour_kwargs=None):
update(ax_side, path_side, sunpath_side) levels = levels if levels else [-1, 0, 1]
update(ax_top, circles_top, topview) contour_kwargs = contour_kwargs if contour_kwargs else {}
update(ax_top, path_top, sunpath_top) for coll in container[0].collections:
coll.remove()
container[0] = ax.contourf(*func(), levels, **contour_kwargs)
for coll in dawnline_map[0].collections: def update_torus(self, rfrac):
coll.remove() self.torus.update(rfrac)
dawnline_map[0] = ax_map.contourf(*contour_map(slider_rf.val, slider_sun.val), [-1, 0, 1], cmap='YlOrBr_r') 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 class ImageStatic(NoInamge):
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) def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
update(ax_side, pos_side, sun_side) super().__init__(rfrac_init, sun_init)
update(ax_top, pos_top, sun_top) plt.show()
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 class ImageInteractive(NoInamge):
slider_sun.on_changed(update_sun)
slider_rf.on_changed(update_torus)
# Create a `matplotlib.widgets.Button` to reset the sliders to initial values. def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
resetax = fig.add_axes([0.8, 0.025, 0.1, 0.04]) super().__init__(rfrac_init, sun_init)
button = Button(resetax, 'Reset', hovercolor='0.975') 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): if __name__ == '__main__':
slider_sun.reset() # ImageStatic()
slider_rf.reset() ImageInteractive()
button.on_clicked(reset)
plt.show()