Peter Müller / NOD3 single dish reduction software

import random
import numpy
from scipy import optimize, signal

def Gaussian(off, dx, dy, height, center_x, center_y, width_x, width_y):

    """Returns a gaussian function with the given parameters"""

    return lambda x,y: off + dx*x + dy*y + height*numpy.exp(
                -(((center_x-x)/width_x)**2+((center_y-y)/width_y)**2)/2)

def gaussian(off, dx, dy, height, x0, y0, wx, wy, rot):

    """Returns a gaussian function with the given parameters"""

    return lambda x,y: off + dx*x + dy*y + height*numpy.exp(
                -((((x-x0)*numpy.cos(rot)-(y-y0)*numpy.sin(rot))/wx)**2
                + (((x-x0)*numpy.sin(rot)+(y-y0)*numpy.cos(rot))/wy)**2)/2)

def gaussmoments(data, x, y):

    """Returns (off, dx, dy, height, x, y, width_x, width_y, rotation)
    the gaussian parameters of a 2D distribution by calculating its
    moments """

    off = 0.0
    dx = 0.0
    dy = 0.0
    rot = 0.0
    width_x = 1.0
    width_y = 1.0
    height = data.max() - data.min()
    return off, dx, dy, height, x, y, width_x, width_y, rot

def fitgaussian(data, l, b, fast=False):

    """Returns (off, dx, dy, height, x, y, width_x, width_y, rot)
    the gaussian parameters of a 2D distribution found by a fit"""

    params = gaussmoments(data, l, b)
    #errorfunction = lambda p: numpy.ravel(gaussian(*p)(*numpy.indices(data.shape)) - data)
    errorfunction = lambda p: numpy.ravel(gaussian(*p)(*numpy.indices(data.shape)) - data)
    if fast:
       try:
          p, cov, info, mesg, success = \
          optimize.leastsq(errorfunction, params, full_output=1, \
                           maxfev=200, ftol=1.e-6, xtol=1.e-5)
       except: 
          return None, None
    else:
       try:
          p, cov, info, mesg, success = \
          optimize.leastsq(errorfunction, params, full_output=1)
       except: 
          return None, None
    #if success > 0 and success < 3: 
    if numpy.shape(cov) != (): 
       chisq = numpy.sum(info["fvec"]*info["fvec"]) 
       dof = len(p)
       rms = numpy.sqrt(chisq/float(dof))
       error = []
       for i in range(dof):
           error.append(numpy.sqrt(numpy.abs(cov[i,i]))*rms)
       return p, error
    else: 
       return None, None

def gauss_kern(size, sizey=None):
    """ Returns a normalized 2D gauss kernel array for convolutions """
    size = int(size)
    if not sizey:
       sizey = size
    else:
       sizey = int(sizey)
    x, y = numpy.mgrid[-size:size+1, -sizey:sizey+1]*0.1
    g = numpy.exp(-(x**2/float(size) + y**2/float(sizey)))
    return g / g.sum()

def smooth_image(im, n, ny=None):
    """ blurs the image by convolving with a gaussian kernel of typical
        size n. The optional keyword argument ny allows for a different
        size in the y direction.
    """
    g = gauss_kern(n, sizey=ny)
    improc = signal.convolve(im, g, mode='same')
    return(improc)

def sinc_kern(size, sizey=None):
    """ Returns a normalized 2D sinc kernel array for convolutions """
    size = int(size)
    if not sizey:
       sizey = size
    else:
       sizey = int(sizey)
    x, y = numpy.mgrid[-size:size+1, -sizey:sizey+1] * numpy.pi/180
    g = numpy.sinc(x * y)
    return g / g.sum()

def pump_image(im, n, ny=None):
    """ resize the image by convolving with a sinc kernel of typical
        size n. The optional keyword argument ny allows for a different
        size in the y direction.
    """
    g = sinc_kern(n, sizey=ny)
    improc = signal.convolve(im, g, mode='same')
    return improc
     

