add raytracing and distance flags to argsparse

This commit is contained in:
Daniel Mevec 2023-02-16 07:49:45 +01:00
parent 22f499fe33
commit 6b47046e36

View file

@ -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')