/* multifit/fdfridge.c

 * 

 * Copyright (C) 2014 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 <gsl/gsl_math.h>

#include <gsl/gsl_errno.h>

#include <gsl/gsl_vector.h>

#include <gsl/gsl_matrix.h>

#include <gsl/gsl_multifit_nlin.h>

#include <gsl/gsl_blas.h>

static int fdfridge_f(const gsl_vector * x, void * params, gsl_vector * f);

static int fdfridge_df(const gsl_vector * x, void * params, gsl_matrix * J);

gsl_multifit_fdfridge *

gsl_multifit_fdfridge_alloc (const gsl_multifit_fdfsolver_type * T,

                             const size_t n, const size_t p)

{

  gsl_multifit_fdfridge * work;

  work = calloc(1, sizeof(gsl_multifit_fdfridge));

  if (work == NULL)

    {

      GSL_ERROR_VAL("failed to allocate workspace",

                    GSL_ENOMEM, 0);

    }

  work->s = gsl_multifit_fdfsolver_alloc (T, n + p, p);

  if (work->s == NULL)

    {

      gsl_multifit_fdfridge_free(work);

      GSL_ERROR_VAL("failed to allocate space for fdfsolver",

                    GSL_ENOMEM, 0);

    }

  work->wts = gsl_vector_alloc(n + p);

  if (work->wts == NULL)

    {

      gsl_multifit_fdfridge_free(work);

      GSL_ERROR_VAL("failed to allocate space for weight vector",

                    GSL_ENOMEM, 0);

    }

  work->f = gsl_vector_alloc(n);

  if (work->f == NULL)

    {

      gsl_multifit_fdfridge_free(work);

      GSL_ERROR_VAL("failed to allocate space for f vector",

                    GSL_ENOMEM, 0);

    }

  work->n = n;

  work->p = p;

  work->lambda = 0.0;

  /* initialize weights to 1 (for augmented portion of vector) */

  gsl_vector_set_all(work->wts, 1.0);

  return work;

} /* gsl_multifit_fdfridge_alloc() */

void

gsl_multifit_fdfridge_free(gsl_multifit_fdfridge *work)

{

  if (work->s)

    gsl_multifit_fdfsolver_free(work->s);

  if (work->wts)

    gsl_vector_free(work->wts);

  if (work->f)

    gsl_vector_free(work->f);

  free(work);

}

const char *

gsl_multifit_fdfridge_name(const gsl_multifit_fdfridge * w)

{

  return gsl_multifit_fdfsolver_name(w->s);

}

gsl_vector *

gsl_multifit_fdfridge_position (const gsl_multifit_fdfridge * w)

{

  return gsl_multifit_fdfsolver_position(w->s);

}

gsl_vector *

gsl_multifit_fdfridge_residual (const gsl_multifit_fdfridge * w)

{

  return gsl_multifit_fdfsolver_residual(w->s);

}

size_t

gsl_multifit_fdfridge_niter (const gsl_multifit_fdfridge * w)

{

  return w->s->niter;

}

int

gsl_multifit_fdfridge_set (gsl_multifit_fdfridge * w,

                           gsl_multifit_function_fdf * f,

                           const gsl_vector * x,

                           const double lambda)

{

  return gsl_multifit_fdfridge_wset(w, f, x, lambda, NULL);

} /* gsl_multifit_fdfridge_set() */

int

gsl_multifit_fdfridge_wset (gsl_multifit_fdfridge * w,

                            gsl_multifit_function_fdf * f,

                            const gsl_vector * x,

                            const double lambda,

                            const gsl_vector * wts)

{

  const size_t n = w->n;

  const size_t p = w->p;

  if (n != f->n || p != f->p)

    {

      GSL_ERROR ("function size does not match solver", GSL_EBADLEN);

    }

  else if (p != x->size)

    {

      GSL_ERROR ("vector length does not match solver", GSL_EBADLEN);

    }

  else if (wts != NULL && n != wts->size)

    {

      GSL_ERROR ("weight vector length does not match solver", GSL_EBADLEN);

    }

  else

    {

      int status;

      gsl_vector_view wv = gsl_vector_subvector(w->wts, 0, n);

      /* save user defined fdf */

      w->fdf = f;

      /* build modified fdf for Tikhonov terms */

      w->fdftik.f = &fdfridge_f;

      w->fdftik.df = &fdfridge_df;

      w->fdftik.n = n + p; /* add p for Tikhonov terms */

      w->fdftik.p = p;

      w->fdftik.params = (void *) w;

      /* store damping parameter */

      w->lambda = lambda;

      w->L = NULL;

      if (wts)

        {

          /* copy weight vector into user portion of w->wts */

          gsl_vector_memcpy(&wv.vector, wts);

          status = gsl_multifit_fdfsolver_wset(w->s, &(w->fdftik), x, w->wts);

        }

      else

        {

          status = gsl_multifit_fdfsolver_wset(w->s, &(w->fdftik), x, NULL);

        }

      /* update function/Jacobian evaluations */

      f->nevalf = w->fdftik.nevalf;

      f->nevaldf = w->fdftik.nevaldf;

      return status;

    }

} /* gsl_multifit_fdfridge_wset() */

int

gsl_multifit_fdfridge_set2 (gsl_multifit_fdfridge * w,

                            gsl_multifit_function_fdf * f,

                            const gsl_vector * x,

                            const gsl_vector * lambda)

{

  return gsl_multifit_fdfridge_wset2(w, f, x, lambda, NULL);

} /* gsl_multifit_fdfridge_set2() */

int

gsl_multifit_fdfridge_wset2 (gsl_multifit_fdfridge * w,

                             gsl_multifit_function_fdf * f,

                             const gsl_vector * x,

                             const gsl_vector * lambda,

                             const gsl_vector * wts)

{

  const size_t n = w->n;

  const size_t p = w->p;

  if (n != f->n || p != f->p)

    {

      GSL_ERROR ("function size does not match solver", GSL_EBADLEN);

    }

  else if (p != x->size)

    {

      GSL_ERROR ("vector length does not match solver", GSL_EBADLEN);

    }

  else if (lambda->size != p)

    {

      GSL_ERROR ("lambda vector size does not match solver", GSL_EBADLEN);

    }

  else if (wts != NULL && n != wts->size)

    {

      GSL_ERROR ("weight vector length does not match solver", GSL_EBADLEN);

    }

  else

    {

      int status;

      gsl_vector_view wv = gsl_vector_subvector(w->wts, 0, n);

      /* save user defined fdf */

      w->fdf = f;

      w->fdf->nevalf = 0;

      w->fdf->nevaldf = 0;

      /* build modified fdf for Tikhonov terms */

      w->fdftik.f = &fdfridge_f;

      w->fdftik.df = &fdfridge_df;

      w->fdftik.n = n + p; /* add p for Tikhonov terms */

      w->fdftik.p = p;

      w->fdftik.params = (void *) w;

      /* store damping matrix */

      w->lambda = 0.0;

      w->L_diag = lambda;

      w->L = NULL;

      if (wts)

        {

          /* copy weight vector into user portion */

          gsl_vector_memcpy(&wv.vector, wts);

          status = gsl_multifit_fdfsolver_wset(w->s, &(w->fdftik), x, w->wts);

        }

      else

        {

          status = gsl_multifit_fdfsolver_wset(w->s, &(w->fdftik), x, NULL);

        }

      /* update function/Jacobian evaluations */

      f->nevalf = w->fdftik.nevalf;

      f->nevaldf = w->fdftik.nevaldf;

      return status;

    }

} /* gsl_multifit_fdfridge_wset2() */

int

gsl_multifit_fdfridge_set3 (gsl_multifit_fdfridge * w,

                            gsl_multifit_function_fdf * f,

                            const gsl_vector * x,

                            const gsl_matrix * L)

{

  return gsl_multifit_fdfridge_wset3(w, f, x, L, NULL);

} /* gsl_multifit_fdfridge_set3() */

int

gsl_multifit_fdfridge_wset3 (gsl_multifit_fdfridge * w,

                             gsl_multifit_function_fdf * f,

                             const gsl_vector * x,

                             const gsl_matrix * L,

                             const gsl_vector * wts)

{

  const size_t n = w->n;

  const size_t p = w->p;

  if (n != f->n || p != f->p)

    {

      GSL_ERROR ("function size does not match solver", GSL_EBADLEN);

    }

  else if (p != x->size)

    {

      GSL_ERROR ("vector length does not match solver", GSL_EBADLEN);

    }

  else if (L->size2 != p)

    {

      GSL_ERROR ("L matrix size2 does not match solver", GSL_EBADLEN);

    }

  else if (wts != NULL && n != wts->size)

    {

      GSL_ERROR ("weight vector length does not match solver", GSL_EBADLEN);

    }

  else

    {

      int status;

      gsl_vector_view wv = gsl_vector_subvector(w->wts, 0, n);

      /* save user defined fdf */

      w->fdf = f;

      w->fdf->nevalf = 0;

      w->fdf->nevaldf = 0;

      /* build modified fdf for Tikhonov terms */

      w->fdftik.f = &fdfridge_f;

      w->fdftik.df = &fdfridge_df;

      w->fdftik.n = n + p; /* add p for Tikhonov terms */

      w->fdftik.p = p;

      w->fdftik.params = (void *) w;

      /* store damping matrix */

      w->lambda = 0.0;

      w->L_diag = NULL;

      w->L = L;

      if (wts)

        {

          /* copy weight vector into user portion */

          gsl_vector_memcpy(&wv.vector, wts);

          status = gsl_multifit_fdfsolver_wset(w->s, &(w->fdftik), x, w->wts);

        }

      else

        {

          status = gsl_multifit_fdfsolver_wset(w->s, &(w->fdftik), x, NULL);

        }

      /* update function/Jacobian evaluations */

      f->nevalf = w->fdftik.nevalf;

      f->nevaldf = w->fdftik.nevaldf;

      return status;

    }

} /* gsl_multifit_fdfridge_wset3() */

