text:
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <time.h>

#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_statistics.h>

#define VERSION "2021-03-16"

#define DEBUG 0

/* Set the defaults; you can override these from the command line.
int    POLYORDER = 2;
int    ITERS     = 10000;
double STEP_SIZE = 0.01;
int    NVERT = 200;
int    KNOTS = 4;
int    COLIDX[3] = {-1, -1, -1};
int    FIND_KNOTS = 0;
int    FIND_STEP = 0;

int    CHAIN_LENGTH = 8;

int    APPLY_PSHIFT = 0;
*/

int    CHAIN_LENGTH;
int    KNOTS;
int    POLYORDER;
int    NVERT;
int    APPLY_PSHIFT;

int    CSV_INPUT = 0;
int    CSV_OUTPUT = 0;
int    EBAI_OUTPUT = 0;
int    SUMMARY_OUTPUT = 0;

char buffer[256];
char *messages = "";

/*
double LLIMIT = -1e20;
double ULIMIT = 1e20;
 */
/* Maximum number of observed data points: */
//#define CHUNK 10

typedef struct triplet {
    double x;
    double y;
    double z;
} triplet;

char * combine( const char *s1, const char *s2 )
{
    char *result = malloc( strlen( s1 ) + strlen( s2 ) + 1 );

    if ( result )
    {
        char *p = result;

        while ( *s1 ) *p++ = *s1++;

        while ( ( *p++ = *s2++ ) );
    }

    return result;
}

int sort_by_phase (const void *a, const void *b)
{
    const triplet *da = (const triplet *) a;
    const triplet *db = (const triplet *) b;

    return (da->x > db->x) - (da->x < db->x);
}

int sort_by_value (const void *a, const void *b)
{
    const double *da = (const double *) a;
    const double *db = (const double *) b;

    return (*da > *db) - (*da < *db);
}

int find_knots (triplet *data, int nobs, double **knots)
{
    int i, j;
    double average = 0.0;
    double std = 0.0;
    int chain_too_short;
    int chain_wrapped = 0;

    int chains = 0;
    struct {
        int len;
        int start;
        int end;
    } chain[2];

    for (i = 0; i < nobs; i++)
        average += data[i].y;
    average /= (double) nobs;

    for (i = 0; i < nobs; i++)
        std += pow(data[i].y - average, 2.0);
    std /= (double) (nobs-1);
    std = sqrt(std);

    /* Find eclipses below this limit */
    average = average - std/2.0;

    if (!SUMMARY_OUTPUT && !EBAI_OUTPUT) {
        sprintf (buffer,"# searching for knots automatically:\n");
        messages = combine(messages ,buffer);
        sprintf (buffer,"# * average-std value of the flux: %lf\n", average);
        messages = combine(messages ,buffer);
    }

    /* To allow wrapping of the phase interval, skip the first chain: */
    i = 0;
    while (data[i].y < average) {
        if (DEBUG)
            sprintf (buffer,"# delaying point  %8.4lf (%3d): %lf < %lf\n", data[i].x, i, data[i].y, average);
        i++;
    }

    for ( ; i < nobs; i++) {
        if (data[i].y > average) {
            if (DEBUG)
                sprintf (buffer,"# skipping point  %8.4lf (%3d): %lf > %lf\n", data[i].x, i, data[i].y, average);
            continue;
        }

        /* Check if the chain is at least CHAIN_LENGTH long: */
        chain_too_short = 0;
        if (DEBUG)
            sprintf (buffer,"# chain starts at %8.4lf (%3d): %lf < %lf\n", data[i].x, i, data[i].y, average);
        for (j = 1; j < CHAIN_LENGTH; j++) {
            if (i+j == nobs) {
                i = -j;
                chain_wrapped = 1;
            }
            if (data[i+j].y > average) {
                if (DEBUG)
                    sprintf (buffer,"# chain broken at %8.4lf (%3d): %lf > %lf\n", data[i+j].x, i+j, data[i+j].y, average);
                i += j;
                chain_too_short = 1;
                break;
            }
            else
                if (DEBUG)
                    sprintf (buffer,"# chain cont'd at %8.4lf (%3d): %lf < %lf\n", data[i+j].x, i+j, data[i+j].y, average);
        }

        if (chain_wrapped && chain_too_short) break;
        if (chain_too_short) continue;

        while (data[i+j].y < average) {
            if (i+j == nobs) {
                i = -j;
                chain_wrapped = 1;
            }
            if (DEBUG)
                sprintf (buffer,"# chain cont'd at %8.4lf (%3d): %lf < %lf\n", data[i+j].x, i+j, data[i+j].y, average);
            j++;
        }

        if (chains < 2) {
            chains++;
            chain[chains-1].len = j-1;
            chain[chains-1].start = chain_wrapped*nobs + i;
            chain[chains-1].end = i+j-1;
        }
        else {
            chains++; /* Just to count all chains */
            if (j-1 > chain[0].len) {
                if (chain[0].len > chain[1].len) {
                    chain[1].len = chain[0].len;
                    chain[1].start = chain[0].start;
                    chain[1].end = chain[0].end;
                }
                chain[0].len = j-1;
                chain[0].start = chain_wrapped*nobs + i;
                chain[0].end = i+j-1;
            }
            else if (j-1 > chain[1].len) {
                chain[1].len = j-1;
                chain[1].start = chain_wrapped*nobs + i;
                chain[1].end = i+j-1;
            }
            /* else drop through without recording it because it is shorter. */
        }

        if (!SUMMARY_OUTPUT && !EBAI_OUTPUT){
            sprintf (buffer,"# * found a chain from %lf (index %d/%d) to %lf (index %d/%d)\n", data[chain_wrapped*nobs+i].x, chain_wrapped*nobs+i, nobs-1, data[i+j-1].x, i+j-1, nobs-1);
            messages = combine(messages ,buffer);
        }
        i += j-1;
        if (chain_wrapped) break;
    }

    if (!SUMMARY_OUTPUT && !EBAI_OUTPUT){
        sprintf (buffer,"# * total number of chains found: %d\n", chains);
        messages = combine(messages , buffer);
    }

    /* If the number of chains is less than 2, the search for knots failed. */
    if (chains < 2)
        return -1;

    KNOTS = 4;
    *knots = malloc (KNOTS * sizeof (**knots));
    (*knots)[0] = data[chain[0].start].x; (*knots)[1] = data[chain[0].end].x;
    (*knots)[2] = data[chain[1].start].x; (*knots)[3] = data[chain[1].end].x;

    return 0;
}

double compute_flux (double phase, double *knots, double **ck)
{
    int j, k;
    double knot, dknot, flux;

    while (phase < -0.5) phase += 1.0;
    while (phase >  0.5) phase -= 1.0;

    if (phase < knots[0]) {
        j = KNOTS-1;
        knot = knots[j]-1.0;
        dknot = knots[0]-knots[KNOTS-1]+1.0;
    }
    else if (phase > knots[KNOTS-1]) {
        j = KNOTS-1;
        knot = knots[j];
        dknot = knots[0]-knots[KNOTS-1]+1.0;
    }
    else {
        j = 0;
        while (knots[j+1] < phase && j < KNOTS-1) j++;
        knot = knots[j];
        dknot = knots[j+1]-knots[j];
    }

    flux = ck[j][POLYORDER];
    for (k = 0; k < POLYORDER; k++) {
        flux += ck[j][k] * pow (phase-knot, POLYORDER-k);
        if (j == KNOTS-1 && k < POLYORDER-1)
            flux -= ck[j][k] * (phase-knot) * pow (dknot, POLYORDER-k-1);
    }

    return flux;
}

double polyfit (triplet *data, int nobs, int KNOTS, double *knots, int VERBOSE, int PRINT, double *outphase, double *outpolyfit)
{
    int i, j, k, intervals, kfinal, cum, int1index;
    double chisq, chi2tot, knot, pshift = 0.0;
    double **ck;

    int *n; /* Array of numbers of data points per interval */
    gsl_vector **x, **y, **w, **c;
    gsl_matrix **A, **cov;
    gsl_multifit_linear_workspace **mw;

    /* Count data points between knots: */
    n = calloc (KNOTS, sizeof (*n));

    /* Fast-forward to the first knot: */
    i = 0; cum = 0;
    while (data[i].x < knots[0]) { i++; cum++; }
    n[KNOTS-1] = i;
    int1index = cum;

    for (j = 0; j < KNOTS-1; j++) {
        while (data[i].x < knots[j+1] && i < nobs) i++;
        n[j] += i-cum; /* we need += because of the wrapped interval */
        if (n[j] <= POLYORDER) {
            /* The intervals between knots don't have enough points; bail out. */
            free (n);
            return 1e10;
        }
        cum += n[j];
    }

    /* Add post-last knot data to the last interval: */
    n[j] += nobs-cum;
    if (n[j] <= POLYORDER) {
        /* The intervals between knots don't have enough points; bail out. */
        free (n);
        return 1e10;
    }

    if (VERBOSE && !SUMMARY_OUTPUT && !EBAI_OUTPUT) {
        sprintf (buffer,"# Phase space partitioning:\n# \n");
        messages = combine(messages , buffer);
        for (k = 0; k < KNOTS-1; k++){
            sprintf (buffer, "#   interval %2d: [% 3.3lf, % 3.3lf), %3d data points\n", k, knots[k], knots[k+1], n[k]);
            messages = combine(messages ,buffer);
        }
        sprintf (buffer,"#   interval %2d: [% 3.3lf, % 3.3lf), %3d data points\n", k, knots[k], knots[0], n[k]);
        messages = combine(messages ,buffer);
    }

    /* Allocate arrays of vectors of interval data: */
    x = malloc (KNOTS * sizeof (*x));
    y = malloc (KNOTS * sizeof (*y));
    w = malloc (KNOTS * sizeof (*w));

    for (j = 0; j < KNOTS; j++) {
        x[j] = gsl_vector_alloc (n[j]);
        y[j] = gsl_vector_alloc (n[j]);
        w[j] = gsl_vector_alloc (n[j]);
    }

    /* Copy the (phase-shifted) data to these vectors: */
    cum = 0;
    for (j = 0; j < KNOTS-1; j++) {
        for (i = 0; i < n[j]; i++) {
            gsl_vector_set (x[j], i, data[int1index+cum+i].x - knots[j]);
            gsl_vector_set (y[j], i, data[int1index+cum+i].y);
            gsl_vector_set (w[j], i, data[int1index+cum+i].z);
            if (DEBUG)
                sprintf (buffer,"%d: %lf\t%lf\t%lf\n", j, data[int1index+cum+i].x - knots[j], data[int1index+cum+i].y, data[int1index+cum+i].z);
        }
        cum += n[j];
    }

    /* Copy the wrapped interval data to these vectors: */
    knot = knots[j];
    for (i = 0; i < n[j]; i++) {
        if (int1index+cum+i == nobs) {
            cum = -int1index-i;
            knot -= 1.0;
        }
        gsl_vector_set (x[j], i, data[int1index+cum+i].x - knot);
        gsl_vector_set (y[j], i, data[int1index+cum+i].y);
        gsl_vector_set (w[j], i, data[int1index+cum+i].z);
        if (DEBUG)
            sprintf (buffer,"%d: %lf\t%lf\t%lf\n", j, data[int1index+cum+i].x - knot, data[int1index+cum+i].y, data[int1index+cum+i].z);
    }

    /* Write out a header: */
    if (VERBOSE && !SUMMARY_OUTPUT && !EBAI_OUTPUT) {
        sprintf (buffer,"# \n# Weighted least-squares solution of the polyfit:\n# \n");
        messages = combine(messages ,buffer);
        sprintf (buffer,"#    knot\t\t");
        messages = combine(messages ,buffer);
        for (k = 0; k < POLYORDER+1; k++){
            sprintf (buffer,"   c[%d]\t\t", k);
            messages = combine(messages ,buffer);
        }
        sprintf (buffer,"\n");
        messages = combine(messages ,buffer);
    }

    /* If polynomial order is 1, the last interval is determined by constraints: */
    if (POLYORDER == 1) intervals = KNOTS-1; else intervals = KNOTS;

    /* Allocate all vectors and matrices: */
    A   = malloc (intervals * sizeof (*A));
    c   = malloc (intervals * sizeof (*c));
    cov = malloc (intervals * sizeof (*cov));
    mw  = malloc (intervals * sizeof (*mw));

    /* The first interval has all polynomial coefficients free. */
    A[0]   = gsl_matrix_alloc (n[0], POLYORDER+1);
    c[0]   = gsl_vector_alloc (POLYORDER+1);
    cov[0] = gsl_matrix_alloc (POLYORDER+1, POLYORDER+1);
    mw[0]  = gsl_multifit_linear_alloc (n[0], POLYORDER+1);

    /* Intervals 1...(N-1) are constrained by the connectivity constraint. */
    for (j = 1; j < KNOTS-1; j++) {
        A[j]   = gsl_matrix_alloc (n[j], POLYORDER);
        c[j]   = gsl_vector_alloc (POLYORDER);
        cov[j] = gsl_matrix_alloc (POLYORDER, POLYORDER);
        mw[j]  = gsl_multifit_linear_alloc (n[j], POLYORDER);
    }

    /* The last interval has two constraints (connectivity and periodicity). */
    if (j < intervals) {
        A[j]   = gsl_matrix_alloc (n[j], POLYORDER-1);
        c[j]   = gsl_vector_alloc (POLYORDER-1);
        cov[j] = gsl_matrix_alloc (POLYORDER-1, POLYORDER-1);
        mw[j]  = gsl_multifit_linear_alloc (n[j], POLYORDER-1);
    }

    /* Allocate polynomial coefficients: */
    ck = malloc (KNOTS * sizeof (*ck));
    for (k = 0; k < KNOTS; k++)
        ck[k] = malloc ((POLYORDER+1) * sizeof (**ck));

    /* Set all elements to all matrices: */
    for (j = 0; j < intervals; j++) {
        if      (j == 0)       kfinal = POLYORDER+1;
        else if (j == KNOTS-1) kfinal = POLYORDER-1;
        else                   kfinal = POLYORDER;
        for (i = 0; i < n[j]; i++)
            for (k = 0; k < kfinal; k++)
                gsl_matrix_set (A[j], i, k, pow (gsl_vector_get (x[j], i), POLYORDER-k));
    }

    /*********************   FITTING THE 1ST INTERVAL:   **********************/

    gsl_multifit_wlinear (A[0], w[0], y[0], c[0], cov[0], &chisq, mw[0]);

    for (k = 0; k < POLYORDER+1; k++)
        ck[0][k] = gsl_vector_get (c[0], k);
    chi2tot = chisq;

    if (VERBOSE && !SUMMARY_OUTPUT && !EBAI_OUTPUT) {
        sprintf (buffer,"# % 12.12lf\t", knots[0]);
        messages = combine(messages ,buffer);
        for (k = 0; k < POLYORDER+1; k++){
            sprintf (buffer,"% 12.12lf\t", ck[0][POLYORDER-k]);
            messages = combine(messages ,buffer);
        }
        sprintf (buffer,"\n");
        messages = combine(messages ,buffer);
    }

    /********************   FITTING INTERVALS 2-(N-1):   **********************/

    for (j = 1; j < KNOTS-1; j++) {
        /* Satisfy the connectivity constraint: */
        ck[j][POLYORDER] = ck[j-1][POLYORDER];
        for (k = 0; k < POLYORDER; k++)
            ck[j][POLYORDER] += ck[j-1][k] * pow (knots[j]-knots[j-1], POLYORDER-k);

        /* Apply the connectivity constraint: */
        for (i = 0; i < n[j]; i++)
            gsl_vector_set (y[j], i, gsl_vector_get (y[j], i) - ck[j][POLYORDER]);

        gsl_multifit_wlinear (A[j], w[j], y[j], c[j], cov[j], &chisq, mw[j]);
        chi2tot += chisq;

        for (k = 0; k < POLYORDER; k++)
            ck[j][k] = gsl_vector_get (c[j], k);

        if (VERBOSE && !SUMMARY_OUTPUT && !EBAI_OUTPUT) {
            sprintf (buffer,"# % 12.12lf\t", knots[j]);
            messages = combine(messages ,buffer);
            for (k = 0; k < POLYORDER+1; k++){
                sprintf (buffer,"% 12.12lf\t", ck[j][POLYORDER-k]);
                messages = combine(messages ,buffer);
            }
            sprintf (buffer,"\n");
            messages = combine(messages ,buffer);
        }
    }

    /********************   FITTING THE LAST INTERVAL:   **********************/

    /* Satisfy the connectivity constraint: */
    ck[j][POLYORDER] = ck[j-1][POLYORDER];
    for (k = 0; k < POLYORDER; k++)
        ck[j][POLYORDER] += ck[j-1][k] * pow (knots[j]-knots[j-1], POLYORDER-k);

    /* Satisfy the periodicity constraint: */
    ck[j][POLYORDER-1] = (ck[0][POLYORDER]-ck[j][POLYORDER])/(knots[0]-knots[j]+1.0);

    /* Apply both constraints: */
    for (i = 0; i < n[j]; i++)
        gsl_vector_set (y[j], i, gsl_vector_get (y[j], i) - ck[j][POLYORDER] - ck[j][POLYORDER-1] * gsl_vector_get (x[j], i));

    if (POLYORDER > 1) {
        for (k = 0; k < POLYORDER-1; k++)
            for (i = 0; i < n[j]; i++)
                gsl_matrix_set (A[j], i, k, gsl_matrix_get (A[j], i, k) - gsl_vector_get (x[j], i) * pow(knots[0]-knots[j]+1, POLYORDER-k-1));

        gsl_multifit_wlinear (A[j], w[j], y[j], c[j], cov[j], &chisq, mw[j]);

        for (k = 0; k < POLYORDER-1; k++)
            ck[j][k] = gsl_vector_get (c[j], k);
        chi2tot += chisq;
    }
    else {
        /*
         * If we are fitting linear polynomials, there's nothing to fit here
         * because both knots are constrained (connectivity and periodicity).
         * However, we still need to traverse data points to get chi2.
         */

        chisq = 0.0;
        for (i = 0; i < n[j]; i++)
            chisq += gsl_vector_get (w[j], i) * gsl_vector_get (y[j], i) * gsl_vector_get (y[j], i);
        chi2tot += chisq;
    }

    if (VERBOSE && !SUMMARY_OUTPUT && !EBAI_OUTPUT) {
        sprintf (buffer,"# % 12.12lf\t", knots[j]);
        messages = combine(messages ,buffer);
        for (k = 0; k < POLYORDER+1; k++){
            sprintf (buffer,"% 12.12lf\t", ck[j][POLYORDER-k]);
            messages = combine(messages ,buffer);
        }
        sprintf (buffer,"\n");
        messages = combine(messages ,buffer);
    }

    if (PRINT) {
        double phase;

        if (POLYORDER == 2) {
            double minval, maxval, xval, yval, xmax = 0.0;

            minval = 1e3;
            maxval = -1.0;

            for (j = 0; j < KNOTS; j++) {
                double yv1 = compute_flux (knots[j], knots, ck);
                double yv2 = compute_flux (j == KNOTS-1 ? knots[0] : knots[j+1], knots, ck);

                /* Concave parabola: check for minimum at end-points, find maximum. */
                if (ck[j][0] < 0) {
                    /* Check for minimum at end-points: */
                    if (yv1 < minval || yv2 < minval) {
                        pshift = yv1 < yv2 ? knots[j] : j == KNOTS-1 ? knots[0] : knots[j+1];
                        minval = yv1 < yv2 ? yv1 : yv2;
                    }

                    /* Find the maximum: */
                    xval = knots[j]-(ck[j][1]-(j == KNOTS-1 ? ck[j][0]*(knots[0]-knots[j]+1.0) : 0))/2.0/ck[j][0];

                    /* If the maximum is not enclosed in the interval, check end-points: */
                    if ((xval < knots[j] || xval > (j == KNOTS-1 ? knots[0]+1.0 : knots[j+1])) && (yv1 > maxval || yv2 > maxval)) {
                        xmax = yv1 > yv2 ? knots[j] : j == KNOTS-1 ? knots[0] : knots[j+1];
                        maxval = yv1 > yv2 ? yv1 : yv2;
                    }
                    /* If the maximum is enclosed in the interval: */
                    else {
                        yval = compute_flux (xval, knots, ck);
                        if (yval > maxval) {
                            xmax = xval;
                            maxval = yval;
                        }
                    }
                }
                /* Convex parabola: check for maximum at end-points, find minimum. */
                else {
                    if (yv1 > maxval || yv2 > maxval) {
                        xmax = yv1 > yv2 ? knots[j] : j == KNOTS-1 ? knots[0] : knots[j+1];
                        maxval = yv1 > yv2 ? yv1 : yv2;
                    }

                    /* Find the minimum: */
                    xval = knots[j]-(ck[j][1]-(j == KNOTS-1 ? ck[j][0]*(knots[0]-knots[j]+1.0) : 0))/2.0/ck[j][0];

                    /* If the minimum is not enclosed in the interval, check end-points: */
                    if ((xval < knots[j] || xval > (j == KNOTS-1 ? knots[0]+1.0 : knots[j+1])) && (yv1 < minval || yv2 < minval)) {
                        pshift = yv1 < yv2 ? knots[j] : j == KNOTS-1 ? knots[0] : knots[j+1];
                        minval = yv1 < yv2 ? yv1 : yv2;
                    }
                    /* If the minimum is enclosed in the interval: */
                    else {
                        yval = compute_flux (xval, knots, ck);
                        if (minval > yval) {
                            minval = yval;
                            pshift = xval;
                        }
                    }
                }
            }

            if (SUMMARY_OUTPUT)
                sprintf (buffer,"%lf\n", pshift);

            else if (!EBAI_OUTPUT) {
                sprintf (buffer,"# \n# Polyfit minimum occurs at: (%lf, %lf)\n", pshift, minval);
                messages = combine(messages ,buffer);
                sprintf (buffer,"# Polyfit maximum occurs at: (%lf, %lf)\n", xmax, maxval);
                messages = combine(messages ,buffer);
                sprintf (buffer,"# Polyfit amplitude is: %lf (%3.3lf%%)\n", maxval-minval, (maxval-minval)/maxval*100.0);
                messages = combine(messages ,buffer);
            }
        }

        if (!SUMMARY_OUTPUT && !EBAI_OUTPUT) {
            //sprintf (buffer,"# \n# Theoretical light curve:\n# \n#   Phase:\t   Flux:\n");
            for (i = 0; i <= NVERT; i++) {
                phase = -0.5 + (double) i / NVERT;
                if (APPLY_PSHIFT) phase += pshift;
                /*
                if (CSV_OUTPUT)
                    sprintf (buffer,"%lf,%12.12lf\n", phase, compute_flux (phase, knots, ck));
                else
                    sprintf (buffer,"  % lf\t% 12.12lf\n", phase, compute_flux (phase, knots, ck));
                */
                outphase[i]   = phase;
                outpolyfit[i] = compute_flux (phase, knots, ck);
            }
        }
        else if (EBAI_OUTPUT) {
            for (i = 0; i <= NVERT; i++) {
                phase = -0.5 + (double) i / NVERT;
                if (APPLY_PSHIFT) phase += pshift;
                sprintf (buffer,"% 12.12lf\n", compute_flux (phase, knots, ck));
            }
        }
    }

    /* Done! Wrap it up: */
    for (i = 0; i < KNOTS-1; i++) {
        gsl_vector_free (x[i]);
        gsl_vector_free (y[i]);
        gsl_vector_free (w[i]);
        free (ck[i]);
    }
    free (x);
    free (y);
    free (w);
    free (n);
    free (ck);

    for (i = 0; i < intervals; i++) {
        gsl_vector_free (c[i]);
        gsl_matrix_free (A[i]);
        gsl_matrix_free (cov[i]);
        gsl_multifit_linear_free (mw[i]);
    }
    free (c);
    free (A);
    free (cov);
    free (mw);

    return chi2tot;
}

double *tokenize(char *line, int *token_no)
{
    double *tokens = NULL;
    int LEN = 0;

    char *token;

    if (CSV_INPUT)
        token = strtok(line, ",");
    else
        token = strtok(line, " \t");

    if (!token)
        return NULL;

    while (token) {
        LEN += 1;
        tokens = realloc(tokens, LEN*sizeof(*tokens));
        tokens[LEN-1] = atof(token);
        if (CSV_INPUT)
            token = strtok(NULL, ",");
        else
            token = strtok(NULL, " \t");
    }

    *token_no = LEN;
    return tokens;
}

