#!/usr/bin/env python
# coding: utf-8

# # Calculating a PSF for Roman Wide Field Instrument (WFI)
# 
# This demonstration notebook offers a graphical interface to the basic functionality of WebbPSF-Roman, as well an example of performing a calculation with the Python scripting interface suited for more advanced calculations.
# 
# ## Background and boilerplate
# 
# Before we can do a calculation, we must set up the notebook by importing the packages we use and setting up logging output so we can follow the progress of the calculations. The cell below imports WebbPSF and standard scientific Python tools, and configures some options to make plots prettier.
# 
# **You shouldn't need to edit anything in the next cell, just go ahead and run it.**
# 
# *(Note: click in a cell and use **Shift + Enter** or click the play button <i class="fa-step-forward fa"></i> above to run it)*

# In[12]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib
matplotlib.rcParams['figure.figsize'] = (16, 7)  # make the default figure size larger
matplotlib.rcParams['image.interpolation'] = 'nearest'  # don't blur/smooth image plots
from matplotlib import pyplot as plt
import webbpsf
import webbpsf.roman


# The next cell tells WebbPSF to log information about what it is doing so that we can watch the progress of the calculation:

# In[13]:


webbpsf.setup_logging()


# **Note:** As you explore in this notebook, you may see certain warnings that look like this:
# 
# <div class="output_subarea output_stderr" style="font-family: monospace; margin: 1em">Warning: something happened! </div>
# 
# For the most part, warnings are safe to ignore. In particular, warnings referencing the matplotlib plotting library or the FITS library in Astropy don't indicate anything that could affect the accuracy of the calculations.

# # Using the WFI model in WebbPSF
# 
# Each instrument in WebbPSF is represented as a Python `class`, and the Wide Field Instrument model is in `webbpsf.roman.WFI`. We can to instantiate one to work with, in the same was as any of the JWST instruments.

# In[3]:


wfi = webbpsf.roman.WFI()


# ## Using the notebook interface
# 
# There's a notebook-friendly interface for the Wide Field Instrument PSF model. Bring it up in your notebook by running the following cell, then experiment with the different options, or read on for more explanation.  
# 
# Note that the calculations will typically take several seconds to run. 

# In[4]:


webbpsf.show_notebook_interface('wfi')


# ### Calculate PSF
# 
# When you click the "Calculate PSF" button, you will see some output as the calculation progresses. When it completes, it will show a plot showing four panels (counting left-to-right, top-to-bottom) representing the optical planes in the model:
# 
#   1. The entrance pupil transmission in black and white, and the phase ranging from red to blue to show the wavefront error. We include here a plausible estimate for wavefront error due to mirror polishing variations. (This is currently approximated by a map of high-frequency errors in the Hubble primary, since such data is not yet available for the Roman primary mirror.)
#   2. The exit pupil, showing the same data but with a change in coordinate system due to passing through focus. This is the pupil orientation as seen by the WFI looking outwards at the sky. 
#   3. The same transmission and phase map with the addition of the phase term due to field-dependent optical aberrations and intermediate instrument optics. (Specifically, this is based on Zernike coefficients derived from the Cycle 5 optical modeling effort at GSFC.)
#   4. The final oversampled detector plane, with log-scaled intensity
# 
# Below that, you will see a side-by-side comparison of the oversampled PSF, and the PSF binned down to detector pixels.
# 
# Also, a button labeled "Download FITS image from last calculation" will appear below the "Calculate PSF" button. Click that to download the oversampled and detector-pixel-binned images as a multi-extension FITS file. (WebbPSF also offers tools to analyze PSFs within the notebook or your own scripts, which are described in the next section.)
# 
# ### Display Optical System
# 
# This shows a 2 x 2 grid of plots. The left hand side shows transmission (e.g. pupil or mask shape), and the right side shows optical path difference (which is converted to phase across the pupil). The first row represents the pupil plane at the primary mirror, the second row is the telescope exit pupil as seen by the instrument, and the third row is the notional pupil plane after all the field dependent aberrations have been applied to the wavefront, but before the final propagation to a detector or image plane.
# 
# ### Clear Output
# 
# The output from the calculation process can be pretty verbose, so this button is here to clear both text output and plots.

# ## Using the Python programming interface
# 
# Alternatively, you can configure the WFI instance yourself in Python. A more detailed example is presented in the [WebbPSF-Roman documentation page](https://pythonhosted.org/webbpsf/roman.html), but here we will just show a simple monochromatic calculation at the default field position.
# 
# The `wfi.calcPSF()` method returns a [FITS HDUList object](http://docs.astropy.org/en/stable/io/fits/index.html), which you can then write out to a file or analyze further in the notebook.

