<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 BGS Archetypes
The goal of this notebook is to derive a set of spectral archetypes from the BGS template set using Guangtun Zhu's SetCoverPy algorithm.
Preliminaries.
Step1:  Initialize the random seed so the results are reproducible, below.
Step2:  Output path and filenames.
Step6:  Read the BGS basis templates.
Read both a set of lower-resolution (1 A/pix) templates sampled over a restricted wavelength range (roughly 3500-7000 A) and the same set at higher resolution (0.2 A/pix) and over the wavelength range 0.12-2 micron.  The lower-resolution templates will be used to determine the archetypes (since speed is an issue) while the full-resolution templates is what we actually write out.
In both cases we (arbitrarily) normalize every template to r=18 and adopt a nominal velocity dispersion of 100 km/s.
Step7:  Compute the NxN chi2 matrix.
We use chi2 as the "distance" matrix for the Set Cover problem.  
Then, we need to determine what threshold chi2 value differentiates "different" templates.  
Note that the threshold chi^2 value can be tuned until the desired number of archetypes is achieved.  However, If we want the archetypes to describe each spectrum in the parent sample to a precision of prec=0.1 (10%) then we we should set chi2min to be approximately npix*prec^2.
Step9:  Compute and plot the number of archetypes vs chi2 threshold.
Step11:  Choose a chi2 threshold value then get the final set of archetypes.
Step12:  Sort by Dn(4000).
Step17:  Generate some QAplots.
<ASSISTANT_TASK:>
Python Code:
import os
import numpy as np
import matplotlib.pyplot as plt
from desispec.io.util import write_bintable, makepath
from desisim.io import write_templates
from desisim.archetypes import compute_chi2, ArcheTypes
import multiprocessing
nproc = multiprocessing.cpu_count() // 2 
plt.style.use('seaborn-talk')
%matplotlib inline
Explanation: BGS Archetypes
The goal of this notebook is to derive a set of spectral archetypes from the BGS template set using Guangtun Zhu's SetCoverPy algorithm.
Preliminaries.
End of explanation
seed = 123
rand = np.random.RandomState(seed)
Explanation: Initialize the random seed so the results are reproducible, below.
End of explanation
version = 'v1.0'
outdir = os.path.join(os.getenv('DESI_ROOT'), 'spectro', 'templates', 'archetypes', 'bgs', version)
print('Setting output directory to {}'.format(outdir))
os.makedirs(outdir, exist_ok=True)
chi2file = os.path.join(outdir, 'bgs_archetypes_chi2_{}.fits'.format(version))
archfile = os.path.join(outdir, 'bgs_archetypes_{}.fits'.format(version))
Explanation: Output path and filenames.
End of explanation
def _build_templates(args):
    Filler function for the multiprocessing.
    build_templates(*args)
def build_templates(bgs, input_meta, verbose=False):
    flux, _, meta = bgs.make_templates(input_meta=input_meta, novdisp=True, 
                                          nocolorcuts=True, verbose=verbose)
    return [flux.astype('f4'), meta]
def read_and_normalize(verbose=False, nproc=1, minwave=1200, maxwave=2e4, 
                       cdelt=0.2, nominal_rmag=18.0, nominal_vdisp=1000, 
                       subset=False):
    Read and normalize the full set of basis templates.
    
    
    from astropy.table import vstack
    from desisim.templates import BGS
    from desisim.io import empty_metatable
    bgs = BGS(minwave=minwave, maxwave=maxwave, cdelt=cdelt)
    bgs.normline = None # no emission line
    
    nspec = len(bgs.basemeta)
    if subset:
        nspec = 1000
        these = rand.choice(len(bgs.basemeta), nspec)
        print('Selecting a subset of {} / {} templates!'.format(nspec, len(bgs.basemeta)))
    else:
        these = np.arange(nspec)
    
    input_meta = empty_metatable(nmodel=nspec, objtype='BGS')
    input_meta['TEMPLATEID'] = these
    input_meta['REDSHIFT'] = 0.0
    input_meta['MAG'] = nominal_rmag
    input_meta['VDISP'] = nominal_vdisp
    input_meta['SEED'] = rand.randint(2**32, size=nspec)
    # Not sure why multiprocessing isn't working in this case.
    if nproc > 1:
        chunk = np.array_split(these, nproc)
    
        tempargs = list()
        for ii in range(nproc):
            tempargs.append((bgs, input_meta[chunk[ii]], verbose))
        
        pool = multiprocessing.Pool(nproc)
        out = pool.map(_build_templates, tempargs)
        flux = np.vstack(out[0])
        meta = vstack(out[1])
    else:
        flux, meta = build_templates(bgs, input_meta, verbose)
    nspec, npix = flux.shape
    print('Generated {} rest-frame BGS spectra with {} pixels.'.format(nspec, npix))
    
    return flux, bgs.wave, meta, bgs.basemeta[these]
