! $Id: ropp_1dvar_levmarq.f90 6490 2020-09-10 17:36:27Z idculv $

!****s* 1DVar/ropp_1dvar_levmarq
!
! NAME
!    ropp_1dvar_levmarq - Solve the 1DVar for background data using the
!                         Levenberg-Marquardt minimiser
!
! SYNOPSIS
!    CALL ropp_1dvar_levmarq(obs, bg, state, config, diag)
!
!
! DESCRIPTION
!    This subroutine evaluates a quadratic cost function for a
!    variational data assimilation procedure.
!
!    More specifically, this routine calculates a cost function
!
!           1  /         |  -1 |         \
!       J = - < y - H(x) | O   | y - H(x) > +
!           2  \         |     |         /
!
!                       1  /       |  -1 |       \
!                       - < x - x  | B   | x - x  >
!                       2  \     b |     |      b/
!
!    where the background state x_b is given by the state vector
!    state.
!
!    A solution for x is obtained by minimising J using the Levenberg-Marquardt
!    minimisation method.
!
! INPUTS
!    obs                Observation data structure.
!    bg                 Background data structure.
!    state              State vector structure.
!    config             Configuration structure.
!    diag               Diagnostics structure.
!
! OUTPUT
!    state
!    diag
!
! REFERENCES
!   W.H. Press, S.A. Teukolsjy, W.T. Vetterling and B.P. Flannery,
!   Numerical Recipes in C - The Art of Scientific Computing.
!   2nd Ed., Cambridge University Press, 1992.
!
! AUTHOR
!   Met Office, Exeter, UK.
!   Any comments on this software should be given via the ROM SAF
!   Helpdesk at http://www.romsaf.org
!
! COPYRIGHT
!   (c) EUMETSAT. All rights reserved.
!   For further details please refer to the file COPYRIGHT
!   which you should have received as part of this distribution.
!
!****

!-------------------------------------------------------------------------------
! 1. Bending angle
!-------------------------------------------------------------------------------

SUBROUTINE ropp_1dvar_levmarq_bangle(obs, bg, state, config, diag)
  
! 1.1 Declarations
! ----------------

  USE typesizes, ONLY: wp => EightByteReal
  USE ropp_utils
  USE ropp_fm
  USE ropp_1dvar, not_this => ropp_1dvar_levmarq_bangle
  USE matrix

  IMPLICIT NONE

  TYPE(Obs1dBangle), INTENT(inout)      :: obs          ! Observation data
  TYPE(State1dFM),   INTENT(inout)      :: bg           ! Background data
  TYPE(State1dFM),   INTENT(inout)      :: state        ! State vector
  TYPE(VarConfig),   INTENT(in)         :: config       ! Configuration options
  TYPE(VarDiag),     INTENT(inout)      :: diag         ! Diagnostic output

  REAL(wp)                              :: J            ! Cost function value
  TYPE(State1dFM)                       :: x            ! Control vector
  TYPE(State1dFM)                       :: x_old        ! Control for LM parameter testing
  TYPE(Obs1dBangle)                     :: y            ! Forward model obs
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: K            ! K-matrix
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_x      ! Change of state
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_y      ! Change of observation
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: dJ_dx        ! Cost function gradient
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: diag_d2J     ! Diagonal d2J/dx2
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: d2J_dx2      ! 2nd derivative cost fn
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: KO           ! K O^-1
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: Bm1          ! B^-1
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: Om1          ! O^-1
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_J      ! Change of cost fn
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: state_last   ! Previous state vector
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: state_sigma  ! State std deviation
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_state  ! Change of state vector
  REAL(wp)                              :: J_last       ! Previous cost function
  REAL(wp)                              :: J_min        ! Minimum cost function
  INTEGER                               :: i_pointer    ! Value index
  INTEGER                               :: i            ! Counter
  INTEGER                               :: n_iter       ! Number of iterations
  INTEGER                               :: n_iter_max   ! Maximum number of iterations
  INTEGER                               :: n_grad       ! Number of gradient updates
  INTEGER                               :: nobs         ! Number of observations
  INTEGER                               :: nstate       ! No. of state elements
  REAL(wp)                              :: lambda       ! Levenberg-Marquardt factor
  REAL(wp)                              :: lambda_max   ! Maximum Levenberg-Marquardt factor

  CHARACTER(len =   4)                  :: it_str
  CHARACTER(len =   4)                  :: gr_str
  CHARACTER(len =  15)                  :: ch_str, co_str
  CHARACTER(len = 256)                  :: routine

