base.py rolled back to old StarlightBase behavior.

base.py

'''
Created on 30/05/2014

@author: andre
'''

import pystarlight.io  # @UnusedImport
import atpy
import numpy as np
from scipy.interpolate.interpolate import interp1d

###############################################################################

class StarlightBase(object):
    '''
    Starlight base interface and utilities.
    
    Parameters
    ----------
    base_file : string
        File describing the base. This is the one specified
        in the starlight grid file.
        
    base_dir : string
        Directory containing the base spectra.
    
    Examples
    --------
    TODO: examples
    '''

    def __init__(self, baseFile, baseDir):
        self._baseTable = atpy.Table(baseFile, type='starlightv4_base', basedir=baseDir, read_basedir=True)
        self.ageBase__t = np.unique(self._baseTable.age_base)
        self.nAges = len(self.ageBase__t)

        self.metBase__t = np.unique(self._baseTable.Z_base)
        self.nMet = len(self.metBase__t)
        
        self.l_ssp__l = self._baseTable.l_ssp[0]
        if (self._baseTable.l_ssp != self.l_ssp__l).any():
            raise 'Error! Base is not well defined on all elements in wavelength!' #Base must be consistent on lambdas! Check this here.
        self.nWavelength = len(self.l_ssp__l)
         
        # Base spectra using the original wavelength sampling.
        self.f_ssp__tZl = self.from_j_to_tZ(self._baseTable.f_ssp, n=self.nWavelength)
        
        # Base files, for whatever reason.
        self.sspfile = self._baseTable.sspfile
        
        # Interpolation function for the spectra.
        self._interp_f_ssp__tZl = self._getInterpSpectra()
    
    
    def from_j_to_tZ(self, a, n=1):
        '''
        Convert an array from 1-D j-notation to 2-D (Z, age)-notation.
        '''
        return np.transpose(a.reshape((self.nMet, self.nAges, n)), axes=(1, 0, 2))
        
        
    def _getInterpSpectra(self):
        '''
        Interpolation function for latter use.
        Note the fill value of 0 when interpolating out of bounds.
        '''
        return interp1d(self.l_ssp__l, self.f_ssp__tZl,
                        axis=2, bounds_error=False, fill_value=0)


    def getLookbacktimeSpectraOperator(self, evoBase__T, wavelength=None):
        '''
        Returns an operator for computing spectra in lookbacktime.
        If wavelength is None, use the base wavelength array.
        
        Examples
        ========
        evoBase__T = 10.0**np.arange(8, 10, 0.1) # log_10 steps
        mini__tZ = (popmu_ini__tZ * Mini_tot / 100.0)
        f_ssp__TtZl = base.getLookbacktimeSpectraOperator(evoBase__T)
        F__Tl = np.tensordot(f_ssp__TtZl, mini__tZ, [(1, 2), (0, 1)])
        
        # plot all the spectra.
        plt.figure()
        for i, age in enumerate(evoBase__T):
            plt.plot(base.l_ssp__l, F__Tl[i], label=age)
        plt.legend()
        plt.show()
        '''
        interpolated_f_ssp__tZl = self._interp_f_ssp__tZl(wavelength)
        f_ssp__TtZl = np.zeros((len(evoBase__T), self.nAges, self.nMet, len(wavelength)))
        for i, T in enumerate(evoBase__T):
            mask = self.ageBase__t >= T
            sel = np.digitize(self.ageBase__t[mask] - T, self.ageBase__t)
            f_ssp__TtZl[i][mask] = interpolated_f_ssp__tZl[sel]
        return f_ssp__TtZl