Designing New 🌈 Features¶

The Rainbow class is the heart of chromatic. We aim for it to be as intuitive and easy-to-use as possible, to enable transparent and repeatable analysis of spectrosopic time-series datasets. This page collects a few explanations that might be useful if you're trying to develop new Rainbow features.

In [1]:

from chromatic import SimulatedRainbow

from chromatic import SimulatedRainbow

How are Rainbow objects organized?¶

The user mainly interacts with quantities like time, wavelength, and flux via the .time, .wavelength, .flux attributes. However, if you more closely at what's happening inside a Rainbow object, you'll see that these are properties that pull data from some core dictionaries. The general user probably doesn't need to interact with the core dictionaries, but if you're developing new features for Rainbow objects you will want to understand what's happening a little more clearly.

The core dictionaries are designed to store quantities that have dimensions like either wavelength, time, or flux:

• .wavelike[...] contains everything that has the same dimensions as wavelength. This dictionary will contain at least a 'wavelength' key, with the actual wavelengths themselves. It might also contain information like the average spectrum of the star, the average S/N at each wavelength, the number of original detector pixels that wound up in this particular wavelength bin, and/or a mask of what wavelengths should be considered good or bad (no matter the time point).
• .timelike[...] contains everything that has the same dimensons as time. This dictionary will contain at least a 'time' key, with the actual times themselves. It might also contain information like a broadband light curve of the transit, the x or y position of the spectrum on the detector, the temperature of the detector, and/or a mask of what times should be considered good or bad (no matter the wavelength).
• .fluxlike[...] contains everything that has the same dimensons as flux. This dictionary will contain at least a 'flux' key, with the actual fluxes themselves. It should also contain an 'uncertainty' keyword with the uncertainties associated with those fluxes (or maybe None). It might also contain information about other quantities that depend on both time and wavelength, such as the centroid of the spectral trace, the maximum fraction of saturation, and/or a mask of what individual points should be considered good or bad.

There is one more core dictionary:

• .metadata contains general information that might be useful to hang onto, to pass along to another derived object, or to save out to a file.

In [2]:

r = SimulatedRainbow().inject_noise()

r = SimulatedRainbow().inject_noise()

In [3]:

r.wavelike.keys()

r.wavelike.keys()

Out[3]:

dict_keys(['wavelength', 'original_wave_index'])

In [4]:

r.timelike.keys()

r.timelike.keys()

Out[4]:

dict_keys(['time', 'original_time_index'])

In [5]:

r.fluxlike.keys()

r.fluxlike.keys()

Out[5]:

dict_keys(['flux', 'model', 'uncertainty'])

In [6]:

r.metadata.keys()

r.metadata.keys()

Out[6]:

dict_keys(['name', 'history', 'R', 'wscale', 'tscale', 'signal_to_noise'])

When you retrieve variables with something like .wavelength, .time, .flux, data is being pulled directly from these dictionaries.

What can Rainbow objects do?¶

We defined a small lexicon of things that Rainbow objects can do. If you want to write a new feature, hopefully it fits into one of these categories. If not, we can certainly discuss adding new ones!

• actions return a new Rainbow object as the result. As such, they can be chained together into commands like r.bin(R=5).normalize().plot().
• wavelike return a wavelike-shaped array from a Rainbow, with one quantity for each wavelength.
• timelike return a timelike-shaped array from a Rainbow, with one quantity for each time.
• visualizations create some graphical representation of the data in a Rainbow.

Each of these categories has its own directory inside chromatic/rainbows/ where the corresponding code should be stored. The .help() method attached to any Rainbow will list everything you can do with it.

In [7]:

r.help()

r.help()

Hooray for you! You asked for help on what you can do
with this 🌈 object. Here's a quick reference of a few
available options for things to try.

-----------
| actions |
-----------

🌈🧮📝 | +-*/                         
   Do basic math operations with two Rainbows.
🌈🗂🔪 | .[:,:]()                     
   Index, slice, or mask a Rainbow to get a subset.
🌈🚧🌊 | .align_wavelengths()         
   Align spectra with different wavelength grids onto one shared axis.
🌈🧺🧱 | .bin()                       
   Bin to a new wavelength or time grid.
🌈🧑‍🤝‍🧑🌈 | .compare()                   
   Connect to other 🌈 objects for easy comparison.