char * polyfitter (
  double* inphase,
  double* influx,
  double* inerror,
  int nobs,
  int order,
  int ITERS,
  double STEP_SIZE,
  int vertices,
  int nKNOTS,
  double* knots,
  int FIND_KNOTS,
  int FIND_STEP,
  int chainlength,
  int applypshift,
  double* outphase,
  double* outpolyfit
)
{
    int i, iter, status;
    triplet *data;
    //FILE *in;
    double chi2, chi2test, u;
    //char line[255];

    //double *knots = NULL;
    double knotsTest[4];
    double *test;
    //int colno, inputs = 0;
    //double *cols;

    gsl_rng *r;

    messages = "";

    CHAIN_LENGTH = chainlength;
    KNOTS = nKNOTS;
    POLYORDER = order;
    NVERT = vertices;
    APPLY_PSHIFT = applypshift;

    /********************   READ DATA:   **********************/
    /* Read the input data: */
    data = malloc (nobs*sizeof(*data));

    for(i=0;i < nobs; i++){
        data[i].x = inphase[i];
        data[i].y = influx[i];
        data[i].z = inerror[i];
    }

    /*
    if (argc < 2) {
        sprintf (buffer,"polyfit %s\n\n", VERSION);
        sprintf (buffer,"Usage: ./polyfit [options] lc.dat\n\n");
        sprintf (buffer,"File lc.dat should have at least 1 column (flux),\n");
        sprintf (buffer,"2 columns (phase and flux), or 3 columns (phase,\n");
        sprintf (buffer,"flux, and standard deviation)\n\n");
        sprintf (buffer,"Options:\n\n");
        sprintf (buffer,"  -o order          ..  fitting polynomial order (default: 2)\n");
        sprintf (buffer,"  -i iters          ..  number of iterations (default: 10000)\n");
        sprintf (buffer,"  -s step           ..  step for random knot displacement (default: 0.01)\n");
        sprintf (buffer,"  -k k1 k2 ... kN   ..  explicit list of knots\n");
        sprintf (buffer,"  -n vertices       ..  number of vertices in the computed fit\n");
        sprintf (buffer,"  -c c1 c2 [c3]     ..  input file columns to be parsed (starting from 0)\n");
        sprintf (buffer,"  --find-knots      ..  attempt to find knots automatically\n");
        sprintf (buffer,"  --find-step       ..  attempt to find step automatically\n");
        sprintf (buffer,"  --chain-length    ..  minimum chain length for automatic knot search\n");
        sprintf (buffer,"  --apply-pshift    ..  shift phase so that the polyfit minimum is at phase 0\n");
        sprintf (buffer,"  --csv-input       ..  take comma-separated-value file as input\n");
        sprintf (buffer,"  --csv-output      ..  print comma-separated-values at output\n");
        sprintf (buffer,"  --limits          ..  set lower/upper phase limit (detault: no filtering)\n");
        sprintf (buffer,"  --ebai-output     ..  write a one-column output in ebai-compatible format\n");
        sprintf (buffer,"  --summary-output  ..  output a single summary line from polyfit\n\n");
        exit (0);
    }

    for (i = 1; i < argc; i++) {
        if (strcmp (argv[i], "-o") == 0)
            POLYORDER = atoi (argv[++i]);
        if (strcmp (argv[i], "-i") == 0)
            ITERS = atoi (argv[++i]);
        if (strcmp (argv[i], "-s") == 0)
            STEP_SIZE = atof (argv[++i]);
        if (strcmp (argv[i], "-n") == 0)
            NVERT = atoi (argv[++i]);
        if (strcmp (argv[i], "-k") == 0) {
            double knot;
            i++;
            KNOTS = 0;
            while (sscanf (argv[i], "%lf", &knot) == 1) {
                if (i == argc-1) break;
                KNOTS++;
                knots = realloc (knots, KNOTS * sizeof (*knots));
                knots[KNOTS-1] = knot;
                i++;
            }
            i--;
        }
        if (strcmp (argv[i], "-c") == 0) {
            int val;

            i++;
            inputs = 0;
            while (sscanf (argv[i], "%d", &val) == 1) {
                if (i == argc-1) break;
                COLIDX[inputs] = val;
                inputs++;
                i++;
            }
            i--;
        }
        if (strcmp(argv[i], "--limits") == 0) {
            sscanf(argv[i+1], "%lf", &LLIMIT);
            sscanf(argv[i+2], "%lf", &ULIMIT);
            i += 2;
        }
        if (strcmp (argv[i], "--find-knots") == 0)
            FIND_KNOTS = 1;
        if (strcmp (argv[i], "--find-step") == 0)
            FIND_STEP = 1;
        if (strcmp (argv[i], "--chain-length") == 0)
            CHAIN_LENGTH = atoi (argv[++i]);
        if (strcmp (argv[i], "--csv-input") == 0)
            CSV_INPUT = 1;
        if (strcmp (argv[i], "--csv-output") == 0)
            CSV_OUTPUT = 1;
        if (strcmp (argv[i], "--summary-output") == 0)
            SUMMARY_OUTPUT = 1;
        if (strcmp (argv[i], "--ebai-output") == 0)
            EBAI_OUTPUT = 1;
        if (strcmp (argv[i], "--apply-pshift") == 0)
            APPLY_PSHIFT = 1;
    }

    if (SUMMARY_OUTPUT)
        sprintf (buffer,"%s\t", argv[argc-1]);

    in = fopen (argv[argc-1], "r");
    if (!in) {
        sprintf (buffer,"file %s not found, aborting.\n", argv[argc-1]);
        exit (0);
    }
    */

    /********************   READ DATA:   **********************/
    /* Read the input data: */
    /*
    data = malloc (CHUNK*sizeof(*data));
    i = 0;
    while (!feof (in)) {
        if (!fgets (line, 255, in)) break;
        if (line[0] == '\n') break;       // break on empty lines
        if (line[0] == '#') continue;     // ignore comments

        // Break up the line into tokens:
        cols = tokenize(line, &colno);

        // Guess the role of columns if not given by the -c switch:
        if (i == 0 && COLIDX[0] == -1 && COLIDX[1] == -1 && COLIDX[2] == -1) {
            if (colno == 1) {
                COLIDX[1] = 0;
            }
            else if (colno == 2) {
                COLIDX[0] = 0;
                COLIDX[1] = 1;
            }
            else {
                COLIDX[0] = 0;
                COLIDX[1] = 1;
                COLIDX[2] = 2;
            }
        }

        inputs = colno <= 3 ? colno : 3;

        if (colno <= COLIDX[0] || colno <= COLIDX[1] || colno <= COLIDX[2]) {
            sprintf(buffer,"The input file has %d columns, but column index %d is requested; aborting.\n", colno, COLIDX[2] <= colno ? COLIDX[2] : (COLIDX[1] <= colno ? COLIDX[1] : COLIDX[0]));
            exit(0);
        }

        // if (i == 0)
        //     sprintf(buffer,"[0]=%d [1]=%d [2]=%d\n", COLIDX[0], COLIDX[1], COLIDX[2]);

        // Check if the data point is within limits:
        if ((COLIDX[0] != -1) && (cols[COLIDX[0]] < LLIMIT || cols[COLIDX[0]] > ULIMIT))
            continue;

        if (inputs == 1) {
            data[i].y = cols[COLIDX[1]];
            data[i].z = 1e-3;
        }
        else if (inputs == 2) {
            data[i].x = cols[COLIDX[0]];
            data[i].y = cols[COLIDX[1]];
            data[i].z = 1e-3;
        }
        else {
            data[i].x = cols[COLIDX[0]];
            data[i].y = cols[COLIDX[1]];
            data[i].z = cols[COLIDX[2]];
        }

        free(cols);

        i++;
        if (i % CHUNK == 0)
            data = realloc(data, (i/CHUNK+1)*(CHUNK)*sizeof(*data));
    }
    fclose (in);
    nobs = i;
    */

    /*
    if (!SUMMARY_OUTPUT && !EBAI_OUTPUT)
        sprintf (buffer,"# %d data points read in from %s.\n", nobs, argv[argc-1]);

    // If only fluxes are provided, add equidistant nobs phases:
    if (inputs == 1) {
        if (!SUMMARY_OUTPUT && !EBAI_OUTPUT)
            sprintf(buffer,"# * no phases given, assuming equidistant sampling between -0.5 and 0.5.\n");
        for (i=0; i<nobs; i++) {
            data[i].x = -0.5 + (double) i/(nobs-1);
        }
    }
    */

    /* Sort data in phase: */
    qsort (data, nobs, sizeof (*data), sort_by_phase);

    /********************   FIND KNOTS:   **********************/
    if (FIND_KNOTS) {
        status = find_knots (data, nobs, &knots);
        if (!SUMMARY_OUTPUT && !EBAI_OUTPUT) {
            if (status != 0){
                sprintf (buffer,"# * automatic search for knots failed, reverting to defaults.\n");
                messages = combine(messages ,buffer);
            }
            else{
                sprintf (buffer,"# * automatic search for knots successful.\n");
                messages = combine(messages ,buffer);
            }
            sprintf (buffer,"# \n");
            messages = combine(messages ,buffer);
        }
    }

    /* The initial phase intervals for knots: */
    if (!knots) {
        KNOTS = 4;
        knots = malloc (KNOTS * sizeof (*knots));
        knots[0] = -0.4; knots[1] = -0.1; knots[2] = 0.1; knots[3] = 0.4;
    }

    /* Sort the knots in ascending order: */
    qsort (knots, KNOTS, sizeof (*knots), sort_by_value);

    /********************   FIND STEP:   **********************/
    if (FIND_STEP) {
        double diff;
        /* Step size would be the minimum width between two knots / 5: */
        diff = fabs (knots[1]-knots[0]);
        for (i = 1; i < KNOTS-1; i++)
            if (fabs (knots[i+1]-knots[i]) < diff)
                diff = fabs (knots[i+1]-knots[i]);
        if (fabs (knots[i+1]-knots[0]) < diff)
            diff = fabs (knots[i+1]-knots[0]);

        STEP_SIZE = diff / 5.0;
    }

    /****************************************************************/
    /*********    This part corrects wrong initial KNOTS       ******/
    /********* if KNOTS search fails use fixed inital values   ******/
    /****************************************************************/

    /********* Secondary is inside primary:   ************/
    /* If knots are all around primary minimum aka secondary is inside primary. -- narrow primary minimum */
    if ( ((knots[KNOTS-1]-knots[0])>0.001) && ((knots[KNOTS-1]-knots[0])<0.49) && ((knots[2]-knots[1])<0.015) ){
        KNOTS = 4;
        knots = malloc (KNOTS * sizeof (*knots));
        knots[0] = -0.4; knots[1] = -0.1; knots[2] = 0.1; knots[3] = 0.4;
    }
    /* If knots are all around primary minimum aka secondary is inside primary. -- wide primary minimum */
    if ( ((knots[KNOTS-1]-knots[0])>0.2) && ((knots[KNOTS-1]-knots[0])<0.49) && ((knots[2]-knots[1])<0.15) ){
        KNOTS = 4;
        knots = malloc (KNOTS * sizeof (*knots));
        knots[0] = -0.4; knots[1] = -0.1; knots[2] = 0.1; knots[3] = 0.4;
    }

    /********* If knots are too close:   ************/
    if ( (knots[KNOTS-1]-knots[0])<0.015 ){
        KNOTS = 4;
        knots = malloc (KNOTS * sizeof (*knots));
        knots[0] = -0.4; knots[1] = -0.1; knots[2] = 0.1; knots[3] = 0.4;
    }

    /********* Secondary is inside primary, which is around -0.5 – 0.5 ************/
    /* Check knots again but handle if minimum is at phase 0.5 (aka around -0.5 and 0.5) */
    for(i=0; i<KNOTS; i++){
        if (knots[i]<0) { knotsTest[i] = knots[i] + 0.5; }
        else {knotsTest[i] = knots[i] - 0.5;}
    }

    /* sort shifted knosts */
    qsort (knotsTest, KNOTS, sizeof (*knotsTest), sort_by_value);

    /* If knots are all around primary minimum aka secondary is inside primary. -- narrow primary minimum */
    if ( ((knotsTest[KNOTS-1]-knotsTest[0])>0.001) && ((knotsTest[KNOTS-1]-knotsTest[0])<0.49) && ((knotsTest[2]-knotsTest[1])<0.069) ){
        KNOTS = 4;
        knots = malloc (KNOTS * sizeof (*knotsTest));
        knots[0] = -0.4; knots[1] = -0.1; knots[2] = 0.1; knots[3] = 0.4;
    }
    /* If knots are all around primary minimum aka secondary is inside primary. -- wide primary minimum */
    if ( ((knotsTest[KNOTS-1]-knotsTest[0])>0.2) && ((knotsTest[KNOTS-1]-knotsTest[0])<0.49) && ((knotsTest[2]-knotsTest[1])<0.15) ){
        KNOTS = 4;
        knots = malloc (KNOTS * sizeof (*knots));
        knots[0] = -0.4; knots[1] = -0.1; knots[2] = 0.1; knots[3] = 0.4;
    }

    /********* Secondary minimum is a sharp peak:   ************/
    /* If secondary minimum is sharp, fix secondary's knots. -- sharp minimum is left to primary */
    if ( ((knots[1]-knots[0])<0.015) && ((knots[3]-knots[2])>0.05) ){
        knots[0] = -0.4; knots[1] = 0.4;
    }
    /* If secondary minimum is sharp, fix secondary's knots. -- sharp minimum is right to primary */
    if ( ((knots[1]-knots[0])>0.05) && ((knots[3]-knots[2])<0.015) ){
        knots[2] = -0.4; knots[3] = 0.4;
    }

    /********* Secondary minimum is a sharp peak, around phase -0.5 – 0.5  ************/
    /* Check knots again but handle if sharp minimum is at phase 0.5 (aka around -0.5 and 0.5) */
    for(i=0; i<KNOTS; i++){
        if (knots[i]<0.25) { knotsTest[i] = knots[i] + 0.25; }
        else { knotsTest[i] = knots[i] - 0.75; }
    }

    /* sort shifted knosts */
    qsort (knotsTest, KNOTS, sizeof (*knotsTest), sort_by_value);

    /* If secondary minimum is sharp, fix secondary's knots. -- sharp minimum is left to primary */
    if ( ((knotsTest[1]-knotsTest[0])<0.005) && ((knotsTest[3]-knotsTest[2])>0.1) ){
        knots[0] = -0.4; knots[1] = 0.4;
    }
    /* If secondary minimum is sharp, fix secondary's knots. -- sharp minimum is right to primary */
    if ( ((knotsTest[1]-knotsTest[0])>0.1) && ((knotsTest[3]-knotsTest[2])<0.005) ){
        knots[2] = -0.4; knots[3] = 0.4;
    }

    /* Sort the knots in ascending order: */
    qsort (knots, KNOTS, sizeof (*knots), sort_by_value);

    /********************   PRINT INFORMATION:   **********************/
    if (!SUMMARY_OUTPUT && !EBAI_OUTPUT) {
        sprintf (buffer,"# Fitting polynomial order: %d\n", POLYORDER);
        messages = combine( messages, buffer );
        sprintf (buffer,"# Initial set of knots: {");
        messages = combine( messages, buffer );
        for (i = 0; i < KNOTS-1; i++){
            sprintf (buffer,"%lf, ", knots[i]);
            messages = combine(messages ,buffer);
        }
        sprintf (buffer,"%lf}\n", knots[i]);
        messages = combine(messages ,buffer);
        sprintf (buffer,"# Number of iterations for knot search: %d\n", ITERS);
        messages = combine(messages ,buffer);
        sprintf (buffer,"# Step size for knot search: %lf\n# \n", STEP_SIZE);
        messages = combine(messages ,buffer);
    }

    /********* If primary minimum is a sharp peak near phase 0  ************/
    /* Sort absolute values -> if there is peak near 0 and 2nd peak far from 0 */
    for(i=0; i<KNOTS; i++){
        knotsTest[i] = fabs(knots[i]);
    }

    /* sort shifted knosts */
    qsort (knotsTest, KNOTS, sizeof (*knotsTest), sort_by_value);

    if ( (knotsTest[0] < 0.1) && ( (knotsTest[1]-knotsTest[0]) > 0.1) ){
        KNOTS = 4;
        knots = malloc (KNOTS * sizeof (*knotsTest));
        knots[0] = -0.4; knots[1] = -0.1; knots[2] = 0.1; knots[3] = 0.4;
    }

    /********* If both minima are very thin -> use 4th order fit  ************/
    /* If there is NO minimum at phase -0.5 - 0.5 */
    /*
    if ( ((knots[1]-knots[0])<0.06) && ((knots[3]-knots[2])<0.06) ){
        POLYORDER = 4;
    }
    // If one minimum is at phase -0.5 - 0.5
    for(i=0; i<KNOTS; i++){
        if (knots[i]<0.25) { knotsTest[i] = knots[i] + 0.25; }
        else { knotsTest[i] = knots[i] - 0.75; }
    }

    // sort shifted knosts
    qsort (knotsTest, KNOTS, sizeof (*knotsTest), sort_by_value);

    // If secondary minimum is sharp, fix secondary's knots. -- sharp minimum is left to primary
    if ( ((knotsTest[1]-knotsTest[0])<0.06) && ((knotsTest[3]-knotsTest[2])<0.06) ){
        POLYORDER = 4;
    }
    */

    /********************************************************************************/
    /********************   MCMC like search for best KNOTS:   **********************/
    /********************************************************************************/

    r = gsl_rng_alloc (gsl_rng_mt19937);
    gsl_rng_set (r, 1);

    chi2 = polyfit (data, nobs, KNOTS, knots, 0, 0, outphase, outpolyfit);
    if (!SUMMARY_OUTPUT && !EBAI_OUTPUT){
        sprintf (buffer,"# Original chi2: %e\n", chi2);
        messages = combine(messages ,buffer);
    }

    test = malloc (KNOTS * sizeof (*test));

    /*
    for (iter = 0; iter < ITERS; iter++) {
        for (i = 0; i < KNOTS; i++) {
            u = gsl_rng_uniform (r);
            test[i] = knots[i] + STEP_SIZE * 2 * u - STEP_SIZE;
            if (test[i] < -0.5) test[i] += 1.0;
            if (test[i] >  0.5) test[i] -= 1.0;
        }

        qsort (test, KNOTS, sizeof (*test), sort_by_value);

        chi2test = polyfit (data, nobs, KNOTS, test, 0, 0);
        if (chi2test < chi2) {
            chi2 = chi2test;
            for (i = 0; i < KNOTS; i++)
                knots[i] = test[i];
        }
    }
    */

    double chi2arr[200];
    double chi2med, chi2MAD;
    double work[200];
    double chi2stored;
    iter = 0;
    while (iter < ITERS) {
        for (i = 0; i < KNOTS; i++) {
            u = gsl_rng_uniform (r);
            test[i] = knots[i] + STEP_SIZE * 2 * u - STEP_SIZE;
            if (test[i] < -0.5) test[i] += 1.0;
            if (test[i] >  0.5) test[i] -= 1.0;
        }

        qsort (test, KNOTS, sizeof (*test), sort_by_value);

        chi2test = polyfit (data, nobs, KNOTS, test, 0, 0, outphase, outpolyfit);
        chi2stored = chi2;
        if (chi2test < chi2) {
            chi2 = chi2test;
            for (i = 0; i < KNOTS; i++)
                knots[i] = test[i];
        }
        if (iter<200){
            chi2arr[iter] = chi2test;
        }
        else{
            if (iter==200){
                chi2med = gsl_stats_median(chi2arr, 1, 200);
                chi2MAD = gsl_stats_mad(chi2arr, 1, 200, work);
                sprintf(buffer,"# Chi2med=%f chi2MAD=%f Threshold=%f\n",chi2med,chi2MAD,chi2med - 5*chi2MAD);
                messages = combine(messages ,buffer);
            }
            if ((chi2test < chi2stored) && (chi2test < chi2med - 5*chi2MAD)){
                sprintf(buffer,"# BREAKING @%d\n",iter);
                messages = combine(messages ,buffer);
                break;
            }
        }

        iter++;
    }

    if (!SUMMARY_OUTPUT && !EBAI_OUTPUT){
        sprintf (buffer,"# Final chi2:    %e\n", chi2);
        messages = combine(messages ,buffer);
    }

    if (SUMMARY_OUTPUT)
        sprintf (buffer,"%lf\t", chi2);

    chi2 = polyfit (data, nobs, KNOTS, knots, 1, 1, outphase, outpolyfit);

    gsl_rng_free (r);
    free (data);
    free (test);
    //free (knots);

    return messages;
    }

max_stars_repo_name: astrobatty/polyfitter

max_issues_repo_name: astrobatty/polyfitter

max_forks_repo_name: astrobatty/polyfitter


text:
#ifndef APPROX2_H
#define APPROX2_H

#include <string>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <stdio.h>      /* printf */
#include <time.h>
#include <fstream>
#include <algorithm>
#include <iomanip>
#include <ctime>
#include <math.h>

template<typename L, typename D>
class APPROX2
{

protected:

  // parameters

  D mu_f;

  D mu_psi;

  L n;   // x\in \R^n

  L tau;  //number of threads on each node/computer

  D sumofLi;

  // variables
  std::vector<D> u;

  std::vector<D> z;

  std::vector<D> x;

  std::vector<D> v;

  std::vector<D> t;

  D gamma;

  D theta;

  // sampling variables

  std::vector<D> proba_vector;

  std::vector<D> S;

  std::vector<D> all_n;

  std::vector<D> sampled;

  D max_p;

  // auxiliary variables

  D tauovern;
  D novertau;
  D novertau2;
  D novertau3;
  D novertau4;
  D taumuovern;

  // variables for print
  D primal_value;

  D dual_value;

  D gradient_norm;

  D subgradient;

  D epsilon;

  L max_nb_loops;

  L evaluation;

  D delta;   //primal dual gap

  D delta2; //distance to subgradient

  ofstream samp;

  L nb_iters;

  L nb_of_iters_per_loop;

  D running_time;

  L print_every_N;

  L mod;

public:

  gsl_rng * rng;
  virtual inline D partial_i_of_f(L i){return D(NULL);}
  virtual inline D compute_prox(D, D, D, L){return D(NULL);}
  virtual inline void compute_primal_value() {}
  virtual inline void compute_dual_value(){}
  virtual inline void compute_gradient_norm(){}
  virtual inline void set_v(){}
  virtual inline void set_p(){}
  virtual inline void update_z_coordinate( L, D){}
  virtual inline void update_u_coordinate( L, D){}
  virtual inline void update_x_coordinate( L, D){}
  APPROX2()
  {

  }

  void set_rng()
  {
    gsl_rng_env_setup();
    const gsl_rng_type * T;
    T = gsl_rng_default;
    rng = gsl_rng_alloc(T);
    gsl_rng_set(rng,time(NULL));
    //gsl_rng_set(rng, 27432042);

  }

// sample i with probability pi=proba_vector[i]
  L sampling()
  {
    //L i=(floor)(gsl_rng_uniform(rng)*n);
    L i=gsl_rng_uniform_int(rng, n);
    if(tau==1)
    {
      D y=gsl_rng_uniform(rng);
      while(y*max_p>proba_vector[i])
      {
        i=(floor)(gsl_rng_uniform(rng)*n);
        y=gsl_rng_uniform(rng);
      }
    }
    return i;
  }

// sample S
  void batch_sampling()
  {
    if(tau<n)
    {
      L i=sampling();
      for(L k=0;k<tau;k++)
      {
        while(sampled[i]==1)
        {
          i=sampling();
        }
        sampled[i]=1;
        S[k]=i;
      }
      for(L k=0;k<tau;k++)
      {
        sampled[S[k]]=0;
      }
    }
    else {
      S=all_n;
      //cout<<"s=all_n"<<endl;
    }
  }

  void compute_x()
  {
    for(L i=0;i<n;i++)
    x[i]=gamma*u[i]+z[i];
  }

  void update_theta()
  {
    if(mod==1){
      D theta2=theta*theta;
      D tmp=mu_f+mu_psi-theta2*novertau2-mu_psi*(theta+1)*novertau;
      D tmp2=theta2*novertau4+theta*mu_psi*novertau3;
      D tmp3=mu_f+mu_psi-novertau2*theta2-novertau*mu_psi*(theta+1);
      theta=0.5/novertau2*(sqrt(tmp*tmp+4*tmp2)+tmp3);
    }

  }

  void initialize(vector<D> & x0, L val_tau, D val_mu_f, D val_mu_psi, L eval, L p_N, L val_mod)
  {
    tau=val_tau;
    theta=tau/(n+0.0);
    tauovern=tau/(n+0.0);
    novertau=n/(tau+0.0);
    novertau2=novertau*novertau;
    novertau3=novertau2*novertau;
    novertau4=novertau2*novertau2;
    all_n.resize(n,0);
    for(L i=0;i<n;i++)
    all_n[i]=i;
    gamma=1;
    mu_f=val_mu_f;
    mu_psi=val_mu_psi;
    evaluation=eval;
    nb_of_iters_per_loop=floor(max(1.,n/(tau+0.0)));
    print_every_N=p_N;
    mod=val_mod;

    taumuovern=tau*mu_f/(n+0.0);
    delta=std::numeric_limits<double>::max();
    gradient_norm=std::numeric_limits<double>::max();

    u.clear();
    u.resize(n,0);
    z.clear();
    z.resize(n,0);
    x.clear();
    x.resize(n,0);
    for(L i=0;i<n;i++)
    {
      z[i]=x0[i];
      x[i]=x0[i];
    }
    sampled.clear();
    sampled.resize(n,0);
    S.clear();
    S.resize(tau,0);
    t.clear();
    t.resize(n,0);

    set_v();
    set_p();
    set_rng();

    if(evaluation==1)
    {
      compute_primal_value();
      compute_dual_value();
      cout<<"initial primal value="<<primal_value<<endl;
      cout<<"initial dual value="<<dual_value<<endl;
    }
    else if(evaluation==2)
    {
      compute_primal_value();
      compute_gradient_norm();
      cout<<"initial primal value="<<primal_value<<endl;
      cout<<"initial gradient norm="<<gradient_norm<<endl;
    }
    running_time=0;
    cout<<"finished APPROX initializing"<<endl;
  }

  void compute_and_record_result()
  {
    if(evaluation==1&&nb_iters%print_every_N==0)
    {
      compute_primal_value();
      compute_dual_value();
      compute_gradient_norm();
      delta=primal_value-dual_value;
      cout<<setprecision(9)<<(0.0+nb_iters)<<";  "<<running_time<<"; primal dual gap="<<delta<<" primal value="<<primal_value<<"; dual value="<<dual_value<<endl;
      samp<<(0.0+nb_iters)<<" "<<delta<<" "<<running_time<<" "<<primal_value<<" "<<dual_value<<" "<<gradient_norm<<endl;
      gradient_norm=std::numeric_limits<double>::max();
    }
    else if(evaluation==2&&nb_iters%print_every_N==0)
    {
      compute_primal_value();
      compute_gradient_norm();
      cout<<setprecision(9)<<(0.0+nb_iters)<<";  "<<running_time<<"; gradient norm="<<gradient_norm<<" primal value="<<primal_value<<endl;
      samp<<(0.0+nb_iters)<<" "<<gradient_norm<<" "<<running_time<<" "<<primal_value<<endl;
    }
    else if(evaluation==3&&nb_iters%print_every_N==0) //running APPROX for solving the subproblem in ALM
    {
      compute_primal_value();
      compute_dual_value();
      compute_gradient_norm();
      delta=primal_value-dual_value;
      cout<<setprecision(9)<<"      "<<(0.0+nb_iters)<<";  "<<running_time<<"; primal dual gap="<<delta<<" primal value="<<primal_value<<"; dual value="<<dual_value<<"; gradient_norm="<<gradient_norm<<endl;
      samp<<(0.0+nb_iters)<<" "<<delta<<" "<<running_time<<" "<<primal_value<<" "<<dual_value<<" "<<gradient_norm<<endl;
      gradient_norm=std::numeric_limits<double>::max();
    }

  }

   void compute_and_record_result_always(){
      if(evaluation==1)
    {
      compute_primal_value();
      compute_dual_value();
      compute_gradient_norm();
      delta=primal_value-dual_value;
      cout<<setprecision(9)<<(0.0+nb_iters)<<";  "<<running_time<<"; primal dual gap="<<delta<<" primal value="<<primal_value<<"; dual value="<<dual_value<<endl;
      samp<<(0.0+nb_iters)<<" "<<delta<<" "<<running_time<<" "<<primal_value<<" "<<dual_value<<" "<<gradient_norm<<endl;
      gradient_norm=std::numeric_limits<double>::max();
    }
    else if(evaluation==2)
    {
      compute_primal_value();
      compute_gradient_norm();
      cout<<setprecision(9)<<(0.0+nb_iters)<<";  "<<running_time<<"; gradient norm="<<gradient_norm<<" primal value="<<primal_value<<endl;
      samp<<(0.0+nb_iters)<<" "<<gradient_norm<<" "<<running_time<<" "<<primal_value<<endl;
    }
    else if(evaluation==3) //running APPROX for solving the subproblem in ALM
    {
      compute_primal_value();
      compute_dual_value();
      compute_gradient_norm();
      delta=primal_value-dual_value;
      cout<<setprecision(9)<<"      "<<(0.0+nb_iters)<<";  "<<running_time<<"; primal dual gap="<<delta<<" primal value="<<primal_value<<"; dual value="<<dual_value<<"; gradient_norm="<<gradient_norm<<endl;
      samp<<(0.0+nb_iters)<<" "<<delta<<" "<<running_time<<" "<<primal_value<<" "<<dual_value<<" "<<gradient_norm<<endl;
      gradient_norm=std::numeric_limits<double>::max();
    }

   }

  void APPROX_MU(vector<D> & x0, L val_tau, D val_mu_f, D val_mu_psi, L eval, L p_N,  L max_nb_epoch, D eps, string filename, L val_mod)
  {
    initialize(x0,  val_tau,  val_mu_f,  val_mu_psi,  eval, p_N, val_mod);
    cout<<"running APPROX MU"<<" ; "<<filename<<" max_nb_epoch "<<max_nb_epoch<<"; eps="<<eps<<endl;

    nb_iters=0;
    compute_and_record_result();
    //srand48(27432042);
    srand(time(NULL));
    D start;
    while(delta>eps && nb_iters<max_nb_epoch)
    {
      start = std::clock();
      D ti=0;
      D si=0;
      D gi=0;
      D Li=0;
      L i=0;
      for(L it=0;it<nb_of_iters_per_loop;it++)
      {
        gamma*=(1-theta)/(1-taumuovern);
        if(theta==1) gamma=1;
        batch_sampling();
        for(L it_S=0;it_S<tau;it_S++)
        {
          i=S[it_S];
          si=z[i]+taumuovern*gamma/theta*u[i];
          gi=partial_i_of_f(i);
          Li=v[i]*novertau*theta;
          t[i]=compute_prox(gi, Li, si,i);
        }
        for(L it_S=0;it_S<tau;it_S++)
        {
          i=S[it_S];
          ti=t[i];
          update_z_coordinate(i, ti+taumuovern*gamma/theta*u[i]);
          update_u_coordinate(i, (-1+novertau*theta)/gamma*ti-taumuovern/theta*u[i]);
        }
        update_theta();
      }
      running_time+=( std::clock() - start ) / (double) CLOCKS_PER_SEC;
      nb_iters++;
      compute_and_record_result();
    }

    compute_and_record_result_always();
  }

  void APPROX_MU2(vector<D> & x0, L val_tau, D val_mu_f, D val_mu_psi, L eval, L p_N,  L max_nb_epoch, D eps, string filename, L val_mod)
  {
    initialize(x0,  val_tau,  val_mu_f,  val_mu_psi,  eval, p_N, val_mod);
    cout<<"running APPROX MU"<<" ; "<<filename<<" max_nb_epoch "<<max_nb_epoch<<"; eps="<<eps<<endl;

    nb_iters=0;
    compute_and_record_result();
    //srand48(27432042);
    srand(time(NULL));
    D start;
    while(nb_iters<max_nb_epoch)
    {
      start = std::clock();
      D ti=0;
      D si=0;
      D gi=0;
      D Li=0;
      L i=0;
      for(L it=0;it<nb_of_iters_per_loop;it++)
      {
        gamma*=(1-theta)/(1-taumuovern);
        if(theta==1) gamma=1;
        batch_sampling();
        for(L it_S=0;it_S<tau;it_S++)
        {
          i=S[it_S];
          si=z[i]+taumuovern*gamma/theta*u[i];
          gi=partial_i_of_f(i);
          Li=v[i]*novertau*theta;
          t[i]=compute_prox(gi, Li, si,i);
        }
        for(L it_S=0;it_S<tau;it_S++)
        {
          i=S[it_S];
          ti=t[i];
          update_z_coordinate(i, ti+taumuovern*gamma/theta*u[i]);
          update_u_coordinate(i, (-1+novertau*theta)/gamma*ti-taumuovern/theta*u[i]);
        }
        update_theta();
      }
      running_time+=( std::clock() - start ) / (double) CLOCKS_PER_SEC;
      nb_iters++;
      compute_and_record_result();
    }

    compute_and_record_result_always();
  }

  void prox_grad_step(){
    D gi;
    for(L i=0;i<n;i++){
      gi=partial_i_of_f(i);
      t[i]=compute_prox(gi, sumofLi, x[i],i);
    }
    for(L i=0;i<n;i++){
      update_x_coordinate(i, t[i]);
    }

  }

  void do_single_step_prox(vector<D> & Tx){
    for(L i=0;i<n;i++){
      D gi=partial_i_of_f(i);
      Tx[i]=compute_prox(gi, sumofLi, x[i],i)+x[i];
    }
  }

};

#endif

ext: h


text:
/**
 *
 * @file core_clarfx_tbrd.c
 *
 *  PLASMA core_blas kernel
 *  PLASMA is a software package provided by Univ. of Tennessee,
 *  Univ. of California Berkeley and Univ. of Colorado Denver
 *
 * @version 2.6.0
 * @author Azzam Haidar
 * @date 2011-05-15
 * @generated c Tue Jan  7 11:44:50 2014
 *
 **/
#include <lapacke.h>
#include "common.h"

/***************************************************************************//**
 *
 * @ingroup CORE_PLASMA_Complex32_t
 *
 *  CORE_clarfx2 applies a complex elementary reflector H to a complex m by n
 *  matrix C, from either the left or the right. H is represented in the
 *  form
 *
 *        H = I - tau * v * v'
 *
 *  where tau is a complex scalar and v is a complex vector.
 *
 *  If tau = 0, then H is taken to be the unit matrix
 *
 *  This version uses inline code if H has order < 11.
 *
 *******************************************************************************
 *
 * @param[in] side
 *          @arg PlasmaLeft : form  H * C
 *          @arg PlasmaRight: form  C * H
 *
 * @param[in] N
 *          The number of columns of C1 and C2 if side = PlasmaLeft.
 *          The number of rows of C1 and C2 if side = PlasmaRight.
 *
 * @param[in] V
 *          The float complex V in the representation of H.
 *
 * @param[in] TAU
 *          The value tau in the representation of H.
 *
 * @param[in,out] C1
 *          dimension (LDC1,N), if side = PlasmaLeft
 *          dimension (LDC1,1), if side = PlasmaRight
 *          On entry, the m by n matrix C1.
 *          On exit, C1 is overwritten by the matrix H * C1 if SIDE = PlasmaLeft,
 *          or C1 * H if SIDE = PlasmaRight.
 *
 * @param[in] LDC1
 *          The leading dimension of the array C1.
 *          LDC1 >= max(1,N), if side == PlasmaRight.
 *          LDC1 >= 1, otherwise.
 *
 * @param[in,out] C2
 *          dimension (LDC2,N), if side = PlasmaLeft
 *          dimension (LDC2,1), if side = PlasmaRight
 *          On entry, the m by n matrix C2.
 *          On exit, C2 is overwritten by the matrix H * C2 if SIDE = PlasmaLeft,
 *          or C2 * H if SIDE = PlasmaRight.
 *
 * @param[in] LDC2
 *          The leading dimension of the array C2.
 *          LDC2 >= max(1,N), if side == PlasmaRight.
 *          LDC2 >= 1, otherwise.
 *
 *******************************************************************************
 *
 * @return
 *          \retval PLASMA_SUCCESS successful exit
 *          \retval <0 if -i, the i-th argument had an illegal value
 *
 ******************************************************************************/
int
CORE_clarfx2(PLASMA_enum side, int N,
             PLASMA_Complex32_t V,
             PLASMA_Complex32_t TAU,
             PLASMA_Complex32_t *C1, int LDC1,
             PLASMA_Complex32_t *C2, int LDC2)
{
    PLASMA_Complex32_t V2, T2, SUM;
    int j;

    if (TAU == (PLASMA_Complex32_t)0.0)
        return PLASMA_SUCCESS;

    /*
     * Special code for 2 x 2 Householder where V1 = I
    */
    if(side==PlasmaLeft){
        V2 = conjf(V);
        T2 = TAU*conjf(V2);
        for (j = 0; j < N ; j++, C1+=LDC1 ) {
            SUM = *C1 + V2 * (*C2);
            *C1 = *C1 - SUM*TAU;
            *C2 = *C2 - SUM*T2;
            C2 += LDC2;
        }
    }
    else {
        V2 = V;
        T2 = TAU*conjf(V2);
        for (j = 0; j < N ; j++, C1++){
            SUM = *C1 + V2 * (*C2);
            *C1 = *C1 - SUM*TAU;
            *C2 = *C2 - SUM*T2;
            C2++;
        }
    }

    return PLASMA_SUCCESS;
}

/***************************************************************************//**
 *
 * @ingroup CORE_PLASMA_Complex32_t
 *
 *  CORE_clarfx2c applies a complex elementary reflector H to a diagonal corner
 *  C=[C1, C2, C3], from both the left and the right side. C = H * C * H.
 *  It is used in the case of Hermetian.
 *  If PlasmaLower, a left apply is followed by a right apply.
 *  If PlasmaUpper, a right apply is followed by a left apply.
 *  H is represented in the form
 *
 *  This routine is a special code for a corner C diagonal block  C1
 *                                                                C2  C3
 *
 *
 *        H = I - tau * v * v'
 *
 *  where tau is a complex scalar and v is a complex vector.
 *
 *  If tau = 0, then H is taken to be the unit matrix
 *
 *  This version uses inline code if H has order < 11.
 *
 *******************************************************************************
 *
 * @param[in] uplo
 *          = PlasmaUpper: Upper triangle of A is stored;
 *          = PlasmaLower: Lower triangle of A is stored.
 *
 * @param[in] V
 *          The float complex V in the representation of H.
 *
 * @param[in] TAU
 *          The value tau in the representation of H.
 *
 * @param[in,out] C1
 *          On entry, the element C1.
 *          On exit, C1 is overwritten by the result H * C * H.
 *
 * @param[in,out] C2
 *          On entry, the element C2.
 *          On exit, C2 is overwritten by the result H * C * H.
  *
 * @param[in,out] C3
 *          On entry, the element C3.
 *          On exit, C3 is overwritten by the result H * C * H.
 *
 *******************************************************************************
 *
 * @return
 *          \retval PLASMA_SUCCESS successful exit
 *          \retval <0 if -i, the i-th argument had an illegal value
 *
 ******************************************************************************/
int
CORE_clarfx2c(PLASMA_enum uplo,
              PLASMA_Complex32_t V,
              PLASMA_Complex32_t TAU,
              PLASMA_Complex32_t *C1,
              PLASMA_Complex32_t *C2,
              PLASMA_Complex32_t *C3)
{
    PLASMA_Complex32_t T2, SUM, TEMP;

    /* Quick return */
    if (TAU == (PLASMA_Complex32_t)0.0)
        return PLASMA_SUCCESS;

    /*
     *        Special code for a diagonal block  C1
     *                                           C2  C3
     */
    if(uplo==PlasmaLower) {
        /*
         *  Do the corner Left then Right (used for the lower case
         *  tridiag) L and R for the 2x2 corner
         *             C(N-1, N-1)  C(N-1,N)        C1  TEMP
         *             C(N  , N-1)  C(N  ,N)        C2  C3
         *  For Left : use conjf(TAU) and V.
         *  For Right: nothing, keep TAU and V.
         *  Left 1 ==> C1
         *             C2
         */
        TEMP = conjf(*C2); /*  copy C2 here before modifying it. */
        T2   = conjf(TAU) * V;
        SUM  = *C1 + conjf(V) * (*C2);
        *C1  = *C1 - SUM * conjf(TAU);
        *C2  = *C2 - SUM * T2;
        /*  Left 2 ==> TEMP */
        /*             C3 */
        SUM  = TEMP + conjf(V) * (*C3);
        TEMP = TEMP - SUM * conjf(TAU);
        *C3  = *C3  - SUM * T2;
        /*  Right 1 ==>  C1 TEMP.  NB: no need to compute corner (2,2)=TEMP */
        T2  = TAU * conjf(V);
        SUM = *C1 + V*TEMP;
        *C1 = *C1 - SUM*TAU;
        /*  Right 2 ==> C2 C3 */
        SUM = *C2   + V*(*C3);
        *C2 = *C2   - SUM*TAU;
        *C3 = *C3   - SUM*T2;
    }
    else {
        /*
         * Do the corner Right then Left (used for the upper case tridiag)
         *             C(N-1, N-1)  C(N-1,N)        C1    C2
         *             C(N  , N-1)  C(N  ,N)        TEMP  C3
         *  For Left : use TAU       and conjf(V).
         *  For Right: use conjf(TAU) and conjf(V).
         *  Right 1 ==> C1 C2
         */
        V    = conjf(V);
        TEMP = conjf(*C2); /*  copy C2 here before modifying it. */
        T2   = conjf(TAU) * conjf(V);
        SUM  = *C1 + V   * (*C2);
        *C1  = *C1 - SUM * conjf(TAU);
        *C2  = *C2 - SUM * T2;
        /*  Right 2 ==> TEMP C3 */
        SUM  = TEMP + V   * (*C3);
        TEMP = TEMP - SUM * conjf(TAU);
        *C3  = *C3  - SUM * T2;
        /*  Left 1 ==> C1 */
        /*             TEMP. NB: no need to compute corner (2,1)=TEMP */
        T2  = TAU * V;
        SUM = *C1 + conjf(V) * TEMP;
        *C1 = *C1 - SUM * TAU;
        /*  Left 2 ==> C2 */
        /*             C3 */
        SUM = *C2 + conjf(V) * (*C3);
        *C2 = *C2 - SUM * TAU;
        *C3 = *C3 - SUM * T2;
    }

    return PLASMA_SUCCESS;
}

/***************************************************************************//**
 *
 * @ingroup CORE_PLASMA_Complex32_t
 *
 *  CORE_clarfx2c applies a complex elementary reflector H to a diagonal corner
 *  C=[C1, C2, C3], from both the left and the right side. C = H * C * H.
 *  It is used in the case of general matrices, where it create a nnz at the
 *  NEW_NNZ position, then it eliminate it and update the reflector V and TAU.
 *  If PlasmaLower, a left apply is followed by a right apply.
 *  If PlasmaUpper, a right apply is followed by a left apply.
 *  H is represented in the form
 *
 *  This routine is a special code for a corner C diagonal block  C1  NEW_NNZ
 *                                                                C2  C3
 *
 *
 *        H = I - tau * v * v'
 *
 *  where tau is a complex scalar and v is a complex vector.
 *
 *  If tau = 0, then H is taken to be the unit matrix
 *
 *  This version uses inline code if H has order < 11.
 *
 *******************************************************************************
 *
 * @param[in] uplo
 *          = PlasmaUpper: Upper triangle of A is stored;
 *          = PlasmaLower: Lower triangle of A is stored.
 *
 * @param[in, out] V
 *          On entry, the float complex V in the representation of H.
 *          On exit, the float complex V in the representation of H,
 *          updated by the elimination of the NEW_NNZ created by the
 *          left apply in case of PlasmaLower or the right apply in
 *          case of PlasmaUpper.
 *
 * @param[in,out] TAU
 *          On entry, the value tau in the representation of H.
 *          On exit, the value tau in the representation of H,
 *          updated by the elimination of the NEW_NNZ created by the
 *          left apply in case of PlasmaLower or the right apply in
 *          case of PlasmaUpper.
 *
 * @param[in,out] C1
 *          On entry, the element C1.
 *          On exit, C1 is overwritten by the result H * C * H.
 *
 * @param[in,out] C2
 *          On entry, the element C2.
 *          On exit, C2 is overwritten by the result H * C * H.
  *
 * @param[in,out] C3
 *          On entry, the element C3.
 *          On exit, C3 is overwritten by the result H * C * H.
 *
 *******************************************************************************
 *
 * @return
 *          \retval PLASMA_SUCCESS successful exit
 *          \retval <0 if -i, the i-th argument had an illegal value
 *
 ******************************************************************************/
int
CORE_clarfx2ce(PLASMA_enum uplo,
               PLASMA_Complex32_t *V,
               PLASMA_Complex32_t *TAU,
               PLASMA_Complex32_t *C1,
               PLASMA_Complex32_t *C2,
               PLASMA_Complex32_t *C3)
{
    PLASMA_Complex32_t T2, SUM, TEMP, VIN, TAUIN;

    /* Quick return */
    if (*TAU == (PLASMA_Complex32_t)0.0)
        return PLASMA_SUCCESS;

    /*
     *        Special code for a diagonal block  C1
     *                                           C2  C3
     */
    if(uplo==PlasmaLower){
        /*
         *  Do the corner for the lower case BIDIAG ==> Left then will
         *  create a new nnz. eliminate it and modify V TAU and then
         *  Right L and R for the 2x2 corner
         *             C(N-1, N-1)  C(N-1,N)        C1  TEMP
         *             C(N  , N-1)  C(N  ,N)        C2  C3
         */
        VIN   = *V;
        TAUIN = conjf(*TAU);
        /*  Left 1 ==> C1 */
        /*             C2 */
        VIN  = conjf(VIN);
        T2   = TAUIN * conjf(VIN);
        SUM  = *C1   + VIN*(*C2);
        *C1  = *C1   - SUM*TAUIN;
        *C2  = *C2   - SUM*T2;
        /*   new nnz at TEMP and update C3 */
        SUM  =        VIN * (*C3);
        TEMP =      - SUM * TAUIN;
        *C3  = *C3  - SUM * T2;
        /*  generate Householder to annihilate the nonzero created at TEMP */
        *V    = TEMP;
        LAPACKE_clarfg_work( 2, C1, V, 1, TAU);
        VIN   = conjf(*V);
        TAUIN = conjf(*TAU);
        /*  Right 1 ==> C2 C3 */
        /* VIN     = VIN */
        T2  = TAUIN * conjf(VIN);
        SUM = *C2   + VIN*(*C3);
        *C2 = *C2   - SUM*TAUIN;
        *C3 = *C3   - SUM*T2;
    }else if(uplo==PlasmaUpper){
        /*
         * Do the corner for the upper case BIDIAG ==> Right then will
        *  create a new nnz. eliminate it and modify V TAU and then
        *  Left
        *             C(N-1, N-1)  C(N-1,N)        C1    C2
        *             C(N  , N-1)  C(N  ,N)        TEMP  C3
        *  For Left : use conjf(TAU) and V.
        *  For Right: use conjf(TAU) and conjf(V) as input.
        */
        VIN   = conjf(*V);
        TAUIN = conjf(*TAU);
        /*  Right 1 ==> C1 C2 */
        /* VIN     = VIN */
        T2   = TAUIN*conjf(VIN);
        SUM  = *C1  + VIN*(*C2);
        *C1  = *C1  - SUM*TAUIN;
        *C2  = *C2  - SUM*T2;
        /*   new nnz at TEMP and update C3 */
        SUM  =        VIN * (*C3);
        TEMP =      - SUM * TAUIN;
        *C3  = *C3  - SUM * T2;
        /*  generate Householder to annihilate the nonzero created at TEMP */
        *V    = TEMP;
        LAPACKE_clarfg_work( 2, C1, V, 1, TAU);
        VIN   = *V;
        TAUIN = conjf(*TAU);
        /*  apply from the Left using the NEW V TAU to the remaining 2 elements [C2 C3] */
        /*  Left 2 ==> C2 */
        /*             C3 */
        VIN = conjf(VIN);
        T2  = TAUIN*conjf(VIN);
        SUM = *C2 + VIN*(*C3);
        *C2 = *C2 - SUM*TAUIN;
        *C3 = *C3 - SUM*T2;
    }
    return PLASMA_SUCCESS;
}

ext: c


text:
/* linalg/rqrc.c
 *
 * Copyright (C) 2020 Patrick Alken
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <stdlib.h>
#include <string.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>

/*
 * this module contains routines for the QR factorization of a matrix
 * using the recursive Level 3 BLAS algorithm of Elmroth and Gustavson with
 * additional modifications courtesy of Julien Langou.
 */

static int unpack_Q1(gsl_matrix_complex * Q);
static int unpack_Q2(const gsl_matrix_complex * QR, const gsl_matrix_complex * T, gsl_matrix_complex * Q);
static int aux_ULH(const gsl_matrix_complex * L, gsl_matrix_complex * U);
static int aux_mLU(gsl_matrix_complex * A);
static int aux_ApUBH(const gsl_matrix_complex * U, const gsl_matrix_complex * B, gsl_matrix_complex * A);

/*
gsl_linalg_complex_QR_decomp_r()
  QR decomposition using Level 3 BLAS recursive algorithm of:

Elmroth, E. and Gustavson, F.G., 2000. Applying recursion to serial and parallel
  QR factorization leads to better performance. IBM Journal of Research and Development,
  44(4), pp.605-624.

Inputs: A - matrix to be factored, M-by-N with M >= N
        T - N-by-N upper triangular factor of block reflector

Return: success/error

Notes:
1) on output, upper triangle of A contains R; elements below the diagonal
are columns of V. The matrix Q is

Q = I - V T V^H

where T is upper triangular. Note that diag(T) = tau
*/

int
gsl_linalg_complex_QR_decomp_r (gsl_matrix_complex * A, gsl_matrix_complex * T)
{
  const size_t M = A->size1;
  const size_t N = A->size2;

  if (M < N)
    {
      GSL_ERROR ("M must be >= N", GSL_EBADLEN);
    }
  else if (T->size1 != T->size2)
    {
      GSL_ERROR ("T matrix must be square", GSL_ENOTSQR);
    }
  else if (T->size1 != N)
    {
      GSL_ERROR ("T matrix does not match dimensions of A", GSL_EBADLEN);
    }
  else if (N == 1)
    {
      /* base case, compute householder transform for single column matrix */

      gsl_complex * T00 = gsl_matrix_complex_ptr(T, 0, 0);
      gsl_vector_complex_view v = gsl_matrix_complex_column(A, 0);

      *T00 = gsl_linalg_complex_householder_transform(&v.vector);
      return GSL_SUCCESS;
    }
  else
    {
      /*
       * partition matrices:
       *
       *       N1  N2              N1  N2
       * N1 [ A11 A12 ] and  N1 [ T11 T12 ]
       * M2 [ A21 A22 ]      N2 [  0  T22 ]
       */
      int status;
      const size_t N1 = N / 2;
      const size_t N2 = N - N1;
      const size_t M2 = M - N1;

      gsl_matrix_complex_view A11 = gsl_matrix_complex_submatrix(A, 0, 0, N1, N1);
      gsl_matrix_complex_view A12 = gsl_matrix_complex_submatrix(A, 0, N1, N1, N2);
      gsl_matrix_complex_view A21 = gsl_matrix_complex_submatrix(A, N1, 0, M2, N1);
      gsl_matrix_complex_view A22 = gsl_matrix_complex_submatrix(A, N1, N1, M2, N2);

      gsl_matrix_complex_view T11 = gsl_matrix_complex_submatrix(T, 0, 0, N1, N1);
      gsl_matrix_complex_view T12 = gsl_matrix_complex_submatrix(T, 0, N1, N1, N2);
      gsl_matrix_complex_view T22 = gsl_matrix_complex_submatrix(T, N1, N1, N2, N2);

      gsl_matrix_complex_view m;

      /*
       * Eq. 2: recursively factor
       *
       * [ A11 ] = Q1 [ R11 ]
       * [ A21 ]      [  0  ]
       *                                          N1
       * Note: Q1 = I - V1 T11 V1^H, where V1 = [ V11 ] N1
       *                                        [ V21 ] M2
       */
      m = gsl_matrix_complex_submatrix(A, 0, 0, M, N1);
      status = gsl_linalg_complex_QR_decomp_r(&m.matrix, &T11.matrix);
      if (status)
        return status;

      /*
       * Eq. 3:
       *
       * [ R12 ] := Q1^H [ A12 ] = [ A12 ] - [ V11 W ]
       * [ A22 ]         [ A22 ]   [ A22 ]   [ V21 W ]
       *
       * where W = T11^H (V11^H A12 + V21^H A22), and using T12 as temporary storage
       */
      gsl_matrix_complex_memcpy(&T12.matrix, &A12.matrix);                                            /* W := A12 */
      gsl_blas_ztrmm(CblasLeft, CblasLower, CblasConjTrans, CblasUnit, GSL_COMPLEX_ONE,
                     &A11.matrix, &T12.matrix);                                                       /* W := V11^H * A12 */
      gsl_blas_zgemm(CblasConjTrans, CblasNoTrans, GSL_COMPLEX_ONE, &A21.matrix, &A22.matrix,
                     GSL_COMPLEX_ONE, &T12.matrix);                                                   /* W := W + V21^H * A22 */
      gsl_blas_ztrmm(CblasLeft, CblasUpper, CblasConjTrans, CblasNonUnit, GSL_COMPLEX_ONE,
                     &T11.matrix, &T12.matrix);                                                       /* W := T11^H * W */
      gsl_blas_zgemm(CblasNoTrans, CblasNoTrans, GSL_COMPLEX_NEGONE, &A21.matrix, &T12.matrix,
                     GSL_COMPLEX_ONE, &A22.matrix);                                                   /* A22 = A22 - V21 * W */
      gsl_blas_ztrmm(CblasLeft, CblasLower, CblasNoTrans, CblasUnit, GSL_COMPLEX_ONE, &A11.matrix,
                     &T12.matrix);                                                                    /* tmp = V11 * W */
      gsl_matrix_complex_sub(&A12.matrix, &T12.matrix);                                               /* R12 := A12 - V11 * W */

      /*
       * Eq. 4: recursively factor
       *
       * A22 = Q2~ R22
       *
       *              N1 M2
       * Note: Q2 = [ I   0  ] N1
       *            [ 0  Q2~ ] M2
       */
      status = gsl_linalg_complex_QR_decomp_r(&A22.matrix, &T22.matrix);
      if (status)
        return status;

      /*
       * Eq. 13: update T12 := -T11 * V1^H * V2 * T22
       *
       * where:
       *
       *        N1                N2
       * V1 = [ V11 ] N1   V2 = [  0  ] N1
       *      [ V21 ] N2        [ V22 ] N2
       *      [ V31 ] M-N       [ V32 ] M-N
       *
       * Note: V1^H V2 = V21^H V22 + V31^H V32
       * Also, V11, V22 are unit lower triangular
       */

      m = gsl_matrix_complex_submatrix(&A21.matrix, 0, 0, N2, N1);                                 /* V21 */
      gsl_matrix_complex_conjtrans_memcpy(&T12.matrix, &m.matrix);                                 /* T12 := V21^H */

      m = gsl_matrix_complex_submatrix(A, N1, N1, N2, N2);                                         /* V22 */
      gsl_blas_ztrmm(CblasRight, CblasLower, CblasNoTrans, CblasUnit, GSL_COMPLEX_ONE,
                     &m.matrix, &T12.matrix);                                                      /* T12 := V21^H * V22 */

      if (M > N)
        {
          gsl_matrix_complex_view V31 = gsl_matrix_complex_submatrix(A, N, 0, M - N, N1);
          gsl_matrix_complex_view V32 = gsl_matrix_complex_submatrix(A, N, N1, M - N, N2);

          gsl_blas_zgemm(CblasConjTrans, CblasNoTrans, GSL_COMPLEX_ONE, &V31.matrix, &V32.matrix,
                         GSL_COMPLEX_ONE, &T12.matrix);                                            /* T12 := T12 + V31^T * V32 */
        }

      gsl_blas_ztrmm(CblasLeft, CblasUpper, CblasNoTrans, CblasNonUnit, GSL_COMPLEX_NEGONE,
                     &T11.matrix, &T12.matrix);                                                    /* T12 := -T11 * T12 */
      gsl_blas_ztrmm(CblasRight, CblasUpper, CblasNoTrans, CblasNonUnit, GSL_COMPLEX_ONE,
                     &T22.matrix, &T12.matrix);                                                    /* T12 := T12 * T22 */

      return GSL_SUCCESS;
    }
}

