Full redrock API Reference¶

This is used to include docstrings from modules. See the autodoc documentation.

If you’re loading a module here, and don’t see some functions, try adding the :imported-members: option.

redrock¶

Redrock redshift fitter.

redrock.archetypes¶

Classes and functions for archetypes.

class redrock.archetypes.All_archetypes(lstfilename=None, archetypes_dir=None)[source]¶

Class to store all different archetypes of all the different spectype.

Parameters:

lstfilename (lst str) – List of file to get the templates from
archetypes_dir (str) – Directory to the archetypes

class redrock.archetypes.Archetype(filename)[source]¶

Class to store all different archetypes from the same spectype.

The archetype data are read from a redrock-format archetype file.

Parameters:: filename (str) – the path to the archetype file

eval(subtype, dwave, coeff, wave, z)[source]¶

get_best_archetype(spectra, weights, flux, wflux, dwave, z, legendre)[source]¶

Get the best archetype for the given redshift and spectype.

Parameters:

spectra (list) – list of Spectrum objects.
weights (array) – concatenated spectral weights (ivar).
flux (array) – concatenated flux values.
wflux (array) – concatenated weighted flux values.
dwave (dic) – dictionary of wavelength grids
z (float) – best redshift
legendre (dic) – legendre polynomial

Returns:

chi2 of best archetype zcoef (array): zcoef of best archetype fulltype (str): fulltype of best archetype

Return type:

chi2 (float)

rebin_template(index, z, dwave, trapz=True)[source]¶

redrock.archetypes.find_archetypes(archetypes_dir=None)[source]¶

Return list of rrarchetype-*.fits archetype files

Search directories in this order, returning results from first one found:

archetypes_dir
$RR_ARCHETYPE_DIR
<redrock_code>/archetypes/

Parameters:: archetypes_dir (str) – optional directory containing the archetypes.
Returns:: a list of archetype files.
Return type:: list

redrock.constants¶

Set constants used by the rest of the package.

redrock.external¶

This provides interfaces to actual data files.

redrock.external.boss¶

redrock wrapper tools for BOSS

redrock.external.boss.read_spectra(spplates_name, targetids=None, use_frames=False, fiberid=None, coadd=False, cache_Rcsr=False, use_andmask=False)[source]¶

Read targets from a list of spectra files

Parameters:

spplates_name (list or str) – input spPlate files or pattern to match files.
targetids (list) – restrict targets to this subset.
use_frames (bool) – if True, use frames.
fiberid (int) – Use this fiber ID.
coadd (bool) – if True, compute and use the coadds.
cache_Rcsr (bool) – pre-calculate and cache sparse CSR format of resolution matrix R
use_andmask (bool) – sets ivar = 0 to pixels with and_mask != 0

Returns:

(targets, meta) where targets is a list of Target objects and meta is a Table of metadata (currently only BRICKNAME).

Return type:

tuple

redrock.external.boss.rrboss(options=None, comm=None)[source]¶

Estimate redshifts for BOSS targets.

This loads targets serially and copies them into a DistTargets class. It then runs redshift fitting and writes the output to a catalog.

Parameters:

options (list) – optional list of commandline options to parse.
comm (mpi4py.Comm) – MPI communicator to use.

redrock.external.boss.write_zbest(outfile, zbest, template_version, archetype_version)[source]¶

Write zbest Table to outfile

Parameters:

outfile (str) – output file.
zbest (Table) – the output best fit results.

redrock.external.desi¶

redrock wrapper tools for DESI

class redrock.external.desi.DistTargetsDESI(spectrafiles, coadd=True, targetids=None, first_target=None, n_target=None, comm=None, cache_Rcsr=False, cosmics_nsig=0, capacities=None)[source]¶

Distributed targets for DESI.

DESI spectral data is grouped by sky location, but is just a random collection of spectra for all targets. Read this into memory while grouping by target ID, preserving order in which each target first appears.

We pass through the spectra files once to compute all the book-keeping associated with regrouping the spectra by target. Then we pass through again and actually read and distribute the data.

Parameters:

spectrafiles (str or list) – a list of input files or pattern match of files.
coadd (bool) – if False, do not compute the coadds.
targetids (list) – (optional) restrict the global set of target IDs to this list.
first_target (int) – (optional) integer offset of the first target to consider in each file. Useful for debugging / testing.
n_target (int) – (optional) number of targets to consider in each file. Useful for debugging / testing.
comm (mpi4py.MPI.Comm) – (optional) the MPI communicator.
cache_Rcsr – pre-calculate and cache sparse CSR format of resolution matrix R
cosmics_nsig (float) – cosmic rejection threshold used in coaddition
capacities (list) – (optional) list of process capacities. If None, use equal capacity per process. A process with higher capacity can handle more work.

redrock.external.desi.rrdesi(options=None, comm=None)[source]¶

Estimate redshifts for DESI targets.

This loads distributed DESI targets from one or more spectra grouping files and computes the redshifts. The outputs are written to a redrock scan file and a DESI redshift catalog.

Parameters:

options (list) – optional list of commandline options to parse.
comm (mpi4py.Comm) – MPI communicator to use.

redrock.external.desi.write_zbest(outfile, zbest, fibermap, exp_fibermap, tsnr2, template_version, archetype_version, spec_header=None)[source]¶

Write zbest and fibermap Tables to outfile

Parameters:

outfile (str) – output path.
zbest (Table) – best fit table.
fibermap (Table) – the coadded fibermap from the original inputs.
tsnr2 (Table) – table of input coadded TSNR2 values
exp_fibermap (Table) – the per-exposure fibermap from the orig inputs.
template_version (str) – template version used
archetype_version (str) – archetype version used

Options:: spec_header (dict-like): header of HDU 0 of input spectra

Modifies input tables.meta[‘EXTNAME’]

redrock.fitz¶

Functions for fitting minima of chi^2 results.

redrock.fitz.find_minima(x)[source]¶

Return indices of local minima of x, including edges.

The indices are sorted small to large.

