Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import unittest 

2 

3import os 

4import numpy as np 

5import json 

6import tempfile 

7import shutil 

8 

9import lsst.utils.tests 

10from lsst.utils import getPackageDir 

11 

12from lsst.sims.catalogs.db import fileDBObject 

13from lsst.sims.catUtils.utils import FastStellarLightCurveGenerator 

14from lsst.sims.catUtils.utils import StellarLightCurveGenerator 

15from lsst.sims.catUtils.utils import AgnLightCurveGenerator 

16from lsst.sims.catUtils.utils import FastAgnLightCurveGenerator 

17 

18from lsst.sims.utils.CodeUtilities import sims_clean_up 

19from lsst.sims.utils import ModifiedJulianDate 

20 

21ROOT = os.path.abspath(os.path.dirname(__file__)) 

22 

23 

24class FastStellar_stellar_lc_gen_case(unittest.TestCase): 

25 

26 longMessage = True 

27 

28 @classmethod 

29 def setUpClass(cls): 

30 """ 

31 Create a fake catalog of RR Lyrae stars and MLT dwarves with flaring 

32 light curves. Store it in cls.stellar_db 

33 """ 

34 cls.scratchDir = tempfile.mkdtemp(dir=ROOT, prefix='FastStellar_stellar_lc_gen_case-') 

35 

36 cls.raRange = (78.0, 85.0) 

37 cls.decRange = (-69.0, -65.0) 

38 

39 rng = np.random.RandomState(88) 

40 cls.n_stars = 20 

41 sed_dir = os.path.join(getPackageDir("sims_sed_library")) 

42 sed_dir = os.path.join(sed_dir, "starSED", "kurucz") 

43 list_of_seds = os.listdir(sed_dir) 

44 

45 lc_dir = os.path.join(getPackageDir("sims_sed_library"), "rrly_lc") 

46 lc_dir = os.path.join(lc_dir, "RRab") 

47 list_of_rrly_lc = ['rrly_lc/RRab/%s' % ww for ww in os.listdir(lc_dir) if "per.txt" in ww] 

48 

49 cls.mlt_lc_file_name = os.path.join(cls.scratchDir, "fast_lc_mlt_file.npz") 

50 if os.path.exists(cls.mlt_lc_file_name): 

51 os.unlink(cls.mlt_lc_file_name) 

52 

53 mlt_lc_files = {} 

54 mlt_lc_files['lc_1_time'] = np.arange(0.0, 3652.51, 0.1) 

55 mlt_lc_files['lc_1_g'] = 2.2e32*np.power(np.cos(mlt_lc_files['lc_1_time']/100.0-5.0),2) 

56 mlt_lc_files['lc_1_r'] = 1.3e32*(1.0+np.sin(mlt_lc_files['lc_1_time']/100.0-3.0)) 

57 

58 mlt_lc_files['lc_2_time'] = np.arange(0.0, 3652.51, 0.1) 

59 mlt_lc_files['lc_2_g'] = 5.1e33*(1.0+np.cos(mlt_lc_files['lc_2_time']/300.0-10.0)) 

60 mlt_lc_files['lc_2_r'] = 4.3e32*(1.0+np.sin(mlt_lc_files['lc_2_time']/50.0-71.0)) 

61 

62 with open(cls.mlt_lc_file_name, 'wb') as file_handle: 

63 np.savez(file_handle, **mlt_lc_files) 

64 

65 cls.dtype = np.dtype([('id', np.int), 

66 ('raDeg', np.float), 

67 ('decDeg', np.float), 

68 ('raJ2000', np.float), 

69 ('decJ2000', np.float), 

70 ('magNorm', np.float), 

71 ('galacticAv', np.float), 

72 ('sedFilename', str, 300), 

73 ('varParamStr', str, 300), 

74 ('parallax', np.float), 

75 ('ebv', np.float)]) 

76 

77 # write the catalog as a text file to be ingested with fileDBObject 

78 cls.txt_name = os.path.join(cls.scratchDir, "fast_stellar_lc_catalog.txt") 

79 with open(cls.txt_name, "w") as output_file: 

80 output_file.write('# a silly header\n') 

