-
Notifications
You must be signed in to change notification settings - Fork 6
/
adaptive_ggm.py
executable file
·116 lines (94 loc) · 3.39 KB
/
adaptive_ggm.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
#!/usr/bin/env python3
import os.path
import argparse
import subprocess
from astropy.io import fits
import numpy as np
def _escapeArgs(args):
"""Escape arguments for printing."""
out = []
for a in args:
a = a.replace("'", r"\'")
wrap = False
for c in ' \t[]<>|()$"#!*?~&;':
if c in a:
wrap = True
break
if wrap:
a = "'%s'" % a
out.append(a)
return ' '.join(out)
def call(*args, **argsv):
"""Shortcut for running external programs."""
if not argsv.get('silent', False):
print("+", _escapeArgs(args))
subprocess.check_call(args)
def calcGradient(inimg, log):
"""Calculate gradient magnitude of input image."""
if log:
inimg = np.log10(inimg)
gy = inimg[1:,:]-inimg[:-1,:]
gx = inimg[:,1:]-inimg[:,:-1]
gyc = 0.5*(gy[1:,:]+gy[:-1,:])
gxc = 0.5*(gx[:,1:]+gx[:,:-1])
gyc = np.pad(gyc, ((1,1),(0,0)))
gxc = np.pad(gxc, ((0,0),(1,1)))
grad = np.sqrt(gxc**2 + gyc**2)
return grad
def main():
parser = argparse.ArgumentParser(
description="Adapive Gaussian gradient magnitude",
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('counts', help='counts image filename')
parser.add_argument('--image', help='optional image filename to apply smoothing to (uses counts if not set)')
parser.add_argument('--sn', type=float, default=32, help='Smoothing S/N ratio')
parser.add_argument('--log', type=bool, default=True, help='Calculate log after smoothing')
parser.add_argument('--mask', help='Mask image (optional)')
parser.add_argument('--scale', help='Intermediate scale map filename', default='scale.fits')
parser.add_argument('--smoothed', help='Intermediate smoothed image filename', default='smoothed.fits')
parser.add_argument('--threads', type=int, default=4, help='Number of threads to use when smoothing')
parser.add_argument('--contbin-dir', help='Override location of contour binning code')
parser.add_argument('output', help='output image filename')
args = parser.parse_args()
# where to find program
program = 'accumulate_counts'
if args.contbin_dir:
program = os.path.join(args.contbin_dir, program)
# first calculate scale map
print('* Calculating scale map')
cargs = [
args.counts,
'--scale=%s' % args.scale,
'--sn=%g' % args.sn,
'--threads=%i' % args.threads,
]
if args.mask:
cargs.append('--mask=%s' % args.mask)
call(program, *cargs)
# now smooth input image (using input image of counts image)
print('* Smoothing input image')
inimage = args.image or args.counts
cargs = [
inimage,
'--apply',
'--scale=%s' % args.scale,
'--applied=%s' % args.smoothed,
'--threads=%i' % args.threads,
'--gaussian',
]
if args.mask:
cargs.append('--mask=%s' % args.mask)
call(program, *cargs)
# get gradient
print('* Calculating gradient')
smf = fits.open(args.smoothed, 'readonly')
smimg = smf[0].data
grad = calcGradient(smimg, args.log)
smf.close()
# write to output filename (based on input image)
inf = fits.open(inimage, 'readonly')
inf[0].data = grad
print('* Writing gradient image', args.output)
inf.writeto(args.output, overwrite=True)
if __name__ == '__main__':
main()