Note

this is somewhat conservative in the case of repeated values: find_minima([1,1,1,2,2,2]) -> [0,1,2,4,5]

Parameters:: x (array-like) – The data array.
Returns:: The indices.
Return type:: (array)

redrock.fitz.fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=False)[source]¶

Refines redshift measurement around up to nminima minima.

Parameters:

zchi2 (array) – chi^2 values for each redshift.
redshifts (array) – the redshift values.
target (Target) – the target for this fit which includes a list of Spectrum objects at different wavelength grids.
template (Template) – the template for this fit.
nminima (int) – the number of minima to consider.
use_gpu (bool) – use GPU or not

Returns:

the fit parameters for the minima.

Return type:

Table

redrock.fitz.get_dv(z, zref)[source]¶

Returns velocity difference in km/s for two redshifts

Parameters:

z (float) – redshift for comparison.
zref (float) – reference redshift.

Returns:

the velocity difference.

Return type:

(float)

redrock.fitz.minfit(x, y)[source]¶

Fits y = y0 + ((x-x0)/xerr)**2

See redrock.zwarning.ZWarningMask.BAD_MINFIT for zwarn failure flags

Parameters:

x (array) – x values.
y (array) – y values.

Returns:

(x0, xerr, y0, zwarn) where zwarn=0 is good fit.

Return type:

(tuple)

redrock.plotspec¶

Visualization tools for plotting spectra.

redrock.priors¶

Classes and functions for priors.

class redrock.priors.Priors(filename)[source]¶

Class to store all different redshift priors.

Parameters:: filename (str) – the path to the redshift prior file.

Note

The file should have at least one HDU with EXTNAME=’PRIORS’, composed of:

TARGETID: the id of the object
Z: the mean of the prior
SIGMA: the sigma (in dz) for the prior

eval(targetid, z)[source]¶

Return prior contribution to the chi2 for the given TARGETID and redshift grid :param targetid: TARGETID of the object :param z: redshift at which to evaluate template flux

Returns:: prior values on the redshift grid

static gaussian(z, z0, s0)[source]¶

Return a Gaussian prior of mean z0 and sigma s0 on the grid z :param z: redshift grid :param z0: mean :param s0: sigma

Returns:: prior values on the redshift grid

static lorentzien(z, z0, s0)[source]¶

Return a Lorentzien prior of mean z0 and sigma s0 on the grid z :param z: redshift grid :param z0: mean :param s0: sigma

Returns:: prior values on the redshift grid

static tophat(z, z0, s0)[source]¶

Return a tophat prior of mean z0 and width s0 on the grid z.

Parameters:

z – redshift grid
z0 – mean
s0 – width

Returns:

prior values on the redshift grid.

Warning

np.NaN <= np.NaN -> False
np.NaN <=/>= 0.0 -> False
np.inf >= 0.0 -> True
np.inf >= np.inf -> True

redrock.rebin¶

Tools for binning data.

redrock.rebin._trapz_rebin_1d(x, y, edges, results)¶

Numba-friendly version of trapezoidal rebinning

See redrock.rebin.trapz_rebin() for input descriptions. results is pre-allocated array of length len(edges)-1 to keep results

redrock.rebin._trapz_rebin_batch(x, y, edges, myz, results, redshifted_x)¶

Numba-friendly version of trapezoidal rebinning for multiple bases and/or redshifts. This is a wrapper to call _trapz_rebin_1d multiple times.

See redrock.rebin.trapz_rebin() for input descriptions. results is pre-allocated array of shape (nz, len(edges)-1, nbasis) to keep results.

redrock.rebin._trapz_rebin_batch_gpu(x, y, edges, myz, result_shape)[source]¶

Rebin y(x) flux density using trapezoidal integration between bin edges GPU algorithm can rebin in batch for multiple redshifts and bases and returns 1-d, 2-d, or 3-d array (n redshifts x n bins x n basis) where n basis is the optional second dimension of the y input array and n redshifts is the length of the optional myz array.

Parameters:

x (1-d array) – input x values.
y (1-d or 2-d array) – input y values (for all bases).
edges (1-d array) – new bin edges.
myz (array) – redshift array to rebin in batch, applying redshifts on-the-fly to x
result_shape (tuple) – output shape of results array - this is required because for instance if an input scalar redshift or no redshift was given to trapz_rebin, then the result shape should be (nbin, nbasis) whereas if redshift array of length 1 was given as input it should be (1, nbin, nbasis). Similarly if the y array given to trapz_rebin is 1-d, the basis dimension is omitted but if it is 2-d with (1, n) shape, it will explicitly be 1.

Returns:

integrated results with shape (nz, nbin, nbasis): where nbin = len(results) = len(edges)-1 if no input redshift is given to trapz_rebin or the redshift is a scalar, the nz dimension will be omitted. If the input y given to trapz_rebin is 1-d then the nbasis dimension will be omitted. e.g. for 1-d input y and nz = 100, the shape will be (100, nbin) while for y with shape (1, n) the result will be (100, nbin, 1). For 1-d input y and scalar or omitted redshift, the shape will be (nbin).

Return type:

cp.array

redrock.rebin.centers2edges(centers)[source]¶

Convert bin centers to bin edges, guessing at what you probably meant

Parameters:: centers (array) – bin centers,
Returns:: bin edges, lenth = len(centers) + 1
Return type:: array

redrock.rebin.rebin_template(template, myz, dwave=None, dedges=None, use_gpu=False, xmin=None, xmax=None)[source]¶

Rebin a template to a set of wavelengths.

Given a template and a single redshift - or an array of redshifts, rebin the template to a set of wavelength arrays. The wavelengths can be passed as centers (dwave) or edges (dedges) of the wavelength bins.

Parameters:

