CODEWITHSUNDEEP
pythonpyqt5pyqtgraph
Jesse Rio
Jesse Rio

Reputation: 95

How to plot 2D data (image) with x and y coordinates

I would like to plot a 2D spectrum where the x coordinate is the spectral dimension (wavelength) and the y coordinate is the spatial dimension (in arcseconds of the sky) using pyqtgraph.

I've been able to do this using an ImageItem() but I cannot seem to figure out how to display the x and y axes in the right coordinates. I don't want to just change the labels or the ticks but indeed the coordinates of the plot because I later need to perform operations using these values (wavelength and arcsecs).

Here is a minimal working example:

import pyqtgraph as pg
import numpy as np


# The fake data
wvlg = np.linspace(300, 600, 5000)
arcsec = np.linspace(-5, 5, 100)
flux = np.ones((wvlg.shape[0], arcsec.shape[0])) * np.exp(-(arcsec)**2/0.1)
flux += np.random.normal(0, 0.1, size=(wvlg.shape[0], arcsec.shape[0]))

# The plotting
win = pg.GraphicsLayoutWidget(show=True)
ax2D = win.addPlot(title='2D spectrum', row=0, col=0)
img = pg.ImageItem()
img.setImage(flux)
ax2D.addItem(img)
# Some line converting the x and y values to wvlg and arcsec

This gives an image where the x and y axis show the index value, whereas I would like to show the corresponding wavelength and arcsec values.

The 2D spectrum plotted but without the coordinates in the x and y axes

Is there an easy way to do this that I just grossly overlooked in the documentation?

Upvotes: 1

Views: 1898

Answers (1)

titusjan
titusjan

Reputation: 5546

You can use the setRect method of the ImageItem class to set the extent of the data. See my example below.

Note that I moved the image by half a pixel so that the pixel centers match the exact coordinates. Otherwise the coordinates would align with one of the pixel's corner points.

import pyqtgraph as pg
import numpy as np

from PyQt5 import QtCore, QtWidgets


def main():
    app = QtWidgets.QApplication([])
    
    # The fake data
    wvlg = np.linspace(300, 600, 5000)
    arcsec = np.linspace(-5, 5, 100)
    flux = np.ones((wvlg.shape[0], arcsec.shape[0])) * np.exp(-(arcsec)**2/0.1)
    flux += np.random.normal(0, 0.1, size=(wvlg.shape[0], arcsec.shape[0]))

    # The plotting
    win = pg.GraphicsLayoutWidget()
    ax2D = win.addPlot(title='2D spectrum', row=0, col=0)
    img = pg.ImageItem()
    img.setImage(flux)
    ax2D.addItem(img)
    
    print(flux.shape)

    # Move the image by half a pixel so that the center of the pixels are
    # located at the coordinate values
    dx = wvlg[1]-wvlg[0]
    dy = arcsec[1]-arcsec[0]
    print("pixel size x: {}, pixel size y: {}".format(dx, dy))
    
    rect = QtCore.QRectF(wvlg[0] - dx/2, arcsec[0] - dy/2, 
                         wvlg[-1] - wvlg[0], arcsec[-1] - arcsec[0])
    print(rect)
    img.setRect(rect)

    ax2D.setLabels(left='arcsec', bottom='wvlg')
    
    win.show()
    win.raise_()
    app.exec_()
    
if __name__ == "__main__":
    main()

Upvotes: 2

Related Questions