def ausgleich(inp, g, EPS=1e-9):

    n = len(inp[0])
    m = len(inp)
    if m < n:
        n = min(n, m)
        for i in range(m):
            inp[i][n-1] = inp[i][-1]

    q = []
    out = []
    mat = []
    sum = []
    for i in range(n):
        q.append(0.0)
        out.append(0.0)
        mat.append([])
        for j in range(n):
            mat[i].append(0.0)

    for i in range(n):
        sum.append([])
        for j in range(n):
            sum[i].append([])
            for k in range(n):
                sum[i][j].append(0.0)
    for i in range(n):
        for j in range(i,n):
            for k in range(m):
                sum[0][i][j] += inp[k][i] * inp[k][j] * g[k]

    for k in range(n-1):
        if abs(sum[k][k][k]) < EPS:
            print "no solution possible", k, '1'
            return None
        for j in range(k, n-1):
            for i in range(j, n-1):
                sum[k+1][j+1][i+1] = sum[k][j+1][i+1] - sum[k][k][j+1] \
                                   * sum[k][k][i+1] / sum[k][k][k]

    for k in range(n-2, -1, -1):
        if abs(sum[k][k][k]) < EPS:
            print "no solution possible", k, '2'
            return None
        for i in range(n-1, k, -1):
            if i == n-1:
                out[k] = sum[k][k][i] / sum[k][k][k]
            else:
                out[k] -= out[i] * sum[k][k][i] / sum[k][k][k]

    return out[:-1]


def _general_function(params, xdata, ydata, function):
    return function(xdata, *params) - ydata

def _weighted_general_function(params, xdata, ydata, function, weights):
    return weights * (function(xdata, *params) - ydata)

def curve_fit(f, xdata, ydata, p0=None, sigma=None, absolute_sigma=False, **kw):
    """
    Use non-linear least squares to fit a function, f, to data.
    Assumes ``ydata = f(xdata, *params) + eps``
    Parameters
    ----------
    f : callable
    The model function, f(x, ...). It must take the independent
    variable as the first argument and the parameters to fit as
    separate remaining arguments.
    xdata : An M-length sequence or an (k,M)-shaped array
    for functions with k predictors.
    The independent variable where the data is measured.
    ydata : M-length sequence
    The dependent data --- nominally f(xdata, ...)
    p0 : None, scalar, or N-length sequence
    Initial guess for the parameters. If None, then the initial
    values will all be 1 (if the number of parameters for the function
    can be determined using introspection, otherwise a ValueError
    is raised).
    sigma : None or M-length sequence, optional
    If not None, these values are used as weights in the
    least-squares problem.
    absolute_sigma : bool, optional
    If False, `sigma` denotes relative weights of the data points.
    The returned covariance matrix `pcov` is based on *estimated*
    errors in the data, and is not affected by the overall
    magnitude of the values in `sigma`. Only the relative
    magnitudes of the `sigma` values matter.
    If True, `sigma` describes one standard deviation errors of
    the input data points. The estimated covariance in `pcov` is
    based on these values.
    Returns
    -------
    popt : array
    Optimal values for the parameters so that the sum of the squared error
    of ``f(xdata, *popt) - ydata`` is minimized
    pcov : 2d array
    The estimated covariance of popt. The diagonals provide the variance
    of the parameter estimate. To compute one standard deviation errors
    on the parameters use ``perr = np.sqrt(np.diag(pcov))``.
    How the `sigma` parameter affects the estimated covariance
    depends on `absolute_sigma` argument, as described above.
    See Also
    --------
    leastsq
    Notes
    -----
    The algorithm uses the Levenberg-Marquardt algorithm through `leastsq`.
    Additional keyword arguments are passed directly to that algorithm.
    Examples
    --------
    >>> import numpy as np
    >>> from scipy.optimize import curve_fit
    >>> def func(x, a, b, c):
    ... return a * np.exp(-b * x) + c
    >>> xdata = np.linspace(0, 4, 50)
    >>> y = func(xdata, 2.5, 1.3, 0.5)
    >>> ydata = y + 0.2 * np.random.normal(size=len(xdata))
    >>> popt, pcov = curve_fit(func, xdata, ydata)
    """

    if p0 is None:
       # determine number of parameters by inspecting the function
       import inspect
       args, varargs, varkw, defaults = inspect.getargspec(f)
       if len(args) < 2:
          msg = "Unable to determine number of fit parameters."
          #raise ValueError(msg)
          return ValueError(msg), []
       if 'self' in args:
          p0 = [1.0] * (len(args)-2)
       else:
          p0 = [1.0] * (len(args)-1)

    # Check input arguments
    if numpy.isscalar(p0):
       p0 = array([p0])

    ydata = numpy.asanyarray(ydata)
    if isinstance(xdata, (list, tuple)):
       # `xdata` is passed straight to the user-defined `f`, so allow
       # non-array_like `xdata`.
       xdata = numpy.asarray(xdata)

    args = (xdata, ydata, f)
    if sigma is None:
       func = _general_function
    else:
       func = _weighted_general_function
       args += (1.0 / numpy.asarray(sigma),)

    # Remove full_output from kw, otherwise we're passing it in twice.
    return_full = kw.pop('full_output', False)
    res = optimize.leastsq(func, p0, args=args, full_output=1, **kw)
    (popt, pcov, infodict, errmsg, ier) = res

    if ier not in [1, 2, 3, 4]:
       msg = "Optimal parameters not found: " + errmsg
       raise RuntimeError(msg)
       #return RuntimeError(msg), []

    if pcov is None:
       # indeterminate covariance
       pcov = numpy.zeros((len(popt), len(popt)), dtype=float)
       pcov.fill(numpy.inf)
       reduced_chi_sq = 99999999.0
    elif not absolute_sigma:
       if len(ydata) > len(p0):
          s_sq = (numpy.asarray(func(popt, *args))**2).sum() / (len(ydata) - len(p0) - 1)
          pcov = pcov * s_sq
       else:
          pcov.fill(numpy.inf)
          reduced_chi_sq = 99999999.0

    if return_full:
       return popt, pcov, infodict, errmsg, ier
    else:
       chi_sq = (numpy.asarray(func(popt, *args))**2).sum()
       return popt, pcov

