def test_lattice_coords(): np.random.seed(5) # Check 1D for _ in np.arange(10): N = np.random.randint(1, 2000) arr = np.ones((N, )) primitiveVector = np.random.uniform(-1.0, 1.0) lattice = batoid.Lattice(arr, primitiveVector) np.testing.assert_allclose( np.squeeze(lattice.coords), np.arange(-(N // 2), -(-N // 2)) * primitiveVector) # Check 2D for _ in np.arange(10): N1 = np.random.randint(1, 200) N2 = np.random.randint(1, 200) arr = np.ones((N1, N2)) pv1 = np.random.uniform(-1.0, 1.0, size=2) pv2 = np.random.uniform(-1.0, 1.0, size=2) lattice = batoid.Lattice(arr, np.vstack([pv1, pv2])) for _ in np.arange(100): i = np.random.randint(0, N1) j = np.random.randint(0, N2) np.testing.assert_allclose(lattice.coords[i, j], (i - N1 // 2) * pv1 + (j - N2 // 2) * pv2) # Check 3D for _ in np.arange(10): N1 = np.random.randint(1, 20) N2 = np.random.randint(1, 20) N3 = np.random.randint(1, 20) arr = np.ones((N1, N2, N3)) pv1 = np.random.uniform(-1.0, 1.0, size=3) pv2 = np.random.uniform(-1.0, 1.0, size=3) pv3 = np.random.uniform(-1.0, 1.0, size=3) lattice = batoid.Lattice(arr, np.vstack([pv1, pv2, pv3])) with np.printoptions(threshold=20**3): do_pickle(lattice) for __ in np.arange(100): i = np.random.randint(0, N1) j = np.random.randint(0, N2) k = np.random.randint(0, N3) np.testing.assert_allclose( lattice.coords[i, j, k], (i - N1 // 2) * pv1 + (j - N2 // 2) * pv2 + (k - N3 // 2) * pv3)
def test_lattice_coords(): np.random.seed(5) # Check 1D for _ in np.arange(10): N = 2 * np.random.randint(2, 1024) arr = np.ones((N, )) primitiveVector = np.random.uniform(-1.0, 1.0) lattice = batoid.Lattice(arr, primitiveVector) np.testing.assert_allclose(np.squeeze(lattice.coords), np.arange(-N / 2, N / 2) * primitiveVector) # Check 2D for _ in np.arange(10): N1 = 2 * np.random.randint(2, 1024) N2 = 2 * np.random.randint(2, 1024) arr = np.ones((N1, N2)) pv1 = np.random.uniform(-1.0, 1.0, size=2) pv2 = np.random.uniform(-1.0, 1.0, size=2) lattice = batoid.Lattice(arr, np.vstack([pv1, pv2])) for _ in np.arange(100): i = np.random.randint(0, N1) j = np.random.randint(0, N2) np.testing.assert_allclose(lattice.coords[i, j], (i - N1 / 2) * pv1 + (j - N2 / 2) * pv2) # Check 3D for _ in np.arange(10): N1 = 2 * np.random.randint(2, 64) N2 = 2 * np.random.randint(2, 64) N3 = 2 * np.random.randint(2, 64) arr = np.ones((N1, N2, N3)) pv1 = np.random.uniform(-1.0, 1.0, size=3) pv2 = np.random.uniform(-1.0, 1.0, size=3) pv3 = np.random.uniform(-1.0, 1.0, size=3) lattice = batoid.Lattice(arr, np.vstack([pv1, pv2, pv3])) coords = lattice.coords for __ in np.arange(100): i = np.random.randint(0, N1) j = np.random.randint(0, N2) k = np.random.randint(0, N3) np.testing.assert_allclose(lattice.coords[i, j, k], (i - N1 / 2) * pv1 + (j - N2 / 2) * pv2 + (k - N3 / 2) * pv3)
def fftPSF(optic, theta_x, theta_y, wavelength, nx=32, projection='postel', pad_factor=2, _addedWF=None): """Compute PSF using FFT. Parameters ---------- optic : batoid.Optic Optic for which to compute wavefront. theta_x, theta_y : float Field angle in radians wavelength : float, optional Wavelength in meters nx : int, optional Size of ray grid to generate to compute wavefront. Default: 32 projection : {'postel', 'zemax', 'gnomonic', 'stereographic', 'lambert', 'orthographic'} Projection used to convert field angle to direction cosines. pad_factor : int, optional Factor by which to pad pupil array. Default: 2 Returns ------- psf : batoid.Lattice A batoid.Lattice object containing the relative PSF values and the primitive lattice vectors of the focal plane grid. """ L = optic.pupilSize * pad_factor # im_dtheta = wavelength / L wf = wavefront(optic, theta_x, theta_y, wavelength, nx, projection=projection, lattice=True, _addedWF=_addedWF) wfarr = wf.array pad_size = nx * pad_factor expwf = np.zeros((pad_size, pad_size), dtype=np.complex128) start = pad_size // 2 - nx // 2 stop = pad_size // 2 + nx // 2 expwf[start:stop, start:stop][~wfarr.mask] = np.exp(2j * np.pi * wfarr[~wfarr.mask]) psf = np.abs(np.fft.fftshift(np.fft.fft2(expwf)))**2 primitiveU = wf.primitiveVectors primitiveK = dkdu(optic, theta_x, theta_y, wavelength, projection=projection).dot(primitiveU) primitiveX = np.vstack( reciprocalLatticeVectors(primitiveK[0], primitiveK[1], pad_size)) return batoid.Lattice(psf, primitiveX)
def test_ne(): rng = np.random.default_rng(57) N1 = rng.integers(1, 5) N2 = rng.integers(1, 5) N3 = rng.integers(1, 5) arr = np.ones((N1, N2, N3)) pv1 = rng.uniform(-1.0, 1.0, size=3) pv2 = rng.uniform(-1.0, 1.0, size=3) pv3 = rng.uniform(-1.0, 1.0, size=3) lattice1 = batoid.Lattice(arr, np.vstack([pv1, pv2, pv3])) lattice2 = batoid.Lattice(arr[..., 0], np.vstack([pv1, pv2, pv3])[:2, :2]) lattice3 = batoid.Lattice(arr, 2 * np.vstack([pv1, pv2, pv3])) lattice4 = batoid.Lattice(2 * arr, np.vstack([pv1, pv2, pv3])) objs = [batoid.CoordSys(), lattice1, lattice2, lattice3, lattice4] all_obj_diff(objs)
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 simulateDonut(self, optic, fieldx, fieldy): """ Simulate a donut image by raytracing photons through optic. Parameters ---------- optic: batoid.Optic The optic to raytrace through. fieldx: float The x field position in degrees. theta_y: float The y field position in degrees. Returns ------- batoid.Lattice The donut image. """ thetax, thetay = np.deg2rad([fieldx, fieldy]) flux = 1 xcos, ycos, zcos = batoid.utils.gnomonicToDirCos(thetax, thetay) rays = batoid.uniformCircularGrid( optic.backDist, optic.pupilSize / 2, optic.pupilSize * optic.pupilObscuration / 2, xcos, ycos, zcos, self.nphot, self.wavelength, flux, optic.inMedium) optic.traceInPlace(rays) rays.trimVignettedInPlace() xcent, ycent = np.mean(rays.x), np.mean(rays.y) width = self.crop * self.pix xedges = np.linspace(xcent - width / 2, xcent + width / 2, self.crop + 1) yedges = np.linspace(ycent - width / 2, ycent + width / 2, self.crop + 1) # flip here because 1st dimension corresponds to y-dimension in bitmap image result, _, _ = np.histogram2d(rays.y, rays.x, bins=[yedges, xedges]) primitiveX = np.array([[self.pix, 0], [0, self.pix]]) return batoid.Lattice(result, primitiveX)
def fftPSF(optic, theta_x, theta_y, wavelength, nx=32, pad_factor=2): """Compute PSF using FFT. 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 pad_factor : int, optional Factor by which to pad pupil array. Default: 2 Returns ------- psf : batoid.Lattice A batoid.Lattice object containing the relative PSF values and the primitive lattice vectors of the focal plane grid. """ L = optic.pupilSize*pad_factor # im_dtheta = wavelength / L wf = wavefront(optic, theta_x, theta_y, wavelength, nx) wfarr = wf.array pad_size = nx*pad_factor expwf = np.zeros((pad_size, pad_size), dtype=np.complex128) start = pad_size//2-nx//2 stop = pad_size//2+nx//2 expwf[start:stop, start:stop][~wfarr.mask] = np.exp(2j*np.pi*wfarr[~wfarr.mask]) psf = np.abs(np.fft.fftshift(np.fft.fft2(expwf)))**2 primitiveU = wf.primitiveVectors primitiveK = dkdu(optic, theta_x, theta_y, wavelength).dot(primitiveU) primitiveX = np.vstack(reciprocalLatticeVectors(primitiveK[0], primitiveK[1], pad_size)) return batoid.Lattice(psf, primitiveX)
def huygensPSF(optic, theta_x=None, theta_y=None, wavelength=None, nx=None, projection='postel', dx=None, dy=None, nxOut=None): r"""Compute a PSF via the Huygens construction. Parameters ---------- optic : batoid.Optic Optical system theta_x, theta_y : float Field angle in radians wavelength : float, optional Wavelength in meters nx : int, optional Size of ray grid to use. projection : {'postel', 'zemax', 'gnomonic', 'stereographic', 'lambert', 'orthographic'} Projection used to convert field angle to direction cosines. dx, dy : float, optional Lattice scales to use for PSF evaluation locations. Default, use fftPSF lattice. Returns ------- psf : batoid.Lattice The PSF. Notes ----- The Huygens construction is to evaluate the PSF as I(x) \propto \Sum_u exp(i phi(u)) exp(i k(u).r) The u are assumed to uniformly sample the entrance pupil, but not include any rays that get vignetted before they reach the focal plane. The phis are the phases of the exit rays evaluated at a single arbitrary time. The k(u) indicates the conversion of the uniform entrance pupil samples into nearly (though not exactly) uniform samples in k-space of the output rays. The output locations where the PSF is evaluated are governed by dx, dy and nx. If dx and dy are None, then the same lattice as in fftPSF will be used. If dx and dy are scalars, then a lattice with primitive vectors [dx, 0] and [0, dy] will be used. If dx and dy are 2-vectors, then those will be the primitive vectors of the output lattice. """ from numbers import Real if dx is None: primitiveU = np.array([[optic.pupilSize / nx, 0], [0, optic.pupilSize / nx]]) primitiveK = dkdu(optic, theta_x, theta_y, wavelength, projection=projection).dot(primitiveU) pad_factor = 2 primitiveX = np.vstack( reciprocalLatticeVectors(primitiveK[0], primitiveK[1], pad_factor * nx)) elif isinstance(dx, Real): if dy is None: dy = dx primitiveX = np.vstack([[dx, 0], [0, dy]]) pad_factor = 1 else: primitiveX = np.vstack([dx, dy]) pad_factor = 1 if nxOut is None: nxOut = nx dirCos = fieldToDirCos(theta_x, theta_y, projection=projection) rays = batoid.rayGrid(optic.backDist, optic.pupilSize, dirCos[0], dirCos[1], dirCos[2], nx, wavelength=wavelength, flux=1, medium=optic.inMedium, lattice=True) amplitudes = np.zeros((nxOut * pad_factor, nxOut * pad_factor), dtype=np.complex128) out = batoid.Lattice( np.zeros((nxOut * pad_factor, nxOut * pad_factor), dtype=float), primitiveX) rays = optic.trace(rays) rays.trimVignetted() # Need transpose to conform to numpy [y,x] ordering convention xs = out.coords[..., 0].T + np.mean(rays.x) ys = out.coords[..., 1].T + np.mean(rays.y) zs = np.zeros_like(xs) points = np.concatenate([aux[..., None] for aux in (xs, ys, zs)], axis=-1) time = rays[0].t for idx in np.ndindex(amplitudes.shape): amplitudes[idx] = rays.sumAmplitude(points[idx], time) out.array = np.abs(amplitudes)**2 return out
def wavefront(optic, theta_x, theta_y, wavelength, nx=32, projection='postel', sphereRadius=None, lattice=False, _addedWF=None): """Compute wavefront. Parameters ---------- optic : batoid.Optic Optic for which to compute wavefront. theta_x, theta_y : float Field angle in radians wavelength : float, optional Wavelength in meters nx : int, optional Size of ray grid to generate to compute wavefront. Default: 32 projection : {'postel', 'zemax', 'gnomonic', 'stereographic', 'lambert', 'orthographic'} Projection used to convert field angle to direction cosines. sphereRadius : float, optional Radius of reference sphere in meters. If None, then use optic.sphereRadius. lattice : bool, optional If true, then decenter the grid so it spans (-N/2, N/2+1), as appropriate for Fourier transforms. 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 = fieldToDirCos(theta_x, theta_y, projection=projection) rays = batoid.rayGrid(optic.backDist, optic.pupilSize, dirCos[0], dirCos[1], dirCos[2], nx, wavelength, 1.0, optic.inMedium, lattice=lattice) if sphereRadius is None: sphereRadius = optic.sphereRadius optic.trace(rays) 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. targetCoordSys = rays.coordSys.shiftLocal(point + np.array([0, 0, sphereRadius])) rays.toCoordSys(targetCoordSys) sphere = batoid.Sphere(-sphereRadius) sphere.intersect(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) if _addedWF is not None: arr += _addedWF primitiveVectors = np.vstack([[optic.pupilSize / nx, 0], [0, optic.pupilSize / nx]]) return batoid.Lattice(arr, primitiveVectors)
def huygensPSF(optic, theta_x, theta_y, wavelength, projection='postel', nx=None, dx=None, dy=None, nxOut=None, reference='mean'): r"""Compute a PSF via the Huygens construction. Parameters ---------- optic : batoid.Optic Optical system theta_x, theta_y : float Field angle in radians wavelength : float Wavelength in meters projection : {'postel', 'zemax', 'gnomonic', 'stereographic', 'lambert', 'orthographic'} Projection used to convert field angle to direction cosines. nx : int, optional Size of ray grid to use. dx, dy : float, optional Lattice scales to use for PSF evaluation locations. Default, use fftPSF lattice. nxOut : int, optional Size of the output lattice. Default is to use nx. reference : {'chief', 'mean'} If 'chief', then center the output lattice where the chief ray intersects the focal plane. If 'mean', then center at the mean non-vignetted ray intersection. Returns ------- psf : batoid.Lattice The PSF. Notes ----- The Huygens construction is to evaluate the PSF as .. math:: I(x) \propto \sum_u \exp(i \phi(u)) \exp(i k(u) \cdot r) The :math:`u` are assumed to uniformly sample the entrance pupil, but not include any rays that get vignetted before they reach the focal plane. The :math:`\phi` s are the phases of the exit rays evaluated at a single arbitrary time. The :math:`k(u)` indicates the conversion of the uniform entrance pupil samples into nearly (though not exactly) uniform samples in k-space of the output rays. The output locations where the PSF is evaluated are governed by ``dx``, ``dy``, and ``nx``. If ``dx`` and ``dy`` are None, then the same lattice as in fftPSF will be used. If ``dx`` and ``dy`` are scalars, then a lattice with primitive vectors ``[dx, 0]`` and ``[0, dy]`` will be used. If ``dx`` and ``dy`` are 2-vectors, then those will be the primitive vectors of the output lattice. """ from numbers import Real if dx is None: if (nx % 2) == 0: primitiveU = np.array([[optic.pupilSize / (nx - 2), 0], [0, optic.pupilSize / (nx - 2)]]) else: primitiveU = np.array([[optic.pupilSize / (nx - 1), 0], [0, optic.pupilSize / (nx - 1)]]) primitiveK = dkdu(optic, theta_x, theta_y, wavelength, projection=projection).dot(primitiveU) pad_factor = 2 primitiveX = np.vstack( reciprocalLatticeVectors(primitiveK[0], primitiveK[1], pad_factor * nx)) elif isinstance(dx, Real): if dy is None: dy = dx primitiveX = np.vstack([[dx, 0], [0, dy]]) pad_factor = 1 else: primitiveX = np.vstack([dx, dy]) pad_factor = 1 if nxOut is None: nxOut = nx dirCos = fieldToDirCos(theta_x, theta_y, projection=projection) rays = batoid.RayVector.asGrid(optic=optic, wavelength=wavelength, dirCos=dirCos, nx=nx) amplitudes = np.zeros((nxOut * pad_factor, nxOut * pad_factor), dtype=np.complex128) out = batoid.Lattice( np.zeros((nxOut * pad_factor, nxOut * pad_factor), dtype=float), primitiveX) optic.traceInPlace(rays) if reference == 'mean': w = np.where(1 - rays.vignetted)[0] point = np.mean(rays.r[w], axis=0) elif reference == 'chief': cridx = (nx // 2) * nx + nx // 2 if (nx % 2) == 0 else (nx * nx - 1) // 2 point = rays[cridx].r rays.trimVignettedInPlace() # Need transpose to conform to numpy [y,x] ordering convention xs = out.coords[..., 0].T + point[0] ys = out.coords[..., 1].T + point[1] zs = np.zeros_like(xs) points = np.concatenate([aux[..., None] for aux in (xs, ys, zs)], axis=-1) time = rays[0].t for idx in np.ndindex(amplitudes.shape): amplitudes[idx] = rays.sumAmplitude(points[idx], time) out.array = np.abs(amplitudes)**2 return out
def fftPSF(optic, theta_x, theta_y, wavelength, projection='postel', nx=32, pad_factor=2, sphereRadius=None, reference='mean', _addedWF=None): """Compute PSF using FFT. Parameters ---------- optic : batoid.Optic Optical system theta_x, theta_y : float Field angle in radians wavelength : float Wavelength in meters projection : {'postel', 'zemax', 'gnomonic', 'stereographic', 'lambert', 'orthographic'} Projection used to convert field angle to direction cosines. nx : int, optional Size of ray grid to use. pad_factor : int, optional Factor by which to pad pupil array. Default: 2 sphereRadius : float, optional The radius of the reference sphere. Nominally this should be set to the distance to the exit pupil, though the calculation is usually not very sensitive to this. Many of the telescopes that come with batoid have values for this set in their yaml files, which will be used if this is None. reference : {'chief', 'mean'} If 'chief', then center the output lattice where the chief ray intersects the focal plane. If 'mean', then center at the mean non-vignetted ray intersection. Returns ------- psf : batoid.Lattice A batoid.Lattice object containing the relative PSF values and the primitive lattice vectors of the focal plane grid. """ wf = wavefront(optic, theta_x, theta_y, wavelength, nx=nx, projection=projection, sphereRadius=sphereRadius, reference=reference) wfarr = wf.array pad_size = nx * pad_factor expwf = np.zeros((pad_size, pad_size), dtype=np.complex128) start = pad_size // 2 - nx // 2 stop = pad_size // 2 + nx // 2 expwf[start:stop, start:stop][~wfarr.mask] = \ np.exp(2j*np.pi*wfarr[~wfarr.mask]) psf = np.abs(np.fft.fftshift(np.fft.fft2(expwf)))**2 primitiveU = wf.primitiveVectors primitiveK = dkdu(optic, theta_x, theta_y, wavelength, projection=projection).dot(primitiveU) primitiveX = np.vstack( reciprocalLatticeVectors(primitiveK[0], primitiveK[1], pad_size)) return batoid.Lattice(psf, primitiveX)
def wavefront(optic, theta_x, theta_y, wavelength, projection='postel', nx=32, sphereRadius=None, reference='mean'): """Compute wavefront. Parameters ---------- optic : batoid.Optic Optical system theta_x, theta_y : float Field angle in radians wavelength : float Wavelength in meters projection : {'postel', 'zemax', 'gnomonic', 'stereographic', 'lambert', 'orthographic'} Projection used to convert field angle to direction cosines. nx : int, optional Size of ray grid to use. sphereRadius : float, optional The radius of the reference sphere. Nominally this should be set to the distance to the exit pupil, though the calculation is usually not very sensitive to this. Many of the telescopes that come with batoid have values for this set in their yaml files, which will be used if this is None. reference : {'chief', 'mean'} If 'chief', then center the output lattice where the chief ray intersects the focal plane. If 'mean', then center at the mean non-vignetted ray intersection. 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 = fieldToDirCos(theta_x, theta_y, projection=projection) rays = batoid.RayVector.asGrid(optic=optic, wavelength=wavelength, nx=nx, dirCos=dirCos) if sphereRadius is None: sphereRadius = optic.sphereRadius optic.traceInPlace(rays) if reference == 'mean': w = np.where(1 - rays.vignetted)[0] point = np.mean(rays.r[w], axis=0) elif reference == 'chief': cridx = (nx // 2) * nx + nx // 2 if (nx % 2) == 0 else (nx * nx - 1) // 2 point = rays[cridx].r # Place vertex of reference sphere one radius length away from the # intersection point. So transform our rays into that coordinate system. targetCoordSys = rays.coordSys.shiftLocal(point + np.array([0, 0, sphereRadius])) rays.toCoordSysInPlace(targetCoordSys) sphere = batoid.Sphere(-sphereRadius) sphere.intersectInPlace(rays) if reference == 'mean': w = np.where(1 - rays.vignetted)[0] t0 = np.mean(rays.t[w]) elif reference == 'chief': t0 = rays[cridx].t arr = np.ma.masked_array((t0 - rays.t) / wavelength, mask=rays.vignetted).reshape(nx, nx) if (nx % 2) == 0: primitiveVectors = np.vstack([[optic.pupilSize / (nx - 2), 0], [0, optic.pupilSize / (nx - 2)]]) else: primitiveVectors = np.vstack([[optic.pupilSize / (nx - 1), 0], [0, optic.pupilSize / (nx - 1)]]) return batoid.Lattice(arr, primitiveVectors)