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
|
||||
from matplotlib.widgets import Slider, Button
|
||||
|
||||
RAYTRACING = False
|
||||
DISTANCE = False
|
||||
|
||||
|
||||
def uproot(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)
|
||||
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_min = rfrac
|
||||
self.sun = None
|
||||
|
|
@ -40,6 +37,9 @@ class TorusWorld:
|
|||
self.put_sun(sun, 0)
|
||||
self.surface_map = None
|
||||
|
||||
self._rt = raytrace
|
||||
self._dist = distance
|
||||
|
||||
def update(self, rfrac=None):
|
||||
rfrac = rfrac if rfrac else self.r_min
|
||||
self.r_min = rfrac
|
||||
|
|
@ -73,7 +73,7 @@ class TorusWorld:
|
|||
sx, sy, sz = self.sun
|
||||
retx, rety, retz = rx-sx, ry-sy, rz-sz
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
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
|
||||
|
||||
|
||||
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.lines = dict()
|
||||
|
|
@ -100,7 +100,7 @@ class Image():
|
|||
self.contour_kwargs = dict(cmap=plt.colormaps['hot'], vmin=0, vmax=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.res = 200
|
||||
self.rectangular_map = True
|
||||
|
|
@ -362,8 +362,8 @@ class SummedImage(Image):
|
|||
|
||||
class InteractiveImage(Image):
|
||||
|
||||
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2):
|
||||
super().__init__(rfrac_init, sun_init)
|
||||
def __init__(self, rfrac_init=0.5, sun_init=np.pi/2, **t_kwargs):
|
||||
super().__init__(rfrac_init, sun_init, **t_kwargs)
|
||||
self.init_interactivity(rfrac_init, sun_init)
|
||||
|
||||
def init_interactivity(self, rfrac_init, sun_init):
|
||||
|
|
@ -406,8 +406,8 @@ class InteractiveImage(Image):
|
|||
|
||||
class AnimatedImage(Image):
|
||||
|
||||
def __init__(self, rfrac_init=0.5, sun_init=-np.pi, runtime_sec=5, fps=30):
|
||||
super().__init__(rfrac_init, sun_init)
|
||||
def __init__(self, rfrac_init=0.5, sun_init=-np.pi, runtime_sec=5, fps=30, **t_kwargs):
|
||||
super().__init__(rfrac_init, sun_init, **t_kwargs)
|
||||
|
||||
self.fig.tight_layout()
|
||||
|
||||
|
|
@ -436,7 +436,7 @@ if __name__ == '__main__':
|
|||
def parse_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):
|
||||
InteractiveImage(args.r_minor, args.sun).run()
|
||||
|
|
@ -465,6 +465,8 @@ if __name__ == '__main__':
|
|||
return '[{0},{1}]'.format(self.start, self.end)
|
||||
|
||||
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()
|
||||
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')
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue