├── chromatic ├── noise │ ├── __init__.py │ └── jwst │ │ ├── __init__.py │ │ ├── simulate.py │ │ ├── extract.py │ │ ├── visualize.py │ │ └── etc.py ├── archives │ ├── __init__.py │ └── mast.py ├── rainbows │ ├── converters │ │ ├── __init__.py │ │ └── for_altair.py │ ├── get │ │ ├── fluxlike │ │ │ ├── __init__.py │ │ │ └── subset.py │ │ ├── __init__.py │ │ ├── timelike │ │ │ ├── __init__.py │ │ │ ├── median_lightcurve.py │ │ │ ├── descriptions.txt │ │ │ ├── average_lightcurve.py │ │ │ └── subset.py │ │ └── wavelike │ │ │ ├── __init__.py │ │ │ ├── median_spectrum.py │ │ │ ├── descriptions.txt │ │ │ ├── average_spectrum.py │ │ │ ├── expected_uncertainty.py │ │ │ ├── spectral_resolution.py │ │ │ ├── measured_scatter.py │ │ │ ├── measured_scatter_in_bins.py │ │ │ └── subset.py │ ├── visualizations │ │ ├── timelike │ │ │ ├── __init__.py │ │ │ ├── descriptions.txt │ │ │ ├── median_lightcurve.py │ │ │ └── average_lightcurve.py │ │ ├── diagnostics │ │ │ ├── __init__.py │ │ │ ├── descriptions.txt │ │ │ ├── paint_quantities.py │ │ │ └── plot_quantities.py │ │ ├── models │ │ │ ├── __init__.py │ │ │ └── descriptions.txt │ │ ├── wavelike │ │ │ ├── __init__.py │ │ │ ├── descriptions.txt │ │ │ ├── median_spectrum.py │ │ │ ├── average_spectrum.py │ │ │ ├── spectral_resolution.py │ │ │ └── noise_comparison_in_bins.py │ │ ├── __init__.py │ │ ├── descriptions.txt │ │ ├── plot.py │ │ ├── paint.py │ │ └── colors.py │ ├── helpers │ │ ├── __init__.py │ │ ├── descriptions.txt │ │ ├── save.py │ │ ├── get.py │ │ └── help.py │ ├── actions │ │ ├── __init__.py │ │ ├── compare.py │ │ ├── shift.py │ │ ├── descriptions.txt │ │ ├── inject_outliers.py │ │ ├── attach_model.py │ │ ├── inject_spectrum.py │ │ ├── flag_outliers.py │ │ ├── inflate_uncertainty.py │ │ └── fold.py │ ├── readers │ │ ├── rainbow_npy.py │ │ ├── rainbow_FITS.py │ │ ├── radica.py │ │ ├── schlawin.py │ │ ├── eureka_channels.py │ │ ├── espinoza.py │ │ ├── feinstein.py │ │ ├── text.py │ │ ├── eureka_txt.py │ │ ├── carter_and_may.py │ │ ├── xarray_stellar_spectra.py │ │ ├── xarray_raw_light_curves.py │ │ ├── nres.py │ │ └── xarray_fitted_light_curves.py │ ├── writers │ │ ├── rainbow_npy.py │ │ ├── rainbow_FITS.py │ │ ├── text.py │ │ ├── __init__.py │ │ └── template.py │ ├── __init__.py │ └── withmodel.py ├── spectra │ ├── __init__.py │ ├── library.py │ └── planck.py ├── version.py ├── tools │ ├── __init__.py │ └── custom_units.py ├── __init__.py └── tests │ ├── test_archives.py │ ├── setup_tests.py │ ├── test_units.py │ ├── test_concatenate.py │ ├── test_trimming.py │ ├── test_time.py │ ├── __init__.py │ ├── test_fold.py │ ├── test_outliers.py │ ├── test_history.py │ ├── test_multi.py │ ├── test_shift.py │ ├── test_remove_trends.py │ ├── test_normalize.py │ ├── test_conversions.py │ ├── test_summaries.py │ ├── test_ok.py │ ├── test_io.py │ ├── test_spectra.py │ ├── test_operations.py │ ├── test_withmodel.py │ ├── test_simulations.py │ └── test_align_wavelengths.py ├── .pre-commit-config.yaml ├── docs ├── index.md ├── documentation.ipynb ├── multi.ipynb.ignore ├── tools │ ├── colormaps.ipynb │ └── transits.ipynb └── api.md ├── LICENSE ├── notes ├── how-to-do-pre-commit.md └── how-to-do-the-docs.md ├── mkdocs.yml ├── README.md └── .gitignore /chromatic/noise/__init__.py: -------------------------------------------------------------------------------- 1 | from .jwst import * 2 | -------------------------------------------------------------------------------- /chromatic/archives/__init__.py: -------------------------------------------------------------------------------- 1 | from .mast import * 2 | -------------------------------------------------------------------------------- /chromatic/rainbows/converters/__init__.py: -------------------------------------------------------------------------------- 1 | from .for_altair import * 2 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/fluxlike/__init__.py: -------------------------------------------------------------------------------- 1 | from .subset import * 2 | -------------------------------------------------------------------------------- /chromatic/spectra/__init__.py: -------------------------------------------------------------------------------- 1 | from .planck import * 2 | from .phoenix import * 3 | -------------------------------------------------------------------------------- /chromatic/version.py: -------------------------------------------------------------------------------- 1 | __version__ = "0.5.0" 2 | 3 | 4 | def version(): 5 | return __version__ 6 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/__init__.py: -------------------------------------------------------------------------------- 1 | from .timelike import * 2 | from .wavelike import * 3 | from .fluxlike import * 4 | -------------------------------------------------------------------------------- /chromatic/tools/__init__.py: -------------------------------------------------------------------------------- 1 | from .custom_units import * 2 | from .resampling import * 3 | from .transits import * 4 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/timelike/__init__.py: -------------------------------------------------------------------------------- 1 | from .average_lightcurve import * 2 | from .median_lightcurve import * 3 | -------------------------------------------------------------------------------- /chromatic/rainbows/helpers/__init__.py: -------------------------------------------------------------------------------- 1 | from .history import * 2 | from .help import * 3 | from .get import * 4 | from .save import * 5 | -------------------------------------------------------------------------------- /chromatic/noise/jwst/__init__.py: -------------------------------------------------------------------------------- 1 | from .etc import * 2 | from .pandexo import * 3 | from .simulate import * 4 | from .visualize import * 5 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/diagnostics/__init__.py: -------------------------------------------------------------------------------- 1 | from .paint_quantities import * 2 | from .plot_quantities import * 3 | from .histogram import * 4 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/timelike/__init__.py: -------------------------------------------------------------------------------- 1 | from .average_lightcurve import * 2 | from .median_lightcurve import * 3 | from .subset import * 4 | from .time import * 5 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .paint_with_models import * 2 | from .plot_one_wavelength_with_models import * 3 | from .plot_with_model import * 4 | -------------------------------------------------------------------------------- /chromatic/__init__.py: -------------------------------------------------------------------------------- 1 | from .version import * 2 | from .rainbows import * 3 | from .imports import * 4 | from .spectra import * 5 | from .archives import * 6 | from .tools import * 7 | -------------------------------------------------------------------------------- /chromatic/tests/test_archives.py: -------------------------------------------------------------------------------- 1 | from ..archives import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_download_from_mast(): 6 | d = download_from_mast() 7 | assert d[0]["Status"] == "COMPLETE" 8 | -------------------------------------------------------------------------------- /chromatic/tests/setup_tests.py: -------------------------------------------------------------------------------- 1 | import os, pytest 2 | import matplotlib.pyplot as plt 3 | 4 | test_directory = "examples/" 5 | 6 | try: 7 | os.mkdir(test_directory) 8 | except FileExistsError: 9 | pass 10 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/timelike/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🎨⏰🐋 | plot_average_lightcurve | Plot the weighted average flux per time. 3 | 🎨⏰🖖 | plot_median_lightcurve | Plot the median flux per time. 4 | -------------------------------------------------------------------------------- /chromatic/spectra/library.py: -------------------------------------------------------------------------------- 1 | from ..imports import * 2 | from ..resampling import bintoR, bintogrid 3 | from ..version import __version__ 4 | import astropy.config.paths 5 | from astropy.utils.data import is_url_in_cache, cache_total_size 6 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/wavelike/__init__.py: -------------------------------------------------------------------------------- 1 | from .spectral_resolution import * 2 | from .noise_comparison import * 3 | from .average_spectrum import * 4 | from .median_spectrum import * 5 | from .noise_comparison_in_bins import * 6 | -------------------------------------------------------------------------------- /chromatic/tests/test_units.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_custom_units(): 6 | u.Unit("(MJy/sr)^2") 7 | u.Unit("(DN/s)^2") 8 | u.Unit("microns") 9 | u.Unit("electrons") 10 | -------------------------------------------------------------------------------- /chromatic/rainbows/helpers/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🙋🌈📄 | help | Get one-line help summaries of available methods. 3 | 📓🌈🪵 | history | Get the history that went into this Rainbow. 4 | 💾🌈📼 | save | Save this Rainbow out to a permanent file. 5 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/__init__.py: -------------------------------------------------------------------------------- 1 | from .spectral_resolution import * 2 | from .average_spectrum import * 3 | from .median_spectrum import * 4 | from .expected_uncertainty import * 5 | from .measured_scatter import * 6 | from .measured_scatter_in_bins import * 7 | from .subset import * 8 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/diagnostics/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🎨🗂📺 | paint_quantities | Show multiple 2D (wavelength and time) quantities as imshow maps. 3 | 🎨🗂🧶 | plot_quantities | Show multiple 1D (wavelength or time) quantities as scatter plots. 4 | 🎨🐍🗑 | histogram | Plot histogram for one wavelength. 5 | -------------------------------------------------------------------------------- /chromatic/tests/test_concatenate.py: -------------------------------------------------------------------------------- 1 | from .setup_tests import * 2 | from ..rainbows import * 3 | 4 | 5 | def test_concatenate_simple(): 6 | s = SimulatedRainbow() 7 | more_wavelengths = s.concatenate_in_wavelength(s) 8 | assert more_wavelengths.nwave == 2 * s.nwave 9 | more_times = s.concatenate_in_time(s) 10 | assert more_times.ntime == 2 * s.ntime 11 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/psf/black 3 | rev: 22.6.0 4 | hooks: 5 | - id: black 6 | - repo: https://github.com/psf/black 7 | rev: 22.6.0 8 | hooks: 9 | - id: black-jupyter 10 | - repo: https://github.com/pre-commit/pre-commit-hooks 11 | rev: v4.3.0 12 | hooks: 13 | - id: trailing-whitespace 14 | - id: end-of-file-fixer 15 | - id: check-added-large-files 16 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/models/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🎨📺🎢 | paint_with_models | Paint a flux map with model components. 3 | 🎨🖌🎢 | plot_one_wavelength_with_models | Plot one wavelength's light curve with model components. 4 | 🎨📽🎢 | animate_with_models | Animate all wavelengths' light curves with model components. 5 | 🎨🧶🎢 | plot_with_model | Plot a sequence of light curves with their models. 6 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/wavelike/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🎨🌊🔭 | plot_average_spectrum | Plot the weighted average flux per wavelength. 3 | 🎨🌊💎 | plot_spectral_resolution | Plot the spectral resolution per wavelength. 4 | 🎨🌊🎧 | plot_noise_comparison | Plot the measured and expected scatter per wavelength. 5 | 🎨🌊🥦 | plot_noise_comparison_in_bins | Plot measured and expected scatter in different size bins. 6 | -------------------------------------------------------------------------------- /chromatic/tests/test_trimming.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_trim(): 6 | r = SimulatedRainbow().inject_noise() 7 | r.fluxlike["flux"][:3, :] = np.nan 8 | r.fluxlike["flux"][:, -4:] = np.nan 9 | original_shape = r.shape 10 | t = r.trim() 11 | new_shape = t.shape 12 | assert new_shape[0] == original_shape[0] - 3 13 | assert new_shape[1] == original_shape[1] - 4 14 | print(r, t) 15 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/__init__.py: -------------------------------------------------------------------------------- 1 | from .timelike import * 2 | from .wavelike import * 3 | from .models import * 4 | from .diagnostics import * 5 | 6 | 7 | from .colors import * 8 | from .imshow import * 9 | from .scatter import * 10 | from .pcolormesh import * 11 | from .paint import * 12 | 13 | from .plot_lightcurves import * 14 | from .animate import * 15 | from .interactive import * 16 | from .plot_spectra import * 17 | from .plot import * 18 | from .utilities import * 19 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/median_spectrum.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_median_spectrum"] 4 | 5 | 6 | def get_median_spectrum(self): 7 | """ 8 | Return a spectrum of the star, medianed over all times. 9 | 10 | Returns 11 | ------- 12 | median_spectrum : array 13 | Wavelike array of fluxes. 14 | """ 15 | with warnings.catch_warnings(): 16 | warnings.simplefilter("ignore") 17 | return np.nanmedian(self.get_ok_data(), axis=1) 18 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/timelike/median_lightcurve.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_median_lightcurve"] 4 | 5 | 6 | def get_median_lightcurve(self): 7 | """ 8 | Return a lightcurve of the star, medianed over all wavelengths. 9 | 10 | Returns 11 | ------- 12 | median_lightcurve : array 13 | Timelike array of fluxes. 14 | """ 15 | with warnings.catch_warnings(): 16 | warnings.simplefilter("ignore") 17 | return np.nanmedian(self.get_ok_data(), axis=0) 18 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/timelike/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | ⏰🐋🌈 | get_average_lightcurve | Get the weighted average light curve. 3 | ⏰🔎🌈 | get_for_time | Get a quantity associated with a time index. 4 | ⏰🖖🌈 | get_median_lightcurve | Get the median light curve. 5 | ⏰👌🌈 | get_ok_data_for_time | Get a quantity associated with a time index. 6 | ⏰🛰🌈 | get_times_as_astropy | Get the times as an astropy Time object. 7 | ⏰🚀🌈 | set_times_from_astropy | Set the times from an astropy Time object (modifies in-place). 8 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/__init__.py: -------------------------------------------------------------------------------- 1 | from .binning import * 2 | from .trim import * 3 | from .normalization import * 4 | from .align_wavelengths import * 5 | from .shift import * 6 | from .inject_transit import * 7 | from .inject_systematics import * 8 | from .inject_noise import * 9 | from .inject_spectrum import * 10 | from .inject_outliers import * 11 | from .fold import * 12 | from .flag_outliers import * 13 | from .compare import * 14 | from .remove_trends import * 15 | from .attach_model import * 16 | from .inflate_uncertainty import * 17 | from .concatenate import * 18 | -------------------------------------------------------------------------------- /chromatic/tests/test_time.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_astropy_times(): 6 | s = SimulatedRainbow(time=(Time.now().jd + np.linspace(0, 1)) * u.day) 7 | format = "jd" 8 | scale = "tdb" 9 | original_times = s.time 10 | astropy_times = s.get_times_as_astropy( 11 | format=format, scale=scale, is_barycentric=True 12 | ) 13 | s.set_times_from_astropy(astropy_times, is_barycentric=True) 14 | 15 | assert np.all(s.time == original_times) 16 | assert np.all(astropy_times == s.get_times_as_astropy()) 17 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/timelike/average_lightcurve.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_average_lightcurve"] 4 | 5 | 6 | def get_average_lightcurve(self): 7 | """ 8 | Return a lightcurve of the star, averaged over all wavelengths. 9 | 10 | This uses `bin`, which is a horribly slow way of doing what is 11 | fundamentally a very simply array calculation, because we 12 | don't need to deal with partial pixels. 13 | 14 | Returns 15 | ------- 16 | lightcurve : array 17 | Timelike array of fluxes. 18 | """ 19 | return self.get_average_lightcurve_as_rainbow().flux[0, :] 20 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🌊🐋🌈 | get_average_spectrum | Get the weighted average spectrum. 3 | 🌊🔎🌈 | get_for_wavelength | Get a quantity associated with a wavelength index. 4 | 🌊🎯🌈 | get_measured_scatter | Get the measured scatter on the time series for each wavelength. 5 | 🌊🖖🌈 | get_median_spectrum | Get the median spectrum. 6 | 🌊👌🌈 | get_ok_data_for_wavelength | Get a quantity associated with a wavelength index. 7 | 🌊💎🌈 | get_spectral_resolution | Get the spectral resolution (R=w/dw). 8 | 🌊🃏🌈 | get_typical_uncertainty | Get the typical uncertainty for each wavelength. 9 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/average_spectrum.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_average_spectrum"] 4 | 5 | 6 | def get_average_spectrum(self): 7 | """ 8 | Return a average_spectrum of the star, averaged over all times. 9 | 10 | This uses `bin`, which is a horribly slow way of doing what is 11 | fundamentally a very simply array calculation, because we 12 | don't need to deal with partial pixels. 13 | 14 | Returns 15 | ------- 16 | average_spectrum : array 17 | Wavelike array of average spectrum. 18 | """ 19 | return self.get_average_spectrum_as_rainbow().flux[:, 0] 20 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🎨📽⏰ | animate_lightcurves | Animate a sequence of light curves across different wavelengths. 3 | 🎨📽🌊 | animate_spectra | Animate a sequence of spectra across different times. 4 | 🎨🖼📺 | imshow | Paint a map of flux across wavelength and time. 5 | 🎨🕹📺 | imshow_interact | Show flux map and lightcurves with interactive wavelength selection. 6 | 🎨🖌📺 | pcolormesh | Paint a map of flux across wavelength and time (with non-uniform grids). 7 | 🎨🖌🧶 | plot | Plot a sequence of light curves with vertical offsets. 8 | -------------------------------------------------------------------------------- /docs/index.md: -------------------------------------------------------------------------------- 1 | # `chromatic` 2 | 3 | The `chromatic-lightcurves` package is a friendly python tool for working with spectroscopic light curves. Inspired by the [`lightkurve`](https://docs.lightkurve.org/) package, we aim for this tool to provide an easy way interact with datasets representing the brightness of a star as a function of both time and wavelength. 4 | 5 | This package defines a spectroscopic light curve object (a `Rainbow` = 🌈) to provide easy access to wavelength, time, flux, and uncertainty attributes. It allows reading and writing these datasets from/to a variety of formats, simplifies many common calculations, and provides a plethora of visualizations. Read on for quick introductions to how these little 🌈s work! 6 | -------------------------------------------------------------------------------- /chromatic/tests/__init__.py: -------------------------------------------------------------------------------- 1 | from .test_io import * 2 | from .test_basics import * 3 | from .test_operations import * 4 | from .test_normalize import * 5 | from .test_binning import * 6 | from .test_simulations import * 7 | from .test_visualizations import * 8 | from .test_multi import * 9 | from .test_resampling import * 10 | from .test_units import * 11 | from .test_align_wavelengths import * 12 | from .test_shift import * 13 | from .test_conversions import * 14 | from .test_summaries import * 15 | from .test_withmodel import * 16 | from .test_ok import * 17 | from .test_history import * 18 | from .test_spectra import * 19 | from .test_remove_trends import * 20 | from .test_time import * 21 | from .test_fold import * 22 | from .test_outliers import * 23 | -------------------------------------------------------------------------------- /chromatic/rainbows/helpers/save.py: -------------------------------------------------------------------------------- 1 | from ..writers import * 2 | 3 | 4 | def save(self, filepath="test.rainbow.npy", format=None, **kw): 5 | """ 6 | Save this `Rainbow` out to a file. 7 | 8 | Parameters 9 | ---------- 10 | filepath : str 11 | The filepath pointing to the file to be written. 12 | (For now, it needs a `.rainbow.npy` extension.) 13 | format : str, optional 14 | The file format of the file to be written. If `None`, 15 | the format will be guessed automatically from the 16 | filepath. 17 | **kw : dict, optional 18 | All other keywords will be passed to the writer. 19 | """ 20 | # figure out the best writer 21 | writer = guess_writer(filepath, format=format) 22 | 23 | # use that writer to save the file 24 | writer(self, filepath, **kw) 25 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/timelike/median_lightcurve.py: -------------------------------------------------------------------------------- 1 | from ..utilities import * 2 | from ....imports import * 3 | 4 | __all__ = ["plot_median_lightcurve"] 5 | 6 | 7 | def plot_median_lightcurve(self, filename=None, **kw): 8 | """ 9 | Plot the weighted median lightcurve as a function of time. 10 | 11 | Parameters 12 | ---------- 13 | **kw : dict 14 | Additional keywords will be passed onward to the helper 15 | function `._scatter_timelike_or_wavelike`. Please see its 16 | docstrings for options about plot appearance and layout. 17 | """ 18 | y = self.get_median_lightcurve() 19 | self._scatter_timelike_or_wavelike( 20 | x=self.time, 21 | y=y, 22 | ylabel=f"Flux ({_get_unit_string(y)})", 23 | **kw, 24 | ) 25 | if filename is not None: 26 | self.savefig(filename) 27 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/wavelike/median_spectrum.py: -------------------------------------------------------------------------------- 1 | from ..utilities import * 2 | from ....imports import * 3 | 4 | __all__ = ["plot_median_spectrum"] 5 | 6 | 7 | def plot_median_spectrum(self, filename=None, **kw): 8 | """ 9 | Plot the weighted median spectrum as a function of wavelength. 10 | 11 | Parameters 12 | ---------- 13 | **kw : dict 14 | Additional keywords will be passed onward to the helper 15 | function `._scatter_timelike_or_wavelike`. Please see its 16 | docstrings for options about plot appearance and layout. 17 | """ 18 | y = self.get_median_spectrum() 19 | self._scatter_timelike_or_wavelike( 20 | x=self.wavelength, 21 | y=y, 22 | ylabel=f"Flux ({_get_unit_string(y)})", 23 | **kw, 24 | ) 25 | if filename is not None: 26 | self.savefig(filename) 27 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/timelike/average_lightcurve.py: -------------------------------------------------------------------------------- 1 | from ..utilities import * 2 | from ....imports import * 3 | 4 | __all__ = ["plot_average_lightcurve"] 5 | 6 | 7 | def plot_average_lightcurve(self, filename=None, **kw): 8 | """ 9 | Plot the weighted average lightcurve as a function of time. 10 | 11 | Parameters 12 | ---------- 13 | **kw : dict 14 | Additional keywords will be passed onward to the helper 15 | function `._scatter_timelike_or_wavelike`. Please see its 16 | docstrings for options about plot appearance and layout. 17 | """ 18 | y = self.get_average_lightcurve() 19 | self._scatter_timelike_or_wavelike( 20 | x=self.time, 21 | y=y, 22 | ylabel=f"Flux ({_get_unit_string(y)})", 23 | **kw, 24 | ) 25 | if filename is not None: 26 | self.savefig(filename) 27 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/wavelike/average_spectrum.py: -------------------------------------------------------------------------------- 1 | from ..utilities import * 2 | from ....imports import * 3 | 4 | __all__ = ["plot_average_spectrum"] 5 | 6 | 7 | def plot_average_spectrum(self, filename=None, **kw): 8 | """ 9 | Plot the weighted average spectrum as a function of wavelength. 10 | 11 | Parameters 12 | ---------- 13 | **kw : dict 14 | Additional keywords will be passed onward to the helper 15 | function `._scatter_timelike_or_wavelike`. Please see its 16 | docstrings for options about plot appearance and layout. 17 | """ 18 | y = self.get_average_spectrum() 19 | self._scatter_timelike_or_wavelike( 20 | x=self.wavelength, 21 | y=y, 22 | ylabel=f"Flux ({_get_unit_string(y)})", 23 | **kw, 24 | ) 25 | if filename is not None: 26 | self.savefig(filename) 27 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/rainbow_npy.py: -------------------------------------------------------------------------------- 1 | """ 2 | Define a reader for chromatic .rainbow.npy files. 3 | """ 4 | 5 | # import the general list of packages 6 | from ...imports import * 7 | 8 | # define list of the only things that will show up in imports 9 | __all__ = ["from_rainbow_npy"] 10 | 11 | 12 | def from_rainbow_npy(rainbow, filepath): 13 | """ 14 | Populate a Rainbow from a file in the .rainbow.npy format. 15 | 16 | Parameters 17 | ---------- 18 | 19 | rainbow : Rainbow 20 | The object to be populated. 21 | 22 | filepath : str 23 | The path to the file to load, which should probably 24 | have an extension of `.rainbow.npy` 25 | """ 26 | 27 | # read in your file, however you like 28 | loaded_core_dictionaries, version_used = np.load(filepath, allow_pickle=True) 29 | 30 | for k in rainbow._core_dictionaries: 31 | vars(rainbow)[k] = loaded_core_dictionaries[k] 32 | -------------------------------------------------------------------------------- /chromatic/rainbows/writers/rainbow_npy.py: -------------------------------------------------------------------------------- 1 | """ 2 | Define a writer for chromatic .rainbow.npy files. 3 | """ 4 | 5 | # import the general list of packages 6 | from ...imports import * 7 | from ...version import version 8 | 9 | # define list of the only things that will show up in imports 10 | __all__ = ["to_rainbow_npy"] 11 | 12 | 13 | def to_rainbow_npy(self, filepath, **kw): 14 | """ 15 | Write a Rainbow to a file in the .rainbow.npy format. 16 | 17 | Parameters 18 | ---------- 19 | 20 | self : Rainbow 21 | The object to be saved. 22 | 23 | filepath : str 24 | The path to the file to write. 25 | """ 26 | 27 | assert ".rainbow.npy" in filepath 28 | 29 | # populate a dictionary containing the four core dictionaries 30 | dictionary_to_save = self._get_core_dictionaries() 31 | 32 | # save that to a file 33 | np.save(filepath, [dictionary_to_save, version()], allow_pickle=True) 34 | -------------------------------------------------------------------------------- /chromatic/tools/custom_units.py: -------------------------------------------------------------------------------- 1 | import astropy.units as u 2 | 3 | MJy_sr = u.def_unit("MJy/sr", u.MJy / u.sr) 4 | MJy_sr_sq = u.def_unit("(MJy/sr)^2", (u.MJy / u.sr) ** 2) 5 | u.add_enabled_units([MJy_sr, MJy_sr_sq]) 6 | 7 | 8 | data_number_per_second = u.def_unit("DN/s") 9 | data_number_per_second_sq = u.def_unit("(DN/s)^2", data_number_per_second**2) 10 | for kludge_DN_unit in [data_number_per_second, data_number_per_second_sq]: 11 | try: 12 | u.add_enabled_units([kludge_DN_unit]) 13 | except ValueError: 14 | pass 15 | 16 | electrons_per_group = u.def_unit("Electrons/group") 17 | u.add_enabled_units([electrons_per_group]) 18 | 19 | # FIXME: some of these feel super kuldgy; how do we make this 20 | # interact more smoothly with units already defined in astropy 21 | # (for example, DN should work, but is always flaky for me...) 22 | 23 | u.add_enabled_aliases( 24 | {"microns": u.micron, "ELECTRONS": u.electron, "electrons": u.electron} 25 | ) 26 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/compare.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | from ..multi import * 3 | 4 | __all__ = ["compare"] 5 | 6 | 7 | def compare(self, rainbows): 8 | """ 9 | Compare this `Rainbow` to others. 10 | 11 | (still in development) This connects the current `Rainbow` 12 | to a collection of other `Rainbow` objects, which can then 13 | be visualized side-by-side in a uniform way. 14 | 15 | Parameters 16 | ---------- 17 | rainbows : list 18 | A list containing one or more other `Rainbow` objects. 19 | If you only want to compare with one other `Rainbow`, 20 | supply it in a 1-element list like `.compare([other])` 21 | 22 | Returns 23 | ------- 24 | rainbow : MultiRainbow 25 | A `MultiRainbow` comparison object including all input `Rainbow`s 26 | """ 27 | try: 28 | rainbows.remove(self) 29 | except (ValueError, IndexError): 30 | pass 31 | return compare_rainbows([self] + rainbows) 32 | -------------------------------------------------------------------------------- /chromatic/rainbows/__init__.py: -------------------------------------------------------------------------------- 1 | from .withmodel import * 2 | from .rainbow import * 3 | from .simulated import * 4 | from .multi import * 5 | from .writers import * 6 | 7 | 8 | def read_rainbow(filepath, **kw): 9 | """ 10 | A friendly wrapper to load time-series spectra and/or 11 | multiwavelength light curves into a `chromatic` Rainbow 12 | object. It will try its best to pick the best reader 13 | and return the most useful kind of object. 14 | 🦋🌅2️⃣🪜🎬👀🇮🇹📕🧑‍🏫🌈 15 | 16 | Parameters 17 | ---------- 18 | filepath : str, list 19 | The file or files to open. 20 | **kw : dict 21 | All other keyword arguments will be passed to 22 | the `Rainbow` initialization. 23 | 24 | Returns 25 | ------- 26 | rainbow : Rainbow, RainbowWithModel 27 | The loaded data! 28 | """ 29 | r = Rainbow(filepath, **kw) 30 | if "model" in r.fluxlike: 31 | return RainbowWithModel(**r._get_core_dictionaries()) 32 | else: 33 | return r 34 | -------------------------------------------------------------------------------- /chromatic/tests/test_fold.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_fold(N=5): 6 | 7 | for i in range(5): 8 | original_time = ( 9 | np.linspace(np.random.uniform(-0.5, -0.2), np.random.uniform(0.2, 0.5), 500) 10 | * u.day 11 | ) 12 | period = np.random.uniform(1, 10) * u.day 13 | t0 = np.random.uniform(-10, 10) * u.day 14 | N = np.random.randint(-10, 10) 15 | s = ( 16 | SimulatedRainbow(time=original_time + t0 + period * N, R=5) 17 | .inject_transit(P=period.to_value("day"), t0=t0.to_value("day")) 18 | .inject_noise(signal_to_noise=1000) 19 | ) 20 | f = s.fold(period=period, t0=t0) 21 | assert np.all(np.isclose(f.time, original_time, atol=1e-12)) 22 | fi, ax = plt.subplots(1, 2, figsize=(8, 3), constrained_layout=True) 23 | s.paint(ax=ax[0]) 24 | f.paint(ax=ax[1]) 25 | plt.savefig( 26 | os.path.join(test_directory, "demonstration-of-folding-to-period-and-t0.pdf") 27 | ) 28 | -------------------------------------------------------------------------------- /chromatic/rainbows/helpers/get.py: -------------------------------------------------------------------------------- 1 | def get(self, key, default=None): 2 | """ 3 | Retrieve an attribute by its string name. 4 | (This is a friendlier wrapper for `getattr()`). 5 | 6 | `r.get('flux')` is identical to `r.flux` 7 | 8 | This is different from indexing directly into 9 | a core dictionary (for example, `r.fluxlike['flux']`), 10 | because it can also be used to get the results of 11 | properties that do calculations on the fly (for example, 12 | `r.residuals` in the `RainbowWithModel` class). 13 | 14 | Parameters 15 | ---------- 16 | key : str 17 | The name of the attribute, property, or core dictionary item to get. 18 | default : any, optional 19 | What to return if the attribute can't be found. 20 | 21 | Returns 22 | ------- 23 | thing : any 24 | The thing you were trying to get. If unavailable, 25 | return the `default` (which by default is `None`) 26 | """ 27 | try: 28 | return getattr(self, key) 29 | except AttributeError: 30 | return default 31 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/expected_uncertainty.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_expected_uncertainty"] 4 | 5 | 6 | def get_expected_uncertainty(self, function=np.nanmedian, *args, **kw): 7 | """ 8 | Get the typical per-wavelength uncertainty. 9 | 10 | Parameters 11 | ---------- 12 | function : function, optional 13 | What function should be used to choose the "typical" 14 | value for each wavelength? Good options are probably 15 | things like `np.nanmedian`, `np.median`, `np.nanmean` 16 | `np.mean`, `np.percentile` 17 | *args : list, optional 18 | Addition arguments will be passed to `function` 19 | **kw : dict, optional 20 | Additional keyword arguments will be passed to `function` 21 | 22 | Returns 23 | ------- 24 | uncertainty_per_wavelength : array 25 | The uncertainty associated with each wavelength. 26 | """ 27 | uncertainty_per_wavelength = function( 28 | self.uncertainty, *args, axis=self.timeaxis, **kw 29 | ) 30 | return uncertainty_per_wavelength 31 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/fluxlike/subset.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = [ 4 | "get_ok_data", 5 | ] 6 | 7 | 8 | def get_ok_data( 9 | self, 10 | y="flux", 11 | minimum_acceptable_ok=1, 12 | ): 13 | """ 14 | A small wrapper to get the good data as a 2D array. 15 | 16 | Extract fluxes as 2D array, marking data that are not `ok` 17 | either as nan or by inflating uncertainties to infinity. 18 | 19 | Parameters 20 | ---------- 21 | y : array 22 | The desired quantity (default is `flux`) 23 | minimum_acceptable_ok : float, optional 24 | The smallest value of `ok` that will still be included. 25 | (1 for perfect data, 1e-10 for everything but terrible data, 0 for all data) 26 | 27 | Returns 28 | ------- 29 | flux : array 30 | The fluxes, but with not-OK data replaced with nan. 31 | """ 32 | # get 1D array of what to keep 33 | ok = self.ok >= minimum_acceptable_ok 34 | 35 | # create copy of flux array 36 | y = self.get(y) * 1 37 | 38 | # set bad values to nan 39 | y[ok == False] = np.nan 40 | 41 | return y 42 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Zach Berta-Thompson 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /chromatic/tests/test_outliers.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_flag_outliers(): 6 | noiseless = SimulatedRainbow(R=10, dt=10 * u.minute).inject_transit() 7 | noiseless_outliers = noiseless.inject_outliers() 8 | noiseless_flagged = noiseless_outliers.flag_outliers() 9 | 10 | noisy = SimulatedRainbow(R=10, dt=10 * u.minute).inject_transit().inject_noise() 11 | noisy_outliers = noisy.inject_outliers() 12 | noisy_flagged = noisy_outliers.flag_outliers() 13 | 14 | kw = dict(vmin=0.9, vmax=1.1) 15 | fi, ax = plt.subplots(2, 3, figsize=(10, 4), dpi=300) 16 | noisy.paint(ax=ax[0, 0], **kw) 17 | noisy_outliers.paint(ax=ax[0, 1], **kw) 18 | noisy_flagged.paint(ax=ax[0, 2], **kw) 19 | 20 | noiseless.paint(ax=ax[1, 0], **kw) 21 | noiseless_outliers.paint(ax=ax[1, 1], **kw) 22 | noiseless_flagged.paint(ax=ax[1, 2], **kw) 23 | 24 | plt.savefig(os.path.join(test_directory, "demonstration-of-flagging-outliers.pdf")) 25 | assert np.all( 26 | (noisy_flagged.fluxlike["injected_outliers"] != 0) 27 | == noisy_flagged.fluxlike["flagged_as_outlier"] 28 | ) 29 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/plot.py: -------------------------------------------------------------------------------- 1 | from .plot_lightcurves import * 2 | from .plot_spectra import * 3 | 4 | __all__ = ["plot"] 5 | 6 | 7 | def plot(self, xaxis="time", **kw): 8 | """ 9 | Plot flux either as a sequence of offset lightcurves (default) 10 | or as a sequence of offset spectra. 11 | 12 | Parameters 13 | ---------- 14 | xaxis : string 15 | What should be plotted on the x-axis of the plot? 16 | 'time' will plot a different light curve for each wavelength 17 | 'wavelength' will plot a different spectrum for each timepoint 18 | **kw : dict 19 | All other keywords will be passed along to either 20 | `.plot_lightcurves` or `.plot_spectra` as appropriate. 21 | Please see the docstrings for either of those functions 22 | to figure out what keyword arguments you might want to 23 | provide here. 24 | """ 25 | 26 | if xaxis.lower()[0] == "t": 27 | return self.plot_lightcurves(**kw) 28 | elif xaxis.lower()[0] == "w": 29 | return self.plot_spectra(**kw) 30 | else: 31 | cheerfully_suggest("Please specify either 'time' or 'wavelength' for `.plot()`") 32 | -------------------------------------------------------------------------------- /chromatic/tests/test_history.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | from ..rainbows.helpers.history import represent_as_copypasteable 4 | 5 | 6 | def test_history(): 7 | np.random.seed(0) 8 | x = SimulatedRainbow().inject_transit().inject_noise().bin(R=5).normalize() 9 | h = x.history() 10 | 11 | for k in ["Rainbow", "inject_transit", "bin", "normalize"]: 12 | assert k in h 13 | 14 | np.random.seed(0) 15 | new = eval(h) 16 | assert x == new 17 | 18 | 19 | def test_history_with_slicing_and_addition(): 20 | a = SimulatedRainbow().inject_transit().normalize()[:10:2, :] 21 | b = a * 2 22 | c = a + b 23 | d = eval(c.history()) 24 | assert c == d 25 | assert d == a * 3 26 | 27 | 28 | def test_represent_as_copypasteable(): 29 | for x_in in [ 30 | np.array([1, 2]) * u.day, 31 | [1, 2], 32 | [1, 2] * u.day, 33 | 1 * u.day, 34 | np.linspace(0, 1, 10), 35 | ]: 36 | string = represent_as_copypasteable(x_in) 37 | x_out = eval(string) 38 | match = np.all(np.isclose(x_in, x_out)) 39 | print(f"{x_in} == {x_out}? {match}") 40 | assert match 41 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/wavelike/spectral_resolution.py: -------------------------------------------------------------------------------- 1 | from ..utilities import * 2 | from ....imports import * 3 | 4 | __all__ = ["plot_spectral_resolution"] 5 | 6 | 7 | def plot_spectral_resolution( 8 | self, 9 | pixels_per_resolution_element=1, 10 | filename=None, 11 | **kw, 12 | ): 13 | """ 14 | Plot the spectral resolution as a function of wavelength. 15 | 16 | Parameters 17 | ---------- 18 | pixels_per_resolution_element : float 19 | How many pixels do we consider as a resolution element? 20 | **kw : dict 21 | Additional keywords will be passed onward to the helper 22 | function `._scatter_timelike_or_wavelike`. Please see its 23 | docstrings for options about plot appearance and layout. 24 | """ 25 | self._scatter_timelike_or_wavelike( 26 | x=self.wavelength, 27 | y=self.get_spectral_resolution( 28 | pixels_per_resolution_element=pixels_per_resolution_element 29 | ), 30 | ylabel=f"$R=\lambda/d\lambda$ ({pixels_per_resolution_element} pixel)", 31 | **kw, 32 | ) 33 | plt.title(self.get("title")) 34 | if filename is not None: 35 | self.savefig(filename) 36 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/spectral_resolution.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_spectral_resolution"] 4 | 5 | 6 | def get_spectral_resolution(self, pixels_per_resolution_element=1): 7 | """ 8 | Estimate the R=w/dw spectral resolution. 9 | 10 | Higher spectral resolutions correspond to more wavelength 11 | points within a particular interval. By default, it's 12 | estimated for the interval between adjacent wavelength 13 | bins. In unbinned data coming directly from a telescope, 14 | there's a good chance that adjacent pixels both sample 15 | the same resolution element as blurred by the telescope 16 | optics, so the `pixels_per_resolution_element` keyword 17 | should likely be larger than 1. 18 | 19 | Parameters 20 | ---------- 21 | pixels_per_resolution_element : float, optional 22 | How many pixels do we consider as a resolution element? 23 | 24 | Returns 25 | ------- 26 | R : array 27 | The spectral resolution at each wavelength. 28 | """ 29 | 30 | # calculate spectral resolution, for this pixels/element 31 | w = self.wavelength 32 | dw = np.gradient(self.wavelength) 33 | R = np.abs(w / dw / pixels_per_resolution_element) 34 | 35 | return R 36 | -------------------------------------------------------------------------------- /chromatic/tests/test_multi.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_multi(): 6 | a = ( 7 | SimulatedRainbow(R=3, dt=20 * u.minute) 8 | .inject_transit() 9 | .inject_noise(signal_to_noise=400) 10 | ) 11 | b = ( 12 | SimulatedRainbow(R=3, dt=20 * u.minute) 13 | .inject_transit() 14 | .inject_noise(signal_to_noise=800) 15 | ) 16 | 17 | m = MultiRainbow([a, b]) 18 | m.plot(spacing=0.02) 19 | plt.savefig(os.path.join(test_directory, "multi-plot-demonstration.png")) 20 | 21 | m.imshow() 22 | plt.savefig(os.path.join(test_directory, "multi-imshow-demonstration.png")) 23 | 24 | m.animate_lightcurves( 25 | filename=os.path.join( 26 | test_directory, "multi-animate-lightcurves-demonstration.gif" 27 | ) 28 | ) 29 | m.animate_spectra( 30 | filename=os.path.join(test_directory, "multi-animate-spectra-demonstration.gif") 31 | ) 32 | 33 | m.normalize() 34 | # m.align_wavelengths().wavelength 35 | # m[:, :] 36 | plt.close("all") 37 | 38 | 39 | def test_compare_wrappers(): 40 | rainbows = [SimulatedRainbow(R=10 ** np.random.uniform(0.1, 2)) for _ in range(3)] 41 | 42 | a = compare_rainbows(rainbows) 43 | b = rainbows[0].compare(rainbows) 44 | 45 | assert a.names == b.names 46 | assert a.rainbows == b.rainbows 47 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/rainbow_FITS.py: -------------------------------------------------------------------------------- 1 | """ 2 | Define a reader for chromatic .rainbow.FITS files. 3 | """ 4 | 5 | # import the general list of packages 6 | from ...imports import * 7 | 8 | # define list of the only things that will show up in imports 9 | __all__ = ["from_rainbow_FITS"] 10 | 11 | 12 | def from_rainbow_FITS(rainbow, filepath): 13 | """ 14 | Populate a Rainbow from a file in the rainbow_FITS format. 15 | 16 | Parameters 17 | ---------- 18 | 19 | rainbow : Rainbow 20 | The object to be populated. 21 | 22 | filepath : str 23 | The path to the file to load. 24 | """ 25 | 26 | # open the FITS file 27 | hdu_list = fits.open(filepath) 28 | 29 | # load the header into metadata 30 | h = hdu_list["primary"].header 31 | for k in h: 32 | if k.lower() in [ 33 | "name", 34 | "wscale", 35 | "tscale", 36 | "signal_to_noise", 37 | "history", 38 | ]: # KLUDGE! 39 | rainbow.metadata[k.lower()] = h[k] 40 | elif k not in ["SIMPLE", "BITPIX", "NAXIS", "EXTEND"]: 41 | rainbow.metadata[k] = h[k] 42 | 43 | # load into the three core dictionaries 44 | for e in ["fluxlike", "wavelike", "timelike"]: 45 | table = Table.read(hdu_list[e]) 46 | for k in table.colnames: 47 | vars(rainbow)[e][k] = table[k].quantity * 1 48 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/shift.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["shift"] 4 | 5 | 6 | def shift(self, velocity=0 * u.km / u.s): 7 | """ 8 | Doppler shift the wavelengths of this `Rainbow`. 9 | 10 | This shifts the wavelengths in a `Rainbow` by 11 | applying a velocity shift. Positive velocities make 12 | wavelengths longer (redshift); negative velocities make 13 | wavelengths shorter (bluesfhit). 14 | 15 | Parameters 16 | ---------- 17 | velocity : Quantity 18 | The systemic velocity by which we should shift, 19 | with units of velocity (for example, u.km/u.s) 20 | """ 21 | 22 | # create a history entry for this action (before other variables are defined) 23 | h = self._create_history_entry("shift", locals()) 24 | 25 | # create a new copy of this rainbow 26 | new = self._create_copy() 27 | 28 | # get the speed of light from astropy constants 29 | lightspeed = con.c.to("km/s") # speed of light in km/s 30 | 31 | # calculate beta and make sure the units cancel 32 | beta = (velocity / lightspeed).decompose() 33 | 34 | # apply wavelength shift 35 | new_wavelength = new.wavelength * np.sqrt((1 + beta) / (1 - beta)) 36 | new.wavelike["wavelength"] = new_wavelength 37 | 38 | # append the history entry to the new Rainbow 39 | new._record_history_entry(h) 40 | 41 | # return the new object 42 | return new 43 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/radica.py: -------------------------------------------------------------------------------- 1 | # import the general list of packages 2 | from ...imports import * 3 | 4 | # define list of the only things that will show up in imports 5 | __all__ = ["from_radica"] 6 | 7 | 8 | def from_radica(self, filepath, order=1): 9 | """ 10 | Populate a Rainbow from a file in the radica format. 11 | 12 | Parameters 13 | ---------- 14 | 15 | self : Rainbow 16 | The object to be populated. 17 | filepath : str 18 | The path to the file to load. 19 | """ 20 | 21 | hdu = fits.open(filepath) 22 | 23 | n_orders = 2 24 | cheerfully_suggest( 25 | f""" 26 | Loading NIRISS spectroscopic `order={order}``. Two orders are available, 27 | and you can set which (1,2) you want to read with the `order=` option. 28 | """ 29 | ) 30 | self.fluxlike["wavelength_2d"] = ( 31 | np.transpose(hdu[f"Wave 2D Order {order}"].data * u.micron) * 1 32 | ) 33 | self.wavelike["wavelength"] = ( 34 | np.nanmedian(self.fluxlike["wavelength_2d"], axis=1) * 1 35 | ) 36 | self.fluxlike["flux"] = ( 37 | np.transpose(hdu[f"Flux Order {order}"].data * u.electron) * 1 38 | ) 39 | self.fluxlike["uncertainty"] = ( 40 | np.transpose(hdu[f"Flux Error Order {order}"].data * u.electron) * 1 41 | ) 42 | astropy_times = Time(hdu["Time"].data, format="jd", scale="tdb") 43 | self.set_times_from_astropy(astropy_times, is_barycentric=True) 44 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/descriptions.txt: -------------------------------------------------------------------------------- 1 | cartoon | name | description 2 | 🌈🧮📝 | +-*/ | Do basic math operations with two Rainbows. 3 | 🌈🗂🔪 | [:,:] | Index, slice, or mask a Rainbow to get a subset. 4 | 🌈🚧🌊 | align_wavelengths | Align spectra with different wavelength grids onto one shared axis. 5 | 🌈🧺🧱 | bin | Bin to a new wavelength or time grid. 6 | 🌈🧑‍🤝‍🧑🌈 | compare | Connect to other 🌈 objects for easy comparison. 7 | 🌈🐈⏰ | concatenate_in_time | Stitch together two Rainbows with identical wavelengths. 8 | 🌈🐈🌊 | concatenate_in_wavelength | Stitch together two Rainbows with identical times. 9 | 🌈🚩👀 | flag_outliers | Flag outlier data points. 10 | 🌈⏲🎞 | fold | Fold times relative to a particular period and epoch. 11 | 🌈🧺⏰ | get_average_lightcurve_as_rainbow | Bin down to a single integrated light curve. 12 | 🌈🧺🌊 | get_average_spectrum_as_rainbow | Bin down to a single integrated spectrum. 13 | 🌈🎧🎲 | inject_noise| Inject (uncorrelated, simple) random noise. 14 | 🌈🎧🎹 | inject_systematics| Inject (correlated, wobbly) systematic noise. 15 | 🌈⭐️👻 | inject_spectrum | Inject a static stellar spectrum. 16 | 🌈🪐🚞 | inject_transit | Inject a transit signal. 17 | 🌈🫓😑 | normalize | Normalize by dividing through by a typical spectrum (and/or light curve). 18 | 🌈🏴‍☠️🛁 | remove_trends | Remove smooth trends in time and/or wavelength. 19 | 🌈🚇🌊 | shift | Doppler shift wavelengths. 20 | 🌈🍱💇 | trim | Trim away wavelengths or times. 21 | -------------------------------------------------------------------------------- /notes/how-to-do-pre-commit.md: -------------------------------------------------------------------------------- 1 | # how to set up pre-commit 2 | 3 | [pre-commit](https://pre-commit.com) seems to be a neat little tool for making sure some tools get run every time someone tries to commit to the repository. Here are some quick notes from Zach trying to get it set up. 4 | 5 | I followed the instructions at https://pre-commit.com, including running `pre-commit sample-config` to get a sample start to the config file. I copied and pasted this into `.pre-commit-config.yaml`. I made some modifications based on the example at [exoplanet](https://github.com/exoplanet-dev/exoplanet/blob/main/.pre-commit-config.yaml) to add `black`, apparently both for the code and for the documentation notebooks. 6 | 7 | I ran `pre-commit autoupdate` to assign the current latest versions of the tools. It sounds like these won't naturally update, so if there are changes to them, we should rerun `pre-commit autoupdate` at some point to update the version numbers. 8 | 9 | I ran `pre-commit install` to connect it to the repository. I iterated `autoupdate` and `install` a few times to get rid of some warnings about mutable versions. 10 | 11 | I tried to run `pre-commit run --all-files` to run it on everything that had already been committed. We should redo this if we change the pre-commit rules at any point. The first time I ran it there were some complaints because by the style-checker thingamabobs, but then I put `black` before them, reran, and everything passed happily. 12 | 13 | Seems to work!? 14 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/schlawin.py: -------------------------------------------------------------------------------- 1 | # import the general list of packages 2 | from ...imports import * 3 | 4 | # define list of the only things that will show up in imports 5 | __all__ = ["from_schlawin"] 6 | 7 | 8 | def from_schlawin(rainbow, filepath): 9 | """ 10 | Populate a Rainbow from a file in Everett Schlawin's tshirt format. 11 | 12 | Parameters 13 | ---------- 14 | 15 | rainbow : Rainbow 16 | The object to be populated. 17 | filepath : str 18 | The path to the file to load. 19 | """ 20 | 21 | # read in your file, however you like 22 | hdu_list = fits.open(filepath) 23 | 24 | # populate a 1D array of wavelengths (with astropy units of length) 25 | rainbow.wavelike["indices"] = hdu_list["disp indices"].data * 1 26 | rainbow.wavelike["wavelength"] = hdu_list["wavelength"].data * u.micron * 1 27 | 28 | # populate a 1D array of times (with astropy units of time) 29 | rainbow.timelike["time"] = hdu_list["time"].data * u.day * 1 30 | 31 | # populate a 2D (row = wavelength, col = array of fluxes 32 | for k in [ 33 | "optimal spec", 34 | "opt spec err", 35 | "sum spec", 36 | "sum spec err", 37 | "background spec", 38 | "refpix", 39 | ]: 40 | rainbow.fluxlike[k] = hdu_list[k].data.T.squeeze() * 1 41 | rainbow.fluxlike["flux"] = rainbow.fluxlike["optimal spec"] * 1 42 | rainbow.fluxlike["uncertainty"] = rainbow.fluxlike["opt spec err"] * 1 43 | 44 | # kludgily pull all extension headers into the metadata 45 | for k in hdu_list: 46 | rainbow.metadata[f"{k}-header"] = hdu_list[k].header 47 | -------------------------------------------------------------------------------- /chromatic/rainbows/writers/rainbow_FITS.py: -------------------------------------------------------------------------------- 1 | """ 2 | Define a writer for chromatic .rainbow.FITS files. 3 | """ 4 | 5 | # import the general list of packages 6 | from ...imports import * 7 | 8 | # define list of the only things that will show up in imports 9 | __all__ = ["to_rainbow_FITS"] 10 | 11 | 12 | def to_rainbow_FITS(self, filepath, overwrite=True): 13 | """ 14 | Write a Rainbow to a FITS file. 15 | 16 | Parameters 17 | ---------- 18 | 19 | self : Rainbow 20 | The object to be saved. 21 | 22 | filepath : str 23 | The path to the file to write. 24 | """ 25 | 26 | # create a header for the metadata 27 | header = fits.Header() 28 | header["comment"] = "Data are stored in three main extensions" 29 | header["comment"] = " [1]=['FLUXLIKE'] contains (nwave, ntime)-shaped arrays" 30 | header["comment"] = " [2]=['WAVELIKE'] contains (nwave)-shaped arrays" 31 | header["comment"] = " [3]=['TIMELIKE'] contains (ntime)-shaped arrays" 32 | header["comment"] = "The primary extension header contains some metadata." 33 | 34 | for k in self.metadata: 35 | try: 36 | header[k] = self.metadata[k] 37 | except ValueError: 38 | cheerfully_suggest(f"metadata item '{k}' cannot be saved to FITS header") 39 | 40 | primary_hdu = fits.PrimaryHDU(header=header) 41 | 42 | # create extensions for the three other core dictionaries 43 | flux_hdu = fits.BinTableHDU(Table(self.fluxlike), name="fluxlike") 44 | wave_hdu = fits.BinTableHDU(Table(self.wavelike), name="wavelike") 45 | time_hdu = fits.BinTableHDU(Table(self.timelike), name="timelike") 46 | 47 | hdu_list = fits.HDUList([primary_hdu, flux_hdu, wave_hdu, time_hdu]) 48 | hdu_list.writeto(filepath, overwrite=overwrite) 49 | -------------------------------------------------------------------------------- /chromatic/rainbows/helpers/help.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["help"] 4 | 5 | 6 | def help(self): 7 | """ 8 | Print a quick reference of key actions available for this `Rainbow`. 9 | """ 10 | print( 11 | textwrap.dedent( 12 | """ 13 | Hooray for you! You asked for help on what you can do 14 | with this 🌈 object. Here's a quick reference of a few 15 | available options for things to try.""" 16 | ) 17 | ) 18 | 19 | base_directory = pkg_resources.resource_filename("chromatic", "rainbows") 20 | descriptions_files = [] 21 | for level in ["*", "*/*"]: 22 | descriptions_files += glob.glob( 23 | os.path.join(base_directory, level, "descriptions.txt") 24 | ) 25 | categories = [ 26 | d.replace(base_directory + "/", "").replace("/descriptions.txt", "") 27 | for d in descriptions_files 28 | ] 29 | for i in np.argsort(categories): 30 | c, d = categories[i], descriptions_files[i] 31 | header = ( 32 | "\n" + "-" * (len(c) + 4) + "\n" + f"| {c} |\n" + "-" * (len(c) + 4) + "\n" 33 | ) 34 | 35 | table = ascii.read(d) 36 | items = [] 37 | for row in table: 38 | name = row["name"] 39 | if hasattr(self, name) or (name in ["+-*/", "[:,:]"]): 40 | if name in "+-*/": 41 | function_call = f"{name}" 42 | else: 43 | function_call = f".{name}()" 44 | 45 | item = ( 46 | f"{row['cartoon']} | {function_call:<28} \n {row['description']}" 47 | ) 48 | items.append(item) 49 | if len(items) > 0: 50 | print(header) 51 | print("\n".join(items)) 52 | -------------------------------------------------------------------------------- /chromatic/tests/test_shift.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_shift(): 6 | N = 1000 7 | w = np.linspace(0.6, 0.7, N) * u.micron 8 | f = u.Quantity(np.zeros(N)) 9 | for i in range(3): 10 | w0 = np.random.uniform(0.6, 0.7) * u.micron 11 | sigma = np.random.uniform(0.001, 0.01) * u.micron 12 | f += np.exp(-0.5 * ((w - w0) / sigma) ** 2) 13 | 14 | for v in [0, -1e4, 1e4] * u.km / u.s: 15 | print(f"Testing for v={v}") 16 | unshifted = SimulatedRainbow(wavelength=w, star_flux=f).inject_noise() 17 | shifted = unshifted.shift(v) 18 | shifted_and_then_shifted_back = shifted.shift(-v) 19 | 20 | fi, ax = plt.subplots( 21 | 1, 3, sharex=True, sharey=True, figsize=(8, 4), constrained_layout=True 22 | ) 23 | unshifted.paint(ax=ax[0]) 24 | shifted.paint(ax=ax[1]) 25 | shifted_and_then_shifted_back.paint(ax=ax[2]) 26 | plt.savefig( 27 | os.path.join( 28 | test_directory, 29 | "demonstration-of-shifting-wavelengths.pdf", 30 | ) 31 | ) 32 | 33 | assert np.all( 34 | np.isclose(unshifted.wavelength, shifted_and_then_shifted_back.wavelength) 35 | ) 36 | 37 | if v.to_value("km/s") > 0: 38 | print(f"The velocity {v} should cause a redshift.") 39 | assert np.all(shifted.wavelength > unshifted.wavelength) 40 | elif v.to_value("km/s") == 0: 41 | print(f"The velocity {v} should cause no shift.") 42 | print(" It does!") 43 | elif v.to_value("km/s") < 0: 44 | print(f"The velocity {v} should cause a blueshift.") 45 | assert np.all(shifted.wavelength < unshifted.wavelength) 46 | print() 47 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/eureka_channels.py: -------------------------------------------------------------------------------- 1 | # import the general list of packages 2 | from ...imports import * 3 | 4 | # define list of the only things that will show up in imports 5 | __all__ = ["from_eureka_channels", "from_eureka_S5"] 6 | 7 | 8 | def from_eureka_channels(rainbow, filepath): 9 | """ 10 | Populate a Rainbow from Eureka! S5 modeling outputs, 11 | which are wavelength-binned light curves with models. 12 | 13 | Parameters 14 | ---------- 15 | 16 | rainbow : Rainbow 17 | The object to be populated. 18 | 19 | filepath : str 20 | The path to the files to load. 21 | 22 | """ 23 | # get list of files 24 | filenames = expand_filenames(filepath) 25 | 26 | # load them into tables 27 | tables = [ascii.read(f) for f in filenames] 28 | 29 | # load the time and wavelength arrays 30 | rainbow.timelike["time"] = tables[0]["time"] * u.day 31 | rainbow.wavelike["wavelength"] = [t["wavelength"][0] for t in tables] * u.micron 32 | rainbow.wavelike["wavelength_lower"] = [ 33 | (t["wavelength"] - t["bin_width"])[0] for t in tables 34 | ] * u.micron 35 | rainbow.wavelike["wavelength_upper"] = [ 36 | (t["wavelength"] + t["bin_width"])[0] for t in tables 37 | ] * u.micron 38 | keys_used = ["time", "wavelength", "bin_width"] 39 | 40 | # populate the fluxlike arrays 41 | for k in ["lcdata", "lcerr", "model", "residuals"]: 42 | rainbow.fluxlike[k] = np.array([t[k] for t in tables]) 43 | keys_used.append(k) 44 | 45 | for k in tables[0].colnames: 46 | if k != keys_used: 47 | rainbow.fluxlike[k] = np.array([t[k] for t in tables]) 48 | 49 | rainbow.fluxlike["flux"] = rainbow.lcdata * 1 50 | rainbow.fluxlike["uncertainty"] = rainbow.lcerr * 1 51 | 52 | 53 | from_eureka_S5 = from_eureka_channels 54 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/espinoza.py: -------------------------------------------------------------------------------- 1 | """ 2 | Define a reader for Néstor's NIRISS extracted spectra .npy files. 3 | 4 | This was kludged together to read the Stage 2 outputs from 5 | https://stsci.app.box.com/s/tyg3qqd85601gkbw5koowrx0obekeg0m/folder/154382588636 6 | """ 7 | 8 | # import the general list of packages 9 | from ...imports import * 10 | 11 | # define list of the only things that will show up in imports 12 | __all__ = ["from_espinoza"] 13 | 14 | 15 | def from_espinoza(rainbow, filepath): 16 | """ 17 | Populate a Rainbow from a file in the espinoza format. 18 | 19 | Parameters 20 | ---------- 21 | 22 | rainbow : Rainbow 23 | The object to be populated. 24 | filepath : str 25 | The path to the file to load. It should be a glob-friendly 26 | string like `.../*_order1.npy` that points to both a 27 | `spectra_order1.npy` and a `wavelengths_order1.npy` file. 28 | """ 29 | 30 | # read in two files, one for the flux, one for the wavelength 31 | spectra_filename, wavelengths_filename = sorted(glob.glob(filepath)) 32 | assert ("spectra") in spectra_filename 33 | assert ("wavelength") in wavelengths_filename 34 | spectra = np.load(spectra_filename) 35 | wavelengths = np.load(wavelengths_filename) 36 | 37 | # populate a 1D array of wavelengths (with astropy units of length) 38 | rainbow.wavelike["wavelength"] = wavelengths * u.micron * 1 39 | 40 | # populate a 1D array of times (with astropy units of time) 41 | times = np.arange(spectra.shape[0]) * u.minute 42 | cheerfully_suggest("The times are totally made up!") 43 | rainbow.timelike["time"] = times * 1 44 | 45 | # populate a 2D (row = wavelength, col = array of fluxes 46 | rainbow.fluxlike["flux"] = spectra[:, 0, :].T * 1 47 | rainbow.fluxlike["uncertainty"] = spectra[:, 1, :].T * 1 48 | -------------------------------------------------------------------------------- /chromatic/tests/test_remove_trends.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_remove_trends(): 6 | s = ( 7 | SimulatedRainbow(dw=0.1 * u.micron) 8 | .inject_transit() 9 | .inject_systematics() 10 | .inject_noise(signal_to_noise=300) 11 | ) 12 | fi, ax = plt.subplots( 13 | 2, 14 | 6, 15 | figsize=(12, 5), 16 | sharey="row", 17 | sharex=True, 18 | dpi=300, 19 | constrained_layout=True, 20 | ) 21 | imkw = dict(vmin=0.98, vmax=1.02, xaxis="wavelength", colorbar=False) 22 | s.paint(ax=ax[0, 0], **imkw) 23 | plt.title("raw data") 24 | s.plot_noise_comparison(ax=ax[1, 0]) 25 | for i, method in enumerate( 26 | ["median_filter", "savgol_filter", "differences", "polyfit"] 27 | ): 28 | 29 | with warnings.catch_warnings(): 30 | warnings.simplefilter("ignore") 31 | x = s.remove_trends(method=method) 32 | x.paint(ax=ax[0, 1 + i], **imkw) 33 | ax[0, i + 1].set_title(method) 34 | x.plot_noise_comparison(ax=ax[1, i + 1]) 35 | 36 | x = s.remove_trends(method="custom", model=s.planet_model) 37 | ax[0, -1].set_title("custom") 38 | 39 | x.paint(ax=ax[0, -1], **imkw) 40 | x.plot_noise_comparison(ax=ax[1, -1]) 41 | plt.ylim(0, 0.02) 42 | plt.title("actual systematics") 43 | plt.savefig( 44 | os.path.join( 45 | test_directory, 46 | "demonstration-of-removing-trends-with-different-methods.pdf", 47 | ) 48 | ) 49 | 50 | 51 | def test_remove_trends_options(): 52 | s = SimulatedRainbow().inject_noise() 53 | for method in ["median_filter", "savgol_filter", "polyfit"]: 54 | with pytest.warns(match="didn't supply all expected keywords for"): 55 | s.remove_trends(method=method) 56 | -------------------------------------------------------------------------------- /chromatic/noise/jwst/simulate.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | from ...rainbows import SimulatedRainbow 3 | 4 | 5 | def make_simulated_jwst_rainbow(table, dt=None, tlim=None, which_snr_to_use=None): 6 | """ 7 | Make a simulated Rainbow object with noise and cadence set by PandExo. 8 | 9 | Parameters 10 | ---------- 11 | table : Table 12 | The 1D tabular output from either PandExo or ETC. 13 | dt : Quantity 14 | The time cadence, with units of time. 15 | (default of None will get integration time from PandExo) 16 | tlim : Quantity 17 | The time span of the observation [tmin, tmax], with units of time. 18 | (default of None will set to two transit durations) 19 | which_snr_to_use : str 20 | Which choice of `snr_per_integration_...` should be used? 21 | """ 22 | 23 | # set the wavelengths 24 | wavelength = table["wavelength"] * u.micron 25 | 26 | # set the cadence (or use from PandExo) 27 | if dt is None: 28 | dt = table.meta["time_per_integration"] * u.s 29 | 30 | # set the duration of the observation 31 | if tlim is None: 32 | transit_duration = table.meta["transit_duration"] * u.hour 33 | tlim = transit_duration * [-1, 1] 34 | 35 | # set the S/N per integration per pixel 36 | if which_snr_to_use is None: 37 | snr = np.inf 38 | for k in table.colnames: 39 | if "snr_per_integration" in k: 40 | this_snr = np.median(table[k]) 41 | if this_snr < snr: 42 | snr = this_snr 43 | which_snr_to_use = k 44 | print( 45 | f"Using {k} because its median S/N ({np.median(snr):.0f}) is most cautious." 46 | ) 47 | 48 | # create the simulated rainbow 49 | s = SimulatedRainbow(tlim=tlim, dt=dt, wavelength=wavelength).inject_noise( 50 | signal_to_noise=table[which_snr_to_use] 51 | ) 52 | return s 53 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/inject_outliers.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["inject_outliers"] 4 | 5 | 6 | def inject_outliers(self, fraction=0.01, amplitude=10): 7 | """ 8 | Inject some random outliers. 9 | 10 | To approximately simulate cosmic rays or other 11 | rare weird outliers, this randomly injects 12 | outliers into a small fraction of pixels. For 13 | this simple method, outliers will have the same 14 | amplitude, either as a ratio above the per-data-point 15 | or as a fixed number (if no uncertainties exist). 16 | 17 | Parameters 18 | ---------- 19 | fraction : float, optional 20 | The fraction of pixels that should get outliers. 21 | (default = 0.01) 22 | amplitude : float, optional 23 | If uncertainty > 0, how many sigma should outliers be? 24 | If uncertainty = 0, what number should be injected? 25 | (default = 10) 26 | 27 | Returns 28 | ------- 29 | rainbow : Rainbow 30 | A new `Rainbow` object with outliers injected. 31 | """ 32 | 33 | # create a history entry for this action (before other variables are defined) 34 | h = self._create_history_entry("inject_outliers", locals()) 35 | 36 | # create a copy of the existing Rainbow 37 | new = self._create_copy() 38 | 39 | # pick some random pixels to inject outliers 40 | outliers = np.random.uniform(0, 1, self.shape) < fraction 41 | 42 | # inject outliers based on uncertainty if possible 43 | if np.any(self.uncertainty > 0): 44 | new.fluxlike["injected_outliers"] = outliers * amplitude * self.uncertainty 45 | else: 46 | new.fluxlike["injected_outliers"] = outliers * amplitude 47 | 48 | # modify the flux 49 | new.fluxlike["flux"] += new.fluxlike["injected_outliers"] 50 | 51 | # append the history entry to the new Rainbow 52 | new._record_history_entry(h) 53 | 54 | # return the new object 55 | return new 56 | -------------------------------------------------------------------------------- /chromatic/rainbows/writers/text.py: -------------------------------------------------------------------------------- 1 | """ 2 | Define a writer for chromatic .rainbow.npy files. 3 | """ 4 | 5 | # import the general list of packages 6 | from ...imports import * 7 | 8 | # define list of the only things that will show up in imports 9 | __all__ = ["to_text"] 10 | 11 | 12 | def to_text(self, filepath, overwrite=True, group_by="wavelength"): 13 | """ 14 | Write a Rainbow to a file in the text format. 15 | 16 | Parameters 17 | ---------- 18 | 19 | self : Rainbow 20 | The object to be saved. 21 | 22 | filepath : str 23 | The path to the file to write. 24 | """ 25 | 26 | # a 1D array of wavelengths (with astropy units of length) 27 | the_1D_array_of_wavelengths = self.wavelike["wavelength"] 28 | 29 | # a 1D array of times (with astropy units of time) 30 | the_1D_array_of_times = self.timelike["time"] 31 | 32 | # a 2D (row = wavelength, col = array of fluxes 33 | the_2D_array_of_fluxes = self.fluxlike["flux"] 34 | 35 | # write out your file, however you like 36 | w, t = np.meshgrid(self.wavelength.to("micron"), self.time.to("day"), indexing="ij") 37 | 38 | def make_into_columns(x): 39 | if group_by == "wavelength": 40 | return x.flatten() 41 | if group_by == "time": 42 | return x.T.flatten() 43 | 44 | # create a table 45 | table = Table( 46 | dict(wavelength=make_into_columns(w), time=make_into_columns(t)), 47 | meta=self.metadata, 48 | ) 49 | table["flux"] = make_into_columns(self.flux) 50 | try: 51 | table["uncertainty"] = make_into_columns(self.uncertainty) 52 | except KeyError: 53 | pass 54 | 55 | other_keys = list(self.fluxlike.keys()) 56 | other_keys.remove("flux") 57 | other_keys.remove("uncertainty") 58 | 59 | for k in other_keys: 60 | table[k] = make_into_columns(self.fluxlike[k]) 61 | 62 | table.write(filepath, format="ascii.ecsv", overwrite=overwrite) 63 | -------------------------------------------------------------------------------- /mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: chromatic 2 | site_url: https://zkbt.github.com/chromatic 3 | nav: 4 | - Welcome: index.md 5 | - installation.ipynb 6 | - quickstart.ipynb 7 | - User Guide: 8 | - basics.ipynb 9 | - io.ipynb 10 | - creating.ipynb 11 | - actions.ipynb 12 | - visualizing.ipynb 13 | - models.ipynb 14 | - Examples: 15 | - example-timeseries-spectra.ipynb 16 | - Developer Guide: 17 | - designing.ipynb 18 | - documentation.ipynb 19 | - github.ipynb 20 | - Related Tools: 21 | - tools/spectra.ipynb 22 | - tools/binning.ipynb 23 | - tools/colormaps.ipynb 24 | - tools/transits.ipynb 25 | - api.md 26 | theme: 27 | name: "material" 28 | features: 29 | - navigation.tracking 30 | repo_url: https://github.com/zkbt/chromatic/ 31 | plugins: 32 | - search 33 | - mkdocs-jupyter: 34 | execute : True 35 | include_source : True 36 | - mkdocstrings: 37 | default_handler: python 38 | handlers: 39 | python: 40 | paths: [../chromatic] 41 | selection: 42 | docstring_style: "numpy" 43 | rendering: 44 | show_source: True 45 | show_root_heading: True 46 | show_root_toc_entry: False 47 | show_root_full_path: False 48 | show_category_heading: False 49 | show_submodules: False 50 | merge_init_into_class: False 51 | show_if_no_docstring: False 52 | heading_level: 3 53 | show_bases: False 54 | - exclude: 55 | glob: 56 | - "example-datasets/*" 57 | - "downloads-for-exoatlas/*" 58 | - "*.pdf" 59 | - "*.fits" 60 | - "*.npy" 61 | markdown_extensions: 62 | - toc: 63 | permalink: "#" 64 | 65 | # this is super borrowed from Christina Hedges' fabulous 66 | # https://christinahedges.github.io/astronomy_workflow/ 67 | -------------------------------------------------------------------------------- /notes/how-to-do-the-docs.md: -------------------------------------------------------------------------------- 1 | # how to use mkdocs 2 | 3 | Zach is new to mkdocs, but it seems easier than sphinx, so he's keen to try it out. Trying to follow christinahedges's very useful tutorial [here](https://christinahedges.github.io/astronomy_workflow/notebooks/3.0-building/mkdocs.html), here's what we're doing: 4 | 5 | First, `poetry` felt a little scary for one day, so I'm trying to do this just with pip installing the things I need: 6 | `pip install mkdocs mkdocs-material mkdocstrings 'pytkdocs[numpy-style]' mkdocs-jupyter` 7 | 8 | Then, we set up the docs by going into the base directory for this repository, and running 9 | `mkdocs new .` 10 | which made a `docs/` directory and a `mkdocs.yml` 11 | 12 | Then, we copied her template into the `docs/index.md` file and made appropriate changes for names. 13 | 14 | Then, we edited the `docs/api.md` file to point to the objects and methods we want to explain. Docstrings for functions should follow the [`numpy` style conventions](https://numpydoc.readthedocs.io/en/latest/format.html). Options for rendering automatically are best described [here](https://mkdocstrings.github.io/python/usage/), although I don't fully understand everything going on in the YAML. 15 | 16 | Then, we used the `mkdocs-jupyter` plugin to be able use jupyter notebooks as the source for writing docs, following the examples on their pages. We added some notebooks to the `docs/` diretory and pointed to them in the `mkdocs.yml` file. 17 | 18 | Then, we ran `mkdocs serve`, and woah, a live version of the docs appeared at http://127.0.0.1:8000/. It was particularly cool (and way better than `sphinx` that I could make a change to any of the files and simply reload the page to see them update live into the docs). Hooray! 19 | 20 | Then, we ran `mkdocs gh-deploy`, and double woah, it deployed a pretty version of the docs up at zkbt.github.io/chromatic! For the sake of not making the deployment `gh-pages` branch annoyingly large, add the `--no-history` option to erase the repository each time. 21 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/measured_scatter.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_measured_scatter"] 4 | 5 | 6 | def get_measured_scatter( 7 | self, quantity="flux", method="standard-deviation", minimum_acceptable_ok=1e-10 8 | ): 9 | """ 10 | Get measured scatter for each wavelength. 11 | 12 | Calculate the standard deviation (or outlier-robust 13 | equivalent) for each wavelength, which can be compared 14 | to the expected per-wavelength uncertainty. 15 | 16 | Parameters 17 | ---------- 18 | quantity : string, optional 19 | The `fluxlike` quantity for which we should calculate the scatter. 20 | method : string, optional 21 | What method to use to obtain measured scatter. 22 | Current options are 'MAD', 'standard-deviation'. 23 | minimum_acceptable_ok : float, optional 24 | The smallest value of `ok` that will still be included. 25 | (1 for perfect data, 1e-10 for everything but terrible data, 0 for all data) 26 | 27 | Returns 28 | ------- 29 | scatter : array 30 | Wavelike array of measured scatters. 31 | """ 32 | 33 | if method not in ["standard-deviation", "MAD"]: 34 | cheerfully_suggest( 35 | f""" 36 | '{method}' is not an available method. 37 | Please choose from ['MAD', 'standard-deviation']. 38 | """ 39 | ) 40 | with warnings.catch_warnings(): 41 | warnings.simplefilter("ignore") 42 | 43 | scatters = np.zeros(self.nwave) 44 | for i in range(self.nwave): 45 | x, y, sigma = self.get_ok_data_for_wavelength( 46 | i, y=quantity, minimum_acceptable_ok=minimum_acceptable_ok 47 | ) 48 | if u.Quantity(y).unit == u.Unit(""): 49 | y_value, y_unit = y, 1 50 | else: 51 | y_value, y_unit = y.value, y.unit 52 | if method == "standard-deviation": 53 | scatters[i] = np.nanstd(y_value) 54 | elif method == "MAD": 55 | scatters[i] = mad_std(y_value, ignore_nan=True) 56 | return scatters * y_unit 57 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/feinstein.py: -------------------------------------------------------------------------------- 1 | # import the general list of packages 2 | from ...imports import * 3 | 4 | # define list of the only things that will show up in imports 5 | __all__ = ["from_feinstein_numpy", "from_feinstein_h5"] 6 | 7 | 8 | def from_feinstein_numpy(rainbow, filepath): 9 | """ 10 | Populate a Rainbow from a file in the feinstein numpy format. 11 | 12 | Parameters 13 | ---------- 14 | 15 | rainbow : Rainbow 16 | The object to be populated. 17 | filepath : str 18 | The path to the file to load. 19 | """ 20 | 21 | # time, wavelength, spectra, err = np.load(filepath, allow_pickle=True) 22 | dat = np.load(filepath, allow_pickle=True).item() 23 | 24 | rainbow.wavelike["wavelength"] = dat["wavelength"] * u.micron * 1 25 | 26 | # populate a 1D array of times (with astropy units of time) 27 | times = dat["time"] * u.day 28 | rainbow.timelike["time"] = times * 1 29 | 30 | # populate a 2D (row = wavelength, col = array of fluxes 31 | flux = np.zeros((len(dat["wavelength"]), len(times))) 32 | uncertainty = np.zeros_like(dat["flux"]) 33 | 34 | rainbow.fluxlike["flux"] = dat["flux"].T * 1 35 | rainbow.fluxlike["uncertainty"] = dat["flux_err"].T * 1 36 | 37 | 38 | def from_feinstein_h5(self, filepath, order=1, version="opt"): 39 | """ 40 | Populate a Rainbow from a file in the feinstein H5 format. 41 | 42 | Parameters 43 | ---------- 44 | 45 | rainbow : Rainbow 46 | The object to be populated. 47 | filepath : str 48 | The path to the file to load. 49 | """ 50 | 51 | f = h5.File(filepath) 52 | 53 | astropy_times = Time(np.array(f["time"]), format="mjd", scale="tdb") 54 | self.set_times_from_astropy(astropy_times, is_barycentric=True) # ??? 55 | self.wavelike["wavelength"] = ( 56 | np.array(f[f"wavelength_order_{order}"]) * u.micron * 1 57 | ) 58 | for k in ["box_flux", "box_var", "opt_flux", "opt_var"]: 59 | self.fluxlike[k] = np.transpose(np.array(f[f"{k}_order_{order}"]) * 1) 60 | 61 | self.fluxlike["flux"] = self.fluxlike[f"{version}_flux"] * 1 62 | self.fluxlike["uncertainty"] = self.fluxlike[f"{version}_var"] * 1 63 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/text.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["from_text"] 4 | 5 | 6 | def from_text(rainbow, filename, **kwargs): 7 | """ 8 | Populate a Rainbow from a generic text file. 9 | This text file must at least contain the columns 10 | `wavelength`, `time`, `flux`, `uncertainty` 11 | 12 | Parameters 13 | ---------- 14 | 15 | rainbow : Rainbow 16 | The object to be populated. (This is intended 17 | to be used as a class method of Rainbow or 18 | of a class derived from Rainbow, as a way of 19 | initializing an object from files.) 20 | filename : str 21 | The path to the file. 22 | """ 23 | 24 | # load the data 25 | data = ascii.read(filename) 26 | 27 | # pull out some variables 28 | t = sorted(np.unique(data["time"])) 29 | w = sorted(np.unique(data["wavelength"])) 30 | 31 | fluxlike = {} 32 | fluxlike_keys = list(data.colnames) 33 | fluxlike_keys.remove("time") 34 | fluxlike_keys.remove("wavelength") 35 | for k in fluxlike_keys: 36 | fluxlike[k] = np.ones(shape=(len(w), len(t))) 37 | for i in tqdm(range(len(data)), leave=False): 38 | # this is slow, but general 39 | i_time = t == data[i]["time"] 40 | i_wavelength = w == data[i]["wavelength"] 41 | for k in fluxlike_keys: 42 | fluxlike[k][i_wavelength, i_time] = data[i][k] 43 | 44 | # do a slightly better guess of the uncertainty column 45 | if "uncertainty" not in fluxlike: 46 | for k in ["error", "flux_error", "sigma", "unc"]: 47 | try: 48 | fluxlike["uncertainty"] = fluxlike[k] * 1 49 | break 50 | except KeyError: 51 | pass 52 | 53 | timelike = {} 54 | try: 55 | t.unit 56 | except AttributeError: 57 | t *= u.day 58 | timelike["time"] = t 59 | 60 | wavelike = {} 61 | try: 62 | w.unit 63 | except AttributeError: 64 | w *= u.micron 65 | wavelike["wavelength"] = w 66 | 67 | # populate the rainbow 68 | rainbow._initialize_from_dictionaries( 69 | wavelike=wavelike, timelike=timelike, fluxlike=fluxlike 70 | ) 71 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/attach_model.py: -------------------------------------------------------------------------------- 1 | """ 2 | Tools for attaching models? 3 | """ 4 | 5 | from ...imports import * 6 | 7 | __all__ = ["attach_model"] 8 | 9 | 10 | def attach_model(self, model, **kw): 11 | """ 12 | Attach a `fluxlike` model, thus making a new `RainbowWithModel.` 13 | 14 | Having a model attached makes it possible to make calculations 15 | (residuals, chi^2) and visualizations comparing data to model. 16 | 17 | The `model` array will be stored in `.fluxlike['model']`. 18 | After running this to make a `RainbowWithModel` it's OK 19 | (and faster) to simply update `.fluxlike['model']` or `.model`. 20 | 21 | Parameters 22 | ---------- 23 | model : array, Quantity 24 | An array of model values, with the same shape as 'flux' 25 | **kw : dict, optional 26 | All other keywords will be interpreted as items 27 | that can be added to a `Rainbow`. You might use this 28 | to attach intermediate model steps or quantities. 29 | Variable names ending with `_model` can be particularly 30 | easily incorporated into multi-part model visualizations 31 | (for example, `'planet_model'` or `'systematics_model'`). 32 | 33 | 34 | Returns 35 | ------- 36 | rainbow : RainbowWithModel 37 | A new `RainbowWithModel` object, with the model attached. 38 | """ 39 | 40 | # create a history entry for this action (before other variables are defined) 41 | h = self._create_history_entry("attach_model", locals()) 42 | 43 | # make sure the shape is reasonable 44 | assert np.shape(model) == np.shape(self.flux) 45 | 46 | # add the model to the fluxlike array 47 | inputs = self._create_copy()._get_core_dictionaries() 48 | inputs["fluxlike"]["model"] = model 49 | 50 | # import here (rather than globally) to avoid recursion? 51 | from ..withmodel import RainbowWithModel 52 | 53 | # create new object 54 | new = RainbowWithModel(**inputs) 55 | 56 | # add other inputs to the model 57 | for k, v in kw.items(): 58 | new.__setattr__(k, v) 59 | 60 | # append the history entry to the new Rainbow 61 | new._record_history_entry(h) 62 | 63 | # return the RainboWithModel 64 | return new 65 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/eureka_txt.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["from_eureka_S3_txt"] 4 | 5 | 6 | def eureadka_txt(filename): 7 | """ 8 | Read eureka's concatenated results table 9 | and parse it into a Rainbow-friendly format 10 | 11 | Parameters 12 | ---------- 13 | filename : str 14 | The path to the file. 15 | """ 16 | 17 | # load the data 18 | data = ascii.read(filename) 19 | 20 | # figure out a time array 21 | for time_key in ["time", "bjdtdb"]: 22 | try: 23 | t = np.unique(data[time_key]) 24 | break 25 | except KeyError: 26 | pass 27 | timelike = {} 28 | timelike["time"] = t * u.day * 1 29 | 30 | # figure out a wavelength array 31 | w = np.unique(data["wave_1d"]) 32 | wavelike = {} 33 | wavelike["wavelength"] = w * u.micron * 1 34 | 35 | # populate the fluxlike quantities 36 | fluxlike = {} 37 | i_time = np.arange(len(t)) 38 | # loop through wavelengths, populating all times for each 39 | for i_wavelength in tqdm(range(len(w)), leave=False): 40 | 41 | for k in data.colnames[2:]: 42 | # if an array for this key doesn't exist, create it 43 | if k not in fluxlike: 44 | fluxlike[k] = np.zeros((len(w), len(t))) 45 | 46 | # figure out all indices for this wavelengh 47 | indices_for_this_wavelength = i_wavelength + i_time * len(w) 48 | fluxlike[k][i_wavelength, i_time] = data[k][indices_for_this_wavelength] 49 | 50 | fluxlike["flux"] = fluxlike["optspec"] * 1 51 | fluxlike["uncertainty"] = fluxlike["opterr"] * 1 52 | 53 | return wavelike, timelike, fluxlike 54 | 55 | 56 | def from_eureka_S3_txt(rainbow, filename, **kwargs): 57 | """ 58 | Populate a Rainbow from a eureka pipeline S3 output. 59 | 60 | Parameters 61 | ---------- 62 | 63 | rainbow : Rainbow 64 | The object to be populated. 65 | filename : str 66 | The path to the file. 67 | """ 68 | 69 | # load the Eureka event 70 | wavelike, timelike, fluxlike = eureadka_txt(filename) 71 | 72 | # populate the rainbow 73 | rainbow._initialize_from_dictionaries( 74 | wavelike=wavelike, timelike=timelike, fluxlike=fluxlike 75 | ) 76 | -------------------------------------------------------------------------------- /chromatic/rainbows/writers/__init__.py: -------------------------------------------------------------------------------- 1 | from .rainbow_npy import * 2 | from .rainbow_FITS import * 3 | from .xarray_stellar_spectra import * 4 | from .xarray_raw_light_curves import * 5 | from .xarray_fitted_light_curves import * 6 | 7 | from .text import * 8 | 9 | 10 | # construct a dictionary of available writers 11 | available_writers = {k: globals()[k] for k in globals() if k[0:3] == "to_"} 12 | 13 | 14 | def guess_writer(filepath, format=None): 15 | """ 16 | A wrapper to guess the appropriate writer from the filename 17 | (and possibily an explicitly-set file format string). 18 | 19 | Parameters 20 | ---------- 21 | filepath : str 22 | The path to the file to be written. 23 | format : str, function, (optional) 24 | The file format with which to write the file. 25 | If None, guess format from filepath. 26 | If str, pull reader from dictionary of writers. 27 | If function, treat as a `to_???` writer function. 28 | """ 29 | from fnmatch import fnmatch 30 | import glob 31 | 32 | # get all the possible filenames (= expand wildcard) 33 | f = filepath 34 | 35 | # if format is a function, return it as the reader 36 | if callable(format): 37 | return format 38 | # if format='abcdefgh', return the `to_abcdefgh` function 39 | elif format is not None: 40 | return available_writers[f"to_{format}"] 41 | # does it look like a .rainbow.npy chromatic file? 42 | elif fnmatch(f, "*.rainbow.npy"): 43 | return to_rainbow_npy 44 | elif fnmatch(f.lower(), "*.rainbow.fits") or fnmatch(f.lower(), "*.rainbow.fit"): 45 | return to_rainbow_FITS 46 | # does it look like an ERS-xarray format? 47 | elif fnmatch(f, "*stellar-spec*.xc"): 48 | return to_xarray_stellar_spectra 49 | # does it look like an ERS-xarray format? 50 | elif fnmatch(f, "*raw-light-curve*.xc"): 51 | return to_xarray_raw_light_curves 52 | # does it look like an ERS-xarray format? 53 | elif fnmatch(f, "*fitted-light-curve*.xc"): 54 | return to_xarray_fitted_light_curves 55 | elif fnmatch(f, "*.txt") or fnmatch(filepath, "*.csv"): 56 | return to_text 57 | else: 58 | raise ValueError( 59 | f""" 60 | We're having trouble guessing the output format from the filename 61 | {f} 62 | Please try specifying a `format=` keyword to your `.save` call. 63 | """ 64 | ) 65 | -------------------------------------------------------------------------------- /chromatic/tests/test_normalize.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_normalize(plot=False): 6 | N = 37 7 | w = np.logspace(0, 1, N) * u.micron 8 | f = np.cos(2 * np.pi * w.value / 3) + 1 9 | 10 | snr = 100 11 | a = SimulatedRainbow(wavelength=w).inject_noise(signal_to_noise=snr) 12 | b = SimulatedRainbow(wavelength=w, star_flux=f).inject_noise(signal_to_noise=snr) 13 | c = ( 14 | SimulatedRainbow(wavelength=w, star_flux=f) 15 | .inject_transit() 16 | .inject_noise(signal_to_noise=snr) 17 | ) 18 | 19 | for x in [a, b, c]: 20 | nw = x.normalize(axis="w") 21 | nt = x.normalize(axis="t") 22 | nwt = x.normalize(axis="w").normalize(axis="t") 23 | ntw = x.normalize(axis="t").normalize(axis="w") 24 | 25 | for r in [nw, nt, nwt, ntw]: 26 | r.fluxlike["relative-uncertainty"] = r.uncertainty / r.flux 27 | assert np.all(np.isclose(r.uncertainty / r.flux, 1 / snr, rtol=0.1)) 28 | if plot: 29 | r.paint_quantities(maxcol=4) 30 | plt.close("all") 31 | 32 | 33 | def test_normalization_negative_warnings(): 34 | snr = 0.1 35 | rainbow = SimulatedRainbow().inject_noise(signal_to_noise=0.1) 36 | 37 | with pytest.warns(match="get_median_spectrum"): 38 | rainbow.normalize(axis="wavelength") 39 | ok = rainbow.get_median_spectrum() > 0 40 | rainbow[ok, :].normalize(axis="wavelength") 41 | 42 | with pytest.warns(match="get_median_lightcurve"): 43 | rainbow.normalize(axis="time") 44 | ok = rainbow.get_median_lightcurve() > 0 45 | rainbow[:, ok].normalize(axis="time") 46 | 47 | 48 | def test_normalization_with_not_ok(): 49 | snr = 0.1 50 | rainbow = SimulatedRainbow().inject_noise(signal_to_noise=0.1) 51 | rainbow.ok = rainbow.flux > 0 52 | with warnings.catch_warnings(): 53 | warnings.simplefilter("error") 54 | rainbow.normalize() 55 | 56 | 57 | def test_is_probably_normalized(): 58 | f = [2] * u.W / u.m**2 59 | kw = dict(star_flux=f, R=10, dt=10 * u.minute) 60 | assert SimulatedRainbow(**kw)._is_probably_normalized() == False 61 | assert SimulatedRainbow(**kw).normalize()._is_probably_normalized() == True 62 | assert SimulatedRainbow(**kw).inject_noise()._is_probably_normalized() == False 63 | assert ( 64 | SimulatedRainbow(**kw).inject_noise().normalize()._is_probably_normalized() 65 | == True 66 | ) 67 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/wavelike/noise_comparison_in_bins.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["plot_noise_comparison_in_bins"] 4 | 5 | 6 | def plot_noise_comparison_in_bins( 7 | self, 8 | ax=None, 9 | cmap=None, 10 | vmin=None, 11 | vmax=None, 12 | expected=True, 13 | measured_errorbarkw={}, 14 | expected_plotkw={}, 15 | filename=None, 16 | ylim=[1e-5, 1e-1], 17 | **kw, 18 | ): 19 | """ 20 | ax : Axes 21 | The axes into which the plot should be drawn. 22 | cmap : str, Colormap 23 | The color map to use for expressing wavelength. 24 | vmin : Quantity 25 | The wavelength at the bottom of the cmap. 26 | vmax : Quantity 27 | The wavelength at the top of the cmap. 28 | errorbar_kw : dict 29 | 30 | **kw : dictionary 31 | Additional keywords are passed to `.get_measured_scatter_in_bins`. 32 | Please see the docstring for that method for options. 33 | 34 | """ 35 | 36 | # figure out where the plot should be drawn 37 | if ax is None: 38 | ax = plt.subplot() 39 | plt.sca(ax) 40 | 41 | # make sure the color map is set up 42 | self._make_sure_cmap_is_defined(cmap=cmap, vmin=vmin, vmax=vmax) 43 | 44 | # calculate the binned down scatters 45 | x = self.get_measured_scatter_in_bins(**kw) 46 | 47 | # loop through wavelengths 48 | for i, w in enumerate(self.wavelength): 49 | color = self.get_wavelength_color(w) 50 | ekw = dict(marker="o", color=color) 51 | ekw.update(**measured_errorbarkw) 52 | plt.errorbar(x["N"], x["scatters"][i], x["uncertainty"][i], **ekw) 53 | if expected: 54 | pkw = dict(color=color, alpha=0.2) 55 | pkw.update(**expected_plotkw) 56 | plt.plot(x["N"], x["expectation"][i], **pkw) 57 | 58 | # set good plot appearance default 59 | plt.yscale("log") 60 | plt.xscale("log") 61 | plt.ylim(*ylim) 62 | plt.xlim(1, np.nanmax(x["N"])) 63 | plt.xlabel("# of Times in Bin") 64 | 65 | t_unit = u.minute 66 | twin_ax = ax.twiny() 67 | 68 | plt.sca(twin_ax) 69 | dt = x["dt"].to(t_unit).value 70 | plt.plot(dt, x["expectation"][i], alpha=0) 71 | plt.xlim(np.nanmin(dt), np.nanmax(dt)) 72 | plt.xscale("log") 73 | plt.xlabel(f"Duration of Bin ({t_unit.to_string()})") 74 | 75 | plt.sca(ax) 76 | # plt.title(self.get("title")) 77 | 78 | if filename is not None: 79 | self.savefig(filename) 80 | return ax 81 | -------------------------------------------------------------------------------- /chromatic/tests/test_conversions.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | 3 | closekw = dict(rtol=1e-10) 4 | 5 | 6 | def test_to_df(): 7 | 8 | r = SimulatedRainbow(dt=1 * u.minute, R=50).inject_noise(signal_to_noise=10) 9 | 10 | r_df = r.to_df() 11 | 12 | # ensure the length of the df is the same as the original rainbow 13 | assert len(r_df) == r.nflux 14 | 15 | # ensure we have the right column names 16 | columnnames = r_df.columns 17 | for colname in ["Time (d)", "Wavelength (micron)", "Flux", "Flux Uncertainty"]: 18 | assert colname in columnnames 19 | 20 | # check the values in the df match the rainbow 21 | assert np.isclose( 22 | r_df["Time (d)"].values[0], r.time.to_value("h")[0] / 24.0, **closekw 23 | ) # default=days 24 | assert r_df["Wavelength (micron)"].values[0] == r.wavelength.to_value("micron")[0] 25 | assert r_df["Flux"].values[0] == r.flux[0, 0] 26 | assert r_df["Flux Uncertainty"].values[0] == r.uncertainty[0, 0] 27 | 28 | # test the timeformat parameter 29 | for t_unit in ["h", "hour", "day", "minute", "second", "s"]: 30 | r_df = r.to_df(t_unit=t_unit) 31 | assert f"Time ({t_unit})" in r_df.columns 32 | 33 | 34 | def test_to_nparray(): 35 | r = SimulatedRainbow(dt=1 * u.minute, R=50).inject_noise(signal_to_noise=100) 36 | 37 | rflux, rfluxu, rtime, rwavel = r.to_nparray() 38 | 39 | assert len(rtime) == r.ntime 40 | assert len(rwavel) == r.nwave 41 | assert len(rtime) * len(rwavel) == r.nflux 42 | assert len(rflux.flatten()) == r.nflux 43 | assert len(rfluxu.flatten()) == r.nflux 44 | assert np.shape(rflux) == r.shape 45 | assert np.shape(rfluxu) == r.shape 46 | 47 | assert np.all(rflux == r.flux) 48 | assert np.all(rfluxu == r.uncertainty) 49 | assert np.all(rwavel == r.wavelength.to_value("micron")) 50 | 51 | # issues with rounding errors: 52 | assert np.all( 53 | np.isclose(rtime, r.time.to_value("h") / 24.0, **closekw) 54 | ) # the default is days 55 | 56 | # test if the hours format works 57 | rflux, rfluxu, rtime, rwavel = r.to_nparray(t_unit="h") 58 | assert np.all(np.isclose(rtime, r.time.to_value("h"), **closekw)) 59 | 60 | # test if the minutes format works 61 | rflux, rfluxu, rtime, rwavel = r.to_nparray(t_unit="min") 62 | assert np.all(np.isclose(rtime, r.time.to_value("h") * 60, **closekw)) 63 | 64 | # test if the minutes format works 65 | rflux, rfluxu, rtime, rwavel = r.to_nparray(t_unit="s") 66 | assert np.all(np.isclose(rtime, r.time.to_value("h") * 3600, **closekw)) 67 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/carter_and_may.py: -------------------------------------------------------------------------------- 1 | # import the general list of packages 2 | from ...imports import * 3 | 4 | # define list of the only things that will show up in imports 5 | __all__ = ["from_carter_and_may"] 6 | 7 | 8 | def from_carter_and_may(self, filepath): 9 | """ 10 | Populate a Rainbow from a file in the ERS benchmark 11 | transmission spectrum of WASP-39b format. 12 | 13 | https://zenodo.org/records/10161743 14 | 15 | Parameters 16 | ---------- 17 | 18 | rainbow : Rainbow 19 | The object to be populated. This function is meant 20 | to be called during the initialization of a Rainbow 21 | object. You can/should assume that the `rainbow` object 22 | being passed will already contain the four core 23 | dictionaries `timelike`, `wavelike`, `fluxlike`, and 24 | `metadata`. This function should, at a minimum, add 25 | the following items 26 | + `self.timelike['time']` 27 | + `self.wavelike['wavelength']` 28 | + `self.fluxlike['flux']` 29 | and optionally, additional entries like 30 | + `self.metadata['some-useful-parameter']` 31 | + `self.fluxlike['uncertainty']` 32 | + `self.fluxlike['ok']` 33 | 34 | filepath : str 35 | The path to the file to load. 36 | """ 37 | 38 | # read in your file, however you like 39 | f = h5.File(filepath) 40 | 41 | # populate a 1D array of wavelengths (with astropy units of length) 42 | self.wavelike["wavelength"] = f["wave_1d"] * u.micron 43 | 44 | # populate a 1D array of times (with astropy units of time) 45 | self.timelike["time"] = (f["time"]) * u.day + 2400000.5 * u.day 46 | 47 | # populate a 2D (row = wavelength, col = time) array of fluxes 48 | self.fluxlike["flux"] = np.transpose(f["optspec"]) 49 | self.fluxlike["uncertainty"] = np.transpose(f["opterr"]) 50 | self.fluxlike["ok"] = np.transpose(f["optmask"]) == 0 51 | 52 | for k in f.keys(): 53 | if k not in ["wave_1d", "time", "optspec", "opterr", "optmask"]: 54 | if np.shape(f[k]) == np.shape(self.time): 55 | self.timelike[k] = f[k] 56 | elif np.shape(f[k]) == np.shape(self.wavelength): 57 | self.wavelike[k] = f[k] 58 | elif np.shape(f[k]) == np.shape(self.flux.T): 59 | self.fluxlike[k] = np.transpose(f[k]) 60 | 61 | ok = self.wavelength != 0 62 | for k in self.wavelike: 63 | self.wavelike[k] = self.wavelike[k][ok] 64 | for k in self.fluxlike: 65 | self.fluxlike[k] = self.fluxlike[k][ok, :] 66 | -------------------------------------------------------------------------------- /chromatic/noise/jwst/extract.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | 4 | def extract_sn_from_image(images, extraction_box=3): 5 | """ 6 | Use a 2D image to estimate per-wavelength S/N. 7 | 8 | (This is tuned for MIRI/LRS. It probably won't work on anything else.) 9 | 10 | Returns 11 | ------- 12 | sn_extracted : array 13 | The 1D extracted S/N. 14 | """ 15 | 16 | # extract S/N and saturation images 17 | sn = images["snr"] 18 | saturation = images["saturation"] 19 | 20 | # define extraction aperture 21 | y_pix = np.arange(sn.shape[0]) - sn.shape[0] / 2 + 0.5 22 | extent = [0, sn.shape[1] - 1, -sn.shape[0] / 2 + 0.5, sn.shape[0] / 2 + 0.5] 23 | extract = np.abs(y_pix) < extraction_box 24 | extract_alpha = 0.5 25 | print(f"{np.sum(extract)} pixels being used in the extraction") 26 | 27 | # display the images 28 | fi, ax = plt.subplots( 29 | 2, 30 | 2, 31 | figsize=(7.5, 3), 32 | gridspec_kw=dict(width_ratios=[8, 1]), 33 | sharey="row", 34 | sharex="col", 35 | ) 36 | 37 | plt.sca(ax[0, 0]) 38 | plt.imshow(np.log10(sn), extent=extent, cmap="gray") 39 | 40 | sat_color = "red" 41 | sat_cmap = one2another(bottom=sat_color, top=sat_color, alpha_bottom=0, alpha_top=1) 42 | plt.imshow(saturation, extent=extent, cmap=sat_cmap, vmax=2) 43 | for sign in [-1, 1]: 44 | plt.axhline(sign * extraction_box, alpha=extract_alpha, color="gray") 45 | 46 | plt.sca(ax[0, 1]) 47 | spatial_profile = np.nanmean(sn, axis=1) 48 | plt.plot(spatial_profile, y_pix, color="black") 49 | for sign in [-1, 1]: 50 | plt.axhline(sign * extraction_box, alpha=extract_alpha, color="gray") 51 | 52 | sn_extracted = np.sqrt(np.sum(sn**2, axis=0)) 53 | plt.sca(ax[1, 0]) 54 | plt.plot(sn_extracted, color="black") 55 | plt.ylabel("S/N per pixel\nper integration") 56 | 57 | for i, label in zip([1, 2], ["partial", "full"]): 58 | is_saturated = np.nonzero(np.max(saturation, axis=0) >= i)[0] 59 | if len(is_saturated) > 0: 60 | left, right = np.min(is_saturated), np.max(is_saturated) 61 | plt.axvspan( 62 | left, 63 | right, 64 | color=sat_color, 65 | alpha=0.25 * i, 66 | ) 67 | print(f"{len(is_saturated)} pixels experience at least {label} saturation") 68 | return sn_extracted 69 | 70 | 71 | def trim_image(image, disperser=None): 72 | if disperser.lower() == "p750l": 73 | return image.T[:, 27:399] 74 | else: 75 | return image 76 | return image 77 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/paint.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["paint"] 4 | 5 | 6 | def paint( 7 | self, 8 | ax=None, 9 | quantity="flux", 10 | xaxis="time", 11 | w_unit="micron", 12 | t_unit="day", 13 | colorbar=True, 14 | mask_ok=True, 15 | color_ok="tomato", 16 | alpha_ok=0.8, 17 | vmin=None, 18 | vmax=None, 19 | filename=None, 20 | **kw, 21 | ): 22 | """ 23 | Paint a 2D image of flux as a function of time and wavelength. 24 | 25 | This is a wrapper that tries to choose the best between `.imshow()` 26 | and `.pcolormesh()` for painting a 2D map. `imshow` is faster and 27 | does a better job handling antialiasing for large datasets, but it 28 | works best for linearly or logarithmically uniform axes. `pcolormesh` 29 | is more flexible about non-uniform axes but is slower and doesn't 30 | do anything to combat antialiasing. 31 | 32 | Parameters 33 | ---------- 34 | ax : Axes, optional 35 | The axes into which to make this plot. 36 | quantity : str, optional 37 | The fluxlike quantity to imshow. 38 | (Must be a key of `rainbow.fluxlike`). 39 | xaxis : str 40 | What to use as the horizontal axis, 'time' or 'wavelength'. 41 | w_unit : str, Unit, optional 42 | The unit for plotting wavelengths. 43 | t_unit : str, Unit, optional 44 | The unit for plotting times. 45 | colorbar : bool, optional 46 | Should we include a colorbar? 47 | mask_ok : bool, optional 48 | Should we mark which data are not OK? 49 | color_ok : str, optional 50 | The color to be used for masking data points that are not OK. 51 | alpha_ok : float, optional 52 | The transparency to be used for masking data points that are not OK. 53 | **kw : dict, optional 54 | All other keywords will be passed on to `plt.imshow` or `plt.pcolormesh`, 55 | so you can have more detailed control over the plot 56 | appearance. Common keyword arguments might include: 57 | `[cmap, norm, interpolation, alpha, vmin, vmax]` (and more) 58 | More details are available at 59 | https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html 60 | """ 61 | # record all keyword inputs in a dictionary 62 | inputs = locals() 63 | inputs.pop("self") 64 | kw = inputs.pop("kw") 65 | inputs.update(**kw) 66 | 67 | imshow_scales = ["linear", "log"] 68 | if (self.wscale in imshow_scales) and (self.tscale in imshow_scales): 69 | ax = self.imshow(**inputs) 70 | else: 71 | ax = self.pcolormesh(**inputs) 72 | 73 | return ax 74 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # chromatic 2 | Tools for visualizing spectroscopic light curves, with flux as a function of wavelength and time. Read the 🌈[documentation](https://zkbt.github.io/chromatic/)🌈 to see how it works! 3 | 4 | It's being developed in support of the JWST Transiting Exoplanet Community Early Release Science Program ([ers-transit](https://ers-transit.github.io/)) and easier multiwavelength observations of transiting exoplanets in general, from telescopes in space or on the ground. This package is still actively being developed. Please submit Issues for bugs you notice, features that aren't clearly explained in the documentation, or functionality you'd like to see us implement. 5 | 6 | ## Installation 7 | If you want to install this code just to use it, you can simply run 8 | 9 | ``` 10 | pip install chromatic-lightcurves 11 | ``` 12 | 13 | and it should install everything, along with all the dependencies needed to run the code. If you previously installed this package and need to grab a newer version, you can run 14 | 15 | ``` 16 | pip install --upgrade chromatic-lightcurves 17 | ``` 18 | For Developer Installation instructions, see the [documentation](https://zkbt.github.io/chromatic/installation/). 19 | 20 | ## Usage 21 | 22 | For an ultra-quick start try 23 | ```python 24 | from chromatic import * 25 | r = SimulatedRainbow().inject_transit().inject_noise() 26 | r.normalize().bin(dw=0.5*u.micron, dt=15*u.minute).paint() 27 | ``` 28 | and then see the 🌈[documentation](https://zkbt.github.io/chromatic/)🌈 for more. 29 | 30 | 31 | ## Contributing 32 | 33 | We welcome contributions from anyone who agrees to follow the `ers-transit` [Code of Conduct](https://ers-transit.github.io/code-of-conduct.html#ers-transit). If you're on the `ers-transit` slack, please join the #hack-chromatic channel there and say hello; otherwise, please contact Zach directly or just dive right in! 34 | 35 | A great initial way to contribute would be to [submit an Issue](https://github.com/zkbt/chromatic/issues) about a bug, question, or suggestion you might have. If you want to contribute code, the [Developer Guide](https://zkbt.github.io/chromatic/designing/) is probably the best place to start. We know it can feel a little scary to try to contribute to a shared code package, so we try our best to be friendly and helpful to new contributors trying to learn how! 36 | 37 | *And for context, Zach is a little new to trying to manage a big collaborative code project, so if there are things we could be doing better, please let him know!* 38 | 39 | The goal is to submit `chromatic-lightcurves` to the [Journal of Open Source Software](https://joss.theoj.org/) before the end of 2022. If you contribute before then, you'll be included on the paper! 40 | -------------------------------------------------------------------------------- /chromatic/rainbows/writers/template.py: -------------------------------------------------------------------------------- 1 | """ 2 | This module serves as a template for creating a new Rainbow 3 | writer. If you want to add the ability to writer chromatic 4 | light curves to a new kind of file format, a good process 5 | would be to do something like the following 6 | 7 | 1. Copy this `template.py` file into a new file in the 8 | `writers/` directory, ideally with a name that's easy 9 | to recognize, such as `writers/abcdefgh.py` (assuming 10 | `abcdefgh` is the name of your format) 11 | 12 | 2. Start by finding and replacing `abcdefgh` in this 13 | template with the name of your format. 14 | 15 | 3. Edit the `to_abcdefgh` function so that it will 16 | write out a Rainbow to whatever format you want. 17 | 18 | 4. Edit the `writers/__init__.py` file to import your 19 | `to_abcdefgh` function to be accessible when people 20 | try to write Rainbows out to files. Add an `elif` statement 21 | to the `guess_writer` function that will help guess which 22 | writer to use from some aspect(s) of the filename. 23 | 24 | (This `guess_writer` function also accepts a `format=` 25 | keyword that allows the user to explicitly specify that 26 | the abcdefgh writer should be used.) 27 | 28 | 5. Submit a pull request to the github repository for 29 | this package, so that other folks can use your handy 30 | new writer too! 31 | """ 32 | 33 | # import the general list of packages 34 | from ...imports import * 35 | 36 | # define list of the only things that will show up in imports 37 | __all__ = ["to_abcdefgh"] 38 | 39 | 40 | def to_abcdefgh(self, filepath): 41 | """ 42 | Write a Rainbow to a file in the abcdefgh format. 43 | 44 | Parameters 45 | ---------- 46 | 47 | self : Rainbow 48 | The object to be saved. 49 | 50 | filepath : str 51 | The path to the file to write. 52 | """ 53 | 54 | # a 1D array of wavelengths (with astropy units of length) 55 | the_1D_array_of_wavelengths = self.wavelike["wavelength"] 56 | 57 | # a 1D array of times (with astropy units of time) 58 | the_1D_array_of_times = self.timelike["time"] 59 | 60 | # a 2D (row = wavelength, col = array of fluxes 61 | the_2D_array_of_fluxes = self.fluxlike["flux"] 62 | 63 | # write out your file, however you like 64 | write_to_abcdefgh( 65 | filepath, 66 | the_1D_array_of_wavelengths, 67 | the_1D_array_of_times, 68 | the_2D_array_of_fluxes, 69 | **some_other_stuff_too_maybe, 70 | ) 71 | 72 | # add some warnings if there's any funny business 73 | if something_goes_wonky(): 74 | cheerfully_suggest( 75 | f""" 76 | Here's a potential problem that the user should know about. 77 | """ 78 | ) 79 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/diagnostics/paint_quantities.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["paint_quantities"] 4 | 5 | 6 | def paint_quantities( 7 | self, quantities=None, maxcol=3, panel_size=(5, 4), filename=None, **kw 8 | ): 9 | """ 10 | imshow fluxlikes as a function of time (x = time, y = wavelength, color = flux). 11 | 12 | Parameters 13 | ---------- 14 | quantities : None, list, optional 15 | The fluxlike quantity to imshow. 16 | maxcol : int, optional 17 | The maximum number of columns to show (Optional). 18 | panel_size : tuple, optional 19 | The (approximate) size of a single panel, which will 20 | be used to set the overall figsize based on the 21 | number of rows and columns (Optional). 22 | **kw : dict, optional 23 | Additional keywords will be passed on to `imshow` 24 | 25 | """ 26 | 27 | # decide which quantities to plot 28 | if quantities is None: 29 | allkeys = self.fluxlike.keys() 30 | else: 31 | allkeys = quantities[:] 32 | 33 | # set up the geometry of the grid 34 | if len(allkeys) > maxcol: 35 | rows = int(np.ceil(len(allkeys) / maxcol)) 36 | cols = maxcol 37 | else: 38 | rows = 1 39 | cols = np.min([len(allkeys), maxcol]) 40 | 41 | # create the figure and grid of axes 42 | fig, axes = plt.subplots( 43 | rows, 44 | cols, 45 | figsize=(cols * panel_size[0], rows * panel_size[1]), 46 | sharex=True, 47 | sharey=True, 48 | constrained_layout=True, 49 | ) 50 | 51 | # make the axes easier to index 52 | if len(allkeys) > 1: 53 | ax = axes.flatten() 54 | else: 55 | ax = [axes] 56 | 57 | # display each quantity 58 | for k, key in enumerate(allkeys): 59 | # make the painting (or an empty box) 60 | if key in self.fluxlike.keys(): 61 | self.paint(quantity=key, ax=ax[k], **kw) 62 | else: 63 | ax[k].text( 64 | 0.5, 65 | 0.5, 66 | f"No {key}", 67 | transform=ax[k].transAxes, 68 | ha="center", 69 | va="center", 70 | ) 71 | # add a title for each box 72 | ax[k].set_title(key) 73 | 74 | # hide xlabel except on the bottom row 75 | if k < (len(allkeys) - cols): 76 | ax[k].set_xlabel("") 77 | else: 78 | ax[k].tick_params(labelbottom=True) 79 | 80 | # hide ylabel except on the left column 81 | if (k % cols) > 0: 82 | ax[k].set_ylabel("") 83 | 84 | # hide any additional axes 85 | if k + 1 <= len(ax): 86 | for axi in ax[k + 1 :]: 87 | axi.axis("Off") 88 | if filename is not None: 89 | self.savefig(filename) 90 | -------------------------------------------------------------------------------- /docs/documentation.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "0c858c15", 6 | "metadata": {}, 7 | "source": [ 8 | "# Writing 🌈 Documentation\n", 9 | "\n", 10 | "If you're contributing a new feature to `chromatic`, please consider also contributing some documentation to explain how your feature works. Here's the very short version of how to add to the documentation:\n", 11 | "\n", 12 | "1. Install in development mode (see [Installation](../installation)), so you have access to `mkdocs` and the various extensions needed to render the documentation.\n", 13 | "1. Decide whether your explanation would fit well within an existing page or whether you need a new one. In the `docs/` directory, find the appropraite `.ipynb` notebook file or create a new one. If you create a new one, add it to the `nav:` section of the `mkdocs.yml` file in the main repository directory so that `mkdocs` will know to include it.\n", 14 | "1. Write your example and explanation in a `.ipynb` file. Your audience should be smart people who want to use the code but don't have much experience with it yet. Be friendly and encouraging!\n", 15 | "1. From the Terminal, run `mkdocs serve`. This will convert all of the source notebooks into a live website, and give you a little address that you can copy and paste into a browser window. While that `mkdocs serve` command is still running, small changes you make to existing `.ipynb` source files will appear (sometimes after a few minutes) on the live locally-hosted webserver. \n", 16 | "1. Once you're happy with your new documentation, before committing it to the repository, please run \"Kernal > Restart & Clear Output\" or something similar to remove all outputs from the source notebook file. The `chromatic` repository will hang onto *all* changes that you commit to it, so it would very quickly get annoyingly large unless we leave the outputs out of committed notebook files. Double check the outputs are all gone, save your notebook, and then commit it to the `git` repository (see [Contributing 🌈 Code with GitHub](../github)).\n", 17 | "\n", 18 | "Periodically, after reviewing and copy-editing the documentation, we'll deploy the newest version up to the web at [zkbt.github.io/chromatic/](https://zkbt.github.io/chromatic/) for all to enjoy." 19 | ] 20 | } 21 | ], 22 | "metadata": { 23 | "kernelspec": { 24 | "display_name": "Python 3", 25 | "language": "python", 26 | "name": "python3" 27 | }, 28 | "language_info": { 29 | "codemirror_mode": { 30 | "name": "ipython", 31 | "version": 3 32 | }, 33 | "file_extension": ".py", 34 | "mimetype": "text/x-python", 35 | "name": "python", 36 | "nbconvert_exporter": "python", 37 | "pygments_lexer": "ipython3", 38 | "version": "3.9.12" 39 | } 40 | }, 41 | "nbformat": 4, 42 | "nbformat_minor": 5 43 | } 44 | -------------------------------------------------------------------------------- /chromatic/tests/test_summaries.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_lightcurve_and_spectrum_summaries(): 6 | fi, ax = plt.subplots(1, 2, sharey=True) 7 | s = SimulatedRainbow().inject_noise().inject_transit() 8 | ax[0].plot(s.time, s.get_average_lightcurve()) 9 | ax[1].plot(s.wavelength, s.get_average_spectrum()) 10 | plt.savefig( 11 | os.path.join( 12 | test_directory, "demonstration-of-average-lightcurve-and-spectrum.pdf" 13 | ) 14 | ) 15 | 16 | 17 | def test_measured_scatter_summary(): 18 | # make a fake Rainbow with known noise 19 | signal_to_noise = 100 20 | s = SimulatedRainbow().inject_noise(signal_to_noise=signal_to_noise) 21 | 22 | # mathematically, what do we expect? 23 | expected_sigma = 1 / signal_to_noise 24 | uncertainty_on_expected_sigma = expected_sigma / np.sqrt(2 * (s.ntime - 1)) 25 | # (see equation 3.48 of Sivia and Skilling for the uncertainty on sigma) 26 | 27 | # loop through methods 28 | methods = ["standard-deviation", "MAD"] 29 | fi, ax = plt.subplots( 30 | 1, 31 | len(methods), 32 | sharey=True, 33 | sharex=True, 34 | figsize=(8, 3), 35 | dpi=300, 36 | constrained_layout=True, 37 | facecolor="white", 38 | ) 39 | for i, method in enumerate(methods): 40 | # calculate the measured scatter 41 | measured_scatter = s.get_measured_scatter(method=method) 42 | 43 | # plot the measured scatters + 1-sigma expectation for its uncertainty 44 | plt.sca(ax[i]) 45 | plt.plot(s.wavelength, measured_scatter, marker="o", color="black") 46 | plt.axhline(expected_sigma, color="black", alpha=0.3) 47 | plt.axhspan( 48 | expected_sigma - uncertainty_on_expected_sigma, 49 | expected_sigma + uncertainty_on_expected_sigma, 50 | color="gray", 51 | alpha=0.3, 52 | ) 53 | 54 | # plot range that shouldn't be exceeded in this test 55 | many_sigma = 10 56 | plt.axhspan( 57 | expected_sigma - uncertainty_on_expected_sigma * many_sigma, 58 | expected_sigma + uncertainty_on_expected_sigma * many_sigma, 59 | color="red", 60 | alpha=0.05, 61 | zorder=-10, 62 | ) 63 | plt.ylim(0, None) 64 | plt.xscale("log") 65 | plt.xlabel(f'Wavelength ({s.wavelength.unit.to_string("latex_inline")})') 66 | plt.ylabel(f"Measured Scatter\n('{method}')") 67 | 68 | # throw an error if the measured scatter gets too large 69 | assert np.all( 70 | np.abs(measured_scatter - expected_sigma) 71 | < uncertainty_on_expected_sigma * many_sigma 72 | ) 73 | 74 | plt.savefig(os.path.join(test_directory, "demonstration-of-measured-scatter.pdf")) 75 | -------------------------------------------------------------------------------- /chromatic/tests/test_ok.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_ok_rows_and_columns(): 6 | s = SimulatedRainbow().inject_noise() 7 | s.wavelike["ok"] = np.arange(s.nwave) > 10 8 | s.timelike["ok"] = np.arange(s.ntime) > 5 9 | s.fluxlike["ok"] = np.ones(s.shape) 10 | assert s.ok[0, 0] == 0 11 | assert s.ok[-1, -1] == 1 12 | assert np.all(s.ok[:10, :] == 0) 13 | assert np.all(s.ok[:, :5] == 0) 14 | 15 | 16 | def test_bin_with_not_ok_data(): 17 | with warnings.catch_warnings(): 18 | warnings.simplefilter("ignore") 19 | for ok_fraction in [0, 0.5, 0.99, 1]: 20 | a = ( 21 | SimulatedRainbow(dt=2 * u.minute, dw=0.2 * u.micron) 22 | .inject_transit() 23 | .inject_noise() 24 | ) 25 | a.ok = np.random.uniform(size=a.shape) < ok_fraction 26 | 27 | if ok_fraction == 0: 28 | with pytest.raises((RuntimeError, IndexError)): 29 | should_fail = a.bin( 30 | dw=0.7 * u.micron, dt=20 * u.minute, minimum_acceptable_ok=1 31 | ) 32 | continue 33 | 34 | cautious = a.bin( 35 | dw=0.4 * u.micron, dt=4 * u.minute, minimum_acceptable_ok=1 36 | ) 37 | carefree = a.bin( 38 | dw=0.4 * u.micron, dt=4 * u.minute, minimum_acceptable_ok=0, trim=False 39 | ) 40 | if np.any(a.ok == 0): 41 | assert np.any((carefree.ok != 1) & (carefree.ok != 0)) 42 | 43 | 44 | def test_get_helpers(): 45 | s = SimulatedRainbow() 46 | assert s.get("flux") is s.flux 47 | 48 | s.get_for_wavelength(0).shape == (s.ntime,) 49 | s.get_for_wavelength(0, "flux").shape == (s.ntime,) 50 | s.get_for_wavelength(0, "time").shape == (s.ntime,) 51 | with pytest.raises(RuntimeError): 52 | s.get_for_wavelength(0, "wavelength") 53 | 54 | s.get_for_time(0).shape == (s.nwave,) 55 | s.get_for_time(0, "flux").shape == (s.nwave,) 56 | s.get_for_time(0, "wavelength").shape == (s.nwave,) 57 | with pytest.raises(RuntimeError): 58 | s.get_for_time(0, "time") 59 | 60 | 61 | def test_get_ok_data_helpers(quantity="flux"): 62 | s = SimulatedRainbow(dw=0.5 * u.micron, dt=20 * u.minute).inject_noise() 63 | s.ok = np.random.uniform(0, 1, s.shape) > 0.5 64 | s.paint() 65 | 66 | fi, ax = plt.subplots(2, 2, sharex="col", sharey=True, constrained_layout=True) 67 | for r, e in enumerate([False, True]): 68 | ax[r, 0].set_title(f"express_badness_with_uncertainty={e}") 69 | 70 | for c, f in enumerate([s.get_ok_data_for_wavelength, s.get_ok_data_for_time]): 71 | for i in range(3): 72 | x, y, sigma = f(i, y=quantity, express_badness_with_uncertainty=e) 73 | ax[r, c].scatter(x, y - i * 0.2, c=np.isfinite(sigma)) 74 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/inject_spectrum.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | from ...spectra import * 3 | 4 | __all__ = ["inject_spectrum"] 5 | 6 | 7 | def inject_spectrum( 8 | self, 9 | temperature=5800 * u.K, 10 | logg=4.43, 11 | metallicity=0.0, 12 | radius=1 * u.Rsun, 13 | distance=10 * u.pc, 14 | phoenix=True, 15 | ): 16 | """ 17 | Inject a stellar spectrum into the flux. 18 | 19 | This injects a constant stellar spectrum into 20 | all times in the `Rainbow`. Injection happens 21 | by multiplying the `.model` flux array, so for 22 | example a model that already has a transit in 23 | it will be scaled up to match the stellar spectrum 24 | in all wavelengths. 25 | 26 | Parameters 27 | ---------- 28 | temperature : Quantity, optional 29 | Temperature, in K (with no astropy units attached). 30 | logg : float, optional 31 | Surface gravity log10[g/(cm/s**2)] (with no astropy units attached). 32 | metallicity : float, optional 33 | Metallicity log10[metals/solar] (with no astropy units attached). 34 | radius : Quantity, optional 35 | The radius of the star. 36 | distance : Quantity, optional 37 | The distance to the star. 38 | phoenix : bool, optional 39 | If `True`, use PHOENIX surface flux. 40 | If `False`, use Planck surface flux. 41 | 42 | Returns 43 | ------- 44 | rainbow : Rainbow 45 | A new `Rainbow` object with the spectrum injected. 46 | """ 47 | 48 | # create a history entry for this action (before other variables are defined) 49 | h = self._create_history_entry("inject_spectrum", locals()) 50 | 51 | # create a copy of the existing Rainbow 52 | new = self._create_copy() 53 | 54 | # warn if maybe we shouldn't inject anything 55 | if np.all(u.Quantity(self.flux).value != 1): 56 | cheerfully_suggest( 57 | f""" 58 | None of the pre-existing flux values were 1, 59 | which hints at the possibility that there 60 | might already be a spectrum in them. Please 61 | watch out for weird units or values! 62 | """ 63 | ) 64 | 65 | if phoenix: 66 | f = get_phoenix_photons 67 | else: 68 | f = get_planck_photons 69 | 70 | # get the spectrum from the surface 71 | _, surface_flux = f( 72 | temperature=u.Quantity(temperature).value, 73 | logg=logg, 74 | metallicity=metallicity, 75 | wavelength=self.wavelength, 76 | ) 77 | 78 | # get the received flux at Earth 79 | received_flux = surface_flux * (radius / distance).decompose() ** 2 80 | 81 | # do math with spectrum 82 | for k in ["flux", "model", "uncertainty"]: 83 | try: 84 | new.fluxlike[k] = self.get(k) * self._broadcast_to_fluxlike(received_flux) 85 | except KeyError: 86 | pass 87 | 88 | # append the history entry to the new Rainbow 89 | new._record_history_entry(h) 90 | 91 | # return the new object 92 | return new 93 | -------------------------------------------------------------------------------- /chromatic/tests/test_io.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_rainbow_npy(): 6 | filename = os.path.join(test_directory, "test.rainbow.npy") 7 | a = SimulatedRainbow().inject_noise() 8 | a.save(filename) 9 | b = Rainbow(filename) 10 | assert a == b 11 | 12 | 13 | def test_rainbow_FITS(): 14 | filename = os.path.join(test_directory, "test.rainbow.fits") 15 | a = SimulatedRainbow().inject_noise() 16 | with warnings.catch_warnings(): 17 | warnings.simplefilter("ignore") 18 | a.save(filename, overwrite=True) 19 | b = read_rainbow(filename) 20 | assert a == b 21 | 22 | 23 | def test_text(): 24 | filename = os.path.join(test_directory, "test.rainbow.txt") 25 | a = SimulatedRainbow().inject_noise() 26 | a.save(filename, overwrite=True) 27 | b = read_rainbow(filename) 28 | assert a == b 29 | 30 | 31 | def test_xarray(): 32 | 33 | for f in [ 34 | "stellar-spec-planettest-modetest-codechromatic-authorzkbt.xc", 35 | "raw-light-curves-planettest-modetest-codechromatic-authorzkbt.xc", 36 | "fitted-light-curves-planettest-modetest-codechromatic-authorzkbt.xc", 37 | ]: 38 | filename = os.path.join(test_directory, f) 39 | print(filename) 40 | a = SimulatedRainbow().inject_transit().inject_systematics().inject_noise() 41 | 42 | with pytest.warns(match="required metadata keyword"): 43 | a.save(filename) 44 | b = read_rainbow(filename) 45 | assert a == b 46 | 47 | 48 | def test_rainbow_npy(): 49 | filename = os.path.join(test_directory, "test.rainbow.npy") 50 | a = SimulatedRainbow().inject_noise() 51 | a.save(filename) 52 | b = Rainbow(filename) 53 | assert a == b 54 | 55 | 56 | def test_guess_readers(): 57 | 58 | assert guess_reader("some-neat-file.rainbow.npy") == from_rainbow_npy 59 | assert guess_reader("some-neat-file-x1dints.fits") == from_x1dints 60 | 61 | # eureka readers 62 | assert guess_reader("S3_wasp39b_ap6_bg7_SpecData.h5") == from_eureka_S3 63 | assert guess_reader("S4_wasp39b_ap6_bg7_LCData.h5") == from_eureka_S4 64 | assert ( 65 | guess_reader( 66 | [ 67 | "S5_wasp39b_ap6_bg7_Table_Save_ch0.txt", 68 | "S5_wasp39b_ap6_bg7_Table_Save_ch1.txt", 69 | ] 70 | ) 71 | == from_eureka_S5 72 | ) 73 | 74 | # xarray common format 75 | assert guess_reader("stellar-spec-wow.xc") == from_xarray_stellar_spectra 76 | assert guess_reader("raw-light-curves-wow.xc") == from_xarray_raw_light_curves 77 | assert guess_reader("fitted-light-curve-wow.xc") == from_xarray_fitted_light_curves 78 | 79 | 80 | def test_guess_writers(): 81 | 82 | assert guess_writer("some-neat-file.rainbow.npy") == to_rainbow_npy 83 | 84 | # xarray common format 85 | assert guess_writer("stellar-spec-wow.xc") == to_xarray_stellar_spectra 86 | assert guess_writer("raw-light-curves-wow.xc") == to_xarray_raw_light_curves 87 | assert guess_writer("fitted-light-curve-wow.xc") == to_xarray_fitted_light_curves 88 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # a few extra things to watch out for 2 | .vscode/settings.json 3 | .DS_Store 4 | *.mp4 5 | *.gif 6 | *.png 7 | *.pdf 8 | *.dat 9 | *.npy 10 | *.txt 11 | *.html 12 | *.h5 13 | *.xc 14 | !descriptions.txt 15 | sandbox/ 16 | docs/example-datasets/ 17 | docs/downloads-for-exoatlas/ 18 | 19 | site/ 20 | examples/ 21 | *.swp 22 | 23 | 24 | # Byte-compiled / optimized / DLL files 25 | __pycache__/ 26 | *.py[cod] 27 | *$py.class 28 | 29 | # C extensions 30 | *.so 31 | 32 | # Distribution / packaging 33 | .Python 34 | build/ 35 | develop-eggs/ 36 | dist/ 37 | downloads/ 38 | eggs/ 39 | .eggs/ 40 | lib/ 41 | lib64/ 42 | parts/ 43 | sdist/ 44 | var/ 45 | wheels/ 46 | pip-wheel-metadata/ 47 | share/python-wheels/ 48 | *.egg-info/ 49 | .installed.cfg 50 | *.egg 51 | MANIFEST 52 | 53 | # PyInstaller 54 | # Usually these files are written by a python script from a template 55 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 56 | *.manifest 57 | *.spec 58 | 59 | # Installer logs 60 | pip-log.txt 61 | pip-delete-this-directory.txt 62 | 63 | # Unit test / coverage reports 64 | htmlcov/ 65 | .tox/ 66 | .nox/ 67 | .coverage 68 | .coverage.* 69 | .cache 70 | nosetests.xml 71 | coverage.xml 72 | *.cover 73 | *.py,cover 74 | .hypothesis/ 75 | .pytest_cache/ 76 | 77 | # Translations 78 | *.mo 79 | *.pot 80 | 81 | # Django stuff: 82 | *.log 83 | local_settings.py 84 | db.sqlite3 85 | db.sqlite3-journal 86 | 87 | # Flask stuff: 88 | instance/ 89 | .webassets-cache 90 | 91 | # Scrapy stuff: 92 | .scrapy 93 | 94 | # Sphinx documentation 95 | docs/_build/ 96 | 97 | # PyBuilder 98 | target/ 99 | 100 | # Jupyter Notebook 101 | .ipynb_checkpoints 102 | 103 | # IPython 104 | profile_default/ 105 | ipython_config.py 106 | 107 | # pyenv 108 | .python-version 109 | 110 | # pipenv 111 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 112 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 113 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 114 | # install all needed dependencies. 115 | #Pipfile.lock 116 | 117 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 118 | __pypackages__/ 119 | 120 | # Celery stuff 121 | celerybeat-schedule 122 | celerybeat.pid 123 | 124 | # SageMath parsed files 125 | *.sage.py 126 | 127 | # Environments 128 | .env 129 | .venv 130 | env/ 131 | venv/ 132 | ENV/ 133 | env.bak/ 134 | venv.bak/ 135 | 136 | # Spyder project settings 137 | .spyderproject 138 | .spyproject 139 | 140 | # Rope project settings 141 | .ropeproject 142 | 143 | # mkdocs documentation 144 | /site 145 | 146 | # mypy 147 | .mypy_cache/ 148 | .dmypy.json 149 | dmypy.json 150 | 151 | # Pyre type checker 152 | .pyre/ 153 | *.fits 154 | chromatic/tests/example-extracted-datasets/eureka-extraction/S3_wasp43b_Table_Save.txt 155 | chromatic/rainbows/readers/Untitled.ipynb 156 | chromatic/tests/test_reader.ipynb 157 | chromatic/tests/eureka_dat_test.ipynb 158 | chromatic/tests/eureka_txt_test.ipynb 159 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/measured_scatter_in_bins.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["get_measured_scatter_in_bins"] 4 | 5 | 6 | def get_measured_scatter_in_bins( 7 | self, ntimes=2, nbins=4, method="standard-deviation", minimum_acceptable_ok=1e-10 8 | ): 9 | """ 10 | Get measured scatter in time bins of increasing sizes. 11 | 12 | For uncorrelated Gaussian noise, the scatter should 13 | decrease as 1/sqrt(N), where N is the number points 14 | in a bin. This function calculates the scatter for 15 | a range of N, thus providing a quick test for 16 | correlated noise. 17 | 18 | Parameters 19 | ---------- 20 | ntimes : int 21 | How many times should be binned together? Binning will 22 | continue recursively until fewer that nbins would be left. 23 | nbins : int 24 | What's the smallest number of bins that should be used to 25 | calculate a scatter? The absolute minimum is 2. 26 | method : string 27 | What method to use to obtain measured scatter. Current options are 'MAD', 'standard-deviation'. 28 | minimum_acceptable_ok : float 29 | The smallest value of `ok` that will still be included. 30 | (1 for perfect data, 1e-10 for everything but terrible data, 0 for all data) 31 | 32 | Returns 33 | ------- 34 | scatter_dictionary : dict 35 | Dictionary with lots of information about scatter in bins per wavelength. 36 | """ 37 | 38 | from ...rainbow import Rainbow 39 | 40 | if "remove_trends" in self.history(): 41 | cheerfully_suggest( 42 | f""" 43 | The `remove_trends` function was applied to this `Rainbow`, 44 | making it very plausible that some long-timescale signals 45 | and/or noise have been suppressed. Be suspicious of binned 46 | scatters on long timescales. 47 | """ 48 | ) 49 | 50 | # create a simplified rainbow so we don't waste time binning 51 | simple = Rainbow( 52 | time=self.time, 53 | wavelength=self.wavelength, 54 | flux=self.flux, 55 | uncertainty=self.uncertainty, 56 | ok=self.ok, 57 | ) 58 | 59 | # loop through binning until done 60 | binnings = [simple] 61 | N = [1] 62 | while binnings[-1].ntime > ntimes * nbins: 63 | binnings.append( 64 | binnings[-1].bin(ntimes=ntimes, minimum_acceptable_ok=minimum_acceptable_ok) 65 | ) 66 | N.append(N[-1] * ntimes) 67 | 68 | scatters = [b.get_measured_scatter(method=method) for b in binnings] 69 | expectation = [b.get_expected_uncertainty() for b in binnings] 70 | uncertainty_on_scatters = ( 71 | scatters 72 | / np.sqrt(2 * (np.array([b.ntime for b in binnings]) - 1))[:, np.newaxis] 73 | ) 74 | dt = [np.median(np.diff(b.time)) for b in binnings] 75 | 76 | return dict( 77 | N=np.array(N), 78 | dt=u.Quantity(dt), 79 | scatters=np.transpose(scatters), 80 | expectation=np.transpose(expectation), 81 | uncertainty=np.transpose(uncertainty_on_scatters), 82 | ) 83 | # (see equation 3.48 of Sivia and Skilling for the uncertainty on sigma) 84 | -------------------------------------------------------------------------------- /docs/multi.ipynb.ignore: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "2d585005", 6 | "metadata": {}, 7 | "source": [ 8 | "# Comparing 🌈 to 🌈\n", 9 | "\n", 10 | "Often, we'll want to directly compare two different Rainbows. A wrapper called `compare_rainbows` tries to make doing so a little simpler, by providing an interface to apply many of the familiar `Rainbow` methods to multiple objects at once." 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": null, 16 | "id": "54e2ace1", 17 | "metadata": {}, 18 | "outputs": [], 19 | "source": [ 20 | "from chromatic import *" 21 | ] 22 | }, 23 | { 24 | "cell_type": "code", 25 | "execution_count": null, 26 | "id": "ba1549ef", 27 | "metadata": {}, 28 | "outputs": [], 29 | "source": [ 30 | "t = SimulatedRainbow().inject_transit()\n", 31 | "a = t.inject_noise(signal_to_noise=1000)\n", 32 | "b = t.inject_noise(signal_to_noise=np.sqrt(100 * 1000))\n", 33 | "c = t.inject_noise(signal_to_noise=100)" 34 | ] 35 | }, 36 | { 37 | "cell_type": "code", 38 | "execution_count": null, 39 | "id": "9e14233e", 40 | "metadata": {}, 41 | "outputs": [], 42 | "source": [ 43 | "c = compare_rainbows([a, b, c], names=[\"noisy\", \"noisier\", \"noisiest\"])" 44 | ] 45 | }, 46 | { 47 | "cell_type": "code", 48 | "execution_count": null, 49 | "id": "9eb4c042", 50 | "metadata": {}, 51 | "outputs": [], 52 | "source": [ 53 | "c.bin(R=4).plot(spacing=0.01)" 54 | ] 55 | }, 56 | { 57 | "cell_type": "code", 58 | "execution_count": null, 59 | "id": "688dd9f0", 60 | "metadata": {}, 61 | "outputs": [], 62 | "source": [ 63 | "c.imshow(cmap=\"gray\")" 64 | ] 65 | }, 66 | { 67 | "cell_type": "code", 68 | "execution_count": null, 69 | "id": "c0f5af80", 70 | "metadata": {}, 71 | "outputs": [], 72 | "source": [ 73 | "c.animate_lightcurves()" 74 | ] 75 | }, 76 | { 77 | "cell_type": "markdown", 78 | "id": "34c31359", 79 | "metadata": {}, 80 | "source": [ 81 | "" 82 | ] 83 | }, 84 | { 85 | "cell_type": "code", 86 | "execution_count": null, 87 | "id": "d2bf4aa5", 88 | "metadata": {}, 89 | "outputs": [], 90 | "source": [ 91 | "c.animate_spectra()" 92 | ] 93 | }, 94 | { 95 | "cell_type": "markdown", 96 | "id": "d6eb5762", 97 | "metadata": {}, 98 | "source": [ 99 | "" 100 | ] 101 | } 102 | ], 103 | "metadata": { 104 | "kernelspec": { 105 | "display_name": "Python 3", 106 | "language": "python", 107 | "name": "python3" 108 | }, 109 | "language_info": { 110 | "codemirror_mode": { 111 | "name": "ipython", 112 | "version": 3 113 | }, 114 | "file_extension": ".py", 115 | "mimetype": "text/x-python", 116 | "name": "python", 117 | "nbconvert_exporter": "python", 118 | "pygments_lexer": "ipython3", 119 | "version": "3.9.12" 120 | } 121 | }, 122 | "nbformat": 4, 123 | "nbformat_minor": 5 124 | } 125 | -------------------------------------------------------------------------------- /chromatic/noise/jwst/visualize.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | 4 | def plot_pandexo(t): 5 | """ 6 | Make a quick summary plot of PandExo tabular results. 7 | """ 8 | plt.figure(figsize=(8, 4)) 9 | plt.title( 10 | f"{t.meta['pandexo_input']['Instrument']}+{t.meta['pandexo_input']['Mode']}+{t.meta['pandexo_input']['Disperser']} | PandExo" 11 | ) 12 | for k in t.colnames: 13 | if "snr" in k: 14 | plt.plot( 15 | t["wavelength"], 16 | t[k], 17 | label=f"{k}\n(median S/N = {np.median(t[k]):.1f})", 18 | ) 19 | plt.yscale("log") 20 | plt.legend(frameon=False, bbox_to_anchor=(1, 1)) 21 | plt.xlabel("Wavelength (microns)") 22 | plt.ylabel("S/N per extracted pixel") 23 | 24 | N_groups = t.meta["number_of_groups_per_integration"] 25 | summary = f""" 26 | {t.meta['number_of_groups_per_integration']} groups/integration 27 | ({N_groups}-1)/({N_groups}+1) = {t.meta['observing_efficiency']:.2%} duty cycle 28 | {t.meta['time_per_group']}s/group 29 | {t.meta['time_per_integration']}s/integration 30 | {t.meta['number_of_integrations_per_transit']} integrations/transit 31 | {len(t)} wavelengths""" 32 | plt.text(1, 0, summary, transform=plt.gca().transAxes, va="bottom") 33 | 34 | 35 | def plot_etc(t): 36 | """ 37 | Make a quick summary plot of ETC tabular results. 38 | """ 39 | plt.figure(figsize=(8, 4)) 40 | i = t.meta["configuration"]["instrument"] 41 | plt.title(f"{i['instrument']}+{i['aperture']}+{i['disperser']} | ETC") 42 | for k in t.colnames: 43 | if "snr" in k: 44 | plt.plot( 45 | t["wavelength"], 46 | t[k], 47 | label=f"{k}\n(median S/N = {np.median(t[k]):.1f})", 48 | ) 49 | plt.yscale("log") 50 | plt.legend(frameon=False, bbox_to_anchor=(1, 1)) 51 | plt.xlabel("Wavelength (microns)") 52 | plt.ylabel("S/N per extracted pixel") 53 | 54 | N_groups = t.meta["configuration"]["detector"]["ngroup"] 55 | summary = f""" 56 | {N_groups} groups/integration 57 | ({N_groups}-1)/({N_groups}+1) = {(N_groups-1)/(N_groups+1):.2%} duty cycle 58 | {'?'}s/group 59 | {'?'}s/integration 60 | {len(t)} wavelengths""" 61 | # {t.meta['configuration']['detector']['nint']} integrations/exposure 62 | plt.text(1, 0, summary, transform=plt.gca().transAxes, va="bottom") 63 | 64 | 65 | def plot_etc_and_pandexo(t_etc, t_pandexo): 66 | """ 67 | Compare S/N estimates between the official ETC and Pandexo. 68 | 69 | Parameters 70 | ---------- 71 | t_etc : Table 72 | The 1D tabular output from `read_etc` 73 | t_pandexo : Table 74 | The 1D tabular output from `read_pandexo` 75 | """ 76 | plt.figure(figsize=(8, 4)) 77 | for t, l in zip([t_etc, t_pandexo], ["ETC", "PandExo"]): 78 | for k in t.colnames: 79 | if "snr_per_integration" in k: 80 | plt.plot(t["wavelength"], t[k], label=f"{k} from {l}") 81 | 82 | plt.xlabel("Wavelength (microns)") 83 | plt.ylabel("S/N per extracted pixel per integration") 84 | plt.legend(frameon=False, bbox_to_anchor=(1, 1)) 85 | # plt.yscale('log') 86 | -------------------------------------------------------------------------------- /chromatic/tests/test_spectra.py: -------------------------------------------------------------------------------- 1 | from .setup_tests import * 2 | from ..imports import * 3 | from ..spectra import * 4 | 5 | 6 | def test_spectral_library_R(cmap=one2another("indigo", "tomato"), N=5): 7 | fi, ax = plt.subplots( 8 | 3, 1, figsize=(8, 6), dpi=300, sharex=True, sharey=True, constrained_layout=True 9 | ) 10 | for R, a in zip([10, 100, 1000], ax): 11 | plt.sca(a) 12 | for i, T in enumerate( 13 | np.round(np.logspace(np.log10(2300), np.log10(12000), N)) 14 | ): 15 | color = cmap(i / (N - 1)) 16 | w, f = get_phoenix_photons(temperature=T, R=R) 17 | plt.loglog(w, f, color=color) 18 | w, f = get_planck_photons(temperature=T, R=R) 19 | plt.plot(w, f, color=color, linestyle=":") 20 | plt.ylim(1e21, 1e25) 21 | plt.xlim(0.05, 5) 22 | plt.text(0.98, 0.92, f"R={R}", transform=a.transAxes, ha="right", va="top") 23 | fi.supxlabel(f"Wavelength ({w.unit.to_string('latex_inline')})") 24 | fi.supylabel(f"Surface Flux ({f.unit.to_string('latex_inline')})") 25 | plt.savefig( 26 | os.path.join( 27 | test_directory, "demonstration-of-spectral-library-with-constant-R.pdf" 28 | ) 29 | ) 30 | 31 | 32 | def test_spectral_library_wavelengths(cmap=one2another("indigo", "tomato"), N=5): 33 | fi, ax = plt.subplots( 34 | 3, 1, figsize=(8, 6), dpi=300, sharex=True, sharey=True, constrained_layout=True 35 | ) 36 | for how_many, a in zip([10, 100, 1000], ax): 37 | plt.sca(a) 38 | wavelength = np.logspace(-1, 0, how_many) * u.micron 39 | for i, T in enumerate( 40 | np.round(np.logspace(np.log10(2300), np.log10(12000), N)) 41 | ): 42 | color = cmap(i / (N - 1)) 43 | w, f = get_phoenix_photons(temperature=T, wavelength=wavelength) 44 | plt.loglog(w, f, color=color) 45 | w, f = get_planck_photons(temperature=T, wavelength=wavelength) 46 | plt.plot(w, f, color=color, linestyle=":") 47 | plt.ylim(1e21, 1e25) 48 | plt.xlim(0.05, 2) 49 | plt.text( 50 | 0.98, 0.92, f"N={how_many}", transform=a.transAxes, ha="right", va="top" 51 | ) 52 | fi.supxlabel(f"Wavelength ({w.unit.to_string('latex_inline')})") 53 | fi.supylabel(f"Surface Flux ({f.unit.to_string('latex_inline')})") 54 | plt.savefig( 55 | os.path.join( 56 | test_directory, 57 | "demonstration-of-spectral-library-with-custom-wavelengths.pdf", 58 | ) 59 | ) 60 | 61 | 62 | def test_spectral_library_loads_correct_R(): 63 | for i in range(2): 64 | w = np.linspace(0.7, 0.71, 1000) * u.micron 65 | wamb, s_amb = get_phoenix_photons( 66 | temperature=int(3900.0), wavelength=w, logg=4.4, metallicity=0.0 67 | ) 68 | wspot, s_spot = get_phoenix_photons( 69 | temperature=int(3100.0), wavelength=w, logg=4.4, metallicity=0.0 70 | ) 71 | plt.figure() 72 | plt.plot(wspot, s_spot) 73 | plt.plot(wamb, s_amb) 74 | plt.savefig( 75 | os.path.join( 76 | test_directory, 77 | "demonstration-of-spectral-library-loading-minimum-necessary-resolution.pdf", 78 | ) 79 | ) 80 | -------------------------------------------------------------------------------- /chromatic/archives/mast.py: -------------------------------------------------------------------------------- 1 | from astroquery.mast import Observations 2 | 3 | __all__ = ["download_from_mast"] 4 | 5 | 6 | def download_from_mast( 7 | proposal_id="2734", 8 | target_name="WASP-96", 9 | obs_collection="JWST", 10 | instrument_name="*", 11 | productSubGroupDescription="X1DINTS", 12 | calib_level=2, 13 | ): 14 | """ 15 | Download (JWST) data products from MAST. 16 | 17 | Given a set of search criteria, find the data products 18 | and download them locally. This has been tested only 19 | minimally for retrieving X1DINTS files for JWST 20 | time-series observations for quick look analyses. 21 | It should hopefully work more broadly? 22 | 23 | If the data Exclusive Access data that are still 24 | in their proprietary period, you will need to 25 | login with `Observations.login()` to authenticate 26 | your connection to the MAST archive. 27 | 28 | For more details on programmatically accessing 29 | data from MAST, it may be useful to explore the 30 | example notebooks available at: 31 | 32 | https://spacetelescope.github.io/mast_notebooks/ 33 | 34 | Parameters 35 | ---------- 36 | 37 | **kw : dict 38 | Remaining keywords will be passed to `Observations.query_critera`. 39 | Available keywords for this initial search likely include: 40 | ['intentType', 'obs_collection', 'provenance_name', 41 | 'instrument_name', 'project', 'filters', 'wavelength_region', 42 | 'target_name', 'target_classification', 'obs_id', 's_ra', 's_dec', 43 | 'proposal_id', 'proposal_pi', 'obs_title', 'dataproduct_type', 44 | 'calib_level', 't_min', 't_max', 't_obs_release', 't_exptime', 45 | 'em_min', 'em_max', 'objID', 's_region', 'jpegURL', 'distance', 46 | 'obsid', 'dataRights', 'mtFlag', 'srcDen', 'dataURL', 47 | 'proposal_type', 'sequence_number'] 48 | 49 | Returns 50 | ------- 51 | downloaded : astropy.table.Table 52 | A table summarizing what was downloaded. The columns 53 | of this table will likely include: 54 | ['Local Path', 'Status', 'Message', 'URL'] 55 | where 'Local Path' indicates to where the file 56 | was downloaded (relative to the local path), 57 | and 'Status' indicates whether the download 58 | was successful. 59 | 60 | """ 61 | 62 | # get a table of all observations matching the search criteria 63 | obs_table = Observations.query_criteria( 64 | obs_collection=obs_collection, 65 | proposal_id=proposal_id, 66 | instrument_name=instrument_name, 67 | target_name=target_name, 68 | ) 69 | 70 | # get a table of the list of products associated with those observations 71 | products = Observations.get_product_list(obs_table["obsid"]) 72 | 73 | # filter those products down to just what's requested 74 | if isinstance(calib_level, int): 75 | calib_level = [calib_level] 76 | filtered_products = Observations.filter_products( 77 | products, 78 | productSubGroupDescription=productSubGroupDescription, 79 | calib_level=calib_level, 80 | ) 81 | 82 | # download the filtered products 83 | downloaded = Observations.download_products(filtered_products) 84 | 85 | # return table summarizing the downloaded files 86 | return downloaded 87 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/xarray_stellar_spectra.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | 3. Edit the `from_xarray_stellar_spectra` function so that it will 4 | load a chromatic light curve file in your format and, 5 | for some Rainbow object `rainbow`, populate at least: 6 | 7 | + rainbow.timelike['time'] 8 | + rainbow.wavelike['wavelength'] 9 | + rainbow.fluxlike['flux'] 10 | 11 | You'll need to replace the cartoon functions on each 12 | line with the actual code needed to load your file. 13 | 14 | (This template assumes that only one file needs to be 15 | loaded. If you need to load multiple segments, or each 16 | time point is stored in its own file or something, then 17 | check out `stsci.py` for an example of loading and 18 | stitching together multiple input files. You'll probably 19 | want to change `filepath` to accept a glob-friendly string 20 | like `my-neato-formatted-files-*.npy` or some such.) 21 | 22 | 4. Edit the `readers/__init__.py` file to import your 23 | `from_xarray_stellar_spectra` function to be accessible when people 24 | are trying to create new Rainbows. Add an `elif` statement 25 | to the `guess_reader` function that will help guess which 26 | reader to use from some aspect(s) of the filename. 27 | 28 | (This `guess_reader` function also accepts a `format=` 29 | keyword that allows the user to explicitly specify that 30 | the xarray_stellar_spectra reader should be used.) 31 | 32 | 5. Submit a pull request to the github repository for 33 | this package, so that other folks can use your handy 34 | new reader too! 35 | """ 36 | 37 | # import the general list of packages 38 | from ...imports import * 39 | from ..writers.xarray_stellar_spectra import xr, json, chromatic_to_ers 40 | 41 | ers_to_chromatic = {v: k for k, v in chromatic_to_ers.items()} 42 | 43 | # define list of the only things that will show up in imports 44 | __all__ = ["from_xarray_stellar_spectra"] 45 | 46 | 47 | def from_xarray_stellar_spectra(self, filepath): 48 | """ 49 | Populate a Rainbow from a file in the xarray_stellar_spectra format. 50 | 51 | Parameters 52 | ---------- 53 | 54 | self : self 55 | The object to be populated. 56 | 57 | filepath : str 58 | The path to the file to load. 59 | """ 60 | 61 | import xarray as xr 62 | 63 | ds = xr.open_dataset(filepath) 64 | 65 | def make_Quantity(da): 66 | """ 67 | Convert a data array into a chromatic quantity (with astropy units). 68 | 69 | """ 70 | unit_string = da.attrs.get("units", "") 71 | if unit_string != "": 72 | unit = u.Unit(unit_string) 73 | else: 74 | unit = 1 75 | 76 | return da.data * unit 77 | 78 | self.wavelike["wavelength"] = make_Quantity(ds["wavelength"]) 79 | self.timelike["time"] = make_Quantity(ds["time"]) 80 | 81 | for key, da in ds.items(): 82 | chromatic_key = ers_to_chromatic.get(key, key) 83 | self._put_array_in_right_dictionary(chromatic_key, make_Quantity(da)) 84 | for k, v in da.attrs.items(): 85 | if k != "units": 86 | metadata_key = f"metadata-for-{chromatic_key}" 87 | try: 88 | self.metadata[metadata_key] 89 | except KeyError: 90 | self.metadata[metadata_key] = dict() 91 | self.metadata[metadata_key][k] = v 92 | 93 | for k, v in ds.attrs.items(): 94 | 95 | self.metadata[k] = json.loads(v) 96 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/timelike/subset.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = [ 4 | "get_for_time", 5 | "get_ok_data_for_time", 6 | ] 7 | 8 | 9 | def get_for_time(self, i, quantity="flux"): 10 | """ 11 | Get `'quantity'` associated with time `'i'`. 12 | 13 | Parameters 14 | ---------- 15 | i : int 16 | The time index to retrieve. 17 | quantity : string 18 | The quantity to retrieve. If it is flux-like, 19 | column 'i' will be returned. If it is wave-like, 20 | the array itself will be returned. 21 | 22 | Returns 23 | ------- 24 | quantity : array, Quantity 25 | The 1D array of 'quantity' corresponding to time 'i'. 26 | """ 27 | z = self.get(quantity) 28 | if np.shape(z) == self.shape: 29 | return z[:, i] 30 | elif len(z) == self.nwave: 31 | return z 32 | else: 33 | raise RuntimeError( 34 | f""" 35 | You tried to retrieve time {i} from '{quantity}', 36 | but this quantity is neither flux-like nor wave-like. 37 | It's not possible to return a wave-like array. Sorry! 38 | """ 39 | ) 40 | 41 | 42 | def get_ok_data_for_time( 43 | self, 44 | i, 45 | x="wavelength", 46 | y="flux", 47 | sigma="uncertainty", 48 | minimum_acceptable_ok=1, 49 | express_badness_with_uncertainty=False, 50 | ): 51 | """ 52 | A small wrapper to get the good data from a time. 53 | 54 | Extract a slice of data, marking data that are not `ok` either 55 | by trimming them out entirely or by inflating their 56 | uncertainties to infinity. 57 | 58 | Parameters 59 | ---------- 60 | i : int 61 | The time index to retrieve. 62 | x : string, optional 63 | What quantity should be retrieved as 'x'? (default = 'time') 64 | y : string, optional 65 | What quantity should be retrieved as 'y'? (default = 'flux') 66 | sigma : string, optional 67 | What quantity should be retrieved as 'sigma'? (default = 'uncertainty') 68 | minimum_acceptable_ok : float, optional 69 | The smallest value of `ok` that will still be included. 70 | (1 for perfect data, 1e-10 for everything but terrible data, 0 for all data) 71 | express_badness_with_uncertainty : bool, optional 72 | If False, data that don't pass the `ok` cut will be removed. 73 | If True, data that don't pass the `ok` cut will have their 74 | uncertainties inflated to infinity (np.inf). 75 | 76 | Returns 77 | ------- 78 | x : array 79 | The time. 80 | y : array 81 | The desired quantity (default is `flux`) 82 | sigma : array 83 | The uncertainty on the desired quantity 84 | """ 85 | 86 | # get 1D independent variable 87 | x_values = self.get_for_time(i, x) * 1 88 | 89 | # get 1D array of what to keep 90 | ok = self.ok[:, i] >= minimum_acceptable_ok 91 | 92 | # get 1D array of the quantity 93 | y_values = self.get_for_time(i, y) * 1 94 | 95 | # get 1D array of uncertainty 96 | sigma_values = self.get_for_time(i, sigma) * 1 97 | 98 | if express_badness_with_uncertainty: 99 | sigma_values[ok == False] = np.inf 100 | return x_values, y_values, sigma_values 101 | else: 102 | return x_values[ok], y_values[ok], sigma_values[ok] 103 | -------------------------------------------------------------------------------- /chromatic/noise/jwst/etc.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | import json 3 | from .extract import * 4 | from .visualize import * 5 | 6 | 7 | def read_etc(directory, t_integration=None, extract=False): 8 | """ 9 | Read the outputs of the JWST ETC into a table and some images. 10 | 11 | Parameters 12 | ---------- 13 | directory : str 14 | The filepath to the directory downloaded from the ETC. 15 | t_integration : float 16 | The integration time in seconds (used to provide a photon-only S/N). 17 | extract : bool 18 | Should we try to extract a S/N from a 2D image? (This seemed to be 19 | necessary, at least a while ago, for MIRI/LRS. If any pixels 20 | saturated it wouldn't return a tabular S/N estimate.) 21 | 22 | Returns 23 | ------- 24 | results : dict 25 | A dictionary containing noise estimates and intermediate ingredients. 26 | The typical keys are: 27 | `1D` = tabular results along the wavelength axis 28 | `2D` = image results along the wavelength axis and one spatial axis 29 | `3D` = cube results along the wavelength axis and two spatial axes 30 | """ 31 | 32 | # find all the FITS files in the directory 33 | fits_filenames = glob.glob(os.path.join(directory, "*/*.fits")) 34 | input_filename = os.path.join(directory, "input.json") 35 | 36 | # load the metadata inputs 37 | with open(input_filename) as f: 38 | metadata = json.load(f) 39 | 40 | spectra = {} 41 | images = {} 42 | cubes = {} 43 | 44 | disperser = metadata["configuration"]["instrument"]["disperser"] 45 | 46 | for f in fits_filenames: 47 | 48 | if "lineplot" in f: 49 | k = os.path.basename(f).split(".fits")[0].replace("lineplot_", "") 50 | if k in ["wave_calc", "target", "fp", "bg", "bg_rate", "total_flux"]: 51 | continue 52 | print(k) 53 | if k == "wave_pix" in k: 54 | spectra["wavelength"] = fits.open(f)[1].data["wavelength"] 55 | else: 56 | spectra[k] = fits.open(f)[1].data[k] 57 | if "image" in f: 58 | k = os.path.basename(f).split(".fits")[0].replace("image_", "") 59 | images[k] = trim_image(fits.open(f)[0].data, disperser=disperser) 60 | if "cube" in f: 61 | k = os.path.basename(f).split(".fits")[0].replace("cube_", "") 62 | cubes[k] = fits.open(f)[0].data 63 | 64 | n_integrations_per_exposure = metadata["configuration"]["detector"]["nint"] 65 | spectra["snr_per_exposure"] = spectra["extracted_flux"] / spectra["extracted_noise"] 66 | spectra["snr_per_integration"] = spectra["snr_per_exposure"] / np.sqrt( 67 | n_integrations_per_exposure 68 | ) 69 | 70 | if t_integration is not None: 71 | signal = spectra["extracted_flux"] * t_int 72 | noise = np.sqrt( 73 | (spectra["extracted_flux"] + spectra["extracted_bg_total"]) * t_int 74 | ) 75 | spectra["snr_per_exposure_from_photons_only"] = signal / noise 76 | spectra["snr_per_integration_from_photons_only"] = ( 77 | signal / noise / np.sqrt(n_integrations_per_exposure) 78 | ) 79 | metadata["time_per_integration"] = t_integration 80 | 81 | t = Table(spectra, meta=metadata) 82 | if extract: 83 | t["snr_extracted_from_image"] = extract_sn_from_image(images) 84 | 85 | return {"1D": t, "2D": images, "3D": cubes} 86 | -------------------------------------------------------------------------------- /chromatic/rainbows/get/wavelike/subset.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = [ 4 | "get_for_wavelength", 5 | "get_ok_data_for_wavelength", 6 | ] 7 | 8 | 9 | def get_for_wavelength(self, i, quantity="flux"): 10 | """ 11 | Get `'quantity'` associated with wavelength `'i'`. 12 | 13 | Parameters 14 | ---------- 15 | i : int 16 | The wavelength index to retrieve. 17 | quantity : string 18 | The quantity to retrieve. If it is flux-like, 19 | row 'i' will be returned. If it is time-like, 20 | the array itself will be returned. 21 | 22 | Returns 23 | ------- 24 | quantity : array, Quantity 25 | The 1D array of 'quantity' corresponding to wavelength 'i'. 26 | """ 27 | z = self.get(quantity) 28 | if np.shape(z) == self.shape: 29 | return z[i, :] 30 | elif len(z) == self.ntime: 31 | return z 32 | else: 33 | raise RuntimeError( 34 | f""" 35 | You tried to retrieve wavelength {i} from '{quantity}', 36 | but this quantity is neither flux-like nor time-like. 37 | It's not possible to return a time-like array. Sorry! 38 | """ 39 | ) 40 | 41 | 42 | def get_ok_data_for_wavelength( 43 | self, 44 | i, 45 | x="time", 46 | y="flux", 47 | sigma="uncertainty", 48 | minimum_acceptable_ok=1, 49 | express_badness_with_uncertainty=False, 50 | ): 51 | """ 52 | A small wrapper to get the good data from a wavelength. 53 | 54 | Extract a slice of data, marking data that are not `ok` either 55 | by trimming them out entirely or by inflating their 56 | uncertainties to infinity. 57 | 58 | Parameters 59 | ---------- 60 | i : int 61 | The wavelength index to retrieve. 62 | x : string, optional 63 | What quantity should be retrieved as 'x'? (default = 'time') 64 | y : string, optional 65 | What quantity should be retrieved as 'y'? (default = 'flux') 66 | sigma : string, optional 67 | What quantity should be retrieved as 'sigma'? (default = 'uncertainty') 68 | minimum_acceptable_ok : float, optional 69 | The smallest value of `ok` that will still be included. 70 | (1 for perfect data, 1e-10 for everything but terrible data, 0 for all data) 71 | express_badness_with_uncertainty : bool, optional 72 | If False, data that don't pass the `ok` cut will be removed. 73 | If True, data that don't pass the `ok` cut will have their 74 | uncertainties inflated to infinity (np.inf). 75 | 76 | Returns 77 | ------- 78 | x : array 79 | The time. 80 | y : array 81 | The desired quantity (default is `flux`) 82 | sigma : array 83 | The uncertainty on the desired quantity 84 | """ 85 | 86 | # get 1D independent variable 87 | x_values = self.get_for_wavelength(i, x) * 1 88 | 89 | # get 1D array of what to keep 90 | ok = self.ok[i, :] >= minimum_acceptable_ok 91 | 92 | # get 1D array of the quantity 93 | y_values = self.get_for_wavelength(i, y) * 1 94 | 95 | # get 1D array of uncertainty 96 | sigma_values = self.get_for_wavelength(i, sigma) * 1 97 | 98 | if express_badness_with_uncertainty: 99 | sigma_values[ok == False] = np.inf 100 | return x_values, y_values, sigma_values 101 | else: 102 | return x_values[ok], y_values[ok], sigma_values[ok] 103 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/xarray_raw_light_curves.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | 3. Edit the `from_xarray_raw_light_curves` function so that it will 4 | load a chromatic light curve file in your format and, 5 | for some Rainbow object `rainbow`, populate at least: 6 | 7 | + rainbow.timelike['time'] 8 | + rainbow.wavelike['wavelength'] 9 | + rainbow.fluxlike['flux'] 10 | 11 | You'll need to replace the cartoon functions on each 12 | line with the actual code needed to load your file. 13 | 14 | (This template assumes that only one file needs to be 15 | loaded. If you need to load multiple segments, or each 16 | time point is stored in its own file or something, then 17 | check out `stsci.py` for an example of loading and 18 | stitching together multiple input files. You'll probably 19 | want to change `filepath` to accept a glob-friendly string 20 | like `my-neato-formatted-files-*.npy` or some such.) 21 | 22 | 4. Edit the `readers/__init__.py` file to import your 23 | `from_xarray_raw_light_curves` function to be accessible when people 24 | are trying to create new Rainbows. Add an `elif` statement 25 | to the `guess_reader` function that will help guess which 26 | reader to use from some aspect(s) of the filename. 27 | 28 | (This `guess_reader` function also accepts a `format=` 29 | keyword that allows the user to explicitly specify that 30 | the xarray_raw_light_curves reader should be used.) 31 | 32 | 5. Submit a pull request to the github repository for 33 | this package, so that other folks can use your handy 34 | new reader too! 35 | """ 36 | 37 | # import the general list of packages 38 | from ...imports import * 39 | from ..writers.xarray_raw_light_curves import xr, json, chromatic_to_ers 40 | 41 | ers_to_chromatic = {v: k for k, v in chromatic_to_ers.items()} 42 | 43 | # define list of the only things that will show up in imports 44 | __all__ = ["from_xarray_raw_light_curves"] 45 | 46 | 47 | def from_xarray_raw_light_curves(self, filepath): 48 | """ 49 | Populate a Rainbow from a file in the xarray_raw_light_curves format. 50 | 51 | Parameters 52 | ---------- 53 | 54 | self : self 55 | The object to be populated. 56 | 57 | filepath : str 58 | The path to the file to load. 59 | """ 60 | 61 | import xarray as xr 62 | 63 | ds = xr.open_dataset(filepath) 64 | 65 | def make_Quantity(da): 66 | """ 67 | Convert a data array into a chromatic quantity (with astropy units). 68 | 69 | """ 70 | unit_string = da.attrs.get("units", "") 71 | if unit_string != "": 72 | unit = u.Unit(unit_string) 73 | else: 74 | unit = 1 75 | 76 | return da.data * unit 77 | 78 | self.wavelike["wavelength"] = make_Quantity( 79 | ds[ers_to_chromatic.get("wavelength", "wavelength")] 80 | ) 81 | self.timelike["time"] = make_Quantity(ds[ers_to_chromatic.get("time", "time")]) 82 | 83 | for key, da in ds.items(): 84 | chromatic_key = ers_to_chromatic.get(key, key) 85 | self._put_array_in_right_dictionary(chromatic_key, make_Quantity(da)) 86 | for k, v in da.attrs.items(): 87 | if k != "units": 88 | metadata_key = f"metadata-for-{chromatic_key}" 89 | try: 90 | self.metadata[metadata_key] 91 | except KeyError: 92 | self.metadata[metadata_key] = dict() 93 | self.metadata[metadata_key][k] = v 94 | 95 | for k, v in ds.attrs.items(): 96 | 97 | self.metadata[k] = json.loads(v) 98 | -------------------------------------------------------------------------------- /chromatic/rainbows/withmodel.py: -------------------------------------------------------------------------------- 1 | from .rainbow import * 2 | 3 | 4 | class RainbowWithModel(Rainbow): 5 | """ 6 | `RainbowWithModel` objects have a fluxlike `model` 7 | attached to them, meaning that they can 8 | 9 | This class definition inherits from `Rainbow`. 10 | """ 11 | 12 | # which fluxlike keys will respond to math between objects 13 | _keys_that_respond_to_math = ["flux", "model"] 14 | 15 | # which keys get uncertainty weighting during binning 16 | _keys_that_get_uncertainty_weighting = ["flux", "model", "uncertainty"] 17 | 18 | @property 19 | def residuals(self): 20 | """ 21 | Calculate the residuals on the fly, 22 | to make sure they're always up to date. 23 | 24 | The residuals are calculated simply 25 | as the `.flux` - `.model`, so they are 26 | in whatever units those arrays have. 27 | 28 | Returns 29 | ------- 30 | residuals : array, Quantity 31 | The 2D array of residuals (nwave, ntime). 32 | """ 33 | return self.flux - self.model 34 | 35 | @property 36 | def chi_squared(self): 37 | """ 38 | Calculate chi-squared. 39 | 40 | This calculates the sum of the squares of 41 | the uncertainty-normalized residuals, 42 | sum(((flux - model)/uncertainty)**2) 43 | 44 | Data points marked as not OK are ignored. 45 | 46 | Returns 47 | ------- 48 | chi_squared : float 49 | The chi-squared value. 50 | """ 51 | r = (self.flux - self.model) / self.uncertainty 52 | return np.sum(r[self.ok] ** 2) 53 | 54 | @property 55 | def residuals_plus_one(self): 56 | """ 57 | A tiny wrapper to get the residuals + 1. 58 | 59 | Returns 60 | ------- 61 | residuals_plus_one : array, Quantity 62 | The 2D array of residuals + 1 (nwave, ntime). 63 | """ 64 | return self.flux - self.model + 1 65 | 66 | @property 67 | def ones(self): 68 | """ 69 | Generate an array of ones that looks like the flux. 70 | (A tiny wrapper needed for `plot_with_model`) 71 | 72 | Returns 73 | ------- 74 | ones : array, Quantity 75 | The 2D array ones (nwave, ntime). 76 | """ 77 | return np.ones_like(self.flux) 78 | 79 | def _validate_core_dictionaries(self): 80 | super()._validate_core_dictionaries() 81 | try: 82 | model = self.get("model") 83 | assert np.shape(model) == np.shape(self.flux) 84 | except (AttributeError, AssertionError): 85 | message = """ 86 | No fluxlike 'model' was found attached to this 87 | `RainbowWithModel` object. The poor thing, 88 | its name is a lie! Please connect a model. 89 | The simplest way to do so might look like... 90 | `rainbow.model = np.ones(rainbow.shape)` 91 | ...or similarly with a more interesting array. 92 | """ 93 | cheerfully_suggest(message) 94 | 95 | from .visualizations import ( 96 | plot_with_model, 97 | plot_with_model_and_residuals, 98 | paint_with_models, 99 | plot_one_wavelength_with_models, 100 | animate_with_models, 101 | ) 102 | 103 | 104 | # REMOVE THE RAINBOW WITH MODEL AND JUST ADD A VALIDATION STEP TO ALL MODEL-DEPENDENT THINGS? 105 | -------------------------------------------------------------------------------- /chromatic/rainbows/converters/for_altair.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["to_nparray", "to_df"] 4 | 5 | 6 | def to_nparray(self, t_unit="d", w_unit="micron"): 7 | """Convert Rainbow object to 1D and 2D numpy arrays 8 | Parameters 9 | ---------- 10 | self : Rainbow object 11 | chromatic Rainbow object to convert into array format 12 | t_unit : str 13 | (optional, default='d') 14 | The time units to use (seconds, minutes, hours, days etc.) 15 | w_unit : str 16 | (optional, default='micron') 17 | The wavelength units to use 18 | Returns 19 | ---------- 20 | rflux : array 21 | flux [n_wavelengths x n_integrations] 22 | rfluxu : array 23 | flux uncertainty [n_wavelengths x n_integrations] 24 | rtime : array 25 | time (t_unit) [n_integrations] 26 | rwavel : array 27 | wavelength (w_unit) [n_wavelengths] 28 | """ 29 | 30 | rflux = np.array( 31 | self.fluxlike["flux"] 32 | ) # flux : [n_wavelengths x n_integrations] 33 | rfluxu = np.array( 34 | self.fluxlike["uncertainty"] 35 | ) # uncertainty : [n_wavelengths x n_integrations] 36 | rtime = self.timelike[ 37 | "time" 38 | ] # time (BJD_TDB, hours) : [n_integrations] 39 | rwavel = self.wavelike[ 40 | "wavelength" 41 | ] # wavelength (microns) : [n_wavelengths] 42 | 43 | try: 44 | # nice bit of code copied from .imshow(), thanks Zach! 45 | # Change time into the units requested by the user 46 | t_unit = u.Unit(t_unit) 47 | except: 48 | cheerfully_suggest("Unrecognised Time Format! Returning day by default") 49 | t_unit = u.Unit("d") 50 | rtime = np.array(rtime.to(t_unit).value) 51 | 52 | try: 53 | # Change wavelength into the units requested by the user 54 | w_unit = u.Unit(w_unit) 55 | except: 56 | cheerfully_suggest( 57 | "Unrecognised Wavelength Format! Returning micron by default" 58 | ) 59 | w_unit = u.Unit("micron") 60 | rwavel = np.array(rwavel.to(w_unit).value) 61 | 62 | return rflux, rfluxu, rtime, rwavel 63 | 64 | 65 | def to_df(self, t_unit="d", w_unit="micron"): 66 | """Convert Rainbow object to pandas dataframe 67 | Parameters 68 | ---------- 69 | self : Rainbow object 70 | chromatic Rainbow object to convert into pandas df format 71 | t_unit : str 72 | (optional, default='d') 73 | The time units to use (seconds, minutes, hours, days etc.) 74 | w_unit : str 75 | (optional, default='micron') 76 | The wavelength units to use 77 | Returns 78 | ---------- 79 | pd.DataFrame 80 | The rainbow object flattened and converted into a pandas dataframe 81 | """ 82 | # extract 1D/2D formats for the flux, uncertainty, time and wavelengths 83 | rflux, rfluxu, rtime, rwavel = self.to_nparray(t_unit=t_unit, w_unit=w_unit) 84 | 85 | # put all arrays onto same dimensions 86 | x, y = np.meshgrid(rtime, rwavel) 87 | rainbow_dict = { 88 | f"Time ({t_unit})": x.ravel(), 89 | f"Wavelength ({w_unit})": y.ravel(), 90 | "Flux": rflux.ravel(), 91 | "Flux Uncertainty": rfluxu.ravel(), 92 | } 93 | 94 | # convert to pandas dataframe 95 | df = pd.DataFrame(rainbow_dict) 96 | return df 97 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/nres.py: -------------------------------------------------------------------------------- 1 | """ 2 | Read multiple exposures of a single order of LCO/NRES spectra. 3 | """ 4 | 5 | # import the general list of packages 6 | from ...imports import * 7 | 8 | # define list of the only things that will show up in imports 9 | __all__ = ["from_nres"] 10 | 11 | 12 | def from_nres(rainbow, filepath, order=52): 13 | """ 14 | Populate a Rainbow from a file in the NRES format. 15 | 16 | Parameters 17 | ---------- 18 | 19 | rainbow : Rainbow 20 | The object to be populated. 21 | 22 | filepath : str 23 | The path to the file to load. 24 | 25 | order : float 26 | The order to extract. 27 | Acceptable values span 52 to 119 28 | """ 29 | 30 | # read in your file, however you like 31 | filenames = glob.glob(filepath) 32 | filenames = np.sort(filenames) 33 | 34 | order_index = order - 52 # the 0th order is order 52 35 | 36 | for i, f in tqdm(enumerate(filenames), leave=False): 37 | 38 | # open this 39 | hdu = fits.open(f) 40 | 41 | # get the time associated with this observation 42 | date = hdu["PRIMARY"].header["MJD-OBS"] 43 | date_bjd = date + 2400000.5 44 | 45 | # get the science spectrum 46 | science_fiber_id = hdu["PRIMARY"].header["SCIFIBER"] 47 | is_science = hdu["SPECTRUM"].data["FIBER"] == science_fiber_id 48 | data = hdu["SPECTRUM"].data[is_science] 49 | this_order = data[order_index] 50 | order = this_order["ORDER"] 51 | ok = this_order["MASK"] == 0 52 | 53 | # these should be the good data for the particular order in question 54 | wavelengths = np.array(this_order["WAVELENGTH"]) * u.angstrom # [mask] 55 | normalized_fluxes = this_order["NORMFLUX"] # [mask] 56 | normalized_uncertainties = np.abs(this_order["NORMUNCERTAINTY"]) # [mask]) 57 | 58 | sort = np.argsort(wavelengths.value) 59 | wavelengths = wavelengths[sort] 60 | fluxes = normalized_fluxes[sort] 61 | errors = normalized_uncertainties[sort] 62 | 63 | if i == 0: 64 | ntimes = len(filenames) 65 | nwaves = len(wavelengths) 66 | for k in ["wavelength_2d", "flux", "uncertainty", "ok"]: 67 | rainbow.fluxlike[k] = np.zeros((nwaves, ntimes)) 68 | rainbow.fluxlike["wavelength_2d"] *= u.micron 69 | rainbow.fluxlike["ok"] = rainbow.fluxlike["ok"].astype(np.bool) 70 | rainbow.timelike["time"] = np.zeros(ntimes) * u.day 71 | 72 | # populate a 1D array of times (with astropy units of time) 73 | rainbow.timelike["time"][i] = date_bjd * u.day * 1 74 | 75 | # populate a 2D (row = wavelength, col = time, value = wavelength) 76 | rainbow.fluxlike["wavelength_2d"][:, i] = wavelengths * 1 77 | 78 | # populate a 2D (row = wavelength, col = time) array of fluxes 79 | rainbow.fluxlike["flux"][:, i] = fluxes * 1 80 | 81 | # populate a 2D (row = wavelength, col = time) array of uncertainties 82 | rainbow.fluxlike["uncertainty"][:, i] = errors * 1 83 | 84 | # populate a 2D (row = wavelength, col = time) array of ok 85 | rainbow.fluxlike["ok"][:, i] = ok * 1 86 | 87 | # populate a 1D array of wavelengths (with astropy units of length) 88 | rainbow.wavelike["wavelength"] = np.nanmedian( 89 | rainbow.fluxlike["wavelength_2d"], axis=rainbow.timeaxis 90 | ) 91 | 92 | # add some warnings if there's any funny business 93 | if len(filenames) == 0: 94 | cheerfully_suggest( 95 | f""" 96 | There are no files of that name in this folder! 97 | """ 98 | ) 99 | -------------------------------------------------------------------------------- /chromatic/spectra/planck.py: -------------------------------------------------------------------------------- 1 | from ..imports import * 2 | 3 | __all__ = ["calculate_planck_flux", "get_planck_photons"] 4 | 5 | 6 | def calculate_planck_flux(wavelength, temperature): 7 | """ 8 | Calculate the surface flux from a thermally emitted surface, 9 | according to Planck function. 10 | 11 | Parameters 12 | ---------- 13 | wavelength : Quantity 14 | The wavelengths at which to calculate, 15 | with units of wavelength. 16 | temperature : Quantity 17 | The temperature of the thermal emitter, 18 | with units of K. 19 | 20 | Returns 21 | ------- 22 | surface_flux : Quantity 23 | The surface flux, evaluated at the wavelengths. 24 | """ 25 | 26 | # define variables as shortcut to the constants we need 27 | h = con.h 28 | k = con.k_B 29 | c = con.c 30 | 31 | # the thing that goes in the exponent (its units better cancel!) 32 | z = h * c / (wavelength * k * temperature) 33 | 34 | # calculate the intensity from the Planck function 35 | intensity = (2 * h * c**2 / wavelength**5 / (np.exp(z) - 1)) / u.steradian 36 | 37 | # calculate the flux assuming isotropic emission 38 | flux = np.pi * u.steradian * intensity 39 | 40 | # return the intensity 41 | return flux.to("W/(m**2*nm)") 42 | 43 | 44 | def get_planck_photons( 45 | temperature=3000, wavelength=None, R=100, wlim=[0.04, 6] * u.micron, **kw 46 | ): 47 | """ 48 | Calculate the surface flux from a thermally emitted surface, 49 | according to Planck function, in units of photons/(s * m**2 * nm). 50 | 51 | Parameters 52 | ---------- 53 | temperature : Quantity 54 | The temperature of the thermal emitter, 55 | with units of K. 56 | wavelength : Quantity, optional 57 | The wavelengths at which to calculate, 58 | with units of wavelength. 59 | R : float, optional 60 | The spectroscopic resolution for creating a log-uniform 61 | grid that spans the limits set by `wlim`, only if 62 | `wavelength` is not defined. 63 | wlim : Quantity, optional 64 | The two-element [lower, upper] limits of a wavelength 65 | grid that would be populated with resolution `R`, only if 66 | `wavelength` is not defined. 67 | **kw : dict, optional 68 | Other keyword arguments will be ignored. 69 | 70 | Returns 71 | ------- 72 | photons : Quantity 73 | The surface flux in photon units 74 | 75 | This evaluates the Planck function at the exact 76 | wavelength values; it doesn't do anything fancy to integrate 77 | over binwidths, so if you're using very wide (R~a few) bins 78 | your integrated fluxes will be messed up. 79 | 80 | """ 81 | 82 | # make sure the temperature unit is good (whether or not it's supplied) 83 | temperature_unit = u.Quantity(temperature).unit 84 | if temperature_unit == u.K: 85 | temperature_with_unit = temperature 86 | elif temperature_unit == u.Unit(""): 87 | temperature_with_unit = temperature * u.K 88 | 89 | # create a wavelength grid if one isn't supplied 90 | if wavelength is None: 91 | wavelength_unit = wlim.unit 92 | wavelength = ( 93 | np.exp(np.arange(np.log(wlim[0].value), np.log(wlim[1].value), 1 / R)) 94 | * wavelength_unit 95 | ) 96 | 97 | energy = calculate_planck_flux( 98 | wavelength=wavelength, temperature=temperature_with_unit 99 | ) 100 | photon_energy = con.h * con.c / wavelength / u.ph 101 | 102 | return wavelength, (energy / photon_energy).to(u.ph / (u.s * u.m**2 * u.nm)) 103 | -------------------------------------------------------------------------------- /docs/tools/colormaps.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "42179ce5", 6 | "metadata": {}, 7 | "source": [ 8 | "# Creating Custom Colormaps\n", 9 | "\n", 10 | "Colormaps provide a way to translate from numbers to colors. The `matplotlib` [colormaps](https://matplotlib.org/stable/tutorials/colors/colormaps.html) are lovely, but sometimes we just want to say \"I'd like a colormap that goes from this color to that color.\" The `one2another` colormap generator does just that." 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": null, 16 | "id": "39a565c9", 17 | "metadata": {}, 18 | "outputs": [], 19 | "source": [ 20 | "from chromatic import one2another, version" 21 | ] 22 | }, 23 | { 24 | "cell_type": "code", 25 | "execution_count": null, 26 | "id": "a695119f", 27 | "metadata": {}, 28 | "outputs": [], 29 | "source": [ 30 | "version()" 31 | ] 32 | }, 33 | { 34 | "cell_type": "markdown", 35 | "id": "f06ba943", 36 | "metadata": {}, 37 | "source": [ 38 | "## How do we make a new colormap? \n", 39 | "\n", 40 | "All we need to specify is the color at the bottom and the color at the top. If we want to get really fancy, we can specify different `alpha` (= opacity) values for either of these limits." 41 | ] 42 | }, 43 | { 44 | "cell_type": "code", 45 | "execution_count": null, 46 | "id": "1643a4d9", 47 | "metadata": {}, 48 | "outputs": [], 49 | "source": [ 50 | "one2another(\"orchid\", \"black\")" 51 | ] 52 | }, 53 | { 54 | "cell_type": "code", 55 | "execution_count": null, 56 | "id": "6926732c", 57 | "metadata": {}, 58 | "outputs": [], 59 | "source": [ 60 | "one2another(\"black\", \"orchid\")" 61 | ] 62 | }, 63 | { 64 | "cell_type": "code", 65 | "execution_count": null, 66 | "id": "41ded69d", 67 | "metadata": {}, 68 | "outputs": [], 69 | "source": [ 70 | "one2another(bottom=\"orchid\", top=\"orchid\", alpha_bottom=0)" 71 | ] 72 | }, 73 | { 74 | "cell_type": "markdown", 75 | "id": "606c2d57", 76 | "metadata": {}, 77 | "source": [ 78 | "That's it! You can use these colormaps anywhere you'd set `cmap=` in a `matplotlib` function." 79 | ] 80 | }, 81 | { 82 | "cell_type": "code", 83 | "execution_count": null, 84 | "id": "2b50f6c0", 85 | "metadata": {}, 86 | "outputs": [], 87 | "source": [ 88 | "import numpy as np, matplotlib.pyplot as plt\n", 89 | "\n", 90 | "# make a custom cmap\n", 91 | "my_fancy_cmap = one2another(\"orchid\", \"black\")\n", 92 | "\n", 93 | "# make some fake data\n", 94 | "x, y = np.random.uniform(-1, 1, [2, 1000])\n", 95 | "z = np.random.normal(0, 1, [10, 10])\n", 96 | "\n", 97 | "# use the cmap\n", 98 | "fi, ax = plt.subplots(1, 2, constrained_layout=True)\n", 99 | "ax[0].scatter(x, y, c=x**2 + y**2, cmap=my_fancy_cmap)\n", 100 | "ax[0].axis(\"scaled\")\n", 101 | "ax[1].imshow(z, cmap=my_fancy_cmap)\n", 102 | "ax[1].axis(\"scaled\");" 103 | ] 104 | } 105 | ], 106 | "metadata": { 107 | "kernelspec": { 108 | "display_name": "Python 3 (ipykernel)", 109 | "language": "python", 110 | "name": "python3" 111 | }, 112 | "language_info": { 113 | "codemirror_mode": { 114 | "name": "ipython", 115 | "version": 3 116 | }, 117 | "file_extension": ".py", 118 | "mimetype": "text/x-python", 119 | "name": "python", 120 | "nbconvert_exporter": "python", 121 | "pygments_lexer": "ipython3", 122 | "version": "3.10.4" 123 | } 124 | }, 125 | "nbformat": 4, 126 | "nbformat_minor": 5 127 | } 128 | -------------------------------------------------------------------------------- /chromatic/tests/test_operations.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | import pytest 4 | 5 | 6 | def test_rainbow_operations(): 7 | nw, nt = 20, 40 8 | a = Rainbow( 9 | wavelength=np.linspace(0.5, 5, nw) * u.micron, 10 | time=np.linspace(-1, 1, nt) * u.hour, 11 | flux=np.ones((nw, nt)), 12 | ) 13 | 14 | b = Rainbow( 15 | wavelength=np.linspace(0.5, 5, nw) * u.micron, 16 | time=np.linspace(-1, 1, nt) * u.hour, 17 | flux=np.zeros((nw, nt)) + 0.1, 18 | ) 19 | 20 | wl_like = np.linspace(-0.01, 0.01, nw) 21 | t_like = np.linspace(-0.01, 0.01, nt) 22 | fx_like = np.ones(a.shape) 23 | 24 | assert (a + b).fluxlike["flux"][0][0] == 1.1 25 | assert (a + wl_like).fluxlike["flux"][0][0] == 0.99 26 | assert (a + t_like).fluxlike["flux"][0][0] == 0.99 27 | assert (a + fx_like).fluxlike["flux"][0][0] == 2 28 | assert (a + 1).fluxlike["flux"][0][0] == 2 29 | 30 | assert (a - b).fluxlike["flux"][0][0] == 0.9 31 | assert (a - wl_like).fluxlike["flux"][0][0] == 1.01 32 | assert (a - t_like).fluxlike["flux"][0][0] == 1.01 33 | assert (a - fx_like).fluxlike["flux"][0][0] == 0 34 | assert (a - 1).fluxlike["flux"][0][0] == 0 35 | 36 | assert (a * b).fluxlike["flux"][0][0] == 0.1 37 | assert (a * wl_like).fluxlike["flux"][0][0] == -0.01 38 | assert (a * t_like).fluxlike["flux"][0][0] == -0.01 39 | assert (a * fx_like).fluxlike["flux"][0][0] == 1 40 | assert (a * 1).fluxlike["flux"][0][0] == 1 41 | 42 | assert (a / b).fluxlike["flux"][0][0] == 1 / 0.1 43 | assert (a / wl_like).fluxlike["flux"][0][0] == 1 / -0.01 44 | assert (a / t_like).fluxlike["flux"][0][0] == 1 / -0.01 45 | assert (a / fx_like).fluxlike["flux"][0][0] == 1 46 | assert (a / 1).fluxlike["flux"][0][0] == 1 47 | 48 | 49 | def test_rainbow_operations_warnings(): 50 | 51 | nw = 20 52 | 53 | # make sure we raise an error if it's not obvious whether we're doing wavelength or time 54 | with pytest.warns(match="reconsider letting them have the same size"): 55 | c = Rainbow( 56 | wavelength=np.linspace(0.5, 5, nw) * u.micron, 57 | time=np.linspace(-1, 1, nw) * u.hour, 58 | flux=np.ones((nw, nw)), 59 | ) 60 | 61 | # with pytest.warns(match="we can't tell which is which"): 62 | wl_like = np.linspace(-0.01, 0.01, nw) 63 | 64 | with pytest.raises(Exception): 65 | c * wl_like 66 | with pytest.raises(Exception): 67 | c / wl_like 68 | with pytest.raises(Exception): 69 | c + wl_like 70 | with pytest.raises(Exception): 71 | c - wl_like 72 | 73 | with pytest.warns(match="do not exactly match."): 74 | d = c._create_copy() 75 | d.wavelength *= 1.01 76 | d.time *= 0.99 77 | c + d 78 | 79 | with pytest.raises(ValueError): 80 | c[1:, :] + d 81 | 82 | 83 | def test_operations_with_uncertainty(): 84 | a = SimulatedRainbow().inject_noise(signal_to_noise=100) + 1 85 | b = SimulatedRainbow().inject_noise(signal_to_noise=50) * 0.5 86 | for x in [ 87 | "a", 88 | "b", 89 | "(a+0)", 90 | "(a+1)", 91 | "(a+b)", 92 | "(a-b)", 93 | "(a*b)", 94 | "(a/b)", 95 | "(a+b)/(a-b)*2", 96 | ]: 97 | print(f"{x:^38}") 98 | r = eval(x) 99 | 100 | print(f" mean(flux) = {np.mean(r.flux)}") 101 | print(f" mean(model) = {np.mean(r.model)}") 102 | print(f"mean(uncertainty) = {np.mean(r.uncertainty)}") 103 | print() 104 | 105 | scaled_sigma = np.std(r.residuals / r.uncertainty) 106 | assert np.isclose(scaled_sigma, 1, rtol=0.02) 107 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/colors.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = [ 4 | "setup_wavelength_colors", 5 | "_make_sure_cmap_is_defined", 6 | "get_wavelength_color", 7 | ] 8 | 9 | 10 | def setup_wavelength_colors(self, cmap=None, vmin=None, vmax=None, log=None): 11 | """ 12 | Set up a color map and normalization function for 13 | colors datapoints by their wavelengths. 14 | 15 | Parameters 16 | ---------- 17 | cmap : str, Colormap 18 | The color map to use. 19 | vmin : Quantity 20 | The wavelength at the bottom of the cmap. 21 | vmax : Quantity 22 | The wavelength at the top of the cmap. 23 | log : bool 24 | If True, colors will scale with log(wavelength). 25 | If False, colors will scale with wavelength. 26 | If None, the scale will be guessed from the internal wscale. 27 | """ 28 | 29 | # populate the cmap object 30 | self.cmap = plt.colormaps.get_cmap(cmap) 31 | 32 | vmin = vmin 33 | if vmin is None: 34 | vmin = np.nanmin(self.wavelength) 35 | vmax = vmax 36 | if vmax is None: 37 | vmax = np.nanmax(self.wavelength) 38 | 39 | if (self.wscale in ["log"]) or (log == True): 40 | self.norm = col.LogNorm( 41 | vmin=vmin.to("micron").value, vmax=vmax.to("micron").value 42 | ) 43 | if (self.wscale in ["?", "linear"]) or (log == False): 44 | self.norm = col.Normalize( 45 | vmin=vmin.to("micron").value, vmax=vmax.to("micron").value 46 | ) 47 | 48 | 49 | def _make_sure_cmap_is_defined(self, cmap=None, vmin=None, vmax=None): 50 | """ 51 | A helper function that can be called at the start of 52 | any plot that that's using wavelength-colors to make 53 | sure that the wavelength-based colormap has been 54 | defined. 55 | 56 | Parameters 57 | ---------- 58 | cmap : str, Colormap 59 | The color map to use for expressing wavelength. 60 | vmin : Quantity 61 | The minimum value to use for the wavelength colormap. 62 | vmax : Quantity 63 | The maximum value to use for the wavelength colormap. 64 | """ 65 | 66 | if hasattr(self, "cmap"): 67 | if (cmap is not None) or (vmin is not None) or (vmax is not None): 68 | if ( 69 | (cmap != self.get("cmap")) 70 | and (vmin != self.get("vmin")) 71 | and (vmax != self.get("vmax")) 72 | ): 73 | cheerfully_suggest( 74 | """ 75 | It looks like you're trying to set up a new custom 76 | cmap and/or wavelength normalization scheme. You 77 | should be aware that a cmap has already been defined 78 | for this object; if you're visualizing the same 79 | rainbow in different ways, we strongly suggest 80 | that you not change the cmap or normalization 81 | between them, for visual consistency. 82 | """ 83 | ) 84 | else: 85 | return 86 | self.setup_wavelength_colors(cmap=cmap, vmin=vmin, vmax=vmax) 87 | 88 | 89 | def get_wavelength_color(self, wavelength): 90 | """ 91 | Determine the color corresponding to one or more wavelengths. 92 | 93 | Parameters 94 | ---------- 95 | wavelength : Quantity 96 | The wavelength value(s), either an individual 97 | wavelength or an array of N wavelengths. 98 | 99 | Returns 100 | ------- 101 | colors : array 102 | An array of RGBA colors [or an (N,4) array]. 103 | """ 104 | w_unitless = wavelength.to("micron").value 105 | normalized_w = self.norm(w_unitless) 106 | return self.cmap(normalized_w) 107 | -------------------------------------------------------------------------------- /docs/tools/transits.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "id": "b9327a8a", 6 | "metadata": {}, 7 | "source": [ 8 | "# Making Model Transits \n", 9 | "\n", 10 | "When making simulated datasets, the `.inject_transit` action can use two different functions for making transits. We make those functions directly available, in case you want to use them for your own wonderful purposes. The options are:\n", 11 | "- `exoplanet_transit` = limb-darkened transits generated with [`exoplanet_core`](https://github.com/exoplanet-dev/exoplanet-core)\n", 12 | "- `trapezoidal_transit` = a very simple non-limb-darkened trapezoidal transit as in [Winn (2010)](https://ui.adsabs.harvard.edu/abs/2010exop.book...55W/abstract)\n", 13 | "\n", 14 | "You should probably use `exoplanet_transit`, unless you have a good reason to want to use a cartoon trapezoid instead." 15 | ] 16 | }, 17 | { 18 | "cell_type": "code", 19 | "execution_count": null, 20 | "id": "68e8e313", 21 | "metadata": {}, 22 | "outputs": [], 23 | "source": [ 24 | "from chromatic import exoplanet_transit, trapezoidal_transit\n", 25 | "import numpy as np, matplotlib.pyplot as plt" 26 | ] 27 | }, 28 | { 29 | "cell_type": "markdown", 30 | "id": "3db25f93", 31 | "metadata": {}, 32 | "source": [ 33 | "The first argument to each of these model function is the array of times for which the flux should be computed, and the remaining keyword arguments allow you to change the planet parameters." 34 | ] 35 | }, 36 | { 37 | "cell_type": "code", 38 | "execution_count": null, 39 | "id": "59a92204", 40 | "metadata": {}, 41 | "outputs": [], 42 | "source": [ 43 | "t = np.linspace(-0.1, 0.1, 1000)\n", 44 | "plt.figure(figsize=(8, 3))\n", 45 | "for rp, c in zip([0.09, 0.1, 0.11], [\"orange\", \"red\", \"purple\"]):\n", 46 | " limb_darkened_model = exoplanet_transit(t, rp=rp)\n", 47 | " plt.plot(t, limb_darkened_model, color=c, label=f\"$R_p$ = {rp:.2f}\")\n", 48 | " trapezoid_model = trapezoidal_transit(t, delta=rp**2)\n", 49 | " plt.plot(t, trapezoid_model, color=c, linestyle=\"--\", label=f\"$R_p$ = {rp:.2f}\")\n", 50 | "plt.xlabel(\"Time (days)\")\n", 51 | "plt.ylabel(\"Relative Flux\")\n", 52 | "plt.legend(frameon=False);" 53 | ] 54 | }, 55 | { 56 | "cell_type": "markdown", 57 | "id": "f7a12e8d", 58 | "metadata": {}, 59 | "source": [ 60 | "## " 61 | ] 62 | }, 63 | { 64 | "cell_type": "markdown", 65 | "id": "c5152191", 66 | "metadata": {}, 67 | "source": [ 68 | "These two functions take different keywords, as explained in their docstrings." 69 | ] 70 | }, 71 | { 72 | "cell_type": "code", 73 | "execution_count": null, 74 | "id": "b8d1b3a4", 75 | "metadata": {}, 76 | "outputs": [], 77 | "source": [ 78 | "exoplanet_transit?" 79 | ] 80 | }, 81 | { 82 | "cell_type": "code", 83 | "execution_count": null, 84 | "id": "b4ab2373", 85 | "metadata": {}, 86 | "outputs": [], 87 | "source": [ 88 | "trapezoidal_transit?" 89 | ] 90 | }, 91 | { 92 | "cell_type": "markdown", 93 | "id": "57e733c7", 94 | "metadata": {}, 95 | "source": [ 96 | "Have fun!" 97 | ] 98 | } 99 | ], 100 | "metadata": { 101 | "kernelspec": { 102 | "display_name": "exoatlas", 103 | "language": "python", 104 | "name": "python3" 105 | }, 106 | "language_info": { 107 | "codemirror_mode": { 108 | "name": "ipython", 109 | "version": 3 110 | }, 111 | "file_extension": ".py", 112 | "mimetype": "text/x-python", 113 | "name": "python", 114 | "nbconvert_exporter": "python", 115 | "pygments_lexer": "ipython3", 116 | "version": "3.13.2" 117 | } 118 | }, 119 | "nbformat": 4, 120 | "nbformat_minor": 5 121 | } 122 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/flag_outliers.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | from scipy.special import erfc 3 | 4 | __all__ = ["flag_outliers"] 5 | 6 | 7 | def flag_outliers(self, how_many_sigma=5, remove_trends=True, inflate_uncertainty=True): 8 | """ 9 | Flag outliers as not `ok`. 10 | 11 | This examines the flux array, identifies significant outliers, 12 | and marks them 0 in the `ok` array. The default procedure is to use 13 | a median filter to remove temporal trends (`remove_trends`), 14 | inflate the uncertainties based on the median-absolute-deviation 15 | scatter (`inflate_uncertainty`), and call points outliers if they 16 | deviate by more than a certain number of sigma (`how_many_sigma`) 17 | from the median-filtered level. 18 | 19 | The returned `Rainbow` object should be identical to the input 20 | one, except for the possibility that some elements in `ok` array 21 | will have been marked as zero. (The filtering or inflation are 22 | not applied to the returned object.) 23 | 24 | Parameters 25 | ---------- 26 | how_many_sigma : float, optional 27 | Standard deviations (sigmas) allowed for individual data 28 | points before they are flagged as outliers. 29 | remove_trends : bool, optional 30 | Should we remove trends from the flux data before 31 | trying to look for outliers? 32 | inflate_uncertainty : bool, optional 33 | Should uncertainties per wavelength be inflated to 34 | match the (MAD-based) standard deviation of the data? 35 | 36 | Returns 37 | ------- 38 | rainbow : Rainbow 39 | A new Rainbow object with the outliers flagged as 0 in `.ok` 40 | """ 41 | 42 | # create a history entry for this action (before other variables are defined) 43 | h = self._create_history_entry("flag_outliers", locals()) 44 | 45 | # create a copy of the existing rainbow 46 | new = self._create_copy() 47 | 48 | # how many outliers are expected from noise alone 49 | outliers_expected_from_normal_distribution = erfc(how_many_sigma) * self.nflux * 2 50 | if outliers_expected_from_normal_distribution >= 1: 51 | cheerfully_suggest( 52 | f""" 53 | When drawing from a normal distribution, an expected {outliers_expected_from_normal_distribution:.1f} out of 54 | the total {self.nflux} datapoints in {self} would be marked 55 | as a >{how_many_sigma} sigma outlier. 56 | 57 | If you don't want to accidentally clip legitimate data points that 58 | might have arisen merely by chance, please consider setting the 59 | outlier flagging threshold (`sigma=`) to a larger value. 60 | """ 61 | ) 62 | 63 | # create a trend-filtered object 64 | if remove_trends: 65 | filtered = new.remove_trends(method="median_filter", size=(3, 5)) 66 | else: 67 | filtered = new._create_copy() 68 | 69 | # update the uncertainties, if need be 70 | if np.all(filtered.uncertainty == 0): 71 | filtered.uncertainty = ( 72 | np.ones(filtered.shape) 73 | * filtered.get_measured_scatter(method="MAD")[:, np.newaxis] 74 | ) 75 | inflate_uncertainty = False 76 | 77 | # inflate the per-wavelength uncertainties, as needed 78 | if inflate_uncertainty: 79 | with warnings.catch_warnings(): 80 | warnings.simplefilter("ignore") 81 | inflated = filtered.inflate_uncertainty(method="MAD", remove_trends=True) 82 | else: 83 | inflated = filtered 84 | 85 | # decide which points are outliers 86 | is_outlier = np.abs(inflated.flux - 1) > how_many_sigma * inflated.uncertainty 87 | 88 | # update the output object 89 | new.fluxlike["flagged_as_outlier"] = is_outlier 90 | new.ok = new.ok * (is_outlier == False) 91 | 92 | # append the history entry to the new Rainbow 93 | new._record_history_entry(h) 94 | 95 | return new 96 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/inflate_uncertainty.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["inflate_uncertainty"] 4 | 5 | 6 | def inflate_uncertainty( 7 | self, 8 | method="MAD", 9 | remove_trends=True, 10 | remove_trends_method="median_filter", 11 | remove_trends_kw={}, 12 | minimum_inflate_ratio=1.0, 13 | ): 14 | """ 15 | Inflate uncertainties to match observed scatter. 16 | 17 | This is a quick and approximate tool for inflating 18 | the flux uncertainties in a `Rainbow` to match the 19 | observed scatter. With defaults, this will estimate 20 | the scatter using a robust median-absolute-deviation 21 | estimate of the standard deviation (`method='MAD'`), 22 | applied to time-series from which temporal trends 23 | have been removed (`remove_trends=True`), and inflate 24 | the uncertainties on a per-wavelength basis. The trend 25 | removal, by default by subtracting off local medians 26 | (`remove_trends_method='median_filter'`), will squash 27 | many types of both astrophysical and systematic trends, 28 | so this function should be used with caution in 29 | applicants where precise and reliable uncertainties 30 | are needed. 31 | 32 | Parameters 33 | ---------- 34 | method : string 35 | What method to use to obtain measured scatter. 36 | Current options are 'MAD', 'standard-deviation'. 37 | remove_trends : bool 38 | Should we remove trends before estimating by how 39 | much we need to inflate the uncertainties? 40 | remove_trends_method : str 41 | What method should be used to remove trends? 42 | See `.remove_trends` for options. 43 | remove_trends_kw : dict 44 | What keyword arguments should be passed to `remove_trends`? 45 | minimum_inflate_ratio : float, optional 46 | the minimum inflate_ratio that can be. We don't want people 47 | to deflate uncertainty unless a very specific case of unstable 48 | pipeline output. 49 | 50 | Returns 51 | ------- 52 | removed : Rainbow 53 | The Rainbow with estimated signals removed. 54 | """ 55 | 56 | # create a history entry for this action (before other variables are defined) 57 | h = self._create_history_entry("inflate_uncertainty", locals()) 58 | 59 | # create a new copy 60 | new = self._create_copy() 61 | 62 | # if desired, remove trends before estimating inflation factor 63 | if remove_trends: 64 | trend_removed = new.remove_trends(**remove_trends_kw) 65 | else: 66 | trend_removed = new 67 | 68 | # estimate the scatter 69 | measured_scatter = trend_removed.get_measured_scatter( 70 | method=method, minimum_acceptable_ok=1e-10 71 | ) 72 | 73 | # get the expected uncertainty 74 | expected_uncertainty = trend_removed.get_expected_uncertainty() 75 | 76 | # calculate the necessary inflation ratio 77 | inflate_ratio = measured_scatter / expected_uncertainty 78 | 79 | # warn if there are some inflation ratios below minimum (usually = 1) 80 | if np.min(inflate_ratio) < minimum_inflate_ratio: 81 | cheerfully_suggest( 82 | f""" 83 | {np.sum(inflate_ratio < minimum_inflate_ratio)} uncertainty inflation ratios would be below 84 | the `minimum_inflate_ratio` of {minimum_inflate_ratio}, so they have not been changed. 85 | """ 86 | ) 87 | inflate_ratio = np.maximum(inflate_ratio, minimum_inflate_ratio) 88 | 89 | # store the inflation ratio 90 | new.wavelike["inflate_ratio"] = inflate_ratio 91 | 92 | # inflate the uncertainties 93 | new.uncertainty = new.uncertainty * inflate_ratio[:, np.newaxis] 94 | 95 | # append the history entry to the new Rainbow 96 | new._record_history_entry(h) 97 | 98 | # return the new Rainbow 99 | return new 100 | -------------------------------------------------------------------------------- /chromatic/rainbows/visualizations/diagnostics/plot_quantities.py: -------------------------------------------------------------------------------- 1 | from ....imports import * 2 | 3 | __all__ = ["plot_quantities"] 4 | 5 | 6 | def plot_quantities( 7 | self, 8 | quantities=None, 9 | xaxis="time", 10 | maxcol=1, 11 | panel_size=(6, 2), 12 | what_is_x="index", 13 | filename=None, 14 | **kw, 15 | ): 16 | """ 17 | Plot {xaxis}-like quantities as a function of {xaxis} index 18 | or any other {xaxis} quantity (such as "time" or "wavelength"). 19 | 20 | Parameters 21 | ---------- 22 | quantities : list like 23 | The {xaxis}-like quantity to plot. 24 | xaxis : string 25 | Whether the quantities are alike to 'time' or 'wave'. Default is 'time'. (Optional) 26 | maxcol : int 27 | The maximum number of columns to show (Optional). 28 | panel_size : tuple 29 | The size in inches dedicated to each panel, default is (6,2). (Optional) 30 | what_is_x : string 31 | The quantity to plot on the what_is_x, default is index. (Optional) 32 | 33 | """ 34 | # decide which dictionary to plot 35 | if xaxis not in ["time", "wave", "wavelength"]: 36 | raise Exception("Unknown xaxis. Choose from [time, wave]") 37 | elif xaxis == "time": 38 | like_dict = self.timelike 39 | else: 40 | like_dict = self.wavelike 41 | 42 | # decide which quantities to plot 43 | if quantities is None: 44 | allkeys = like_dict.keys() 45 | else: 46 | allkeys = quantities[:] 47 | 48 | # set up the geometry of the grid 49 | if len(allkeys) > maxcol: 50 | rows = int(np.ceil(len(allkeys) / maxcol)) 51 | cols = maxcol 52 | else: 53 | rows = 1 54 | cols = np.min([len(allkeys), maxcol]) 55 | 56 | # create the figure and grid of axes 57 | fig, axes = plt.subplots( 58 | rows, 59 | cols, 60 | figsize=(cols * panel_size[0], rows * panel_size[1]), 61 | sharex=True, 62 | constrained_layout=True, 63 | ) 64 | # make the axes easier to index 65 | if len(allkeys) > 1: 66 | ax = axes.flatten() 67 | else: 68 | ax = [axes] 69 | 70 | # set what_is_x variable 71 | if what_is_x.lower() == "index": 72 | x_values = np.arange(0, len(like_dict[list(like_dict.keys())[0]])) 73 | xlab = f"{xaxis} Index" 74 | else: 75 | if what_is_x in like_dict.keys(): 76 | x_values = like_dict[what_is_x] 77 | xlab = what_is_x 78 | else: 79 | raise Exception("Desired `what_is_x` quantity is not in given dictionary") 80 | 81 | # display each quantity 82 | for k, key in enumerate(allkeys): 83 | # make the plot (or an empty box) 84 | if key in like_dict.keys(): 85 | ax[k].plot( 86 | x_values, like_dict[key], color=plt.cm.viridis(k / len(allkeys)), **kw 87 | ) 88 | ax[k].set_xlabel(xlab) 89 | ax[k].set_ylabel( 90 | f"{key} ({u.Quantity(like_dict[key]).unit.to_string('latex_inline')})" 91 | ) 92 | else: 93 | ax[k].text( 94 | 0.5, 95 | 0.5, 96 | f"No {key}", 97 | transform=ax[k].transAxes, 98 | ha="center", 99 | va="center", 100 | ) 101 | 102 | # add a title for each box 103 | ax[k].set_title(key) 104 | 105 | # hide xlabel except on the bottom row 106 | if k < (len(allkeys) - cols): 107 | ax[k].set_xlabel("") 108 | else: 109 | ax[k].tick_params(labelbottom=True) 110 | 111 | # hide any additional axes 112 | if k + 1 <= len(ax): 113 | for axi in ax[k + 1 :]: 114 | axi.axis("Off") 115 | if filename is not None: 116 | self.savefig(filename) 117 | -------------------------------------------------------------------------------- /chromatic/rainbows/actions/fold.py: -------------------------------------------------------------------------------- 1 | from ...imports import * 2 | 3 | __all__ = ["fold", "mask_transit"] 4 | 5 | 6 | def fold(self, period=None, t0=None, event="Mid-Transit"): 7 | """ 8 | Fold this `Rainbow` to a period and reference epoch. 9 | 10 | This changes the times from some original time into 11 | a phased time, for example the time within an orbital 12 | period, relative to the time of mid-transit. This 13 | is mostly a convenience function for plotting data 14 | relative to mid-transit and/or trimming data based 15 | on orbital phase. 16 | 17 | Parameters 18 | ---------- 19 | period : Quantity 20 | The orbital period of the planet (with astropy units of time). 21 | t0 : Quantity 22 | Any mid-transit epoch (with astropy units of time). 23 | event : str 24 | A description of the event that happens periodically. 25 | For example, you might want to switch this to 26 | 'Mid-Eclipse' (as well as offsetting the `t0` by the 27 | appropriate amount relative to transit). This description 28 | may be used in plot labels. 29 | 30 | Returns 31 | ------- 32 | folded : Rainbow 33 | The folded `Rainbow`. 34 | """ 35 | 36 | # create a history entry for this action (before other variables are defined) 37 | h = self._create_history_entry("fold", locals()) 38 | 39 | # warn 40 | if (period is None) or (t0 is None): 41 | message = """ 42 | Folding to a transit period requires both 43 | `period` and `t0` be specified. Please try again. 44 | """ 45 | cheerfully_suggest(message) 46 | return self 47 | 48 | # create a copy of the existing rainbow 49 | new = self._create_copy() 50 | 51 | # calculate predicted time from transit 52 | new.time = (((self.time - t0) + 0.5 * period) % period) - 0.5 * period 53 | # (the nudge by 0.5 period is to center on -period/2 to period/2) 54 | 55 | # change the default time label 56 | new.metadata["time_label"] = f"Time from {event}" 57 | 58 | # append the history entry to the new Rainbow 59 | new._record_history_entry(h) 60 | 61 | return new 62 | 63 | 64 | def mask_transit(self, period, t0, duration, event="Mid-Transit"): 65 | """ 66 | Mask a transit in this `Rainbow`, to focus on out-of-transit. 67 | 68 | Parameters 69 | ---------- 70 | period : Quantity 71 | The orbital period of the planet (with astropy units of time). 72 | t0 : Quantity 73 | Any mid-transit epoch (with astropy units of time). 74 | duration : Quantity 75 | The total duration of the transit to remove. 76 | event : str 77 | A description of the event that happens periodically. 78 | For example, you might want to switch this to 79 | 'Mid-Eclipse' (as well as offsetting the `t0` by the 80 | appropriate amount relative to transit). This description 81 | may be used in plot labels. 82 | 83 | Returns 84 | ------- 85 | masked : Rainbow 86 | The `Rainbow` with the transit masked as not `ok`. 87 | """ 88 | # create a history entry for this action (before other variables are defined) 89 | h = self._create_history_entry("mask_transit", locals()) 90 | 91 | if "fold" in self.history(): 92 | raise ValueError( 93 | f""" 94 | It looks like this Rainbow has already been folded. 95 | Please run `.mask_transit` only on unfolded data. 96 | (It will refold while masking.) 97 | """ 98 | ) 99 | 100 | # fold and mask transit 101 | folded = self.fold(period=period, t0=t0, event=event) 102 | folded.timelike["ok"] = np.abs(folded.time) >= duration / 2 103 | 104 | # append the history entry to the new Rainbow 105 | folded._record_history_entry(h) 106 | 107 | return folded 108 | -------------------------------------------------------------------------------- /chromatic/rainbows/readers/xarray_fitted_light_curves.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 | 3. Edit the `from_xarray_fitted_light_curves` function so that it will 4 | load a chromatic light curve file in your format and, 5 | for some Rainbow object `rainbow`, populate at least: 6 | 7 | + rainbow.timelike['time'] 8 | + rainbow.wavelike['wavelength'] 9 | + rainbow.fluxlike['flux'] 10 | 11 | You'll need to replace the cartoon functions on each 12 | line with the actual code needed to load your file. 13 | 14 | (This template assumes that only one file needs to be 15 | loaded. If you need to load multiple segments, or each 16 | time point is stored in its own file or something, then 17 | check out `stsci.py` for an example of loading and 18 | stitching together multiple input files. You'll probably 19 | want to change `filepath` to accept a glob-friendly string 20 | like `my-neato-formatted-files-*.npy` or some such.) 21 | 22 | 4. Edit the `readers/__init__.py` file to import your 23 | `from_xarray_fitted_light_curves` function to be accessible when people 24 | are trying to create new Rainbows. Add an `elif` statement 25 | to the `guess_reader` function that will help guess which 26 | reader to use from some aspect(s) of the filename. 27 | 28 | (This `guess_reader` function also accepts a `format=` 29 | keyword that allows the user to explicitly specify that 30 | the xarray_fitted_light_curves reader should be used.) 31 | 32 | 5. Submit a pull request to the github repository for 33 | this package, so that other folks can use your handy 34 | new reader too! 35 | """ 36 | 37 | # import the general list of packages 38 | from ...imports import * 39 | from ..writers.xarray_fitted_light_curves import xr, json, chromatic_to_ers 40 | 41 | ers_to_chromatic = {v: k for k, v in chromatic_to_ers.items()} 42 | 43 | # define list of the only things that will show up in imports 44 | __all__ = ["from_xarray_fitted_light_curves"] 45 | 46 | 47 | def from_xarray_fitted_light_curves(self, filepath): 48 | """ 49 | Populate a Rainbow from a file in the xarray_fitted_light_curves format. 50 | 51 | Parameters 52 | ---------- 53 | 54 | self : self 55 | The object to be populated. 56 | 57 | filepath : str 58 | The path to the file to load. 59 | """ 60 | 61 | import xarray as xr 62 | 63 | ds = xr.open_dataset(filepath) 64 | 65 | def make_Quantity(da): 66 | """ 67 | Convert a data array into a chromatic quantity (with astropy units). 68 | 69 | """ 70 | unit_string = da.attrs.get("units", "") 71 | if unit_string != "": 72 | unit = u.Unit(unit_string) 73 | else: 74 | unit = 1 75 | 76 | return da.data * unit 77 | 78 | self.wavelike["wavelength"] = make_Quantity( 79 | ds[ers_to_chromatic.get("wavelength", "wavelength")] 80 | ) 81 | self.timelike["time"] = make_Quantity(ds[ers_to_chromatic.get("time", "time")]) 82 | 83 | for key, da in ds.items(): 84 | chromatic_key = ers_to_chromatic.get(key, key) 85 | self._put_array_in_right_dictionary(chromatic_key, make_Quantity(da)) 86 | for k, v in da.attrs.items(): 87 | if k != "units": 88 | metadata_key = f"metadata-for-{chromatic_key}" 89 | try: 90 | self.metadata[metadata_key] 91 | except KeyError: 92 | self.metadata[metadata_key] = dict() 93 | self.metadata[metadata_key][k] = v 94 | 95 | # kludge to convert to corrected flux 96 | k = "systematics_model" 97 | if k not in self.fluxlike: 98 | self.fluxlike[k] = self.flux / self.get("corrected_flux", 1) 99 | 100 | # kludge to convert to corrected flux 101 | # k = "corrected_flux_error" 102 | # self.fluxlike[k] = self.get(k, self.uncertainty / self.get("systematics_model", 1)) 103 | 104 | for k, v in ds.attrs.items(): 105 | 106 | self.metadata[k] = json.loads(v) 107 | -------------------------------------------------------------------------------- /chromatic/tests/test_withmodel.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_attach_model(): 6 | plt.close("all") 7 | 8 | simulated = SimulatedRainbow().inject_noise().inject_transit().inject_systematics() 9 | original = simulated._create_copy() 10 | 11 | # pull out the model 12 | inputs = {} 13 | for k in ["model", "planet_model", "systematics_model"]: 14 | inputs[k] = simulated.fluxlike[k] 15 | del simulated.fluxlike[k] 16 | 17 | # reattach the model 18 | withmodel = simulated.attach_model(**inputs) 19 | assert simulated != withmodel 20 | assert withmodel == original 21 | 22 | 23 | def test_paint_with_models(): 24 | plt.close("all") 25 | 26 | s = ( 27 | SimulatedRainbow() 28 | .inject_transit() 29 | .inject_systematics() 30 | .inject_noise() 31 | .bin(R=50, dt=5 * u.minute) 32 | ) 33 | s.paint_with_models(cmap="gray") 34 | plt.savefig( 35 | os.path.join(test_directory, "demonstration-of-imshow-data-with-model.pdf") 36 | ) 37 | s.paint_with_models(models=["systematics_model", "planet_model"], cmap="gray") 38 | plt.savefig( 39 | os.path.join( 40 | test_directory, "demonstration-of-imshow-data-with-model-components.pdf" 41 | ) 42 | ) 43 | 44 | s.paint_with_models(models=["systematics_model", "planet_model"], cmap="gray") 45 | s.paint_with_models( 46 | models=["systematics_model", "planet_model"], cmap="gray", label=False 47 | ) 48 | s.paint_with_models( 49 | models=["systematics_model", "planet_model"], 50 | cmap="gray", 51 | label_textkw=dict(color="red"), 52 | ) 53 | s.paint_with_models( 54 | models=["systematics_model", "planet_model"], cmap="gray", label="outside" 55 | ) 56 | 57 | 58 | def test_plot_with_model_and_residuals(): 59 | plt.close("all") 60 | 61 | s = SimulatedRainbow(R=3, dt=3 * u.minute) 62 | r = s.inject_transit( 63 | limb_dark="quadratic", 64 | u=np.transpose( 65 | [np.linspace(1.0, 0.0, s.nwave), np.linspace(0.5, 0.0, s.nwave)] 66 | ), 67 | method="exoplanet", 68 | ).inject_noise(signal_to_noise=1000) 69 | for i, options in enumerate( 70 | [ 71 | dict(), 72 | dict(errorbar=True), 73 | dict(cmap=one2another("skyblue", "sienna")), 74 | dict(data_plotkw=dict(alpha=0.5), cmap="magma_r"), 75 | dict(figsize=(8, 4), cmap=one2another("orchid", "indigo")), 76 | ] 77 | ): 78 | r.plot_with_model_and_residuals(**options) 79 | plt.savefig( 80 | os.path.join( 81 | test_directory, 82 | f"demonstration-of-rainbow-of-lightcurves-and-residuals-example{i}.png", 83 | ) 84 | ) 85 | 86 | 87 | def test_plot_and_animate_with_models(output="gif"): 88 | plt.close("all") 89 | s = ( 90 | SimulatedRainbow(R=5) 91 | .inject_transit() 92 | .inject_systematics() 93 | .inject_noise(signal_to_noise=500) 94 | ) 95 | s.setup_wavelength_colors( 96 | cmap="magma", vmin=0.1 * u.micron, vmax=10 * u.micron, log=True 97 | ) 98 | for o in ["vertical", "horizontal"]: 99 | for e in [True, False]: 100 | s.plot_one_wavelength_with_models(0, errorbar=e, orientation=o) 101 | error_string = {True: "with", False: "without"}[e] + "-errorbars" 102 | filename = f"data-with-models-{o}-{error_string}" 103 | plt.savefig( 104 | os.path.join(test_directory, f"demonstration-of-plot-{filename}.png") 105 | ) 106 | s.animate_with_models( 107 | os.path.join( 108 | test_directory, f"demonstration-of-animate-{filename}.{output}" 109 | ), 110 | errorbar=e, 111 | orientation=o, 112 | ) 113 | -------------------------------------------------------------------------------- /docs/api.md: -------------------------------------------------------------------------------- 1 | # Reference 2 | 3 | ::: chromatic.rainbows.read_rainbow 4 | 5 | ::: chromatic.rainbows.rainbow.Rainbow 6 | selection: 7 | members: 8 | - Rainbow 9 | - __init__ 10 | - __getattr__ 11 | - __setattr__ 12 | - __getitem__ 13 | ::: chromatic.rainbows.withmodel.RainbowWithModel 14 | ::: chromatic.rainbows.simulated.SimulatedRainbow 15 | 16 | ## 🌈 Helpers 17 | ::: chromatic.rainbows.helpers.get 18 | ::: chromatic.rainbows.helpers.help 19 | ::: chromatic.rainbows.helpers.history 20 | ::: chromatic.rainbows.helpers.save 21 | 22 | ## 🌈 Actions 23 | ::: chromatic.rainbows.actions.align_wavelengths 24 | ::: chromatic.rainbows.actions.attach_model 25 | ::: chromatic.rainbows.actions.binning 26 | options: 27 | show_root_heading: False 28 | ::: chromatic.rainbows.actions.compare 29 | ::: chromatic.rainbows.actions.flag_outliers 30 | ::: chromatic.rainbows.actions.fold 31 | ::: chromatic.rainbows.actions.inflate_uncertainty 32 | ::: chromatic.rainbows.actions.inject_noise 33 | ::: chromatic.rainbows.actions.inject_outliers 34 | ::: chromatic.rainbows.actions.inject_spectrum 35 | ::: chromatic.rainbows.actions.inject_systematics 36 | ::: chromatic.rainbows.actions.inject_transit 37 | ::: chromatic.rainbows.actions.normalization 38 | options: 39 | show_root_heading: False 40 | ::: chromatic.rainbows.actions.operations 41 | options: 42 | show_root_heading: False 43 | selection: 44 | members: 45 | - __add__ 46 | - __sub__ 47 | - __mul__ 48 | - __truediv__ 49 | - __eq__ 50 | ::: chromatic.rainbows.actions.remove_trends 51 | ::: chromatic.rainbows.actions.shift 52 | ::: chromatic.rainbows.actions.trim 53 | 54 | 55 | ## 🌈 Get/Timelike 56 | ::: chromatic.rainbows.get.timelike.average_lightcurve 57 | options: 58 | show_root_heading: False 59 | ::: chromatic.rainbows.get.timelike.median_lightcurve 60 | options: 61 | show_root_heading: False 62 | ::: chromatic.rainbows.get.timelike.subset 63 | options: 64 | show_root_heading: False 65 | ::: chromatic.rainbows.get.timelike.time 66 | options: 67 | show_root_heading: False 68 | 69 | ## 🌈 Get/Wavelike 70 | ::: chromatic.rainbows.get.wavelike.average_spectrum 71 | options: 72 | show_root_heading: False 73 | ::: chromatic.rainbows.get.wavelike.expected_uncertainty 74 | options: 75 | show_root_heading: False 76 | ::: chromatic.rainbows.get.wavelike.measured_scatter_in_bins 77 | options: 78 | show_root_heading: False 79 | ::: chromatic.rainbows.get.wavelike.measured_scatter 80 | options: 81 | show_root_heading: False 82 | ::: chromatic.rainbows.get.wavelike.median_spectrum 83 | options: 84 | show_root_heading: False 85 | ::: chromatic.rainbows.get.wavelike.spectral_resolution 86 | options: 87 | show_root_heading: False 88 | ::: chromatic.rainbows.get.wavelike.subset 89 | options: 90 | show_root_heading: False 91 | 92 | ## 🌈 Visualizations 93 | ::: chromatic.rainbows.visualizations.animate 94 | options: 95 | show_root_heading: False 96 | ::: chromatic.rainbows.visualizations.colors 97 | options: 98 | show_root_heading: False 99 | ::: chromatic.rainbows.visualizations.imshow 100 | ::: chromatic.rainbows.visualizations.interactive 101 | options: 102 | show_root_heading: False 103 | ::: chromatic.rainbows.visualizations.pcolormesh 104 | ::: chromatic.rainbows.visualizations.plot_lightcurves 105 | ::: chromatic.rainbows.visualizations.plot_spectra 106 | ::: chromatic.rainbows.visualizations.plot 107 | 108 | ## 🔨 Tools 109 | ::: chromatic.spectra.planck.get_planck_photons 110 | ::: chromatic.spectra.phoenix.get_phoenix_photons 111 | ::: chromatic.tools.resampling 112 | options: 113 | show_root_heading: False 114 | selection: 115 | members: 116 | - bintoR 117 | - bintogrid 118 | - resample_while_conserving_flux 119 | ::: chromatic.imports 120 | options: 121 | show_root_heading: False 122 | -------------------------------------------------------------------------------- /chromatic/tests/test_simulations.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from .setup_tests import * 3 | 4 | 5 | def test_simulated_basics(): 6 | 7 | a = SimulatedRainbow().inject_noise() 8 | b = SimulatedRainbow(tlim=[-1, 1] * u.hour, dt=1 * u.minute).inject_noise() 9 | c = SimulatedRainbow(time=np.linspace(-1, 1) * u.hour).inject_noise() 10 | d = SimulatedRainbow(wlim=[0.1, 1] * u.micron, R=100).inject_noise() 11 | assert d.wscale == "log" 12 | e = SimulatedRainbow(wlim=[0.1, 1] * u.micron, dw=1 * u.nm).inject_noise() 13 | assert e.wscale == "linear" 14 | f = SimulatedRainbow(wavelength=np.logspace(0, 1) * u.micron).inject_noise() 15 | 16 | 17 | def test_photon_noise(N=10000): 18 | 19 | # inject noise in two different ways 20 | noiseless = SimulatedRainbow().inject_transit() 21 | gaussian = noiseless.inject_noise(signal_to_noise=np.sqrt(N)).normalize() 22 | poisson = noiseless.inject_noise(number_of_photons=N).normalize() 23 | 24 | # imshow the two ways 25 | fi, ax = plt.subplots(1, 2, figsize=(8, 3), dpi=300, constrained_layout=True) 26 | gaussian.paint(ax=ax[0]) 27 | ax[0].set_title(r"Gaussian ($\sigma=$" + f"{1/np.sqrt(N):.3f}=" + r"$1/\sqrt{N}$)") 28 | poisson.paint(ax=ax[1]) 29 | ax[1].set_title(f"Poisson (N={N})") 30 | 31 | # make sure the standard deviation is about right 32 | # sigma = np.std(poisson.residuals / poisson.model) 33 | # assert np.isclose(sigma, 1 / np.sqrt(N), rtol=0.1) 34 | plt.savefig( 35 | os.path.join( 36 | test_directory, 37 | "demonstration-of-injecting-noise-as-poisson-or-gaussian.pdf", 38 | ) 39 | ) 40 | 41 | 42 | def test_star_flux(): 43 | g = SimulatedRainbow( 44 | wavelength=np.logspace(0, 1) * u.micron, star_flux=np.logspace(0, 1) 45 | ).inject_noise() 46 | 47 | 48 | def test_inject_transit(): 49 | s = SimulatedRainbow() 50 | x = (s.wavelength / max(s.wavelength)).value 51 | 52 | # test that injected parameters are stored appropriately 53 | delta = 0.5 * np.sin(x * 10) 54 | t0 = 0.01 55 | t = s.inject_transit(delta=delta, t0=t0, method="trapezoid") 56 | assert t.metadata["injected_transit_parameters"]["t0"] == t0 57 | assert np.all(t.wavelike["injected_transit_delta"] == delta) 58 | 59 | # test that all methods work 60 | fi, ax = plt.subplots(3, 2, figsize=(8, 6), constrained_layout=True, dpi=300) 61 | 62 | method = "trapezoid" 63 | s.inject_transit(method=method).paint(ax=ax[0, 0]) 64 | plt.title(f"{method} | default") 65 | s.inject_transit( 66 | planet_radius=np.sqrt(x) / 10, 67 | tau=0.02, 68 | t0=x * 0.03, 69 | T=x * 0.05, 70 | method=method, 71 | ).paint(ax=ax[0, 1]) 72 | plt.title(f"{method} | wacky") 73 | 74 | method = "exoplanet" 75 | s.inject_transit(method=method).paint(ax=ax[1, 0]) 76 | plt.title(f"{method} | default") 77 | s.inject_transit( 78 | planet_radius=np.sqrt(x) / 10, 79 | t0=x * 0.03, 80 | inc=89, 81 | a=1 / x * 10, 82 | u=np.random.uniform(0, 0.5, [s.nwave, 2]), 83 | method=method, 84 | ).paint(ax=ax[1, 1]) 85 | plt.title(f"{method} | wacky") 86 | 87 | method = "exoplanet" 88 | s.inject_transit(method=method).imshow(ax=ax[2, 0]) 89 | plt.title(f"{method} | default") 90 | s.inject_transit( 91 | planet_radius=np.sqrt(x) / 10, 92 | t0=x * 0.03, 93 | inc=89, 94 | a=1 / x * 10, 95 | limb_dark="nonlinear", 96 | u=np.random.uniform(0, 0.5, [s.nwave, 4]), 97 | method=method, 98 | ).paint(ax=ax[2, 1]) 99 | plt.title(f"{method} | wacky") 100 | 101 | plt.savefig(os.path.join(test_directory, "demonstration-of-injecting-transit.pdf")) 102 | plt.close("all") 103 | 104 | 105 | def test_inject_systematics(): 106 | SimulatedRainbow().inject_transit().inject_noise().inject_systematics().bin( 107 | R=10, dt=10 * u.minute 108 | ).paint_with_models() 109 | plt.savefig( 110 | os.path.join(test_directory, "demonstration-of-injecting-fake-systematics.pdf") 111 | ) 112 | -------------------------------------------------------------------------------- /chromatic/tests/test_align_wavelengths.py: -------------------------------------------------------------------------------- 1 | from ..rainbows import * 2 | from ..rainbows.actions.align_wavelengths import _create_shared_wavelength_axis 3 | from .setup_tests import * 4 | 5 | 6 | def create_simulation_with_wobbly_wavelengths( 7 | fractional_shift=0.002, 8 | signal_to_noise=100, 9 | dw=0.001 * u.micron, 10 | wlim=[0.95, 1.05] * u.micron, 11 | dt=30 * u.minute, 12 | **kw 13 | ): 14 | # set up a function to create a fake absorption line spectrum 15 | N = 10 16 | centers = ( 17 | np.random.uniform(wlim[0].to_value("micron"), wlim[1].to_value("micron"), N) 18 | * u.micron 19 | ) 20 | sigmas = np.random.uniform(0.001, 0.003, N) * u.micron 21 | amplitudes = np.random.uniform(0, 0.5, N) 22 | 23 | def fake_spectrum(w): 24 | f = u.Quantity(np.ones(np.shape(w))) 25 | for c, s, a in zip(centers, sigmas, amplitudes): 26 | f *= 1 - a * np.exp(-0.5 * (w - c) ** 2 / s**2) 27 | return f 28 | 29 | # create 30 | r = SimulatedRainbow(dw=dw, wlim=wlim, dt=dt, **kw).inject_noise( 31 | signal_to_noise=signal_to_noise 32 | ) 33 | wobbly_wavelengths = r.wavelength[:, np.newaxis] * np.random.normal( 34 | 1, fractional_shift, r.ntime 35 | ) 36 | r.fluxlike["wavelength_2d"] = wobbly_wavelengths 37 | r.fluxlike["model"] = fake_spectrum(r.fluxlike["wavelength_2d"]) 38 | r.fluxlike["flux"] = r.fluxlike["flux"] * r.fluxlike["model"] 39 | 40 | return r 41 | 42 | 43 | def test_create_shared_wavelength_axis(fractional_shift=0.002, dw=0.01 * u.micron): 44 | r = create_simulation_with_wobbly_wavelengths( 45 | fractional_shift=fractional_shift, dw=dw 46 | ) 47 | _create_shared_wavelength_axis(r, wscale="linear", visualize=True) 48 | plt.savefig( 49 | os.path.join( 50 | test_directory, 51 | "demonstration-of-creating-shared-wavelength-axis.pdf", 52 | ) 53 | ) 54 | 55 | 56 | def test_align_wavelengths(fractional_shift=0.002, dw=0.001 * u.micron, **kw): 57 | r = create_simulation_with_wobbly_wavelengths( 58 | fractional_shift=fractional_shift, dw=dw, **kw 59 | ) 60 | a = r.align_wavelengths() 61 | plt.close("all") 62 | fi, ax = plt.subplots(2, 2, dpi=300, figsize=(8, 6), constrained_layout=True) 63 | for i, x in enumerate([r, a]): 64 | plt.sca(ax[0, i]) 65 | plt.imshow(remove_unit(x.flux), aspect="auto", cmap="gray") 66 | plt.title(["original", "aligned"][i] + " flux") 67 | plt.xlabel("time index") 68 | plt.ylabel("wavelength index") 69 | plt.colorbar() 70 | 71 | plt.sca(ax[1, i]) 72 | plt.imshow(remove_unit(x.fluxlike["wavelength_2d"]), aspect="auto") 73 | plt.title(["original", "aligned"][i] + " wavelength") 74 | plt.xlabel("time index") 75 | plt.ylabel("wavelength index") 76 | plt.colorbar() 77 | 78 | plt.savefig( 79 | os.path.join( 80 | test_directory, 81 | "demonstration-of-aligning-2D-wavelengths-to-shared-axis.pdf", 82 | ) 83 | ) 84 | 85 | 86 | def test_align_wavelengths_with_not_ok_data(visualize=False): 87 | with warnings.catch_warnings(): 88 | warnings.simplefilter("ignore") 89 | for ok_fraction in [0.25, 0.5, 0.75, 1.0]: 90 | r = create_simulation_with_wobbly_wavelengths( 91 | fractional_shift=0.0005, 92 | wlim=[0.99, 1.01] * u.micron, 93 | dw=0.001 * u.micron, 94 | dt=20 * u.minute, 95 | ) 96 | r.fluxlike["ok"] = np.random.uniform(size=r.shape) < ok_fraction 97 | cautious = r.align_wavelengths(minimum_acceptable_ok=1) 98 | carefree = r.align_wavelengths(minimum_acceptable_ok=0) 99 | if visualize: 100 | cautious.paint_quantities() 101 | plt.suptitle(ok_fraction) 102 | carefree.paint_quantities() 103 | plt.suptitle(ok_fraction) 104 | 105 | # if np.any(r.ok == 0): 106 | # assert np.any((carefree.ok != 1) & (carefree.ok != 0)) 107 | --------------------------------------------------------------------------------