def __init__(self, exp, *, doTweakCentroid=True, doForceCoM=False, savePlots=None, centroid=None, boxHalfSize=50): self.exp = exp self.savePlots = savePlots self.doTweakCentroid = doTweakCentroid self.doForceCoM = doForceCoM self.boxHalfSize = boxHalfSize if centroid is None: qfmTaskConfig = QuickFrameMeasurementTaskConfig() qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig) result = qfmTask.run(exp) if not result.success: msg = ("Failed to automatically find source in image. " "Either provide a centroid manually or use a new image") raise RuntimeError(msg) self.centroid = result.brightestObjCentroid else: self.centroid = centroid self.imStats = getImageStats(self.exp) # need the background levels now self.data = self.getStarBoxData() if self.doTweakCentroid: self.tweakCentroid(self.doForceCoM) self.data = self.getStarBoxData() self.xx, self.yy = self.getMeshGrid(self.data) self.imStats.centroid = self.centroid self.imStats.intCentroid = self.intCoords(self.centroid) self.imStats.intCentroidRounded = self.intRoundCoords(self.centroid) self.imStats.nStatPixInBox = self.nSatPixInBox self.radialAverageAndFit()
class SpectrumExaminer(): """Task for the QUICK spectral extraction of single-star dispersed images. For a full description of how this tasks works, see the run() method. """ # ConfigClass = SummarizeImageTaskConfig # _DefaultName = "summarizeImage" def __init__(self, exp, display=None, debug=False, savePlotAs=None, **kwargs): super().__init__(**kwargs) self.exp = exp self.display = display self.debug = debug self.savePlotAs = savePlotAs qfmTaskConfig = QuickFrameMeasurementTaskConfig() self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig) pstConfig = ProcessStarTask.ConfigClass() pstConfig.offsetFromMainStar = 400 self.processStarTask = ProcessStarTask(config=pstConfig) self.imStats = getImageStats(exp) self.init() @staticmethod def bboxToAwfDisplayLines(box): """Takes a bbox, returns a list of lines such that they can be plotted: for line in lines: display.line(line, ctype='red') """ x0 = box.beginX x1 = box.endX y0 = box.beginY y1 = box.endY return [[(x0, y0), (x1, y0)], [(x0, y0), (x0, y1)], [(x1, y0), (x1, y1)], [(x0, y1), (x1, y1)]] def eraseDisplay(self): if self.display: self.display.erase() def displaySpectrumBbox(self): if self.display: lines = self.bboxToAwfDisplayLines(self.spectrumbbox) for line in lines: self.display.line(line, ctype='red') else: print("No display set") def displayStarLocation(self): if self.display: self.display.dot('x', *self.qfmResult.brightestObjCentroid, size=50) self.display.dot('o', *self.qfmResult.brightestObjCentroid, size=50) else: print("No display set") def calcGoodSpectrumSection(self, threshold=5, windowSize=5): length = len(self.ridgeLineLocations) chunks = length // windowSize stddevs = [] for i in range(chunks + 1): stddevs.append( np.std(self.ridgeLineLocations[i * windowSize:(i + 1) * windowSize])) goodPoints = np.where(np.asarray(stddevs) < threshold)[0] minPoint = (goodPoints[2] - 2) * windowSize maxPoint = (goodPoints[-3] + 3) * windowSize minPoint = max(minPoint, 0) maxPoint = min(maxPoint, length) if self.debug: plt.plot(range(0, length + 1, windowSize), stddevs) plt.hlines(threshold, 0, length, colors='r', ls='dashed') plt.vlines(minPoint, 0, max(stddevs) + 10, colors='k', ls='dashed') plt.vlines(maxPoint, 0, max(stddevs) + 10, colors='k', ls='dashed') plt.title(f'Ridgeline scatter, windowSize={windowSize}') return (minPoint, maxPoint) def fit(self): def gauss(x, a, x0, sigma): return a * np.exp(-(x - x0)**2 / (2 * sigma**2)) data = self.spectrumData[self.goodSlice] nRows, nCols = data.shape # don't subtract the row median or even a percentile - seems bad # fitting a const also seems bad - needs some better thought parameters = np.zeros((nRows, 3)) pCovs = [] xs = np.arange(nCols) for rowNum, row in enumerate(data): peakPos = self.ridgeLineLocations[rowNum] amplitude = row[peakPos] width = 7 try: pars, pCov = curve_fit(gauss, xs, row, [amplitude, peakPos, width], maxfev=100) pCovs.append(pCov) except RuntimeError: pars = [np.nan] * 3 if not np.all([p < 1e7 for p in pars]): pars = [np.nan] * 3 parameters[rowNum] = pars parameters[:, 0] = np.abs(parameters[:, 0]) parameters[:, 2] = np.abs(parameters[:, 2]) self.parameters = parameters def plot(self, saveAs=None): fig = plt.figure(figsize=(10, 10)) # spectrum ax0 = plt.subplot2grid((4, 4), (0, 0), colspan=3) ax0.tick_params(axis='x', top=True, bottom=False, labeltop=True, labelbottom=False) d = self.spectrumData[self.goodSlice].T vmin = np.percentile(d, 1) vmax = np.percentile(d, 99) pos = ax0.imshow(self.spectrumData[self.goodSlice].T, vmin=vmin, vmax=vmax, origin='lower') div = make_axes_locatable(ax0) cax = div.append_axes("bottom", size="7%", pad="8%") fig.colorbar(pos, cax=cax, orientation="horizontal", label="Counts") # spectrum histogram axHist = plt.subplot2grid((4, 4), (0, 3)) data = self.spectrumData histMax = np.nanpercentile(data, 99.99) histMin = np.nanpercentile(data, 0.001) axHist.hist(data[(data >= histMin) & (data <= histMax)].flatten(), bins=100) underflow = len(data[data < histMin]) overflow = len(data[data > histMax]) axHist.set_yscale('log', nonpositive='clip') axHist.set_title('Spectrum pixel histogram') text = f"Underflow = {underflow}" text += f"\nOverflow = {overflow}" anchored_text = AnchoredText(text, loc=1, pad=0.5) axHist.add_artist(anchored_text) # peak fluxes ax1 = plt.subplot2grid((4, 4), (1, 0), colspan=3) ax1.plot(self.ridgeLineValues[self.goodSlice], label='Raw peak value') ax1.plot(self.parameters[:, 0], label='Fitted amplitude') ax1.axhline(self.continuumFlux98, ls='dashed', color='g') ax1.set_ylabel('Peak amplitude (ADU)') ax1.set_xlabel('Spectrum position (pixels)') ax1.legend(title=f"Continuum flux = {self.continuumFlux98:.0f} ADU", loc="center right", framealpha=0.2, facecolor="black") ax1.set_title('Ridgeline plot') # FWHM ax2 = plt.subplot2grid((4, 4), (2, 0), colspan=3) ax2.plot(self.parameters[:, 2] * 2.355, label="FWHM (pix)") fwhmValues = self.parameters[:, 2] * 2.355 amplitudes = self.parameters[:, 0] minVal, maxVal = self.getStableFwhmRegion(fwhmValues, amplitudes) medianFwhm, bestFwhm = self.getMedianAndBestFwhm( fwhmValues, minVal, maxVal) ax2.axhline(medianFwhm, ls='dashed', color='k', label=f"Median FWHM = {medianFwhm:.1f} pix") ax2.axhline(bestFwhm, ls='dashed', color='r', label=f"Best FWHM = {bestFwhm:.1f} pix") ax2.axvline(minVal, ls='dashed', color='k', alpha=0.2) ax2.axvline(maxVal, ls='dashed', color='k', alpha=0.2) ymin = max(np.nanmin(fwhmValues) - 5, 0) ymax = medianFwhm * 2 ax2.set_ylim(ymin, ymax) ax2.set_ylabel('FWHM (pixels)') ax2.set_xlabel('Spectrum position (pixels)') ax2.legend(loc="upper right", framealpha=0.2, facecolor="black") ax2.set_title('Spectrum FWHM') # row fluxes ax3 = plt.subplot2grid((4, 4), (3, 0), colspan=3) ax3.plot(self.rowSums[self.goodSlice], label="Sum across row") ax3.set_ylabel('Total row flux (ADU)') ax3.set_xlabel('Spectrum position (pixels)') ax3.legend(framealpha=0.2, facecolor="black") ax3.set_title('Row sums') # textbox top # ax4 = plt.subplot2grid((4, 4), (1, 3)) ax4 = plt.subplot2grid((4, 4), (1, 3), rowspan=2) text = "short text" text = self.generateStatsTextboxContent(0) text += self.generateStatsTextboxContent(1) text += self.generateStatsTextboxContent(2) text += self.generateStatsTextboxContent(3) stats_text = AnchoredText(text, loc="center", pad=0.5, prop=dict(size=10.5, ma="left", backgroundcolor="white", color="black", family='monospace')) ax4.add_artist(stats_text) ax4.axis('off') # textbox middle if self.debug: ax5 = plt.subplot2grid((4, 4), (2, 3)) text = self.generateStatsTextboxContent(-1) stats_text = AnchoredText(text, loc="center", pad=0.5, prop=dict(size=10.5, ma="left", backgroundcolor="white", color="black", family='monospace')) ax5.add_artist(stats_text) ax5.axis('off') plt.tight_layout() plt.show() if self.savePlotAs: fig.savefig(self.savePlotAs) def init(self): pass def generateStatsTextboxContent(self, section, doPrint=True): x, y = self.qfmResult.brightestObjCentroid exptime = self.exp.getInfo().getVisitInfo().getExposureTime() info = self.exp.getInfo() vi = info.getVisitInfo() fullFilterString = info.getFilterLabel().physicalLabel filt = fullFilterString.split(FILTER_DELIMITER)[0] grating = fullFilterString.split(FILTER_DELIMITER)[1] airmass = vi.getBoresightAirmass() rotangle = vi.getBoresightRotAngle().asDegrees() azAlt = vi.getBoresightAzAlt() az = azAlt[0].asDegrees() el = azAlt[1].asDegrees() md = self.exp.getMetadata() obsInfo = ObservationInfo(md, subset={'object'}) obj = obsInfo.object lines = [] if section == 0: lines.append("----- Star stats -----") lines.append(f"Star centroid @ {x:.0f}, {y:.0f}") lines.append(f"Star max pixel = {self.starPeakFlux:,.0f} ADU") lines.append( f"Star Ap25 flux = {self.qfmResult.brightestObjApFlux25:,.0f} ADU" ) lines.extend(["", ""]) # section break return '\n'.join([line for line in lines]) if section == 1: lines.append("------ Image stats ---------") imageMedian = np.median(self.exp.image.array) lines.append(f"Image median = {imageMedian:.2f} ADU") lines.append(f"Exposure time = {exptime:.2f} s") lines.extend(["", ""]) # section break return '\n'.join([line for line in lines]) if section == 2: lines.append("------- Rate stats ---------") lines.append( f"Star max pixel = {self.starPeakFlux/exptime:,.0f} ADU/s") lines.append( f"Spectrum contiuum = {self.continuumFlux98/exptime:,.1f} ADU/s" ) lines.extend(["", ""]) # section break return '\n'.join([line for line in lines]) if section == 3: lines.append("----- Observation info -----") lines.append(f"object = {obj}") lines.append(f"filter = {filt}") lines.append(f"grating = {grating}") lines.append(f"rotpa = {rotangle:.1f}") lines.append(f"az = {az:.1f}") lines.append(f"el = {el:.1f}") lines.append(f"airmass = {airmass:.3f}") return '\n'.join([line for line in lines]) if section == -1: # special -1 for debug lines.append("---------- Debug -----------") lines.append(f"spectrum bbox: {self.spectrumbbox}") lines.append( f"Good range = {self.goodSpectrumMinY},{self.goodSpectrumMaxY}" ) return '\n'.join([line for line in lines]) return def run(self): self.qfmResult = self.qfmTask.run(self.exp) self.intCentroidX = int( np.round(self.qfmResult.brightestObjCentroid)[0]) self.intCentroidY = int( np.round(self.qfmResult.brightestObjCentroid)[1]) self.starPeakFlux = self.exp.image.array[self.intCentroidY, self.intCentroidX] self.spectrumbbox = self.processStarTask.calcSpectrumBBox( self.exp, self.qfmResult.brightestObjCentroid, 200) self.spectrumData = self.exp.image[self.spectrumbbox].array self.ridgeLineLocations = np.argmax(self.spectrumData, axis=1) self.ridgeLineValues = self.spectrumData[ range(self.spectrumbbox.getHeight()), self.ridgeLineLocations] self.rowSums = np.sum(self.spectrumData, axis=1) coords = self.calcGoodSpectrumSection() self.goodSpectrumMinY = coords[0] self.goodSpectrumMaxY = coords[1] self.goodSlice = slice(coords[0], coords[1]) self.continuumFlux90 = np.percentile(self.ridgeLineValues, 90) # for emission stars self.continuumFlux98 = np.percentile(self.ridgeLineValues, 98) # for most stars self.fit() self.plot() return @staticmethod def getMedianAndBestFwhm(fwhmValues, minIndex, maxIndex): with warnings.catch_warnings( ): # to supress nan warnings, which are fine warnings.simplefilter("ignore") clippedValues = sigma_clip(fwhmValues[minIndex:maxIndex]) # cast back with asArray needed becase sigma_clip returns # masked array which doesn't play nice with np.nan<med/percentile> clippedValues = np.asarray(clippedValues) medianFwhm = np.nanmedian(clippedValues) bestFocusFwhm = np.nanpercentile(np.asarray(clippedValues), 2) return medianFwhm, bestFocusFwhm def getStableFwhmRegion(self, fwhmValues, amplitudes, smoothing=1, maxDifferential=4): # smooth the fwhmValues values # differentiate # take the longest contiguous region of 1s # check section corresponds to top 25% in ampl to exclude 2nd order # if not, pick next longest run, etc # walk out from ends of that list over bumps smaller than maxDiff smoothFwhm = np.convolve(fwhmValues, np.ones(smoothing) / smoothing, mode='same') diff = np.diff(smoothFwhm, append=smoothFwhm[-1]) indices = np.where(1 - np.abs(diff) < 1)[0] diffIndices = np.diff(indices) # [list(g) for k, g in groupby('AAAABBBCCD')] -->[['A', 'A', 'A', 'A'], # ... ['B', 'B', 'B'], ['C', 'C'], ['D']] indexLists = [list(g) for k, g in groupby(diffIndices)] listLengths = [len(lst) for lst in indexLists] amplitudeThreshold = np.nanpercentile(amplitudes, 75) sortedListLengths = sorted(listLengths) for listLength in sortedListLengths[::-1]: longestListLength = listLength longestListIndex = listLengths.index(longestListLength) longestListStartTruePosition = int( np.sum(listLengths[0:longestListIndex])) longestListStartTruePosition += int(longestListLength / 2) # we want the mid-run value if amplitudes[longestListStartTruePosition] > amplitudeThreshold: break startOfLongList = np.sum(listLengths[0:longestListIndex]) endOfLongList = startOfLongList + longestListLength endValue = endOfLongList for lst in indexLists[longestListIndex + 1:]: value = lst[0] if value > maxDifferential: break endValue += len(lst) startValue = startOfLongList for lst in indexLists[longestListIndex - 1::-1]: value = lst[0] if value > maxDifferential: break startValue -= len(lst) startValue = int(max(0, startValue)) endValue = int(min(len(fwhmValues), endValue)) if not self.debug: return startValue, endValue medianFwhm, bestFocusFwhm = self.getMedianAndBestFwhm( fwhmValues, startValue, endValue) xlim = (-20, len(fwhmValues)) plt.figure(figsize=(10, 6)) plt.plot(fwhmValues) plt.vlines(startValue, 0, 50, 'r') plt.vlines(endValue, 0, 50, 'r') plt.hlines(medianFwhm, xlim[0], xlim[1]) plt.hlines(bestFocusFwhm, xlim[0], xlim[1], 'r', ls='--') plt.vlines(startOfLongList, 0, 50, 'g') plt.vlines(endOfLongList, 0, 50, 'g') plt.ylim(0, 200) plt.xlim(xlim) plt.show() plt.figure(figsize=(10, 6)) plt.plot(diffIndices) plt.vlines(startValue, 0, 50, 'r') plt.vlines(endValue, 0, 50, 'r') plt.vlines(startOfLongList, 0, 50, 'g') plt.vlines(endOfLongList, 0, 50, 'g') plt.ylim(0, 30) plt.xlim(xlim) plt.show() return startValue, endValue
class SpectralFocusAnalyzer(): """Analyze a focus sweep taken for spectral data. Take slices across the spectrum for each image, fitting a Gaussian to each slice, and perform a parabolic fit to these widths. The number of slices and their distances can be customized by calling setSpectrumBoxOffsets(). Nominal usage is something like: %matplotlib inline dayObs = 20210101 seqNums = [100, 101, 102, 103, 104] focusAnalyzer = SpectralFocusAnalyzer() focusAnalyzer.setSpectrumBoxOffsets([500, 750, 1000, 1250]) focusAnalyzer.getFocusData(dayObs, seqNums, doDisplay=True) focusAnalyzer.fitDataAndPlot() focusAnalyzer.run() can be used instead of the last two lines separately. """ def __init__(self, location, **kwargs): self.butler = makeDefaultLatissButler(location) repoDir = LATISS_REPO_LOCATION_MAP[location] self._bestEffort = BestEffortIsr(repoDir, **kwargs) qfmTaskConfig = QuickFrameMeasurementTaskConfig() self._quickMeasure = QuickFrameMeasurementTask(config=qfmTaskConfig) self.spectrumHalfWidth = 100 self.spectrumBoxLength = 20 self._spectrumBoxOffsets = [882, 1170, 1467] self._setColors(len(self._spectrumBoxOffsets)) def setSpectrumBoxOffsets(self, offsets): """Set the current spectrum slice offsets. Parameters ---------- offsets : `list` of `float` The distance at which to slice the spectrum, measured in pixels from the main star's location. """ self._spectrumBoxOffsets = offsets self._setColors(len(offsets)) def getSpectrumBoxOffsets(self): """Get the current spectrum slice offsets. Returns ------- offsets : `list` of `float` The distance at which to slice the spectrum, measured in pixels from the main star's location. """ return self._spectrumBoxOffsets def _setColors(self, nPoints): self.COLORS = cm.rainbow(np.linspace(0, 1, nPoints)) @staticmethod def _getFocusFromHeader(exp): return float(exp.getMetadata()["FOCUSZ"]) def _getBboxes(self, centroid): x, y = centroid bboxes = [] for offset in self._spectrumBoxOffsets: bbox = geom.Box2I( geom.Point2I(x - self.spectrumHalfWidth, y + offset), geom.Point2I(x + self.spectrumHalfWidth, y + offset + self.spectrumBoxLength)) bboxes.append(bbox) return bboxes def _bboxToMplRectangle(self, bbox, colorNum): xmin = bbox.getBeginX() ymin = bbox.getBeginY() xsize = bbox.getWidth() ysize = bbox.getHeight() rectangle = Rectangle((xmin, ymin), xsize, ysize, alpha=1, facecolor='none', lw=2, edgecolor=self.COLORS[colorNum]) return rectangle @staticmethod def gauss(x, *pars): amp, mean, sigma = pars return amp * np.exp(-(x - mean)**2 / (2. * sigma**2)) def run(self, dayObs, seqNums, doDisplay=False, hideFit=False, hexapodZeroPoint=0): """Perform a focus sweep analysis for spectral data. For each seqNum for the specified dayObs, take a slice through the spectrum at y-offsets as specified by the offsets (see get/setSpectrumBoxOffsets() for setting these) and fit a Gaussian to the spectrum slice to measure its width. For each offset distance, fit a parabola to the fitted spectral widths and return the hexapod position at which the best focus was achieved for each. Parameters ---------- dayObs : `int` The dayObs to use. seqNums : `list` of `int` The seqNums for the focus sweep to analyze. doDisplay : `bool` Show the plots? Designed to be used in a notebook with %matplotlib inline. hideFit : `bool`, optional Hide the fit and just return the result? hexapodZeroPoint : `float`, optional Add a zeropoint offset to the hexapod axis? Returns ------- bestFits : `list` of `float` A list of the best fit focuses, one for each spectral slice. """ self.getFocusData(dayObs, seqNums, doDisplay=doDisplay) bestFits = self.fitDataAndPlot(hideFit=hideFit, hexapodZeroPoint=hexapodZeroPoint) return bestFits def getFocusData(self, dayObs, seqNums, doDisplay=False): """Perform a focus sweep analysis for spectral data. For each seqNum for the specified dayObs, take a slice through the spectrum at y-offsets as specified by the offsets (see get/setSpectrumBoxOffsets() for setting these) and fit a Gaussian to the spectrum slice to measure its width. Parameters ---------- dayObs : `int` The dayObs to use. seqNums : `list` of `int` The seqNums for the focus sweep to analyze. doDisplay : `bool` Show the plots? Designed to be used in a notebook with %matplotlib inline. Notes ----- Performs the focus analysis per-image, holding the data in the class. Call fitDataAndPlot() after running this to perform the parabolic fit to the focus data itself. """ fitData = {} filters = set() objects = set() for seqNum in seqNums: fitData[seqNum] = {} dataId = {'day_obs': dayObs, 'seq_num': seqNum, 'detector': 0} exp = self._bestEffort.getExposure(dataId) # sanity checking filt = exp.getFilter().getName() expRecord = getExpRecordFromDataId(self.butler, dataId) obj = expRecord.target_name objects.add(obj) filters.add(filt) assert isDispersedExp( exp), f"Image is not dispersed! (filter = {filt})" assert len(filters) == 1, "You accidentally mixed filters!" assert len(objects) == 1, "You accidentally mixed objects!" quickMeasResult = self._quickMeasure.run(exp) centroid = quickMeasResult.brightestObjCentroid spectrumSliceBboxes = self._getBboxes( centroid) # inside the loop due to centroid shifts if doDisplay: fig, axes = plt.subplots(1, 2, figsize=(18, 9)) exp.image.array[exp.image.array <= 0] = 0.001 axes[0].imshow(exp.image.array, norm=LogNorm(), origin='lower', cmap='gray_r') plt.tight_layout() arrowy, arrowx = centroid[0] - 400, centroid[ 1] # numpy is backwards dx, dy = 0, 300 arrow = Arrow(arrowy, arrowx, dy, dx, width=200., color='red') circle = Circle(centroid, radius=25, facecolor='none', color='red') axes[0].add_patch(arrow) axes[0].add_patch(circle) for i, bbox in enumerate(spectrumSliceBboxes): rect = self._bboxToMplRectangle(bbox, i) axes[0].add_patch(rect) for i, bbox in enumerate(spectrumSliceBboxes): data1d = np.mean(exp[bbox].image.array, axis=0) # flatten data1d -= np.median(data1d) xs = np.arange(len(data1d)) # get rough estimates for fit # can't use sigma from quickMeasResult due to SDSS shape # failing on saturated starts, and fp.getShape() is weird amp = np.max(data1d) mean = np.argmax(data1d) sigma = 20 p0 = amp, mean, sigma try: coeffs, var_matrix = curve_fit(self.gauss, xs, data1d, p0=p0) except RuntimeError: coeffs = (np.nan, np.nan, np.nan) fitData[seqNum][i] = FitResult(amp=abs(coeffs[0]), mean=coeffs[1], sigma=abs(coeffs[2])) if doDisplay: axes[1].plot(xs, data1d, 'x', c=self.COLORS[i]) highResX = np.linspace(0, len(data1d), 1000) if coeffs[0] is not np.nan: axes[1].plot(highResX, self.gauss(highResX, *coeffs), 'k-') if doDisplay: # show all color boxes together plt.title(f'Fits to seqNum {seqNum}') plt.show() focuserPosition = self._getFocusFromHeader(exp) fitData[seqNum]['focus'] = focuserPosition self.fitData = fitData self.filter = filters.pop() self.object = objects.pop() return def fitDataAndPlot(self, hideFit=False, hexapodZeroPoint=0): """Fit a parabola to each series of slices and return the best focus. For each offset distance, fit a parabola to the fitted spectral widths and return the hexapod position at which the best focus was achieved for each. Parameters ---------- hideFit : `bool`, optional Hide the fit and just return the result? hexapodZeroPoint : `float`, optional Add a zeropoint offset to the hexapod axis? Returns ------- bestFits : `list` of `float` A list of the best fit focuses, one for each spectral slice. """ data = self.fitData filt = self.filter obj = self.object bestFits = [] titleFontSize = 18 legendFontSize = 12 labelFontSize = 14 arcminToPixel = 10 sigmaToFwhm = 2.355 f, axes = plt.subplots(2, 1, figsize=[10, 12]) focusPositions = [ data[k]['focus'] - hexapodZeroPoint for k in sorted(data.keys()) ] fineXs = np.linspace(np.min(focusPositions), np.max(focusPositions), 101) seqNums = sorted(data.keys()) nSpectrumSlices = len(data[list(data.keys())[0]]) - 1 pointsForLegend = [0.0 for offset in range(nSpectrumSlices)] for spectrumSlice in range( nSpectrumSlices ): # the blue/green/red slices through the spectrum # for scatter plots, the color needs to be a single-row 2d array thisColor = np.array([self.COLORS[spectrumSlice]]) amps = [data[seqNum][spectrumSlice].amp for seqNum in seqNums] widths = [ data[seqNum][spectrumSlice].sigma / arcminToPixel * sigmaToFwhm for seqNum in seqNums ] pointsForLegend[spectrumSlice] = axes[0].scatter(focusPositions, amps, c=thisColor) axes[0].set_xlabel('Focus position (mm)', fontsize=labelFontSize) axes[0].set_ylabel('Height (ADU)', fontsize=labelFontSize) axes[1].scatter(focusPositions, widths, c=thisColor) axes[1].set_xlabel('Focus position (mm)', fontsize=labelFontSize) axes[1].set_ylabel('FWHM (arcsec)', fontsize=labelFontSize) quadFitPars = np.polyfit(focusPositions, widths, 2) if not hideFit: axes[1].plot(fineXs, np.poly1d(quadFitPars)(fineXs), c=self.COLORS[spectrumSlice]) fitMin = -quadFitPars[1] / (2.0 * quadFitPars[0]) bestFits.append(fitMin) axes[1].axvline(fitMin, color=self.COLORS[spectrumSlice]) msg = f"Best focus offset = {np.round(fitMin, 2)}" axes[1].text(fitMin, np.mean(widths), msg, horizontalalignment='right', verticalalignment='center', rotation=90, color=self.COLORS[spectrumSlice], fontsize=legendFontSize) titleText = f"Focus curve for {obj} w/ {filt}" plt.suptitle(titleText, fontsize=titleFontSize) legendText = self._generateLegendText(nSpectrumSlices) axes[0].legend(pointsForLegend, legendText, fontsize=legendFontSize) axes[1].legend(pointsForLegend, legendText, fontsize=legendFontSize) f.tight_layout(rect=[0, 0.03, 1, 0.95]) plt.show() for i, bestFit in enumerate(bestFits): print(f"Best fit for spectrum slice {i} = {bestFit:.4f}mm") return bestFits def _generateLegendText(self, nSpectrumSlices): if nSpectrumSlices == 1: return ['m=+1 spectrum slice'] if nSpectrumSlices == 2: return ['m=+1 blue end', 'm=+1 red end'] legendText = [] legendText.append('m=+1 blue end') for i in range(nSpectrumSlices - 2): legendText.append('m=+1 redder...') legendText.append('m=+1 red end') return legendText