🌈🐈⏰ | .concatenate_in_time()       
   Stitch together two Rainbows with identical wavelengths.
🌈🐈🌊 | .concatenate_in_wavelength() 
   Stitch together two Rainbows with identical times.
🌈🚩👀 | .flag_outliers()             
   Flag outlier data points.
🌈⏲🎞 | .fold()                      
   Fold times relative to a particular period and epoch.
🌈🧺⏰ | .get_average_lightcurve_as_rainbow() 
   Bin down to a single integrated light curve.
🌈🧺🌊 | .get_average_spectrum_as_rainbow() 
   Bin down to a single integrated spectrum.
🌈🎧🎲 | .inject_noise()              
   Inject (uncorrelated, simple) random noise.
🌈🎧🎹 | .inject_systematics()        
   Inject (correlated, wobbly) systematic noise.
🌈⭐️👻 | .inject_spectrum()           
   Inject a static stellar spectrum.
🌈🪐🚞 | .inject_transit()            
   Inject a transit signal.
🌈🫓😑 | .normalize()                 
   Normalize by dividing through by a typical spectrum (and/or light curve).
🌈🏴‍☠️🛁 | .remove_trends()             
   Remove smooth trends in time and/or wavelength.
🌈🚇🌊 | .shift()                     
   Doppler shift wavelengths.
🌈🍱💇 | .trim()                      
   Trim away wavelengths or times.

----------------
| get/timelike |
----------------

⏰🐋🌈 | .get_average_lightcurve()    
   Get the weighted average light curve.
⏰🔎🌈 | .get_for_time()              
   Get a quantity associated with a time index.
⏰🖖🌈 | .get_median_lightcurve()     
   Get the median light curve.
⏰👌🌈 | .get_ok_data_for_time()      
   Get a quantity associated with a time index.
⏰🛰🌈 | .get_times_as_astropy()      
   Get the times as an astropy Time object.
⏰🚀🌈 | .set_times_from_astropy()    
   Set the times from an astropy Time object (modifies in-place).

----------------
| get/wavelike |
----------------

🌊🐋🌈 | .get_average_spectrum()      
   Get the weighted average spectrum.
🌊🔎🌈 | .get_for_wavelength()        
   Get a quantity associated with a wavelength index.
🌊🎯🌈 | .get_measured_scatter()      
   Get the measured scatter on the time series for each wavelength.
🌊🖖🌈 | .get_median_spectrum()       
   Get the median spectrum.
🌊👌🌈 | .get_ok_data_for_wavelength() 
   Get a quantity associated with a wavelength index.
🌊💎🌈 | .get_spectral_resolution()   
   Get the spectral resolution (R=w/dw).

-----------
| helpers |
-----------

🙋🌈📄 | .help()                      
   Get one-line help summaries of available methods.
📓🌈🪵 | .history()                   
   Get the history that went into this Rainbow.
💾🌈📼 | .save()                      
   Save this Rainbow out to a permanent file.

------------------
| visualizations |
------------------

🎨📽⏰ | .animate_lightcurves()       
   Animate a sequence of light curves across different wavelengths.
🎨📽🌊 | .animate_spectra()           
   Animate a sequence of spectra across different times.
🎨🖼📺 | .imshow()                    
   Paint a map of flux across wavelength and time.
🎨🕹📺 | .imshow_interact()           
   Show flux map and lightcurves with interactive wavelength selection.
🎨🖌📺 | .pcolormesh()                
   Paint a map of flux across wavelength and time (with non-uniform grids).
🎨🖌🧶 | .plot()                      
   Plot a sequence of light curves with vertical offsets.

------------------------------
| visualizations/diagnostics |
------------------------------

🎨🗂📺 | .imshow_quantities()         
   Show multiple 2D (wavelength and time) quantities as imshow maps.
🎨🗂🧶 | .plot_quantities()           
   Show multiple 1D (wavelength or time) quantities as scatter plots.

-------------------------
| visualizations/models |
-------------------------

🎨📺🎢 | .imshow_with_models()        
   Paint a flux map with model components.
🎨🖌🎢 | .plot_one_wavelength_with_models() 
   Plot one wavelength's light curve with model components.
🎨📽🎢 | .animate_with_models()       
   Animate all wavelengths' light curves with model components.
🎨🧶🎢 | .plot_with_model()           
   Plot a sequence of light curves with their models.

---------------------------
| visualizations/timelike |
---------------------------

🎨⏰🐋 | .plot_average_lightcurve()   
   Plot the weighted average flux per time.
🎨⏰🖖 | .plot_median_lightcurve()    
   Plot the median flux per time.

---------------------------
| visualizations/wavelike |
---------------------------

🎨🌊🔭 | .plot_average_spectrum()     
   Plot the weighted average flux per wavelength.
🎨🌊💎 | .plot_spectral_resolution()  
   Plot the spectral resolution per wavelength.
🎨🌊🎧 | .plot_noise_comparison()     
   Plot the measured and expected scatter per wavelength.
🎨🌊🥦 | .plot_noise_comparison_in_bins() 
   Plot measured and expected scatter in different size bins.

How do we add new abilities to Rainbow objects?¶

If you want to add a new method to a Rainbow object, there are a few general steps you'll probably want to follow.

1. Install in development mode (see Installation), so you can modify the code package directly and test the code in place.
2. Decide a category of "things a Rainbow can do" in which it belongs.
3. Find the directory for that category in the chromatic code package. For example, "actions" that return new Rainbow objects will be in chromatic/rainbows/actions/.
4. Look at another example in that directory to get a sense for the general layout. For example, you'll notice that "actions" generally create a copy of self, make some changes to that copy, and then return it as the new object.
5. Add your new function, either to an existing .py file where it would make sense or to its own new .py file. The first argument to your function should be self, which is the Rainbow object itself, and then any additional arguments should follow afterward. A good way to test out your new function in a notebook or a small isolated script might look something like the following:

from chromatic import *
r = SimulatedRainbow().inject_noise()
def snazzy_new_action(self, x=2):
    
    # create a copy of the original object
    new = self._create_copy()
    
    # do something
    new.fluxlike['flux'] = new.flux**x
    
    # return the modified copy
    return new
output = snazzy_new_action(r)

This allows you to develop and test your new function without having to worry about it being imported properly into the core Rainbow definition.

1. Connect your new function into the Rainbow class definition. Normally, we might define a new method directly in the same file as the class definition itself, but we wanted to split the method definitions into multiple files and directories, to make things easier to find. Importing your new function to become a Rainbow method takes a few steps. For example, let's imagine you're making a new "action" called snazzy_new_action and it's located in chromatic/rainbows/actions/snazzy.py:
   • In rainbows/actions/snazzy.py, include line of code like __all__ = ['snazzy_new_action'] at the top. The __all__ list defines what would get imported via a line like from chromatic.rainbows.actions.snazzy import * (things not in __all__ will need to be explicitly imported).
   • In rainbows/actions/__init__.py, include a line of code like from .snazzy import *, so that imports from chromatic.rainbows.actions will know how to find things in snazzy.py.
   • In rainbows/rainbow.py, down at the very bottom of the class definition for Rainbow, add snazzy_new_action to the list of from .actions import (...). This, finally, will mean that all Rainbow objects will have access to your new method, and we can do things like r.snazzy_new_action() from any Rainbow.
   • In rainbows/actions/descriptions.txt, add a row describing your snazzy new action. This table defines what appears in the .help() method.
2. Write a test for your new feature. This is a function that somehow tests whether your new feature works; the simplest form would be "does this function run", a slightly fancier version would be "are its outputs accurate." To write a test:
   • Look in chromatic/tests/ to find a bunch of test_*.py files with examples of automated tests in them.
   • Create a function that has *test* somewhere in its name and store it somewhere in the tests directory. At a minimum, this function should test that the bit of code you wrote runs without breaking.
   • Run pytest from the command line within the main repository directory. This will run your test function (along with all the other tests) and tell you whether it passed or failed. Make sure it passes!

How do we add new readers/writers for Rainbow objects?¶

You might want to create a new reader or writer, to allow chromatic to interact with your own datasets or tools. To facilitate this, templates are available with human-friendly instructions for how to add a new reader or writer. If you want to try to incorporate a new reader or writer, please:

1. Install in development mode (see Installation), so you can modify the code package directly and test the code in place.
2. Navigate to chromatic/rainbows/[readers|writers]/template.py.
3. Follow the instructions in the comments in that file.

Good luck!