Skip to content

Implementation of extractions at different wavelengths per fiber #85

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
segasai wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Implementation of extractions at different wavelengths per fiber #85

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b3b0420
first implementation of the evalcoeff supporting multi-d wavelength
segasai Jan 9, 2025
befd581
add gpu part
segasai Jan 9, 2025
0558008
remove print
segasai Jan 9, 2025
0440558
prepare another f-n
segasai Jan 9, 2025
6fe0cb9
add get_spots tests
segasai Jan 9, 2025
511a0c6
fix docstrings
segasai Jan 9, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 38 additions & 25 deletions py/gpu_specter/extract/cpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@ def evalcoeffs(psfdata, wavelengths, specmin=0, nspec=None):

Args:
psfdata: PSF data from io.read_psf() of Gauss Hermite PSF file
wavelengths: 1D array of wavelengths
wavelengths: 1D or 2D array of wavelengths. if 2d the shape should be
(nspec0, nwave) where nspec0 is the number of
fibers in psfdata

Options:
specmin: first spectrum to include
Expand All @@ -34,43 +36,51 @@ def evalcoeffs(psfdata, wavelengths, specmin=0, nspec=None):
'''
if nspec is None:
nspec = psfdata['PSF']['COEFF'].shape[1]

p = dict(WAVE=wavelengths)
if wavelengths.ndim == 2:
wave2d = True
else:
wave2d = False
p = dict()

#- Evaluate X and Y which have different dimensionality from the
#- PSF coefficients (and might have different WAVEMIN, WAVEMAX)
meta = psfdata['XTRACE'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
p['X'] = legval(ww, psfdata['XTRACE']['X'][specmin:specmin+nspec].T)

meta = psfdata['YTRACE'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
p['Y'] = legval(ww, psfdata['YTRACE']['Y'][specmin:specmin+nspec].T)
for k in ['X', 'Y']:
meta = psfdata[k + 'TRACE'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
if wave2d:
ww = ww[specmin:specmin+nspec]
p[k] = np.zeros((nspec, ww.shape[-1]))
for i in range(nspec):
p[k][i] = legval(ww[i], psfdata[k+'TRACE'][k][specmin+i])
else:
p[k] = legval(ww.T, psfdata[k+'TRACE'][k][specmin:specmin+nspec].T)

#- Evaluate the remaining PSF coefficients with a shared dimensionality
#- and WAVEMIN, WAVEMAX
meta = psfdata['PSF'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
if wave2d:
ww = ww[specmin:specmin+nspec]
L = np.polynomial.legendre.legvander(ww, meta['LEGDEG'])

nparam = psfdata['PSF']['COEFF'].shape[0]
ndeg = psfdata['PSF']['COEFF'].shape[2]

nwave = L.shape[0]
nghx = meta['GHDEGX']+1
nghy = meta['GHDEGY']+1
nwave = wavelengths.shape[-1]
nghx = meta['GHDEGX'] + 1
nghy = meta['GHDEGY'] + 1
p['GH'] = np.zeros((nghx, nghy, nspec, nwave))
for name, coeff in zip(psfdata['PSF']['PARAM'], psfdata['PSF']['COEFF']):
name = name.strip()
coeff = coeff[specmin:specmin+nspec]
if wave2d:
curv = np.einsum('kji,ki->kj', L, coeff)
else:
curv = np.einsum('ji,ki->kj', L, coeff)
# L.dot(coeff.T).T
if name.startswith('GH-'):
i, j = map(int, name.split('-')[1:3])
p['GH'][i,j] = L.dot(coeff.T).T
p['GH'][i, j] = curv
else:
p[name] = L.dot(coeff.T).T
p[name] = curv

#- Include some additional keywords that we'll need
for key in ['HSIZEX', 'HSIZEY', 'GHDEGX', 'GHDEGY']:
Expand All @@ -85,7 +95,8 @@ def calc_pgh(ispec, wavelengths, psfparams):

Args:
ispec : integer spectrum number
wavelengths : array of wavelengths to evaluate
wavelengths : array of wavelengths to evaluate
either 1d (nwave,) or 2d (nspec,nwave)
psfparams : dictionary of PSF parameters returned by evalcoeffs

returns pGHx, pGHy
Expand All @@ -104,7 +115,7 @@ def calc_pgh(ispec, wavelengths, psfparams):
#- spot size (ny,nx)
nx = 2*p['HSIZEX']+1
ny = 2*p['HSIZEY']+1
nwave = len(wavelengths)
nwave = wavelengths.shape[-1]
# print('Spot size (ny,nx) = {},{}'.format(ny, nx))
# print('nwave = {}'.format(nwave))

Expand Down Expand Up @@ -192,15 +203,17 @@ def get_spots(specmin, nspec, wavelengths, psfdata):
Args:
specmin: first spectrum to include
nspec: number of spectra to evaluate spots for
wavelengths: 1D array of wavelengths
wavelengths: 1D or 2D array of wavelengths. if 2D the wavelength shape
needs to be (nspec0, nwave) where nspec0 is the number
of spectra in psfdata
psfdata: PSF data from io.read_psf() of Gauss Hermite PSF file

Returns:
spots: 4D array[ispec, iwave, ny, nx] of PSF spots
corners: (xc,yc) where each is 2D array[ispec,iwave] lower left corner of spot

'''
nwave = len(wavelengths)
nwave = wavelengths.shape[-1]
p = evalcoeffs(psfdata, wavelengths, specmin, nspec)
nx = 2*p['HSIZEX']+1
ny = 2*p['HSIZEY']+1
Expand Down
70 changes: 46 additions & 24 deletions py/gpu_specter/extract/gpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@ def evalcoeffs(psfdata, wavelengths, specmin=0, nspec=None):