81 sed_dex = rng.randint(0, len(list_of_seds), size=cls.n_stars//2) 

82 lc_dex = rng.randint(0, len(list_of_rrly_lc), size=cls.n_stars//2) 

83 mjd0 = rng.random_sample(cls.n_stars//2)*10000.0+40000.0 

84 raList = rng.random_sample(cls.n_stars//2)*(cls.raRange[1]-cls.raRange[0])+cls.raRange[0] 

85 decList = cls.decRange[0] + rng.random_sample(cls.n_stars//2)*(cls.decRange[1]-cls.decRange[1]) 

86 magNormList = rng.random_sample(cls.n_stars//2)*3.0+14.0 

87 AvList = rng.random_sample(cls.n_stars//2)*0.2+0.1 

88 pxList = rng.random_sample(cls.n_stars//2)*0.1 

89 for ix in range(cls.n_stars//2): 

90 varparams = {'varMethodName': 'applyRRly', 

91 'pars': {'tStartMjd': mjd0[ix], 

92 'filename': list_of_rrly_lc[lc_dex[ix]]}} 

93 varparamstr = json.dumps(varparams) 

94 output_file.write("%d;%lf;%lf;%lf;%lf;%lf;%lf;%s;%s;%lf;%lf\n" 

95 % (ix, raList[ix], decList[ix], 

96 np.radians(raList[ix]), 

97 np.radians(decList[ix]), 

98 magNormList[ix], AvList[ix], 

99 list_of_seds[sed_dex[ix]], 

100 varparamstr,pxList[ix], 

101 AvList[ix]/3.1)) 

102 

103 sed_dex = rng.randint(0, len(list_of_seds), size=cls.n_stars//2) 

104 lc_dex = rng.randint(1, 3, size=cls.n_stars//2) 

105 mjd0 = rng.random_sample(cls.n_stars//2)*10000.0+40000.0 

106 raList = rng.random_sample(cls.n_stars//2)*(cls.raRange[1]-cls.raRange[0])+cls.raRange[0] 

107 decList = cls.decRange[0] + rng.random_sample(cls.n_stars//2)*(cls.decRange[1]-cls.decRange[1]) 

108 magNormList = rng.random_sample(cls.n_stars//2)*3.0+14.0 

109 AvList = rng.random_sample(cls.n_stars//2)*0.2+0.1 

110 pxList = rng.random_sample(cls.n_stars//2)*0.1 

111 for ix in range(cls.n_stars//2): 

112 varparams = {'m':'MLT', 'p':{'lc':'lc_%d' % lc_dex[ix], 't0':rng.random_sample()*1000.0}} 

113 varparamstr = json.dumps(varparams) 

114 output_file.write("%d;%lf;%lf;%lf;%lf;%lf;%lf;%s;%s;%lf;%lf\n" 

115 % (ix+cls.n_stars/2, raList[ix], decList[ix], 

116 np.radians(raList[ix]), 

117 np.radians(decList[ix]), 

118 magNormList[ix], AvList[ix], 

119 list_of_seds[sed_dex[ix]], 

120 varparamstr,pxList[ix], 

121 AvList[ix]/3.1)) 

122 

123 cls.stellar_db = fileDBObject(cls.txt_name, delimiter=';', 

124 runtable='test', dtype=cls.dtype, 

125 idColKey='id') 

126 

127 cls.stellar_db.raColName = 'raDeg' 

128 cls.stellar_db.decColName = 'decDeg' 

129 cls.stellar_db.objectTypeId = 32 

130 

131 cls.opsimDb = os.path.join(getPackageDir("sims_data"), "OpSimData") 

132 cls.opsimDb = os.path.join(cls.opsimDb, "opsimblitz1_1133_sqlite.db") 

133 

134 @classmethod 

135 def tearDownClass(cls): 

136 sims_clean_up() 

137 if os.path.exists(cls.txt_name): 

138 os.unlink(cls.txt_name) 

139 if os.path.exists(cls.mlt_lc_file_name): 

140 os.unlink(cls.mlt_lc_file_name) 

141 if os.path.exists(cls.scratchDir): 

142 shutil.rmtree(cls.scratchDir) 

143 

144 def test_fast_stellar_lc_gen(self): 

145 

146 bandpass = ('r', 'g') 

147 

148 lc_slow = StellarLightCurveGenerator(self.stellar_db, self.opsimDb) 

149 lc_slow._lightCurveCatalogClass._mlt_lc_file = self.mlt_lc_file_name 

150 ptngs = lc_slow.get_pointings((68.0, 95.0), (-69.0, -55.0), bandpass=bandpass) 

151 slow_lc, slow_truth = lc_slow.light_curves_from_pointings(ptngs, chunk_size=10) 

152 self.assertEqual(len(slow_truth), self.n_stars) 

153 self.assertEqual(len(slow_lc), self.n_stars) 

154 

155 lc_fast = FastStellarLightCurveGenerator(self.stellar_db, self.opsimDb) 

156 lc_fast._lightCurveCatalogClass._mlt_lc_file = self.mlt_lc_file_name 

157 ptngs = lc_fast.get_pointings((68.0, 95.0), (-69.0, -55.0), bandpass=bandpass) 

158 fast_lc, fast_truth = lc_fast.light_curves_from_pointings(ptngs, chunk_size=10) 

159 

160 self.assertEqual(len(fast_lc), len(slow_lc)) 

161 self.assertEqual(len(fast_truth), len(slow_truth)) 

162 

163 for obj_id in fast_lc: 

164 self.assertEqual(len(fast_lc[obj_id]), len(slow_lc[obj_id])) 

165 for bp in fast_lc[obj_id]: 

166 self.assertEqual(len(fast_lc[obj_id][bp]), len(slow_lc[obj_id][bp])) 

167 for data_key in fast_lc[obj_id][bp]: 

168 self.assertEqual(len(slow_lc[obj_id][bp][data_key]), len(fast_lc[obj_id][bp][data_key])) 

169 self.assertEqual(slow_lc[obj_id][bp][data_key].shape, slow_lc[obj_id][bp][data_key].shape) 

170 self.assertLess(np.abs(fast_lc[obj_id][bp][data_key]-slow_lc[obj_id][bp][data_key]).max(), 

171 1.0e-10) 

172 

173 for obj_id in fast_truth: 

174 self.assertEqual(fast_truth[obj_id], slow_truth[obj_id]) 

175 

176 

177class Fast_agn_lc_gen_test_case(unittest.TestCase): 

178 

179 longMessge = True 

180 

181 @classmethod 

182 def setUpClass(cls): 

183 rng = np.random.RandomState(119) 

184 

185 n_galaxies = 20 

186 

187 sed_dir = os.path.join(getPackageDir("sims_sed_library"), "galaxySED") 

188 list_of_seds = os.listdir(sed_dir) 

189 disk_sed_dexes = rng.randint(0, len(list_of_seds), size=n_galaxies) 

190 bulge_sed_dexes = rng.randint(0, len(list_of_seds), size=n_galaxies) 

191 

192 avBulge = rng.random_sample(n_galaxies)*0.3+0.1 

193 avDisk = rng.random_sample(n_galaxies)*0.3+0.1 

194 

195 mjdList = rng.random_sample(n_galaxies)*10.0+49330.0 

196 redshiftList = rng.random_sample(n_galaxies)*1.5+0.01 

197 

198 tauList = rng.random_sample(n_galaxies)*1.0+1.0 

199 sfuList = rng.random_sample(n_galaxies)*2.0+1.0 

200 sfgList = rng.random_sample(n_galaxies)*2.0+1.0 

201 sfrList = rng.random_sample(n_galaxies)*2.0+1.0 

202 sfiList = rng.random_sample(n_galaxies)*2.0+1.0 

203 sfzList = rng.random_sample(n_galaxies)*2.0+1.0 

204 sfyList = rng.random_sample(n_galaxies)*2.0+1.0 

205 

206 raList = rng.random_sample(n_galaxies)*7.0+78.0 

207 decList = rng.random_sample(n_galaxies)*4.0-69.0 

208 

209 normDisk = rng.random_sample(n_galaxies)*5.0+20.0 

210 normBulge = rng.random_sample(n_galaxies)*5.0+20.0 

211 normAgn = rng.random_sample(n_galaxies)*5.0+20.0 

212 

213 with lsst.utils.tests.getTempFilePath('.txt') as txt_cat_name: 

214 with open(txt_cat_name, "w") as output_file: 

215 for ix in range(n_galaxies): 

216 varParam = {'varMethodName': 'applyAgn', 

217 'pars': {'agn_tau': tauList[ix], 'agn_sfu': sfuList[ix], 

218 'agn_sfg': sfgList[ix], 'agn_sfr': sfrList[ix], 

219 'agn_sfi': sfiList[ix], 'agn_sfz': sfzList[ix], 

220 'agn_sfy': sfyList[ix], 't0_mjd': mjdList[ix], 

221 'seed': rng.randint(0, 200000)}} 

222 

223 paramStr = json.dumps(varParam) 

224 

225 output_file.write("%d;%f;%f;" % (ix, raList[ix], decList[ix]) 

226 + "%f;%f;" % (np.radians(raList[ix]), np.radians(decList[ix])) 

227 + "%f;" % (redshiftList[ix]) 

228 + "%s;%f;%f;" % (list_of_seds[disk_sed_dexes[ix]], 

229 avDisk[ix], normDisk[ix]) 

230 + "%s;%f;%f;" % (list_of_seds[bulge_sed_dexes[ix]], 

231 avBulge[ix], normBulge[ix]) 

232 + "agn.spec;%s;%f\n" % (paramStr, normAgn[ix])) 

233 

234 dtype = np.dtype([ 

235 ('galid', np.int), 

236 ('raDeg', np.float), ('decDeg', np.float), 

237 ('raJ2000', np.float), ('decJ2000', np.float), 

238 ('redshift', np.float), 

239 ('sedFilenameDisk', str, 300), ('internalAvDisk', np.float), 

240 ('magNormDisk', np.float), 

241 ('sedFilenameBulge', str, 300), ('internalAvBulge', np.float), 

242 ('magNormBulge', np.float), 

243 ('sedFilenameAgn', str, 300), ('varParamStr', str, 600), 

244 ('magNormAgn', np.float) 

245 ]) 

246 

247 cls.agn_db = fileDBObject(txt_cat_name, delimiter=';', 

248 runtable='test', dtype=dtype, 

249 idColKey='galid') 

250 

251 cls.agn_db.raColName = 'raDeg' 

252 cls.agn_db.decColName = 'decDeg' 

253 cls.agn_db.objectTypeId = 112 

254 

255 # what follows is a hack to deal with the fact thar 

256 # our varParamStr values are longer than 256 characters 

257 # which is the default maximum length that a 

258 # CatalogDBObject expects a string to be 

259 # 

260 cls.agn_db.dbTypeMap['STRING'] = (str, 600) 

261 cls.agn_db.columns = None 

262 cls.agn_db._make_default_columns() 

263 cls.agn_db._make_column_map() 

264 cls.agn_db._make_type_map() 

265 

266 cls.opsimDb = os.path.join(getPackageDir("sims_data"), "OpSimData") 

267 cls.opsimDb = os.path.join(cls.opsimDb, "opsimblitz1_1133_sqlite.db") 

268 

269 @classmethod 

270 def tearDownClass(cls): 

271 sims_clean_up() 

272 

273 def test_fast_agn_light_curves(self): 

274 raRange = (78.0, 85.0) 

275 decRange = (-69.0, -65.0) 

276 bandpass = ('g', 'r') 

277 

278 slow_lc_gen = AgnLightCurveGenerator(self.agn_db, self.opsimDb) 

279 pointings = slow_lc_gen.get_pointings(raRange, decRange, bandpass=bandpass) 

280 for row in pointings: 

281 for obs in row: 

282 mjd = ModifiedJulianDate(TAI=obs.mjd.TAI-49000.0+59580.0) 

283 obs.mjd = mjd 

284 

285 slow_lc, slow_truth = slow_lc_gen.light_curves_from_pointings(pointings) 

286 

287 self.assertGreater(len(slow_lc), 2) # make sure we got some light curves 

288 

289 fast_lc_gen = FastAgnLightCurveGenerator(self.agn_db, self.opsimDb) 

290 pointings = fast_lc_gen.get_pointings(raRange, decRange, bandpass=bandpass) 

291 for row in pointings: 

292 for obs in row: 

293 mjd = ModifiedJulianDate(TAI=obs.mjd.TAI-49000.0+59580.0) 

294 obs.mjd = mjd 

295 

296 fast_lc, fast_truth = fast_lc_gen.light_curves_from_pointings(pointings) 

297 

298 self.assertEqual(len(slow_lc), len(fast_lc)) 

299 self.assertEqual(len(slow_truth), len(fast_truth)) 

300 

301 for obj_id in slow_lc: 

302 self.assertEqual(len(fast_lc[obj_id]), len(slow_lc[obj_id])) 

303 for bp in fast_lc[obj_id]: 

304 self.assertEqual(len(fast_lc[obj_id][bp]), len(slow_lc[obj_id][bp])) 

305 for data_key in fast_lc[obj_id][bp]: 

306 self.assertEqual(fast_lc[obj_id][bp][data_key].shape, slow_lc[obj_id][bp][data_key].shape) 

307 self.assertLess(np.abs(fast_lc[obj_id][bp][data_key]-slow_lc[obj_id][bp][data_key]).max(), 

308 2.0e-10, msg='failed on %d, %s, %s' % (obj_id, bp, data_key)) 

309 

310 

311if __name__ == "__main__": 311 ↛ 312line 311 didn't jump to line 312, because the condition on line 311 was never true

312 lsst.utils.tests.init() 

313 unittest.main()