/* Solves the square system A x = b for x using the QR factorisation,
 *
 *  R x = Q^H b
 *
 * where Q = I - V T V^H
 */

int
gsl_linalg_complex_QR_solve_r (const gsl_matrix_complex * QR, const gsl_matrix_complex * T,
                               const gsl_vector_complex * b, gsl_vector_complex * x)
{
  const size_t N = QR->size2;

  if (QR->size1 != N)
    {
      GSL_ERROR ("QR matrix must be square", GSL_ENOTSQR);
    }
  else if (T->size1 != QR->size1 || T->size2 != QR->size2)
    {
      GSL_ERROR ("T matrix must be N-by-N", GSL_EBADLEN);
    }
  else if (N != b->size)
    {
      GSL_ERROR ("matrix size must match b size", GSL_EBADLEN);
    }
  else if (N != x->size)
    {
      GSL_ERROR ("matrix size must match solution size", GSL_EBADLEN);
    }
  else
    {
      size_t i;

      /* compute Q^H b = [I - V T^H V^H] b */

      /* x := V^H b */
      gsl_vector_complex_memcpy(x, b);
      gsl_blas_ztrmv(CblasLower, CblasConjTrans, CblasUnit, QR, x);

      /* x = T^H * x */
      gsl_blas_ztrmv(CblasUpper, CblasConjTrans, CblasNonUnit, T, x);

      /* x = V * x */
      gsl_blas_ztrmv(CblasLower, CblasNoTrans, CblasUnit, QR, x);

      /* x = b - V * x */
      for (i = 0; i < N; ++i)
        {
          gsl_complex * xi = gsl_vector_complex_ptr(x, i);
          gsl_complex bi = gsl_vector_complex_get(b, i);
          GSL_REAL(*xi) = GSL_REAL(bi) - GSL_REAL(*xi);
          GSL_IMAG(*xi) = GSL_IMAG(bi) - GSL_IMAG(*xi);
        }

      /* Solve R x = Q^H b, storing x in-place */
      gsl_blas_ztrsv (CblasUpper, CblasNoTrans, CblasNonUnit, QR, x);

      return GSL_SUCCESS;
    }
}

/* Find the least squares solution to the overdetermined system
 *
 *   A x = b
 *
 * for M >= N using the QR factorization A = Q R.
 *
 * Inputs: QR   - [R; V] matrix, M-by-N
 *         T    - upper triangular block reflector, N-by-N
 *         b    - right hand side, size M
 *         x    - (output) solution, size M
 *                x(1:N) = least squares solution vector
 *                x(N+1:M) = vector whose norm equals ||b - Ax||
 *         work - workspace, size N
 */

int
gsl_linalg_complex_QR_lssolve_r (const gsl_matrix_complex * QR, const gsl_matrix_complex * T,
                                 const gsl_vector_complex * b, gsl_vector_complex * x, gsl_vector_complex * work)
{
  const size_t M = QR->size1;
  const size_t N = QR->size2;

  if (M < N)
    {
      GSL_ERROR ("QR matrix must have M >= N", GSL_EBADLEN);
    }
  else if (T->size1 != N || T->size2 != N)
    {
      GSL_ERROR ("T matrix must be N-by-N", GSL_EBADLEN);
    }
  else if (M != b->size)
    {
      GSL_ERROR ("matrix size must match b size", GSL_EBADLEN);
    }
  else if (M != x->size)
    {
      GSL_ERROR ("matrix size must match solution size", GSL_EBADLEN);
    }
  else if (N != work->size)
    {
      GSL_ERROR ("matrix size must match work size", GSL_EBADLEN);
    }
  else
    {
      gsl_matrix_complex_const_view R = gsl_matrix_complex_const_submatrix (QR, 0, 0, N, N);
      gsl_vector_complex_view x1 = gsl_vector_complex_subvector(x, 0, N);

      /* compute x = Q^H b */
      gsl_vector_complex_memcpy(x, b);
      gsl_linalg_complex_QR_QHvec_r (QR, T, x, work);

      /* Solve R x = Q^H b */
      gsl_blas_ztrsv (CblasUpper, CblasNoTrans, CblasNonUnit, &R.matrix, &x1.vector);

      return GSL_SUCCESS;
    }
}

/*
gsl_linalg_complex_QR_unpack_r()
  Unpack matrices Q and R

Inputs: QR - packed QR format, M-by-N
        T  - block reflector matrix, N-by-N
        Q  - (output) Q matrix, M-by-M
        R  - (output) R matrix, N-by-N

Return: success/error

Notes:
1) Implementation provided by Julien Langou
*/

int
gsl_linalg_complex_QR_unpack_r(const gsl_matrix_complex * QR, const gsl_matrix_complex * T,
                               gsl_matrix_complex * Q, gsl_matrix_complex * R)
{
  const size_t M = QR->size1;
  const size_t N = QR->size2;

  if (M < N)
    {
      GSL_ERROR ("M must be >= N", GSL_EBADLEN);
    }
  else if (Q->size1 != M || Q->size2 != M)
    {
      GSL_ERROR ("Q matrix must be M-by-M", GSL_EBADLEN);
    }
  else if (R->size1 != N || R->size2 != N)
    {
      GSL_ERROR ("R matrix must be N-by-N", GSL_EBADLEN);
    }
  else if (T->size1 != N || T->size2 != N)
    {
      GSL_ERROR ("T matrix must be N-by-N", GSL_EBADLEN);
    }
  else
    {
      gsl_matrix_complex_const_view RV = gsl_matrix_complex_const_submatrix(QR, 0, 0, N, N);
      gsl_matrix_complex_view Q1 = gsl_matrix_complex_submatrix(Q, 0, 0, M, N);
      gsl_matrix_complex_view m;

      /*
       * set Q1 = [ T ]
       *          [ V ]
       */
      m = gsl_matrix_complex_submatrix(Q, 0, 0, N, N);
      gsl_matrix_complex_tricpy(CblasUpper, CblasNonUnit, &m.matrix, T);
      gsl_matrix_complex_tricpy(CblasLower, CblasUnit, &m.matrix, &RV.matrix);

      if (M > N)
        {
          gsl_matrix_complex_const_view tmp = gsl_matrix_complex_const_submatrix(QR, N, 0, M - N, N);
          m = gsl_matrix_complex_submatrix(Q, N, 0, M - N, N);
          gsl_matrix_complex_memcpy(&m.matrix, &tmp.matrix);
        }

      unpack_Q1(&Q1.matrix);

      if (M > N)
        {
          gsl_matrix_complex_view Q2 = gsl_matrix_complex_submatrix(Q, 0, N, M, M - N);
          unpack_Q2(QR, T, &Q2.matrix);
        }

      /* copy R */
      gsl_matrix_complex_tricpy(CblasUpper, CblasNonUnit, R, &RV.matrix);

      return GSL_SUCCESS;
    }
}

/*
gsl_linalg_QR_QHvec_r()
  Apply M-by-M Q^H to the M-by-1 vector b

Inputs: QR   - [R; V] matrix encoded by gsl_linalg_complex_QR_decomp_r
        T    - block reflector matrix
        b    - M-by-1 vector replaced by Q^H b on output
        work - workspace, length N

Notes:
1) Q^H b = (I - V T^H V^H) b
         = b - V T^H [ V1^H V2^H ] [ b1 ]
                                   [ b2 ]
         = b - V T^H [ V1^H b1 + V2^H b2 ]
         = [ b1 ] - [ V1 w ]
           [ b2 ]   [ V2 w ]

where w = T^H ( V1^H b1 + V2^H b2 )
*/

int
gsl_linalg_complex_QR_QHvec_r(const gsl_matrix_complex * QR, const gsl_matrix_complex * T,
                              gsl_vector_complex * b, gsl_vector_complex * work)
{
  const size_t M = QR->size1;
  const size_t N = QR->size2;

  if (M < N)
    {
      GSL_ERROR ("M must be >= N", GSL_EBADLEN);
    }
  else if (T->size1 != N || T->size2 != N)
    {
      GSL_ERROR ("T matrix must be N-by-N", GSL_EBADLEN);
    }
  else if (b->size != M)
    {
      GSL_ERROR ("b vector must have length M", GSL_EBADLEN);
    }
  else if (work->size != N)
    {
      GSL_ERROR ("workspace must be length N", GSL_EBADLEN);
    }
  else
    {
      gsl_matrix_complex_const_view V1 = gsl_matrix_complex_const_submatrix(QR, 0, 0, N, N);
      gsl_vector_complex_view b1 = gsl_vector_complex_subvector(b, 0, N);
      gsl_vector_complex_view b2;

      /* work := V1^H b1 */
      gsl_vector_complex_memcpy(work, &b1.vector);
      gsl_blas_ztrmv(CblasLower, CblasConjTrans, CblasUnit, &V1.matrix, work);

      if (M > N)
        {
          gsl_matrix_complex_const_view V2 = gsl_matrix_complex_const_submatrix(QR, N, 0, M - N, N);

          /* work = work + V2^H b2 */
          b2 = gsl_vector_complex_subvector(b, N, M - N);
          gsl_blas_zgemv(CblasConjTrans, GSL_COMPLEX_ONE, &V2.matrix, &b2.vector, GSL_COMPLEX_ONE, work);
        }

      /* work = T^H * work */
      gsl_blas_ztrmv(CblasUpper, CblasConjTrans, CblasNonUnit, T, work);

      if (M > N)
        {
          /* b2 = b2 - V2 * work */
          gsl_matrix_complex_const_view V2 = gsl_matrix_complex_const_submatrix(QR, N, 0, M - N, N);
          gsl_blas_zgemv(CblasNoTrans, GSL_COMPLEX_NEGONE, &V2.matrix, work, GSL_COMPLEX_ONE, &b2.vector);
        }

      /* b1 = b1 - V1 * work */
      gsl_blas_ztrmv(CblasLower, CblasNoTrans, CblasUnit, &V1.matrix, work);
      gsl_vector_complex_sub(&b1.vector, work);

      return GSL_SUCCESS;
    }
}

/*
unpack_Q1()
  Compute Q_1

Inputs: Q  - on input, contains T in upper triangle and V in lower trapezoid
             on output, contains Q_1
             M-by-N

Return: success/error

Notes:
1)                                                         N
Q1 = [ Q1 Q2 ] [ I_n ] = (I - V T V^H) [ I; 0 ] = [ I - V1 T V1^H ] N
               [  0  ]                            [   - V2 T V1^H ] M - N
*/

static int
unpack_Q1(gsl_matrix_complex * Q)
{
  int status;
  const size_t M = Q->size1;
  const size_t N = Q->size2;
  gsl_matrix_complex_view Q1 = gsl_matrix_complex_submatrix(Q, 0, 0, N, N);
  gsl_vector_complex_view diag = gsl_matrix_complex_diagonal(&Q1.matrix);

  /* Q1 := T V1^H */
  status = aux_ULH(&Q1.matrix, &Q1.matrix);
  if (status)
    return status;

  if (M > N)
    {
      /* compute Q2 := - V2 T V1^H */
      gsl_matrix_complex_view V2 = gsl_matrix_complex_submatrix(Q, N, 0, M - N, N);
      gsl_blas_ztrmm(CblasRight, CblasUpper, CblasNoTrans, CblasNonUnit, GSL_COMPLEX_NEGONE, &Q1.matrix, &V2.matrix);
    }

  /* Q1 := - V1 T V1^H */
  status = aux_mLU(&Q1.matrix);
  if (status)
    return status;

  /* Q1 := I - V1 T V1^H */
  gsl_vector_complex_add_constant(&diag.vector, GSL_COMPLEX_ONE);

  return GSL_SUCCESS;
}

