# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
"""
This module is for optics simulation. This is exposed automatically with::
from k1lib.imports import *
kop.Rays.parallel(...) # exposed
A comprehensive introduction is available here: https://mlexps.com/optics/1-kop-intro/"""
import k1lib, time, math, html, copy; from k1lib import cli, p5; import numpy as np; from collections import deque
__all__ = ["Drawable", "RandomPoints", "Drawables",
"RaysFocus", "Rays", "Wavelengths", "sellmeier", "gps",
"Surface", "LineSurface", "ArcSurface", "polySolver", "PolySurface", "RaysPath",
"OpticElement", "Polygon", "Lens", "OpticSystem"]
colormap_rgb = np.array([[6, 1, 31],[12, 0, 40],[14, 0, 51], [16, 1, 60], [17, 1, 76], [23, 0, 90], [26, 1, 105], [28, 0, 119], [28, 0, 136], [34, 0, 151],[36, 1, 165], [37, 0, 176], [37, 1, 187], [36, 0, 194], [37, 0, 202], [34, 0, 209], [31, 0, 217], [28, 1, 220], [25, 0, 224], [18, 1, 227], [16, 0, 229], [14, 0, 233], [10, 0, 237], [9, 0, 237], [7, 0, 240], [3, 0, 242], [0, 0, 244], [0, 0, 244], [2, 5, 244], [1, 8, 244], [0, 13, 242], [0, 18, 242], [2, 22, 239], [0, 28, 236], [0, 33, 236], [0, 37, 232], [0, 44, 229], [2, 49, 227], [0, 55, 220], [0, 60, 218], [0, 66, 214], [1, 73, 209], [0, 77, 205], [0, 84, 200], [0, 91, 194], [0, 96, 193], [0, 101, 189], [0, 106, 182], [0, 111, 177], [1, 118, 172], [0, 120, 165], [0, 122, 159], [0, 128, 153], [1, 131, 147], [1, 132, 140], [1, 135, 134], [0, 140, 131], [0, 145, 126], [0, 148, 124], [0, 152, 122], [0, 158, 118], [1, 162, 118], [1, 168, 116], [0, 172, 114], [0, 178, 113],[0, 182, 112], [0, 186, 111], [1, 188, 109], [2, 191, 107], [0, 194, 107], [1, 195, 108], [0, 198, 101], [1, 200, 99], [0, 204, 96], [1, 209, 97], [2, 211, 94], [1, 217, 90], [0, 220, 88], [0, 225, 81], [1, 228, 77], [1, 231, 71], [1, 232, 68], [0, 230, 60], [0, 230, 52], [0, 230, 43], [0, 230, 33], [0, 228, 21], [0, 228, 11], [2, 229, 0], [16, 229, 0], [28, 229, 0], [40, 230, 0], [56, 232, 0], [72, 232, 2], [84, 230, 1], [98, 231, 0],[111, 230, 0],[124, 230, 0], [137, 230, 1], [151, 228, 0], [162, 227, 0], [173, 229, 0], [186, 227, 0], [198, 224, 1], [211, 226, 0], [221, 221, 0], [227, 216, 0], [230, 210, 1], [237, 201, 1], [240, 193, 1],[242, 184, 0], [245, 173, 0], [248, 165, 1], [250, 155, 0], [251, 145, 0], [252, 136, 1], [254, 126, 1], [255, 115, 0], [255, 104, 3], [254, 95, 1], [255, 83, 1], [255, 72, 2], [255, 61, 0], [253, 49, 0], [255, 39, 2], [253, 28, 0], [255, 17, 4], [255, 8, 1], [254, 2, 1], [254, 0, 10], [255, 0, 14], [255, 0, 18], [251, 0, 24], [250, 0, 30], [250, 0, 30], [248, 0, 35], [246, 0, 41], [246, 0, 41], [242, 0, 40], [242, 0, 40], [240, 0, 45], [237, 0, 46], [233, 0, 45], [230, 1, 44], [226, 0, 42], [222, 0, 41], [218, 0, 39], [214, 0, 38], [206, 0, 36], [200, 1, 34], [195, 0, 32], [189, 0, 30], [185, 0, 31], [177, 0, 28], [169, 0, 26], [162, 0, 24], [152, 0, 23],[144, 1, 21], [136, 1, 18], [128, 1, 20], [121, 0, 19], [111, 0, 16], [104, 0, 14], [96, 0, 12], [88, 1, 10], [83, 0, 12], [73, 0, 9], [67, 0, 9], [62, 1, 9], [57, 0, 7], [51, 0, 7], [46, 0, 5], [42, 0, 4], [39, 0, 5], [33, 1, 4], [30, 0, 4],[25, 0, 3], [25, 0, 3], [22, 0, 2],[21, 0, 1], [16, 0, 0], [15, 1, 1], [14, 0, 0], [12, 0, 0], [9, 0, 1], [9, 0, 1], [8, 0, 0]])
colormap_nm = np.linspace(400,700,len(colormap_rgb)) # courtesy of https://stackoverflow.com/questions/71977306/how-to-convert-rgb-to-wavelength-in-python
def wavelength_to_rgb(wavlgth): return np.asarray(colormap_rgb[min(range(len(colormap_nm)), key=lambda i: abs(colormap_nm[i]-wavlgth))]) # wavelength_to_rgb
[docs]class Wavelengths: # vectorized version of wavelength_to_rgb # Wavelengths
[docs] def __init__(self, wavs:"List[int]"): # Wavelengths
"""Container for wavelengths in nm. Pretty much only used to convert to rgb in vectorized way""" # Wavelengths
self.wavs = np.array(wavs | cli.deref()) # Wavelengths
def _toRgb(self): # Wavelengths
wavs = self.wavs; single = len(wavs.shape) == 0 # Wavelengths
if single: wavs = wavs[None] # Wavelengths
ans = colormap_rgb[abs(colormap_nm[None]-wavs[:,None]) | cli.toArgmin().all()] # Wavelengths
return ans[0] if single else ans # Wavelengths
def __repr__(self): return f"<Wavelengths #={len(self.wavs)} min={round(self.wavs.min(), 2)}nm max={round(self.wavs.max(), 2)}nm std={round(self.wavs.std(), 5)}nm>" # Wavelengths
pi = 3.141592653589793; inf = float("inf"); settings = k1lib.Settings(); k1lib.settings.add("kop", settings, "from k1lib.kop module, for optics-related tools"); # Wavelengths
settings.add("display", k1lib.Settings(), "display settings").add("consts", k1lib.Settings(), "magic constants throughout the sim. By default works pretty well, but you can tweak these if you need unrealistic setups, like super big focal length, etc") # Wavelengths
settings.display.add("drawable", k1lib.Settings() # Wavelengths
.add("axes", True, "whether to add x and y axes to the sketch") # Wavelengths
.add("maxWh", 800, "when drawing a sketch, it will be rescaled so that the maximum of width and height of the final image is this number. Increase to make the sketch bigger") # Wavelengths
.add("grid", True, "whether to add grid lines to the sketch") # Wavelengths
.add("gridColor", (255, 255, 255)), "generic draw settings") # Wavelengths
[docs]class Drawable: # Drawable
"""Base class to do common drawing routines""" # Drawable
[docs] def bounds(self): # Drawable
"""Should return (x1, y1, x2, y2). Base classes should always implement this method""" # Drawable
return NotImplemented # Drawable
def _draw(self, config): return NotImplemented # Drawable
def _preprocess(self): # Drawable
xmin, ymin, xmax, ymax = bounds = self.bounds(); dx = xmax-xmin; dy = ymax-ymin; s = settings.display.drawable.maxWh/max(dx, dy) # Drawable
c = {"depth": 0, "bounds": bounds}; config = {**self.config, **c} if isinstance(self, Drawables) else c # Drawable
p5.newSketch(dx, dy, pad=20, scale=s, xoff=xmin, yoff=ymin); p5.background(222); p5.textSize(12/s) # Drawable
xm = xmin-20/s; ym = ymin-20/s # xmin and ymin without padding # Drawable
try: xticks = k1lib.ticks(xmin, xmax); bsx = max(int(len(xticks)/(dx*s/30)), 1) # this is for when the width or height is too small to draw, then this will only draw a tick for every bsx ticks available # Drawable
except: xticks = [xmin]; bsx = 1 # if xmin == xmax # Drawable
try: yticks = k1lib.ticks(ymin, ymax); bsy = max(int(len(yticks)/(dy*s/30)), 1) # Drawable
except: yticks = [ymin]; bsy = 1 # Drawable
if settings.display.drawable.grid: # Drawable
with p5.context(): # Drawable
p5.stroke(*settings.display.drawable.gridColor) # Drawable
for x in xticks | cli.batched(bsx, True) | cli.item().all(): p5.line(x, ymin-20/s, x, ymax+20/s); # Drawable
for y in yticks | cli.batched(bsy, True) | cli.item().all(): p5.line(xmin-20/s, y, xmax+20/s, y); # Drawable
if settings.display.drawable.axes: # Drawable
for x in xticks | cli.batched(bsx, True) | cli.item().all(): p5.line(x, ymin, x, ym); p5.text(f"{x}", x+5/s, ym+5/s) # Drawable
for y in yticks | cli.batched(bsy, True) | cli.item().all(): p5.line(xmin, y, xm, y); p5.text(f"{y}", xm+5/s, y+5/s) # Drawable
self._draw(config) # Drawable
def _repr_extra(self): return "" # Drawable
def _repr_html_(self): self._preprocess(); r = html.escape(f"{self}"); return f"{r}{self._repr_extra()}<br>{p5.svg()}" # Drawable
[docs] def img(self) -> "PIL.Image.Image": # Drawable
"""Returns PIL image""" # Drawable
self._preprocess(); return p5.img() # Drawable
[docs] def svg(self) -> str: # Drawable
"""Returns svg string""" # Drawable
self._preprocess(); return p5.svg() # Drawable
[docs]class RandomPoints(Drawable): # displays random points # RandomPoints
[docs] def __init__(self, data): # RandomPoints
"""Container to place random points to be graphed either alone or inside a :class:`Drawables`""" # RandomPoints
self.data = data | cli.deref() | cli.toNdArray() # RandomPoints
def _draw(self, config): [p5.ellipse(x,y,0.5) for x,y in self.data]; return p5.svg() # RandomPoints
[docs] def bounds(self): return min(self.data[0]), min(self.data[1]), max(self.data[0]), max(self.data[1]) # RandomPoints
def __copy__(self): return RandomPoints(np.copy(self.data)) # RandomPoints
[docs]class Drawables(Drawable): # a container to just draw everything out # Drawables
[docs] def __init__(self, drawables:"List[Drawable]", config=None): # Drawables
"""A container with multiple :class:`Drawable`s, to draw out everything possible within a single coordinate frame""" # Drawables
self.drawables = drawables; self.config = config or {} # Drawables
def __copy__(self): return Drawables([copy.copy(d) for d in self.drawables]) # Drawables
[docs] def bounds(self): return self.drawables | cli.op().bounds().all() | cli.T() | cli.toMin() + cli.toMin() + cli.toMax() + cli.toMax() | cli.deref() # Drawables
def _draw(self, config): [d._draw({**config, **self.config}) for d in self.drawables]; return p5.svg() # Drawables
def __repr__(self): return f"<Drawables #n={len(self.drawables)}>" # Drawables
settings.add("colorD", {"rainbow": [400, 650], "red": [620, 750], "orange": [590, 620], "yellow": [570, 590], "green": [495, 570], "blue": [450, 495], "purple": [400, 450]}, "color wavelength ranges to be used in constructing Rays") # Drawables
def color_interpreter(col, N): # color_interpreter
if isinstance(col, str): # color_interpreter
if col not in settings.colorD: raise Exception(f"Does not understand the color '{col}'. Current options are: {', '.join(settings.colorD.keys())}. You can set your own custom colors using `settings.kop.colorD`") # color_interpreter
return np.linspace(*settings.colorD[col], N) # color_interpreter
if isinstance(col, (int, float, np.number)): return np.ones(N)*col # color_interpreter
if isinstance(col, (list, tuple)) and len(col) == 2: return np.linspace(*col, N) # color_interpreter
raise Exception(f"Does not understand the color/wavelength '{col}'. It can be a string like 'red', a specific wavelength in nm (450), or a list/tuple with 2 floats specifying start and end wavelengths") # color_interpreter
settings.display.add("rays", k1lib.Settings().add("showOrigin", True, "whether to add a small red dot to the beginning of the mean (x,y) of a Rays or not").add("infLength", 100, "length in mm to display Rays if their length is infinite"), "display settings of kop.Rays") # color_interpreter
[docs]class RaysFocus(Drawable): # RaysFocus
def __init__(self, rays:"Rays"): # RaysFocus
r = rays; x1 = r.data[:,0][None]; x2 = r.data[:,0][:,None]; y1 = r.data[:,1][None]; y2 = r.data[:,1][:,None]; t1 = r.data[:,2][None]; t2 = r.data[:,2][:,None] # RaysFocus
with k1lib.ignoreWarnings(): det = 1/(-np.cos(t1)*np.sin(t2)+np.sin(t1)*np.cos(t2)); self.rays = rays # RaysFocus
invalids = det==inf; det[invalids] = 0; a = det*-np.sin(t2); b = det*np.cos(t2); c = det*-np.sin(t1); d = det*np.cos(t1) # RaysFocus
alpha = a*(x2-x1) + b*(y2-y1); self.xs = (x1+alpha*np.cos(t1))[~invalids]; self.ys = (y1+alpha*np.sin(t1))[~invalids]; self.xy = np.stack([self.xs, self.ys]) # RaysFocus
xmin, ymin, xmax, ymax = self.bounds(); self.xmin = xmin; self.xmax = xmax; self.ymin = ymin; self.ymax = ymax; self.xmedian = np.median(self.xs); self.ymedian = np.median(self.ys) # RaysFocus
self.dx = xmax - xmin; self.dy = ymax - ymin; self.x = (xmin+xmax)/2; self.y = (ymin+ymax)/2; self.xstd = self.xy[0].std(); self.ystd = self.xy[1].std() # RaysFocus
[docs] def bounds(self): return self.xy[0].min(), self.xy[1].min(), self.xy[0].max(), self.xy[1].max() # RaysFocus
def _draw(self, config): # RaysFocus
bounds = config["bounds"]; pointSize = math.sqrt((bounds[2]-bounds[0])*(bounds[3]-bounds[1]))/100 # RaysFocus
for x, y in self.xy.T: p5.ellipse(x, y, pointSize) # RaysFocus
def __repr__(self): # RaysFocus
xdigits = max(int(-math.log10(self.dx))+2, 0); ydigits = max(int(-math.log10(self.dy))+2, 0) # RaysFocus
return f"<RaysFocus N={len(self.xy[0])} xrange=({round(self.xmin, xdigits)}, {round(self.xmedian, xdigits)}, {round(self.xmax, xdigits)}) xstd={round(self.xstd, 6)} yrange=({round(self.ymin, ydigits)}, {round(self.ymedian, ydigits)}, {round(self.ymax, ydigits)}) ystd={round(self.ystd, 6)}>" # RaysFocus
[docs]class Rays(Drawable): # Rays
[docs] def __init__(self, data, prevRays:"Rays"=None, ogSurface:"Surface"=None): # Rays
"""Represents a bunch of rays, stored in a numpy array so that all
operations are fast
`data` format: array of shape (N, 6)
- 0) x coordinate of starting point
- 1) y coordinate of starting point
- 2) angle counter clockwise from positive x axis
- 3) length (inf for forever)
- 4) wavelength in nm
- 5) #transforms
- 6) power in Watts, defaulted to 1W
- 7) intersected?: used for outgoing rays exiting out of a surface. If surface._cast(prevRay) shows that it does intersect, then add that to the outgoing rays
The #transforms is a little complicated:
- Original ray is 0, then each time a glass/mirror surface does something interesting,
the outgoing ray is incremented by 1. If it doesn't touch the optic element, then
it keeps the same number as before
:param prevRays: a reference to the previous Rays object, in order to limit the length of all previous rays!
:param ogSurface: the surface that generates this Rays""" # Rays
self.data = data; self.prevRays = prevRays; self.ogSurface = ogSurface # Rays
def __copy__(self): return Rays(np.copy(self.data), None if self.prevRays is None else copy.copy(self.prevRays), self.ogSurface) # Rays
@staticmethod # Rays
def _genElse(array, color=700, power=1): N = len(array); return Rays(array | cli.insertColumn([inf]*N, False) | cli.insertColumn(color_interpreter(color, N), False) | cli.insertColumn([0]*N, False) | cli.insertColumn([power]*N, False) | cli.insertColumn([0]*N, False)) # Rays
[docs] @staticmethod # Rays
def parallel(x=0, y=0, theta=0, N=10, height=10, color=700, power=1): # Rays
"""Creates parallel rays starting from a particular point with a particular angle
:param theta: angle from positive x axis counterclockwise
:param N: how many rays to create total?
:param height: the height of the parallel rays if it were to have no angle""" # Rays
a = theta + pi/2; b = [x+math.cos(a)*height/2, y+math.sin(a)*height/2] # Rays
a = theta - pi/2; c = [x+math.cos(a)*height/2, y+math.sin(a)*height/2] # Rays
return Rays._genElse([np.linspace(b[0], c[0], N), np.linspace(b[1], c[1], N)] | cli.toNdArray() | cli.T() | cli.insertColumn([theta]*N, False), color=color, power=power) # Rays
[docs] @staticmethod # Rays
def parallelToBounds(x=0, y=0, bounds=None, N=10, color=700, power=1, coverage=0.9, angleOffset=0): # Rays
"""Creates parallel rays starting from a particular point to the center of some bounds.
The bounds should have the format (xmin, ymin, xmax, ymax). :class:`Surface`s, :class:`OpticElement`s
all have the .bounds() method, so you can use them
See also: :meth:`parallel`
:param coverage: number from 0 to 1. 1 meaning the parallel lines should cover the entire (rotated)
height of the bounds, 0.5 means it covers only half of the height, and so on. Think of this as %height""" # Rays
if bounds is None: raise Exception(".bounds parameter is required") # Rays
x1, y1, x2, y2 = bounds; xs = np.array([x1, x2, x2, x1]); ys = np.array([y1, y1, y2, y2]); angles = np.arctan2(ys-y, xs-x)%(2*pi) # Rays
amin = angles.argmin(); x1 = xs[amin]; y1 = ys[amin]; amax = angles.argmax(); x2 = xs[amax]; y2 = ys[amax] # Rays
dx = x2-x1; mx = (x1+x2)/2; dy = y2-y1; my = (y1+y2)/2; angle = math.atan2(my-y, mx-x) # Rays
h = math.sin(angle+pi-math.atan2(y2-y1, x2-x1))*math.sqrt(dx*dx+dy*dy) # Rays
return Rays.parallel(x, y, angle+angleOffset, N, h*coverage, color, power) # Rays
[docs] @staticmethod # Rays
def pointToBounds(x=0, y=0, bounds=None, N=10, color=700, power=1, coverage=0.9, angleOffset=0): # Rays
"""Creates rays starting from a particular point to the bounds.
See also: :meth:`parallelToBounds`
:param coverage: instead of %height like :meth:`parallelToBounds`, this is the anglular coverage,
1 for covering the entire object.
:param angleOffset: offset to the angle from the starting point to the center of the bounds""" # Rays
if bounds is None: raise Exception(".bounds parameter is required") # Rays
x1, y1, x2, y2 = bounds; xs = np.array([x1, x2, x2, x1]); ys = np.array([y1, y1, y2, y2]); angles = np.arctan2(ys-y, xs-x)#%(2*pi) # Rays
amin = angles.argmin(); x1 = xs[amin]; y1 = ys[amin]; amax = angles.argmax(); x2 = xs[amax]; y2 = ys[amax] # Rays
startAngle = math.atan2(y1-y, x1-x) % (2*pi); endAngle = math.atan2(y2-y, x2-x) % (2*pi) # Rays
if endAngle < startAngle: endAngle += 2*pi # Rays
angleDiff = (endAngle - startAngle)*(1-coverage) # Rays
return Rays._genElse(np.array([[x]*N, [y]*N, np.linspace(startAngle+angleDiff/2+angleOffset, endAngle-angleDiff/2+angleOffset, N)]).T, color, power) # Rays
[docs] def bounds(self): # returns bounds for plotting by p5 # Rays
with k1lib.ignoreWarnings(): l = np.copy(self.data[:,3]); l[l==inf] = settings.display.rays.infLength # Rays
data = self.data; a = data[:,0] + np.cos(data[:,2])*l; b = data[:,1] + np.sin(data[:,2])*l # Rays
xmin = min(data[:,0].min(), a.min()); xmax = max(data[:,0].max(), a.max()) # Rays
ymin = min(data[:,1].min(), b.min()); ymax = max(data[:,1].max(), b.max()) # Rays
return xmin, ymin, xmax, ymax # Rays
[docs] def hitbox(self, spread=2, maxWH=400, mode="simple", intersectOnly=True) -> "PIL.Image.Image": # Rays
"""Visualizes the Ray's starting point, with colors and whatnot.
Example::
r = Rays.parallel(theta=pi/8)
r.hitbox()
:param spread: will color this much nearby pixels
:param maxWH: size of the returned image. Either width or height will be around this big, whichever is bigger
:param mode: several plotting modes, each with its own performance and accuracy characteristics
:param intersectOnly: if True (default), will only plot the hitbox of the rays that actually hits
the Surface that generates the ray and not some previous surface
""" # Rays
if mode == "simple": # Rays
if self.data[:,7].sum()<0.5: raise Exception("No ray intersected the Surface that generates this Rays") # Rays
data = self.data[self.data[:,7]>0.5][:,[0, 1, 4, 6]]; xy = data[:,:2] # Rays
xmin = min(xy[:,0]); xmax = max(xy[:,0]); xmid = (xmin+xmax)/2; dx = xmax-xmin # Rays
ymin = min(xy[:,1]); ymax = max(xy[:,1]); ymid = (ymin+ymax)/2; dy = ymax-ymin # Rays
N = len(data); dist = np.sum((xy[None] - xy[:,None])**2, 2)**0.5; dist[range(N), range(N)] = inf; dm = dist.min() # Rays
if maxWH is None: maxWH = min(max(dx/dm, dy/dm), 500); maxWH # number of pixels of the longest length. This is calculated automatically if user does not give a particular size # Rays
pixelSize = max(dx, dy)/maxWH # how much real distance is between 2 adjacent pixels? Ideally exactly the same as dm # Rays
pixelH = int(max(dy/pixelSize*1.1, 20)); pixelW = int(max(dx/pixelSize*1.1, 20)) # Rays
im = np.ones((pixelH, pixelW, 3))*200; rgb = Wavelengths(data[:,2]) | cli.toRgb() # Rays
ys = (-(xy[:,1]-ymid)/pixelSize+pixelH/2).astype(int); xs = ((xy[:,0]-xmid)/pixelSize+pixelW/2).astype(int) # Rays
for i in np.arange(1+spread*2)-spread: # Rays
for j in np.arange(1+spread*2)-spread: im[ys+i, xs+j] = rgb # Rays
return im | cli.toImg() # Rays
else: raise Exception(f"Currently doesn't support mode '{mode}'") # Rays
[docs] def focus(self) -> "RaysFocus": return RaysFocus(self) # Rays
def __add__(self, rays): # Rays
if not isinstance(rays, Rays): raise Exception(f"Can only add Rays to another Rays! The other value has type {type(rays)}") # Rays
return Rays(np.concatenate([self.data, rays.data])) # Rays
def __repr__(self): return f"<Rays N={len(self.data)} avgTheta={self.data | cli.cut(2) | cli.toMean() | cli.op()/pi*180 | cli.aS(round, 2)}deg intersectedSurface?={self.data[:,7].sum()>0.5} ogSurface={self.ogSurface}>" # Rays
def _draw(self, config={}, ignored=None): # Rays
with k1lib.ignoreWarnings(): data = self.data; l = np.copy(self.data[:,3]); l[l==inf] = settings.display.rays.infLength; bounds = config["bounds"] # Rays
a = data[:,0] + np.cos(data[:,2])*l; b = data[:,1] + np.sin(data[:,2])*l; xmin, ymin, xmax, ymax = self.bounds() # Rays
with p5.context(): # Rays
for params in [data[:,0], data[:,1], a, b, data[:,4], ignored if ignored is not None else [False]*len(data)] | cli.T() | ~cli.filt("x", 5): # Rays
p5.stroke(*wavelength_to_rgb(params[4])); p5.line(*params[:4]) # Rays
if settings.display.rays.showOrigin: p5.fill(255, 0, 0); p5.noStroke(); p5.ellipse(*data[:, :2] | cli.T() | cli.toMean().all(), math.sqrt((bounds[2]-bounds[0])*(bounds[3]-bounds[1]))/100) # Rays
r = html.escape(f"{self}"); return p5.svg() # Rays
gps = {"BK7": (1.03961212, 0.231792344, 1.01046945, 0.00600069867, 0.0200179144, 103.560653), # borosilicate crown glass # Rays
"sapphire": (1.4313493, 0.65054713, 5.3414021, 0.0052799261, 0.0142382647, 325.017834), # Rays
"fusedSilica": (0.6961663, 0.4079426, 0.8974794, 0.00467914826, 0.0135120631, 97.9340025), # Rays
"MgF": (0.48755108, 0.39875031, 2.3120353, 0.001882178, 0.008951888, 566.13559), # Rays
"air": (0,0,0,0,0,0)} # Rays
settings.add("gps", gps, "All builtin glass parameters of the system. All have 6 floats, for B1,B2,B3,C1,C2,C3 parameters used in the sellmeier equation: https://en.wikipedia.org/wiki/Sellmeier_equation") # Rays
_invGlassParams = [None, 0] # Rays
def invGlassParams(): # invGlassParams
global gps, _invGlassParams # invGlassParams
if time.time() - _invGlassParams[1] > 1: gps = {**settings.gps, **gps}; settings.gps = gps; _invGlassParams = [{v:k for k, v in gps.items()}, time.time()] # invGlassParams
return _invGlassParams[0] # invGlassParams
[docs]def sellmeier(wavelengths:"np.ndarray", glassParam): # sellmeier
"""Runs the sellmeier equation to get the index of refraction for multiple wavelengths
for a particular glass material""" # sellmeier
if isinstance(glassParam, (int, float, np.number)): return np.array([glassParam]*len(wavelengths)) # sellmeier
else: # sellmeier
la2 = (wavelengths/1000)**2; b1, b2, b3, c1, c2, c3 = glassParam # sellmeier
return (1 + b1*la2/(la2-c1) + b2*la2/(la2-c2) + b3*la2/(la2-c3))**0.5 # sellmeier
sellmeier(np.linspace(400, 700, 10), gps["BK7"]) # sellmeier
_surface_autoInc = k1lib.AutoIncrement(prefix="_surface_") # sellmeier
def snell(rtheta, ln, n1, n2): # this takes into account whether we've encountered total internal reflection or not # snell
a = n1*np.sin(rtheta-ln)/n2; b = (ln + np.arcsin(a)).clearNan(); inRange = (-1 <= a) * (a <= 1) # snell
f = -2*np.cos(rtheta-ln); c = np.arctan2(np.sin(rtheta)+f*np.sin(ln), np.cos(rtheta)+f*np.cos(ln)) # snell
return inRange * b + (1-inRange) * c # snell
settings.display.add("surface", k1lib.Settings().add("showIndex", True, "whether to show the index of the Surface in an OpticSystem or not"), "display settings for kop.Surface class") # snell
settings.consts.add("inchForward", 1e-6, "after new Rays have been built, inch forward the origin of the new Rays by a this tiny amount so that it 'clears' the last Surface") # snell
[docs]class Surface(Drawable): # Surface
[docs] def __init__(self, gp1=None, gp2=None, mode="glass", capture=False, angleStd=0, opticElement=None): # Surface
"""Represents a geometry, like line, arc, or parabola. Rays can interact with a
surface, resulting in new Rays. An optic element can have multiple of these surfaces and
react differently to incoming rays.
The convension is that gp1 is the environment's glassParam, while gp2 is the Surface's
glassParam. A surface needs both in order to calculate the deflection angles.
These modes are available:
- glass: transparent (or semi-transparent) surface, bends incoming rays
- mirror: perfectly reflective (or allow some percentage of light to pass through)
- sensor: absorbes all incoming rays, can retrieve the image that's absorbed
- opaque: reflects incoming rays in a random direction outward
:param gp1: glass params 1, usually this is of the environment (air)
:param gp2: glass param 2, usually this is of the object (borosilicate glass, sapphire, water)
:param mode: explained above
:param capture: if True, captures all rays that comes in through .cast() function. Captures
will stay the same across copy.copy() instances to save perf! The captures will be
available at ``.captures``
:param angleStd: if specified, when casting outgoing rays, will add this much standard deviation
in the rays' angle, useful to quickly simulate a bumpy piece of glass
:param opticElement: the OpticElement that this Surface is a part of""" # Surface
self.idx = _surface_autoInc() # Surface
self.gp1 = (self.gp1 or gp1) if hasattr(self, "gp1") else gp1 # Surface
self.gp2 = (self.gp2 or gp2) if hasattr(self, "gp2") else gp2 # Surface
self.mode = (self.mode or mode) if hasattr(self, "mode") else mode # Surface
self.capture = capture; self.captures = []; self.angleStd = angleStd; self.opticElement = opticElement # Surface
def __copy__(self): res = super().__new__(type(self)); res.idx = self.idx; res.gp1 = self.gp1; res.gp2 = self.gp2; res.mode = self.mode; res.capture = self.capture; res.captures = self.captures; res.angleStd = self.angleStd; res.opticElement = self.opticElement; return res # Surface
def _cast(self, rays:"Rays"): # Surface
"""Expected to be implemented by all subclasses.
Given N rays, should return these arrays:
- (N,) bool array: whether the rays intersect the surface or not
- (N,) float array: length of the previous ray if intersect. Else can be any number, but please no whacky values like inf or nan
- (N,) float array: angle of vector normal to the surface. Say it points up. Then top area is considered gp1, bottom area is gp2. If don't intersect then can be any number
This should be enough to totally and completely describe the surface properties, and the refraction
code only needs to be written once""" # Surface
return NotImplemented # Surface
def _cast2(self, rays): # Surface
intersected, beta, ln = self._cast(rays) # Surface
beta.clearNan(); ln.clearNan(); beta[beta==inf] = 0; ln[ln==inf] = 0 # Surface
return intersected, beta, ln # Surface
[docs] def castS(self, rays) -> "Rays": # Surface
"""This will truncate the lengths of the incoming Rays object, and return
a new rays, as well as the lengths of the old rays, for easy comparison. If the
rays don't interact with the lens at all, then just copy the old rays over, and the
length is the same.""" # Surface
if self.capture: self.captures.append([copy.copy(rays)]) # Surface
if self.mode == "glass": # Surface
with k1lib.ignoreWarnings(): # Surface
if self.gp1 is None: raise Exception(f"gp1 not specified in {self}!") # Surface
if self.gp2 is None: raise Exception(f"gp2 not specified in {self}!") # Surface
r = rays; intersected, beta, ln = self._cast2(rays); rtheta = r.data[:,2]; N = len(r.data) # Surface
n1 = sellmeier(r.data[:,4], self.gp1); n2 = sellmeier(r.data[:,4], self.gp2) # index of refractions # Surface
lns = np.stack([ln, pi+ln]).T; ln_version = abs((rtheta[:,None]-lns+pi)%(2*pi)-pi) | cli.toArgmin().all() # Surface
d1to2 = ln_version # whether the ray is going from gp1 (glass param 1, top of line) to gp2. Completely coincidentally, this is exactly equal to ln_version (line normal version, either as-is, or flipped 180) # Surface
chosenLn = lns[range(N), ln_version] # actual normal that it's going to use # Surface
t1 = snell(rtheta, chosenLn, n1, n2); t2 = snell(rtheta, chosenLn, n2, n1) # Surface
rthetap = d1to2*t1 + (1-d1to2)*t2; intersections = [r.data[:,0]+np.cos(r.data[:,2])*(beta+1e-8), r.data[:,1]+np.sin(r.data[:,2])*(beta+1e-8)] | cli.toNdArray() | cli.T() # Surface
elif self.mode == "mirror": # Surface
with k1lib.ignoreWarnings(): # Surface
r = rays; intersected, beta, ln = self._cast2(rays); rtheta = r.data[:,2]; N = len(r.data); lns = np.stack([ln, pi+ln]).T # Surface
angleDiff = (rtheta[:,None]-lns+pi)%(2*pi)-pi; ln_version = abs(angleDiff) | cli.toArgmin().all() # Surface
f = 2*np.cos(angleDiff[range(N), ln_version]) # factor to multiply the normal vector by # Surface
chosenLn = lns[range(N), 1-ln_version] # actual normal that it's going to use # Surface
rthetap = np.arctan2(np.sin(rtheta)+f*np.sin(chosenLn), np.cos(rtheta)+f*np.cos(chosenLn)) # Surface
intersections = [r.data[:,0]+np.cos(r.data[:,2])*(beta-1e-8), r.data[:,1]+np.sin(r.data[:,2])*(beta-1e-8)] | cli.toNdArray() | cli.T() # Surface
else: raise Exception(f"Doesn't support mode '{self.mode}' yet") # Surface
with k1lib.ignoreWarnings(): # make sure intersected, intersections, rthetap variables are defined upon reaching this part # Surface
if abs(self.angleStd) > 1e-11: rthetap = rthetap + np.random.randn(rthetap.shape)*self.angleStd # adding variations # Surface
intersections[:,0] += np.cos(rthetap)*settings.consts.inchForward; intersections[:,1] += np.sin(rthetap)*settings.consts.inchForward # inch forward a little to clear the last surface # Surface
outr = np.concatenate([intersections | cli.insertColumn(rthetap, False), r.data[:,3:]], 1) # has not corrected for whether it has intersected the line # Surface
outr[~intersected] = np.copy(r.data[~intersected]); outr[:,3] = float("inf"); outr[intersected, 5] = r.data[intersected, 5] + 1; outr[~intersected, 5] = r.data[~intersected, 5]; outr[:,7] = intersected; outr = Rays(outr, r, self) # Surface
# tries to limit the length of previous rays # Surface
r.data[:,3] = beta * intersected + (r.data[:,3] * (1 - intersected)).clearNan(); newR = r; r = r.prevRays # Surface
while r is not None: # for the rays that does not have its depth incremented, set the prevLength to this lengths # Surface
idxs = np.abs(newR.data[:,5] - r.data[:,5]) < 1e-9; r.data[idxs,3] = newR.data[idxs,3]; newR = r; r = r.prevRays # Surface
if self.capture: self.captures[-1].append(copy.copy(outr)) # Surface
return outr # Surface
[docs] def cast(self, rays, **kwargs) -> "RaysPath": # convenience function, not core mechanism # Surface
return OpticSystem().add(self).cast(rays, **kwargs) # Surface
def _draw(self, config): # Surface
if settings.display.surface.showIndex and "surfaceIdx2Id" in config: # Surface
with p5.context(): xmin, ymin, xmax, ymax = self.bounds(); p5.fill(100); p5.noStroke(); p5.text(config["surfaceIdx2Id"][self.idx], (xmin+xmax)/2, (ymin+ymax)/2) # Surface
def _repr(self): # helper sections at the end of every surface # Surface
gp1 = f" gp1=None" if self.gp1 is None else (f" gp1='{invGlassParams()[self.gp1]}'" if self.gp1 in invGlassParams() else "") # Surface
gp2 = f" gp2=None" if self.gp2 is None else (f" gp2='{invGlassParams()[self.gp2]}'" if self.gp2 in invGlassParams() else "") # Surface
oe = "" if self.opticElement is None else f' opticElement={self.opticElement}' # Surface
return f"{gp1}{gp2}{' capture=True' if self.capture else ''}{oe}" # Surface
def ray_intersection(line:"list[float]", rays:Rays): # ray_intersection
with k1lib.ignoreWarnings(): # ray_intersection
lx1, ly1, lx2, ly2 = line; dlx = lx2 - lx1; dly = ly2 - ly1 # ray_intersection
if abs(dlx) < 1e-11: dlx += 1e-10 # surprisingly, it still works for such tiny numbers! # ray_intersection
R = dly/dlx; rx = rays.data[:,0]; ry = rays.data[:,1]; rtheta = rays.data[:,2] # ray_intersection
beta = (ry-ly1+R*(lx1-rx))/(np.cos(rtheta)*R-np.sin(rtheta)) # vector length of ray # ray_intersection
alpha = (rx+beta*np.cos(rtheta)-lx1)/dlx; return alpha, beta # vector length in line surface # ray_intersection
[docs]class LineSurface(Surface): # LineSurface
[docs] def __init__(self, x1, y1, x2, y2, **kwargs): # LineSurface
"""Surface defined using a straight line.
If x1,y1 is on the left, x2,y2 on the right, then gp1 (glass param) is on top, gp2 is
at the bottom.""" # LineSurface
if not hasattr(self, "_superAlreadyInit"): super().__init__(**kwargs) # LineSurface
self.x1 = x1; self.y1 = y1; self.x2 = x2; self.y2 = y2 # LineSurface
def __copy__(self): res = super().__copy__(); res._superAlreadyInit = True; res.__init__(self.x1, self.y1, self.x2, self.y2); return res # LineSurface
def _cast(self, rays): # LineSurface
l = self; alpha, beta = ray_intersection([l.x1, l.y1, l.x2, l.y2], rays) # LineSurface
ln = (math.atan2(l.y2 - l.y1, l.x2 - l.x1) + pi/2) % (2*pi) # line normal angle # LineSurface
intersected = (0 <= alpha) * (alpha <= 1) * (beta > 0); return intersected, beta, np.ones(len(rays.data))*ln # LineSurface
[docs] def bounds(self): return min(self.x1, self.x2), min(self.y1, self.y2), max(self.x1, self.x2), max(self.y1, self.y2) # LineSurface
def _draw(self, config): super()._draw(config); p5.line(self.x1, self.y1, self.x2, self.y2); return p5.svg() # LineSurface
def __repr__(self): return f"<Line(Surface) ({round(self.x1, 2)}, {round(self.y1, 2)}, {round(self.x2, 2)}, {round(self.y2, 2)}){self._repr()}>" # LineSurface
[docs]class ArcSurface(Surface): # ArcSurface
[docs] def __init__(self, x, y, r, startAngle, endAngle, **kwargs): # ArcSurface
"""Surface defined using a circular arc""" # ArcSurface
if not hasattr(self, "_superAlreadyInit"): super().__init__(**kwargs) # ArcSurface
self.x = x; self.y = y; self.r = r # ArcSurface
startAngle = startAngle%(2*pi); endAngle = endAngle-2*pi # ensuring startAngle is always positive, and endAngle is always greater than startAngle. This defines the angle sweeped counterclockwise from startAngle to endAngle # ArcSurface
while endAngle < startAngle: endAngle += 2*pi # ArcSurface
self.startAngle = startAngle; self.endAngle = endAngle # ArcSurface
[docs] @staticmethod # ArcSurface
def from2Points(x1, y1, x2, y2, r, **kwargs) -> "ArcSurface": # ArcSurface
"""Creates an :class:`ArcSurface` from 2 points and a radius.
Example::
kop.ArcSurface.from2Points(1, 2, 3, 4, 10)
""" # ArcSurface
if r < 0: x1,y1,x2,y2 = x2,y2,x1,y1; r = -r # ArcSurface
mx = (x1+x2)/2; dx = x2-x1; my = (y1+y2)/2; dy = y2-y1; dist = math.sqrt(dx*dx+dy*dy); a = math.atan2(dy, dx)-pi/2; h2 = r*r-dist*dist/4 # ArcSurface
if h2 < 0: raise Exception(f"Arc radius {r} too small. No arc will pass through both of the specified points") # ArcSurface
h = math.sqrt(h2); cx = mx + h*math.cos(a); cy = my + h*math.sin(a); return ArcSurface(cx, cy, r, math.atan2(y2-cy, x2-cx), math.atan2(y1-cy, x1-cx), **kwargs) # ArcSurface
def __copy__(self): res = super().__copy__(); res._superAlreadyInit = True; res.__init__(self.x, self.y, self.r, self.startAngle, self.endAngle); return res # ArcSurface
def _cast(self, rays): # ArcSurface
rx, ry, rtheta = rays.data[:,:3].T; c = np.cos(rtheta); s = np.sin(rtheta); ax = self.x; ay = self.y; ar = self.r; N = len(rx) # ArcSurface
A = rx-ax; B = ry-ay; a = 1; b = 2*(A*c+B*s); c = A*A+B*B-ar*ar; delta = b*b-4*a*c; intersected = delta > 0 # quadratic formula # ArcSurface
with k1lib.ignoreWarnings(): delta = np.sqrt(delta).clearNan(); v1 = (-b-delta)/(2*a); v2 = (-b+delta)/(2*a) # ArcSurface
v1[v1<0]=inf; v1[~intersected]=inf; v2[v2<0]=inf; v2[~intersected]=inf # ArcSurface
vs = np.stack([v1,v2]).T; tx = rx[:,None]+vs*np.cos(rtheta[:,None]); ty = ry[:,None]+vs*np.sin(rtheta[:,None]) # ArcSurface
angles = np.arctan2(ty-ay, tx-ax); angles[angles<self.startAngle] += 2*pi; angles[angles<self.startAngle] += 2*pi; oldAngles = angles # ArcSurface
selected_version = (self.startAngle < angles) * (angles < self.endAngle) * (-inf < vs) * (vs < inf) | cli.toArgmax().all() # ArcSurface
angles = angles[range(N), selected_version]; intersected[~((self.startAngle < angles) * (angles < self.endAngle))] = False # ArcSurface
beta = vs[range(N), selected_version]; intersected[beta==inf] = False; beta[beta==inf] = 0; return intersected, beta, angles%(2*pi) # ArcSurface
[docs] def bounds(self): # step through n=50 points from startAngle to endAngle, then calc the bounds # ArcSurface
angles = np.linspace(self.startAngle, self.endAngle); a = self.x+self.r*np.cos(angles); b = self.y+self.r*np.sin(angles); return a.min(), b.min(), a.max(), b.max() # ArcSurface
def _draw(self, config): # ArcSurface
super()._draw(config); # ArcSurface
with p5.context(): p5.noFill(); p5.arc(self.x,self.y,self.r,self.startAngle,self.endAngle); return p5.svg() # ArcSurface
def __repr__(self): return f"<Arc(Surface) x={round(self.x, 2)} y={round(self.y, 2)} r={round(self.r, 2)} startAngle(deg)={round(self.startAngle/pi*180, 2)} endAngle(deg)={round(self.endAngle/pi*180, 2)}{self._repr()}>" # ArcSurface
def buildF(coeff): # compiles function of the passed in coefficients # buildF
s = " " # buildF
for i, c in enumerate(coeff): # buildF
if c != 0: s += f" + {c}*x**{i}" # buildF
s += " "; s = s.replace("*x**0 ", " ").replace("**1 ", " ").replace("**2 ", "*x ").strip(" +") # buildF
return eval(f"lambda x: {s if s else '0'}") # buildF
def buildDF(coeff): # compiles the derivative function of the passed in coefficients # buildDF
s = " " # buildDF
for i, c in enumerate(coeff[1:]): # buildDF
if c != 0: s += f" + {c*(i+1)}*x**{i}" # buildDF
s += " "; s; s = s.replace("*x**0 ", " ").replace("**1 ", " ").replace("**2 ", "*x ").strip(" +") # buildDF
return eval(f"lambda x: {s if s else '0'}") # buildDF
[docs]def polySolver(polys): # polySolver
"""Expected to take in a (N,F) array of N different polynomials with order F-1 and return 2 arrays:
- intersected (N,) bool array: whether there're any roots
- beta (N,) float array: the root of each polynomial that's >0 and is the minimum out of all roots
Right now, only supports vectorized solvers with order 2 polynomials. Any higher and it
falls back to using un-vectorized np.roots(), or other methods like it. This is separated out
so that it can be swapped out for a better implementation
Potential candidates for >2 order polynomials:
- scipy: fsolve, root, newton, brentq
""" # polySolver
order = polys.shape[1]-1; N = polys.shape[0] # polySolver
if order == 2: # polySolver
c,b,a = polys.T # polySolver
intersected = np.array([True]*N); delta = b*b-4*a*c; intersected[delta < 0] = False # polySolver
delta[delta < 0] = 0; delta = np.sqrt(delta); x1 = (-b-delta)/(2*a); x2 = (-b+delta)/(2*a); xs = np.stack([x1, x2]).T; xs[xs<0] = inf # polySolver
beta_version = xs | cli.toArgmin().all(); beta = xs[range(N), beta_version]; intersected[beta==inf] = False; # polySolver
beta = beta.clearNan(); beta[beta==inf] = 0; beta = (abs(a)<1e-8)*(-c/b).clearNan() + (abs(a)>=1e-8)*beta; intersected[beta<0] = False # polySolver
return intersected, beta # polySolver
else: raise Exception("Currently only support polynomials of order 2, aka a parabola. Why aren't higher dimensions supported? Cause I'm lazy. You can copy this function (polySolver), create your own robust version that supports higher dimensions and insert it into PolySurface.solver") # polySolver
[docs]class PolySurface(Surface): # PolySurface
[docs] def __init__(self, x=0, y=0, angle=0, scale=1, xmin=-10, xmax=10, coeff=(0,0,0.1), solver=polySolver, **kwargs): # PolySurface
"""Surface defined using a polynomial.
Let's say there's a parabola that you'd like to input into the simulation, say ``x^2/10``. There
are several dials and knobs that you can change to position your Surface right where you want it.
There are 2 frame of reference I'll be using. First is world frame, where all your elements sits in,
and second is poly frame, where your polynomial coefficients will define (x, y) in.
In the poly frame, the surface is defined as all x points in (xmin, xmax), together with all y points
calculated straight from your coefficients. So if ``coeff = [1, 2, 3]``, then the equation will be
``py = f(px) = 3px^2 + 2px + 1``. Now pairs of (px, py) points in poly frame is obtained.
Then (px, py) will go through several transformations to get to the world frame:
- Rotated by ``angle`` radians counterclockwise
- Scaled by ``scale``
- Translated by ``(x, y)``
Then to calculate reflections and whatnot, internally, all rays will be translated from world to poly
frame, then the intersection points are solved and then results will be translated from poly back
to world frame.
:param x: translation applied to surface
:param angle: angle to rotate the surface by
:param scale: scaling factor to scale the surface by
:param xmin: to define the domain of the given polynomial in poly frame
:param coeff: polynomial coefficients""" # PolySurface
if not hasattr(self, "_superAlreadyInit"): super().__init__(**kwargs) # PolySurface
if len(coeff) < 2: raise Exception(f"This only works with polynomial of order 2 or higher. Use the simpler :class:`LineSurface` for order 1!") # PolySurface
self.x = x; self.y = y; self.angle = angle; self.scale = scale; self.xmin = xmin; self.xmax = xmax; self.coeff = tuple(coeff); self.solver = solver # PolySurface
self.f = buildF(coeff); self.df = buildDF(coeff) # PolySurface
[docs] @staticmethod # PolySurface
def parabola(x, y, r=5, width=16, angle=0, **kwargs): return PolySurface(x, y, angle, 1, -width/2, width/2, (0, 0, 1/(2*r)), **kwargs) # PolySurface
[docs] @staticmethod # PolySurface
def parabolaFrom2Points(x1, y1, x2, y2, r=5, **kwargs): # PolySurface
dx = x2-x1; mx = (x1+x2)/2; dy = y2-y1; my = (y1+y2)/2 # PolySurface
dist = math.sqrt(dx*dx+dy*dy); a = -1/(2*r); angle = math.atan2(dy, dx); height = a*(dist/2)**2 # PolySurface
x = mx + height*math.cos(angle-pi/2); y = my + height*math.sin(angle-pi/2) # PolySurface
return PolySurface(x, y, angle, 1, -dist/2, dist/2, (0, 0, a), **kwargs) # PolySurface
def __copy__(self): res = super().__copy__(); res._superAlreadyInit = True; res.__init__(self.x, self.y, self.angle, self.scale, self.xmin, self.xmax, self.coeff, self.solver); return res # PolySurface
[docs] def world2poly(self, wxs, wys): c = math.cos(-self.angle); s = math.sin(-self.angle); xs = (wxs-self.x)/self.scale; ys = (wys-self.y)/self.scale; return xs*c-ys*s, xs*s+ys*c # PolySurface
[docs] def poly2world(self, pxs, pys): c = math.cos(self.angle); s = math.sin(self.angle); return self.x+self.scale*(pxs*c-pys*s), self.y+self.scale*(pxs*s+pys*c) # PolySurface
def _cast(self, rays): # PolySurface
rx, ry = self.world2poly(rays.data[:,0], rays.data[:,1]); rtheta = rays.data[:,2]-self.angle; N = len(rx); coeff = self.coeff # PolySurface
# this clusterfuck piece of code is trying to solve f(rx+beta*cos(rtheta))-(ry+beta*sin(rtheta))=0. I don't want to use numerical # PolySurface
# methods as those are not vectorized, so I have to rewrite the whole equation in terms of beta and then do the vectorized # PolySurface
# analytical quadratic formula later on. This code builds up that polynomial array where the variable is beta # PolySurface
a = np.cos(rtheta); b = rx; n = order = len(coeff)-1; polys = np.zeros((N, order+1)) # first column is x^0, last row is x^n # PolySurface
for ex in range(n+1): # trying to expand c0*(ax+b)^0 + c1*(ax+b)^1 + ... + cn*(ax+b)^n # PolySurface
for k in range(ex+1): polys[:,ex-k] += coeff[ex]*math.comb(ex, k)*a**(n-k)*b**k # PolySurface
polys[:,0] -= ry; polys[:,1] -= np.sin(rtheta) # then adding some correction term. Once polynomials are formed, then beta can be solved, and normal angles can be found # PolySurface
intersected, beta = self.solver(polys); tx = rx+beta*np.cos(rtheta); intersected[((tx < self.xmin) | (tx > self.xmax))] = False # PolySurface
tx = np.maximum(np.minimum(tx, self.xmax), self.xmin); angles = np.arctan(self.df(tx))+pi/2+self.angle; return intersected, beta*self.scale, angles%(2*pi) # PolySurface
[docs] def bounds(self): pxs = np.linspace(self.xmin, self.xmax); xs, ys = self.poly2world(pxs, self.f(pxs)); return min(xs), min(ys), max(xs), max(ys) # PolySurface
def _draw(self, config): # PolySurface
super()._draw(config); pxs = np.linspace(self.xmin, self.xmax); xs, ys = self.poly2world(pxs, self.f(pxs)); # PolySurface
[xs, ys] | cli.T() | cli.window(2) | cli.joinSt().all() | ~cli.apply(p5.line) | cli.ignore(); return p5.svg() # PolySurface
def __repr__(self): return f"<Poly(Surface) coeff={self.coeff} x={round(self.x, 2)} y={round(self.y, 2)} angle(deg)={round(self.angle/pi*180, 2)} scale={round(self.scale, 2)} xmin={round(self.xmin, 2)} xmax={round(self.xmax, 2)}{self._repr()}>" # PolySurface
[docs]class RaysPath(Drawable): # RaysPath
[docs] def __init__(self, rayss:"List[Rays]"): # RaysPath
"""Container for a list of :class:`Rays`, and have mechanism to
draw them out with a slight path correction. Expected to contain
output of a ray tracing session.""" # RaysPath
self._og = rayss; self.rayss = [copy.copy(rays) for rays in rayss] # RaysPath
def _draw(self, config): # RaysPath
oldRays = None # RaysPath
for rays in self.rayss: rays._draw(config) if oldRays is None else rays._draw(config, np.abs(rays.data[:,5] - oldRays.data[:,5]) < 1e-9); oldRays = rays # RaysPath
return p5.svg() # RaysPath
[docs] def bounds(self): return self.rayss | cli.op().bounds().all() | cli.T() | cli.toMin() + cli.toMin() + cli.toMax() + cli.toMax() | cli.deref() # RaysPath
def __copy__(self): return RaysPath([rays for rays in self._og]) # RaysPath
def __len__(self): return len(self.rayss) # RaysPath
def __getitem__(self, idx): return self.rayss[idx] # RaysPath
def __repr__(self): return f"<RaysPath #rays={len(self.rayss)}>" # RaysPath
[docs]class OpticElement(Drawable): # OpticElement
surfaces:"List[Surface]" # OpticElement
[docs] def __init__(self): # OpticElement
"""Represents a solid object with several Surfaces""" # OpticElement
pass # OpticElement
def __len__(self): return len(self.surfaces) # OpticElement
def __getitem__(self, idx): return self.surfaces[idx] # OpticElement
[docs] def cast(self, rays, *args, **kwargs): return OpticSystem().add(self).cast(rays, *args, **kwargs) # OpticElement
def signed_area(points): # signed_area
n = len(points); area = 0.0 # signed_area
for i in range(n): j = (i + 1) % n; area += (points[j][0] + points[i][0]) * (points[j][1] - points[i][1]) # signed_area
return area / 2.0 # signed_area
[docs]class Polygon(OpticElement): # Polygon
[docs] def __init__(self, points=None, **kw): # Polygon
"""Creates an OpticElement that has the shape of any polygon you want.
Example::
p = Polygon([[1, 2], [3, 4], [5, 6]])
This will create a prism with verticies at the specified points""" # Polygon
if points and signed_area(points) > 0: points = points | cli.reverse() | cli.deref() # Polygon
self.points = points; self.surfaces = ([*points, points[0]] | cli.window(2) | ~cli.apply(lambda p1, p2: LineSurface(p1[0], p1[1], p2[0], p2[1], **kw, opticElement=self)) | cli.deref()) if points is not None else None # Polygon
def __copy__(self): ans = Polygon(); ans.points = self.points; ans.surfaces = [copy.copy(s) for s in self.surfaces]; return ans # Polygon
[docs] def bounds(self): return self.surfaces | cli.op().bounds().all() | cli.T() | cli.toMin() + cli.toMin() + cli.toMax() + cli.toMax() | cli.deref() # Polygon
def _draw(self, config): [s._draw(config) for s in self.surfaces]; return p5.svg() # Polygon
def __repr__(self): return f"<Polygon(OpticElement) #points={len(self.points)} position=({self.points | cli.T() | cli.toMean().all() | cli.aS(round, 2).all() | cli.join(', ')})" # Polygon
[docs]class Lens(OpticElement): # Lens
[docs] def __init__(self, x=50, y=0, R1=300, R2=300, thickness=3, height=20, angle=0, surface="arc", **kw): # Lens
"""Represents a Lens with center at (x, y), and with 2 radiuses, R1 for left and R2 for right.
:param R1: radius of the left side, works for both arc and parabola mode
:param angle: angle offset of the Lens
:param surface: either "arc" or "parabola", which will make a lens of that geometry""" # Lens
self.x = x; self.y = y; self.R1 = R1; self.R2 = R2; self.thickness = thickness; self.height = height # Lens
dx1 = (-thickness/2 + R1*(1-math.cos(math.asin(height/2/R1)))) if R1 > 0 else -thickness/2 # Lens
dx2 = ( thickness/2 - R2*(1-math.cos(math.asin(height/2/R2)))) if R1 > 0 else thickness/2 # Lens
c = math.cos(angle); s = math.sin(angle); cx = x; cy = y; trans = lambda x,y: [cx+x*c-y*s, cy+x*s+y*c] # Lens
kw = {"gp1": None, "gp2": None, **kw}; fkw = dict(kw); fkw["gp1"] = kw["gp2"]; fkw["gp2"] = kw["gp1"] # Lens
if surface == "arc": meth = ArcSurface.from2Points # Lens
elif surface == "parabola": meth = PolySurface.parabolaFrom2Points # Lens
else: raise Exception(f"Does not recognize the surface '{surface}'. Only 'arc' and 'parabola' is valid right now") # Lens
self.surfaces = [ # Lens
LineSurface(*trans(dx1, height/2), *trans(dx2, height/2), **kw), # Lens
meth(*trans(dx1, -height/2), *trans(dx1, height/2), R1, **(kw if R1 > 0 else fkw)), # Lens
LineSurface(*trans(dx2, -height/2), *trans(dx1, -height/2), **kw), # Lens
meth(*trans(dx2, height/2), *trans(dx2, -height/2), R2, **(kw if R2 > 0 else fkw)) # Lens
] # Lens
def __copy__(self): ans = Lens(); ans.surfaces = [copy.copy(s) for s in self.surfaces]; return ans # Lens
[docs] def bounds(self): return self.surfaces | cli.op().bounds().all() | cli.T() | cli.toMin() + cli.toMin() + cli.toMax() + cli.toMax() | cli.deref() # Lens
def _draw(self, config): [s._draw(config) for s in self.surfaces]; return p5.svg() # Lens
def __repr__(self): return f"<Lens R1={self.R1} R2={self.R2} thickness={self.thickness} height={self.height}>" # Lens
[docs]class OpticSystem(Drawable): # OpticSystem
def __init__(self, elements:"List[OpticElement | Surface]"=None, envGlassParam=gps["air"]): # OpticSystem
self.envGlassParam = envGlassParam; self.opticElements = []; self.rawSurfaces = []; self._debug = []; self.add(*elements or []) # OpticSystem
[docs] def add(self, *elements): # OpticSystem
for element in elements: # OpticSystem
if isinstance(element, Surface): self.rawSurfaces.append(copy.copy(element)) # element is copied so that this OpticSystem can freely modify all Surfaces params without worrying that it will affect other setups # OpticSystem
elif isinstance(element, OpticElement): self.opticElements.append(copy.copy(element)) # OpticSystem
else: raise Exception(f"Can only add either Surface or OpticElement. Object passed in is of type {type(element)}") # OpticSystem
return self # OpticSystem
def _details(self) -> str: # OpticSystem
count = 0; s = "" # OpticSystem
for e in self.opticElements: # OpticSystem
s += f"{e}\n" # OpticSystem
for sur in e.surfaces: s += f"{count}) {sur}\n"; count += 1 # OpticSystem
s += "\n" # OpticSystem
if self.rawSurfaces: # OpticSystem
s += "Raw surfaces:\n" # OpticSystem
for sur in self.rawSurfaces: s += f"{count}) {sur}\n"; count += 1 # OpticSystem
s += "\n" # OpticSystem
return s # OpticSystem
@property # OpticSystem
def surfaces(self): # OpticSystem
"""Grabs all surfaces from all optic elements and raw surfaces""" # OpticSystem
return [*self.opticElements | cli.op().surfaces.all(), self.rawSurfaces] | cli.joinSt() | cli.deref() # OpticSystem
[docs] def cast(self, rays:Rays, passes=20, verbose=False) -> RaysPath: # OpticSystem
rays = copy.copy(rays); surfaces = self.surfaces # OpticSystem
# if verbose: print(f"OpticElements ({len(self.opticElements)} total):"); print(self.opticElements | cli.insId() | ~cli.apply(lambda i,e: f"{i}) {e}") | cli.join("\n")) # OpticSystem
for i, s in enumerate(surfaces): # OpticSystem
a = s.gp1 is None; b = s.gp2 is None # OpticSystem
if a or b: # OpticSystem
if a and not b: s.gp1 = gps["air"] # OpticSystem
elif b and not a: s.gp2 = gps["air"] # OpticSystem
else: raise Exception(f"Surface #{i} {s} have to have at least one specified glass params") # OpticSystem
if verbose: print(self._details()) # OpticSystem
# if verbose: print(f"\nSurfaces ({len(surfaces)} total):"); print(surfaces | cli.insId() | ~cli.apply(lambda i,s: f"{i}) {s}") | cli.join("\n"), end="\n\n") # OpticSystem
r = rays; path = [r]; data = []; self._debug = [] # OpticSystem
for i in range(passes): # OpticSystem
raysBefore = []; raysAfter = []; data.append([]); self._debug.append([]) # OpticSystem
for s in surfaces: rp = copy.copy(r); raysBefore.append(rp); raysAfter.append(s.castS(rp)) # casting to all available surfaces # OpticSystem
raysBefore = np.stack([re.data for re in raysBefore]); raysAfter = np.stack([re.data for re in raysAfter]); self._debug[-1].append([np.copy(raysBefore), np.copy(raysAfter)]) # OpticSystem
intersected = raysAfter[:,:,5] - raysBefore[:,:,5]; intersected = (0.9 < intersected) * (intersected < 1.1); self._debug[-1].append(intersected) # OpticSystem
raysBefore[~intersected,3] = float("inf") # If there're no intersections, then set the raysBefore's lengths to inf. Checks for intersection by noting the #transforms # OpticSystem
argmin = raysBefore[:,:,3] | cli.T() | cli.toArgmin().all(); data[-1].append(argmin); self._debug[-1].append(argmin) # then checks for closest interaction # OpticSystem
r.data[:] = raysBefore[argmin, range(raysBefore.shape[1])]; newR = Rays(raysAfter[argmin, range(raysAfter.shape[1])], ogSurface=surfaces[argmin | cli.count() | ~cli.sort() | cli.cut(1) | cli.item()]) # OpticSystem
if np.allclose(r.data[:,5], newR.data[:,5]): break # trim away at the paths that no longer change in #transforms # OpticSystem
r = newR; path.append(r) # OpticSystem
if verbose: print(data | cli.cut(0) | cli.pretty() | cli.insId() | ~cli.apply(lambda x,y: "".join([f"Pass: {x}; ", f"chosen surfaces: {y}"])) | cli.join("\n")) # OpticSystem
return RaysPath(path) # OpticSystem
def __getitem__(self, idx): return self.surfaces[idx] # OpticSystem
def __len__(self): return len(self.surfaces) # OpticSystem
def _draw(self, config): config["surfaceIdx2Id"] = {s.idx:f"{i}" for i,s in enumerate(self.surfaces)}; [e._draw(config) for e in self.opticElements]; [s._draw(config) for s in self.surfaces]; return p5.svg() # OpticSystem
def __repr__(self): return f"<OpticSystem #opticElements={len(self.opticElements)} #surfaces={len(self.surfaces)}>" # OpticSystem
def _repr_extra(self): return html.escape(self._details()).replace("\n", "<br>") # OpticSystem
[docs] def bounds(self): return [self.opticElements, self.surfaces] | cli.op().bounds().all(2) | cli.joinSt() | cli.T() | cli.toMin() + cli.toMin() + cli.toMax() + cli.toMax() | cli.deref() # OpticSystem