QCaml/Utils/Davidson.ml

open Lacaml.D

type t

let make
    ?guess
    ?(n_states=1)
    ?(n_iter=10)
    ?(threshold=1.e-6)
    diagonal
    matrix_prod
  =

  (* Size of the matrix to diagonalize *)
  let n = Vec.dim diagonal in


  let m = (* Number of requested states *)
    match guess with
    | Some vectors -> (Mat.dim2 vectors) * n_states
    | None -> n_states
  in

  (* Create guess vectors u, with randomly initialized unknown vectors. *)
  let init_vectors =
    let init_vector k =
      Vector.sparse_of_assoc_list n [ (k,1.0) ]
    in
    Array.init m (fun i -> init_vector (i+1)) 
    |> Matrix.sparse_of_vector_array
    |> Matrix.to_mat
  in

  let pick_new u =
    Mat.to_col_vecs_list u
    |> Util.list_pack m
    |> List.rev
    |> List.hd
  in

  let u_new =
   match guess with
   | Some vectors -> Mat.to_col_vecs_list vectors 
   | None -> Mat.to_col_vecs_list init_vectors
  in

  let rec iteration u u_new w iter =
    (* u is a list of orthonormal vectors, on which the operator has
       been applied : w = op.u
       u_new is a list of vector which will increase the size of the
       space.
    *)

    (* Orthonormalize input vectors u_new *)
    let u_new_ortho =
      List.concat [u ; u_new]
      |> Mat.of_col_vecs_list
      |> Util.qr_ortho 
      |> pick_new
    in

    (* Apply the operator the m last vectors *)
    let w_new =
      matrix_prod (
        u_new_ortho
        |> Mat.of_col_vecs_list
        |> Matrix.dense_of_mat )
      |> Matrix.to_mat
      |> Mat.to_col_vecs_list
    in             

    (* Data for the next iteration *)
    let u_next =
      List.concat [ u ; u_new_ortho ]
    and w_next = 
      List.concat [ w ; w_new ]
    in

    (* Build the small matrix h = <U_k | W_l> *)
    let m_U =
      Mat.of_col_vecs_list u_next
    and m_W =
      Mat.of_col_vecs_list w_next
    in
    let m_h =
      gemm ~transa:`T m_U m_W
    in

    (* Diagonalize h *)
    let y, lambda =
      Util.diagonalize_symm m_h
    in

    (* Express m lowest eigenvectors of h in the large basis *)
    let m_new_U =
      gemm  ~n:m  m_U  y
    and m_new_W =
      gemm  ~n:m  m_W  y
    in

    (* Compute the residual as proposed new vectors *)
    let u_proposed = 
      Mat.init_cols n m (fun i k ->
        (lambda.{k} *. m_new_U.{i,k} -. m_new_W.{i,k}) /. 
          (max (diagonal.{i} -. lambda.{k}) 0.01) )
      |> Mat.to_col_vecs_list
    in


    let residual_norms = List.map nrm2 u_proposed in
    let residual_norm  = List.fold_left (fun accu i -> max accu i) 0. residual_norms in

    if Parallel.master then
      Printf.printf "%3d %16.10f  %16.8e%!\n" iter lambda.{1} residual_norm;

    (* Make new vectors sparse *)
    let u_proposed = 
      Mat.of_col_vecs_list u_proposed
    in
    let maxu = lange u_proposed ~norm:`M in
    let thr = maxu *. 0.01 in
    let u_proposed = 
      Mat.map (fun x -> if abs_float x < thr then 0. else x) u_proposed
      |> Mat.to_col_vecs_list
    in


    if residual_norm > threshold then
      let u_next, w_next, iter =
        if iter = n_iter then
          m_new_U |> pick_new, 
          m_new_W |> pick_new,
          0
        else
          u_next, w_next, iter
      in
      iteration u_next u_proposed w_next (iter+1)
    else
      (m_new_U |> pick_new |> Mat.of_col_vecs_list), lambda

  in
  iteration [] u_new [] 1