From 6b47046e36f33004f16f83fc3cd6f63e4b6f5ef2 Mon Sep 17 00:00:00 2001 From: Daniel Mevec Date: Thu, 16 Feb 2023 07:49:45 +0100 Subject: [PATCH] add raytracing and distance flags to argsparse --- torus.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/torus.py b/torus.py index e951da2..0412c44 100644 --- a/torus.py +++ b/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')