/* multifit_nlinear/qr.c

 * 

 * Copyright (C) 2015, 2016 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.

 */

/*

 * This module handles the solution of the linear least squares

 * system:

 *

 * [     J~      ] p~ = - [ f ]

 * [ sqrt(mu)*D~ ]        [ 0 ]

 *

 * using a QR approach. Quantities are scaled according to:

 *

 * J~ = J S

 * D~ = D S

 * p~ = S^{-1} p

 *

 * where S is a diagonal matrix and S_jj = || J_j || and J_j is column

 * j of the Jacobian. This balancing transformation seems to be more

 * numerically stable for some Jacobians.

 */

#include <config.h>

#include <gsl/gsl_math.h>

#include <gsl/gsl_vector.h>

#include <gsl/gsl_matrix.h>

#include <gsl/gsl_linalg.h>

#include <gsl/gsl_multifit_nlinear.h>

#include <gsl/gsl_errno.h>

#include <gsl/gsl_blas.h>

#include "common.c"

#include "qrsolv.c"

typedef struct

{

  size_t p;

  gsl_matrix *QR;            /* QR factorization of J */

  gsl_vector *tau_Q;         /* Householder scalars for Q */

  gsl_matrix *T;             /* workspace matrix for qrsolv, p-by-p */

  gsl_permutation *perm;     /* permutation matrix */

  size_t rank;               /* rank of J */

  gsl_vector *residual;      /* residual of LS problem [ J; sqrt(mu) D ] p = - [ f; 0 ] */

  gsl_vector *qtf;           /* Q^T f */

  gsl_vector *workn;         /* workspace, length n */

  gsl_vector *workp;         /* workspace, length p */

  gsl_vector *work3p;        /* workspace, length 3*p */

  double mu;                 /* LM parameter */

} qr_state_t;

static int qr_init(const void * vtrust_state, void * vstate);

static int qr_presolve(const double mu, const void * vtrust_state, void * vstate);

static int qr_solve(const gsl_vector * f, gsl_vector *x,

                    const void * vtrust_state, void *vstate);

static int qr_rcond(double * rcond, void * vstate);

static void *

qr_alloc (const size_t n, const size_t p)

{

  qr_state_t *state;

  (void)n;

  state = calloc(1, sizeof(qr_state_t));

  if (state == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate qr state", GSL_ENOMEM);

    }

  state->QR = gsl_matrix_alloc(n, p);

  if (state->QR == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for QR", GSL_ENOMEM);

    }

  state->tau_Q = gsl_vector_alloc(p);

  if (state->tau_Q == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for tau_Q",

                      GSL_ENOMEM);

    }

  state->T = gsl_matrix_alloc(p, p);

  if (state->T == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for T", GSL_ENOMEM);

    }

  state->qtf = gsl_vector_alloc(n);

  if (state->qtf == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for qtf",

                      GSL_ENOMEM);

    }

  state->residual = gsl_vector_alloc(n);

  if (state->residual == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for residual",

                      GSL_ENOMEM);

    }

  state->perm = gsl_permutation_calloc(p);

  if (state->perm == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for perm",

                      GSL_ENOMEM);

    }

  state->workn = gsl_vector_alloc(n);

  if (state->workn == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for workn",

                      GSL_ENOMEM);

    }

  state->workp = gsl_vector_alloc(p);

  if (state->workp == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for workp",

                      GSL_ENOMEM);

    }

  state->work3p = gsl_vector_alloc(3 * p);

  if (state->work3p == NULL)

    {

      GSL_ERROR_NULL ("failed to allocate space for work3p",

                      GSL_ENOMEM);

    }

  state->p = p;

  state->mu = 0.0;

  state->rank = 0;

  return state;

}

static void

qr_free(void *vstate)

{

  qr_state_t *state = (qr_state_t *) vstate;

  if (state->QR)

    gsl_matrix_free(state->QR);

  if (state->tau_Q)

    gsl_vector_free(state->tau_Q);

  if (state->T)

    gsl_matrix_free(state->T);

  if (state->qtf)

    gsl_vector_free(state->qtf);

  if (state->residual)

    gsl_vector_free(state->residual);

  if (state->perm)

    gsl_permutation_free(state->perm);

  if (state->workn)

    gsl_vector_free(state->workn);

  if (state->workp)

    gsl_vector_free(state->workp);

  if (state->work3p)

    gsl_vector_free(state->work3p);

  free(state);

}

/* compute J = Q R PT */

static int

qr_init(const void * vtrust_state, void * vstate)

{

  const gsl_multifit_nlinear_trust_state *trust_state =

    (const gsl_multifit_nlinear_trust_state *) vtrust_state;

  qr_state_t *state = (qr_state_t *) vstate;

  int signum;

  /* perform QR decomposition of J */

  gsl_matrix_memcpy(state->QR, trust_state->J);

  gsl_linalg_QRPT_decomp(state->QR, state->tau_Q, state->perm,

                         &signum, state->workp);

  return GSL_SUCCESS;

}

static int

qr_presolve(const double mu, const void * vtrust_state, void * vstate)

{

  qr_state_t *state = (qr_state_t *) vstate;

  state->mu = mu;

  (void) vtrust_state;

  return GSL_SUCCESS;

}

static int

qr_solve(const gsl_vector * f, gsl_vector *x,

         const void * vtrust_state, void *vstate)

{

  qr_state_t *state = (qr_state_t *) vstate;

  int status;

  if (state->mu == 0.0)

    {

      /*

       * compute Gauss-Newton direction by solving

       * J x = f

       * with an attempt to identify rank deficiency in J

       */

      size_t rank = gsl_linalg_QRPT_rank(state->QR, -1.0);

      status = gsl_linalg_QRPT_lssolve2(state->QR, state->tau_Q, state->perm,

                                        f, rank, x, state->residual);

    }

  else

    {

      /*

       * solve:

       *

       * [     J      ] x = [ f ]

       * [ sqrt(mu) D ]     [ 0 ]

       *

       * using QRPT factorization of J

       */

      const gsl_multifit_nlinear_trust_state *trust_state =

        (const gsl_multifit_nlinear_trust_state *) vtrust_state;

      double sqrt_mu = sqrt(state->mu);

      /* compute qtf = Q^T f */

      gsl_vector_memcpy(state->qtf, f);

      gsl_linalg_QR_QTvec(state->QR, state->tau_Q, state->qtf);

      status = qrsolv(state->QR, state->perm, sqrt_mu, trust_state->diag,

                      state->qtf, state->T, x, state->workn);

    }

  /* reverse step to go downhill */

  gsl_vector_scale(x, -1.0);

  return status;

}

static int

qr_rcond(double * rcond, void * vstate)

{

  int status;

  qr_state_t *state = (qr_state_t *) vstate;

  status = gsl_linalg_QRPT_rcond(state->QR, rcond, state->work3p);

  return status;

}

static const gsl_multifit_nlinear_solver qr_type =

{

  "qr",

  qr_alloc,

  qr_init,

  qr_presolve,

  qr_solve,

  qr_rcond,

  qr_free

};

const gsl_multifit_nlinear_solver *gsl_multifit_nlinear_solver_qr = &qr_type;