template (Template) – the template object
myz (float or array of float) – the redshift(s)
dwave (dict) – the keys are the “wavehash” and the values are a 1D array containing the wavelength grid.
dedges (dict) – the keys are the “wavehash” and the values can be either a 1D array containing the bin edges of the wavelength grid (computed from centers2edges) or a 3-element tuple with this 1D array as a CuPy array and the min and max values as scalars in CPU memory.
use_gpu (bool) – whether or not to use the GPU algorithm
xmin (float) – (optional) minimum x-value - x[0]*(1+myz.max()) will be used if omitted - this is useful to avoid GPU to CPU copying in the case where x is a CuPy array and providing the scalar value results in a large speed gain
xmax (float) (optional) maximum x-value - x[-1]*(1+myz.min()) – used if omitted - this is useful to avoid GPU to CPU copying in the case where x is a CuPy array and providing the scalar value results in a large speed gain

Returns:

The rebinned template for every basis function and wavelength: grid in dwave. This is a dict of np or cp 3d arrays (nz x nlambda x nbasis)

Return type:

dict

redrock.rebin.trapz_rebin(x, y, xnew=None, edges=None, myz=None, use_gpu=False, xmin=None, xmax=None, edge_min=None, edge_max=None)[source]¶

Rebin y(x) flux density using trapezoidal integration between bin edges Optionally use GPU helper method to rebin in batch, see trapz_rebin_batch_gpu Note - current return array shape is (nz, nbins, nbasis). Changing to (nz, nbasis, nbins) would intuitively make sense but the former shape is needed by zscan. Flagging this for possible changes down the road.

Notes

y is interpreted as a density, as is the output, e.g.

>>> x = np.arange(10)
>>> y = np.ones(10)
>>> trapz_rebin(x, y, edges=[0,2,4,6,8])  #- density still 1, not 2
array([ 1.,  1.,  1.,  1.])
>>> y = np.ones((2,10)) #nbasis = 2
>>> y[1,:] = np.arange(10)
>>> trapz_rebin(x, y, edges=[0,2,4,6,8], use_gpu=True) #nbasis=2, GPU mode
array([[1., 1.],
       [1., 3.],
       [1., 5.],
       [1., 7.]])
>>> trapz_rebin(x, y, edges=[0,2,4,6,8], myz=[0,0.5],use_gpu=True) #nbasis=2, multiple redshifts, GPU mode
array([[[1.        , 1.        ],
        [1.        , 3.        ],
        [1.        , 5.        ],
        [1.        , 7.        ]],
       [[1.        , 0.66666667],
        [1.        , 2.        ],
        [1.        , 3.33333333],
        [1.        , 4.66666667]]])
>>> trapz_rebin(x, y, edges=[0,2,4,6,8], myz=[0.0,0.5], use_gpu=False) #nbasis=2, CPU mode
array([[[1.        , 1.        ],
        [1.        , 3.        ],
        [1.        , 5.        ],
        [1.        , 7.        ]],
       [[1.        , 0.66666667],
        [1.        , 2.        ],
        [1.        , 3.33333333],
        [1.        , 4.66666667]]])

Parameters:

x (array) – input x values.
y (1-d or 2-d array) – input y values (batch mode allows 2-d array with multiple bases).
edges (array) – (optional) new bin edges.
myz (array) – (optional) redshift array to rebin in batch, applying redshifts on-the-fly to x
use_gpu (boolean) – whether or not to use GPU algorithm
xmin (float) – (optional) minimum x-value - x[0]*(1+myz.max()) will be used if omitted - this is useful to avoid GPU to CPU copying in the case where x is a CuPy array and providing the scalar value results in a large speed gain
xmax (float) (optional) maximum x-value - x[-1]*(1+myz.min()) – used if omitted - this is useful to avoid GPU to CPU copying in the case where x is a CuPy array and providing the scalar value results in a large speed gain
edge_min (float) – (optional) minimum new bin edge - edges[0] will be used if omitted - this is useful to avoid GPU to CPU copying in the case where edges is precomputed on the GPU as a CuPy array and providing the scalar value results in a large speed gain
edge_max (float) – (optional) maximum new bin edge - edges[-1] will be used if omitted - this is useful to avoid GPU to CPU copying in the case where edges is precomputed on the GPU as a CuPy array and providing the scalar value results in a large speed gain

Returns:

integrated results with len(results) = len(edges)-1: In batch mode, returns np.array (use_gpu=False) or cp.array (use_gpu=True) with shape (nz, nbin, nbasis) where nbin = len(results) = len(edges)-1 if myz is None or myz is a scalar, the nz dimension will be omitted. If y is 1-d then the nbasis dimension will be omitted. e.g. for 1-d y and nz = 100, the shape will be (100, nbin) while for y with shape (1, n) the result will be (100, nbin, 1). For 1-d input y and scalar or omitted myz, the shape will be (nbin).

Return type:

array

Raises:

ValueError – if edges are outside the range of x or if len(x) != len(y)

redrock.results¶

Functions for reading and writing full redrock results to HDF5.

redrock.results.read_zscan(filename)[source]¶

Read redrock.zfind results from a file.

Returns:

(zbest, results) where zbest is a Table with keys TARGETID, Z,

ZERR, ZWARN and results is a nested dictionary results[targetid][templatetype] with keys:

z: array of redshifts scanned

zchi2: array of chi2 fit at each z

penalty: array of chi2 penalties for unphysical fits at each z

zbest: best fit redshift (finer resolution fit around zchi2
min)

minchi2: chi2 at zbest

zerr: uncertainty on zbest

zwarn: 0=good, non-0 is a warning flag

Return type:

tuple

redrock.results.read_zscan_redrock(filename)[source]¶

Read redrock.zfind results from a file to be reused by redrock itself.

Returns:

dictionary of results for each local target ID.: dic.keys() are TARGETID dic[tg].keys() are TEMPLATE dic[tg][ft].keys() are [‘penalty’, ‘zcoeff’, ‘zchi2’, ‘redshifts’]

Return type:

dict

redrock.results.write_zscan(filename, zscan, zfit, clobber=False)[source]¶

Writes redrock.zfind results to a file.

The nested dictionary structure of results is mapped into a nested group structure of the HDF5 file:

