In [ ]:
import batoid
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

In [ ]:
telescope = batoid.Optic.fromYaml("HSC.yaml")

In [ ]:
def pupil(thx, thy, nside=512):
    rays = batoid.RayVector.asGrid(
        optic=telescope, lx=10, theta_x=thx, theta_y=thy,
        nx=nside, wavelength=600e-9
    )
    tf = telescope.traceFull(rays)
    w = np.logical_not(tf['D']['out'].vignetted)
    # Project to pupil plane
    p = batoid.Plane()
    p.intersect(tf['POPT2']['in'])
    return tf['POPT2']['in'].x[w], tf['POPT2']['in'].y[w]

In [ ]:
fig = plt.figure(figsize=(12, 12))
ax = fig.add_subplot(111)
ax.scatter(*pupil(np.deg2rad(0.75),0), s=0.1)
ax.set_xlim(-4.2, 4.2)
ax.set_ylim(-4.2, 4.2)
ax.set_aspect(1)
fig.show()

In [ ]:
def spanRange(x, nside=512):
    xmin, xmax = np.min(x), np.max(x)
    xspan = xmax - xmin
    xmin = xmin - 0.8*xspan
    xmax = xmax + 0.8*xspan
    return np.linspace(xmin, xmax, nside)

In [ ]:
def pinhole(thx, thy, nside=256):
    # reset skips
    for item in telescope.itemDict:
        telescope[item].skip = False

    # First, need to determine where on the filter to constrain rays.  We'll use the average position of the 
    # pupil beam that would have intersected the filter.
    rays = batoid.RayVector.asGrid(
        optic=telescope, lx=10, theta_x=thx, theta_y=thy,
        nx=nside, wavelength=600e-9
    )
    tf = telescope.traceFull(rays)
    surface = tf['F_entrance']
    rs = surface['out'].trimVignetted()
    xmean, ymean = np.mean(rs.x), np.mean(rs.y)
    # Now we need to generate a bunch of rays that all pass through the above part of the filter, but over 
    # a range of angles.
    # What is the range of angles for the pupil beam?  
    vx = spanRange(rs.vx, nside=nside)
    vy = spanRange(rs.vy, nside=nside)
    vx, vy = np.meshgrid(vx, vy)
    vz = np.sqrt(1-vx*vx+vy*vy)
    # Now need to make a RayVector with appropriate x,y,vx,vy,...
#     rv = batoid.RayVector([
#         batoid.Ray([xmean, ymean, 0], [vx_, vy_, vz_], 0, 600e-9)
#         for vx_, vy_, vz_ in zip(vx.ravel(), vy.ravel(), vz.ravel())])
    rv = batoid.RayVector.fromArrays(
        xmean*np.ones(nside*nside, dtype=float),
        ymean*np.ones(nside*nside, dtype=float),
        np.zeros(nside*nside, dtype=float),
        vx.ravel(), vy.ravel(), vz.ravel(),
        np.zeros(nside*nside, dtype=float),
        600e-9*np.ones(nside*nside, dtype=float),
        coordSys = surface['out'].coordSys
    )
    # trace forward from filter.  So temporarily skip everything before the filter.
    before_items = ['SubaruHSC.POPT2', 
                    'SubaruHSC.FEU',
                    'SubaruHSC.TopRing',
                    'SubaruHSC.BottomRing',
                    'SubaruHSC.TertiarySpiderFirstPass',
                    'SubaruHSC.PM',
                    'SubaruHSC.TertiarySpiderSecondPass',
                    'SubaruHSC.HSC.WFC.G1',
                    'SubaruHSC.HSC.WFC.G2',
                    'SubaruHSC.HSC.WFC.ADC',
                    'SubaruHSC.HSC.WFC.G3',
                    'SubaruHSC.HSC.WFC.G4',
                    'SubaruHSC.HSC.WFC.G5',
                   ]
    for item in before_items:
        telescope[item].skip = True
    forward_rays = telescope.trace(rv.copy())
    # reset skips
    for item in telescope.itemDict:
        telescope[item].skip = False
    # Now skip everything that happens *after* and including the filter
    after_items = ['SubaruHSC.HSC.CAM.F',
                   'SubaruHSC.HSC.CAM.W',
                   'SubaruHSC.HSC.CAM.D',
                  ]
    for item in after_items:
        telescope[item].skip = True
    rv = batoid.RayVector.fromArrays(
        rv.x, rv.y, rv.z,
        -rv.vx, -rv.vy, -rv.vz,
        rv.t, rv.wavelength,
        coordSys = rv.coordSys
    )
    reverse_rays = telescope.trace(rv.copy(), reverse=True)

    # reset skips
    for item in telescope.itemDict:
        telescope[item].skip = False
        
    w = np.where(np.logical_not(reverse_rays.vignetted))[0]
    return forward_rays.x[w], forward_rays.y[w]

In [ ]:
fig = plt.figure(figsize=(12, 12))
ax = fig.add_subplot(111)
ax.scatter(*pinhole(0,0), s=1)
fig.show()

In [ ]:
def plot(thx, thy):
    fig = plt.figure(figsize=(6, 6))
    ax = fig.add_subplot(111)

    pux, puy = pupil(thx, thy)
    xspan = np.max(pux) - np.min(pux)
    yspan = np.max(puy) - np.min(puy)
    span = max(xspan, yspan)
    pux = (pux - np.mean(pux))/span
    puy = (puy - np.mean(puy))/span
    
    phx, phy = pinhole(thx, thy)
    xspan = np.max(phx) - np.min(phx)
    yspan = np.max(phy) - np.min(phy)
    span = max(xspan, yspan)

    phx = -(phx - np.mean(phx))/span
    phy = -(phy - np.mean(phy))/span

    ax.scatter(pux, puy, s=2, alpha=0.1, c='r', label='pupil')
    ax.scatter(phx, phy, s=2, alpha=0.2, c='b', label='pinhole')
    ax.legend()
    fig.show()

In [ ]:
plot(0, np.deg2rad(0.75))

In [ ]:
def both(thx, thy):
    pux, puy = pupil(thx, thy)
    xspan = np.max(pux) - np.min(pux)
    yspan = np.max(puy) - np.min(puy)
    span = max(xspan, yspan)
    pux = (pux - np.mean(pux))/span
    puy = (puy - np.mean(puy))/span

    phx, phy = pinhole(thx, thy)
    xspan = np.max(phx) - np.min(phx)
    yspan = np.max(phy) - np.min(phy)
    span = max(xspan, yspan)
    phx = -(phx - np.mean(phx))/span
    phy = -(phy - np.mean(phy))/span

    return pux, puy, phx, phy

In [ ]:
fig, axes = plt.subplots(nrows=2, ncols=2, figsize=(8,8))
for ax, thy in zip(axes.ravel(), [0.0, 0.25, 0.5, 0.75]):
    pux, puy, phx, phy = both(0.0, np.deg2rad(thy))

    ax.scatter(pux, puy, s=2, alpha=0.1, c='r', label='pupil')
    ax.scatter(phx, phy, s=2, alpha=0.2, c='b', label='pinhole')
    ax.set_title(r"$\theta_y$ = {:5.2f}".format(thy))
    ax.legend(loc="upper right")
fig.show()

In [ ]: