def modelMatch(targetFitsPath, templateDir, waveRange=np.arange(4000, 7600, 1), isSave=True): keywords, flux = astroUtil.getFitsHeaderAndFlux(targetFitsPath) wave = astroUtil.calcWaveWithKeywords(keywords) wave, flux = formatWaveAndFlux(wave, flux, waveRange) templatePath, minCoeff, minChi2 = calcMinChi2(flux, templateDir) # print "%s,%s,%s" % (targetFitsPath, templatePath, minChi2) if isSave: if templatePath == "": return templateData = np.loadtxt(templatePath) wave_template, flux_template = formatWaveAndFlux(templateData[:, 0], templateData[:, 1], waveRange) plt.plot(wave, flux, 'b') plt.plot(wave_template, flux_template * (minCoeff[0] * np.power(wave_template, 2) + minCoeff[1] * wave_template + minCoeff[2]), 'r--') plt.ylabel('flux') plt.xlabel('wavelength') plt.title(templatePath) plt.savefig(targetFitsPath + '.png') plt.close() return templatePath, minChi2
classCount = {"O": 0, "B": 0, "A": 0, "F": 0, "G": 0, "K": 0, "M": 0, "N": 0} inputPath = "C:\dr2" outFile = "data.txt" outFilePath = util.getRelativePath(outFile) output = open(outFilePath, "w") dirs = os.listdir(inputPath) count = 0 for dir in dirs: files = os.listdir(os.path.join(inputPath, dir)) for fileName in files: filePath = os.path.join(inputPath, dir, fileName) try: keywords, flux = astroUtil.getFitsHeaderAndFlux(filePath) record = [] for fea in features: record.append(keywords[fea]) for wl in waveLens: record.append(astroUtil.checkLine(flux, wl)) subClass = keywords[subClassKey][0:1] subClassTag = classMap[subClass] classCount[subClass] += 1 record.append(subClassTag) output.write(util.listToString(record)) output.write("\n") count += 1