/targetids[nt] /zscan/{spectype}/redshifts[nz] /zscan/{spectype}/zchi2[nt, nz] /zscan/{spectype}/penalty[nt, nz] /zscan/{spectype}/zcoeff[nt, nz, nc] or zcoeff[nt, nc, nz] ? /zfit/{targetid}/zfit table…

Parameters:

filename (str) – the output file path.
zscan (dict) – the full set of fit results.
zfit (Table) – the best fit redshift results.
clobber (bool) – if True, delete the file if it exists.

redrock.targets¶

Classes and functions for targets and their spectra.

class redrock.targets.DistTargets(targetids, comm=None)[source]¶

Base class for distributed targets.

Target objects are distributed across the processes in an MPI communicator, but the details of how this data is loaded from disk is specific to a given project. Each project should inherit from this base class and create an appropriate class for the data files being used.

This class defines some general methods and the API that should be followed by these derived classes.

Parameters:

targetids (list) – the global set of target IDs.
comm (mpi4py.MPI.Comm) – (optional) the MPI communicator.

local()[source]¶: Return the local list of Target objects.

local_target_ids()[source]¶: Return the local list of target IDs.

wavegrids()[source]¶: Return the global dictionary of wavelength grids for each wave hash.

class redrock.targets.DistTargetsCopy(targets, comm=None, root=0)[source]¶

Distributed targets built from a copy.

This class is a simple wrapper that distributes targets located on one process to the processes in a communicator.

Parameters:

targets (list) – list of Target objects on one process.
comm (mpi4py.MPI.Comm) – (optional) the MPI communicator.
root (int) – the process which has the input targets locally.

class redrock.targets.Spectrum(wave, flux, ivar, R, Rcsr=None)[source]¶

Simple container class for an individual spectrum.

Parameters:

wave (array) – the wavelength grid.
flux (array) – the flux values.
ivar (array) – the inverse variance.
R (scipy.sparse.dia_matrix) – the resolution matrix in band diagonal format.
Rcsr (scipy.sparse.csr_matrix) – the resolution matrix in CSR format.

sharedmem_pack()[source]¶: Pack spectral data into multiprocessing shared memory.

sharedmem_unpack()[source]¶: Unpack spectral data from multiprocessing shared memory.

class redrock.targets.Target(targetid, spectra, coadd=False, cosmics_nsig=0.0, meta=None)[source]¶

A single target.

This represents the data for a single target, including a unique identifier and the individual spectra observed for this object (or a coadd).

Parameters:

targetid (int or str) – unique targetid
spectra (list) – list of Spectrum objects
coadd (bool) – compute and store the coadd at construction time. The coadd can always be recomputed with the compute_coadd() method.
cosmics_nsig (float) – cosmic rejection threshold in compute_coadd
meta (dict) – optional metadata dictionary for this Target.

compute_coadd(cache_Rcsr=False, cosmics_nsig=0.0)[source]¶

Compute the coadd from the current spectra list.

Parameters:

cache_Rcsr – pre-calculate and cache sparse CSR format of resolution matrix R
cosmics_nsig (float) – number of sigma for cosmic rejection.

This method REPLACES the list of individual spectra with coadds.

sharedmem_pack()[source]¶: Pack all spectra into multiprocessing shared memory.

sharedmem_unpack()[source]¶: Unpack all spectra from multiprocessing shared memory.

redrock.targets.distribute_targets(targets, nproc)[source]¶

Distribute a list of targets among processes.

Given a list of Target objects, compute the load balanced distribution of those targets among a set of processes.

This function is used when one already has a list of Target objects that need to be distributed. This happens, for example, when creating a DistTargetsCopy object from pre-existing Targets, or when using multiprocessing to do operations on the MPI-local list of targets.

Parameters:

targets (list) – list of Target objects.
nproc (int) – number of processes.

Returns:

A list (one element for each process) with each element: being a list of the target IDs assigned to that process.

Return type:

list

redrock.templates¶

Classes and functions for templates.

class redrock.templates.DistTemplate(template, dwave, mp_procs=1, comm=None, use_gpu=False, gpu_mode=False)[source]¶

Distributed template data interpolated to all redshifts.

For a given template, the redshifts are distributed among the processes in the communicator. Then each process will rebin the template to those redshifts for the wavelength grids specified by dwave.

Parameters:

template (Template) – the template to distribute
dwave (dict) – the keys are the “wavehash” and the values are a 1D array containing the wavelength grid.
mp_procs (int) – if not using MPI, restrict the number of multiprocesses to this.
comm (mpi4py.MPI.Comm) – (optional) the MPI communicator.
use_gpu (bool) – (optional) If this process uses GPU
gpu_mode (bool) – (optional) If ANY process uses GPU - the reason we need both is that in GPU mode, the rebinning of all redshifts is done on all GPUs (redundantly but much faster than using MPI allgather) and no redshifts are done on any CPU only proc.

cycle()[source]¶

Pass our piece of data to the next process.

If we have returned to our original data, then return True, otherwise return False.

Parameters:: Nothing –

Returns (bool):: Whether we have finished (True) else False.

class redrock.templates.DistTemplatePiece(index, redshifts, data)[source]¶

One piece of the distributed template data.

This is a simple container for storing interpolated templates for a set of redshift values. It is used for communicating the interpolated templates between processes.

In the MPI case, each process will store at most two of these simultaneously. This is the data that is computed on a single process and passed between processes.

Parameters:

index (int) – the chunk index of this piece- this corresponds to the process rank that originally computed this piece.
redshifts (array) – the redshift range contained in this piece.
data (list) – a list of dictionaries, one for each redshift, and each containing the 2D interpolated template values for all “wavehash” keys.

class redrock.templates.ReDistTemplate(template, dwave, mp_procs=1, comm=None, use_gpu=False, gpu_mode=False)[source]¶

Distributed template data interpolated to all redshifts.

For a given template, the redshifts are distributed among the processes in the communicator. Then each process will rebin the template to those redshifts for the wavelength grids specified by dwave. After rebinning, the full redshift ranges are redistributed to each process in the communicator.

