disperse.pyx 5.52 KB
Newer Older
Fang Yuedong's avatar
Fang Yuedong committed
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

from __future__ import division

import numpy as np
cimport numpy as np

DTYPE = np.double
ITYPE = np.int64

ctypedef np.double_t DTYPE_t

ctypedef np.uint_t UINT_t
ctypedef np.int_t INT_t
ctypedef np.int64_t LINT_t
ctypedef np.int32_t FINT_t
ctypedef np.float32_t FTYPE_t

import cython

cdef extern from "math.h":
    double sqrt(double x)
    double exp(double x)
    
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.embedsignature(True)
def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam, 
                          np.ndarray[LINT_t, ndim=1] idxl, 
                          np.ndarray[DTYPE_t, ndim=1] yfrac, 
                          np.ndarray[DTYPE_t, ndim=1] ysens, 
                          np.ndarray[DTYPE_t, ndim=1] full, 
                          np.ndarray[LINT_t, ndim=1] x0, 
                          np.ndarray[LINT_t, ndim=1] shd, 
34
35
36
                          np.ndarray[LINT_t, ndim=1] shg,
                          np.ndarray[DTYPE_t, ndim=2] flat,
                          np.ndarray[DTYPE_t, ndim=1] wlambda):
Fang Yuedong's avatar
Fang Yuedong committed
37
    """Compute a dispersed 2D spectrum
38

Fang Yuedong's avatar
Fang Yuedong committed
39
40
41
42
    Parameters
    ----------
    flam : direct image matrix, 2 dim [y,x]
    idxl: grating disperse light to pixel, pixel index, 1 dim, length = ysens, yfrac
43
    yfrac:
44
    ysens: sensitivity  use pixel describe
Fang Yuedong's avatar
Fang Yuedong committed
45
    full: output result ,1 dim, y_beam * x_beam
46
    x0: the center of gal in image thumbnail (y, x)
Fang Yuedong's avatar
Fang Yuedong committed
47
48
49
50
51
52
    shd:  shape of direct image
    shg:  shape of grating image
    """
    cdef int i,j,k1,k2
    cdef unsigned int nk,nl,k,shx,shy
    cdef double fl_ij
53

Fang Yuedong's avatar
Fang Yuedong committed
54
55
    nk = len(idxl)
    nl = len(full)
56

57
    if (flat is not None):
Zhang Xin's avatar
Zhang Xin committed
58
        nlamb = len(wlambda)
xin's avatar
xin committed
59
        nflat = len(flat)
Zhang Xin's avatar
Zhang Xin committed
60
61
62
        flat_lamb_min = 2500
        flat_lamb_max = 10000
        lambda_co = np.zeros([len(flat[0]),nlamb])
xin's avatar
xin committed
63
64
65
66
67
68
69
70
71
        lambda_co[0] = lambda_co[0] + ysens
        lambda_co[1] = (wlambda-flat_lamb_min)/(flat_lamb_max-flat_lamb_min)*ysens
        lambda_co[2] = lambda_co[1]*lambda_co[1]*ysens
        lambda_co[3] = lambda_co[1]*lambda_co[2]*ysens

        flat_eff_all = np.zeros(nflat*nlamb)

        for i in range(0, nflat):
            flat_eff_all[i*nlamb:(i+1)*nlamb]=np.dot(flat[i], lambda_co)
Zhang Xin's avatar
Zhang Xin committed
72
73
74

        for i in range(0-x0[1], x0[1]):
            if (x0[1]+i < 0) | (x0[1]+i >= shd[1]):
Fang Yuedong's avatar
Fang Yuedong committed
75
76
                continue

Zhang Xin's avatar
Zhang Xin committed
77
78
79
80
81
82
83
            for j in range(0-x0[0], x0[0]):
                if (x0[0]+j < 0) | (x0[0]+j >= shd[0]):
                    continue

                fl_ij = flam[x0[0]+j, x0[1]+i] #/1.e-17
                if (fl_ij == 0):
                    continue
84

