def test_sag():
np.random.seed(57)
for _ in range(100):
s1 = batoid.Sphere(np.random.uniform(1, 3))
s2 = batoid.Paraboloid(np.random.uniform(1, 3))
sum = batoid.Sum([s1, s2])
x = np.random.normal(size=5000)
y = np.random.normal(size=5000)
np.testing.assert_allclose(
sum.sag(x, y),
s1.sag(x, y) + s2.sag(x, y),
rtol=1e-12,
atol=1e-12
)
s3 = batoid.Quadric(np.random.uniform(3, 5), np.random.uniform(-0.1, 0.1))
sum2 = batoid.Sum([s1, s2, s3])
np.testing.assert_allclose(
sum2.sag(x, y),
s1.sag(x, y) + s2.sag(x, y) + s3.sag(x, y),
rtol=1e-12,
atol=1e-12
)
def test_ne():
objs = [
batoid.Sum([batoid.Plane(), batoid.Plane()]),
batoid.Sum([batoid.Plane(), batoid.Sphere(1.0)]),
batoid.Sum([batoid.Plane(), batoid.Plane(), batoid.Plane()]),
batoid.Plane()
]
all_obj_diff(objs)
def test_properties():
np.random.seed(5)
for _ in range(100):
s1 = batoid.Sphere(np.random.uniform(1, 3))
s2 = batoid.Paraboloid(np.random.uniform(1, 3))
sum = batoid.Sum([s1, s2])
do_pickle(sum)
# check commutativity
assert sum == batoid.Sum([s2, s1])
# order of sum.surfaces is not guaranteed
assert s1 in sum.surfaces
assert s2 in sum.surfaces
s3 = batoid.Quadric(np.random.uniform(3, 5), np.random.uniform(-0.1, 0.1))
sum2 = batoid.Sum([s1, s2, s3])
do_pickle(sum2)
# check commutativity
assert sum2 == batoid.Sum([s2, s3, s1])
assert sum2 == batoid.Sum([s3, s1, s2])
assert sum2 == batoid.Sum([s3, s2, s1])
assert sum2 == batoid.Sum([s2, s1, s3])
assert sum2 == batoid.Sum([s1, s3, s2])
assert s1 in sum2.surfaces
assert s2 in sum2.surfaces
assert s3 in sum2.surfaces
do_pickle(sum)
def test_intersect():
np.random.seed(57721)
rv0 = batoid.RayVector([
batoid.Ray(
np.random.normal(scale=0.1),
np.random.normal(scale=0.1),
10,
np.random.normal(scale=1e-4),
np.random.normal(scale=1e-4),
-1
)
for _ in range(100)
])
for _ in range(100):
s1 = batoid.Sphere(np.random.uniform(3, 10))
s2 = batoid.Paraboloid(np.random.uniform(3, 10))
sum = batoid.Sum([s1, s2])
rv = batoid.RayVector(rv0)
rv1 = sum.intersect(rv)
rv2 = batoid.RayVector([sum.intersect(r) for r in rv])
rv3 = batoid.RayVector(rv)
sum.intersectInPlace(rv)
for r in rv3:
sum.intersectInPlace(r)
assert rays_allclose(rv1, rv)
assert rays_allclose(rv2, rv)
assert rays_allclose(rv3, rv)
def test_refract():
rng = np.random.default_rng(577215)
size = 10_000
for _ in range(100):
s1 = batoid.Sphere(1. / rng.normal(0., 0.2))
s2 = batoid.Paraboloid(rng.uniform(1, 3))
sum = batoid.Sum([s1, s2])
m0 = batoid.ConstMedium(rng.normal(1.2, 0.01))
m1 = batoid.ConstMedium(rng.normal(1.3, 0.01))
x = rng.uniform(-1, 1, size=size)
y = rng.uniform(-1, 1, size=size)
z = np.full_like(x, -10.0)
vx = rng.uniform(-1e-5, 1e-5, size=size)
vy = rng.uniform(-1e-5, 1e-5, size=size)
vz = np.sqrt(1 - vx * vx - vy * vy) / m0.n
rv = batoid.RayVector(x, y, z, vx, vy, vz)
rvr = batoid.refract(sum, rv.copy(), m0, m1)
rvr2 = sum.refract(rv.copy(), m0, m1)
rays_allclose(rvr, rvr2)
# print(f"{np.sum(rvr.failed)/len(rvr)*100:.2f}% failed")
normal = sum.normal(rvr.x, rvr.y)
# Test Snell's law
s0 = np.sum(np.cross(normal, rv.v * m0.n)[~rvr.failed], axis=-1)
s1 = np.sum(np.cross(normal, rvr.v * m1.n)[~rvr.failed], axis=-1)
np.testing.assert_allclose(m0.n * s0, m1.n * s1, rtol=0, atol=1e-9)
# Test that rv.v, rvr.v and normal are all in the same plane
np.testing.assert_allclose(np.einsum("ad,ad->a",
np.cross(normal, rv.v),
rv.v)[~rvr.failed],
0.0,
rtol=0,
atol=1e-12)
def test_reflect():
rng = np.random.default_rng(57721)
size = 10_000
for _ in range(100):
s1 = batoid.Sphere(1. / rng.normal(0., 0.2))
s2 = batoid.Paraboloid(rng.uniform(1, 3))
sum = batoid.Sum([s1, s2])
x = rng.uniform(-1, 1, size=size)
y = rng.uniform(-1, 1, size=size)
z = np.full_like(x, -10.0)
vx = rng.uniform(-1e-5, 1e-5, size=size)
vy = rng.uniform(-1e-5, 1e-5, size=size)
vz = np.full_like(x, 1)
rv = batoid.RayVector(x, y, z, vx, vy, vz)
rvr = batoid.reflect(sum, rv.copy())
rvr2 = sum.reflect(rv.copy())
rays_allclose(rvr, rvr2)
# print(f"{np.sum(rvr.failed)/len(rvr)*100:.2f}% failed")
normal = sum.normal(rvr.x, rvr.y)
# Test law of reflection
a0 = np.einsum("ad,ad->a", normal, rv.v)[~rvr.failed]
a1 = np.einsum("ad,ad->a", normal, -rvr.v)[~rvr.failed]
np.testing.assert_allclose(a0, a1, rtol=0, atol=1e-12)
# Test that rv.v, rvr.v and normal are all in the same plane
np.testing.assert_allclose(np.einsum("ad,ad->a",
np.cross(normal, rv.v),
rv.v)[~rvr.failed],
0.0,
rtol=0,
atol=1e-12)
def test_sum_bicubic():
import os
import yaml
fn = os.path.join(batoid.datadir, "LSST", "LSST_i.yaml")
config = yaml.safe_load(open(fn))
telescope = batoid.parse.parse_optic(config['opticalSystem'])
xcos, ycos, zcos = batoid.utils.gnomonicToDirCos(np.deg2rad(0.8),
np.deg2rad(0.8))
rays = batoid.circularGrid(
telescope.dist, telescope.pupilSize / 2,
telescope.pupilSize * telescope.pupilObscuration / 2, xcos, ycos,
-zcos, 50, 50, 750e-9, 1.0, telescope.inMedium)
out, _ = telescope.trace(rays)
m2 = telescope.itemDict['LSST.M2']
xs = np.linspace(-m2.outRadius, m2.outRadius, 200)
ys = xs
zs = np.zeros((200, 200), dtype=float)
bicubic = batoid.Bicubic(xs, ys, zs)
m2.surface = batoid.Sum([m2.surface, bicubic])
out2, _ = telescope.trace(rays)
# Don't expect exact equality, but should be very similar
assert rays_allclose(out, out2, atol=1e-13)
def test_add_plane():
np.random.seed(577)
for _ in range(100):
# Adding a plane should have zero effect on sag or normal vector
s1 = batoid.Sphere(np.random.uniform(1, 3))
s2 = batoid.Plane()
sum = batoid.Sum([s1, s2])
x = np.random.normal(size=5000)
y = np.random.normal(size=5000)
np.testing.assert_allclose(
sum.sag(x, y),
s1.sag(x, y),
rtol=1e-12,
atol=1e-12
)
for _x, _y in zip(x[:100], y[:100]):
np.testing.assert_allclose(
sum.normal(_x, _y),
s1.normal(_x, _y),
rtol=1e-12,
atol=1e-12
)
def test_sum_paraboloid():
# para_sag = r^2/(2*R^2)
# so two paraboloids yields r^2 * (1/(2*R1) + 1/(2*R2))
# so (1/(2*R1) + 1/(2*R2)) = 1/(2*R)
# implies
# 0.5/(1/(2*R1) + 1/(2*R2)) = R
np.random.seed(5772)
for _ in range(100):
R1 = np.random.uniform(1, 2)
R2 = np.random.uniform(2, 3)
Rsum = 0.5/(1/(2*R1) + 1/(2*R2))
para1 = batoid.Paraboloid(R1)
para2 = batoid.Paraboloid(R2)
paraSum = batoid.Paraboloid(Rsum)
paraSum2 = batoid.Sum([para1, para2])
x = np.random.normal(size=5000)
y = np.random.normal(size=5000)
np.testing.assert_allclose(
paraSum.sag(x, y),
paraSum2.sag(x, y),
rtol=1e-12,
atol=1e-12
)
for _x, _y in zip(x[:100], y[:100]):
np.testing.assert_allclose(
paraSum.normal(_x, _y),
paraSum2.normal(_x, _y),
rtol=1e-12,
atol=1e-12
)
def test_sum_paraboloid():
# para_sag = r^2/(2*R^2)
# so two paraboloids yields r^2 * (1/(2*R1) + 1/(2*R2))
# so (1/(2*R1) + 1/(2*R2)) = 1/(2*R)
# implies
# 0.5/(1/(2*R1) + 1/(2*R2)) = R
rng = np.random.default_rng(57721566)
for _ in range(100):
R1 = rng.uniform(1, 2)
R2 = rng.uniform(2, 3)
Rsum = 0.5 / (1 / (2 * R1) + 1 / (2 * R2))
para1 = batoid.Paraboloid(R1)
para2 = batoid.Paraboloid(R2)
paraSum = batoid.Paraboloid(Rsum)
paraSum2 = batoid.Sum([para1, para2])
x = rng.normal(size=5000)
y = rng.normal(size=5000)
np.testing.assert_allclose(paraSum.sag(x, y),
paraSum2.sag(x, y),
rtol=0,
atol=1e-12)
np.testing.assert_allclose(
paraSum.normal(x, y),
paraSum2.normal(x, y),
rtol=0,
atol=1e-12,
)
def test_sum_bicubic():
telescope = batoid.Optic.fromYaml("LSST_i.yaml")
xcos, ycos, zcos = batoid.utils.gnomonicToDirCos(
np.deg2rad(0.8), np.deg2rad(0.8)
)
rays = batoid.circularGrid(
telescope.backDist,
telescope.pupilSize/2,
telescope.pupilSize*telescope.pupilObscuration/2,
xcos, ycos, -zcos,
50, 50, 750e-9, 1.0, telescope.inMedium
)
out = telescope.trace(rays)
m2 = telescope['LSST.M2']
xs = np.linspace(-m2.outRadius, m2.outRadius, 200)
ys = xs
zs = np.zeros((200, 200), dtype=float)
bicubic = batoid.Bicubic(xs, ys, zs)
m2.surface = batoid.Sum([m2.surface, bicubic])
out2 = telescope.trace(rays)
# Don't expect exact equality, but should be very similar
assert rays_allclose(out, out2, atol=1e-13)
def test_fail():
sum = batoid.Sum([batoid.Plane(), batoid.Sphere(1.0)])
rv = batoid.RayVector(0, 10, 0, 0, 0, -1) # Too far to side
rv2 = batoid.intersect(sum, rv.copy())
np.testing.assert_equal(rv2.failed, np.array([True]))
# This one passes
rv = batoid.RayVector(0, 0, -1, 0, 0, +1)
rv2 = batoid.intersect(sum, rv.copy())
np.testing.assert_equal(rv2.failed, np.array([False]))
def test_fail():
sum = batoid.Sum([batoid.Plane(), batoid.Sphere(1.0)])
ray = batoid.Ray([0,0,sum.sag(0,0)-1], [0,0,-1])
ray = sum.intersect(ray)
assert ray.failed
ray = batoid.Ray([0,0,sum.sag(0,0)-1], [0,0,-1])
sum.intersectInPlace(ray)
assert ray.failed
def test_properties():
rng = np.random.default_rng(5)
for i in range(100):
s1 = batoid.Sphere(rng.uniform(1, 3))
s2 = batoid.Paraboloid(rng.uniform(1, 3))
sum = batoid.Sum([s1, s2])
do_pickle(sum)
assert s1 is sum.surfaces[0]
assert s2 is sum.surfaces[1]
s3 = batoid.Quadric(rng.uniform(3, 5), rng.uniform(-0.1, 0.1))
sum2 = batoid.Sum([s1, s2, s3])
do_pickle(sum2)
assert s1 is sum2.surfaces[0]
assert s2 is sum2.surfaces[1]
assert s3 is sum2.surfaces[2]
def test_LSST_M1_zernike():
"""See how much a ~100 nm zernike perturbation to M1 affects wavefront zernikes
"""
np.random.seed(5772156)
telescope = batoid.Optic.fromYaml("LSST_r.yaml")
theta_x = np.deg2rad(1.185)
theta_y = np.deg2rad(0.45)
fiducialZernikes = batoid.psf.zernike(telescope, theta_x, theta_y, 750e-9)
N = 256
xs = np.linspace(-8.36/2, 8.36/2, N)
ys = np.linspace(-8.36/2, 8.36/2, N)
jmax = 22
for _ in range(10):
coef = np.random.normal(size=jmax+1)*1e-7/np.sqrt(jmax) # aim for ~100 nm rms
R_inner = np.random.uniform(0.0, 0.65)
zsurf = batoid.Zernike(coef, R_outer=8.36/2, R_inner=0.61*8.36/2)
zs = zsurf.sag(*np.meshgrid(xs, ys))
bc = batoid.Bicubic(xs, ys, zs)
# Add Zernike perturbation to M1
zTelescope = batoid.Optic.fromYaml("LSST_r.yaml")
zPerturbedM1 = batoid.Sum([
zTelescope['LSST.M1'].surface,
zsurf
])
zTelescope['LSST.M1'].surface = zPerturbedM1
zZernikes = batoid.psf.zernike(zTelescope, theta_x, theta_y, 750e-9)
# Repeat with bicubic perturbation
bcTelescope = batoid.Optic.fromYaml("LSST_r.yaml")
bcPerturbedM1 = batoid.Sum([
bcTelescope['LSST.M1'].surface,
bc
])
bcTelescope['LSST.M1'].surface = bcPerturbedM1
bcZernikes = batoid.psf.zernike(bcTelescope, theta_x, theta_y, 750e-9)
np.testing.assert_allclose(zZernikes, bcZernikes, rtol=0, atol=1e-3)
def test_properties():
np.random.seed(5)
for _ in range(100):
s1 = batoid.Sphere(np.random.uniform(1, 3))
s2 = batoid.Paraboloid(np.random.uniform(1, 3))
sum = batoid.Sum([s1, s2])
do_pickle(sum)
assert s1 is sum.surfaces[0]
assert s2 is sum.surfaces[1]
s3 = batoid.Quadric(np.random.uniform(3, 5), np.random.uniform(-0.1, 0.1))
sum2 = batoid.Sum([s1, s2, s3])
do_pickle(sum2)
assert s1 is sum2.surfaces[0]
assert s2 is sum2.surfaces[1]
assert s3 is sum2.surfaces[2]
do_pickle(sum)
def update(self, deltax):
"""
Update the telescope based on the provided update to the optical state.
Parameters
----------
deltax: aos.state.ZernikeState
The change in the optical state to apply to the telescope.
"""
super().update(deltax)
m2surf = self.optic.itemDict['LSST.M2']
m2residual = batoid.Zernike(deltax.m2zer, R_outer=m2surf.outRadius, R_inner=m2surf.inRadius)
if isinstance(m2surf, batoid.Sum):
m2nominal = m2surf.surfaces[0]
else:
m2nominal = m2surf.surface
self.optic.itemDict['LSST.M2'].surface = batoid.Sum([m2nominal, m2residual])
m1surf = self.optic.itemDict['LSST.M1']
m3surf = self.optic.itemDict['LSST.M3']
# spread across all M1M3
m1m3residual = batoid.Zernike(deltax.m1m3zer, R_outer=m1surf.outRadius, R_inner=m3surf.inRadius)
if isinstance(m1surf, batoid.Sum):
m1nominal = m1surf.surfaces[0]
else:
m1nominal = m1surf.surface
if isinstance(m3surf, batoid.Sum):
m3nominal = m3surf.surfaces[0]
else:
m3nominal = m3surf.surface
self.optic.itemDict['LSST.M1'].surface = batoid.Sum([m1nominal, m1m3residual])
self.optic.itemDict['LSST.M3'].surface = batoid.Sum([m3nominal, m1m3residual])
def test_normal():
rng = np.random.default_rng(577)
for _ in range(100):
s1 = batoid.Sphere(rng.uniform(1, 3))
s2 = batoid.Paraboloid(rng.uniform(1, 3))
sum = batoid.Sum([s1, s2])
x = rng.normal(size=5000)
y = rng.normal(size=5000)
n1 = s1.normal(x, y)
n2 = s2.normal(x, y)
nx = n1[:, 0] / n1[:, 2] + n2[:, 0] / n2[:, 2]
ny = n1[:, 1] / n1[:, 2] + n2[:, 1] / n2[:, 2]
nz = 1. / np.sqrt(nx * nx + ny * ny + 1)
nx *= nz
ny *= nz
normal = np.array([nx, ny, nz]).T
np.testing.assert_allclose(sum.normal(x, y),
normal,
rtol=0,
atol=1e-12)
s3 = batoid.Quadric(rng.uniform(3, 5), rng.uniform(-0.1, 0.1))
sum2 = batoid.Sum([s1, s2, s3])
n3 = s3.normal(x, y)
nx = n1[:, 0] / n1[:, 2] + n2[:, 0] / n2[:, 2] + n3[:, 0] / n3[:, 2]
ny = n1[:, 1] / n1[:, 2] + n2[:, 1] / n2[:, 2] + n3[:, 1] / n3[:, 2]
nz = 1. / np.sqrt(nx * nx + ny * ny + 1)
nx *= nz
ny *= nz
normal = np.array([nx, ny, nz]).T
np.testing.assert_allclose(sum2.normal(x, y),
normal,
rtol=0,
atol=1e-12)
def update(self, deltax):
"""
Update the telescope based on the provided update to the optical state.
Parameters
----------
deltax: aos.state.BendingState
The change in the optical state to apply to the telescope.
"""
super().update(deltax)
self.m1m3res.applyBending(deltax.m1m3modes)
self.m2res.applyBending(deltax.m2modes)
m1m3bicubic = batoid.Bicubic(self.m1m3res.x, self.m1m3res.y, self.m1m3res.surfResidual)
m2bicubic = batoid.Bicubic(self.m2res.x, self.m2res.y, self.m2res.surfResidual)
m1surf = self.optic.itemDict['LSST.M1']
if isinstance(m1surf, batoid.Sum):
m1nominal = m1surf.surfaces[0]
else:
m1nominal = m1surf.surface
self.optic.itemDict['LSST.M1'].surface = batoid.Sum([m1nominal, m1m3bicubic])
m3surf = self.optic.itemDict['LSST.M3']
if isinstance(m1surf, batoid.Sum):
m3nominal = m3surf.surfaces[0]
else:
m3nominal = m3surf.surface
self.optic.itemDict['LSST.M3'].surface = batoid.Sum([m3nominal, m1m3bicubic])
m2surf = self.optic.itemDict['LSST.M2']
if isinstance(m1surf, batoid.Sum):
m2nominal = m2surf.surfaces[0]
else:
m2nominal = m2surf.surface
self.optic.itemDict['LSST.M2'].surface = batoid.Sum([m2nominal, m2bicubic])
def test_withSurface():
telescope = batoid.Optic.fromYaml("HSC.yaml")
rays = batoid.RayVector.asPolar(telescope,
wavelength=620e-9,
theta_x=np.deg2rad(0.1),
theta_y=0.0,
nrad=10,
naz=60)
trays = telescope.trace(rays.copy())
for key, item in telescope.itemDict.items():
if not isinstance(item, batoid.Interface):
continue
# Do a trivial surface replacement
surf2 = batoid.Sum([batoid.Plane(), item.surface])
telescope2 = telescope.withSurface(key, surf2)
assert telescope != telescope2
trays2 = telescope.trace(rays.copy())
rays_allclose(trays, trays2)
def test_add_plane():
rng = np.random.default_rng(5772156)
for _ in range(100):
# Adding a plane should have zero effect on sag or normal vector
s1 = batoid.Sphere(rng.uniform(1, 3))
s2 = batoid.Plane()
sum = batoid.Sum([s1, s2])
x = rng.normal(size=5000)
y = rng.normal(size=5000)
np.testing.assert_allclose(sum.sag(x, y),
s1.sag(x, y),
rtol=0,
atol=1e-12)
np.testing.assert_allclose(
sum.normal(x, y),
s1.normal(x, y),
rtol=0,
atol=1e-12,
)
def test_intersect():
rng = np.random.default_rng(5772)
size = 10_000
for _ in range(100):
s1 = batoid.Sphere(1. / rng.normal(0., 0.2))
s2 = batoid.Paraboloid(rng.uniform(1, 3))
sum = batoid.Sum([s1, s2])
sumCoordSys = batoid.CoordSys(origin=[0, 0, -1])
x = rng.uniform(-1, 1, size=size)
y = rng.uniform(-1, 1, size=size)
z = np.full_like(x, -10.0)
# If we shoot rays straight up, then it's easy to predict the intersection
vx = np.zeros_like(x)
vy = np.zeros_like(x)
vz = np.ones_like(x)
rv = batoid.RayVector(x, y, z, vx, vy, vz)
np.testing.assert_allclose(rv.z, -10.0)
rv2 = batoid.intersect(sum, rv.copy(), sumCoordSys)
assert rv2.coordSys == sumCoordSys
rv2 = rv2.toCoordSys(batoid.CoordSys())
np.testing.assert_allclose(rv2.x, x)
np.testing.assert_allclose(rv2.y, y)
np.testing.assert_allclose(rv2.z,
sum.sag(x, y) - 1,
rtol=0,
atol=1e-12)
# Check default intersect coordTransform
rv2 = rv.copy().toCoordSys(sumCoordSys)
batoid.intersect(sum, rv2)
assert rv2.coordSys == sumCoordSys
rv2 = rv2.toCoordSys(batoid.CoordSys())
np.testing.assert_allclose(rv2.x, x)
np.testing.assert_allclose(rv2.y, y)
np.testing.assert_allclose(rv2.z,
sum.sag(x, y) - 1,
rtol=0,
atol=1e-12)
def parallel_trace_timing(nside=1024, nthread=None, minChunk=None):
if nthread is not None:
print("setting to nthread to {}".format(nthread))
batoid._batoid.setNThread(nthread)
print("Using {} threads".format(batoid._batoid.getNThread()))
if minChunk is not None:
print("setting to minChunk to {}".format(minChunk))
batoid._batoid.setMinChunk(minChunk)
print("Using minChunk of {}".format(batoid._batoid.getMinChunk()))
# 0.3, 0.3 should be in bounds for current wide-field telescopes
theta_x = np.deg2rad(0.3)
theta_y = np.deg2rad(0.3)
dirCos = np.array([theta_x, theta_y, -1.0])
dirCos = batoid.utils.normalized(dirCos)
rays = batoid.circularGrid(
20, 4.2, 0.5,
dirCos[0], dirCos[1], dirCos[2],
nside, nside, 700e-9, 1.0, batoid.ConstMedium(1.0)
)
nrays = len(rays)
print("Tracing {} rays.".format(nrays))
print()
if args.lsst:
fn = os.path.join(batoid.datadir, "LSST", "LSST_r.yaml")
pm = 'LSST.M1'
else:
fn = os.path.join(batoid.datadir, "HSC", "HSC.yaml")
pm = 'SubaruHSC.PM'
config = yaml.load(open(fn))
telescope = batoid.parse.parse_optic(config['opticalSystem'])
# Optionally perturb the primary mirror using Zernike polynomial
if args.perturbZ != 0:
orig = telescope.itemDict[pm].surface
coefs = np.random.normal(size=args.perturbZ+1)*1e-6 # micron perturbations
perturbation = batoid.Zernike(coefs, R_outer=telescope.pupilSize)
telescope.itemDict[pm].surface = batoid.Sum([orig, perturbation])
# Optionally perturb primary mirror using bicubic spline
if args.perturbBC != 0:
orig = telescope.itemDict[pm].surface
xs = np.linspace(-5, 5, 100)
ys = np.linspace(-5, 5, 100)
def f(x, y):
return args.perturbBC*(np.cos(x) + np.sin(y))
zs = f(*np.meshgrid(xs, ys))
bc = batoid.Bicubic(xs, ys, zs)
telescope.itemDict[pm].surface = batoid.Sum([orig, bc])
print("Immutable trace")
t0 = time.time()
rays_out, _ = telescope.trace(rays)
t1 = time.time()
print("{} rays per second".format(int(nrays/(t1-t0))))
print()
print("Trace in place")
t0 = time.time()
telescope.traceInPlace(rays)
t1 = time.time()
print("{} rays per second".format(int(nrays/(t1-t0))))
assert rays == rays_out
if args.plot:
import matplotlib.pyplot as plt
rays.trimVignettedInPlace()
x = rays.x
y = rays.y
x -= np.mean(x)
y -= np.mean(y)
x *= 1e6
y *= 1e6
plt.scatter(x, y, s=1, alpha=0.01)
plt.xlim(np.std(x)*np.r_[-3,3])
plt.ylim(np.std(y)*np.r_[-3,3])
plt.xlabel("x (microns)")
plt.ylabel("y (microns)")
plt.show()
def parallel_trace_timing(nside=1024, nthread=None, minChunk=None):
if nthread is not None:
print("setting nthread to {}".format(nthread))
batoid._batoid.setNThread(nthread)
print("Using {} threads".format(batoid._batoid.getNThread()))
if minChunk is not None:
print("setting minChunk to {}".format(minChunk))
batoid._batoid.setMinChunk(minChunk)
print("Using minChunk of {}".format(batoid._batoid.getMinChunk()))
# 0.3, 0.3 should be in bounds for current wide-field telescopes
dirCos = batoid.utils.gnomicToDirCos(np.deg2rad(0.3), np.deg2rad(0.3))
if args.lsst:
fn = os.path.join(batoid.datadir, "LSST", "LSST_i.yaml")
pm = 'LSST.M1'
elif args.decam:
fn = os.path.join(batoid.datadir, "DECam", "DECam.yaml")
pm = 'BlancoDECam.PM'
else:
fn = os.path.join(batoid.datadir, "HSC", "HSC.yaml")
pm = 'SubaruHSC.PM'
config = yaml.load(open(fn))
telescope = batoid.parse.parse_optic(config['opticalSystem'])
rays = batoid.circularGrid(
telescope.dist, 0.5 * telescope.pupilSize,
0.5 * telescope.pupilObscuration * telescope.pupilSize, dirCos[0],
dirCos[1], -dirCos[2], nside, nside, 750e-9, 1.0, telescope.inMedium)
nrays = len(rays)
print("Tracing {} rays.".format(nrays))
print()
# Optionally perturb the primary mirror using Zernike polynomial
if args.perturbZ != 0:
orig = telescope.itemDict[pm].surface
coefs = np.random.normal(size=args.perturbZ +
1) * 1e-6 # micron perturbations
perturbation = batoid.Zernike(coefs, R_outer=telescope.pupilSize)
telescope.itemDict[pm].surface = batoid.Sum([orig, perturbation])
# Optionally perturb primary mirror using bicubic spline
if args.perturbBC != 0:
orig = telescope.itemDict[pm].surface
rad = telescope.pupilSize / 2 * 1.1
xs = np.linspace(-rad, rad, 100)
ys = np.linspace(-rad, rad, 100)
def f(x, y):
return args.perturbBC * (np.cos(x) + np.sin(y))
zs = f(*np.meshgrid(xs, ys))
bc = batoid.Bicubic(xs, ys, zs)
telescope.itemDict[pm].surface = batoid.Sum([orig, bc])
if args.immutable:
print("Immutable trace")
t0 = time.time()
for _ in range(args.nrepeat):
rays_in = batoid.RayVector(rays)
rays_out, _ = telescope.trace(rays_in)
t1 = time.time()
print("{} rays per second".format(int(nrays * args.nrepeat /
(t1 - t0))))
print()
else:
print("Trace in place")
t0 = time.time()
for _ in range(args.nrepeat):
rays_out = batoid.RayVector(rays)
telescope.traceInPlace(rays_out)
t1 = time.time()
print("{} rays per second".format(int(nrays * args.nrepeat /
(t1 - t0))))
if args.plot:
import matplotlib.pyplot as plt
rays_out.trimVignettedInPlace()
x = rays_out.x
y = rays_out.y
x -= np.mean(x)
y -= np.mean(y)
x *= 1e6
y *= 1e6
plt.scatter(x, y, s=1, alpha=0.01)
plt.xlim(np.std(x) * np.r_[-5, 5])
plt.ylim(np.std(y) * np.r_[-5, 5])
plt.xlabel("x (microns)")
plt.ylabel("y (microns)")
plt.show()
def parallel_trace_timing(args):
print("Using nrad of {:_d}".format(args.nrad))
if args.lsst:
print("Tracing through LSST optics")
telescope = batoid.Optic.fromYaml("LSST_r.yaml")
pm = 'M1'
elif args.decam:
print("Tracing through DECam optics")
telescope = batoid.Optic.fromYaml("DECam.yaml")
pm = 'PM'
else:
print("Tracing through HSC optics")
telescope = batoid.Optic.fromYaml("HSC.yaml")
pm = 'PM'
building = []
for _ in range(args.nrepeat):
t0 = time.time()
rays = batoid.RayVector.asPolar(optic=telescope,
wavelength=620e-9,
theta_x=np.deg2rad(0.3),
theta_y=np.deg2rad(0.3),
inner=0.5 * telescope.pupilSize *
telescope.pupilObscuration,
nrad=args.nrad,
naz=int(2 * np.pi * args.nrad))
t1 = time.time()
building.append(t1 - t0)
building = np.array(building)
nrays = len(rays)
print("Tracing {:_d} rays.".format(nrays))
print(f"Minimum CPU RAM: {2*nrays*74/1024**3:.2f} GB")
print(f"Minimum GPU RAM: {nrays*74/1024**3:.2f} GB")
print()
print()
if args.nrepeat > 1:
print("Generating: {:_} +/- {:_} rays per second".format(
int(np.mean(nrays / building)),
int(np.std(nrays / building) / np.sqrt(args.nrepeat))))
else:
print("Generating: {:_} rays per second".format(
int(nrays / building[0])))
# Optionally perturb the primary mirror using Zernike polynomial
if args.perturbZ != 0:
orig = telescope[pm].surface
coefs = np.random.normal(size=args.perturbZ +
1) * 1e-6 # micron perturbations
perturbation = batoid.Zernike(coefs, R_outer=telescope.pupilSize)
telescope[pm].surface = batoid.Sum([orig, perturbation])
# Optionally perturb primary mirror using bicubic spline
if args.perturbBC != 0:
orig = telescope[pm].surface
rad = telescope.pupilSize / 2 * 1.1
xs = np.linspace(-rad, rad, 100)
ys = np.linspace(-rad, rad, 100)
def f(x, y):
return args.perturbBC * (np.cos(x) + np.sin(y))
zs = f(*np.meshgrid(xs, ys))
bc = batoid.Bicubic(xs, ys, zs)
telescope[pm].surface = batoid.Sum([orig, bc])
copying = []
tracing = []
overall = []
for _ in range(args.nrepeat):
t1 = time.time()
rays_out = rays.copy()
t2 = time.time()
telescope.trace(rays_out)
rays_out.r # force copy back to host if doing gpu
rays_out.v
rays_out.t
rays_out.flux
rays_out.vignetted
rays_out.failed
t3 = time.time()
copying.append(t2 - t1)
tracing.append(t3 - t2)
overall.append(t3 - t1)
copying = np.array(copying)
tracing = np.array(tracing)
overall = np.array(overall)
if args.nrepeat > 1:
print()
print("copying: {:_} +/- {:_} rays per second".format(
int(np.mean(nrays / copying)),
int(np.std(nrays / copying) / np.sqrt(args.nrepeat))))
print()
print("tracing: {:_} +/- {:_} rays per second".format(
int(np.mean(nrays / tracing)),
int(np.std(nrays / tracing) / np.sqrt(args.nrepeat))))
print()
print("overall")
print("-------")
print("{:_} +/- {:_} rays per second".format(
int(np.mean(nrays / overall)),
int(np.std(nrays / overall) / np.sqrt(args.nrepeat))))
print()
else:
print()
print("copying: {:_} rays per second".format(int(nrays / copying)))
print()
print("tracing: {:_} rays per second".format(int(nrays / tracing)))
print()
print("overall")
print("-------")
print("{:_} rays per second".format(int(nrays / overall)))
print()
if args.plot or args.show:
import matplotlib.pyplot as plt
w = ~rays_out.vignetted
x = rays_out.x[w]
y = rays_out.y[w]
x -= np.mean(x)
y -= np.mean(y)
x *= 1e6
y *= 1e6
plt.scatter(x, y, s=1, alpha=0.01)
plt.xlim(np.std(x) * np.r_[-5, 5])
plt.ylim(np.std(y) * np.r_[-5, 5])
plt.xlabel("x (microns)")
plt.ylabel("y (microns)")
plt.savefig("parallel_trace_timing.png")
if args.show:
plt.show()
