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