int

gsl_multifit_fdfridge_iterate (gsl_multifit_fdfridge * w)

{

  int status = gsl_multifit_fdfsolver_iterate(w->s);

  /* update function/Jacobian evaluations */

  w->fdf->nevalf = w->fdftik.nevalf;

  w->fdf->nevaldf = w->fdftik.nevaldf;

  return status;

}

int

gsl_multifit_fdfridge_driver (gsl_multifit_fdfridge * w,

                              const size_t maxiter,

                              const double xtol,

                              const double gtol,

                              const double ftol,

                              int *info)

{

  int status = gsl_multifit_fdfsolver_driver(w->s, maxiter, xtol,

                                             gtol, ftol, info);

  return status;

} /* gsl_multifit_fdfridge_driver() */

/*

fdfridge_f()

  Callback function to provide residuals, including extra p

Tikhonov terms. The residual vector will have the form:

f~ = [     f     ]

     [ \lambda x ]

where f is the user supplied residuals, x are the model

parameters, and \lambda is the Tikhonov damping parameter

Inputs: x      - model parameters (size p)

        params - pointer to fdfridge workspace

        f      - (output) (n+p) vector to store f~

*/

static int

fdfridge_f(const gsl_vector * x, void * params, gsl_vector * f)

{

  int status;

  gsl_multifit_fdfridge *w = (gsl_multifit_fdfridge *) params;

  const size_t n = w->n;

  const size_t p = w->p;

  gsl_vector_view f_user = gsl_vector_subvector(f, 0, n);

  gsl_vector_view f_tik = gsl_vector_subvector(f, n, p);

  /* call user callback function to get residual vector f */

  status = gsl_multifit_eval_wf(w->fdf, x, NULL, &f_user.vector);

  if (status)

    return status;

  if (w->L_diag)

    {

      /* store diag(L_diag) x in Tikhonov portion of f~ */

      gsl_vector_memcpy(&f_tik.vector, x);

      gsl_vector_mul(&f_tik.vector, w->L_diag);

    }

  else if (w->L)

    {

      /* store Lx in Tikhonov portion of f~ */

      gsl_blas_dgemv(CblasNoTrans, 1.0, w->L, x, 0.0, &f_tik.vector);

    }

  else

    {

      /* store \lambda x in Tikhonov portion of f~ */

      gsl_vector_memcpy(&f_tik.vector, x);

      gsl_vector_scale(&f_tik.vector, w->lambda);

    }

  return GSL_SUCCESS;

} /* fdfridge_f() */

static int

fdfridge_df(const gsl_vector * x, void * params, gsl_matrix * J)

{

  int status;

  gsl_multifit_fdfridge *w = (gsl_multifit_fdfridge *) params;

  const size_t n = w->n;

  const size_t p = w->p;

  gsl_matrix_view J_user = gsl_matrix_submatrix(J, 0, 0, n, p);

  gsl_matrix_view J_tik = gsl_matrix_submatrix(J, n, 0, p, p);

  gsl_vector_view diag = gsl_matrix_diagonal(&J_tik.matrix);

  /* compute Jacobian */

  if (w->fdf->df)

    status = gsl_multifit_eval_wdf(w->fdf, x, NULL, &J_user.matrix);

  else

    {

      /* compute f(x) and then finite difference Jacobian */

      status = gsl_multifit_eval_wf(w->fdf, x, NULL, w->f);

      status += gsl_multifit_fdfsolver_dif_df(x, NULL, w->fdf, w->f,

                                              &J_user.matrix);

    }

  if (status)

    return status;

  if (w->L_diag)

    {

      /* store diag(L_diag) in Tikhonov portion of J */

      gsl_matrix_set_zero(&J_tik.matrix);

      gsl_vector_memcpy(&diag.vector, w->L_diag);

    }

  else if (w->L)

    {

      /* store L in Tikhonov portion of J */

      gsl_matrix_memcpy(&J_tik.matrix, w->L);

    }

  else

    {

      /* store \lambda I_p in Tikhonov portion of J */

      gsl_matrix_set_zero(&J_tik.matrix);

      gsl_vector_set_all(&diag.vector, w->lambda);

    }

  return GSL_SUCCESS;

} /* fdfridge_df() */