%time flux, wave, meta, basemeta = read_and_normalize(nproc=1, cdelt=1.0, minwave=3600, maxwave=7000, subset=True)
nspec, npix = flux.shape
%time hiresflux, hireswave, _, _ = read_and_normalize(nproc=1, cdelt=0.2, minwave=1200, maxwave=2e4, subset=True)
_, hiresnpix = hiresflux.shape
def plot_subset(nplot=25, ncol=5):
    Plot a random sampling of the basis templates.
    nspec, npix = flux.shape
    nrow = np.ceil(nplot / ncol).astype('int')
    these = rand.choice(nspec, nplot, replace=False)
    fig, ax = plt.subplots(nrow, ncol, figsize=(2.2*ncol, 2.2*nrow), sharey=True, sharex=True)
    for thisax, indx in zip(ax.flat, these):
        thisax.plot(wave, flux[indx, :])
        thisax.text(0.95, 0.93, '{:0d}'.format(indx), ha='right', 
             va='top', transform=thisax.transAxes, fontsize=11)
        thisax.xaxis.set_major_locator(plt.MaxNLocator(3))
    fig.subplots_adjust(wspace=0.05, hspace=0.05)
plot_subset()
Explanation: Read the BGS basis templates.
Read both a set of lower-resolution (1 A/pix) templates sampled over a restricted wavelength range (roughly 3500-7000 A) and the same set at higher resolution (0.2 A/pix) and over the wavelength range 0.12-2 micron.  The lower-resolution templates will be used to determine the archetypes (since speed is an issue) while the full-resolution templates is what we actually write out.
In both cases we (arbitrarily) normalize every template to r=18 and adopt a nominal velocity dispersion of 100 km/s.
End of explanation
def write_chi2(chi2):
    from astropy.io import fits
    print('Writing {}'.format(chi2file))
    hdu = fits.PrimaryHDU(chi2)
    hdu.writeto(chi2file, overwrite=True)
%time chi2, amp = compute_chi2(flux)
write_chi2(chi2)
prec = 0.1
chi2min_nominal = npix*prec**2
print(chi2min_nominal, np.log10(chi2min_nominal)) # seems high...
with np.errstate(divide='ignore'):
    logchi2 = np.log10(chi2)
logchi2[chi2 == 0] = -1
fig, ax = plt.subplots(1, 2, figsize=(12, 4))
im = ax[0].imshow(logchi2, origin='lower', interpolation='nearest', 
                  vmin=-1.0, cmap='viridis')
ax[0].set_xlabel('Spectrum Number')
ax[0].set_ylabel('Spectrum Number')
plt.colorbar(im, label='$log_{10}(\chi^{2})$', ax=ax[0])
_ = ax[1].hist(logchi2.reshape(nspec * nspec), bins=30, range=(-1.2, np.max(logchi2)))
ax[1].set_ylabel('Number')
ax[1].set_xlabel('$log_{10}(\chi^{2}$)')
Explanation: Compute the NxN chi2 matrix.
We use chi2 as the "distance" matrix for the Set Cover problem.  
Then, we need to determine what threshold chi2 value differentiates "different" templates.  
Note that the threshold chi^2 value can be tuned until the desired number of archetypes is achieved.  However, If we want the archetypes to describe each spectrum in the parent sample to a precision of prec=0.1 (10%) then we we should set chi2min to be approximately npix*prec^2.
End of explanation
def narch_vs_chi2min(Arch):
    Determine the number of archtypes vs chi2 threshold.
    
    
    cost = np.ones(nspec) # uniform cost
    chi2min = np.logspace(1, 5, 10)
    print(chi2min)
    narch = np.zeros_like(chi2min)
    
    for ii, cmin in enumerate(chi2min):
        iarch = Arch.get_archetypes(chi2_thresh=cmin)
        narch[ii] = len(iarch)
        
    return narch, chi2min
def qa_narch_vs_chi2min():
    fig, ax = plt.subplots()
    ax.scatter(np.log10(chi2min), narch)
    ax.set_xlabel('$log_{10}(\chi^{2})$ Threshold')
    ax.set_ylabel('Number of Archetypes')
    ax.axvline(x=np.log10(chi2min_nominal), color='red', ls='-')
    ax.grid(True)
Arch = ArcheTypes(chi2)
narch, chi2min = narch_vs_chi2min(Arch)
qa_narch_vs_chi2min()
Explanation: Compute and plot the number of archetypes vs chi2 threshold.
End of explanation
def write_archetypes():
    ToDo: Write out the responsibility indices for each archetype.
    from astropy.table import Column
    outmeta = meta[iarch]
    outmeta.add_column(Column(name='RESPONSIBILITY', length=len(iarch), dtype='int8'))
    outmeta['RESPONSIBILITY'] = resp
    print('Writing {}'.format(archfile))
    write_templates(archfile, hiresflux[iarch, :], hireswave, outmeta, objtype='BGS Archetypes')    
