def traceSplitReverse(self, r, inCoordSys=globalCoordSys, forwardCoordSys=None, reverseCoordSys=None, minFlux=1e-3, verbose=False): """Assume Rays are coming in from the reverse direction. returns forwardRays, forwardCoordSys, reverseRays, reverseCoordSys """ if verbose: strtemplate = "traceSplitReverse {:15s} flux = {:18.8f} nphot = {:10d}" print(strtemplate.format(self.name, np.sum(r.flux), len(r))) if self.skip: return r, None, inCoordSys, None transform = batoid.CoordTransform(inCoordSys, self.coordSys) r = transform.applyForward(r) rForward, rReverse = self.rSplitReverse(r) # For now, apply obscuration equally forwards and backwards if self.obscuration is not None: rForward = self.obscuration.obscure(rForward) rReverse = self.obscuration.obscure(rReverse) if forwardCoordSys is None: forwardCoordSys = self.coordSys else: transform = batoid.CoordTransform(self.coordSys, forwardCoordSys) rForward = transform.applyForward(rForward) if reverseCoordSys is None: reverseCoordSys = self.coordSys else: transform = batoid.CoordTransform(self.coordSys, reverseCoordSys) rReverse = transform.applyForward(rReverse) return rForward, rReverse, forwardCoordSys, reverseCoordSys
def test_ne(): objs = [ batoid.CoordSys(), batoid.CoordTransform(batoid.CoordSys(), batoid.CoordSys()), batoid.CoordTransform(batoid.CoordSys(), batoid.CoordSys(origin=(0, 0, 1))) ] all_obj_diff(objs)
def getGlobalRays(traceFull, start=None, end=None, globalSys=globalCoordSys): """Calculate an array of ray vertices in global coordinates. Parameters ---------- traceFull : array Array of per-surface ray-tracing output from traceFull() start : str or None Name of the first surface to include in the output, or use the first surface in the model when None. end : str or None Name of the last surface to include in the output, or use the last surface in the model when None. globalSys : batoid.CoordSys Global coordinate system to use. Returns ------- tuple Tuple (xyz, raylen) of arrays with shapes (nray, 3, nsurf + 1) and (nray,). The xyz array contains the global coordinates of each ray vertex, with raylen giving the number of visible (not vignetted) vertices for each ray. """ names = [trace['name'] for trace in traceFull] try: start = 0 if start is None else names.index(start) except ValueError: raise ValueError('No such start surface "{0}".'.format(start)) try: end = len(names) if end is None else names.index(end) + 1 except ValueError: raise ValueError('No such end surface "{0}".'.format(end)) nsurf = end - start if nsurf <= 0: raise ValueError('Expected start < end.') nray = len(traceFull[start]['in']) # Allocate an array for all ray vertices in global coords. xyz = np.empty((nray, 3, nsurf + 1)) # First point on each ray is where it enters the start surface. transform = batoid.CoordTransform(traceFull[start]['inCoordSys'], globalSys) xyz[:, :, 0] = np.stack(transform.applyForward(*traceFull[start]['in'].r.T), axis=1) # Keep track of the number of visible points on each ray. raylen = np.ones(nray, dtype=int) for i, surface in enumerate(traceFull[start:end]): # Add a point for where each ray leaves this surface. transform = batoid.CoordTransform(surface['outCoordSys'], globalSys) xyz[:, :, i + 1] = np.stack(transform.applyForward(*surface['out'].r.T), axis=1) # Keep track of rays which are still visible. visible = ~surface['out'].vignetted raylen[visible] += 1 return xyz, raylen
def test_simple_transform(): rng = np.random.default_rng(5) size = 10_000 # Worked a few examples out manually # Example 1 coordSys1 = batoid.CoordSys() coordSys2 = batoid.CoordSys().shiftGlobal([0, 0, 1]) rv = randomRayVector(rng, size, coordSys1) transform = batoid.CoordTransform(coordSys1, coordSys2) rv2 = batoid.applyForwardTransform(transform, rv.copy()) np.testing.assert_allclose(rv.x, rv2.x) np.testing.assert_allclose(rv.y, rv2.y) np.testing.assert_allclose(rv.z - 1, rv2.z) # Repeat using toCoordSys rv3 = rv.copy().toCoordSys(coordSys2) np.testing.assert_allclose(rv.x, rv3.x) np.testing.assert_allclose(rv.y, rv3.y) np.testing.assert_allclose(rv.z - 1, rv3.z) # Transform of numpy array x, y, z = transform.applyForwardArray(rv.x, rv.y, rv.z) np.testing.assert_allclose(rv2.x, x) np.testing.assert_allclose(rv2.y, y) np.testing.assert_allclose(rv2.z, z) # Example 2 # First for a single specific point I worked out coordSys1 = batoid.CoordSys() coordSys2 = batoid.CoordSys(origin=[1, 1, 1], rot=batoid.RotY(np.pi / 2)) x = y = z = np.array([2]) vx = vy = vz = np.array([0]) rv = batoid.RayVector(x, y, z, vx, vy, vz) transform = batoid.CoordTransform(coordSys1, coordSys2) rv2 = batoid.applyForwardTransform(transform, rv.copy()) np.testing.assert_allclose(rv2.r, [[-1, 1, 1]]) # Transform of numpy array x, y, z = transform.applyForwardArray(rv.x, rv.y, rv.z) np.testing.assert_allclose(rv2.x, x) np.testing.assert_allclose(rv2.y, y) np.testing.assert_allclose(rv2.z, z) # Here's the generalization # Also using alternate syntax for applyForward here. rv = randomRayVector(rng, size, coordSys1) rv2 = transform.applyForward(rv.copy()) np.testing.assert_allclose(rv2.x, 1 - rv.z) np.testing.assert_allclose(rv2.y, rv.y - 1) np.testing.assert_allclose(rv2.z, rv.x - 1) rv3 = rv.copy().toCoordSys(coordSys2) np.testing.assert_allclose(rv3.x, 1 - rv.z) np.testing.assert_allclose(rv3.y, rv.y - 1) np.testing.assert_allclose(rv3.z, rv.x - 1) # Transform of numpy array x, y, z = transform.applyForwardArray(rv.x, rv.y, rv.z) np.testing.assert_allclose(rv2.x, x) np.testing.assert_allclose(rv2.y, y) np.testing.assert_allclose(rv2.z, z)
def traceSplitReverse(self, r, inCoordSys=globalCoordSys, forwardCoordSys=None, reverseCoordSys=None, minFlux=1e-3, verbose=False): if verbose: strtemplate = "traceSplitReverse {:15s} flux = {:18.8f} nphot = {:10d}" print(strtemplate.format(self.name, np.sum(r.flux), len(r))) if self.skip: return r, None, inCoordSys, None if forwardCoordSys is None: forwardCoordSys = self.items[-1].coordSys if reverseCoordSys is None: reverseCoordSys = self.items[0].coordSys workQueue = [(r, inCoordSys, "reverse", len(self.items)-1)] outRForward = [] outRReverse = [] while workQueue: rays, inCoordSys, direction, itemIndex = workQueue.pop() item = self.items[itemIndex] if direction == "forward": rForward, rReverse, tmpForwardCoordSys, tmpReverseCoordSys = \ item.traceSplit(rays, inCoordSys, minFlux=minFlux, verbose=verbose) elif direction == "reverse": rForward, rReverse, tmpForwardCoordSys, tmpReverseCoordSys = \ item.traceSplitReverse(rays, inCoordSys, minFlux=minFlux, verbose=verbose) else: raise RuntimeError("Shouldn't get here!") rForward.trimVignettedInPlace(minFlux) rReverse.trimVignettedInPlace(minFlux) if itemIndex == 0: if len(rReverse) > 0: if tmpReverseCoordSys != reverseCoordSys: transform = batoid.CoordTransform(tmpReverseCoordSys, reverseCoordSys) transform.applyForwardInPlace(rReverse) outRReverse.append(rReverse) else: if len(rReverse) > 0: workQueue.append((rReverse, tmpReverseCoordSys, "reverse", itemIndex-1)) if itemIndex == len(self.items)-1: if len(rForward) > 0: if tmpForwardCoordSys != forwardCoordSys: transform = batoid.CoordTransform(tmpForwardCoordSys, forwardCoordSys) transform.applyForwardInPlace(rForward) outRForward.append(rForward) else: if len(rForward) > 1: workQueue.append((rForward, tmpForwardCoordSys, "forward", itemIndex+1)) rForward = batoid.concatenateRayVectors(outRForward) rReverse = batoid.concatenateRayVectors(outRReverse) return rForward, rReverse, forwardCoordSys, reverseCoordSys
def test_ne(): objs = [ batoid.CoordSys(), batoid.CoordSys([0,0,1]), batoid.CoordSys([0,1,0]), batoid.CoordSys(batoid.RotX(0.1)), batoid.CoordTransform(batoid.CoordSys(), batoid.CoordSys()), batoid.CoordTransform(batoid.CoordSys(), batoid.CoordSys([0,0,1])), batoid.CoordTransform(batoid.CoordSys(), batoid.CoordSys(batoid.RotX(0.1))) ] all_obj_diff(objs)
def test_roundtrip(): rng = np.random.default_rng(57) size = 10_000 for i in range(10): coordSys1 = randomCoordSys(rng) coordSys2 = randomCoordSys(rng) transform = batoid.CoordTransform(coordSys1, coordSys2) rv0 = randomRayVector(rng, size, coordSys1) rv1 = transform.applyForward(rv0.copy()) rv2 = transform.applyReverse(rv1.copy()) rv3 = rv0.copy().toCoordSys(coordSys2).toCoordSys(coordSys1) assert rv0.coordSys == coordSys1 assert rv1.coordSys == coordSys2 assert rv2.coordSys == coordSys1 assert rv3.coordSys == coordSys1 rays_allclose(rv0, rv2) rays_allclose(rv0, rv3) x, y, z = transform.applyForwardArray(rv0.x, rv0.y, rv0.z) x, y, z = transform.applyReverseArray(x, y, z) np.testing.assert_allclose(x, rv0.x) np.testing.assert_allclose(y, rv0.y) np.testing.assert_allclose(z, rv0.z)
def test_HSC_trace(): fn = os.path.join(batoid.datadir, "HSC", "HSC_old.yaml") config = yaml.load(open(fn)) telescope = batoid.parse.parse_optic(config['opticalSystem']) # Zemax has a number of virtual surfaces that we don't trace in batoid. Also, the HSC.yaml # above includes Baffle surfaces not in Zemax. The following lists select out the surfaces in # common to both models. HSC_surfaces = [ 3, 6, 7, 8, 9, 11, 12, 13, 14, 16, 17, 18, 19, 20, 21, 24, 25, 28, 29, 31 ] surface_names = [ 'PM', 'G1_entrance', 'G1_exit', 'G2_entrance', 'G2_exit', 'ADC1_entrance', 'ADC1_exit', 'ADC2_entrance', 'ADC2_exit', 'G3_entrance', 'G3_exit', 'G4_entrance', 'G4_exit', 'G5_entrance', 'G5_exit', 'F_entrance', 'F_exit', 'W_entrance', 'W_exit', 'D' ] for fn in [ "HSC_raytrace_1.txt", "HSC_raytrace_2.txt", "HSC_raytrace_3.txt" ]: filename = os.path.join(directory, "testdata", fn) with open(filename) as f: arr = np.loadtxt(f, skiprows=22, usecols=list(range(0, 12))) arr0 = arr[0] ray = batoid.Ray(arr0[1] / 1000, arr0[2] / 1000, 16.0, arr0[4], arr0[5], -arr0[6], t=0, wavelength=750e-9) tf = telescope.traceFull(ray) i = 0 for surface in tf: if surface['name'] != surface_names[i]: continue s = surface['out'] v = s.v / np.linalg.norm(s.v) transform = batoid.CoordTransform(surface['outCoordSys'], batoid.CoordSys()) s = transform.applyForward(s) bt_isec = np.array([s.x, s.y, s.z - 16.0]) zx_isec = arr[HSC_surfaces[i] - 1][1:4] / 1000 np.testing.assert_allclose(bt_isec, zx_isec, rtol=0, atol=1e-9) # nanometer agreement bt_angle = np.array([v[0], v[1], v[2]]) zx_angle = arr[HSC_surfaces[i] - 1][4:7] # direction cosines agree to 1e-9 np.testing.assert_allclose(bt_angle, zx_angle, rtol=0, atol=1e-9) i += 1
def traceInPlace(self, r, inCoordSys=globalCoordSys, outCoordSys=None): if self.skip: return r, inCoordSys transform = batoid.CoordTransform(inCoordSys, self.coordSys) transform.applyForwardInPlace(r) # refract, reflect, pass-through - depending on subclass self.interactInPlace(r) if self.obscuration is not None: self.obscuration.obscureInPlace(r) if outCoordSys is None: return r, self.coordSys else: transform = batoid.CoordTransform(self.coordSys, outCoordSys) transform.applyForwardInPlace(r) return r, outCoordSys
def test_composition(): rng = np.random.default_rng(577) size = 10_000 for i in range(10): coordSys1 = randomCoordSys(rng) coordSys2 = randomCoordSys(rng) coordSys3 = randomCoordSys(rng) assert coordSys1 != coordSys2 assert coordSys1 != coordSys3 do_pickle(coordSys1) transform1to2 = batoid.CoordTransform(coordSys1, coordSys2) transform1to3 = batoid.CoordTransform(coordSys1, coordSys3) transform2to3 = batoid.CoordTransform(coordSys2, coordSys3) do_pickle(transform1to2) for j in range(10): rv = randomRayVector(rng, size, coordSys1) rv_a = transform1to3.applyForward(rv.copy()) assert rv_a != rv assert rv_a.coordSys == coordSys3 rv_b = transform2to3.applyForward( transform1to2.applyForward(rv.copy())) assert rv_b != rv assert rv_b.coordSys == coordSys3 rays_allclose(rv_a, rv_b), "error with composite transform of RayVector" rv = randomRayVector(rng, size, coordSys3) rv_a = transform1to3.applyReverse(rv.copy()) assert rv_a != rv assert rv_a.coordSys == coordSys1 rv_b = transform1to2.applyReverse( transform2to3.applyReverse(rv.copy())) assert rv_b != rv assert rv_b.coordSys == coordSys1 rays_allclose( rv_a, rv_b), "error with reverse composite transform of RayVector"
def traceReverse(self, r, inCoordSys=globalCoordSys, outCoordSys=None): """Trace through optic(s) in reverse. Note, you may need to reverse the direction of rays for this to work. """ if self.skip: return r, inCoordSys transform = batoid.CoordTransform(inCoordSys, self.coordSys) r = transform.applyForward(r) r = self.interactReverse(r) if self.obscuration is not None: r = self.obscuration.obscure(r) if outCoordSys is None: return r, self.coordSys else: transform = batoid.CoordTransform(self.coordSys, outCoordSys) return transform.applyForward(r), outCoordSys
def draw2d(self, ax, **kwargs): """ Draw this interface on a 2d matplotlib axis. May not work if elements are non-circular or not axis-aligned. """ if self.outRadius is None: return # Drawing in the x-z plane. x = np.linspace(-self.outRadius, -self.inRadius) y = np.zeros_like(x) z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, z, **kwargs) x = np.linspace(self.inRadius, self.outRadius) y = np.zeros_like(x) z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, z, **kwargs)
def wavefront(optic, theta_x, theta_y, wavelength, nx=32, sphereRadius=None): """Compute wavefront. Parameters ---------- optic : batoid.Optic Optic for which to compute wavefront. theta_x, theta_y : float Field of incoming rays (gnomic projection) wavelength : float Wavelength of incoming rays nx : int, optional Size of ray grid to generate to compute wavefront. Default: 32 sphereRadius : float, optional Radius of reference sphere in meters. If None, then use optic.sphereRadius. Returns ------- wavefront : batoid.Lattice A batoid.Lattice object containing the wavefront values in waves and the primitive lattice vectors of the entrance pupil grid in meters. """ dirCos = gnomicToDirCos(theta_x, theta_y) rays = batoid.rayGrid( optic.dist, optic.pupilSize, dirCos[0], dirCos[1], -dirCos[2], nx, wavelength, 1.0, optic.inMedium ) if sphereRadius is None: sphereRadius = optic.sphereRadius outCoordSys = batoid.CoordSys() optic.traceInPlace(rays, outCoordSys=outCoordSys) w = np.where(1-rays.vignetted)[0] point = np.mean(rays.r[w], axis=0) # We want to place the vertex of the reference sphere one radius length away from the # intersection point. So transform our rays into that coordinate system. transform = batoid.CoordTransform( outCoordSys, batoid.CoordSys(point+np.array([0,0,sphereRadius]))) transform.applyForwardInPlace(rays) sphere = batoid.Sphere(-sphereRadius) sphere.intersectInPlace(rays) w = np.where(1-rays.vignetted)[0] # Should potentially try to make the reference time w.r.t. the chief ray instead of the mean # of the good (unvignetted) rays. t0 = np.mean(rays.t[w]) arr = np.ma.masked_array((t0-rays.t)/wavelength, mask=rays.vignetted).reshape(nx, nx) primitiveVectors = np.vstack([[optic.pupilSize/nx, 0], [0, optic.pupilSize/nx]]) return batoid.Lattice(arr, primitiveVectors)
def drawTrace2d(ax, traceFull, start=None, end=None, **kwargs): if start is None: start = traceFull[0]['name'] if end is None: end = traceFull[-1]['name'] doPlot = False for surface in traceFull: if surface['name'] == start: doPlot = True if doPlot: inTransform = batoid.CoordTransform(surface['inCoordSys'], globalCoordSys) outTransform = batoid.CoordTransform(surface['outCoordSys'], globalCoordSys) for inray, outray in zip(surface['in'], surface['out']): if not outray.vignetted: inray = inTransform.applyForward(inray) outray = outTransform.applyForward(outray) ax.plot([inray.x, outray.x], [inray.z, outray.z], **kwargs) if surface['name'] == end: break
def draw3d(self, ax, **kwargs): """ Draw this interface on a mplot3d axis. Parameters ---------- ax : mplot3d.Axis Axis on which to draw this optic. """ if self.outRadius is None: return # Going to draw 4 objects here: inner circle, outer circle, sag along x=0, sag along y=0 # inner circle if self.inRadius != 0.0: th = np.linspace(0, 2 * np.pi, 100) cth, sth = np.cos(th), np.sin(th) x = self.inRadius * cth y = self.inRadius * sth z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, y, z, **kwargs) #outer circle th = np.linspace(0, 2 * np.pi, 100) cth, sth = np.cos(th), np.sin(th) x = self.outRadius * cth y = self.outRadius * sth z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, y, z, **kwargs) #next, a line at X=0 y = np.linspace(-self.outRadius, -self.inRadius) x = np.zeros_like(y) z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, y, z, **kwargs) y = np.linspace(self.inRadius, self.outRadius) x = np.zeros_like(y) z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, y, z, **kwargs) #next, a line at Y=0 x = np.linspace(-self.outRadius, -self.inRadius) y = np.zeros_like(x) z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, y, z, **kwargs) x = np.linspace(self.inRadius, self.outRadius) y = np.zeros_like(x) z = self.surface.sag(x, y) transform = batoid.CoordTransform(self.coordSys, globalCoordSys) x, y, z = transform.applyForward(x, y, z) ax.plot(x, y, z, **kwargs)
def trace(self, r, inCoordSys=globalCoordSys, outCoordSys=None): """ Trace a ray through this optical element. Parameters ---------- r : batoid.Ray or batoid.RayVector input ray to trace inCoordSys : batoid.CoordSys Coordinate system of incoming ray(s). Default: the global coordinate system. outCoordSys : batoid.CoordSys Coordinate system into which to project the output ray(s). Default: None, which means use the coordinate system of the optic. Returns ------- Ray or RayVector, output CoordSys. """ if self.skip: return r, inCoordSys transform = batoid.CoordTransform(inCoordSys, self.coordSys) r = transform.applyForward(r) # refract, reflect, pass-through - depending on subclass r = self.interact(r) if self.obscuration is not None: r = self.obscuration.obscure(r) if outCoordSys is None: return r, self.coordSys else: transform = batoid.CoordTransform(self.coordSys, outCoordSys) return transform.applyForward(r), outCoordSys
def wavefront(optic, wavelength, theta_x=0, theta_y=0, nx=32, rays=None, saveRays=False, sphereRadius=None): if rays is None: xcos = np.sin(theta_x) ycos = np.sin(theta_y) zcos = -np.sqrt(1.0 - xcos**2 - ycos**2) rays = batoid.rayGrid(optic.dist, optic.pupilSize, xcos, ycos, zcos, nx, wavelength, optic.inMedium) if saveRays: rays = batoid.RayVector(rays) if sphereRadius is None: sphereRadius = optic.sphereRadius outCoordSys = batoid.CoordSys() optic.traceInPlace(rays, outCoordSys=outCoordSys) goodRays = batoid._batoid.trimVignetted(rays) point = np.array( [np.mean(goodRays.x), np.mean(goodRays.y), np.mean(goodRays.z)]) # We want to place the vertex of the reference sphere one radius length away from the # intersection point. So transform our rays into that coordinate system. transform = batoid.CoordTransform( outCoordSys, batoid.CoordSys(point + np.array([0, 0, sphereRadius]))) transform.applyForwardInPlace(rays) sphere = batoid.Sphere(-sphereRadius) sphere.intersectInPlace(rays) goodRays = batoid._batoid.trimVignetted(rays) # Should potentially try to make the reference time w.r.t. the chief ray instead of the mean # of the good (unvignetted) rays. t0 = np.mean(goodRays.t0) ts = rays.t0[:] isV = rays.isVignetted[:] ts -= t0 ts /= wavelength wf = np.ma.masked_array(ts, mask=isV) return wf
def test_shift(): import random random.seed(5) globalCoordSys = batoid.CoordSys() for i in range(30): x = random.gauss(0.1, 2.3) y = random.gauss(0.1, 2.3) z = random.gauss(0.1, 2.3) newCoordSys = globalCoordSys.shiftGlobal([x, y, z]) do_pickle(newCoordSys) np.testing.assert_array_equal(newCoordSys.xhat, [1,0,0]) np.testing.assert_array_equal(newCoordSys.yhat, [0,1,0]) np.testing.assert_array_equal(newCoordSys.zhat, [0,0,1]) np.testing.assert_array_equal(newCoordSys.origin, [x,y,z]) np.testing.assert_array_equal(newCoordSys.rot, np.eye(3)) coordTransform = batoid.CoordTransform(globalCoordSys, newCoordSys) do_pickle(coordTransform) for j in range(30): x2 = random.gauss(0.1, 2.3) y2 = random.gauss(0.1, 2.3) z2 = random.gauss(0.1, 2.3) newNewCoordSys = newCoordSys.shiftGlobal([x2, y2, z2]) np.testing.assert_array_equal(newNewCoordSys.xhat, [1,0,0]) np.testing.assert_array_equal(newNewCoordSys.yhat, [0,1,0]) np.testing.assert_array_equal(newNewCoordSys.zhat, [0,0,1]) np.testing.assert_array_equal(newNewCoordSys.origin, [x+x2, y+y2, z+z2]) np.testing.assert_array_equal(newNewCoordSys.rot, np.eye(3)) newNewCoordSys = newCoordSys.shiftLocal([x2, y2, z2]) np.testing.assert_array_equal(newNewCoordSys.xhat, [1,0,0]) np.testing.assert_array_equal(newNewCoordSys.yhat, [0,1,0]) np.testing.assert_array_equal(newNewCoordSys.zhat, [0,0,1]) np.testing.assert_array_equal(newNewCoordSys.origin, [x+x2, y+y2, z+z2]) np.testing.assert_array_equal(newNewCoordSys.rot, np.eye(3))
def test_composition(): import random random.seed(5) for i in range(10): coordSys1 = randomCoordSys() coordSys2 = randomCoordSys() coordSys3 = randomCoordSys() assert coordSys1 != coordSys2 assert coordSys1 != coordSys3 do_pickle(coordSys1) transform1to2 = batoid.CoordTransform(coordSys1, coordSys2) transform1to3 = batoid.CoordTransform(coordSys1, coordSys3) transform2to3 = batoid.CoordTransform(coordSys2, coordSys3) do_pickle(transform1to2) for i in range(10): vec3 = randomVec3() vec3_a = transform1to3.applyForward(vec3) vec3_b = transform2to3.applyForward(transform1to2.applyForward(vec3)) np.testing.assert_allclose(vec3_a, vec3_b) vec3_ra = transform1to3.applyReverse(vec3) vec3_rb = transform1to2.applyReverse(transform2to3.applyReverse(vec3)) np.testing.assert_allclose(vec3_ra, vec3_rb) ray = randomRay() ray_a = transform1to3.applyForward(ray) ray_b = transform2to3.applyForward(transform1to2.applyForward(ray)) assert ray_isclose(ray_a, ray_b), "error with composite transform of Ray" ray_ra = transform1to3.applyReverse(ray) ray_rb = transform1to2.applyReverse(transform2to3.applyReverse(ray)) assert ray_isclose(ray_ra, ray_rb), "error with reverse composite transform of Ray" rv = randomRayVector() rv_a = transform1to3.applyForward(rv) rv_b = transform2to3.applyForward(transform1to2.applyForward(rv)) assert rays_allclose(rv_a, rv_b), "error with composite transform of RayVector" rv_ra = transform1to3.applyReverse(rv) rv_rb = transform1to2.applyReverse(transform2to3.applyReverse(rv)) assert rays_allclose(rv_ra, rv_rb), "error with reverse composite transform of RayVector" # Test with numpy arrays xyz = rv.x, rv.y, rv.z xyz_a = transform1to3.applyForward(*xyz) xyz_b = transform2to3.applyForward(*transform1to2.applyForward(*xyz)) xyz_c = [transform2to3.applyForward(transform1to2.applyForward(r.r)) for r in rv] np.testing.assert_allclose(xyz_a, xyz_b) np.testing.assert_allclose(xyz_a, np.transpose(xyz_c)) # Should still work if we reshape. xyz2 = rv.x.reshape((2, 5)), rv.y.reshape((2, 5)), rv.z.reshape((2, 5)) xyz2_a = transform1to3.applyForward(*xyz2) xyz2_b = transform2to3.applyForward(*transform1to2.applyForward(*xyz2)) np.testing.assert_allclose(xyz2_a, xyz2_b) # And also work if we reverse np.testing.assert_allclose(xyz, transform1to3.applyReverse(*xyz_a)) np.testing.assert_allclose(xyz, transform1to2.applyReverse(*transform2to3.applyReverse(*xyz_b))) # Test in-place on Ray ray = randomRay() ray_copy = ray.copy() transform1to2.applyForwardInPlace(ray) transform2to3.applyForwardInPlace(ray) transform1to3.applyForwardInPlace(ray_copy) assert ray_isclose(ray, ray_copy) # in-place reverse on Ray ray = randomRay() ray_copy = ray.copy() transform2to3.applyReverseInPlace(ray) transform1to2.applyReverseInPlace(ray) transform1to3.applyReverseInPlace(ray_copy) assert ray_isclose(ray, ray_copy) # Test in-place on RayVector rv = randomRayVector() rv_copy = rv.copy() transform1to2.applyForwardInPlace(rv) transform2to3.applyForwardInPlace(rv) transform1to3.applyForwardInPlace(rv_copy) assert rays_allclose(rv, rv_copy) # in-place reverse on RayVector rv = randomRayVector() rv_copy = rv.copy() transform2to3.applyReverseInPlace(rv) transform1to2.applyReverseInPlace(rv) transform1to3.applyReverseInPlace(rv_copy) assert rays_allclose(rv, rv_copy)
def getXZSlice(self, nslice=0): """Calculate global coordinates for an (x,z) slice through this interface. The calculation is split into two half slices: xlocal <= 0 and xlocal >= 0. When the inner radius is zero, these half slices are merged into one. Otherwise, the two half slices are returned separately. If the local coordinate system involves any rotation the resulting slice may not be calculated correctly since we are really slicing in (xlocal, zlocal) then transforming these to (xglobal, zglobal). Parameters ---------- nslice : int Use the specified number of points on each half slice. When zero, the value will be calculated automatically (and will be 2 for planar surfaces). Returns ------- tuple Tuple (xz1, xz2) of 1D arrays where xz1=[x1, z1] is the xlocal <= 0 half slice and xz2=[x2, z2] is the xlocal >= 0 half slice. """ slice = [] if self.outRadius is None: return slice if nslice <= 0: if isinstance(self.surface, batoid.surface.Plane): nslice = 2 else: nslice = 50 # Calculate (x,z) slice in local coordinates for x <= 0. x = np.linspace(-self.outRadius, -self.inRadius, nslice) y = np.zeros_like(x) z = self.surface.sag(x, y) # Transform slice to global coordinates. transform = batoid.CoordTransform(self.coordSys, batoid.globalCoordSys) xneg, yneg, zneg = transform.applyForward(x, y, z) if np.any(yneg != 0): print('WARNING: getXZSlice used for rotated surface "{0}".'.format( self.name)) # Calculate (x,z) slice in local coordinates for x >= 0. x *= -1 x = x[::-1] z[:] = self.surface.sag(x, y) # Transform slice to global coordinates. xpos, ypos, zpos = transform.applyForward(x, y, z) if np.any(ypos != 0): print('WARNING: getXZSlice used for rotated surface "{0}".'.format( self.name)) slice.append(np.stack((xpos, zpos), axis=0)) # Combine x <= 0 and x >= 0 half slices when inner = 0. if self.inRadius == 0: assert xneg[-1] == xpos[0] and zneg[-1] == zpos[0] return (np.stack((np.hstack( (xneg, xpos[1:])), np.hstack((zneg, zpos[1:]))), axis=0), ) else: return (np.stack((xneg, zneg), axis=0), np.stack((xpos, zpos), axis=0))