Skip to content

Legendre nnls #287

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
craigwarner-ufastro wants to merge 7 commits into
base: main
Choose a base branch
from
Open

Legendre nnls #287

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c4132a5
Modified archetypes using BVLS to add -1*all legendre terms to the basis
craigwarner-ufastro Feb 27, 2024
a660106
Modify prior array in prior_on_coeffs to include cross-terms for
craigwarner-ufastro Mar 12, 2024
7911c5b
Only use additional negative legendre terms and NNLS when in GPU mode
craigwarner-ufastro Mar 12, 2024
bbf187b
Merge branch 'main' into legendre_nnls
sbailey Apr 12, 2024
1ac311a
Moved as much logic as possible related to using NNLS with extra
craigwarner-ufastro Apr 21, 2024
d1b7b8f
Compress output COEFF when using BVLS-NNLS trick to remove legendre
craigwarner-ufastro May 9, 2024
fd76617
Refactor to bring more in-line with main where the conflicts mostly
craigwarner-ufastro May 9, 2024
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
74 changes: 44 additions & 30 deletions py/redrock/archetypes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,8 @@ def __init__(self, filename):

# TODO: Allow Archetype files to specify their IGM model
self.igm_model = 'Inoue14'
# TODO: Allow Archetype files to specify bvls or nnls or pca solver method
self._solver_method = 'bvls'

return

Expand Down Expand Up @@ -165,7 +167,7 @@ def eval(self, subtype, dwave, coeff, wave, z):

return flux

def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=None, binned=None, use_gpu=False, prior=None, ncam=None):
def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=None, binned=None, use_gpu=False, prior_sigma=None, ncam=None):

"""Nearest neighbour archetype approach; fitting with a combinating of nearest archetypes in chi2 space

Expand All @@ -183,7 +185,7 @@ def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest,
dedges (dict): in GPU mode, use pre-computed dict of wavelength bin edges, already on GPU
binned (dict): already computed dictionary of rebinned fluxes
use_gpu (bool): use GPU or not
prior (2d array): prior matrix on coefficients (1/sig**2)
prior_sigma (float): prior to add in the final solution matrix: added as 1/(prior_sigma**2) only for per-camera mode
ncam (int): Number of camera for given Instrument

Returns:
Expand All @@ -197,11 +199,11 @@ def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest,
#is multiplied by trans in get_best_archetype
spectra = target.spectra
nleg = target.nleg
bands = None
legendre = target.legendre(nleg=nleg, use_gpu=False) #Get previously calculated legendre

nleg = legendre[list(legendre.keys())[0]].shape[0]
iBest = np.argsort(zzchi2)[0:n_nearest]
tdata = dict()
if (binned is not None):
if (use_gpu):
binned = { hs:binned[hs][:,iBest].get() for hs in binned }
Expand All @@ -216,17 +218,23 @@ def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest,
#Only multiply if trans[hs] is not None
#Both arrays are on CPU so no need to wrap with asarray
binned[hs] *= trans[hs][:,None]
for hs, w in dwave.items():
tdata[hs] = binned[hs][None,:,:]
if (nleg > 0):
tdata[hs] = np.append(tdata[hs], legendre[hs].transpose()[None,:,:], axis=2)
nbasis = tdata[hs].shape[2]

if per_camera:
#Use CPU mode since small tdata
(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, tdata, weights, flux, wflux, nleg, 1, method='bvls', n_nbh=n_nearest, prior=prior, use_gpu=False, ncam=ncam)
#CW - 4/19/24 - pass use_gpu and solver method to per_camera_coeff_with_least_square_batch
#it will decide there if narch == 1 to use CPU and it will calculate correct prior array
#CW - 5/2/24 - pass binned instead of tdata to put all BVLS/NNLS code into per_camera_coeff_with_least_square_batch
(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, binned, weights, flux, wflux, nleg, 1, method=self._solver_method, n_nbh=n_nearest, prior_sigma=prior_sigma, use_gpu=use_gpu, bands=bands)
#(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, binned, weights, flux, wflux, nleg, 1, method=self._solver_method, n_nbh=n_nearest, prior_sigma=prior_sigma, use_gpu=use_gpu, ncam=ncam)
else:
#Use CPU mode for calc_zchi2 since small tdata
#Calculate tdata here because binned is passed to per_camera_coeff_with_least_square_batch
tdata = dict()
for hs, w in dwave.items():
tdata[hs] = binned[hs][None,:,:]
if (nleg > 0):
tdata[hs] = np.append(tdata[hs], legendre[hs].transpose()[None,:,:], axis=2)
nbasis = tdata[hs].shape[2]
(zzchi2, zzcoeff) = calc_zchi2_batch(spectra, tdata, weights, flux, wflux, 1, nbasis, use_gpu=False)

sstype = ['%s'%(self._subtype[k]) for k in iBest] # subtypes of best archetypes
Expand All @@ -235,7 +243,7 @@ def nearest_neighbour_model(self, target,weights,flux,wflux,dwave,z, n_nearest,
#print(z, zzchi2, zzcoeff, fsstype)
return zzchi2[0], zzcoeff[0], self._rrtype+':::%s'%(fsstype)

def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nearest, trans=None, solve_method='bvls', prior=None, use_gpu=False):
def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nearest, trans=None, prior_sigma=None, use_gpu=False):

"""Get the best archetype for the given redshift and spectype.

