Ticket #482: psPolynomials.c

# include "psphot.h"

// write out the terms of the given 1D polynomial
void psPolynomial1DDump (psPolynomial1D *poly) {

    for (int i = 0; i < poly->n + 1; i++) {
        fprintf (stderr, "x^%d : %g +/- %g\n", i, poly->coeff[i], poly->coeffErr[i]);
    }
}    

psF32 Polynomial1DEval_EAM(psF32 x, const psPolynomial1D* myPoly)
{
    psS32 loop_x = 0;
    psF32 polySum = 0.0;
    psF32 xSum = 1.0;

    for (loop_x = 0; loop_x < myPoly->n + 1; loop_x++) {
        if (myPoly->mask[loop_x] == 0) {
            polySum += xSum * myPoly->coeff[loop_x];
        }
        xSum *= x;
    }

    return(polySum);
}

psVector *Polynomial1DEvalVector_EAM(const psPolynomial1D *myPoly,
                                       const psVector *x)
{
    PS_POLY_CHECK_NULL(myPoly, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_TYPE(x, PS_TYPE_F64, NULL);

    psVector *tmp;

    tmp = psVectorAlloc(x->n, PS_TYPE_F64);
    for (psS32 i=0;i<x->n;i++) {
        tmp->data.F64[i] = Polynomial1DEval_EAM(x->data.F64[i], myPoly);
    }

    return(tmp);
}

// XXX EAM : use Nterm = Norder + 1 definition  
// XXX EAM : should we provide both order and nterms in struct?
psPolynomial1D* Polynomial1DAlloc(psS32 nOrder,
                                  psPolynomialType type)
{
    PS_INT_CHECK_NON_NEGATIVE(nOrder, NULL);

    psS32 i = 0;
    psS32 nTerm = nOrder + 1;
    psPolynomial1D* newPoly = NULL;

    newPoly = (psPolynomial1D* ) psAlloc(sizeof(psPolynomial1D));
    // p_psMemSetDeallocator(newPoly, (psFreeFcn) polynomial1DFree);
    // XXX EAM : me, being lazy

    newPoly->type = type;
    newPoly->n = nOrder;
    newPoly->coeff = (psF32 *)psAlloc(nTerm * sizeof(psF32));
    newPoly->coeffErr = (psF32 *)psAlloc(nTerm * sizeof(psF32));
    newPoly->mask = (psU8 *)psAlloc(nTerm * sizeof(psU8));
    for (i = 0; i < nTerm; i++) {
        newPoly->coeff[i] = 0.0;
        newPoly->coeffErr[i] = 0.0;
        newPoly->mask[i] = 0;
    }
    return(newPoly);
}

// XXX EAM : my alternate BuildSums1D
static psVector *BuildSums1D(psVector* sums, 
                             psF64 x,
                             psS32 nTerm)
{
    psS32 nSum = 0;
    psF64 xSum = 0.0;

    nSum = 2*nTerm;
    if (sums == NULL) {
        sums = psVectorAlloc(nSum, PS_TYPE_F64);
    }
    if (nSum > sums->n) {
        sums = psVectorRealloc(sums, nSum);
    }

    xSum = 1.0;
    for (int i = 0; i < nSum; i++) {
        sums->data.F64[i] = xSum;
        xSum *= x;
    }
    return (sums);
}

// XXX EAM : test version of 1d fitting
psPolynomial1D* VectorFitPolynomial1DOrd_EAM(psPolynomial1D* myPoly,
                                             psVector *mask,
                                             const psVector *x,
                                             const psVector *y,
                                             const psVector *yErr)
{
    // I think this is 1 dimension down
    psImage*      A = NULL;
    psVector*     B = NULL;
    psVector* xSums = NULL;
    psS32 nTerm;
    psF64 wt;

    psTrace(".psLib.dataManip.VectorFitPolynomial1DOrd", 4,
            "---- VectorFitPolynomial1DOrd() begin ----\n");

    // dump minutiae
# ifndef PS_NO_TRACE
    if (psTraceGetLevel (".psLib.dataManip.VectorFitPolynomial1DOrd") >= 5) {
      FILE *f = psTraceGetDestination ();
      fprintf (f, "VectorFitPolynomial1D()\n");
      for (int i = 0; i < x->n; i++) {
        fprintf (f, "(x, y, yErr) is (%f, %f, %f)\n", x->data.F64[i], y->data.F64[i], yErr->data.F64[i]);
      }
    }
# endif
      
    nTerm = myPoly->n + 1;
    A     = psImageAlloc(nTerm, nTerm, PS_TYPE_F64);
    B     = psVectorAlloc(nTerm, PS_TYPE_F64);

    // Initialize data structures (why is this not a function!)
    for (int i = 0; i < nTerm; i++) {
        B->data.F64[i] = 0.0;
        for (int j = 0; j < nTerm; j++) {
            A->data.F64[i][j] = 0.0;
        }
    }

    // xSums look like: 1, x, x^2, ... x^(2n+1)

    // Build the B and A data structs.
    for (int k = 0; k < x->n; k++) {
        if ((mask != NULL) && mask->data.U8[k]) continue;
        xSums = BuildSums1D(xSums, x->data.F64[k], nTerm);
      
        if (yErr == NULL) {
            wt = 1.0;
        } else {
            // this should probably by yErr^2 !!
            // this should filter yErr == 0 values
            wt = 1.0 / PS_SQR(yErr->data.F64[k]);
        }
        for (int i = 0; i < nTerm; i++) {
            B->data.F64[i] += y->data.F64[k] * xSums->data.F64[i] * wt;
        }

        // we could skip half of the array and assign at the end
        // we must handle masked orders
        for (int i = 0; i < nTerm; i++) {
            for (int j = 0; j < nTerm; j++) {
                A->data.F64[i][j] += xSums->data.F64[i + j] * wt;
            }
        }
    }
 
    // GaussJordan version
    if (0) {
        // does the solution in place
        psGaussJordan (A, B);
    
        // the first nTerm entries in B correspond directly to the desired 
        // polynomial coefficients.  this is only true for the 1D case
        for (int k = 0; k < nTerm; k++) {
            myPoly->coeff[k] = B->data.F64[k];
        }
    } 
    else 
    // LUD version of the fit
    {
        psImage *ALUD = NULL;
        psVector* outPerm = NULL;
        psVector* coeffs = NULL;
    
        ALUD = psImageAlloc(nTerm, nTerm, PS_TYPE_F64);
        ALUD = psMatrixLUD(ALUD, &outPerm, A);
        coeffs = psMatrixLUSolve(coeffs, ALUD, B, outPerm);
        for (int k = 0; k < nTerm; k++) {
            myPoly->coeff[k] = coeffs->data.F64[k];
        }
    }

    psFree(A);
    psFree(B);
    psFree(xSums);

    psTrace(".psLib.dataManip.VectorFitPolynomial1DOrd", 4,
            "---- VectorFitPolynomial1DOrd() begin ----\n");
    return (myPoly);
}

// ********************** 2D polynomial functions ******************

// XXX EAM : this version uses myPoly->nX as Norder, not Nterms
psVector *Polynomial2DEvalVector(const psPolynomial2D *myPoly,
                                 const psVector *x,
                                 const psVector *y)

{
    PS_POLY_CHECK_NULL(myPoly, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_TYPE(x, PS_TYPE_F32, NULL);
    PS_VECTOR_CHECK_NULL(y, NULL);
    PS_VECTOR_CHECK_TYPE(y, PS_TYPE_F32, NULL);

    psVector *tmp;
    psS32 vecLen=x->n;

    // Determine the length of the output vector to by the minimum of the x,y vectors
    if (y->n < vecLen) {
        vecLen = y->n;
    }

    // Create output vector to return
    tmp = psVectorAlloc(vecLen, PS_TYPE_F32);

    // Evaluate the polynomial at the specified points
    for (psS32 i=0; i<vecLen; i++) {
        tmp->data.F32[i] = Polynomial2DEval(myPoly,x->data.F32[i],y->data.F32[i]);
    }

    // Return output vector
    return(tmp);
}

// XXX EAM : this version uses the F64 vectors
psVector *Polynomial2DEvalVectorD(const psPolynomial2D *myPoly,
                                  const psVector *x,
                                  const psVector *y)

{
    PS_POLY_CHECK_NULL(myPoly, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_TYPE(x, PS_TYPE_F64, NULL);
    PS_VECTOR_CHECK_NULL(y, NULL);
    PS_VECTOR_CHECK_TYPE(y, PS_TYPE_F64, NULL);

    psVector *tmp;
    psS32 vecLen=x->n;

    // Determine the length of the output vector to by the minimum of the x,y vectors
    if (y->n < vecLen) {
        vecLen = y->n;
    }

    // Create output vector to return
    tmp = psVectorAlloc(vecLen, PS_TYPE_F64);

    // Evaluate the polynomial at the specified points
    for (psS32 i=0; i<vecLen; i++) {
        tmp->data.F64[i] = Polynomial2DEval(myPoly,x->data.F64[i],y->data.F64[i]);
    }

    // Return output vector
    return(tmp);
}

// XXX EAM : use Nterm = Norder + 1 definition  
// the user requests a polynomial of order Norder
psPolynomial2D* Polynomial2DAlloc(psS32 nXorder, psS32 nYorder,
                                    psPolynomialType type)
{
    PS_INT_CHECK_NON_NEGATIVE(nXorder, NULL);
    PS_INT_CHECK_NON_NEGATIVE(nYorder, NULL);

    psS32 x = 0;
    psS32 y = 0;
    psS32 nXterm = nXorder + 1;
    psS32 nYterm = nYorder + 1;
    psPolynomial2D* newPoly = NULL;

    newPoly = (psPolynomial2D* ) psAlloc(sizeof(psPolynomial2D));
    // p_psMemSetDeallocator(newPoly, (psFreeFcn) polynomial2DFree);
    // XXX EAM : me, being lazy

    newPoly->type = type;
    newPoly->nX = nXorder;
    newPoly->nY = nYorder;

    newPoly->coeff = (psF32 **)psAlloc(nXterm * sizeof(psF32 *));
    newPoly->coeffErr = (psF32 **)psAlloc(nXterm * sizeof(psF32 *));
    newPoly->mask = (psU8 **)psAlloc(nXterm * sizeof(psU8 *));
    for (x = 0; x < nXterm; x++) {
        newPoly->coeff[x] = (psF32 *)psAlloc(nYterm * sizeof(psF32));
        newPoly->coeffErr[x] = (psF32 *)psAlloc(nYterm * sizeof(psF32));
        newPoly->mask[x] = (psU8 *)psAlloc(nYterm * sizeof(psU8));
    }
    for (x = 0; x < nXterm; x++) {
        for (y = 0; y < nYterm; y++) {
            newPoly->coeff[x][y] = 0.0;
            newPoly->coeffErr[x][y] = 0.0;
            newPoly->mask[x][y] = 0;
        }
    }
    return(newPoly);
}

// XXX EAM : BuildSums2D in analogy with BuildSums1D
static psImage *BuildSums2D(psImage* sums,
                            psF64 x,      psF64 y,
                            psS32 nXterm, psS32 nYterm)
{
    psS32 nXsum = 0;
    psS32 nYsum = 0;
    psF64 xSum = 1.0;
    psF64 ySum = 1.0;

    nXsum = 2*nXterm;
    nYsum = 2*nYterm;
    if (sums == NULL) {
        sums = psImageAlloc(nXsum, nYsum, PS_TYPE_F64);
    }
    if ((nXsum != sums->numCols) || (nYsum != sums->numRows)) {
        psFree (sums);
        sums = psImageAlloc(nXsum, nYsum, PS_TYPE_F64);
    }

    ySum = 1.0;
    for (int j = 0; j < nYsum; j++) {
        xSum = ySum;
        for (int i = 0; i < nXsum; i++) {
            sums->data.F64[i][j] = xSum;
            xSum *= x;
        }
        ySum *= y;
    }
    return (sums);
}

// XXX EAM : test version of 2d fitting
psPolynomial2D* VectorFitPolynomial2DOrd_EAM(psPolynomial2D* myPoly,
                                             psVector* mask,
                                             const psVector* x,
                                             const psVector* y,
                                             const psVector* z,
                                             const psVector* zErr)
{
    // I think this is 1 dimension down
    psImage*      A = NULL;
    psVector*     B = NULL;
    psImage*   Sums = NULL;
    psF64 wt;
    psS32 nTerm, nXterm, nYterm;

    nXterm = myPoly->nX + 1;
    nYterm = myPoly->nY + 1;
    nTerm = nXterm * nYterm;

    A = psImageAlloc(nTerm, nTerm, PS_TYPE_F64);
    B = psVectorAlloc(nTerm, PS_TYPE_F64);

    // Initialize data structures (why is this not a function!)
    for (int i = 0; i < nTerm; i++) {
        B->data.F64[i] = 0.0;
        for (int j = 0; j < nTerm; j++) {
            A->data.F64[i][j] = 0.0;
        }
    }

    // Sums look like: 1, x, x^2, ... x^(2n+1), y, xy, x^2y, ... x^(2n+1)

    // Build the B and A data structs.
    for (int k  = 0; k < x->n; k++) {
      if ((mask != NULL) && mask->data.U8[k]) continue;
        Sums = BuildSums2D(Sums, x->data.F64[k], y->data.F64[k], nXterm, nYterm);
      
        if (zErr == NULL) {
            wt = 1.0;
        } else {
            // this should probably by zErr^2 !!
            // this should filter zErr == 0 values
            wt = 1.0 / zErr->data.F64[k];
        }

        // we could skip half of the array and assign at the end
        // we must handle masked orders
        for (int n = 0; n < nXterm; n++) {
            for (int m = 0; m < nYterm; m++) {
                B->data.F64[n+m*nXterm] += z->data.F64[k] * Sums->data.F64[n][m] * wt;
            }
        }

        for (int i = 0; i < nXterm; i++) {
            for (int j = 0; j < nYterm; j++) {
                for (int n = 0; n < nXterm; n++) {
                    for (int m = 0; m < nYterm; m++) {
                        A->data.F64[i+j*nXterm][n+m*nXterm] += Sums->data.F64[i+n][j+m] * wt;
                    }
                }
            }
        }
    }
 
    // does the solution in place
    psGaussJordan (A, B);
    
    // XXX: How do we know if these routines were successful?
    // ALUD = psMatrixLUD(ALUD, &outPerm, A);
    // coeffs = psMatrixLUSolve(coeffs, ALUD, B, outPerm);

    for (int n = 0; n < nXterm; n++) {
        for (int m = 0; m < nYterm; m++) {
            myPoly->coeff[n][m] = B->data.F64[n+m*nXterm];
        }
    }

    psFree(A);
    psFree(B);
    psFree(Sums);

    psTrace(".psLib.dataManip.VectorFitPolynomial2DOrd", 4,
            "---- VectorFitPolynomial2DOrd() begin ----\n");
    return (myPoly);
}

// write out the terms of the given 2D polynomial
void psPolynomial2DDump (psPolynomial2D *poly) {

    for (int i = 0; i < poly->nX + 1; i++) {
        for (int j = 0; j < poly->nY + 1; j++) {
            fprintf (stderr, "x^%d y^%d : %g +/- %g\n", i, j, poly->coeff[i][j], poly->coeffErr[i][j]);
        }
    }
}    

psPolynomial2D* RobustFit2D_nomask(psPolynomial2D* poly,
                            const psVector* x,
                            const psVector* y,
                            const psVector* z,
                            const psVector* dz)
{
    psVector *X;
    psVector *Y;
    psVector *Z;
    psVector *dZ;

    psVector *zFit   = NULL;
    psVector *zResid = NULL;
    psStats  *stats  = NULL;

    X  = psVectorCopy (NULL, x, PS_TYPE_F64);
    Y  = psVectorCopy (NULL, y, PS_TYPE_F64);
    Z  = psVectorCopy (NULL, z, PS_TYPE_F64);
    dZ = psVectorCopy (NULL, dz, PS_TYPE_F64);

    for (int N = 0; N < 3; N++) {
        // XXX EAM : this would be better defined with an element mask
        poly   = VectorFitPolynomial2DOrd_EAM (poly, NULL, X, Y, Z, dZ);
        zFit   = Polynomial2DEvalVectorD (poly, x, y);
        zResid = (psVector *) psBinaryOp (NULL, (void *) z, "-", (void *) zFit);

        stats = psStatsAlloc (PS_STAT_CLIPPED_MEAN | PS_STAT_CLIPPED_STDEV);
        stats  = psVectorStats (stats, zResid, NULL, NULL, 0);
        psTrace (".psphot.RobustFit", 4, "residual stats for robust fit:  %g +/- %g (%d pts)\n", stats->clippedMean, stats->clippedStdev, stats->clippedNvalues);

        // re-create X, Y, Z, dZ if pts are valid
        int n = 0;
        for (int i = 0; i < zResid->n; i++) {
            if (fabs(zResid->data.F64[i] - stats->clippedMean) > 3*stats->clippedStdev) {
              continue;
            } 
            X->data.F64[n]  =  x->data.F64[i];
            Y->data.F64[n]  =  y->data.F64[i];
            Z->data.F64[n]  =  z->data.F64[i];
            dZ->data.F64[n] = dz->data.F64[i];
            n++;
        }
        X->n = n;
        Y->n = n;
        Z->n = n;
        dZ->n = n;
    }
    return (poly);
}

// XXX EAM : be careful here with F32 vs F64 vectors 
psPolynomial2D* RobustFit2D(psPolynomial2D* poly,
                            psVector* mask,
                            const psVector* x,
                            const psVector* y,
                            const psVector* z,
                            const psVector* dz)
{
    PS_VECTOR_CHECK_NULL(mask, NULL);
    PS_VECTOR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_NULL(y, NULL);
    PS_VECTOR_CHECK_NULL(z, NULL);
    PS_VECTOR_CHECK_NULL(dz, NULL);

    psVector *zFit   = NULL;
    psVector *zResid = psVectorAlloc (x->n, PS_TYPE_F64);
    psStats  *stats  = psStatsAlloc (PS_STAT_SAMPLE_MEAN | PS_STAT_SAMPLE_STDEV);

    for (int N = 0; N < 3; N++) {
        poly   = VectorFitPolynomial2DOrd_EAM (poly, mask, x, y, z, dz);
        zFit   = Polynomial2DEvalVectorD (poly, x, y);
        zResid = (psVector *) psBinaryOp (zResid, (void *) z, "-", (void *) zFit);

        stats  = psVectorStats (stats, zResid, NULL, mask, 1);
        psTrace (".psphot.RobustFit", 4, "residual stats for robust fit:  %g +/- %g\n", 
                 stats->sampleMean, stats->sampleStdev);

        // set mask if pts are not valid
        // we are masking out any point which is out of range 
        // recovery is not allowed with this scheme
        for (int i = 0; i < zResid->n; i++) {
          if (mask->data.U8[i]) continue;
          if (fabs(zResid->data.F64[i] - stats->sampleMean) > 3*stats->sampleStdev) {
            mask->data.U8[i] = 1;
            continue;
          }       
        }
        psFree (zFit);
    }
    psFree (zResid);
    psFree (stats);
    return (poly);
}

// XXX EAM : VectorFitPolynomial2DOrd and Polynomial2DEvalVector require different types (F32 vs F64)

# if (0) // moved to psLib
// XXX EAM : this version uses myPoly->nX as Norder, not Nterms
psF32 Polynomial2DEval(const psPolynomial2D* myPoly,
                       psF32 x,
                       psF32 y)
{
    PS_POLY_CHECK_NULL(myPoly, NAN);

    psS32 loop_x = 0;
    psS32 loop_y = 0;
    psF32 polySum = 0.0;
    psF32 xSum = 1.0;
    psF32 ySum = 1.0;

    // XXX EAM : nX is order, not nTerms
    for (loop_x = 0; loop_x < myPoly->nX + 1; loop_x++) {
        ySum = xSum;
        // XXX EAM : nX is order, not nTerms
        for (loop_y = 0; loop_y < myPoly->nY + 1; loop_y++) {
            if (myPoly->mask[loop_x][loop_y] == 0) {
                polySum += ySum * myPoly->coeff[loop_x][loop_y];
            }
            ySum *= y;
        }
        xSum *= x;
    }

    return(polySum);
}
# endif