pycode.py

'''
Author: Zhang Xin zhangx@bao.ac.cn
Date: 2024-01-02 13:34:39
LastEditors: Zhang Xin zhangx@bao.ac.cn
LastEditTime: 2024-03-25 13:51:33
FilePath: /csst-simulation/Users/zhangxin/Work/SlitlessSim/CSST_SIM/Star_spec/csst_spec_interp_clean/code/pycode.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''
##运行下面在mac/linux下执行 
# cc -fPIC -shared main_singlestar.c -I/usr/include -I/usr/include/cfitsio -lcfitsio -lm -o test.dylib

import os
import numpy as np
from astropy.io import fits
import ctypes
from astropy.table import Table
import produceSED
import astropy.constants as atcons
import sys
# from ctypes import *

# struct STAR
# {
#   float logte;
#   float logg;
#   //float logL;
#   float Mass;
#   float Av;
#   float mu0;
#   float Z;
#   //float mbolmag;
# };

class Star(ctypes.Structure):
    _fields_ = [
               ('logte',ctypes.c_float),
               ('logg',ctypes.c_float),
               ('Mass',ctypes.c_float),
               ('Av', ctypes.c_float),
               ('mu0', ctypes.c_float),
               ('Z', ctypes.c_float)]

#CHANGE
#topdir='/run/media/chen/1TS/dupe/gitlab/csst_spec_interp/'
topdir='/home/zhangxin/CSST_SIM/star_spec/csst_spec_interp_clean/code/'
code_name='test.so'

d = ctypes.CDLL(os.path.join(topdir,code_name))
d.loadSpecLibs.argtypes=[ctypes.c_char_p, ctypes.c_char_p]
d.loadExts.argtypes=[ctypes.c_char_p]


#######################################################################################
#CHANGE
#nwv = d.loadSpecLibs(str.encode(os.path.join(topdir,'file_CK04.par')),str.encode(topdir))
#d.loadExts(str.encode(os.path.join(topdir,"spec/Ext_odonnell94_R3.1_CK04W.fits")))
#TO
nwv = d.loadSpecLibs(str.encode(os.path.join(topdir,'file_BT-Settl_CSST_wl1000-24000_R1000.par')),str.encode(topdir))
d.loadExts(str.encode(os.path.join(topdir,"spec/Ext_odonnell94_R3.1_CSST_wl1000-24000_R1000.fits")))
########################################################################################

print("Done loading Exts")
spec = (ctypes.c_float*nwv)()
wave = (ctypes.c_float*nwv)()
d.interpSingleStar.argtypes=[ctypes.Structure, ctypes.POINTER(ctypes.c_float)]


'''
#example for a single star
#s=Star(3.620752, 4.701155, 0.599979, 0.067540, 11.200001, 0.008501)
s=Star(3.9739766, 8.108173, 0.6525841, 0.077022046, 9.05, 0.024376422)
#s=Star(3.9739766, 4.99, 0.6525841, 0.077022046, 9.05, 0.024376422)
d.interpSingleStar(s, spec)
spec_ = spec[:]
print(spec[500:509])
'''

from astropy.table import Table

if len(sys.argv) < 2:
	print('usage:\n'+sys.argv[0]+' survey_result.txt')
	sys.exit(0)

catalogFn = str(sys.argv[1])

# catalogFn = "/nfsdata/share/CSSTsimInputCat_TH/C9_RA240_DECp30.fits"
# catalogFn = "/home/zhangxin/CSST_SIM/star_spec/csst_spec_interp_clean/code/data/catalog/trilegal.fits"
outdir = '/nfsdata/share/CSSOSDataProductsSims/trilegalCat/'
# outFn = "C9_RA240_DECp30_refCat_1.fits"

outFn = catalogFn.strip().split('/')[-1][0:-5] + "_refCat_1.fits"
cat = Table.read(catalogFn)

filters = ['nuv','u','g','r','i','z','y']
filters_other = ['2MASS_H','2MASS_J','2MASS_Ks','GAIA_GAIA3.G','GAIA_GAIA3.Gbp','GAIA_GAIA3.Grp','GALEX_GALEX.FUV','GALEX_GALEX.NUV','LSST_u','LSST_g','LSST_r','LSST_i','LSST_z','LSST_y','PAN-STARRS_PS1.g','PAN-STARRS_PS1.r','PAN-STARRS_PS1.i','PAN-STARRS_PS1.z','PAN-STARRS_PS1.y','SLOAN_SDSS.u','SLOAN_SDSS.g','SLOAN_SDSS.r','SLOAN_SDSS.i','SLOAN_SDSS.z']
res = {}

nrows = len(cat)
for fi in filters:
    res[fi] = np.zeros(nrows)
for fi in filters_other:
    res[fi] = np.zeros(nrows)

parallaxs = np.zeros(nrows)

from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
rank_size = comm.Get_size()
print("------------------",rank_size, rank, nrows)
iterNum = 0    
for star in cat:
    if iterNum%10000==0:
        print(iterNum)
    # if iterNum > 100:
    #     iterNum = iterNum + 1
    #     continue
    if iterNum % rank_size != rank:
        iterNum = iterNum + 1
        continue
    # specTable = np.zeros([nwv,2])
    s=Star(star['mwmsc_logte'], star['mwmsc_logg'], star['mwmsc_mass'], star['mwmsc_av'], star['mwmsc_mu0'], star['mwmsc_z'])
    # av stellarmass dm teff logg feh
    # 0.0464   0.7512  10.8000   3.6914   4.5952   0.0122
    #  s=StarParm(obj.param['teff'], obj.param['logg'], obj.param['stellarMass'], obj.param['av'], obj.param['DM'], obj.param['feh'])
    # s=Star(3.6914, 4.5952, 0.7512, 0.0464, 10.8000, 0.0122) 
    #print("star[logTe], star[logg], star[Mass], star[Av], star[mu0], star[Z]: ", star['logTe'], star['logg'], star['Mass'], star['Av'], star['mu0'], star['Z'])
    d.interpSingleStar(s, spec, wave)

    rv_c = star['mwmsc_vrad']/(atcons.c.value/1000.)
    Doppler_factor = np.sqrt((1+rv_c)/(1-rv_c))
    wave_RV = wave*Doppler_factor
    # specTable[:,0] = wave[:]
    # specTable[:,1] = spec[:]
    # print(spec[500:509])
    spec_out = Table(np.array([wave_RV, np.power(10,spec[:])]).T, names=('WAVELENGTH', 'FLUX'))
    spec_norm_phot = produceSED.produceNormSED_photon(inputSED = spec_out,mag_norm = star['mwmsc_gsmag'], norm_filter_thr_fn= 'data/throughputs/SDSS/SLOAN_SDSS.g.fits',ws = 1050, we = 23950)
    # spec_norm_phot = produceSED.produceNormSED_photon(inputSED = spec_out,mag_norm = 17.8360, norm_filter_thr_fn= 'data/throughputs/SDSS/SLOAN_SDSS.g.fits',ws = 1000, we = 24000)
    mags = produceSED.calculatCSSTMAG(spec = spec_norm_phot, throughput_dir = 'data/throughputs/CSST_n/',ws= 2000, we = 11000)

    mags_others = produceSED.calculatCSSTMAG(spec = spec_norm_phot, throughput_dir = 'data/throughputs/filter_transp/',ws= 1050, we = 23950, filelist=filters_other, band_instr='other')

    for fi in filters:
        res[fi][iterNum] = mags[fi]
    for fi in filters_other:
        res[fi][iterNum] = mags_others[fi]
    parallaxs[iterNum] = 1/(10**(star['mwmsc_mu0']*0.2)/100)
    iterNum = iterNum + 1
    # print(mags, mags_others)

# send_data1  = res
# send_data2 = parallaxs
# if rank == 0:
#     total_res = comm.gather(send_data1, root=0)
#     total_parall = comm.gather(send_data2, root=0)

#     for i in np.arange(1, rank_size, 1):
#         for fi in filters:
#             res[fi] = res[fi] + total_res[i][fi]
#         for fi in filters_other:
#             res[fi] = res[fi] + total_res[i][fi]
#         parallaxs = parallaxs + total_parall[i]
#     cat.add_column(np.round(parallaxs,5),name='parallax')
#     for fi in filters:
#         cat.add_column(np.round(res[fi],5),name='interSpec_'+fi)
#     for fi in filters_other:
#         cat.add_column(np.round(res[fi],5),name='interSpec_'+fi)
#     outdir = '/nfsdata/share/CSSOSDataProductsSims/trilegalCat/'
#     outFn = "C9_RA300_DECm60_refCat.fits"
#     cat.write(os.path.join(outdir,outFn),overwrite=True)

# else:
#     comm.gather(send_data1, root=0)
#     comm.gather(send_data2, root=0)


total_res=comm.gather(res, root=0)
total_parall = comm.gather(parallaxs, root=0)

if rank == 0:
    # new_res = res
    # total_res=comm.gather(res, root=0)
    # print("gather data size is ", size(total_res))
    for i in np.arange(1, rank_size, 1):
        for fi in filters:
            res[fi] = res[fi] + total_res[i][fi]
        for fi in filters_other:
            res[fi] = res[fi] + total_res[i][fi]
        parallaxs = parallaxs + total_parall[i]
    cat.add_column(np.round(parallaxs,5),name='parallax')
    for fi in filters:
        cat.add_column(np.round(res[fi],5),name='interSpec_'+fi)
    for fi in filters_other:
        cat.add_column(np.round(res[fi],5),name='interSpec_'+fi)
    # outdir = '/nfsdata/share/CSSOSDataProductsSims/trilegalCat/'
    # outFn = "C9_RA240_DECp30_refCat_1.fits"
    cat.write(os.path.join(outdir,outFn),overwrite=True)
    # print('--------------------------')
    # print(mags['nuv']-star['mwmsc_nuvmag'], mags['u']-star['mwmsc_umag'], mags['g']-star['mwmsc_gmag'], mags['r']-star['mwmsc_rmag'], mags['i']-star['mwmsc_imag'],mags['z']-star['mwmsc_zmag'],mags['y']-star['mwmsc_ymag'])
#exmple of handling a fits catalouge with many stars
#wvs=fits.open('spec/fp00k2odfnew.fits')[2].data['wave'] #in Angstrom
# cat_list=np.loadtxt(topdir+'cat.list', dtype='U')
# for listfile in cat_list:
#     hdu = fits.open(topdir+listfile)
#     cat=hdu[1].data
#     for star in cat:
#         specTable = np.zeros([nwv,2])
#         s=Star(star['logTe'], star['logg'], star['Mass'], star['Av'], star['mu0'], star['Z'])
#         #print("star[logTe], star[logg], star[Mass], star[Av], star[mu0], star[Z]: ", star['logTe'], star['logg'], star['Mass'], star['Av'], star['mu0'], star['Z'])
#         d.interpSingleStar(s, spec, wave)
#         specTable[:,0] = wave[:]
#         specTable[:,1] = spec[:]
#         # print(spec[500:509])
#         spec_out = Table(np.array([wave[:], np.power(10,spec[:])]).T, names=('WAVELENGTH', 'FLUX'))
#         spec_norm_phot = produceSED.produceNormSED_photon(inputSED = spec_out,mag_norm = 24.0, norm_filter_thr_fn= 'data/throughputs/SDSS/SLOAN_SDSS.g.fits',ws = 2500, we = 10000)
#         mags = produceSED.calculatCSSTMAG(spec = spec_norm_phot, throughput_dir = 'data/throughputs/CSST/',ws= 2500, we = 10000)

#         # print("MAG is ",mags)
#         print('--------------------------')
#         print(mags['nuv']-star['mwmsc_nuvmag'], mag['u']-star['mwmsc_umag'], mag['g']-star['mwmsc_gmag'], mag['r']-star['mwmsc_rmag'], mag['i']-star['mwmsc_imag'],mag['z']-star['mwmsc_zmag'],mag['y']-star['mwmsc_ymag'])