Expand All @@ -249,9 +257,8 @@ def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nea
per_camera (bool): True if fitting needs to be done in each camera
n_nearest (int): number of nearest neighbours to be used in chi2 space (including best archetype)
trans (dict): pass previously calcualated Lyman transmission instead of recalculating
solve_method (string): bvls or pca
use_gpu (bool): use GPU or not
prior (2d array): prior matrix on coefficients (1/sig**2)
prior_sigma (float): prior to add in the final solution matrix: added as 1/(prior_sigma**2) only for per-camera mode

Returns:
chi2 (float): chi2 of best archetype
Expand All @@ -261,7 +268,9 @@ def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nea
"""
spectra = target.spectra
nleg = target.nleg
bands = None
legendre = target.legendre(nleg=nleg, use_gpu=use_gpu) #Get previously calculated legendre
solve_method = self._solver_method #get solve method from archetype class instead of passing as arg

#Select np or cp for operations as arrtype
if (use_gpu):
Expand All @@ -287,9 +296,7 @@ def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nea

obs_wave = np.concatenate([dwave[key] for key in new_keys])

nleg = legendre[list(legendre.keys())[0]].shape[0]
zzchi2 = np.zeros(self._narch, dtype=np.float64)
zzcoeff = np.zeros((self._narch, 1+ncam*(nleg)), dtype=np.float64)
#nleg = legendre[list(legendre.keys())[0]].shape[0]

#TODO: return best fit model as well
#zzmodel = np.zeros((self._narch, obs_wave.size), dtype=np.float64)
Expand All @@ -306,40 +313,47 @@ def get_best_archetype(self,target,weights,flux,wflux,dwave,z, per_camera, n_nea
#Rebin in batch
binned = self.rebin_template_batch(z, dwave, trapz=True, dedges=dedges, use_gpu=use_gpu)

tdata = dict()
nbasis = 1

## Prior must be redefined to remove nearest neighbour approach,
# because prior was defined based on n_nearest argument..
# this logic is needed because the first fitting is done with just one archetype
#and then nearest neighbour approach is implemented
if n_nearest is not None:
nnearest_prior = prior.copy() # prior corresponding to nearest_nbh method
prior = prior[n_nearest-1:,][:,n_nearest-1:] # removing first rows/columns corresponding to the nearest_archetypes, and keeping just one row for one archetype

##CW 04/19/24 - no need to redefine prior since it is now not calculated until per_camera_coeff_with_least_square_batch
#if n_nearest is not None:
# nnearest_prior = prior.copy() # prior corresponding to nearest_nbh method
# prior = prior[n_nearest-1:,][:,n_nearest-1:] # removing first rows/columns corresponding to the nearest_archetypes, and keeping just one row for one archetype

for hs, wave in dwave.items():
if (trans[hs] is not None):
#Only multiply if trans[hs] is not None
binned[hs] *= arrtype.asarray(trans[hs][:,None])
#Create 3-d tdata with narch x nwave x nbasis where nbasis = 1+nleg
if nleg > 0:
tdata[hs] = arrtype.append(binned[hs].transpose()[:,:,None], arrtype.tile(arrtype.asarray(legendre[hs]).transpose()[None,:,:], (self._narch, 1, 1)), axis=2)
else:
tdata[hs] = binned[hs].transpose()[:,:,None]
nbasis = tdata[hs].shape[2]

if per_camera:
#Use per_camera_coeff_with_least_square_batch which has all logic associated with BVLS/NNLS solver methods
if (use_gpu):
(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, tdata, gpuweights, gpuflux, gpuwflux, nleg, self._narch, method=solve_method, n_nbh=1, prior=prior, use_gpu=use_gpu, ncam=ncam)
(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, binned, gpuweights, gpuflux, gpuwflux, nleg, self._narch, method=solve_method, n_nbh=1, prior_sigma=prior_sigma, use_gpu=use_gpu, bands=bands)
#(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, binned, gpuweights, gpuflux, gpuwflux, nleg, self._narch, method=solve_method, n_nbh=1, prior_sigma=prior_sigma, use_gpu=use_gpu, ncam=ncam)
else:
(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, tdata, weights, flux, wflux, nleg, self._narch, method=solve_method, n_nbh=1, prior=prior, use_gpu=use_gpu, ncam=ncam)
(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, binned, weights, flux, wflux, nleg, self._narch, method=solve_method, n_nbh=1, prior_sigma=prior_sigma, use_gpu=use_gpu, bands=bands)
#(zzchi2, zzcoeff) = per_camera_coeff_with_least_square_batch(target, binned, weights, flux, wflux, nleg, self._narch, method=solve_method, n_nbh=1, prior_sigma=prior_sigma, use_gpu=use_gpu, ncam=ncam)
else:
#Calculate tdata here because binned is passed to per_camera_coeff_with_least_square_batch
tdata = dict()
nbasis = 1
for hs, wave in dwave.items():
#Create 3-d tdata with narch x nwave x nbasis where nbasis = 1+nleg
if nleg > 0:
tdata[hs] = arrtype.append(binned[hs].transpose()[:,:,None], arrtype.tile(arrtype.asarray(legendre[hs]).transpose()[None,:,:], (self._narch, 1, 1)), axis=2)
else:
tdata[hs] = binned[hs].transpose()[:,:,None]
nbasis = tdata[hs].shape[2]
if (use_gpu):
(zzchi2, zzcoeff) = calc_zchi2_batch(spectra, tdata, gpuweights, gpuflux, gpuwflux, self._narch, nbasis, use_gpu=use_gpu)
else:
(zzchi2, zzcoeff) = calc_zchi2_batch(spectra, tdata, weights, flux, wflux, self._narch, nbasis, use_gpu=use_gpu)

if n_nearest is not None:
best_chi2, best_coeff, best_fulltype = self.nearest_neighbour_model(target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=dedges, binned=binned, use_gpu=use_gpu, prior=nnearest_prior, ncam=ncam)
best_chi2, best_coeff, best_fulltype = self.nearest_neighbour_model(target,weights,flux,wflux,dwave,z, n_nearest, zzchi2, trans, per_camera, dedges=dedges, binned=binned, use_gpu=use_gpu, prior_sigma=prior_sigma, ncam=ncam)
#print(best_chi2, best_coeff, best_fulltype)
return best_chi2, best_coeff, best_fulltype
else:
Expand Down
31 changes: 5 additions & 26 deletions py/redrock/fitz.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,26 +116,6 @@ def legendre_calculate(nleg, dwave):

return legendre

def prior_on_coeffs(n_nbh, deg_legendre, sigma, ncamera):

"""
Args:
n_nbh (int): number of dominant archetypes
deg_legendre (int): number of Legendre polynomials
sigma (int): prior sigma to be used for archetype fitting
ncamera (int): number of cameras for given instrument
Returns:
2d array to be added while solving for archetype fitting

