1#!/usr/local/bin/python3.8
2# This file is part of the Astrometry.net suite.
3# Licensed under a 3-clause BSD style license - see LICENSE
4
5import sys
6from optparse import OptionParser
7
8import matplotlib
9matplotlib.use('Agg')
10
11from pylab import *
12from numpy import *
13
14from astrometry.libkd import spherematch
15
16def plotshift(ixy, rxy, dcell=50, ncells=18, outfn=None, W=None, H=None, hist=False,
17              nhistbins=21):
18              #histbinsize=None):
19    # correspondences we could have hit...
20    radius = dcell * sqrt(2.)
21    #print 'ixy', ixy.shape
22    #print 'rxy', rxy.shape
23    assert((len(rxy) == 0) or (rxy.shape[1] == 2))
24    assert((len(ixy) == 0) or (ixy.shape[1] == 2))
25
26    ix = ixy[:,0]
27    iy = ixy[:,1]
28
29    if W is None:
30        W = max(ix)
31    if H is None:
32        H = max(iy)
33
34    if len(rxy):
35        keep = (rxy[:,0] > -dcell) * (rxy[:,0] < W+dcell) * (rxy[:,1] > -dcell) * (rxy[:,1] < H+dcell)
36        rxy = rxy[keep]
37    #print 'Cut to %i ref sources in range' % len(rxy)
38
39    cellsize = sqrt(W * H / ncells)
40    nw = int(round(W / cellsize))
41    nh = int(round(H / cellsize))
42    #print 'Grid cell size', cellsize
43    #print 'N cells', nw, 'x', nh
44    edgesx = linspace(0, W, nw+1)
45    edgesy = linspace(0, H, nh+1)
46    #print 'Edges:'
47    #print '     x:', edgesx
48    #print '     y:', edgesy
49    if len(rxy) == 0:
50        binx = array([])
51        biny = array([])
52    else:
53        binx = digitize(rxy[:,0], edgesx)
54        biny = digitize(rxy[:,1], edgesy)
55    binx = clip(binx - 1, 0, nw-1)
56    biny = clip(biny - 1, 0, nh-1)
57    bin = biny * nw + binx
58
59    clf()
60
61    for i in range(nh):
62        for j in range(nw):
63            thisbin = i * nw + j
64            R = (bin == thisbin)
65            #print 'cell %i, %i' % (j, i)
66            #print '%i ref sources' % sum(R)
67            matchdx = []
68
69            if sum(R) > 0:
70                (inds,dists) = spherematch.match(rxy[R,:], ixy, radius)
71                #print 'Found %i matches within %g pixels' % (len(dists), radius)
72                ri = inds[:,0]
73                # un-cut ref inds...
74                ri = (flatnonzero(R))[ri]
75                ii = inds[:,1]
76                matchx    = rxy[ri,0]
77                matchy    = rxy[ri,1]
78                matchdx = ix[ii] - matchx
79                matchdy = iy[ii] - matchy
80                #print 'All matches:'
81                #for dx,dy in zip(matchdx,matchdy):
82                #    print '     %.1f, %.1f' % (dx,dy)
83                ok = (matchdx >= -dcell) * (matchdx <= dcell) * (matchdy >= -dcell) * (matchdy <= dcell)
84                matchdx = matchdx[ok]
85                matchdy = matchdy[ok]
86                #print 'Cut to %i within %g x %g square' % (sum(ok), dcell*2, dcell*2)
87                #print 'Cut matches:'
88                #for dx,dy in zip(matchdx,matchdy):
89                #    print '     %.1f, %.1f' % (dx,dy)
90
91            # Subplot places plots left-to-right, TOP-to-BOTTOM.
92            subplot(nh, nw, 1 + ((nh - i - 1)*nw + j))
93
94            if len(matchdx) > 0:
95                #plot(matchdx, matchdy, 'ro', mec='r', mfc='r', ms=5, alpha=0.2)
96                #plot(matchdx, matchdy, 'ro', mec='r', mfc='none', ms=5, alpha=0.2)
97                if hist:
98                    #if histbinsize is None:
99                    #    histbinsize = dcell / 10.
100                    edges = linspace(-dcell, dcell, nhistbins+1)
101                    (H,xe,ye) = histogram2d(matchdx, matchdy, bins=(edges,edges))
102                    imshow(H.T, extent=(min(xe), max(xe), min(ye), max(ye)),
103                           aspect='auto', origin='lower', interpolation='nearest')
104                    text(dcell, dcell, '%i' % H.max(), color='y',
105                         horizontalalignment='right', verticalalignment='top')
106
107                else:
108                    plot(matchdx, matchdy, 'r.', alpha=0.3)
109
110            if hist:
111                axhline(0, color='b', alpha=0.8)
112                axvline(0, color='b', alpha=0.8)
113            else:
114                axhline(0, color='k', alpha=0.5)
115                axvline(0, color='k', alpha=0.5)
116            if i == 0 and j == 0:
117                xticks([-dcell,0,dcell])
118                yticks([-dcell,0,dcell])
119            else:
120                xticks([],[])
121                yticks([],[])
122            axis('scaled')
123            axis([-dcell, dcell, -dcell, dcell])
124    if outfn is not None:
125        #print 'Saving', outfn
126        savefig(outfn)
127
128
129if __name__ == '__main__':
130    from astrometry.util.fits import fits_table
131
132    parser = OptionParser('usage: %prog [options] <image xy> <reference xy> <plot name>')
133    parser.add_option('-X', dest='xcol', help='Name of X column in image table')
134    parser.add_option('-Y', dest='ycol', help='Name of Y column in image table')
135    parser.add_option('-N', dest='nimage', type='int', help='Cut to the first N image sources')
136    parser.add_option('-x', dest='rxcol', help='Name of X column in reference table')
137    parser.add_option('-y', dest='rycol', help='Name of Y column in reference table')
138    parser.add_option('-n', dest='nref', type='int', help='Cut to the first N reference sources')
139    parser.add_option('-c', dest='cells', type='int', help='Approx. number of pieces to cut image into (default:18)')
140    parser.add_option('-s', dest='cellsize', type='int', help='Search radius, in pixels (default 50)')
141    parser.set_defaults(xcol='X', ycol='Y', nimage=0, cells=0, cellsize=0, rxcol='X', rycol='Y', nref=0)
142    opt,args = parser.parse_args()
143
144    if len(args) != 3:
145        parser.print_help()
146        sys.exit(-1)
147
148    imxy = fits_table(args[0])
149    refxy = fits_table(args[1])
150    outfn = args[2]
151
152    kwargs = {}
153    if opt.cells:
154        kwargs['ncells'] = opt.cells
155    if opt.cellsize:
156        kwargs['dcell'] = opt.cellsize
157
158    ix = imxy.getcolumn(opt.xcol)
159    iy = imxy.getcolumn(opt.ycol)
160    ixy = vstack((ix,iy)).T
161    if opt.nimage:
162        ixy = ixy[:opt.nimage,:]
163
164    rx = refxy.getcolumn(opt.rxcol)
165    ry = refxy.getcolumn(opt.rycol)
166    rxy = vstack((rx,ry)).T
167    if opt.nref:
168        rxy = rxy[:opt.nref,:]
169
170    plotshift(ixy, rxy, outfn=outfn, **kwargs)
171