Parameters:

template (Template) – the template to distribute
dwave (dict) – the keys are the “wavehash” and the values are a 1D array containing the wavelength grid.
mp_procs (int) – if not using MPI, restrict the number of multiprocesses to this.
comm (mpi4py.MPI.Comm) – (optional) the MPI communicator.

cycle()[source]¶

This function is a no-op since redshift ranges have been redistributed.

Parameters:: Nothing –

Returns (bool):: Always returns True

class redrock.templates.Template(filename=None, spectype=None, redshifts=None, wave=None, flux=None, subtype=None, zscan_galaxy=None, zscan_qso=None, zscan_star=None)[source]¶

A spectral Template PCA object.

The template data is read from a redrock-format template file. Alternatively, the data can be specified in the constructor.

Parameters:: filename (str) – the path to the template file, either absolute or relative to the RR_TEMPLATE_DIR environment variable.

eval(coeff, wave, z)[source]¶

Return template for given coefficients, wavelengths, and redshift

Parameters:

coeff – array of coefficients length self.nbasis
wave – wavelengths at which to evaluate template flux
z – redshift at which to evaluate template flux

Returns:

template flux array

Notes

A single factor of (1+z)^-1 is applied to the resampled flux to conserve integrated flux after redshifting.

property full_type¶

subtype string.

Type:: Return formatted type

property solve_matrices_algorithm¶: Return a string representing the algorithm to be used in zscan.solve_matrices. Possible values are: PCA NMF Logic can be added here to select a default algorithm based on header keywords etc so that different templates seamlessly use different algorithms.

redrock.templates._mp_rebin_template(template, dwave, zlist, qout, iproc, use_gpu)[source]¶: Function for multiprocessing version of rebinning. With rebinning now done in batch mode, use process index, iproc to keep track of order of redshifts instead of keying dict by individual redshifts.

redrock.templates.eval_model(data, wave, R=None, templates=None)[source]¶

Evaluate model spectra.

Given a bunch of fits with coefficients COEFF, redshifts Z, and types SPECTYPE, SUBTYPE in data, evaluate the redrock model fits at the wavelengths wave using resolution matrix R.

The wavelength and resolution matrices may be dictionaries including for multiple cameras.

Parameters:

data (table-like, [nspec]) – table with information on each model to evaluate. Must contain at least Z, COEFF, SPECTYPE, and SUBTYPE fields.
wave (array [nwave] or dictionary thereof) – array of wavelengths in angstrom at which to evaluate the models.
R (list of [nwave, nwave] arrays of floats or dictionary thereof) – resolution matrices for evaluating spectra.
templates (dictionary of Template) – dictionary with (SPECTYPE, SUBTYPE) giving the template corresponding to each type.

Returns:

model fluxes, array [nspec, nwave]. If wave and R are dict, then a dictionary of model fluxes, one for each camera.

redrock.templates.find_templates(template_dir=None)[source]¶

Return list of redrock-*.fits template files

Search directories in this order, returning results from first one found:

template_dir
$RR_TEMPLATE_DIR
<redrock_code>/templates/

Parameters:: template_dir (str) – optional directory containing the templates.
Returns:: a list of template files.
Return type:: list

redrock.templates.load_dist_templates(dwave, templates=None, comm=None, mp_procs=1, zscan_galaxy=None, zscan_qso=None, zscan_star=None, redistribute=False, use_gpu=False, gpu_mode=False)[source]¶

Read and distribute templates from disk.

This reads one or more template files from disk and distributes them among an MPI communicator. Each process will locally store interpolated data for a redshift slice of each template. For a single redshift, the template is interpolated to the wavelength grids specified by “dwave”.

As an example, imagine 3 templates with independent redshift ranges. Also imagine that the communicator has 2 processes. This function would return a list of 3 DistTemplate objects. Within each of those objects, the 2 processes store the interpolated data for a subset of the redshift range:

DistTemplate #1: zmin1 <—- p0 —-> | <—- p1 —-> zmax1 DistTemplate #2: zmin2 <– p0 –> | <– p1 –> zmax2 DistTemplate #3: zmin3 <— p0 —> | <— p1 —> zmax3

Parameters:

dwave (dict) – the dictionary of wavelength grids. Keys are the “wavehash” and values are an array of wavelengths.
templates (str or None) – if None, find all templates from the redrock template directory. If a path to a file is specified, load that single template. If a path to a directory is given, load all templates in that directory.
comm (mpi4py.MPI.Comm) – (optional) the MPI communicator.
mp_procs (int) – if not using MPI, restrict the number of multiprocesses to this.
redistribute (bool) – (optional) allgather rebinned templates after distributed rebinning so each process has the full redshift range for the template.

Returns:

a list of DistTemplate objects.

Return type:

list

redrock.utils¶

Redrock utility functions.

redrock.utils.distribute_work(nproc, ids, weights=None, capacities=None)[source]¶

Helper function to distribute work among processes with varying capacities.

Parameters:

nproc (int) – the number of processes.
ids (list) – list of work unit IDs
weights (dict) – dictionary of weights for each ID. If None, use equal weighting.
capacities (list) – list of process capacities. If None, use equal capacity per process. A process with higher capacity can handle more work.

Returns:

A list (one element for each process) with each element: being a list of the IDs assigned to that process.

Return type:

list

redrock.utils.elapsed(timer, prefix, comm=None)[source]¶

Get and print the elapsed time.

If timer is None, compute the start time and return. Otherwise, find the elapsed time and print a message before returning the new start time.

Parameters:

timer (float) – time in seconds for some arbitrary epoch. If “None”, get the current time and return.
prefix (str) – string to print before the elapsed time.
comm (mpi4py.MPI.Comm) – optional communicator.

Returns:

the new start time in seconds.

Return type:

float

redrock.utils.encode_column(c)[source]¶

Returns a bytes column encoded into a string column.

Parameters:: c (Table column) – a column of a Table.
Returns:: an array of strings.
Return type:: array

redrock.utils.getGPUCountMPI(comm)[source]¶

Determine the number of GPUs available when running in MPI mode

In MPI mode, cupy.cuda.runtime.getDeviceCount() always returns 1 so this method determines the actual number of GPUs.

Parameters:: comm – the MPI communicator
Returns:: the number of GPUs
Return type:: int

redrock.utils.get_mp(requested)[source]¶

Returns a reasonable number of multiprocessing processes.

This checks whether the requested value makes sense, and also whether we are running on a NERSC login node (and hence would get in trouble trying to use all cores).

Parameters:: requested (int) – the requested number of processes.
Returns:: the number of processes to use.
Return type:: int

redrock.utils.getch()[source]¶: Return a single character from stdin.

redrock.utils.mp_array(original)[source]¶

Allocate a raw shared memory buffer and wrap it in an ndarray.

This allocates a multiprocessing.RawArray and wraps the buffer with an ndarray.

Parameters:

typcode (str) – the type code of the array.
size_or_init – passed to the RawArray constructor.

Returns;: ndarray: the wrapped data.

redrock.utils.native_endian(data)[source]¶

Convert numpy array data to native endianness if needed.

Returns new array if endianness is swapped, otherwise returns input data

By default, FITS data from astropy.io.fits.getdata() are not Intel native endianness and scipy 0.14 sparse matrices have a bug with non-native endian data.

Parameters:

data (array) – input array

Returns:

original array if input in native endianness, otherwise a copy: with the bytes swapped.

Return type:

array

redrock.utils.nersc_login_node()[source]¶: Returns True if we are on a NERSC login node, else False.

redrock.utils.transmission_IGM(zObj, lObs, use_gpu=False)[source]¶

Calculate the transmitted flux fraction from the Lyman series and due to the IGM. This returns the transmitted flux fraction: 1 -> everything is transmitted (medium is transparent) 0 -> nothing is transmitted (medium is opaque)

This method will handle 3 options: 1 -> For GPU mode, zObj is an array of all z for a template and the return value will be a cupy array (nz x nlambda) 2-> In CPU mode, it can auto-detect if zObj is a numpy array and if so, again, all z will be processed as a vector and the return value will be a numpy array (nz x nlambda) 3-> For legacy, it is still supported to pass zObj as a float and in this case, the return value will be a 1-d numpy array (nlambda).

Parameters:

zObj (float or array of float) – Redshift(s) of object
lObs (array of float) – wavelength grid
use_gpu (boolean) – whether to use CUPY

Returns:

transmitted flux fraction (nlambda in case of scalar input; nz x nlambda in case of array input)

Return type:

array of float

redrock.utils.transmission_IGM_old(zObj, lObs, use_gpu=False)[source]¶

Calculate the transmitted flux fraction from the Lyman series This returns the transmitted flux fraction: 1 -> everything is transmitted (medium is transparent) 0 -> nothing is transmitted (medium is opaque)

This method will handle 3 options: 1 -> For GPU mode, zObj is an array of all z for a template and the return value will be a cupy array (nz x nlambda) 2-> In CPU mode, it can auto-detect if zObj is a numpy array and if so, again, all z will be processed as a vector and the return value will be a numpy array (nz x nlambda) 3-> For legacy, it is still supported to pass zObj as a float and in this case, the return value will be a 1-d numpy array (nlambda).

Parameters:

zObj (float or array of float) – Redshift(s) of object
lObs (array of float) – wavelength grid
use_gpu (boolean) – whether to use CUPY

Returns:

transmitted flux fraction (nlambda in case of scalar input; nz x nlambda in case of array input)

Return type:

array of float

redrock.utils.transmission_Lyman(zObj, lObs, use_gpu=False)[source]¶

Calculate the transmitted flux fraction from the Lyman series This returns the transmitted flux fraction: 1 -> everything is transmitted (medium is transparent) 0 -> nothing is transmitted (medium is opaque)

This method will handle 3 options: 1 -> For GPU mode, zObj is an array of all z for a template and the return value will be a cupy array (nz x nlambda) 2-> In CPU mode, it can auto-detect if zObj is a numpy array and if so, again, all z will be processed as a vector and the return value will be a numpy array (nz x nlambda) 3-> For legacy, it is still supported to pass zObj as a float and in this case, the return value will be a 1-d numpy array (nlambda).

Parameters:

zObj (float or array of float) – Redshift(s) of object
lObs (array of float) – wavelength grid
use_gpu (boolean) – whether to use CUPY

Returns:

transmitted flux fraction (nlambda in case of scalar input; nz x nlambda in case of array input)

Return type:

array of float

redrock.zfind¶

Redshift finding algorithms.

redrock.zfind._mp_fitz(chi2, target_data, t, nminima, qout, archetype, use_gpu)[source]¶: Wrapper for multiprocessing version of fitz.

redrock.zfind._rebalance_after_scan(targets, results)[source]¶: Helper for rebalancing targets and results after lopsided zscan

redrock.zfind.calc_deltachi2(chi2, z, zwarn, dvlimit=1000.0)[source]¶

Calculate chi2 differences, excluding candidates with close z or bad fits

Parameters:

chi2 – array of chi2 values
z – array of redshifts
zwarn – array of zwarn values

Options:: dvlimit: exclude candidates that are closer than dvlimit [km/s]
Returns (deltachi2, setzwarn) where deltachi2 is array of chi2 differences: to next best good fit, and setzwarn is boolean array of whether a SMALL_DELTACHI2 zwarn bit should be set.
Note: The final target always has deltachi2=0.0 because we don’t know: what the next chi2 would have been. This can also occur for the last N targets if all N of them are within dvlimit of each other.

redrock.zfind.sort_dict_by_col(d, colname)[source]¶: Sort a dict of np.ndarrays by one key. Replacement for astropy.Table.sort

redrock.zfind.sort_dict_by_cols(d, colnames, sort_first_column_first=True)[source]¶

Sort a dict of np.ndarrays by multiple keys Replacement for astropy.Table.sort

Parameters:

d – a dictionary
colnames – a tuple or list of column names
sort_first_column_first – boolean - np.lexsort((a,b)) will sort by b first and then sort by a. This is the opposite of astropy.Table.sort behavior. Setting this to true will ensure that sort_dict_by_cols(d, (‘a’,’b’)) will result in sorting by a first and then b.

redrock.zfind.sort_zfit(zfit)[source]¶

Sorts zfit table by goodness of fit, using ‘zwarn’ and ‘chi2’ columns

Parameters:: zfit – astropy Table with columns ‘zwarn’ and ‘chi2’

Modifies zfit in-place by sorting it

redrock.zfind.sort_zfit_dict(zfit)[source]¶

Sorts zfit dict by goodness of fit, using ‘zwarn’ and ‘chi2’ columns

Parameters:: zfit – dict of numpy arrays with columns ‘zwarn’ and ‘chi2’

Modifies zfit in-place by sorting it

redrock.zfind.zfind(targets, templates, mp_procs=1, nminima=3, archetypes=None, priors=None, chi2_scan=None, use_gpu=False)[source]¶

Compute all redshift fits for the local set of targets and collect.

Given targets and templates distributed across a set of MPI processes, compute the redshift fits for all redshifts and our local set of targets. Each process computes the fits for a slice of redshift range and then cycles through redshift slices by passing the interpolated templates along to the next process in order.

Note

If using MPI, only the rank 0 process will return results- all other processes with return a tuple of (None, None).

Parameters:

targets (DistTargets) – distributed targets.
templates (list) – list of DistTemplate objects.
mp_procs (int, optional) – if not using MPI, this is the number of multiprocessing processes to use.
nminima (int, optional) – number of chi^2 minima to consider. Passed to fitz().
archetypes (str, optional) – file or directory containing archetypes to use for final fitz choice of best chi2 vs. z minimum.
priors (str, optional) – file containing redshift priors
chi2_scan (str, optional) – file containing already computed chi2 scan
use_gpu (bool, optional) – use gpu for calc_zchi2

Returns:

(allresults, allzfit), where “allresults” is a dictionary of the: full chi^2 fit information, suitable for writing to a redrock scan file. “allzfit” is an astropy Table of only the best fit parameters for a limited set of minima.

Return type:

tuple

redrock.zscan¶

Algorithms for scanning redshifts.

redrock.zscan._batch_dot_product_sparse_gpu(spectra, tdata)[source]¶

GPU implementation. Calculate a batch dot product of the 3 sparse matrices in spectra with every template in tdata. A CUDA kernel replicates the functionality of the scipy sparse matrix dot product but done for every template in parallel so that the kernel is only called once per spectrum.

Parameters:

spectra (list) – list of Spectrum objects.
tdata (dict) – dictionary of interpolated template values for each wavehash.

Returns:

dot products of these 3 spectra with all templates

Return type:

Tbs (cp.array)

redrock.zscan._calc_batch_dot_product_3d2d_gpu(Tbs, zc)[source]¶

GPU implementation. Calculate a batch dot product of the 3d array Tbs with the 2d array zc. The 3-d array shape is A x B x C and the 2-d array shape is A x C. The resulting output 2-d array shape is A x B. E.g., for all A a dot product of a 2d array of shape B x C is performed with a 1-d array of shape C. These are non-sparse numpy arrays.

Parameters:

Tbs (array) – the stacked output from all 3 filters from batch_dot_product_sparse, for all redshift templates (nz x nrows x nbasis)
zc (array) – zcoeffs, the 2-d array output of zc = linalg.solve(all_M, all_y) (nz x nbasis)

Returns:

the output of the dot product, (nz x nrows)

Return type:

model (array)

redrock.zscan._mp_calc_zchi2(indx, target_ids, target_data, t, use_gpu, qout, qprog)[source]¶: Wrapper for multiprocessing version of calc_zchi2.

redrock.zscan._zchi2_one(Tb, weights, flux, wflux, zcoeff, solve_matrices_algorithm='PCA')[source]¶

Calculate a single chi2.

For one redshift and a set of spectral data, compute the chi2 for template data that is already on the correct grid.

redrock.zscan.batch_dot_product_sparse(spectra, tdata, nz, use_gpu)[source]¶

Calculate a batch dot product of the 3 sparse matrices in spectra with every template in tdata. Sparse matrix libraries are used to perform the dot products.

Parameters:

spectra (list) – list of Spectrum objects.
tdata (dict) – dictionary of interpolated template values for each wavehash.
nz (int) – number of templates
use_gpu (bool) – use GPU or not

Returns:

dot products of these 3 spectra with all templates

Return type:

Tbs (list)

redrock.zscan.calc_batch_dot_product_3d2d(Tbs, zc, use_gpu)[source]¶

Calculate a batch dot product of the 3d array Tbs with the 2d array zc. The 3-d array shape is A x B x C and the 2-d array shape is A x C. The resulting output 2-d array shape is A x B. E.g., for all A a dot product of a 2d array of shape B x C is performed with a 1-d array of shape C. These are non-sparse numpy arrays.

Parameters:

Tbs (array) – the stacked output from all 3 filters from batch_dot_product_sparse, for all redshift templates (nz x nrows x nbasis)
zc (array) – zcoeffs, the 2-d array output of zc = linalg.solve(all_M, all_y) (nz x nbasis)
use_gpu (bool) – use GPU or not

Returns:

the output of the dot product, (nz x nrows)

Return type:

model (array)

redrock.zscan.calc_batch_dot_product_3d3d_gpu(a, b, transpose_a=False, fullprecision=True)[source]¶

GPU implementation. Calculate a batch dot product of the 3d array a with the 3d array b. The 3-d array shape is A x B x C and the 2-d array shape is A x C. The resulting output 2-d array shape is A x B. E.g., for all A a dot product of a 2d array of shape B x C is performed with a 1-d array of shape C. These are non-sparse numpy arrays.