Zhang Xin's avatar
Zhang Xin committed
85
86
87
                for k in range(nk):
                    k1 = idxl[k]+j*shg[1]+i
                    if (k1 >= 0) & (k1 < nl):
xin's avatar
xin committed
88
                        flat_ids = k1*nlamb+k
89
                        full[k1] += fl_ij*(1-yfrac[k])*flat_eff_all[flat_ids]
90

91
                    k2 = idxl[k]+(j+1)*shg[1]+i
Zhang Xin's avatar
Zhang Xin committed
92
                    if (k2 >= 0) & (k2 < nl):
xin's avatar
xin committed
93
                        flat_ids = k2*nlamb+k
94
                        full[k2] += fl_ij*yfrac[k]*flat_eff_all[flat_ids]
xin's avatar
xin committed
95

Zhang Xin's avatar
Zhang Xin committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    else:
        for i in range(0-x0[1], x0[1]):
            if (x0[1]+i < 0) | (x0[1]+i >= shd[1]):
                continue

            for j in range(0-x0[0], x0[0]):
                if (x0[0]+j < 0) | (x0[0]+j >= shd[0]):
                    continue

                fl_ij = flam[x0[0]+j, x0[1]+i] #/1.e-17
                if (fl_ij == 0):
                    continue

                for k in range(nk):
                    k1 = idxl[k]+j*shg[1]+i
                    if (k1 >= 0) & (k1 < nl):
112
                        full[k1] += ysens[k]*fl_ij*(1-yfrac[k])
Zhang Xin's avatar
Zhang Xin committed
113

114
                    k2 = idxl[k]+(j+1)*shg[1]+i
Zhang Xin's avatar
Zhang Xin committed
115
                    if (k2 >= 0) & (k2 < nl):
116
                        full[k2] += ysens[k]*fl_ij*yfrac[k]
Zhang Xin's avatar
Zhang Xin committed
117

Fang Yuedong's avatar
Fang Yuedong committed
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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
    return True

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.embedsignature(True)
def compute_segmentation_limits(np.ndarray[FTYPE_t, ndim=2] segm, int seg_id, np.ndarray[FTYPE_t, ndim=2] flam, np.ndarray[LINT_t, ndim=1] shd):
    """Find pixel limits of a segmentation region
    
    Parameters
    ----------
    segm: ndarray (np.float32)
        segmentation array
    
    seg_id: int
        ID to test
    
    flam: ndarray (float)
        Flux array to compute weighted centroid within segmentation region
        
    shd: [int, int]
        Shape of segm
    """
    cdef int i, j, imin, imax, jmin, jmax, area
    cdef double inumer, jnumer, denom, wht_ij
    
    area = 0
    
    imin = shd[0]
    imax = 0
    jmin = shd[1]
    jmax = 0
    
    inumer = 0.
    jnumer = 0.
    denom = 0.
    
    for i in range(shd[0]):
        for j in range(shd[1]):
            if segm[i,j] != seg_id:
                continue
            
            area += 1
            wht_ij = flam[i,j]
            inumer += i*wht_ij
            jnumer += j*wht_ij
            denom += wht_ij
            
            if i < imin:
                imin = i
            if i > imax:
                imax = i
            
            if j < jmin: 
                jmin = j
            if j > jmax:
                jmax = j
    
    ### No matched pixels
    if denom == 0:
        denom = -99
        
    return imin, imax, inumer/denom, jmin, jmax, jnumer/denom, area, denom
            
            
            
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.embedsignature(True)
def seg_flux(np.ndarray[FTYPE_t, ndim=2] flam, np.ndarray[LINT_t, ndim=1] idxl, np.ndarray[DTYPE_t, ndim=1] yfrac, np.ndarray[DTYPE_t, ndim=1] ysens, np.ndarray[DTYPE_t, ndim=1] full, np.ndarray[LINT_t, ndim=1] x0, np.ndarray[LINT_t, ndim=1] shd, np.ndarray[LINT_t, ndim=1] shg):
    pass