chi2_thresh = 10**2.5
print('Choosing a log10(chi2) threshold value of {:.1f}.'.format(np.log10(chi2_thresh)))
_iarch, _resp, _respindx = Arch.get_archetypes(chi2_thresh=chi2_thresh, responsibility=True)
print('Generated {} archetypes.'.format(len(_iarch)))
Explanation: Choose a chi2 threshold value then get the final set of archetypes.
End of explanation
srt = np.argsort(meta['D4000'][_iarch])
iarch = _iarch[srt]
resp = _resp[srt]
respindx = []
for ss in srt:
    respindx.append(_respindx[ss])
write_archetypes()
Explanation: Sort by Dn(4000).
End of explanation
def _markers():
    d4000 = meta['D4000']
    size = 110 * (1+(resp - resp.min()) / resp.ptp())
    shade = (d4000[iarch] - d4000[iarch].min()) / d4000[iarch].ptp()
    col = plt.cm.coolwarm(shade)
    return size, col
def qa_responsibility():
    Generate a color-color plot with the symbol size scaled by the responsibility.
    
    
    rz = -2.5 * np.log10(meta['FLUX_R'] / meta['FLUX_Z'])
    d4000 = meta['D4000']
    size, col = _markers()
    
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4), sharey=True)
    ax1.scatter(rz[iarch], resp, c=col, marker='o', s=size, 
                edgecolor='k')
    ax1.set_xlabel('r - z')
    ax1.set_ylabel('Responsibility')
    ax1.grid(True)
    ax2.scatter(d4000[iarch], resp, c=col, marker='o', s=size, 
               edgecolor='k')
    ax2.set_xlabel('$D_{n}(4000)$')
    ax2.grid(True)
    fig.subplots_adjust(wspace=0.05)
qa_responsibility()
def qa_colorcolor():
    Generate a color-color plot with the symbol size scaled by the responsibility.
    
    
    gr = -2.5 * np.log10(meta['FLUX_G'] / meta['FLUX_R'])
    rz = -2.5 * np.log10(meta['FLUX_R'] / meta['FLUX_Z'])
    d4000 = meta['D4000']
    size, col = _markers()
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    ax1.scatter(rz, gr, s=30, c='lightgray', edgecolor='k')
    ax1.scatter(rz[iarch], gr[iarch], c=col, marker='o', s=size, 
                alpha=1, edgecolor='k')
    ax1.set_xlabel(r'$r - z$')
    ax1.set_ylabel(r'$g - r$')
    ax1.grid(True)
    ax2.scatter(d4000, rz, s=30, c='lightgray', edgecolor='k')
    ax2.scatter(d4000[iarch], rz[iarch], c=col, marker='o', s=size, 
                alpha=1, edgecolor='k')
    ax2.set_xlabel('$D_{n}(4000)$')
    ax2.set_ylabel(r'$g - r$')
    ax2.grid(True)
    fig.subplots_adjust(wspace=0.3)
qa_colorcolor()
def qa_archetypes(ncol=5, nfilter=11):
    Plot the archetypes and the spectra for which they're responsible.
    from scipy.signal import medfilt
    
    _, col = _markers()    
    narch = len(iarch)
    nrow = np.ceil(narch / ncol).astype('int')
    fig, ax = plt.subplots(nrow, ncol, figsize=(2.5*ncol, 2.5*nrow), sharey=True, sharex=True)
    
    ww = (hireswave > 3000) * (hireswave < 1e4)
    for jj, (thisax, indx, rindx, rr) in enumerate(zip(ax.flat, iarch, respindx, resp)):
        if rr > 1:
            for ii in rindx:
                thisax.plot(hireswave[ww], hiresflux[ii, ww], color='lightgrey')
        smoothflux = medfilt(hiresflux[indx, ww], 11)
        thisax.plot(hireswave[ww], smoothflux, color=col[jj])
        thisax.xaxis.set_major_locator(plt.MaxNLocator(2))
        thisax.text(0.95, 0.93, '{:04d}\nResp={}'.format(indx, rr), ha='right', 
             va='top', transform=thisax.transAxes, fontsize=11)
    fig.subplots_adjust(wspace=0.05, hspace=0.05)    
qa_archetypes()
def qa_ages_colormag():
    Generate a color-magnitude plot for the original AGES sample.
    
    
    Mr = basemeta['SDSS_UGRIZ_ABSMAG_Z01'][:, 2]
    gr = basemeta['SDSS_UGRIZ_ABSMAG_Z01'][:, 1]-basemeta['SDSS_UGRIZ_ABSMAG_Z01'][:, 2]
    size, col = _markers()
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.scatter(Mr, gr, s=30, c='lightgray', edgecolor='k')
    ax.scatter(Mr[iarch], gr[iarch], c=col, marker='o', s=size, 
                alpha=1, edgecolor='k')
    ax.set_xlabel(r'$M_{0.1r}$')
    ax.set_ylabel(r'$^{0.1}(g - r)$')
    ax.set_xlim(-16, -23)
    ax.set_ylim(0, 1.3)
    
    ax.grid(True)
qa_ages_colormag()
Explanation: Generate some QAplots.
End of explanation
<END_TASK>