Corrected bug when writing pseudo to EZFIO

ocaml/Qputils.ml

@@ -1,4 +1,4 @@
-open Core.Std;;
+open Core.Std
 
 (*
 let rec transpose = function
@@ -24,5 +24,21 @@ let input_to_sexp s =
     print_endline ("("^result^")");
     "("^result^")"
     |> Sexp.of_string
-;;
+
+let rmdir dirname =
+  let rec remove_one dir =
+    Sys.chdir dir;
+    Sys.readdir "."
+    |> Array.iter ~f:(fun x ->
+      match (Sys.is_directory x, Sys.is_file x) with
+      | (`Yes, _) -> remove_one x
+      | (_, `Yes) -> Sys.remove x
+      | _ -> failwith ("Unable to remove file "^x^".")
+    );
+    Sys.chdir "..";
+    Unix.rmdir dir
+  in
+  remove_one dirname
+  
+    
 

ocaml/qp_create_ezfio_from_xyz.ml

@@ -1,6 +1,6 @@
-open Qputils;;
+open Qputils
-open Qptypes;;
+open Qptypes
-open Core.Std;;
+open Core.Std
 
 let spec =
   let open Command.Spec in
@@ -316,6 +316,7 @@ let run ?o b c d m p cart xyz_file =
   if Sys.file_exists_exn ezfio_file then
     failwith (ezfio_file^" already exists");
 
+  let write_file () = 
       (* Create EZFIO *)
       Ezfio.set_file ezfio_file;
 
@@ -398,14 +399,6 @@ let run ?o b c d m p cart xyz_file =
           if (x > accu) then x
           else accu
         )
-  and kmax = 
-    List.fold pseudo ~init:0 ~f:(fun accu x -> 
-      let x = 
-        List.length x.Pseudo.non_local
-      in
-      if (x > accu) then x
-      else accu
-    )
       and lmax = 
         List.fold pseudo ~init:0 ~f:(fun accu x -> 
           let x = 
@@ -422,6 +415,20 @@ let run ?o b c d m p cart xyz_file =
         )
       in
 
+     let kmax = 
+         Array.init (lmax+1) ~f:(fun i->
+           List.map pseudo ~f:(fun x -> 
+             List.filter x.Pseudo.non_local ~f:(fun (y,_) -> 
+               (Positive_int.to_int y.Pseudo.Primitive_non_local.proj) = i)
+             |> List.length )
+           |> List.fold ~init:0 ~f:(fun accu x ->
+             if accu > x then accu else x)
+           )
+         |> Array.fold ~init:0 ~f:(fun accu i ->
+            if i > accu then i else accu)
+      in
+ 
+              
       let () = 
         Ezfio.set_pseudo_pseudo_klocmax klocmax;
         Ezfio.set_pseudo_pseudo_kmax kmax;
@@ -458,23 +465,30 @@ let run ?o b c d m p cart xyz_file =
         |> Ezfio.set_pseudo_pseudo_n_k;
 
         let tmp_array_v_kl, tmp_array_dz_kl, tmp_array_n_kl = 
-      Array.create ~len:(lmax+1) 
+          Array.init (lmax+1) ~f:(fun _ ->
-       (Array.make_matrix ~dimx:kmax ~dimy:nucl_num 0. ),
+          (Array.make_matrix ~dimx:kmax ~dimy:nucl_num 0. )),
-      Array.create ~len:(lmax+1) 
+          Array.init (lmax+1) ~f:(fun _ ->
-       (Array.make_matrix ~dimx:kmax ~dimy:nucl_num 0. ),
+          (Array.make_matrix ~dimx:kmax ~dimy:nucl_num 0. )),
-      Array.create ~len:(lmax+1) 
+          Array.init (lmax+1) ~f:(fun _ ->
-       (Array.make_matrix ~dimx:kmax ~dimy:nucl_num 0 )
+          (Array.make_matrix ~dimx:kmax ~dimy:nucl_num 0 ))
         in
         List.iteri pseudo ~f:(fun j x -> 
-      List.iteri x.Pseudo.non_local ~f:(fun i (y,c) ->
+          let last_idx =
+            Array.create ~len:(lmax+1) 0
+          in
+          List.iter x.Pseudo.non_local ~f:(fun (y,c) -> 
             let k, y, z =
               Positive_int.to_int y.Pseudo.Primitive_non_local.proj,
               AO_expo.to_float y.Pseudo.Primitive_non_local.expo,
               R_power.to_int y.Pseudo.Primitive_non_local.r_power
             in
+            let i = 
+              last_idx.(k) 
+            in
             tmp_array_v_kl.(k).(i).(j)  <- AO_coef.to_float c;
             tmp_array_dz_kl.(k).(i).(j) <- y;
             tmp_array_n_kl.(k).(i).(j)  <- z;
+            last_idx.(k) <- i+1;
            )
         );    
         let concat_3d tmp_array = 
@@ -599,6 +613,19 @@ let run ?o b c d m p cart xyz_file =
       match Input.Ao_basis.read () with
       | None -> failwith "Error in basis"
       | Some x -> Input.Ao_basis.write x
+  in
+  let () =
+    try write_file () with
+    | ex -> 
+      begin
+        begin
+          match Sys.is_directory ezfio_file with
+          | `Yes ->  rmdir ezfio_file
+          | _ -> ()
+        end; 
+        raise ex;
+      end
+  in ()