Trying to emulate IRAF's noao.onedspec.continuum

continuum.py

import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import LSQUnivariateSpline


def continuum(wave, flux, type='ratio', order=1, low_reject=2, high_reject=0,
              niter=10):
    
    m1 = np.ones_like(wave, dtype=np.bool) # use all points at first
    m1 &= flux!=0 # but remove those where flux = 0
    
    knots = np.linspace(wave[m1][1], wave[m1][-2], order)
    ff = LSQUnivariateSpline(wave[m1], flux[m1], t=knots)

    for _ in range(niter):
        r = flux - ff(wave) # residuals
        sig = r.std() # sigma
        m2 = r < - low_reject * sig
        ff = LSQUnivariateSpline(wave[m1 & ~m2], flux[m1 & ~m2], t=knots)
    
    if type == 'ratio':
        return wave[m1], flux[m1] / ff(wave[m1]), ff
    elif type == 'fit':
        return wave[m1], ff(wave[m1]), ff
    elif type == 'diff':
        return wave[m1], flux[m1] - ff(wave[m1]), ff