def correlate(data1, data2, fmax=0.0, out=False, sort=False):
    d1 = numpy.ravel(1.0*data1)
    d2 = numpy.ravel(1.0*data2)
    mask1 = ~numpy.isnan(d1) * ~numpy.isnan(d2)
    mask2 = (d1 > fmax) * (d2 > fmax)
    mask = mask1*mask2
    d1 = d1[mask]
    d2 = d2[mask]
    if sort:
       d1.sort()
       d2.sort()
    d = numpy.array([d1, numpy.ones(len(d1))])
    x = numpy.linalg.lstsq(d.T, d2)
    b, a = x[0]
    if out:
       return a, b, d1, d2
    else:
       return a, b


def main():

    import pylab

    r = numpy.pi/180
    # Create the gaussian data
    Xin, Yin = numpy.mgrid[0:201, 0:201]
    #data = gaussian(0, 0.0133333, 0.01, 3, 100, 100, 50, 40)(Xin, Yin) #+ random.random()
    params = (0, 0.0, 0.0, 1, 101, 101, 15, 25, 30*r)
    data = gaussian(*params)(Xin, Yin)

    #pylab.matshow(data, cmap=pylab.cm.gist_earth_r)
    pylab.matshow(data)

    dmax = data.argmax()
    ydim, xdim = data.shape
    l = dmax/xdim
    b = dmax - l*xdim
    params = fitgaussian(data, l, b)
    bparams = [0.0, 0.0, 0.0] + list(params[0])[3:]
    fit = gaussian(*bparams)

    #pylab.contour(data-fit(*numpy.indices(data.shape)), cmap=pylab.cm.copper)
    pylab.contour(fit(*numpy.indices(data.shape)))
    ax = pylab.gca()
    rot = 0
    (off, a, b, height, x, y, width_x, width_y, rot) = bparams

    pylab.text(0.95, 0.05, """
    x : %.1f
    y : %.1f
    rot : %.1f
    width_x : %.1f
    width_y : %.1f""" %(x, y, (rot/r % 360.0), width_x, width_y),
            fontsize=16, horizontalalignment='right',
            verticalalignment='bottom', transform=ax.transAxes)
    pylab.show()


if __name__ == '__main__':

    main()