add raytracing and distance flags to argsparse
This commit is contained in:
parent
22f499fe33
commit
6b47046e36
1 changed files with 15 additions and 13 deletions
28
torus.py
28
torus.py
|
|
@ -4,9 +4,6 @@ import matplotlib.pyplot as plt
|
||||||
import matplotlib.animation as animation
|
import matplotlib.animation as animation
|
||||||
from matplotlib.widgets import Slider, Button
|
from matplotlib.widgets import Slider, Button
|
||||||
|
|
||||||
RAYTRACING = False
|
|
||||||
DISTANCE = False
|
|
||||||
|
|
||||||
|
|
||||||
def uproot(arr):
|
def uproot(arr):
|
||||||
mask = np.logical_and(arr > 0, np.isreal(arr))
|
mask = np.logical_and(arr > 0, np.isreal(arr))
|
||||||
|
|
@ -32,7 +29,7 @@ class TorusWorld:
|
||||||
c0 = k3**2 - 4*(rmaj**2)*(rmin**2 - o[2]**2)
|
c0 = k3**2 - 4*(rmaj**2)*(rmin**2 - o[2]**2)
|
||||||
return [c0, c1, c2, c3, c4]
|
return [c0, c1, c2, c3, c4]
|
||||||
|
|
||||||
def __init__(self, rfrac, sun=np.pi/2):
|
def __init__(self, rfrac, sun=np.pi/2, raytrace=True, distance=False):
|
||||||
self.r_maj = 1
|
self.r_maj = 1
|
||||||
self.r_min = rfrac
|
self.r_min = rfrac
|
||||||
self.sun = None
|
self.sun = None
|
||||||
|
|
@ -40,6 +37,9 @@ class TorusWorld:
|
||||||
self.put_sun(sun, 0)
|
self.put_sun(sun, 0)
|
||||||
self.surface_map = None
|
self.surface_map = None
|
||||||
|
|
||||||
|
self._rt = raytrace
|
||||||
|
self._dist = distance
|
||||||
|
|
||||||
def update(self, rfrac=None):
|
def update(self, rfrac=None):
|
||||||
rfrac = rfrac if rfrac else self.r_min
|
rfrac = rfrac if rfrac else self.r_min
|
||||||
self.r_min = rfrac
|
self.r_min = rfrac
|
||||||
|
|
@ -73,7 +73,7 @@ class TorusWorld:
|
||||||
sx, sy, sz = self.sun
|
sx, sy, sz = self.sun
|
||||||
retx, rety, retz = rx-sx, ry-sy, rz-sz
|
retx, rety, retz = rx-sx, ry-sy, rz-sz
|
||||||
length = np.linalg.norm([retx, rety, retz])
|
length = np.linalg.norm([retx, rety, retz])
|
||||||
norm = length*length if DISTANCE else length
|
norm = length*length if self._dist else length
|
||||||
return np.array([retx, rety, retz])/norm
|
return np.array([retx, rety, retz])/norm
|
||||||
|
|
||||||
def is_illuminated(self, phi, theta):
|
def is_illuminated(self, phi, theta):
|
||||||
|
|
@ -85,13 +85,13 @@ class TorusWorld:
|
||||||
return np.allclose(self.surface_point(phi, theta), self.sun+roots[0]*ray)
|
return np.allclose(self.surface_point(phi, theta), self.sun+roots[0]*ray)
|
||||||
|
|
||||||
def illumination(self, phi, theta):
|
def illumination(self, phi, theta):
|
||||||
rt = self.is_illuminated(phi, theta) if RAYTRACING else 1
|
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
|
return np.dot(self.ray_vector(phi, theta), -self.normal_vector(phi, theta)) * rt
|
||||||
|
|
||||||
|
|
||||||
class Image():
|
class Image():
|
||||||
|
|
||||||
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
|
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2, **t_kwargs):
|
||||||
self.fig, self.ax = plt.subplots(figsize=(10, 16))
|
self.fig, self.ax = plt.subplots(figsize=(10, 16))
|
||||||
|
|
||||||
self.lines = dict()
|
self.lines = dict()
|
||||||
|
|
@ -100,7 +100,7 @@ class Image():
|
||||||
self.contour_kwargs = dict(cmap=plt.colormaps['hot'], vmin=0, vmax=1)
|
self.contour_kwargs = dict(cmap=plt.colormaps['hot'], vmin=0, vmax=1)
|
||||||
self.levels = 25 # [-1, 0, 1]
|
self.levels = 25 # [-1, 0, 1]
|
||||||
|
|
||||||
self.torus = TorusWorld(rfrac_init, sun_init)
|
self.torus = TorusWorld(rfrac_init, sun_init, **t_kwargs)
|
||||||
self.illumination = None
|
self.illumination = None
|
||||||
self.res = 200
|
self.res = 200
|
||||||
self.rectangular_map = True
|
self.rectangular_map = True
|
||||||
|
|
@ -362,8 +362,8 @@ class SummedImage(Image):
|
||||||
|
|
||||||
class InteractiveImage(Image):
|
class InteractiveImage(Image):
|
||||||
|
|
||||||
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
|
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2, **t_kwargs):
|
||||||
super().__init__(rfrac_init, sun_init)
|
super().__init__(rfrac_init, sun_init, **t_kwargs)
|
||||||
self.init_interactivity(rfrac_init, sun_init)
|
self.init_interactivity(rfrac_init, sun_init)
|
||||||
|
|
||||||
def init_interactivity(self, rfrac_init, sun_init):
|
def init_interactivity(self, rfrac_init, sun_init):
|
||||||
|
|
@ -406,8 +406,8 @@ class InteractiveImage(Image):
|
||||||
|
|
||||||
class AnimatedImage(Image):
|
class AnimatedImage(Image):
|
||||||
|
|
||||||
def __init__(self, rfrac_init=0.5, sun_init=-np.pi, runtime_sec=5, fps=30):
|
def __init__(self, rfrac_init=0.5, sun_init=-np.pi, runtime_sec=5, fps=30, **t_kwargs):
|
||||||
super().__init__(rfrac_init, sun_init)
|
super().__init__(rfrac_init, sun_init, **t_kwargs)
|
||||||
|
|
||||||
self.fig.tight_layout()
|
self.fig.tight_layout()
|
||||||
|
|
||||||
|
|
@ -436,7 +436,7 @@ if __name__ == '__main__':
|
||||||
def parse_args():
|
def parse_args():
|
||||||
|
|
||||||
def create_static(args):
|
def create_static(args):
|
||||||
Image(args.r_minor, args.sun).run(args.output)
|
Image(args.r_minor, args.sun, raytrace=args.raytrace, distance=args.distance).run(args.output)
|
||||||
|
|
||||||
def create_interactive(args):
|
def create_interactive(args):
|
||||||
InteractiveImage(args.r_minor, args.sun).run()
|
InteractiveImage(args.r_minor, args.sun).run()
|
||||||
|
|
@ -465,6 +465,8 @@ if __name__ == '__main__':
|
||||||
return '[{0},{1}]'.format(self.start, self.end)
|
return '[{0},{1}]'.format(self.start, self.end)
|
||||||
|
|
||||||
p = argparse.ArgumentParser()
|
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()
|
sp = p.add_subparsers()
|
||||||
p_static = sp.add_parser('static', aliases=['s'], help='create a static image')
|
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_interactive = sp.add_parser('interactive', aliases=['i'], help='create an interactive image')
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue