pep8

tools/get_PSF_SLS.py

@@ -7,13 +7,15 @@ import astropy.io.fits as fitsio
 from observation_sim.instruments import Chip, FilterParam, Filter
 import galsim
 
-## 计算 0级或1级光谱在某个波长位置的PSF, 返回值是过采样的PSF，像元大小有CSST图像像元大小的1/2
-## chipId 必须为[1,2,3,4,5,10,21,26,27,28,29,30]
-## order 只有 0 或 1
-## pos_img 为直接成像在图像上的位置,[x, y]
-## wave: 波长，单位A
-def get_SLS_PSF(chipID = 1,pos_img = [6000,4000], order = 1, wave = 8000, sls_psf_dir = '/nfsdata/share/CSSOSDataProductsSims/data/SLS_PSF_PCA_fp_cd/'):
-    orders = {0:'B',1:'A'}
+# 计算 0级或1级光谱在某个波长位置的PSF, 返回值是过采样的PSF，像元大小有CSST图像像元大小的1/2
+# chipId 必须为[1,2,3,4,5,10,21,26,27,28,29,30]
+# order 只有 0 或 1
+# pos_img 为直接成像在图像上的位置,[x, y]
+# wave: 波长，单位A
+
+
+def get_SLS_PSF(chipID=1, pos_img=[6000, 4000], order=1, wave=8000, sls_psf_dir='/nfsdata/share/CSSOSDataProductsSims/data/SLS_PSF_PCA_fp_cd/'):
+    orders = {0: 'B', 1: 'A'}
     chip = Chip(chipID)
     filter_id, filter_type = chip.getChipFilter()
     filt = Filter(
@@ -23,30 +25,31 @@ def get_SLS_PSF(chipID = 1,pos_img = [6000,4000], order = 1, wave = 8000, sls_ps
 
     psf_model = PSFInterpSLS(chip, filt, PSF_data_prefix=sls_psf_dir)
     bandNo = 1
-    for i,brange in enumerate(psf_model.bandranges):
-        if wave>=brange[0] and wave<brange[1]:
+    for i, brange in enumerate(psf_model.bandranges):
+        if wave >= brange[0] and wave < brange[1]:
             bandNo = i + 1
             break
 
-
     psf1, _ = psf_model.get_PSF(
-                            chip, pos_img_local=pos_img, bandNo=bandNo, galsimGSObject=False, g_order=orders[order])
+        chip, pos_img_local=pos_img, bandNo=bandNo, galsimGSObject=False, g_order=orders[order])
     fn = "psf.chip{:}.order{:}.wave{:}.fits".format(chipID, order, wave)
     hdu = fitsio.PrimaryHDU()
     hdu.data = psf1
     hdu.header.set('pixScale', 5)
-    hdu.writeto(fn, overwrite = True)
+    hdu.writeto(fn, overwrite=True)
+
+# 获得成像的PSF,PSF在波长区间分四段采样，返回4段的PSF，每个PSF代表不同band波长范围不同的四段psf，波长从小到大
+# chipId 必须为[6,7,8,9,11,12,13,14,15,16,17,18,19,20,22,23,24,25]
+# pos_img 为直接成像在图像上的位置,[x, y]
+
 
-## 获得成像的PSF,PSF在波长区间分四段采样，返回4段的PSF，每个PSF代表不同band波长范围不同的四段psf，波长从小到大
-## chipId 必须为[6,7,8,9,11,12,13,14,15,16,17,18,19,20,22,23,24,25]
-## pos_img 为直接成像在图像上的位置,[x, y]
-def getPSFImage(chipID = 6, pos_img = [6000, 4500], psf_pho_dir = "/nfsdata/share/CSSOSDataProductsSims/data/psfcube/set1_dynamic/"):
+def getPSFImage(chipID=6, pos_img=[6000, 4500], psf_pho_dir="/nfsdata/share/CSSOSDataProductsSims/data/psfcube/set1_dynamic/"):
     chip = Chip(chipID=chipID)
     print('chip.bound::', chip.bound.xmin, chip.bound.xmax,
-        chip.bound.ymin, chip.bound.ymax)
+          chip.bound.ymin, chip.bound.ymax)
 
     # Setup Position on focalplane
-        # try get the PSF at some location (1234, 1234) on the chip
+    # try get the PSF at some location (1234, 1234) on the chip
     pos_img = pos_img
     x, y = pos_img
     x = x+chip.bound.xmin
@@ -71,23 +74,23 @@ def getPSFImage(chipID = 6, pos_img = [6000, 4500], psf_pho_dir = "/nfsdata/shar
 
         # Get corresponding PSF model
         psf_model = PSFInterp(chip=chip, npsf=100,
-                                PSF_data_file=psf_pho_dir)
+                              PSF_data_file=psf_pho_dir)
         psf = psf_model.get_PSF(
             chip=chip, pos_img=pos_img, bandpass=bandpass, galsimGSObject=False)
         fn = "psf_chip{:}.wave{:}.fits".format(chipID, i)
         hdu = fitsio.PrimaryHDU()
         hdu.data = psf
         hdu.header.set('pixScale', 5)
-        hdu.writeto(fn, overwrite = True)
+        hdu.writeto(fn, overwrite=True)
+
+
 if __name__ == "__main__":
     chipid = 2
     pos_img = [6000, 4500]
     order = 1
     w = 5000
-    get_SLS_PSF(chipID = chipid,pos_img = pos_img, order = order, wave = w)
-
+    get_SLS_PSF(chipID=chipid, pos_img=pos_img, order=order, wave=w)
 
     chipid = 7
     pos_img = [6000, 4500]
-    getPSFImage(chipID = chipid, pos_img = pos_img)
-
+    getPSFImage(chipID=chipid, pos_img=pos_img)