! 1.2 Message handling
! --------------------

  CALL message_get_routine(routine)
  CALL message_set_routine('ropp_1dvar_levmarq_bangle')

! 1.3 Initialise variables
! ------------------------

  i_pointer  = 0
  n_grad     = 0
  n_iter     = 0
  n_iter_max = 50

  lambda     = 1.0E-4_wp
  IF (bg%direct_ion) lambda = 1.0E-2_wp
  lambda_max = 1.0E10_wp

  J          = 0.0_wp
  J_last     = 1.0E30_wp
  J_min      = J_last

  nstate = SIZE(bg%state)
  nobs   = SIZE(obs%bangle)

  ALLOCATE(K(nobs, nstate))
  ALLOCATE(delta_x(nstate))
  ALLOCATE(delta_y(nobs))
  ALLOCATE(dJ_dx(nstate))
  ALLOCATE(diag_d2J(nstate))
  ALLOCATE(d2J_dx2(nstate,nstate))
  ALLOCATE(KO(nstate,nobs))
  ALLOCATE(Bm1(nstate,nstate))
  ALLOCATE(Om1(nobs,nobs))

  IF (ALLOCATED(delta_J)) DEALLOCATE(delta_J)
  ALLOCATE(delta_J(config%conv_check_n_previous))
  delta_J(:) = 0.0_wp

  IF (ALLOCATED(state_last)) DEALLOCATE(state_last)
  ALLOCATE(state_last(SIZE(bg%state)))
  state_last(:) = 0.0_wp

  IF (ALLOCATED(state_sigma)) DEALLOCATE(state_sigma)
  ALLOCATE(state_sigma(SIZE(bg%state)))
  DO i = 1, SIZE(state_sigma)
    state_sigma(i) = SQRT(bg%cov%d(i + i*(i-1)/2)) ! Direct read from matrix_pp
  ENDDO

  IF (ALLOCATED(delta_state)) DEALLOCATE(delta_state)
  ALLOCATE(delta_state(config%conv_check_n_previous))
  delta_state(:) = 0.0_wp

  ! 1.4 Inverse error covariances
  ! -----------------------------

  Bm1 = matrix_invert(bg%cov)
  Om1 = matrix_invert(obs%cov)

  ! 1.5 First guess and pseudo-observations
  ! ---------------------------------------

  x = bg
  y = obs

  ! 1.6 Initial cost function
  ! -------------------------

  IF (ASSOCIATED(x%ak)) THEN
    CALL ropp_fm_state2state_ecmwf(x)
  ELSE
    CALL ropp_fm_state2state_meto(x)
  END IF

  CALL ropp_fm_bangle_1d(x, y)

  delta_x = (x%state - bg%state)
  delta_y = (y%bangle - obs%bangle) * y%weights

  J = 0.5_wp * DOT_PRODUCT(delta_y, MATMUL(Om1, delta_y)) + &
      0.5_wp * DOT_PRODUCT(delta_x, MATMUL(Bm1, delta_x))

  diag%J_init = J
  J_last = J
  J_min  = J

  IF (config%minropp%impres == 0) THEN
    WRITE(it_str, '(i4)') n_iter
    WRITE(ch_str, '(g15.5)') J
    ch_str = ADJUSTL(ch_str)
    co_str = ' -             '

    CALL message(msg_cont, &
         '   n_iter = ' // it_str // '   J = ' // ch_str //  &
         '   max(relative change in state) = ' // co_str)
  END IF

  ! 1.7 Main minimisation iteration loop
  ! ------------------------------------

  DO WHILE (n_iter < n_iter_max)

    ! 1.7.1 Bookkeeping

    n_grad = n_grad + 1

    ! 1.7.2 Compute gradient and Hessian of cost function

    CALL ropp_fm_bangle_1d_grad(x, y, K)

    !TODO: (a) Do we need this?
    !      (b) Does it change the cost function?
    !      (c) When would a change in bending angle ever exceed 50 rad?
    !          Was this lifted from the refrac code? (delta_N = 50 N-units just about plausible.)
    WHERE (ABS(delta_y) > 50.0_wp)
      delta_y = 0.0_wp
    END WHERE

    KO      = MATMUL(TRANSPOSE(K), Om1)
    dJ_dx   = MATMUL(KO, delta_y) + MATMUL(Bm1, delta_x)
    d2J_dx2 = Bm1 + MATMUL(KO, K)

    ! 1.7.3 Levenberg-Marquardt adjustment of Hessian diagonal

    DO i = 1, nstate
      d2J_dx2(i,i) = d2J_dx2(i,i) * (1.0_wp + lambda)
    END DO

    ! 1.7.4 Solve matrix equation d2J_dx2 . dx = -dJ_dx

    delta_x = matrix_solve(d2J_dx2, - dJ_dx)

    ! 1.7.5 Update test state vector

    x_old = x

    IF (x%use_logq) THEN
      x%state = x%state + SIGN(MIN(ABS(delta_x), ABS(x%state/2.0_wp)), delta_x)
    ELSE
      x%state = x%state + delta_x
    END IF

    ! 1.7.6 Special handling of ionospheric state vector elements

    IF (bg%direct_ion) THEN

      i = nstate - 2
      IF (x%state(i) < ropp_ZERO) THEN
        CALL message(msg_warn, "Levenberg-Marquardt solver returns " // &
          "Ne_max < 0 ... suggest examining final value. \n")
      END IF

      i = nstate - 1
      IF (x%state(i) < 0.01_wp*bg%state(i)) THEN
        CALL message(msg_warn, "Levenberg-Marquardt solver returns " // &
          "H_peak < 1% of background ... resetting to background value. \n")
          x%state(i) = bg%state(i)
      END IF

      i = nstate
      IF (x%state(i) < 0.01_wp*bg%state(i)) THEN
        CALL message(msg_warn, "Levenberg-Marquardt solver returns " // &
          "H_width < 1% of background ... resetting to background value. \n")
          x%state(i) = bg%state(i)
      END IF

    END IF

    ! 1.7.7 Compute cost function

    IF (ASSOCIATED(x%ak)) THEN
       CALL ropp_fm_state2state_ecmwf(x)
    ELSE
       CALL ropp_fm_state2state_meto(x)
    END IF

    CALL ropp_fm_bangle_1d(x, y)

    delta_x = (x%state - bg%state)
    delta_y = (y%bangle - obs%bangle) * y%weights

    J  = 0.5_wp * DOT_PRODUCT(delta_y, MATMUL(Om1, delta_y)) + &
         0.5_wp * DOT_PRODUCT(delta_x, MATMUL(Bm1, delta_x))

    ! 1.8 Levenberg-Marquardt update and convergence check
    ! ----------------------------------------------------

    IF ( J >= (J_last + config%conv_check_max_delta_J) ) THEN ! Keep previous state vector, increase lambda, repeat

      ! 1.8.1 Keep old state, increase lambda, bail out if it's too big

      x      = x_old
      lambda = 10.0_wp * lambda

      IF (ABS(lambda) > lambda_max) THEN

        WRITE(ch_str, '(g15.5)') lambda_max
        ch_str = ADJUSTL(ch_str)
        CALL message(msg_cont, '')
        CALL message(msg_info,                                               &
          'Levenberg-Marquardt parameter exceeded max (=' // TRIM(ch_str) // &
          '); convergence is unlikely.\n   ' //                              &
          'Check convergence parameters; conv_check_max_delta_J might ' //   &
          'be too demanding.\n')
         EXIT

      END IF

      ! 1.8.2 Logging

      !TODO: Add a configuration option for lambda update outputs
      IF (config%minropp%impres == 0) THEN

        WRITE(it_str, '(i4)') n_iter
        WRITE(ch_str, '(g15.5)') J
        ch_str = ADJUSTL(ch_str)
        IF (n_iter > 0) THEN
          WRITE(co_str, '(g15.5)') lambda
          co_str = ADJUSTL(co_str)
        ELSE
          co_str = ' -             '
        END IF

        CALL message(msg_cont, &
                     '   n_iter = ' // it_str // '   J = ' // ch_str //  &
                     '   lambda -> ' // co_str)

      END IF

      ! 1.8.3 Next iteration

      CYCLE

    ELSE                 ! Decrease lambda, check convergence, continue iteration

      ! 1.8.4 Decrease lambda

      lambda = 0.1_wp *lambda

      ! 1.8.5 Keep track of current state

      n_iter = n_iter + 1

      IF (config % conv_check_apply) THEN

        i_pointer = MOD(i_pointer + 1, config % conv_check_n_previous)
        IF (i_pointer == 0) i_pointer = config % conv_check_n_previous

        delta_J(i_pointer)     = J_last - J
        delta_state(i_pointer) = MAXVAL(ABS(state_last - x%state)/state_sigma)

        state_last = x%state

        J_last     = J

      END IF

      IF (J < J_min) THEN
        J_min  = J
      END IF

      ! 1.8.6 Check for convergence

      IF (config % conv_check_apply) THEN

        IF (config%minropp%impres == 0) THEN
          WRITE(it_str, '(i4)') n_iter
          WRITE(ch_str, '(g15.5)') J
          ch_str = ADJUSTL(ch_str)
          IF (n_iter > 0) THEN
            WRITE(co_str, '(g15.5)') delta_state(i_pointer)
            co_str = ADJUSTL(co_str)
          ELSE
            co_str = ' -             '
          END IF

          CALL message(msg_cont, &
               '   n_iter = ' // it_str // '   J = ' // ch_str //  &
               '   max(relative change in state) = ' // co_str)
        END IF

        IF (MAXVAL(delta_state) < config%conv_check_max_delta_state    &
            .AND. n_iter > config % conv_check_n_previous) THEN

          WRITE(ch_str, '(g15.5)') config%conv_check_max_delta_state
          ch_str = ADJUSTL(ch_str)
          WRITE(it_str, '(i2)')    config%conv_check_n_previous
          it_str = ADJUSTL(it_str)
          CALL message(msg_cont, '')
          CALL message(msg_info,                                             &
            'Convergence assumed to be achieved as the state vector did ' // &
            'not change by more\n   ' // 'than ' // TRIM(ch_str) // ' ' //   &
            'relative to the assumed background errors for the last ' //     &
             TRIM(it_str) // ' iterations.\n')
          EXIT

        ELSE IF (MAXVAL(ABS(delta_J)) < config%conv_check_max_delta_J   &
                 .AND. n_iter > config % conv_check_n_previous) THEN

          WRITE(ch_str, '(g15.5)') config%conv_check_max_delta_J
          ch_str = ADJUSTL(ch_str)
          WRITE(it_str, '(i2)')    config%conv_check_n_previous
          it_str = ADJUSTL(it_str)
          CALL message(msg_cont, '')
          CALL message(msg_info,                                              &
            'Convergence assumed to be achieved as the cost function did ' // &
            'not change by more\n   ' // 'than ' // TRIM(ch_str) //           &
            ' for the last ' // TRIM(it_str) // ' iterations.\n')
          EXIT

        ELSE IF (ABS(lambda) > lambda_max) THEN

          WRITE(ch_str, '(g15.5)') lambda_max
          ch_str = ADJUSTL(ch_str)
          CALL message(msg_cont, '')
          CALL message(msg_info,                                               &
            'Levenberg-Marquardt parameter exceeded max (=' // TRIM(ch_str) // &
            '); convergence is unlikely.\n   ' //                              &
            'Check convergence parameters; conv_check_max_delta_J might ' //   &
            'be too demanding.\n')
           EXIT

        END IF

      END IF    ! Convergence check

    END IF    ! Levenberg-Marquardt update

  END DO    ! End main iteration loop

  IF (n_iter < n_iter_max) THEN
    WRITE(it_str, '(i4)') n_iter  ;  it_str = ADJUSTL(it_str)
    WRITE(gr_str, '(i4)') n_grad  ;  gr_str = ADJUSTL(gr_str)
    CALL message(msg_info, 'Finished after ' // TRIM(it_str) // ' iterations (' // &
         TRIM(gr_str) // ' forward model / gradient evaluations).\n')
  ELSE
    WRITE(it_str, '(i4)') n_iter_max  ;  it_str = ADJUSTL(it_str)
    CALL message(msg_warn, &
      'Iteration ended after n_iter_max = ' // TRIM(it_str) // &
      ' iterations without achieving convergence.\n')
  END IF

  ! 1.9 Copy solution back to state
  ! -------------------------------

  state = x

  ! 1.10 Diagnostic data
  ! --------------------

  diag%J = J_min

  IF (COUNT(obs%weights > 0.0_wp) > 0) THEN
    diag%J_scaled = 2.0_wp * J_min / REAL(COUNT(obs%weights > 0.0_wp), wp)
  ENDIF

  diag%n_iter = n_iter

  ALLOCATE (diag%J_bgr(SIZE(state%state)))
  delta_x = state%state - bg%state
  diag%J_bgr = 0.5_wp * delta_x * matrix_solve(bg%cov, delta_x)

  ! 1.11 Clean up
  ! -------------

  DEALLOCATE(K)
  DEALLOCATE(delta_x)
  DEALLOCATE(delta_y)
  DEALLOCATE(dJ_dx)
  DEALLOCATE(diag_d2J)
  DEALLOCATE(d2J_dx2)
  DEALLOCATE(KO)
  DEALLOCATE(Bm1)
  DEALLOCATE(Om1)

  CALL message_set_routine(routine)

END SUBROUTINE ropp_1dvar_levmarq_bangle

!***

!-------------------------------------------------------------------------------
! 2. Refractivity
!-------------------------------------------------------------------------------

SUBROUTINE ropp_1dvar_levmarq_refrac(obs, bg, state, config, diag)
  
! 2.1 Declarations
! ----------------

  USE typesizes, ONLY: wp => EightByteReal
  USE ropp_utils
  USE ropp_fm
  USE ropp_1dvar, not_this => ropp_1dvar_levmarq_refrac
  USE matrix

  IMPLICIT NONE

  TYPE(Obs1dRefrac), INTENT(inout)      :: obs          ! Observation data
  TYPE(State1dFM),   INTENT(inout)      :: bg           ! Background data
  TYPE(State1dFM),   INTENT(inout)      :: state        ! State vector
  TYPE(VarConfig),   INTENT(in)         :: config       ! Configuration options
  TYPE(VarDiag),     INTENT(inout)      :: diag         ! Diagnostic output

  REAL(wp)                              :: J            ! Cost function value
  TYPE(State1dFM)                       :: x            ! Control vector
  TYPE(State1dFM)                       :: x_old        ! Control for LM parameter testing
  TYPE(Obs1dRefrac)                     :: y            ! Forward model obs
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: K            ! K-matrix
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_x      ! Change of state
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_y      ! Change of observation
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: dJ_dx        ! Cost function gradient
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: diag_d2J     ! Diagonal d2J/dx2
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: d2J_dx2      ! 2nd derivative cost fn
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: KO           ! K O^-1
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: Bm1          ! B^-1
  REAL(wp), DIMENSION(:,:), ALLOCATABLE :: Om1          ! O^-1
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_J      ! Change of cost fn
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: state_last   ! Previous state vector
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: state_sigma  ! State std deviation
  REAL(wp), DIMENSION(:),   ALLOCATABLE :: delta_state  ! Change of state vector
  REAL(wp)                              :: J_last       ! Previous cost function
  REAL(wp)                              :: J_min        ! Minimum cost function
  INTEGER                               :: i_pointer    ! Value index
  INTEGER                               :: i            ! Counter
  INTEGER                               :: n_iter       ! Number of iterations
  INTEGER                               :: n_iter_max   ! Maximum number of iterations
  INTEGER                               :: n_grad       ! Number of gradient updates
  INTEGER                               :: nobs         ! Number of observations
  INTEGER                               :: nstate       ! No. of state elements
  REAL(wp)                              :: lambda       ! Levenberg-Marquardt factor
  REAL(wp)                              :: lambda_max   ! Maximum Levenberg-Marquardt factor

  CHARACTER(len =   4)                  :: it_str
  CHARACTER(len =   4)                  :: gr_str
  CHARACTER(len =  15)                  :: ch_str, co_str
  CHARACTER(len = 256)                  :: routine

! 2.2 Message handling
! --------------------

  CALL message_get_routine(routine)
  CALL message_set_routine('ropp_1dvar_levmarq_refrac')

! 2.3 Initialise variables
! ------------------------

  i_pointer  = 0
  n_grad     = 0
  n_iter     = 0
  n_iter_max = 50

  lambda     = 1.0E-4_wp
  lambda_max = 1.0E10_wp

  J          = 0.0_wp
  J_last     = 1.0E30_wp
  J_min      = J_last

  nstate = SIZE(bg%state)
  nobs   = SIZE(obs%refrac)

  ALLOCATE(K(nobs, nstate))
  ALLOCATE(delta_x(nstate))
  ALLOCATE(delta_y(nobs))
  ALLOCATE(dJ_dx(nstate))
  ALLOCATE(diag_d2J(nstate))
  ALLOCATE(d2J_dx2(nstate,nstate))
  ALLOCATE(KO(nstate,nobs))
  ALLOCATE(Bm1(nstate,nstate))
  ALLOCATE(Om1(nobs,nobs))

  IF (ALLOCATED(delta_J)) DEALLOCATE(delta_J)
  ALLOCATE(delta_J(config%conv_check_n_previous))
  delta_J(:) = 0.0_wp

  IF (ALLOCATED(state_last)) DEALLOCATE(state_last)
  ALLOCATE(state_last(SIZE(bg%state)))
  state_last(:) = 0.0_wp

  IF (ALLOCATED(state_sigma)) DEALLOCATE(state_sigma)
  ALLOCATE(state_sigma(SIZE(bg%state)))
  DO i = 1, SIZE(state_sigma)
    state_sigma(i) = SQRT(bg%cov%d(i + i*(i-1)/2)) ! Direct read from matrix_pp
  ENDDO

  IF (ALLOCATED(delta_state)) DEALLOCATE(delta_state)
  ALLOCATE(delta_state(config%conv_check_n_previous))
  delta_state(:) = 0.0_wp

  ! 2.4 Inverse error covariances
  ! -----------------------------

  Bm1 = matrix_invert(bg%cov)
  Om1 = matrix_invert(obs%cov)

  ! 2.5 First guess and pseudo-observations
  ! ---------------------------------------

  x = bg
  y = obs

  ! 2.6 Initial cost function
  ! -------------------------

  IF (ASSOCIATED(x%ak)) THEN
    CALL ropp_fm_state2state_ecmwf(x)
  ELSE
    CALL ropp_fm_state2state_meto(x)
  END IF

  IF (x%new_ref_op) THEN
    CALL ropp_fm_refrac_1d_new(x, y)
  ELSE
    CALL ropp_fm_refrac_1d(x, y)
  END IF

  delta_x = (x%state - bg%state)
  delta_y = (y%refrac - obs%refrac) * y%weights

  J = 0.5_wp * DOT_PRODUCT(delta_y, MATMUL(Om1, delta_y)) + &
      0.5_wp * DOT_PRODUCT(delta_x, MATMUL(Bm1, delta_x))

  diag%J_init = J
  J_last = J
  J_min  = J

  IF (config%minropp%impres == 0) THEN
    WRITE(it_str, '(i4)') n_iter
    WRITE(ch_str, '(g15.5)') J
    ch_str = ADJUSTL(ch_str)
    co_str = ' -             '

    CALL message(msg_cont, &
         '   n_iter = ' // it_str // '   J = ' // ch_str //  &
         '   max(relative change in state) = ' // co_str)
  END IF

  ! 2.7 Main minimisation iteration loop
  ! ------------------------------------

  DO WHILE (n_iter < n_iter_max)

    ! 2.7.1 Bookkeeping

    n_grad = n_grad + 1

    ! 2.7.2 Compute gradient and Hessian of cost function

    CALL ropp_fm_refrac_1d_grad(x, y, K)

    !TODO: (a) Do we need this?
    !      (b) Does it change the cost function?
    WHERE (ABS(delta_y) > 50.0_wp)
      delta_y = 0.0_wp
    END WHERE

    KO      = MATMUL(TRANSPOSE(K), Om1)
    dJ_dx   = MATMUL(KO, delta_y) + MATMUL(Bm1, delta_x)
    d2J_dx2 = Bm1 + MATMUL(KO, K)

    ! 2.7.3 Levenberg-Marquardt adjustment of Hessian diagonal

    DO i = 1, nstate
      d2J_dx2(i,i) = d2J_dx2(i,i) * (1.0_wp + lambda)
    END DO

    ! 2.7.4 Solve matrix equation d2J_dx2 . dx = -dJ_dx

    delta_x = matrix_solve(d2J_dx2, - dJ_dx)

    ! 2.7.5 Update test state vector

    x_old = x

    IF (x%use_logq) THEN
      x%state = x%state + SIGN(MIN(ABS(delta_x), ABS(x%state/2.0_wp)), delta_x)
    ELSE
      x%state = x%state + delta_x
    END IF

    ! 2.7.6 Compute cost function

    IF (ASSOCIATED(x%ak)) THEN
       CALL ropp_fm_state2state_ecmwf(x)
    ELSE
       CALL ropp_fm_state2state_meto(x)
    END IF

    IF (x%new_ref_op) THEN
      CALL ropp_fm_refrac_1d_new(x, y)
    ELSE
      CALL ropp_fm_refrac_1d(x, y)
    END IF

    delta_x = (x%state - bg%state)
    delta_y = (y%refrac - obs%refrac) * y%weights

    J  = 0.5_wp * DOT_PRODUCT(delta_y, MATMUL(Om1, delta_y)) + &
         0.5_wp * DOT_PRODUCT(delta_x, MATMUL(Bm1, delta_x))

    ! 2.8 Levenberg-Marquardt update and convergence check
    ! ----------------------------------------------------

    IF ( J >= (J_last + config%conv_check_max_delta_J) ) THEN ! Keep previous state vector, increase lambda, repeat

      ! 2.8.1 Keep old state, increase lambda, bail out if it's too big

      x      = x_old
      lambda = 10.0_wp * lambda

      IF (ABS(lambda) > lambda_max) THEN

        WRITE(ch_str, '(g15.5)') lambda_max
        ch_str = ADJUSTL(ch_str)
        CALL message(msg_cont, '')
        CALL message(msg_info,                                               &
          'Levenberg-Marquardt parameter exceeded max (=' // TRIM(ch_str) // &
          '); convergence is unlikely.\n   ' //                              &
          'Check convergence parameters; conv_check_max_delta_J might ' //   &
          'be too demanding.\n')
         EXIT
      END IF

      ! 2.8.2 Logging

      !TODO: Add a configuration option for lambda update outputs
      IF (config%minropp%impres == 0) THEN

        WRITE(it_str, '(i4)') n_iter
        WRITE(ch_str, '(g15.5)') J
        ch_str = ADJUSTL(ch_str)
        IF (n_iter > 0) THEN
          WRITE(co_str, '(g15.5)') lambda
          co_str = ADJUSTL(co_str)
        ELSE
          co_str = ' -             '
        END IF

        CALL message(msg_cont, &
                     '   n_iter = ' // it_str // '   J = ' // ch_str //  &
                     '   lambda -> ' // co_str)

      END IF

      ! 2.8.3 Next iteration

      CYCLE

    ELSE                 ! Decrease lambda, check convergence, continue iteration

      ! 2.8.4 Decrease lambda

      lambda = 0.1_wp *lambda

      ! 2.8.5 Keep track of current state

      n_iter = n_iter + 1

      IF (config % conv_check_apply) THEN

        i_pointer = MOD(i_pointer + 1, config % conv_check_n_previous)
        IF (i_pointer == 0) i_pointer = config % conv_check_n_previous

        delta_J(i_pointer)     = J_last - J
        delta_state(i_pointer) = MAXVAL(ABS(state_last - x%state)/state_sigma)

        state_last = x%state

        J_last     = J

      END IF

      IF (J < J_min) THEN
        J_min  = J
      END IF

      ! 2.8.6 Check for convergence

      IF (config % conv_check_apply) THEN

        IF (config%minropp%impres == 0) THEN
          WRITE(it_str, '(i4)') n_iter
          WRITE(ch_str, '(g15.5)') J
          ch_str = ADJUSTL(ch_str)
          IF (n_iter > 0) THEN
            WRITE(co_str, '(g15.5)') delta_state(i_pointer)
            co_str = ADJUSTL(co_str)
          ELSE
            co_str = ' -             '
          END IF

          CALL message(msg_cont, &
               '   n_iter = ' // it_str // '   J = ' // ch_str //  &
               '   max(relative change in state) = ' // co_str)
        END IF

        IF (MAXVAL(delta_state) < config%conv_check_max_delta_state    &
            .AND. n_iter > config % conv_check_n_previous) THEN

          WRITE(ch_str, '(g15.5)') config%conv_check_max_delta_state
          ch_str = ADJUSTL(ch_str)
          WRITE(it_str, '(i2)')    config%conv_check_n_previous
          it_str = ADJUSTL(it_str)
          CALL message(msg_cont, '')
          CALL message(msg_info,                                             &
            'Convergence assumed to be achieved as the state vector did ' // &
            'not change by more\n   ' // 'than ' // TRIM(ch_str) // ' ' //   &
            'relative to the assumed background errors for the last ' //     &
             TRIM(it_str) // ' iterations.\n')
          EXIT

        ELSE IF (MAXVAL(ABS(delta_J)) < config%conv_check_max_delta_J   &
                 .AND. n_iter > config % conv_check_n_previous) THEN

          WRITE(ch_str, '(g15.5)') config%conv_check_max_delta_J
          ch_str = ADJUSTL(ch_str)
          WRITE(it_str, '(i2)')    config%conv_check_n_previous
          it_str = ADJUSTL(it_str)
          CALL message(msg_cont, '')
          CALL message(msg_info,                                              &
            'Convergence assumed to be achieved as the cost function did ' // &
            'not change by more\n   ' // 'than ' // TRIM(ch_str) //           &
            ' for the last ' // TRIM(it_str) // ' iterations.\n')
          EXIT

        ELSE IF (ABS(lambda) > lambda_max) THEN

          WRITE(ch_str, '(g15.5)') lambda_max
          ch_str = ADJUSTL(ch_str)
          CALL message(msg_cont, '')
          CALL message(msg_info,                                               &
            'Levenberg-Marquardt parameter exceeded max (=' // TRIM(ch_str) // &
            '); convergence is unlikely.\n   ' //                              &
            'Check convergence parameters; conv_check_max_delta_J might ' //   &
            'be too demanding.\n')
           EXIT

        END IF

      END IF    ! Convergence check

    END IF    ! Levenberg-Marquardt update

  END DO    ! End main iteration loop

  IF (n_iter < n_iter_max) THEN
    WRITE(it_str, '(i4)') n_iter  ;  it_str = ADJUSTL(it_str)
    WRITE(gr_str, '(i4)') n_grad  ;  gr_str = ADJUSTL(gr_str)
    CALL message(msg_info, 'Finished after ' // TRIM(it_str) // ' iterations (' // &
         TRIM(gr_str) // ' forward model / gradient evaluations).\n')
  ELSE
    WRITE(it_str, '(i4)') n_iter_max  ;  it_str = ADJUSTL(it_str)
    CALL message(msg_warn, &
      'Iteration ended after n_iter_max = ' // TRIM(it_str) // &
      ' iterations without achieving convergence.\n')
  END IF

  ! 2.9 Copy solution back to state
  ! -------------------------------

  state = x

  ! 2.10 Diagnostic data
  ! --------------------

  diag%J = J_min

  IF (COUNT(obs%weights > 0.0_wp) > 0) THEN
    diag%J_scaled = 2.0_wp * J_min / REAL(COUNT(obs%weights > 0.0_wp), wp)
  ENDIF

  diag%n_iter = n_iter

  ALLOCATE (diag%J_bgr(SIZE(state%state)))
  delta_x = state%state - bg%state
  diag%J_bgr = 0.5_wp * delta_x * matrix_solve(bg%cov, delta_x)

  ! 2.11 Clean up
  ! -------------

  DEALLOCATE(K)
  DEALLOCATE(delta_x)
  DEALLOCATE(delta_y)
  DEALLOCATE(dJ_dx)
  DEALLOCATE(diag_d2J)
  DEALLOCATE(d2J_dx2)
  DEALLOCATE(KO)
  DEALLOCATE(Bm1)
  DEALLOCATE(Om1)

  CALL message_set_routine(routine)

END SUBROUTINE ropp_1dvar_levmarq_refrac