# In[5]:


plt.figure(figsize=(8,10))
mono_psf = wfi.calc_psf(monochromatic=1.2e-6, display=True)


# Now you have the calculation result in the `mono_psf` variable, and can use various utility functions in WebbPSF to analyze it. The FITS object has an extension called `OVERSAMP` with each pixel split according to the default oversampling factor (4), and an extension called `DET_SAMP` with that image binned down to detector pixels.

# In[6]:


mono_psf.info()


# Let's plot the PSF in detector pixels:

# In[7]:


webbpsf.display_psf(mono_psf, ext='DET_SAMP')


# WebbPSF also includes functions for measuring EE, profiles, and centroids (described in the [WebbPSF documentation](http://pythonhosted.org/webbpsf/api_reference.html#functions) and the [POPPY documentation](http://pythonhosted.org/poppy/api.html#functions)). Below we measure the radial profile and encircled energy curve for the monochromatic PSF. (Note that the FWHM is also computed and labeled on the radial profile plot.)

# In[8]:


plt.figure(figsize=(8, 6))
webbpsf.display_profiles(mono_psf)


# If you want, the FITS object containing the PSF can be written out to a file and downloaded to your computer. This can be useful if you need it as an input to another tool.

# In[10]:


mono_psf.writeto('./mono_psf_1.2um.fits', overwrite=True)


# After you run the previous cell, this link will take you to download the FITS image: <a href="files/mono_psf_1.2um.fits">Download mono_psf_1.2um.fits</a>
# 
# How that works is a little tricky: when you write `./mono_psf_1.2um.fits`, you're saying you want to save the file in the current working directory for the *Python* process. If you're working locally, that's just the directory where you started the `jupyter notebook` command. If you're working on a remote server, files saved from the notebook will be available at the URL `files/your_filename.fits` relative to this page.
# 
# For example, if you're viewing this notebook at `https://example.com/user/janedoe/notebooks/WebbPSF-Roman_Tutorial.ipynb`, your file will be at `https://example.com/user/janedoe/notebooks/`**files/mono_psf_1.2um.fits**.

# ## Some more options

# A list of available WFI filters can be output:

# In[9]:


wfi.filter_list


# In[10]:


wfi = webbpsf.roman.WFI()

wfi.filter = 'GRISM1'
wfi.detector = 'SCA14'

plt.figure(figsize=(8,10))
poly_psf =  wfi.calc_psf(nlambda=10, display=True)


# In[11]:


webbpsf.display_psf(poly_psf, ext='DET_SAMP')


# # What next?
# 
# Keep working in this notebook, if you like! For reference, there's always a [pristine copy of this notebook](https://github.com/spacetelescope/webbpsf/blob/master/notebooks/WebbPSF-Roman_Tutorial.ipynb) to refer back to in the  [WebbPSF GitHub repository](https://github.com/spacetelescope/webbpsf).  If you have not previously used this notebook interface to Python, the Help menu available above has a tutorial and a useful list of keyboard shortcuts. 
# 
#   * **Review the [WebbPSF documentation](https://pythonhosted.org/webbpsf/) and the [POPPY documentation](https://pythonhosted.org/poppy/)**
#   * **Report any issues to us on GitHub** — 
#     WebbPSF and POPPY are developed on GitHub: [spacetelescope/webbpsf](https://github.com/spacetelescope/webbpsf) and [spacetelescope/poppy](https://github.com/spacetelescope/poppy) respectively.
#     The best way to report bugs is through the GitHub issue trackers: [WebbPSF](https://github.com/spacetelescope/webbpsf/issues) or [POPPY](https://github.com/spacetelescope/poppy/issues). (We also welcome pull requests from the community, if there's functionality you think should be included!)
#   * **Contact us through the STScI helpdesk** — You can always email your questions to <a href="mailto:help@stsci.edu">help@stsci.edu</a>, and our helpdesk people will make sure your request gets to the right person
#   * **Sign up for WebbPSF update announcements** —
#     This is entirely optional, but you may wish to sign up to the mailing list webbpsf-users@stsci.edu. This is a low-traffic moderated announce-only list, to which we will periodically post announcements of updates to this software. To subscribe, visit the [maillist.stsci.edu](http://maillist.stsci.edu) server

# In[ ]: