Added jetspec code to repo

specialsoss/__init__.py

@@ -5,3 +5,6 @@
 __author__ = """Joe Filippazzo"""
 __email__ = 'jfilippazzo@stsci.edu'
 __version__ = '0.1.0'
+
+from . import crossdispersion
+from . import specialsoss

specialsoss/crossdispersion.py

@@ -0,0 +1,166 @@
+# -*- coding: utf-8 -*-
+"""A module of functions that can be use to approximate the SOSS psf in the
+cross dispersion direction"""
+
+import os
+
+import numpy as np
+
+
+def batman(x, mu1, sigma1, A1, sigma2, A2, sep):
+    """
+    Generate a batman function of the given parameters
+
+    Parameters
+    ----------
+    x: array-like
+        The x-axis on which to generate the function
+    mu1: float
+        The x-position of the first peak center
+    sigma1: float
+        The stanfard deviation of the two peak distributions
+    A1: float
+        The amplitude of the two peaks
+    sigma2: float
+        The stanfard deviation of the center peak
+    A2: float
+        The amplitude of the center peak
+    sep: float
+        The separation between the two peaks
+
+    Returns
+    -------
+    np.ndarray
+        The y-axis values of the mixed Gaussian
+    """
+    peak1 = gaussian(x, mu1, sigma1, A1)
+    peak2 = gaussian(x, mu1+(sep/2.), sigma2, A2)
+    peak3 = gaussian(x, mu1+sep, sigma1, A1)
+
+    return peak1 + peak2 + peak3
+
+
+def batmen(x, mu1, sigma1, A1, sigma2, A2, sep, mu3, sigma3, A3, sigma4, A4):
+    """
+    Generate two batman functions of the given parameters
+
+    Parameters
+    ----------
+    x: array-like
+        The x-axis on which to generate the function
+    mu1: float
+        The x-position of the first peak center
+    sigma1: float
+        The stanfard deviation of the two peak distributions
+    A1: float
+        The amplitude of the two peaks
+    sigma2: float
+        The stanfard deviation of the center peak
+    A2: float
+        The amplitude of the center peak
+    sep: float
+        The separation between the two peaks
+
+    Returns
+    -------
+    np.ndarray
+        The y-axis values of the mixed Gaussian
+    """
+    batman1 = batman(x, mu1, sigma1, A1, sigma2, A2, sep)
+    batman2 = batman(x, mu3, sigma3, A3, sigma4, A4, sep)
+
+    return batman1 + batman2
+
+
+def bimodal(x, mu1, sigma1, A1, mu2, sigma2, A2):
+    """
+    Generate a bimodal function of two Gaussians of the given parameters
+
+    Parameters
+    ----------
+    x: array-like
+        The x-axis on which to generate the function
+    mu1: float
+        The x-position of the first peak center
+    sigma1: float
+        The stanfard deviation of the first distribution
+    A1: float
+        The amplitude of the first peak
+    mu2: float
+        The x-position of the second peak center
+    sigma2: float
+        The stanfard deviation of the second distribution
+    A2: float
+        The amplitude of the second peak
+
+    Returns
+    -------
+    np.ndarray
+        The y-axis values of the mixed Gaussian
+    """
+    g1 = gaussian(x, mu1, sigma1, A1)
+    g2 = gaussian(x, mu2, sigma2, A2)
+
+    return g1 + g2
+
+
+def gaussian(x, mu, sigma, A):
+    """
+    Generate a Gaussian function of the given parameters
+
+    Parameters
+    ----------
+    x: array-like
+        The x-axis on which to generate the Gaussian
+    mu: float
+        The x-position of the peak center
+    sigma: float
+        The stanfard deviation of the distribution
+    A: float
+        The amplitude of the peak
+
+    Returns
+    -------
+    np.ndarray
+        The y-axis values of the Gaussian
+    """
+    return A*np.exp(-(x-mu)**2/2/sigma**2)
+
+
+def trimodal(x, mu1, sigma1, A1, mu2, sigma2, A2, mu3, sigma3, A3):
+    """
+    Generate a trimodal function of three Gaussians of the given parameters
+
+    Parameters
+    ----------
+    x: array-like
+        The x-axis on which to generate the function
+    mu1: float
+        The x-position of the first peak center
+    sigma1: float
+        The stanfard deviation of the first distribution
+    A1: float
+        The amplitude of the first peak
+    mu2: float
+        The x-position of the second peak center
+    sigma2: float
+        The stanfard deviation of the second distribution
+    A2: float
+        The amplitude of the second peak
+    mu3: float
+        The x-position of the third peak center
+    sigma3: float
+        The stanfard deviation of the third distribution
+    A3: float
+        The amplitude of the third peak
+
+    Returns
+    -------
+    np.ndarray
+        The y-axis values of the mixed Gaussian
+    """
+    g1 = gaussian(x, mu1, sigma1, A1)
+    g2 = gaussian(x, mu2, sigma2, A2)
+    g3 = gaussian(x, mu3, sigma3, A3)
+
+    return g1 + g2 + g3

specialsoss/setup_package.py

@@ -1,4 +1,4 @@
 from distutils.extension import Extension
 
 def get_package_data():
-    return {'specialsoss': ['figures/*']}
+    return {'specialsoss': ['files/*']}

specialsoss/specialsoss.py

@@ -1,12 +1,19 @@
 # -*- coding: utf-8 -*-
 
 """A module to perform optimal spectral extraction of SOSS time series observations"""
+
+import os
+
+from astropy.io import fits
+# import jwst
 import numpy as np
 
+from . import crossdispersion as xdisp
+
 
 class SossObs:
     """A class object to extract and manipulate SOSS spectra"""
-    def __init__(self, path, name='My SOSS Observations'):
+    def __init__(self, filepath, name='My SOSS Observations', **kwargs):
         """Initialize the SOSS extraction object
 
         Parameters
@@ -17,5 +24,62 @@ def __init__(self, path, name='My SOSS Observations'):
             The name of the observation set
         """
         self.name = name
-        self.path = path
+        self.filepath = filepath
+
+        # Extract the data and header info
+        self.run_pipeline(**kwargs)
+
+    def run_pipeline(self, filepath=None, **kwargs):
+        """Run the file through the JWST pipeline, storing the data and header
+
+        Parameters
+        ----------
+        filepath: str
+            The path to the file. If None, uses the filepath given at
+            initialization
+        """
+        # Check if new filepath
+        if filepath is None:
+            filepath = self.filepath
+
+        # Make sure the file exists
+        if not os.path.exists(filepath) or not filepath.endswith('.fits'):
+            raise FileError(filepath, ": Invalid file")
+
+        # Get the data
+        self.raw_data = fits.getdata(filepath)
+
+        # Get the header and store cards as attributes
+        self.header = fits.getheader(filepath)
+        for card in self.header.cards:
+            setattr(self, card[0], card[1])
+
+        # Run the pipeline
+        # self.datacube = jwst.run(filepath, **kwargs)
+
+    def locate_trace(self, func=xdisp.bimodal, bounds=([15,0,15]*2, [110,5,np.inf]*2)):
+        """Locate the traces in the data by convolving a function in each
+        column
+
+        Parameters
+        ----------
+        func: function
+            The function to fit
+        bounds: sequence
+            The lower and upper bounds for the function input parameters
+        """
+        # Compose a median image from the stack
+        median = np.median(self.datacube, axis=0)
+
+        # Get a trace mask by fitting the given function
+        # to each column in the median image
+        self.mask = self.frame_mask(median, func, bounds)
+
+        # Apply the median trace mask to each image to mask the whole cube
+        mask_arr = np.array([self.mask]*len(self.datacube))
+        self.maskedcube = ma.array(self.datacube, mask=mask_arr)
 
+    @static_method
+    def frame_mask(frame, func, bounds):
+        """Generate a mask for the given frame"""
+        pass

tests/test_specialsoss.py

@@ -4,11 +4,13 @@
 """Tests for `specialsoss` package."""
 
 import pytest
+from pkg_resources import resource_filename
 
 from specialsoss import specialsoss
 
 
 def test_SossObs():
     """Test SossObs object"""
-    obs = specialsoss.SossObs(path='.')
+    file = resource_filename('specialsoss', 'files/soss_sample.fits')
+    obs = specialsoss.SossObs(filepath=file)
     assert obs.name == 'My SOSS Observations'