"""

nbasis = n_nbh+deg_legendre*ncamera # 3 desi cameras
prior = np.zeros((nbasis, nbasis), dtype='float64');np.fill_diagonal(prior, 1/(sigma**2))
for i in range(n_nbh):
prior[i][i]=0. ## Do not add prior to the archetypes, added only to the Legendre polynomials
return prior


def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=False, deg_legendre=None, zminfit_npoints=15, per_camera=False, n_nearest=None, prior_sigma=None):
"""Refines redshift measurement around up to nminima minima.

Expand Down Expand Up @@ -293,6 +273,8 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
coeff = np.zeros(template.nbasis)
zwarn |= ZW.Z_FITLIMIT
zwarn |= ZW.BAD_MINFIT
#Set trans to None so as not to pass an empty dict to archetypes! CW 2/26/24
trans = None
else:
#- Unknown problem; re-raise error
raise err
Expand Down Expand Up @@ -331,12 +313,9 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
else:
ncamera = 1
if n_nearest is None:
prior = prior_on_coeffs(1, deg_legendre, prior_sigma, ncamera)
else:
prior = prior_on_coeffs(n_nearest, deg_legendre, prior_sigma, ncamera)
else:
prior=None
chi2min, coeff, fulltype = archetype.get_best_archetype(target,weights,flux,wflux,dwave,zbest, per_camera, n_nearest, trans=trans, use_gpu=use_gpu, prior=prior)
n_nearest = 1
#Pass prior_sigma as a scalar and move calculation of prior array to zscan.py in per_camera_coeff_with_least_square_batch
chi2min, coeff, fulltype = archetype.get_best_archetype(target,weights,flux,wflux,dwave,zbest, per_camera, n_nearest, trans=trans, use_gpu=use_gpu, prior_sigma=prior_sigma)
del trans

results.append(dict(z=zbest, zerr=zerr, zwarn=zwarn,
Expand Down
Loading