9
1
mirror of https://github.com/QuantumPackage/qp2.git synced 2024-12-30 15:15:38 +01:00
qp2/scripts/ezfio_interface/qp_edit_template
2019-01-25 11:39:31 +01:00

376 lines
9.3 KiB
OCaml

(*
vim::syntax=ocaml
*)
open Qputils
open Qptypes
open Core
(** Interactive editing of the input.
WARNING
This file is automatically generated by
`${{QP_ROOT}}/scripts/ezfio_interface/ei_handler.py`
*)
(** Keywords used to define input sections *)
type keyword =
| Ao_basis
| Determinants_by_hand
| Electrons
| Mo_basis
| Nuclei_by_hand
{keywords}
let keyword_to_string = function
| Ao_basis -> "AO basis"
| Determinants_by_hand -> "Determinants_by_hand"
| Electrons -> "Electrons"
| Mo_basis -> "MO basis"
| Nuclei_by_hand -> "Molecule"
{keywords_to_string}
(** Create the header of the temporary file *)
let file_header filename =
Printf.sprintf "
==================================================================
Quantum Package
==================================================================
Editing file `%s`
" filename
(** Creates the header of a section *)
let make_header kw =
let s = keyword_to_string kw in
let l = String.length s in
"\n\n"^s^"\n"^(String.init l ~f:(fun _ -> '='))^"\n\n"
(** Returns the rst string of section [s] *)
let get s =
let header = (make_header s) in
let f (read,to_rst) =
match read () with
| Some text -> header ^ (Rst_string.to_string (to_rst text))
| None -> ""
in
let rst =
try
begin
let open Input in
match s with
| Mo_basis ->
f Mo_basis.(read, to_rst)
| Electrons ->
f Electrons.(read, to_rst)
| Nuclei_by_hand ->
f Nuclei_by_hand.(read, to_rst)
| Ao_basis ->
f Ao_basis.(read, to_rst)
| Determinants_by_hand ->
f Determinants_by_hand.(read_maybe, to_rst)
{section_to_rst}
end
with
| Sys_error msg -> (Printf.eprintf "Info: %s\n%!" msg ; "")
in
rst
(** Applies the changes from the string [str] corresponding to section [s] *)
let set str s =
let header = (make_header s) in
match String.substr_index ~pos:0 ~pattern:header str with
| None -> ()
| Some idx ->
begin
let index_begin = idx + (String.length header) in
let index_end =
match ( String.substr_index ~pos:(index_begin+(String.length header)+1)
~pattern:"==" str) with
| Some i -> i
| None -> String.length str
in
let l = index_end - index_begin in
let str = String.sub ~pos:index_begin ~len:l str
|> Rst_string.of_string
in
let write (of_rst,w) s =
try
match of_rst str with
| Some data -> w data
| None -> ()
with
| _ -> (Printf.eprintf "Info: Read error in %s\n%!"
(keyword_to_string s); ignore (of_rst str) )
in
let open Input in
match s with
{write}
| Electrons -> write Electrons.(of_rst, write) s
| Determinants_by_hand -> write Determinants_by_hand.(of_rst, write) s
| Nuclei_by_hand -> write Nuclei_by_hand.(of_rst, write) s
| Ao_basis -> () (* TODO *)
| Mo_basis -> () (* TODO *)
end
(** Creates the temporary file for interactive editing *)
let create_temp_file ezfio_filename fields =
let temp_filename = Filename.temp_file "qp_edit_" ".rst" in
begin
Out_channel.with_file temp_filename ~f:(fun out_channel ->
(file_header ezfio_filename) :: (List.map ~f:get fields)
|> String.concat ~sep:"\n"
|> Out_channel.output_string out_channel
);
at_exit (fun () -> Sys.remove temp_filename);
temp_filename
end
let run check_only ?ndet ?state ezfio_filename =
(* Set check_only if the arguments are not empty *)
let check_only =
match ndet, state with
| None, None -> check_only
| _ -> true
in
(* Open EZFIO *)
if (not (Sys.file_exists_exn ezfio_filename)) then
failwith (ezfio_filename^" does not exists");
Ezfio.set_file ezfio_filename;
(* Clean qp_stop status *)
[ "qpstop" ; "qpkill" ]
|> List.iter ~f:(fun f ->
let stopfile =
Filename.concat (Qpackage.ezfio_work ezfio_filename) f
in
match Sys.file_exists stopfile with
| `Yes -> Sys.remove stopfile
| _ -> ()
);
(* Reorder basis set *)
begin
let aos =
match Input.Ao_basis.read() with
| Some x -> x
| _ -> assert false
in
let ordering = Input.Ao_basis.ordering aos in
let test = Array.copy ordering in
Array.sort ~compare test ;
if test <> ordering then
begin
Printf.eprintf "Warning: Basis set is not properly ordered. Redordering.\n";
let new_aos = Input.Ao_basis.reorder aos in
Input.Ao_basis.write new_aos;
match Input.Mo_basis.read() with
| None -> ()
| Some mos ->
let new_mos = Input.Mo_basis.reorder mos ordering in
Input.Mo_basis.write new_mos
end
end;
begin
match ndet with
| None -> ()
| Some n -> Input.Determinants_by_hand.update_ndet (Det_number.of_int n)
end;
begin
match state with
| None -> ()
| Some range ->
begin
Input.Determinants_by_hand.extract_states range
end
end;
(*
let output = (file_header ezfio_filename) :: (
List.map ~f:get [
Ao_basis ;
Mo_basis ;
])
in
String.concat output
|> print_string
*)
let tasks = [
Nuclei_by_hand ;
Ao_basis;
Electrons ;
{tasks}
Mo_basis;
Determinants_by_hand ;
]
in
(* Create the temp file *)
let temp_filename =
create_temp_file ezfio_filename tasks
in
(* Open the temp file with external editor *)
let editor =
match Sys.getenv "EDITOR" with
| Some editor -> editor
| None -> "vi"
in
match check_only with
| true -> ()
| false ->
Printf.sprintf "%s %s" editor temp_filename
|> Sys.command_exn
;
(* Re-read the temp file *)
let temp_string =
In_channel.with_file temp_filename ~f:(fun in_channel ->
In_channel.input_all in_channel)
in
List.iter ~f:(fun x -> set temp_string x) tasks
(** Remove the backup file *)
let remove_backup ezfio_filename =
let backup_filename =
Printf.sprintf "%s/work/backup.tar" ezfio_filename
in
try Sys.remove backup_filename
with _ -> ()
(** Create a backup file in case of an exception *)
let create_backup ezfio_filename =
remove_backup ezfio_filename;
let backup_filename =
Printf.sprintf "%s/work/backup.tar" ezfio_filename
in
try
Printf.sprintf "
mkdir -p %s/work ;
tar -cf .backup.tar --exclude=\"work/*\" %s && (mv .backup.tar %s || rm .backup.tar)
"
ezfio_filename ezfio_filename backup_filename
|> Sys.command_exn
with _ -> ()
(** Restore the backup file when an exception occuprs *)
let restore_backup ezfio_filename =
let filename =
Printf.sprintf "%s/work/backup.tar" ezfio_filename
in
if Sys.file_exists_exn filename then
begin
Printf.sprintf "tar -xf %s" filename
|> Sys.command_exn;
remove_backup ezfio_filename
end
let () =
let open Command_line in
begin
set_header_doc (Sys.argv.(0) ^ " - quantum_package command");
set_description_doc
"Opens a text editor to edit the parameters of an EZFIO directory.";
[ {{ short='c'; long="check"; opt=Optional;
doc="Checks the input data";
arg=Without_arg; }};
{{ short='n'; long="ndet"; opt=Optional;
doc="Truncates the wavefunction to the target number of determinants";
arg=With_arg "<int>"; }};
{{
short='s'; long="state"; opt=Optional;
doc="Extracts selected states, for example \"[1,3-5]\"";
arg=With_arg "<range>"; }};
anonymous "EZFIO_DIR" Mandatory "EZFIO directory";
]
|> set_specs ;
end;
(* Handle options *)
let ndet =
match Command_line.get "ndet" with
| None -> None
| Some s -> (try Some (int_of_string s)
with _ -> failwith "[-n|--ndet] expects an integer")
in
let state =
match Command_line.get "state" with
| None -> None
| Some s -> (try Some (Range.of_string s)
with _ -> failwith "[-s|--state] expects a range")
in
let c = Command_line.get_bool "check" in
let ezfio_filename =
match Command_line.anon_args () with
| [x] -> x
| _ -> (Command_line.help () ; failwith "EZFIO_DIR is missing")
in
at_exit (fun () -> remove_backup ezfio_filename);
(* Run the program *)
try
if (not c) then create_backup ezfio_filename;
run c ?ndet ?state ezfio_filename
with
| Failure exc
| Invalid_argument exc ->
begin
Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
Printf.eprintf "%s\n\n" exc;
Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
restore_backup ezfio_filename;
ignore @@ exit 1
end
| Assert_failure (file, line, ch) ->
begin
Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
Printf.eprintf "Assert error in file $QP_ROOT/ocaml/%s, line %d, character %d\n\n" file line ch;
Printf.eprintf "=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-\n\n";
restore_backup ezfio_filename;
ignore @@ exit 2
end;
exit 0