Fitting a partially linear model

gistfile1.py

def fit1T(pNL0,t,f):
    """
    Fit Single transit
    """
    dpNL0 = np.array([0.2,0.2])

    objp = lambda p,t,f : obj1Tlin(p,t,f) + (((p-pNL0)/dpNL0)**2).sum()
    pNL = optimize.fmin(objp,pNL0,args=(t,f),disp=False)
    pL = linfit1T(pNL,t,f)
    pFULL = np.hstack( (pNL[0],pL[0],pNL[1],pL[1:]) )
    if pFULL[1] < 0:
        pFULL[1] = 0

    return pFULL


def obj1Tlin(pNL,t,f):
    """
    Single Transit Objective Function.  For each trial value of epoch
    and width, we determine the best fit by linear fitting.

    Parameters
    ----------
    pNL  - The non-linear parameters [epoch,tdur]
   
    """
    pL = linfit1T(pNL,t,f)
    pFULL = np.hstack( (pNL[0],pL[0],pNL[1],pL[1:]) )
    model = keptoy.P051T(pFULL,t)

    resid  = ((model - f)/1e-4 )
    obj = (resid**2).sum() 
    return obj


def linfit1T(p,t,f):
    """
    Linear fit to 1 Transit.

    Depth and polynomial cofficents are linear

    Parameters
    ----------
    p : [epoch,tdur]
    t : time
    f : flux

    Returns
    -------
    p1 : Best fit [df,pleg0,pleg1...] from linear fitting.
    """
    
    epoch  = p[0]
    tdur   = p[1]
    tDS = trendDS(t)
    ndeg = 3
    # Construct lightcurve design matrix
    plc = np.hstack(( epoch,1.,tdur,list(np.zeros(ndeg+1)) ))
    lcDS = keptoy.P051T(plc,t)

    DS = np.vstack((lcDS,tDS))
    p1 = np.linalg.lstsq(DS.T,f)[0]
    return p1


def trendDS(t):
    ndeg=3
    # Construct polynomial design matrix
    tDS = [] 
    for i in range(ndeg+1):
        pleg = np.zeros(ndeg+1)
        pleg[i] = 1
        tDS.append( keptoy.trend(pleg,t) )
    tDS = np.vstack(tDS)
    return tDS