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, np.ndarray[LINT_t, ndim=1] shg, np.ndarray[DTYPE_t, ndim=2] flat, np.ndarray[DTYPE_t, ndim=1] wlambda): """Compute a dispersed 2D spectrum Parameters ---------- flam : direct image matrix, 2 dim [y,x] idxl: grating disperse light to pixel, pixel index, 1 dim, length = ysens, yfrac yfrac: ysens: sensitivity use pixel describe full: output result ,1 dim, y_beam * x_beam x0: the center of gal in image thumbnail (y, x) 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 nk = len(idxl) nl = len(full) if (flat is not None): nlamb = len(wlambda) nflat = len(flat) flat_lamb_min = 2500 flat_lamb_max = 10000 lambda_co = np.zeros([len(flat[0]),nlamb]) 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) 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): flat_ids = k1*nlamb+k full[k1] += fl_ij*(1-yfrac[k])*flat_eff_all[flat_ids] k2 = idxl[k]+(j+1)*shg[1]+i if (k2 >= 0) & (k2 < nl): flat_ids = k2*nlamb+k full[k2] += fl_ij*yfrac[k]*flat_eff_all[flat_ids] 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): full[k1] += ysens[k]*fl_ij*(1-yfrac[k]) k2 = idxl[k]+(j+1)*shg[1]+i if (k2 >= 0) & (k2 < nl): full[k2] += ysens[k]*fl_ij*yfrac[k] 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