open Util open Lacaml.D open Constants let max_index = 1 lsl 14 type index_pair = { first : int ; second : int } type array2 = (float, Bigarray.float64_elt, Bigarray.fortran_layout) Bigarray.Array2.t type storage_t = | Dense of array2 | Svd of Vec.t * array2 * array2 | Sparse of (int, float) Hashtbl.t type t = { size : int ; two_index : array2; two_index_anti : array2; three_index : array2; three_index_anti : array2; four_index : storage_t ; } let key_of_indices ~r1 ~r2 = let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in let f i k = let p, r = if i <= k then i, k else k, i in p + (r*(r-1))/2 in let p = f i k and q = f j l in f p q let check_bounds r1 r2 t = let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in let size = t.size in assert ( (i lor j lor k lor l) > 0 ); assert ( i <= size && j <= size && k <= size && l <= size ) let dense_index i j size = (j-1)*size + i let sym_index i j = if i < j then (j*(j-1))/2 + i else (i*(i-1))/2 + j let unsafe_get_four_index ~r1 ~r2 t = let open Bigarray.Array2 in let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in if i=k then if j=l then unsafe_get t.two_index i j else unsafe_get t.three_index (dense_index j l t.size) i else if j=l then unsafe_get t.three_index (dense_index i k t.size) j else if i=l then if k=j then unsafe_get t.two_index_anti i j else unsafe_get t.three_index_anti (dense_index j k t.size) i else if j=k then unsafe_get t.three_index_anti (dense_index i l t.size) j else if i=j then if k=l then unsafe_get t.two_index_anti i k else unsafe_get t.three_index_anti (dense_index k l t.size) i else if k=l then (* *) unsafe_get t.three_index_anti (dense_index i j t.size) k else match t.four_index with | Dense a -> unsafe_get a (dense_index i k t.size) (sym_index j l) | Svd _ -> assert false | Sparse a -> let key = key_of_indices ~r1 ~r2 in try Hashtbl.find a key with Not_found -> 0. let get_four_index ~r1 ~r2 t = check_bounds r1 r2 t; unsafe_get_four_index ~r1 ~r2 t let unsafe_set_four_index ~r1 ~r2 ~value t = let open Bigarray.Array2 in let { first=i ; second=k } = r1 and { first=j ; second=l } = r2 in let () = if i=k then begin if j=l then begin unsafe_set t.two_index i j value; unsafe_set t.two_index j i value; unsafe_set t.three_index (dense_index i i t.size) j value; end; unsafe_set t.three_index (dense_index j l t.size) i value; unsafe_set t.three_index (dense_index l j t.size) i value; end else if j=l then begin unsafe_set t.three_index (dense_index i k t.size) j value; unsafe_set t.three_index (dense_index k i t.size) j value; end else if i=l then begin if j=k then begin unsafe_set t.two_index_anti i j value; unsafe_set t.two_index_anti j i value; unsafe_set t.three_index_anti (dense_index i i t.size) j value; end; unsafe_set t.three_index_anti (dense_index j k t.size) i value; unsafe_set t.three_index_anti (dense_index k j t.size) i value; end else if j=k then begin unsafe_set t.three_index_anti (dense_index i l t.size) j value; unsafe_set t.three_index_anti (dense_index l i t.size) j value; end else if i=j then begin if k=l then begin unsafe_set t.two_index_anti i k value; unsafe_set t.two_index_anti k i value; unsafe_set t.three_index_anti (dense_index i i t.size) k value; end; unsafe_set t.three_index_anti (dense_index k l t.size) i value; unsafe_set t.three_index_anti (dense_index l k t.size) i value; end else if k=l then (* *) begin unsafe_set t.three_index_anti (dense_index i j t.size) k value; unsafe_set t.three_index_anti (dense_index j i t.size) k value; end in match t.four_index with | Dense a -> let ik = (dense_index i k t.size) and jl = (dense_index j l t.size) and ki = (dense_index k i t.size) and lj = (dense_index l j t.size) and ik_s = (sym_index i k) and jl_s = (sym_index j l) in begin unsafe_set a ik jl_s value; unsafe_set a ki jl_s value; unsafe_set a jl ik_s value; unsafe_set a lj ik_s value; end | Sparse a -> let key = key_of_indices ~r1 ~r2 in Hashtbl.replace a key value | Svd _ -> assert false let set_four_index ~r1 ~r2 ~value t = check_bounds r1 r2 t; unsafe_set_four_index ~r1 ~r2 ~value t let unsafe_increment_four_index ~r1 ~r2 ~value t = let updated_value = value +. unsafe_get_four_index ~r1 ~r2 t in unsafe_set_four_index ~r1 ~r2 ~value:updated_value t let increment_four_index ~r1 ~r2 ~value t = check_bounds r1 r2 t; unsafe_increment_four_index ~r1 ~r2 ~value t let get ~r1 ~r2 a = get_four_index ~r1 ~r2 a let set ~r1 ~r2 ~value = match classify_float value with | FP_normal -> set_four_index ~r1 ~r2 ~value | FP_zero | FP_subnormal -> fun _ -> () | FP_infinite | FP_nan -> let msg = Printf.sprintf "FourIdxStorage.ml : set : r1 = (%d,%d) ; r2 = (%d,%d)" r1.first r1.second r2.first r2.second in raise (Invalid_argument msg) let increment ~r1 ~r2 = increment_four_index ~r1 ~r2 let create ~size ?(temp_dir="/dev/shm") sparsity = assert (size < max_index); let two_index = SharedMemory.create ~temp_dir Float64 [| size ; size |] |> Bigarray.array2_of_genarray in let two_index_anti = SharedMemory.create ~temp_dir Float64 [| size ; size |] |> Bigarray.array2_of_genarray in let three_index = SharedMemory.create ~temp_dir Float64 [| size * size ; size |] |> Bigarray.array2_of_genarray in let three_index_anti = SharedMemory.create ~temp_dir Float64 [| size * size ; size |] |> Bigarray.array2_of_genarray in let four_index = match sparsity with | `Dense -> let result = SharedMemory.create ~temp_dir Float64 [| size*size ; (size*(size+1))/2 |] |> Bigarray.array2_of_genarray in Dense result | `Sparse -> let result = Hashtbl.create (size*size+13) in Sparse result in { size ; two_index ; two_index_anti ; three_index ; three_index_anti ; four_index } let size t = t.size (** TODO : remove epsilons *) let get_chem t i j k l = get ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } t let get_phys t i j k l = get ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } t let set_chem t i j k l value = set ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t let set_phys t i j k l value = set ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t let increment_chem t i j k l value = increment ~r1:{ first=i ; second=j } ~r2:{ first=k ; second=l } ~value t let increment_phys t i j k l value = increment ~r1:{ first=i ; second=k } ~r2:{ first=j ; second=l } ~value t type element = (** Element for the stream *) { i_r1: int ; j_r2: int ; k_r1: int ; l_r2: int ; value: float } let get_phys_all_i d ~j ~k ~l = Vec.init d.size (fun i -> get_phys d i j k l) let get_chem_all_i d ~j ~k ~l = Vec.init d.size (fun i -> get_chem d i j k l) let get_phys_all_ij d ~k ~l = Mat.init_cols d.size d.size (fun i j -> get_phys d i j k l) let get_chem_all_ij d ~k ~l = if k = l then let result = Mat.col d.three_index k |> Bigarray.genarray_of_array1 in Bigarray.reshape_2 result d.size d.size else match d.four_index with | Dense a -> let kl = sym_index k l in let result = Mat.col a kl |> Bigarray.genarray_of_array1 in Bigarray.reshape_2 result d.size d.size | Sparse a -> Mat.init_cols d.size d.size (fun i j -> get_chem d i j k l) | Svd _ -> assert false let to_stream d = let i = ref 0 and j = ref 1 and k = ref 1 and l = ref 1 in let f_dense _ = incr i; if !i > !k then begin i := 1; incr j; if !j > !l then begin j := 1; incr k; if !k > !l then begin k := 1; incr l; end; end; end; if !l <= d.size then Some { i_r1 = !i ; j_r2 = !j ; k_r1 = !k ; l_r2 = !l ; value = get_phys d !i !j !k !l } else None in Stream.from f_dense (** Write all integrals to a file with the convention *) let to_file ?(cutoff=Constants.epsilon) ~filename data = let oc = open_out filename in to_stream data |> Stream.iter (fun {i_r1 ; j_r2 ; k_r1 ; l_r2 ; value} -> if (abs_float value > cutoff) then Printf.fprintf oc " %5d %5d %5d %5d%20.15f\n" i_r1 j_r2 k_r1 l_r2 value); close_out oc let of_file ~size ~sparsity filename = let result = create ~size sparsity in let ic = Scanf.Scanning.open_in filename in let rec read_line () = let result = try Some (Scanf.bscanf ic " %d %d %d %d %f" (fun i j k l v -> set_phys result i j k l v)) with End_of_file -> None in match result with | Some () -> (read_line [@tailcall]) () | None -> () in read_line (); Scanf.Scanning.close_in ic; result let to_list data = let s = to_stream data in let rec append accu = let d = try Some (Stream.next s) with | Stream.Failure -> None in match d with | None -> List.rev accu | Some d -> (append [@tailcall]) (d :: accu) in append [] let broadcast t = t (* let size = Parallel.broadcast (lazy t.size) in let bufsize = size * size * size in let stream = to_stream t in let rec iterate () = let buffer = Parallel.broadcast (lazy ( if Stream.peek stream = None then None else Some (Array.init bufsize (fun _ -> try Some (Stream.next stream) with _ -> None )) ) ) in match buffer with | None -> () | Some buffer -> begin if not Parallel.master then Array.iter (fun x -> match x with | Some {i_r1 ; j_r2 ; k_r1 ; l_r2 ; value} -> set_phys t i_r1 j_r2 k_r1 l_r2 value | None -> () ) buffer; (iterate [@tailcall]) () end in iterate (); t *) let four_index_transform_dense_sparse ds coef source = let mo_num = Mat.dim2 coef in let destination = create ~size:mo_num ds in let ao_num = Mat.dim1 coef in let mo_num_2 = mo_num * mo_num in let ao_num_2 = ao_num * ao_num in let ao_mo_num = ao_num * mo_num in let range_mo = list_range 1 mo_num in let range_ao = list_range 1 ao_num in let u = Mat.create mo_num_2 mo_num and o = Mat.create ao_num ao_num_2 and p = Mat.create ao_num_2 mo_num and q = Mat.create ao_mo_num mo_num in if Parallel.master then Printf.eprintf "4-idx transformation \n%!"; let task delta = Mat.fill u 0.; List.iter (fun l -> if abs_float coef.{l,delta} > epsilon then begin let jk = ref 1 in List.iter (fun k -> get_chem_all_ij source ~k ~l |> lacpy ~b:o ~bc:!jk |> ignore; jk := !jk + ao_num; ) range_ao; (* o_i_jk *) let p = gemm ~transa:`T ~c:p o coef (* p_jk_alpha = \sum_i o_i_jk c_i_alpha *) in let p' = Bigarray.reshape_2 (Bigarray.genarray_of_array2 p) ao_num ao_mo_num (* p_j_kalpha *) in let q = gemm ~transa:`T ~c:q p' coef (* q_kalpha_beta = \sum_j p_j_kalpha c_j_beta *) in let q' = Bigarray.reshape_2 (Bigarray.genarray_of_array2 q) ao_num mo_num_2 (* q_k_alphabeta = \sum_j p_j_kalpha c_j_beta *) in ignore @@ gemm ~transa:`T ~beta:1. ~alpha:coef.{l,delta} ~c:u q' coef ; (* u_alphabeta_gamma = \sum_k q_k_alphabeta c_k_gamma *) end ) range_ao; let u = Bigarray.reshape (Bigarray.genarray_of_array2 u) [| mo_num ; mo_num ; mo_num |] |> Bigarray.array3_of_genarray in let result = ref [] in List.iter (fun gamma -> List.iter (fun beta -> List.iter (fun alpha -> let x = u.{alpha,beta,gamma} in if x <> 0. then result := (alpha, beta, gamma, delta, x) :: !result; ) (list_range 1 beta) ) range_mo ) (list_range 1 delta); Array.of_list !result in let n = ref 0 in Stream.of_list range_mo |> Farm.run ~f:task ~ordered:true ~comm:Parallel.Node.comm |> Stream.iter (fun l -> if Parallel.master then (incr n ; Printf.eprintf "\r%d / %d%!" !n mo_num); Array.iter (fun (alpha, beta, gamma, delta, x) -> set_chem destination alpha beta gamma delta x) l); if Parallel.master then Printf.eprintf "\n%!"; broadcast destination let svd_of_dense t = let four_index = let f = match t.four_index with | Dense a -> a | _ -> invalid_arg "svd_of_dense expects a Dense storage" in let size = t.size in let b = Mat.create (size*size) (size*size) in for k = 1 to size do for l = 1 to size do let ac = sym_index k l in lacpy f ~n:1 ~ac ~b ~bc:(dense_index k l size) |> ignore; lacpy f ~n:1 ~ac ~b ~bc:(dense_index l k size) |> ignore done done; let d, u, v = gesdd b in let rank = let w = copy d in scal (1. /. d.{1}) w; let rec aux i = if i = Vec.dim w then i else if w.{i} < 1.e-15 then i else aux (i+1) in aux 1 in let d = copy d ~n:rank in let u = lacpy ~n:rank u in let v = lacpy ~m:rank v in Svd (d,u,v) in { t with four_index } let dense_of_svd t = let four_index = let d, u, v = match t.four_index with | Svd (d, u, v) -> d, u, v | _ -> invalid_arg "dense_of_svd expects a Svd storage" in let f = gemm u @@ gemm (Mat.of_diag d) v in let size = t.size in let b = Mat.create (size*size) ((size*(size+1))/2) in for k = 1 to size do for l = 1 to k do let bc = sym_index k l in lacpy f ~ac:(dense_index l k size) ~n:1 ~bc ~b |> ignore done done; Dense b in { t with four_index } let four_index_transform_svd coef source = let t0 = Unix.gettimeofday () in let svd = svd_of_dense source in Printf.printf "Computed SVD in %f seconds\n%!" (Unix.gettimeofday () -. t0); let u0, d0, v0 = match svd.four_index with | Svd (d, u, v) -> u, d, (Mat.transpose_copy v) | _ -> assert false in let t0 = Unix.gettimeofday () in let ao_num = Mat.dim1 coef in let mo_num = Mat.dim2 coef in let svd_num = Vec.dim d0 in let destination = create ~size:mo_num `Dense in let mo_num_2 = mo_num * mo_num in if Parallel.master then Printf.eprintf "4-idx transformation \n%!"; let u_vecs = Mat.to_col_vecs u0 and v_vecs = Mat.to_col_vecs v0 in Printf.printf "%d %d %d %d\n" (Mat.dim1 u0) (Mat.dim2 u0) (Mat.dim1 v0) (Mat.dim2 v0); let task a = let o = Bigarray.reshape_2 (Bigarray.genarray_of_array1 u_vecs.(a-1)) ao_num ao_num in let u = let p = gemm ~transa:`T coef @@ gemm o coef in Bigarray.reshape_1 (Bigarray.genarray_of_array2 p) mo_num_2 in let o = Bigarray.reshape_2 (Bigarray.genarray_of_array1 v_vecs.(a-1)) ao_num ao_num in let v = let p = gemm ~transa:`T coef @@ gemm o coef in Bigarray.reshape_1 (Bigarray.genarray_of_array2 p) mo_num_2 in (u, v) in (* let four_index = Svd (d,u'',v'') in let t = make_from_four_index { source with size = mo_num; two_index; two_index_anti; three_index; three_index_anti; four_index } *) let n = ref 0 in let u_list = ref [] and v_list = ref [] in Util.list_range 1 svd_num |> Stream.of_list |> Farm.run ~f:task ~ordered:true ~comm:Parallel.Node.comm |> Stream.iter (fun (u,v) -> if Parallel.master then (incr n ; Printf.eprintf "\r%d / %d%!" !n mo_num); u_list := u :: !u_list; v_list := v :: !v_list); let u = Mat.of_col_vecs_list (List.rev !u_list) and v = Mat.of_col_vecs_list (List.rev !v_list) in Printf.printf "Computed 4-idx in %f seconds\n%!" (Unix.gettimeofday () -. t0); let result = let p = gemm u @@ gemm ~transb:`T (Mat.of_diag d0) v in Bigarray.reshape (Bigarray.genarray_of_array2 p) [| mo_num ; mo_num ; mo_num ; mo_num |] in for a=1 to mo_num do for b=a to mo_num do for c=1 to mo_num do for d=c to mo_num do let x = result.{a,b,c,d} in if abs_float x > epsilon then set_chem destination a b c d x; done done done done; Printf.printf "Computed 4-idx in %f seconds\n%!" (Unix.gettimeofday () -. t0); if Parallel.master then Printf.eprintf "\n%!"; broadcast destination let four_index_transform coef source = match source.four_index with | Sparse _ -> four_index_transform_dense_sparse `Sparse coef source | Svd _ -> assert false (* four_index_transform_svd coef source *) | Dense _ -> four_index_transform_dense_sparse `Dense coef source