<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Include photometry with Grizli grism fits
grizli has the ability to include constraints from broad-band photometry in the redshift fitting.  Using this capability requires some of the tools provided by the eazy-py module (https
Step1:  Pull out the photometry of the nearest source matched in the photometric catalog.
Step2:  Scaling the spectrum to the photometry
For the grism spectra, the total flux of an object is the flux integrated over the segmentation polygon, which corresponds to some isophotal level set by the SExtractor/SEP threshold parameter.   While this should often be similar to the total flux definition in the photometric catalog (e.g., FLUX_AUTO), they won't necessarily be the same and there could be small offsets in the normalization between photometry and spectra.
Below we demonstrate a quick way to scale the spectrum to the photometry for more reliable combined fits.
Step3:  An offset is clearly visible betwen the spectra and photometry.  In the next step, we use an internal function to compute a correction factor to bring the spectrum in line with the photometry.  
The scale_to_photometry function integrates the observed 1D spectrum (combining all available grisms) through the photometry filter bandpasses and derives a polynomial correction to make the two agree.  
NB Because we use the observed spectrum directly, scale_to_photometry requires that at least one observed filter be entirely covered by the grism spectrum.  And fitting for a higher order correction requires multiple filters across the grism bandpass (generally at least one filter per polynomial order).
Step4:  The function computed a multiplicative factor 10/8.9=1.1 to apply to the spectrum, which now falls nicely on top of the photometry.  Next we fit for a first-order linear scaling (normalization and slope).
Step5:  A bit of an improvement.  But be careful with corrections with order&gt;0 if there is limited photometry available in bands that overlap the spectrum.
Step6:  Redshift PDF
Now compare the redshift PDF for the grism+photometry and photometry-only (photo-z) fits.
Step7:  Grizli internal photometry
The archive query functions shown in the Grizli-Pipeline notebook can optionally query for any HST imaging that overlaps with the grism exposures.  The automatic preparation script will process and align these along with the grism+direct data and produce a simple forced-aperture photometric catalog.  Currently this catalog just has matched aperture photometry and does not make any effort at PSF matching.  
For cases where no external photometric catalog is available, this photometry can be used to help constrain the redshift fits.  The catalog and eazy-py products are automatically generated by the preparation scripts.
The grizli.pipeline.photoz scripts can also try to fetch additional photometry from published Vizier catalogs, e.g., GALEX, SDSS, PanSTARRS, WISE, CFHT-LS.
Step8:  Compare chi-squared
In this case of this bright spectrum with strong features, the grism dominates the fit even with the extensive 3D-HST catalog
<ASSISTANT_TASK:>
Python Code:
%matplotlib inline
import os
import numpy as np
import matplotlib.pyplot as plt
import astropy.io.fits as pyfits
import drizzlepac
import grizli
from grizli.pipeline import photoz
from grizli import utils, prep, multifit, fitting
utils.set_warnings()
print('\n Grizli version: ', grizli.__version__)
# Requires eazy-py:  https://github.com/gbrammer/eazy-py
import eazy
# Run in the directory where you ran the Grizli-Pipeline notebook and 
# extracted spectra of two objects
root = 'j0332m2743'
os.chdir('{0}/Extractions/'.format(root))
# Fetch 3D-HST catalogs
if not os.path.exists('goodss_3dhst.v4.1.cats.tar.gz'):
    os.system('wget https://archive.stsci.edu/missions/hlsp/3d-hst/RELEASE_V4.0/Photometry/GOODS-S/goodss_3dhst.v4.1.cats.tar.gz')
    os.system('tar xzvf goodss_3dhst.v4.1.cats.tar.gz')
# Preparation for eazy-py
eazy.symlink_eazy_inputs(path=os.path.dirname(eazy.__file__)+'/data', 
                         path_is_env=False)
### Initialize **eazy.photoz** object
field = 'goodss'
version = 'v4.1'
params = {}
params['CATALOG_FILE'] = '{0}_3dhst.{1}.cats/Catalog/{0}_3dhst.{1}.cat'.format(field, version)
params['Z_STEP'] = 0.002
params['Z_MAX'] = 10
params['MAIN_OUTPUT_FILE'] = '{0}_3dhst.{1}.eazypy'.format(field, version)
params['PRIOR_FILTER'] = 205
# Galactic extinction
params['MW_EBV'] = {'aegis':0.0066, 'cosmos':0.0148, 'goodss':0.0069, 
                    'uds':0.0195, 'goodsn':0.0103}[field]
params['TEMPLATES_FILE'] = 'templates/fsps_full/tweak_fsps_QSF_12_v3.param'
translate_file = '{0}_3dhst.{1}.cats/Eazy/{0}_3dhst.{1}.translate'.format(field, version)
ez = eazy.photoz.PhotoZ(param_file=None, translate_file=translate_file, 
                        zeropoint_file=None, params=params, 
                        load_prior=True, load_products=False)
## Grism fitting arguments created in Grizli-Pipeline
args = np.load('fit_args.npy')[0]
## First-pass redshift templates, similar to the eazy templates but 
## with separate emission lines
t0 = args['t0']
#############
## Make a helper object for generating photometry in a format that grizli 
## understands. 
## Passing the parameters precomputes a function to quickly interpolate
## the templates through the broad-band filters.  It's not required, 
## but makes the fitting much faster.
## 
## `zgrid` defaults to ez.zgrid, be explicit here to show you can 
## change it.  
phot_obj = photoz.EazyPhot(ez, grizli_templates=t0, zgrid=ez.zgrid) 
### Find IDs of specific objects to extract, same ones from the notebook
import astropy.units as u
tab = utils.GTable()
tab['ra'] = [53.0657456, 53.0624459]
tab['dec'] = [-27.720518, -27.707018]
# Internal grizli catalog
gcat = utils.read_catalog('{0}_phot.fits'.format(root))
idx, dr = gcat.match_to_catalog_sky(tab)
source_ids = gcat['number'][idx]
tab['id'] = source_ids
tab['dr'] = dr.to(u.mas)
tab['dr'].format='.1f'
tab.show_in_notebook()
## Find indices in the 3D-HST photometric catalog
idx3, dr3 = ez.cat.match_to_catalog_sky(tab)
## Run the photozs just for comparison.  Not needed for the grism fitting 
## but the photozs and SEDs give you a check that the photometry looks 
## reasonable
ez.param['VERBOSITY'] = 1.
ez.fit_parallel(idx=idx3, verbose=False) 
# or could run on the whole catalog by not specifying `idx`
# Show SEDs with best-fit templates and p(z)
for ix in idx3:
    ez.show_fit(ix, id_is_idx=True)
### Spline templates for dummy grism continuum fits
wspline = np.arange(4200, 2.5e4)
Rspline = 50
df_spl = len(utils.log_zgrid(zr=[wspline[0], wspline[-1]], dz=1./Rspline))
tspline = utils.bspline_templates(wspline, df=df_spl+2, log=True, clip=0.0001)
i=1 # red galaxy
id=tab['id'][i]
ix = idx3[i]
## This isn't necessary for general fitting, but 
## load the grism spectrum here for demonstrating the grism/photometry scaling
beams_file = '{0}_{1:05d}.beams.fits'.format(args['group_name'], id)
mb = multifit.MultiBeam(beams_file, MW_EBV=args['MW_EBV'],
                        fcontam=args['fcontam'], sys_err=args['sys_err'],
                        group_name=args['group_name'])
Explanation: Include photometry with Grizli grism fits
grizli has the ability to include constraints from broad-band photometry in the redshift fitting.  Using this capability requires some of the tools provided by the eazy-py module (https://github.com/gbrammer/eazy-py).
The example below is intended to be run in the same directory where you've already run the Grizli-Pipeline notebook to process the GOODS-South ERS grism data.  
Here we take photometry from the 3D-HST photometric catalog (Skelton et al.) of the GOODS-South field.  The 3D-HST products are available at https://archive.stsci.edu/prepds/3d-hst/.
End of explanation
# Generate the `phot` dictionary
phot, ii, dd = phot_obj.get_phot_dict(mb.ra, mb.dec)
label = "3DHST Catalog ID: {0}, dr={1:.2f}, zphot={2:.3f}"
print(label.format(ez.cat['id'][ii], dd, ez.zbest[ii]))
print('\n`phot` keys:', list(phot.keys()))
for k in phot:
    print('\n'+k+':\n', phot[k])
    
# Initialize photometry for the MultiBeam object
mb.set_photometry(**phot)
# parametric template fit to get reasonable background
sfit = mb.template_at_z(templates=tspline, fit_background=True, 
                        include_photometry=False)
fig = mb.oned_figure(tfit=sfit)
ax = fig.axes[0]
ax.errorbar(mb.photom_pivot/1.e4, mb.photom_flam/1.e-19, 
            mb.photom_eflam/1.e-19, 
            marker='s', color='k', alpha=0.4, linestyle='None',
            label='3D-HST photometry')
ax.legend(loc='upper left', fontsize=8)
Explanation: Pull out the photometry of the nearest source matched in the photometric catalog.
End of explanation
## First example:  no rescaling
z_phot = ez.zbest[ix]
# Reset scale parameter
if hasattr(mb,'pscale'):
    delattr(mb, 'pscale')
    
t1 = args['t1']
tfit = mb.template_at_z(z=z_phot)
print('No rescaling, chi-squared={0:.1f}'.format(tfit['chi2']))
fig = fitting.full_sed_plot(mb, tfit, zfit=None, bin=4)
Explanation: Scaling the spectrum to the photometry
For the grism spectra, the total flux of an object is the flux integrated over the segmentation polygon, which corresponds to some isophotal level set by the SExtractor/SEP threshold parameter.   While this should often be similar to the total flux definition in the photometric catalog (e.g., FLUX_AUTO), they won't necessarily be the same and there could be small offsets in the normalization between photometry and spectra.
Below we demonstrate a quick way to scale the spectrum to the photometry for more reliable combined fits.
End of explanation
# Reset scale parameter
if hasattr(mb,'pscale'):
    delattr(mb, 'pscale')
# Template rescaling, simple multiplicative factor
scl = mb.scale_to_photometry(order=0)
# has funny units of polynomial coefficients times 10**power, 
# see `grizli.fitting.GroupFitter.compute_scale_array`
# Scale value is the inverse, so, e.g., 
# scl.x = [8.89] means scale the grism spectrum by 10/8.89=1.12
print(scl.x) 
mb.pscale = scl.x 
# Redo template fit
tfit = mb.template_at_z(z=z_phot)
print('Simple scaling, chi-squared={0:.1f}'.format(tfit['chi2']))
fig = fitting.full_sed_plot(mb, tfit, zfit=None, bin=4)
Explanation: An offset is clearly visible betwen the spectra and photometry.  In the next step, we use an internal function to compute a correction factor to bring the spectrum in line with the photometry.  
The scale_to_photometry function integrates the observed 1D spectrum (combining all available grisms) through the photometry filter bandpasses and derives a polynomial correction to make the two agree.  
NB Because we use the observed spectrum directly, scale_to_photometry requires that at least one observed filter be entirely covered by the grism spectrum.  And fitting for a higher order correction requires multiple filters across the grism bandpass (generally at least one filter per polynomial order).
End of explanation
# Reset scale parameter
if hasattr(mb,'pscale'):
    delattr(mb, 'pscale')
# Template rescaling, linear fit
scl = mb.scale_to_photometry(order=1)
# has funny units of polynomial coefficients times 10**power, 
# see `grizli.fitting.GroupFitter.compute_scale_array`
# Scale value is the inverse, so, e.g., 
# scl.x = [8.89] means scale the grism spectrum by 10/8.89=1.12
print(scl.x) 
mb.pscale = scl.x 
# Redo template fit
tfit = mb.template_at_z(z=z_phot)
print('Simple scaling, chi-squared={0:.1f}'.format(tfit['chi2']))
fig = fitting.full_sed_plot(mb, tfit, zfit=None, bin=4)
Explanation: The function computed a multiplicative factor 10/8.9=1.1 to apply to the spectrum, which now falls nicely on top of the photometry.  Next we fit for a first-order linear scaling (normalization and slope).
End of explanation
# Now run the full redshift fit script with the photometry, which will also do the scaling
order=1
fitting.run_all_parallel(id, phot=phot, verbose=False, 
                         scale_photometry=order+1, zr=[1.5, 2.4])
Explanation: A bit of an improvement.  But be careful with corrections with order&gt;0 if there is limited photometry available in bands that overlap the spectrum.
End of explanation
zfit = pyfits.open('{0}_{1:05d}.full.fits'.format(root, id))
z_grism = zfit['ZFIT_STACK'].header['Z_MAP']
print('Best redshift: {0:.4f}'.format(z_grism))
# Compare PDFs
pztab = utils.GTable.gread(zfit['ZFIT_STACK'])
plt.plot(pztab['zgrid'], pztab['pdf'], label='grism+3D-HST')
plt.plot(ez.zgrid, ez.pz[ix,:], label='photo-z')
plt.semilogy()
plt.xlim(z_grism-0.05, z_grism+0.05); plt.ylim(1.e-2, 1000)
plt.xlabel(r'$z$'); plt.ylabel(r'$p(z)$')
plt.grid()
plt.legend()
Explanation: Redshift PDF
Now compare the redshift PDF for the grism+photometry and photometry-only (photo-z) fits.
End of explanation
import grizli.pipeline.photoz
# The catalog is automatically generated with a number of aperture sizes.  A total 
# correction is computed in the detection band, usually a weighted sum of all available 
# WFC3/IR filters, with the correction as the ratio between the aperture flux and the 
# flux over the isophotal segment region, the 'flux' column in the SEP catalog.
aper_ix = 1
total_flux_column = 'flux'
# Get external photometry from Vizier
get_external_photometry = True
# Set `object_only=True` to generate the `eazy.photoz.Photoz` object from the 
# internal photometric catalog without actually running the photo-zs on the catalog 
# with few bands.   
int_ez = grizli.pipeline.photoz.eazy_photoz(root, force=False, object_only=True, 
                  apply_background=True, aper_ix=aper_ix, 
                  apply_prior=True, beta_prior=True, 
                  get_external_photometry=get_external_photometry, 
                  external_limits=3, external_sys_err=0.3, external_timeout=300, 
                  sys_err=0.05, z_step=0.01, z_min=0.01, z_max=12, 
                  total_flux=total_flux_column)
# Available apertures
for k in int_ez.cat.meta:
    if k.startswith('APER'):
        print('Aperture {0}: R={1:>4.1f} pix = {2:>4.2f}"'.format(k, int_ez.cat.meta[k], 
                                                                 int_ez.cat.meta[k]*0.06))
#
k = 'APER_{0}'.format(aper_ix)
print('\nAperture used {0}: R={1:>4.1f} pix = {2:>4.2f}"'.format(k, int_ez.cat.meta[k], 
                                                                 int_ez.cat.meta[k]*0.06))
# Integrate the grism templates through the filters on the redshift grid. 
int_phot_obj = photoz.EazyPhot(int_ez, grizli_templates=t0, zgrid=int_ez.zgrid)
# Reset scale parameter
if hasattr(mb,'pscale'):
    delattr(mb, 'pscale')
# Show the SED
int_phot, ii, dd = int_phot_obj.get_phot_dict(mb.ra, mb.dec)
mb.unset_photometry()
mb.set_photometry(**int_phot)
tfit = mb.template_at_z(z=z_phot)
print('No rescaling, chi-squared={0:.1f}'.format(tfit['chi2']))
fig = fitting.full_sed_plot(mb, tfit, zfit=None, bin=4)
fig.axes[0].set_ylim(-5,80)
# Run the grism fit with the direct image photometry
# Note that here we show that you can just pass the full photometry object 
# and the script will match the nearest photometric entry to the grism object.
order=0
fitting.run_all_parallel(id, phot_obj=int_phot_obj, verbose=False, 
                         scale_photometry=order+1, zr=[1.5, 2.4])
Explanation: Grizli internal photometry
The archive query functions shown in the Grizli-Pipeline notebook can optionally query for any HST imaging that overlaps with the grism exposures.  The automatic preparation script will process and align these along with the grism+direct data and produce a simple forced-aperture photometric catalog.  Currently this catalog just has matched aperture photometry and does not make any effort at PSF matching.  
For cases where no external photometric catalog is available, this photometry can be used to help constrain the redshift fits.  The catalog and eazy-py products are automatically generated by the preparation scripts.
The grizli.pipeline.photoz scripts can also try to fetch additional photometry from published Vizier catalogs, e.g., GALEX, SDSS, PanSTARRS, WISE, CFHT-LS.
End of explanation
zfit2 = pyfits.open('{0}_{1:05d}.full.fits'.format(root, id))
pztab2 = utils.GTable.gread(zfit2['ZFIT_STACK'])
plt.plot(ez.zgrid, ez.fit_chi2[ix,:] - ez.fit_chi2[ix,:].min(), 
         label='3D-HST photo-z')
plt.plot(pztab['zgrid'],  pztab['chi2'] - pztab['chi2'].min(), 
         label='grism + 3D-HST photometry')
plt.plot(pztab2['zgrid'], pztab2['chi2'] - pztab2['chi2'].min(), 
         label='grism + direct image photometry')
plt.legend()
plt.xlabel('z'); plt.ylabel(r'$\chi^2$')
plt.xlim(1.5, 2.4); plt.ylim(-200, 4000); plt.grid()
Explanation: Compare chi-squared
In this case of this bright spectrum with strong features, the grism dominates the fit even with the extensive 3D-HST catalog
End of explanation
<END_TASK>