Args:
psfdata: PSF data from io.read_psf() of Gauss Hermite PSF file
wavelengths: 1D array of wavelengths
wavelengths: 1D or 2D array of wavelengths. if 2d the shape should be
(nspec0, nwave) where nspec0 is the number of
fibers in psfdata

Options:
specmin: first spectrum to include
Expand All @@ -58,45 +60,62 @@ def evalcoeffs(psfdata, wavelengths, specmin=0, nspec=None):
if nspec is None:
nspec = psfdata['PSF']['COEFF'].shape[1]

p = dict(WAVE=wavelengths)
if wavelengths.ndim == 2:
wave2d = True
else:
wave2d = False
nwave = wavelengths.shape[-1]

p = dict()

#- Evaluate X and Y which have different dimensionality from the
#- PSF coefficients (and might have different WAVEMIN, WAVEMAX)
meta = psfdata['XTRACE'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
# TODO: Implement cuda legval
p['X'] = cp.asarray(numpy.polynomial.legendre.legval(ww, psfdata['XTRACE']['X'][specmin:specmin+nspec].T))

meta = psfdata['YTRACE'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
# TODO: Implement cuda legval
p['Y'] = cp.asarray(numpy.polynomial.legendre.legval(ww, psfdata['YTRACE']['Y'][specmin:specmin+nspec].T))
for k in ['X', 'Y']:
meta = psfdata[k + 'TRACE'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
# TODO: Implement cuda legval
if wave2d:
ww = ww[specmin:specmin+nspec]
cur_pk = np.zeros((nspec, ww.shape[-1]))
for i in range(nspec):
cur_pk[i] = numpy.polynomial.legendre.legval(ww[i], psfdata[k+'TRACE'][k][specmin+i])
p[k] = cp.asarray(cur_pk)
else:
p[k] = cp.asarray(numpy.polynomial.legendre.legval(ww.T,
psfdata[k+'TRACE'][k][specmin:specmin+nspec].T))

#- Evaluate the remaining PSF coefficients with a shared dimensionality
#- and WAVEMIN, WAVEMAX
meta = psfdata['PSF'].meta
wavemin, wavemax = meta['WAVEMIN'], meta['WAVEMAX']
ww = (wavelengths - wavemin) * (2.0 / (wavemax - wavemin)) - 1.0
L = legvander(ww, meta['LEGDEG'])

nparam = psfdata['PSF']['COEFF'].shape[0]
ndeg = psfdata['PSF']['COEFF'].shape[2]

nwave = L.shape[0]
if wave2d:
L = cp.zeros((nspec, nwave, meta['LEGDEG']+1))
for i in range(nspec):
L[i] = legvander(ww[specmin+i], meta['LEGDEG'])
else:
L = legvander(ww, meta['LEGDEG'])
# L has a shape of either nspec,nwave,ndeg or nwave, ndeg
nghx = meta['GHDEGX']+1
nghy = meta['GHDEGY']+1
p['GH'] = cp.zeros((nghx, nghy, nspec, nwave))
coeff_gpu = cp.array(native_endian(psfdata['PSF']['COEFF']))
for name, coeff in zip(psfdata['PSF']['PARAM'], coeff_gpu):
name = name.strip()
coeff = coeff[specmin:specmin+nspec]
if wave2d:
curv = cp.einsum('kji,ki->kj', L, coeff)
else:
curv = cp.einsum('ji,ki->kj', L, coeff)
# L.dot(coeff.T).T

if name.startswith('GH-'):
i, j = map(int, name.split('-')[1:3])
p['GH'][i,j] = L.dot(coeff.T).T
p['GH'][i, j] = curv
else:
p[name] = L.dot(coeff.T).T
p[name] = curv

#- Include some additional keywords that we'll need
for key in ['HSIZEX', 'HSIZEY', 'GHDEGX', 'GHDEGY']:
Expand All @@ -110,7 +129,8 @@ def calc_pgh(ispec, wavelengths, psfparams):
Calculate the pixelated Gauss Hermite for all wavelengths of a single spectrum

ispec : integer spectrum number
wavelengths : array of wavelengths to evaluate
wavelengths : array of wavelengths to evaluate
either 1d (nwave,) or 2d (nspec,nwave)
psfparams : dictionary of PSF parameters returned by evalcoeffs

returns pGHx, pGHy
Expand All @@ -128,7 +148,7 @@ def calc_pgh(ispec, wavelengths, psfparams):
#- spot size (ny,nx)
nx = 2*p['HSIZEX'] + 1
ny = 2*p['HSIZEY'] + 1
nwave = len(wavelengths)
nwave = wavelengths.shape[-1]
#- convert to cupy arrays
for k in ['X', 'Y', 'GHSIGX', 'GHSIGY']:
p[k] = cp.asarray(p[k])
Expand Down Expand Up @@ -222,15 +242,17 @@ def get_spots(specmin, nspec, wavelengths, psfdata):
Args:
specmin: first spectrum to include
nspec: number of spectra to evaluate spots for
wavelengths: 1D array of wavelengths
wavelengths: 1D or 2D array of wavelengths. if 2D the wavelength shape
needs to be (nspec0, nwave) where nspec0 is the number
of spectra in psfdata
psfdata: PSF data from io.read_psf() of Gauss Hermite PSF file

Returns:
spots: 4D array[ispec, iwave, ny, nx] of PSF spots
corners: (xc,yc) where each is 2D array[ispec,iwave] lower left corner of spot

'''
nwave = len(wavelengths)
nwave = wavelengths.shape[-1]
p = evalcoeffs(psfdata, wavelengths, specmin, nspec)
nx = 2*p['HSIZEX']+1
ny = 2*p['HSIZEY']+1
Expand Down
8 changes: 8 additions & 0 deletions py/gpu_specter/test/test_psfcoeff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,10 @@ def test_basics(self):
self.assertEqual(psfparams['X'].shape, (25, nwave))
psfparams = evalcoeffs(psfdata, wavelengths, specmin=25, nspec=5)
self.assertEqual(psfparams['Y'].shape, (5, nwave))
wavelengths = np.linspace(meta['WAVEMIN']-10, meta['WAVEMAX']+10, nwave)
psfparams = evalcoeffs(psfdata, wavelengths[None, :] + np.zeros(nspec)[:,
None], specmin=0) #, nspec=25)
self.assertEqual(psfparams['X'].shape, (nspec, nwave))

@unittest.skipIf(not gpu_available, 'gpu not available')
def test_gpu_basics(self):
Expand All @@ -72,6 +76,10 @@ def test_gpu_basics(self):
self.assertEqual(psfparams['X'].shape, (25, nwave))
psfparams = gpu_evalcoeffs(psfdata, wavelengths, specmin=25, nspec=5)
self.assertEqual(psfparams['Y'].shape, (5, nwave))
wavelengths = np.linspace(meta['WAVEMIN']-10, meta['WAVEMAX']+10, nwave)
psfparams = gpu_evalcoeffs(psfdata, wavelengths[None, :] + np.zeros(nspec)[:,
None], specmin=0) #, nspec=25)
self.assertEqual(psfparams['X'].shape, (nspec, nwave))

@unittest.skipIf(not gpu_available, 'gpu not available')
def test_compare_gpu(self):
Expand Down
Loading