quantum_package/ocaml/Pseudo.ml

open Core.Std
open Qptypes


module Primitive_local : sig

  type t = { 
    expo    : AO_expo.t ;
    r_power : R_power.t ;
  } with sexp

  val of_expo_r_power : AO_expo.t -> R_power.t -> t
  val to_string : t -> string
  
end = struct

  type t = { 
    expo    : AO_expo.t ;
    r_power : R_power.t ;
  } with sexp

  let of_expo_r_power dz n = 
     { expo = dz ; r_power = n }

  let to_string p =
     Printf.sprintf "(%d, %f)"
     (R_power.to_int p.r_power)
     (AO_expo.to_float p.expo)
end


module Primitive_non_local : sig

  type t = { 
    expo    : AO_expo.t ;
    r_power : R_power.t ;
    proj    : Positive_int.t
  } with sexp

  val of_proj_expo_r_power : Positive_int.t -> AO_expo.t -> R_power.t -> t
  val to_string : t -> string
  
end = struct

  type t = { 
    expo    : AO_expo.t ;
    r_power : R_power.t ;
    proj    : Positive_int.t
  } with sexp

  let of_proj_expo_r_power p dz n = 
     { expo = dz ; r_power = n ; proj = p }

  let to_string p =
     Printf.sprintf "(%d, %f, %d)" 
       (R_power.to_int p.r_power)
       (AO_expo.to_float p.expo)
       (Positive_int.to_int p.proj)
end




type t = {
  element   : Element.t ;
  n_elec    : Positive_int.t ;
  local     : (Primitive_local.t * AO_coef.t ) list ;
  non_local : (Primitive_non_local.t * AO_coef.t ) list 
} with sexp

let empty e =
  { element   = e;
    n_elec    = Positive_int.of_int 0;
    local     = [];
    non_local = [];
  }

(** Transform the local component of the pseudopotential to a string *)
let to_string_local = function
| [] -> ""
| t  ->
  "Local component:" ::
  ( Printf.sprintf "%20s %8s %20s" "Coeff." "r^n" "Exp." ) ::
  ( List.map t ~f:(fun (l,c) -> Printf.sprintf "%20f %8d %20f"
      (AO_coef.to_float c) 
      (R_power.to_int   l.Primitive_local.r_power)
      (AO_expo.to_float l.Primitive_local.expo)
  ) )
  |> String.concat ~sep:"\n" 


(** Transform the non-local component of the pseudopotential to a string *)
let to_string_non_local = function
| [] -> ""
| t  ->
  "Non-local component:" ::
  ( Printf.sprintf "%20s %8s %20s %8s" "Coeff." "r^n" "Exp." "Proj") ::
  ( List.map t ~f:(fun (l,c) ->
      let p = 
        Positive_int.to_int l.Primitive_non_local.proj
      in
      Printf.sprintf "%20f %8d %20f   |%d><%d|"
      (AO_coef.to_float c) 
      (R_power.to_int   l.Primitive_non_local.r_power)
      (AO_expo.to_float l.Primitive_non_local.expo)
      p p
  ) )
  |> String.concat ~sep:"\n" 

(** Transform the Pseudopotential to a string *)
let to_string t =

  Printf.sprintf "%s   %d electrons removed"
    (Element.to_string t.element)
    (Positive_int.to_int t.n_elec)
  :: to_string_local t.local
  :: to_string_non_local t.non_local
  :: []
  |> List.filter ~f:(fun x -> x <> "")
  |> String.concat ~sep:"\n"
  
  
(** Find an element in the file *)
let find in_channel element =
  In_channel.seek in_channel 0L;
  
  let element_read, old_pos =
     ref Element.X, 
     ref (In_channel.pos in_channel)
  in
  while !element_read <> element
  do
    let buffer = 
      old_pos := In_channel.pos in_channel;
      match In_channel.input_line in_channel with
      | Some line -> String.split ~on:' ' line
         |> List.hd_exn
      | None -> ""
    in
    try
      element_read := Element.of_string buffer
    with
    | Element.ElementError _ -> ()
  done ;
  In_channel.seek in_channel !old_pos;
  !element_read


(** Read the Pseudopotential in GAMESS format *)
let read_element in_channel element =
  ignore (find in_channel element);

  let rec read result =
    match In_channel.input_line in_channel with
    | None -> result
    | Some line ->
      if (String.strip line = "") then
        result
      else
        read (line::result)
  in

  let data = 
    read []
    |> List.rev
  in

  let debug_data = 
    String.concat ~sep:"\n" data
  in

  let decode_first_line = function
  | first_line :: rest ->
    begin
      let first_line_split = 
        String.split first_line ~on:' '
        |> List.filter ~f:(fun x -> (String.strip x) <> "")
      in
      match first_line_split with
      | e :: "GEN" :: n :: p ->
        {  element = Element.of_string e ;
           n_elec  = Int.of_string n |> Positive_int.of_int ;
           local = [] ;
           non_local = []
        }, rest
      | _ -> failwith (
        Printf.sprintf "Unable to read Pseudopotential : \n%s\n"
        debug_data )
    end
  | _ -> failwith ("Error reading pseudopotential\n"^debug_data)
  in  

  let rec loop create_primitive accu = function
  | (0,rest) -> List.rev accu, rest
  | (n,line::rest) -> 
    begin
      match
        String.split line ~on:' '
        |> List.filter ~f:(fun x -> String.strip x <> "")
      with
      | c :: i :: e :: [] ->
        let i =
          Int.of_string i
        in
        let elem =
          ( create_primitive 
            (Float.of_string e |> AO_expo.of_float)
            (i-2 |> R_power.of_int),
            Float.of_string c |> AO_coef.of_float 
          )
        in
        loop create_primitive (elem::accu) (n-1, rest)
      | _ -> failwith ("Error reading pseudopotential\n"^debug_data)
    end
  | _ -> failwith ("Error reading pseudopotential\n"^debug_data)
  in

  let decode_local (pseudo,data) = 
    let decode_local_n n rest =
      let result, rest =
        loop Primitive_local.of_expo_r_power [] (Positive_int.to_int n,rest)
      in
      { pseudo with local = result }, rest
    in
    match data with
    | n :: rest ->
        let n =
          String.strip n
          |> Int.of_string
          |> Positive_int.of_int
        in
        decode_local_n n rest
    | _ -> failwith ("Unable to read (non-)local pseudopotential\n"^debug_data)
  in

  let decode_non_local (pseudo,data) = 
    let decode_non_local_n proj n (pseudo,data) =
      let result, rest = 
        loop (Primitive_non_local.of_proj_expo_r_power proj)
          [] (Positive_int.to_int n, data)
      in
      { pseudo with non_local = pseudo.non_local @ result }, rest
    in
    let rec new_proj (pseudo,data) proj = 
      match data with
      | n :: rest ->
          let n =
            String.strip n
            |> Int.of_string
            |> Positive_int.of_int
          in
          let result = 
            decode_non_local_n proj n (pseudo,rest)
          and proj_next =
            (Positive_int.to_int proj)+1
            |> Positive_int.of_int
          in
          new_proj result proj_next
      | _ -> pseudo
    in
    new_proj (pseudo,data) (Positive_int.of_int 0)
  in

  decode_first_line data
  |> decode_local
  |> decode_non_local 

     
 

include To_md5
let to_md5 = to_md5 sexp_of_t