Parameters:

a (array) – a 3-d array (nz x ncols x nrows) In practice, the Tb array, the stacked output from all 3 filters from batch_dot_product_sparse, for all redshift templates (nz x nrows x nbasis) which should have its transpose performed yielding shape (nz x nbasis x nrows).
b (array) – another 3-d array (nz x nrows x ncols) In practice, the Tb array multiplied y weights, shape (nz x nrows x nbasis)
transpose_a (bool) – Whether or not to transpose the a array before performing the dot product
fullprecision (bool) – Whether or not to ensure reproducibly identical results. The context is that it can be faster to use many parallel threads to compute a single output array element of the dot product result, but due to floating point rounding issues, the random order of the summation can change the result by order 1e-16 (e.g. a+b+c != c+a+b). When set to true, the number of parallel threads is decreased to ensure reproducibility at a trade-off of speed.

Returns:

the output of the dot product, (nz x ncols x ncols): such that M[i,:,:] = a[i,:,:].dot(b[i,:,:])

Return type:

M (array)

redrock.zscan.calc_zchi2(target_ids, target_data, dtemplate, progress=None, use_gpu=False)[source]¶

Calculate chi2 vs. redshift for a given PCA template.

New CPU/GPU algorithms June 2022

Parameters:

target_ids (list) – targets IDs.
target_data (list) – list of Target objects.
dtemplate (DistTemplate) – distributed template data
progress (multiprocessing.Queue) – optional queue for tracking progress, only used if MPI is disabled.

Returns:

(zchi2, zcoeff, zchi2penalty) with:

zchi2[ntargets, nz]: array with one element per target per
redshift
zcoeff[ntargets, nz, ncoeff]: array of best fit template
coefficients for each target at each redshift
zchi2penalty[ntargets, nz]: array of penalty priors per target
and redshift, e.g. to penalize unphysical fits

Return type:

tuple

redrock.zscan.calc_zchi2_batch(spectra, tdata, weights, flux, wflux, nz, nbasis, solve_matrices_algorithm='PCA', use_gpu=False, fullprecision=True)[source]¶

Calculate a batch of chi2. For many redshifts and a set of spectral data, compute the chi2 for template data that is already on the correct grid.

Parameters:

spectra (list) – list of Spectrum objects.
tdata (dict) – dictionary of interpolated template values for each wavehash - arrays are 3d (nz x nlambda x nbasis)
weights (array) – concatenated spectral weights (ivar).
flux (array) – concatenated flux values.
wflux (array) – concatenated weighted flux values.
nz (int) – number of templates
nbasis (int) – nbasis
use_gpu (bool) – use GPU or not
fullprecision (bool) – Whether or not to ensure reproducibly identical results. See calc_batch_dot_product_3d3d_gpu.

Returns:

array with one element per redshift for this target zcoeff (array): array of best fit template coefficients

Return type:

zchi2 (array)

redrock.zscan.calc_zchi2_one(spectra, weights, flux, wflux, tdata)[source]¶

Calculate a single chi2.

For one redshift and a set of spectra, compute the chi2 for template data that is already on the correct grid.

Parameters:

spectra (list) – list of Spectrum objects.
weights (array) – concatenated spectral weights (ivar).
flux (array) – concatenated flux values.
wflux (array) – concatenated weighted flux values.
tdata (dict) – dictionary of interpolated template values for each wavehash.

Returns:

chi^2 and coefficients.

Return type:

tuple

redrock.zscan.calc_zchi2_targets(targets, templates, mp_procs=1, use_gpu=False)[source]¶

Compute all chi2 fits for the local set of targets and collect.

Given targets and templates distributed across a set of MPI processes, compute the coarse-binned chi^2 fit for all redshifts and our local set of targets. Each process computes the fits for a slice of redshift range and then cycles through redshift slices by passing the interpolated templates along to the next process in order.

Parameters:

targets (DistTargets) – distributed targets.
templates (list) – list of DistTemplate objects.
mp_procs (int) – if not using MPI, this is the number of multiprocessing processes to use.
gpu (bool) – (optional) use gpu for calc_zchi2

Returns:

dictionary of results for each local target ID.

Return type:

dict

redrock.zscan.dot_product_sparse_one(spectra, tdata, i)[source]¶

Calculate a dot product of the 3 sparse matrices in spectra with ONE template in tdata. Sparse matrix libraries are used to perform the dot products.

Parameters:

spectra (list) – list of Spectrum objects.
tdata (dict) – dictionary of interpolated template values for each wavehash - arrays are 3d (nz x nlambda x nbasis)
i (int) – index of this redshift

Returns:

dot products of these 3 spectra with ONE templates

Return type:

Tb (array)

redrock.zscan.solve_matrices(M, y, solve_algorithm='PCA', use_gpu=False)[source]¶

Solve the matrix equation y = M*x for the unknown x using the specified algorithm. The default is to use PCA via numpy or cupy linalg.solve. But non-negative matrix factorization (NMF) and other algorithms may be integrated here and selected based on the template.

Parameters:

M (array) – 2d array on CPU; 3d array on GPU for all redshifts
y (array) – 1d array on CPU; 2d array on GPU for all redshifts
solve_algorithm (string) – which algorithm to use
use_gpu (bool) – use GPU or not

Returns:

solution to y = M*x: (1-d array on CPU, 2-d array on GPU)

Return type:

x (array)

Raises:

LinAlgError is passed up if raised by np.linalg.solve –
NotImplementedError if algorithm is not implemented or undefined –

redrock.zscan.spectral_data(spectra)[source]¶

Compute concatenated spectral data products.

This helper function builds full length array quantities needed for the chi2 fit.

Parameters:

spectra (list) – list of Spectrum objects.

Returns:

(weights, flux, wflux) concatenated values used for single: redshift chi^2 fits.

Return type:

tuple

redrock.zwarning¶

Mask bit definitions for zwarning.

WARNING on the warnings: not all of these are implemented yet.