/
getmask.py
158 lines (136 loc) · 4.82 KB
/
getmask.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
import numpy as np
import astropy.units as auni
import astropy.coordinates as acoo
import sqlutilpy
import astropy.wcs as pywcs
import healpy
from cffi import FFI
import _filler
import astropy.io.fits as pyfits
import multiprocessing as mp
ffi = FFI()
WSDB_HOST = open('WSDB_HOST').read()
def getrad(Gmag):
""" formula approximately matching the radius where
fibermag ~=20
"""
return 600. * 1.4**(-Gmag)
def getangle(x, y):
""" for 4 points x,y determine
the rotation angle to approximately align
the coordinate system with one of the axes
"""
pos = np.argmax(y)
pos1 = (pos + 3) % 4
angle = np.arctan2(y[pos] - y[pos1], x[pos] - x[pos1])
return float(angle)
def getwcs0(ra0, dec0, crpix1, crpix2, angle, step):
""" get the wcs dictionary"""
return dict(CRVAL1=ra0,
CRVAL2=dec0,
CRPIX1=crpix1,
CRPIX2=crpix2,
CD1_1=np.cos(angle) * step,
CD2_2=np.cos(angle) * step,
CD1_2=-np.sin(angle) * step,
CD2_1=np.sin(angle) * step,
CTYPE1='RA---TAN',
CTYPE2='DEC--TAN')
def getwcs(nside, hpx, step_asec=1):
"""
get the angle and pixel size for a given healpix
"""
expand = 1.05 # extra expansion
step = step_asec / 3600
ra0, dec0 = healpy.pix2ang(nside, hpx, lonlat=True, nest=True)
xyz = healpy.boundaries(nside, hpx, step=1, nest=True)
edges = np.array([healpy.vec2ang(_, lonlat=True) for _ in xyz.T])
wc = pywcs.WCS(getwcs0(ra0, dec0, 1, 1, 0, 1))
x, y = wc.world_to_pixel(
acoo.SkyCoord(ra=edges[:, 0] * auni.deg, dec=edges[:, 1] * auni.deg))
angle = getangle(x, y)
wc = pywcs.WCS(getwcs0(ra0, dec0, 1, 1, angle, step))
x, y = wc.world_to_pixel(
acoo.SkyCoord(ra=edges[:, 0] * auni.deg, dec=edges[:, 1] * auni.deg))
npix = 2 * int(np.ceil(expand * max(np.max(np.abs(x)), np.max(np.abs(y)))))
# ensure rectangular
return angle, npix
def getmask_hpx(nside, hpx, ofname):
"""
Create the fits sky mask
Parameters
----------
nside: int
hpx: int
ofname: string
"""
ra0, dec0 = healpy.pix2ang(nside, hpx, lonlat=True, nest=True)
step_asec = 1
angle, npix = getwcs(nside, hpx, step_asec=step_asec)
ima, wcdict = getmask(ra0, dec0, step_asec, npix, angle=angle)
header = pyfits.Header()
for k, v in wcdict.items():
header[k] = v
header['IMAGEW'] = npix
header['IMAGEH'] = npix
pyfits.PrimaryHDU(data=ima, header=header).writeto(ofname, overwrite=True)
def doall(nthreads=36):
"""
Process all the hpxes
"""
nside = 64
pool = mp.Pool(nthreads)
res = []
for hpx in range(12 * nside * nside):
res.append((hpx,
pool.apply_async(
getmask_hpx,
(nside, hpx, 'data/gaia_skymap-%05d.fits.gz' % hpx))))
for r in res:
hpx, r = r
r.get()
print(hpx)
def getmask(ra0, dec0, step_asec, npix, angle=0, binning=4, epoch=2022):
"""
Return the integer mask and WCS object for a gaia based sky mask
Arguments
ra0,dec0 center
pixel size in arcsec
number of pixels
angle is the prescribed rotation angle of the projection
"""
minrad_asec, maxrad_asec = 1.5, 1000
step = step_asec / 3600.
wcdict = getwcs0(ra0, dec0, (npix + 1) / 2, (npix + 1) / 2, angle, step)
assert (binning in [1, 2, 4])
wc = pywcs.WCS(wcdict)
pad = maxrad_asec / 3600. # padding in degrees for bright stars
maxrad = npix * np.sqrt(2) * step + pad
ra, dec, pmra, pmdec, refep, g = sqlutilpy.get(
'''select ra,dec,pmra,pmdec,ref_epoch,phot_g_mean_mag from gaia_edr3.gaia_source
where q3c_radial_query(ra,dec,%f,%f,%f)''' % (ra0, dec0, maxrad),
host=WSDB_HOST)
bad = ~np.isfinite(pmra)
pmra[bad] = 0
pmdec[bad] = 0
ra = ra + pmra / 3600e3 / np.cos(np.deg2rad(dec)) * (epoch - refep)
dec = dec + pmdec / 3600e3 * (epoch - refep)
rads = np.clip(getrad(g), minrad_asec, maxrad_asec)
rads = rads / step_asec # in pixels now
nstars = len(ra)
ys, xs = wc.world_to_pixel(
acoo.SkyCoord(ra=ra * auni.deg, dec=dec * auni.deg))
ima = np.zeros((npix, npix), dtype=np.int32)
ima = np.require(ima, dtype=np.int32, requirements='AOC')
args = {'xs': xs, 'ys': ys, 'rads': rads}
Cargs = {}
Xargs = {}
for k in args.keys():
Xargs[k] = np.require(args[k],
dtype=np.float32,
requirements=["A", "O", "C"])
Cargs[k] = ffi.cast('float*', _filler.ffi.from_buffer(Xargs[k]))
Cargs['ima'] = ffi.cast('int*', _filler.ffi.from_buffer(ima))
_filler.lib.procima(Cargs['ima'], npix, Cargs['xs'], Cargs['ys'],
Cargs['rads'], nstars, binning)
return ima, wcdict