open Qptypes;; open Qputils;; open Core.Std;; module Nuclei : sig type t = { nucl_num : Nucl_number.t ; nucl_label : Element.t array; nucl_charge : Charge.t array; nucl_coord : Point3d.t array; } with sexp ;; val read : unit -> t val write : t -> unit val to_string : t -> string val to_rst : t -> Rst_string.t val of_rst : Rst_string.t -> t end = struct type t = { nucl_num : Nucl_number.t ; nucl_label : Element.t array; nucl_charge : Charge.t array; nucl_coord : Point3d.t array; } with sexp ;; let get_default = Qpackage.get_ezfio_default "nuclei";; let read_nucl_num () = let nmax = Nucl_number.get_max () in Nucl_number.of_int ~max:nmax nmax ;; let write_nucl_num n = Nucl_number.to_int n |> Ezfio.set_nuclei_nucl_num ;; let read_nucl_label () = Ezfio.get_nuclei_nucl_label () |> Ezfio.flattened_ezfio |> Array.map ~f:Element.of_string ;; let write_nucl_label ~nucl_num labels = let nucl_num = Nucl_number.to_int nucl_num in let labels = Array.to_list labels |> List.map ~f:Element.to_string in Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| nucl_num |] ~data:labels |> Ezfio.set_nuclei_nucl_label ;; let read_nucl_charge () = Ezfio.get_nuclei_nucl_charge () |> Ezfio.flattened_ezfio |> Array.map ~f:Charge.of_float ;; let write_nucl_charge ~nucl_num charges = let nucl_num = Nucl_number.to_int nucl_num in let charges = Array.to_list charges |> List.map ~f:Charge.to_float in Ezfio.ezfio_array_of_list ~rank:1 ~dim:[| nucl_num |] ~data:charges |> Ezfio.set_nuclei_nucl_charge ;; let read_nucl_coord () = let nucl_num = Nucl_number.to_int (read_nucl_num ()) in let raw_data = Ezfio.get_nuclei_nucl_coord() |> Ezfio.flattened_ezfio in let zero = Point3d.of_string Units.Bohr "0. 0. 0." in let result = Array.create nucl_num zero in for i=0 to (nucl_num-1) do result.(i) <- Point3d.({ x=raw_data.(i); y=raw_data.(nucl_num+i); z=raw_data.(2*nucl_num+i); }); done; result ;; let write_nucl_coord ~nucl_num coord = let nucl_num = Nucl_number.to_int nucl_num in let coord = Array.to_list coord in let coord = (List.map ~f:(fun x-> x.Point3d.x) coord) @ (List.map ~f:(fun x-> x.Point3d.y) coord) @ (List.map ~f:(fun x-> x.Point3d.z) coord) in Ezfio.ezfio_array_of_list ~rank:2 ~dim:[| nucl_num ; 3 |] ~data:coord |> Ezfio.set_nuclei_nucl_coord ;; let read () = { nucl_num = read_nucl_num (); nucl_label = read_nucl_label () ; nucl_charge = read_nucl_charge (); nucl_coord = read_nucl_coord (); } ;; let write { nucl_num ; nucl_label ; nucl_charge ; nucl_coord ; } = write_nucl_num nucl_num ; write_nucl_label ~nucl_num:nucl_num nucl_label; write_nucl_charge ~nucl_num:nucl_num nucl_charge; write_nucl_coord ~nucl_num:nucl_num nucl_coord; ;; let to_string b = Printf.sprintf " nucl_num = %s nucl_label = %s nucl_charge = %s nucl_coord = %s " (Nucl_number.to_string b.nucl_num) (b.nucl_label |> Array.to_list |> List.map ~f:(Element.to_string) |> String.concat ~sep:", " ) (b.nucl_charge |> Array.to_list |> List.map ~f:(Charge.to_string) |> String.concat ~sep:", " ) (b.nucl_coord |> Array.to_list |> List.map ~f:(Point3d.to_string Units.Bohr) |> String.concat ~sep:"\n" ) ;; let to_rst b = let nucl_num = Nucl_number.to_int b.nucl_num in let text = ( Printf.sprintf " %d\n " nucl_num ) :: ( List.init nucl_num ~f:(fun i-> Printf.sprintf " %-3s %d %s" (b.nucl_label.(i) |> Element.to_string) (b.nucl_charge.(i) |> Charge.to_int ) (b.nucl_coord.(i) |> Point3d.to_string Units.Angstrom) ) ) |> String.concat ~sep:"\n" in Printf.sprintf " Nuclear coordinates in xyz format (Angstroms) :: %s " text |> Rst_string.of_string ;; let of_rst s = let l = Rst_string.to_string s |> String.split ~on:'\n' in (* Find lines containing the xyz data *) let rec extract_begin = function | [] -> raise Not_found | line::tail -> let line = String.strip line in if (String.length line > 3) && (String.sub line ~pos:((String.length line)-2) ~len:2 = "::") then tail else extract_begin tail in (* Create a list of Atom.t *) let nmax = Nucl_number.get_max () in let atom_list = match (extract_begin l) with | _ :: nucl_num :: title :: lines -> begin let nucl_num = nucl_num |> String.strip |> Int.of_string |> Nucl_number.of_int ~max:nmax and lines = Array.of_list lines in List.init (Nucl_number.to_int nucl_num) ~f:(fun i -> Atom.of_string Units.Angstrom lines.(i)) end | _ -> failwith "Error in xyz format" in (* Create the Nuclei.t data structure *) { nucl_num = List.length atom_list |> Nucl_number.of_int ~max:nmax; nucl_label = List.map atom_list ~f:(fun x -> x.Atom.element) |> Array.of_list ; nucl_charge = List.map atom_list ~f:(fun x -> x.Atom.charge ) |> Array.of_list ; nucl_coord = List.map atom_list ~f:(fun x -> x.Atom.coord ) |> Array.of_list ; } ;; end