/*
unpack_Q2()
  Compute Q_2

Inputs: QR - [R; V] from QR_decomp_r, M-by-N
        T  - upper triangular T factor, N-by-N
        Q  - (output) Q_2 factor, M-by-(M-N)

Return: success/error

Notes:                         N   M-N
1) Since Q = I - V T V^H = M [ Q1  Q2  ], we have

           M-N
Q2 = Q [    0    ] N
       [ I_{M-N} ] M-N

So, Q2 = Q [ 0; I ] = (I - V T V^H) [ 0; I ] = [   - V1 T V2^H ]
                                               [ I - V2 T V2^H ]
*/

static int
unpack_Q2(const gsl_matrix_complex * QR, const gsl_matrix_complex * T, gsl_matrix_complex * Q)
{
  const size_t M = QR->size1;
  const size_t N = QR->size2;

  if (M <= N)
    {
      GSL_ERROR ("M must be > N", GSL_EBADLEN);
    }
  else if (T->size1 != N || T->size2 != N)
    {
      GSL_ERROR ("T matrix must be N-by-N", GSL_EBADLEN);
    }
  else if (Q->size1 != M || Q->size2 != (M - N))
    {
      GSL_ERROR ("Q matrix must be M-by-(M-N)", GSL_EBADLEN);
    }
  else
    {
      gsl_matrix_complex_const_view V1 = gsl_matrix_complex_const_submatrix(QR, 0, 0, N, N);
      gsl_matrix_complex_const_view V2 = gsl_matrix_complex_const_submatrix(QR, N, 0, M - N, N);
      gsl_matrix_complex_view Q1 = gsl_matrix_complex_submatrix(Q, 0, 0, N, M - N);
      gsl_matrix_complex_view Q2 = gsl_matrix_complex_submatrix(Q, N, 0, M - N, M - N);
      gsl_vector_complex_view diag = gsl_matrix_complex_diagonal(&Q2.matrix);

      /* Q1 := V2^H */
      gsl_matrix_complex_conjtrans_memcpy(&Q1.matrix, &V2.matrix);

      /* Q1 := - T V2^H */
      gsl_blas_ztrmm(CblasLeft, CblasUpper, CblasNoTrans, CblasNonUnit, GSL_COMPLEX_NEGONE, T, &Q1.matrix);

      /* Q2 := - V2 T V2^H */
      gsl_blas_zgemm(CblasNoTrans, CblasNoTrans, GSL_COMPLEX_ONE, &V2.matrix, &Q1.matrix, GSL_COMPLEX_ZERO, &Q2.matrix);

      /* Q2 := I - V2 T V2^H */
      gsl_vector_complex_add_constant(&diag.vector, GSL_COMPLEX_ONE);

      /* Q1 := - V1 T V2^H */
      gsl_blas_ztrmm(CblasLeft, CblasLower, CblasNoTrans, CblasUnit, GSL_COMPLEX_ONE, &V1.matrix, &Q1.matrix);

      return GSL_SUCCESS;
    }
}

/* U := U L^H for triangular matrices L and U; L is unit lower triangular */
static int
aux_ULH(const gsl_matrix_complex * L, gsl_matrix_complex * U)
{
  const size_t N = L->size1;

  if (N != L->size2)
    {
      GSL_ERROR ("L matrix must be square", GSL_ENOTSQR);
    }
  else if (U->size1 != N || U->size2 != N)
    {
      GSL_ERROR ("U matrix must be same size as L", GSL_EBADLEN);
    }
  else if (N == 1)
    {
      /* nothing to do */
      return GSL_SUCCESS;
    }
  else
    {
      int status;
      const size_t N1 = N / 2;
      const size_t N2 = N - N1;

      gsl_matrix_complex_const_view L11 = gsl_matrix_complex_const_submatrix(L, 0, 0, N1, N1);
      gsl_matrix_complex_const_view L21 = gsl_matrix_complex_const_submatrix(L, N1, 0, N2, N1);
      gsl_matrix_complex_const_view L22 = gsl_matrix_complex_const_submatrix(L, N1, N1, N2, N2);

      gsl_matrix_complex_view U11 = gsl_matrix_complex_submatrix(U, 0, 0, N1, N1);
      gsl_matrix_complex_view U12 = gsl_matrix_complex_submatrix(U, 0, N1, N1, N2);
      gsl_matrix_complex_view U22 = gsl_matrix_complex_submatrix(U, N1, N1, N2, N2);

      /* U12 = U12 * L22^H */
      gsl_blas_ztrmm(CblasRight, CblasLower, CblasConjTrans, CblasUnit, GSL_COMPLEX_ONE, &L22.matrix, &U12.matrix);

      /* U12 = U12 + U11 * L21^H */
      status = aux_ApUBH(&U11.matrix, &L21.matrix, &U12.matrix);
      if (status)
        return status;

      status = aux_ULH(&L11.matrix, &U11.matrix);
      if (status)
        return status;

      status = aux_ULH(&L22.matrix, &U22.matrix);
      if (status)
        return status;

      return GSL_SUCCESS;
    }
}

/* store -L*U in A */
static int
aux_mLU(gsl_matrix_complex * A)
{
  const size_t N = A->size1;

  if (N != A->size2)
    {
      GSL_ERROR ("matrix must be square", GSL_ENOTSQR);
    }
  else if (N == 1)
    {
      gsl_complex *A00 = gsl_matrix_complex_ptr(A, 0, 0);
      GSL_REAL(*A00) = -GSL_REAL(*A00);
      GSL_IMAG(*A00) = -GSL_IMAG(*A00);
      return GSL_SUCCESS;
    }
  else
    {
      int status;
      const size_t N1 = N / 2;
      const size_t N2 = N - N1;

      gsl_matrix_complex_view A11 = gsl_matrix_complex_submatrix(A, 0, 0, N1, N1);
      gsl_matrix_complex_view A12 = gsl_matrix_complex_submatrix(A, 0, N1, N1, N2);
      gsl_matrix_complex_view A21 = gsl_matrix_complex_submatrix(A, N1, 0, N2, N1);
      gsl_matrix_complex_view A22 = gsl_matrix_complex_submatrix(A, N1, N1, N2, N2);

      /* A22 = - L22 U22 */
      status = aux_mLU(&A22.matrix);
      if (status)
        return status;

      /* A22 = A22 - L21 U12 */
      gsl_blas_zgemm(CblasNoTrans, CblasNoTrans, GSL_COMPLEX_NEGONE, &A21.matrix, &A12.matrix, GSL_COMPLEX_ONE, &A22.matrix);

      /* A12 - -L11 U12 */
      gsl_blas_ztrmm(CblasLeft, CblasLower, CblasNoTrans, CblasUnit, GSL_COMPLEX_NEGONE, &A11.matrix, &A12.matrix);

      /* A21 = -L21 U11 */
      gsl_blas_ztrmm(CblasRight, CblasUpper, CblasNoTrans, CblasNonUnit, GSL_COMPLEX_NEGONE, &A11.matrix, &A21.matrix);

      /* A11 = - L11 U11 */
      status = aux_mLU(&A11.matrix);
      if (status)
        return status;

      return GSL_SUCCESS;
    }
}

/* A := A + U B^H where U is upper triangular */
static int
aux_ApUBH(const gsl_matrix_complex * U, const gsl_matrix_complex * B, gsl_matrix_complex * A)
{
  const size_t M = A->size1;
  const size_t N = A->size2;

  if (U->size1 != M || U->size2 != M)
    {
      GSL_ERROR ("U matrix has wrong dimensions", GSL_EBADLEN);
    }
  else if (B->size1 != N || B->size2 != M)
    {
      GSL_ERROR ("B matrix has wrong dimensions", GSL_EBADLEN);
    }
  else if (M == 1 && N == 1)
    {
      gsl_complex *aptr = gsl_matrix_complex_ptr(A, 0, 0);
      const gsl_complex U00 = gsl_matrix_complex_get(U, 0, 0);
      const gsl_complex B00_conj = gsl_complex_conjugate(gsl_matrix_complex_get(B, 0, 0));
      const gsl_complex prod = gsl_complex_mul(U00, B00_conj);
      GSL_REAL(*aptr) += GSL_REAL(prod);
      GSL_IMAG(*aptr) += GSL_IMAG(prod);
      return GSL_SUCCESS;
    }
  else if (M == 1)
    {
      gsl_complex U00 = gsl_matrix_complex_get(U, 0, 0);
      size_t i;

      for (i = 0; i < N; ++i)
        {
          gsl_complex * ai = gsl_matrix_complex_ptr(A, 0, i);
          gsl_complex bi = gsl_matrix_complex_get(B, i, 0);
          gsl_complex prod = gsl_complex_mul(U00, gsl_complex_conjugate(bi));
          GSL_REAL(*ai) += GSL_REAL(prod);
          GSL_IMAG(*ai) += GSL_IMAG(prod);
        }

      return GSL_SUCCESS;
    }
  else if (N == 1)
    {
      /*
       * partition:
       *
       *         M1  M2
       * B = 1 [ B11 B12 ]
       *
       *       1        M1  M2       1
       * M1 [ A11 ] + [ U11 U12 ] [ B11^H ] M1
       * M2 [ A21 ]   [  0  U22 ] [ B12^H ] M2
       */
      int status;
      const size_t M1 = M / 2;
      const size_t M2 = M - M1;

      gsl_matrix_complex_view A11 = gsl_matrix_complex_submatrix(A, 0, 0, M1, 1);
      gsl_matrix_complex_view A21 = gsl_matrix_complex_submatrix(A, M1, 0, M2, 1);

      gsl_matrix_complex_const_view U11 = gsl_matrix_complex_const_submatrix(U, 0, 0, M1, M1);
      gsl_matrix_complex_const_view U12 = gsl_matrix_complex_const_submatrix(U, 0, M1, M1, M2);
      gsl_matrix_complex_const_view U22 = gsl_matrix_complex_const_submatrix(U, M1, M1, M2, M2);

      gsl_matrix_complex_const_view B11 = gsl_matrix_complex_const_submatrix(B, 0, 0, 1, M1);
      gsl_matrix_complex_const_view B12 = gsl_matrix_complex_const_submatrix(B, 0, M1, 1, M2);

      /* A11 += U12 * B12^H */
      gsl_blas_zgemm(CblasNoTrans, CblasConjTrans, GSL_COMPLEX_ONE, &U12.matrix, &B12.matrix, GSL_COMPLEX_ONE, &A11.matrix);

      /* A11 := A11 + U11 B11^H */
      status = aux_ApUBH(&U11.matrix, &B11.matrix, &A11.matrix);
      if (status)
        return status;

      /* A21 := A21 + U22 B12^H */
      status = aux_ApUBH(&U22.matrix, &B12.matrix, &A21.matrix);
      if (status)
        return status;

      return GSL_SUCCESS;
    }
  else
    {
      int status;
      const size_t M1 = M / 2;
      const size_t M2 = M - M1;
      const size_t N1 = N / 2;
      const size_t N2 = N - N1;

      gsl_matrix_complex_view A11 = gsl_matrix_complex_submatrix(A, 0, 0, M1, N1);
      gsl_matrix_complex_view A12 = gsl_matrix_complex_submatrix(A, 0, N1, M1, N2);
      gsl_matrix_complex_view A21 = gsl_matrix_complex_submatrix(A, M1, 0, M2, N1);
      gsl_matrix_complex_view A22 = gsl_matrix_complex_submatrix(A, M1, N1, M2, N2);

      gsl_matrix_complex_const_view U11 = gsl_matrix_complex_const_submatrix(U, 0, 0, M1, M1);
      gsl_matrix_complex_const_view U12 = gsl_matrix_complex_const_submatrix(U, 0, M1, M1, M2);
      gsl_matrix_complex_const_view U22 = gsl_matrix_complex_const_submatrix(U, M1, M1, M2, M2);

      gsl_matrix_complex_const_view B11 = gsl_matrix_complex_const_submatrix(B, 0, 0, N1, M1);
      gsl_matrix_complex_const_view B12 = gsl_matrix_complex_const_submatrix(B, 0, M1, N1, M2);
      gsl_matrix_complex_const_view B21 = gsl_matrix_complex_const_submatrix(B, N1, 0, N2, M1);
      gsl_matrix_complex_const_view B22 = gsl_matrix_complex_const_submatrix(B, N1, M1, N2, M2);

      /* A11 := A11 + U11 B11^H */
      status = aux_ApUBH(&U11.matrix, &B11.matrix, &A11.matrix);
      if (status)
        return status;

      /* A11 := A11 + U12 B12^H */
      gsl_blas_zgemm(CblasNoTrans, CblasConjTrans, GSL_COMPLEX_ONE, &U12.matrix, &B12.matrix, GSL_COMPLEX_ONE, &A11.matrix);

      /* A12 := A12 + U11 B21^H */
      status = aux_ApUBH(&U11.matrix, &B21.matrix, &A12.matrix);
      if (status)
        return status;

      /* A12 := A12 + U12 B22^H */
      gsl_blas_zgemm(CblasNoTrans, CblasConjTrans, GSL_COMPLEX_ONE, &U12.matrix, &B22.matrix, GSL_COMPLEX_ONE, &A12.matrix);

      /* A21 := A21 + U22 B12^H */
      status = aux_ApUBH(&U22.matrix, &B12.matrix, &A21.matrix);
      if (status)
        return status;

      /* A22 := A22 + U22 B22^H */
      status = aux_ApUBH(&U22.matrix, &B22.matrix, &A22.matrix);
      if (status)
        return status;

      return GSL_SUCCESS;
    }
}

max_stars_repo_name: zzpwahaha/Chimera-Control-Trim

max_issues_repo_name: zzpwahaha/Chimera-Control-Trim

max_forks_repo_name: zzpwahaha/Chimera-Control-Trim