From 7415033b64fc3b787e555ae7240375e29152a229 Mon Sep 17 00:00:00 2001 From: Anthony Scemama Date: Sat, 10 Feb 2018 22:27:00 +0100 Subject: [PATCH] Less allocations --- Basis/TwoElectronRRVectorized.ml | 205 +++++++++++++++++++------------ 1 file changed, 125 insertions(+), 80 deletions(-) diff --git a/Basis/TwoElectronRRVectorized.ml b/Basis/TwoElectronRRVectorized.ml index f552385..a80e578 100644 --- a/Basis/TwoElectronRRVectorized.ml +++ b/Basis/TwoElectronRRVectorized.ml @@ -161,110 +161,157 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) (* let result = Array.make_matrix np nq 0. in *) - let f1 = + let do_compute = ref false in + let v1 = let f = (Coordinate.coord center_cd xyz) in - Array.init nq (fun k -> expo_d.(k) *. expo_inv_q.(k) *. f) - in - let v1 = - if (at_least_one_valid f1) then - vrr_v m angMom_a cm totAngMom_a (totAngMom_c-1) + let f1 = + Array.init nq (fun k -> + let x = expo_d.(k) *. expo_inv_q.(k) *. f in + if (!do_compute) then x + else (if abs_float x > cutoff then do_compute := true ; x) + ) + in + if (!do_compute) then + match vrr_v m angMom_a cm totAngMom_a (totAngMom_c-1) with + | None -> None + | Some v1 -> + begin + let result = Array.make_matrix np nq 0. in + for l=0 to np-1 do + for k=0 to nq-1 do + result.(l).(k) <- -. v1.(l).(k) *. f1.(k) + done + done; + Some result + end else None in - let f2 = - Array.init np (fun l -> - Array.init nq (fun k -> - expo_inv_q.(k) *. center_pq.(xyz).(l).(k) ) ) - in + let v2 = - if (at_least_one_valid (Array.to_list f2 |> Array.concat)) then - vrr_v (m+1) angMom_a cm totAngMom_a (totAngMom_c-1) + let f2 = + Array.init np (fun l -> + Array.init nq (fun k -> + let x = expo_inv_q.(k) *. center_pq.(xyz).(l).(k) in + if (!do_compute) then x + else (if abs_float x > cutoff then do_compute := true ; x) + ) ) + in + if (!do_compute) then + match vrr_v (m+1) angMom_a cm totAngMom_a (totAngMom_c-1) with + | None -> None + | Some v2 -> + begin + for l=0 to np-1 do + for k=0 to nq-1 do + f2.(l).(k) <- -. v2.(l).(k) *. f2.(l).(k) + done + done; + Some f2 + end else None in let p1 = match v1, v2 with - | Some v1, Some v2 -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> - -. v1.(l).(k) *. f1.(k) -. v2.(l).(k) *. f2.(l).(k)) ) ) - | None, Some v2 -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> -. v2.(l).(k) *. f2.(l).(k)) ) ) - | Some v1, None -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> -. v1.(l).(k) *. f1.(k) ) ) ) | None, None -> None + | Some v1, Some v2 -> + Some (Array.init np (fun l -> Array.init nq (fun k -> + v1.(l).(k) +. v2.(l).(k) ) ) ) + | None, Some v2 -> Some v2 + | Some v1, None -> Some v1 in let p2 = if cxyz < 2 then p1 else - let fcm = - (float_of_int (cxyz-1)) *. 0.5 - in + let fcm = (float_of_int (cxyz-1)) *. 0.5 in let f1 = - Array.init nq (fun k -> fcm *. expo_inv_q.(k)) - in - let f2 = - Array.mapi (fun k x -> x *. expo_inv_q.(k)) f1 + Array.init nq (fun k -> + let x = fcm *. expo_inv_q.(k) in + if (!do_compute) then x + else (if abs_float x > cutoff then do_compute := true ; x) + ) in let v1 = - if (at_least_one_valid f1) then - vrr_v m angMom_a cmm totAngMom_a (totAngMom_c-2) + if (!do_compute) then + match vrr_v m angMom_a cmm totAngMom_a (totAngMom_c-2) with + | None -> None + | Some v1 -> + begin + let result = Array.make_matrix np nq 0. in + for l=0 to np-1 do + for k=0 to nq-1 do + result.(l).(k) <- v1.(l).(k) *. f1.(k) + done + done; + Some result + end else None in + let v2 = - if (at_least_one_valid f2) then - vrr_v (m+1) angMom_a cmm totAngMom_a (totAngMom_c-2) + let f2 = + Array.init nq (fun k -> + let x = expo_inv_q.(k) *. f1.(k) in + if (!do_compute) then x + else (if abs_float x > cutoff then do_compute := true ; x) + ) + in + if (!do_compute) then + match vrr_v (m+1) angMom_a cmm totAngMom_a (totAngMom_c-2) with + | None -> None + | Some v2 -> + begin + let result = Array.make_matrix np nq 0. in + for l=0 to np-1 do + for k=0 to nq-1 do + result.(l).(k) <- v2.(l).(k) *. f2.(k) + done + done; + Some result + end else None in match p1, v1, v2 with - | Some p1, Some v1, Some v2 -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> - p1.(l).(k) +. f1.(k) *. v1.(l).(k) +. f2.(k) *. v2.(l).(k)) ) ) - | Some p1, Some v1, None -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> p1.(l).(k) +. f1.(k) *. v1.(l).(k) ) ) ) - | Some p1, None, Some v2 -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> p1.(l).(k) +. f2.(k) *. v2.(l).(k)) ) ) - | None , Some v1, Some v2 -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> - f1.(k) *. v1.(l).(k) +. f2.(k) *. v2.(l).(k)) ) ) - | Some p1, None, None -> Some p1 - | None , Some v1, None -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> f1.(k) *. v1.(l).(k) ) ) ) - | None, None, Some v2 -> - Some (Array.init np (fun l -> - Array.init nq (fun k -> f2.(k) *. v2.(l).(k)) ) ) | None, None, None -> None + | Some p1, Some v1, Some v2 -> + Some (Array.init np (fun l -> Array.init nq (fun k -> + p1.(l).(k) +. v1.(l).(k) +. v2.(l).(k)) ) ) + | Some p1, Some v1, None -> + Some (Array.init np (fun l -> Array.init nq (fun k -> + p1.(l).(k) +. v1.(l).(k) ) ) ) + | Some p1, None, Some v2 -> + Some (Array.init np (fun l -> Array.init nq (fun k -> + p1.(l).(k) +. v2.(l).(k)) ) ) + | None , Some v1, Some v2 -> + Some (Array.init np (fun l -> Array.init nq (fun k -> + v1.(l).(k) +. v2.(l).(k)) ) ) + | Some p1, None, None -> Some p1 + | None, Some v1, None -> Some v1 + | None, None, Some v2 -> Some v2 in if (axyz < 1) || (cxyz < 1) then p2 else let v = vrr_v (m+1) am cm (totAngMom_a-1) (totAngMom_c-1) in - begin match (p2, v) with - | Some p2, Some v -> Some ( - Array.init np (fun l -> - let fa = - (float_of_int axyz) *. expo_inv_p.(l) *. 0.5 - in - Array.init nq (fun k -> p2.(l).(k) -. expo_inv_q.(k) *. fa *. v.(l).(k)) - ) ) - | Some p2, None -> Some p2 - | None, Some v -> Some ( - Array.init np (fun l -> - let fa = - (float_of_int axyz) *. expo_inv_p.(l) *. 0.5 - in - Array.init nq (fun k -> -. fa *. expo_inv_q.(k) *. v.(l).(k)) - ) ) | None, None -> None - end + | Some p2, None -> Some p2 + | _, Some v -> + begin + let p2 = + match p2 with + | None -> Array.make_matrix np nq 0. + | Some p2 -> p2 + in + for l=0 to np-1 do + let fa = (float_of_int axyz) *. expo_inv_p.(l) *. 0.5 in + for k=0 to nq-1 do + p2.(l).(k) <- p2.(l).(k) -. fa *. expo_inv_q.(k) *. v.(l).(k) + done + done; + Some p2 + end end in Zmap.add map_2d.(m) key result; result @@ -415,15 +462,13 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q totAngMom shell_c, totAngMom shell_d) with | Angular_momentum.(S,S,S,S) -> contracted_class.(0) <- - let expo_inv_p = - Array.map (fun shell_ab -> shell_ab.ShellPair.expo_inv) sp - |> Vec.of_array - and expo_inv_q = - Array.map (fun shell_cd -> shell_cd.ShellPair.expo_inv) sq - |> Vec.of_array - in let np, nq = - Vec.dim expo_inv_p, Vec.dim expo_inv_q + Array.length sp, Array.length sq + in + let expo_inv_p = + Vec.init np (fun ab -> sp.(ab-1).ShellPair.expo_inv) + and expo_inv_q = + Vec.init nq (fun cd -> sq.(cd-1).ShellPair.expo_inv) in let coef =