From 8a851eb4ed5403aacbf3f628e1c61bb106608b0d Mon Sep 17 00:00:00 2001 From: Anthony Scemama Date: Thu, 25 Jan 2018 13:59:31 +0100 Subject: [PATCH] Working on vectorized code --- Basis/TwoElectronRRVectorized.ml | 236 +++++++++++++++---------------- 1 file changed, 118 insertions(+), 118 deletions(-) diff --git a/Basis/TwoElectronRRVectorized.ml b/Basis/TwoElectronRRVectorized.ml index efe23f3..85f0a19 100644 --- a/Basis/TwoElectronRRVectorized.ml +++ b/Basis/TwoElectronRRVectorized.ml @@ -10,10 +10,13 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) (maxm, zero_m_array) (expo_b, expo_d) (expo_inv_p, expo_inv_q) - (center_ab, center_cd, center_pq) coef_prod - map + (center_ab, center_cd, center_pq) + coef_prod map = + let k = 0 in + let getk a = a.(k) in + let totAngMom_a = Angular_momentum.to_int totAngMom_a and totAngMom_b = Angular_momentum.to_int totAngMom_b and totAngMom_c = Angular_momentum.to_int totAngMom_c @@ -22,11 +25,14 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) (** Vertical recurrence relations *) let rec vrr0 m angMom_a = function - | 0 -> Array.mapi (fun j c -> c *. zero_m_array.(j).(m)) coef_prod + | 0 -> Array.mapi (fun k c -> c *. zero_m_array.(k).(m)) coef_prod + |> getk | 1 -> let i = if angMom_a.(0) = 1 then 0 else if angMom_a.(1) = 1 then 1 else 2 - in Array.mapi (fun j c -> - c *. expo_inv_p *.( (Coordinate.coord center_pq.(j) i) *. zero_m_array.(j).(m+1) - -. expo_b *. (Coordinate.coord center_ab i) *. zero_m_array.(j).(m) ) ) coef_prod + in + Array.mapi (fun k c -> c *. expo_inv_p *. + ( (Coordinate.coord center_pq.(k) i) *. zero_m_array.(k).(m+1) + -. expo_b *. (Coordinate.coord center_ab i) *. zero_m_array.(k).(m) ) ) coef_prod + |> getk | totAngMom_a -> let key = Zkey.of_int_tuple (Zkey.Three @@ -46,27 +52,20 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) in am.(xyz) <- am.(xyz) - 1; amm.(xyz) <- amm.(xyz) - 2; - if am.(xyz) < 0 then Array.map (fun _ -> 0.) coef_prod - else - let result = - vrr0 (m+1) am (totAngMom_a-1) - |> Array.mapi (fun j x -> x *. expo_inv_p *. (Coordinate.coord center_pq.(j) xyz)) - |> Array.map2 (+.) - (vrr0 m am (totAngMom_a-1) - |> Array.map (fun x -> -. x *. expo_b *. expo_inv_p *. (Coordinate.coord center_ab xyz)) - ) - in - if amm.(xyz) < 0 then result else - vrr0 (m+1) amm (totAngMom_a-2) - |> Array.map (fun x -> x *. expo_inv_p) - |> Array.map2 (+.) ( - vrr0 m amm (totAngMom_a-2) - |> Array.map (fun x -> x *. (float_of_int am.(xyz)) *. expo_inv_p *. 0.5)) - |> Array.map2 (+.) result - ) - +if am.(xyz) < 0 then 0. else + Array.mapi (fun k _ -> + (-. expo_b *. expo_inv_p *. (Coordinate.coord center_ab xyz)) + *. (vrr0 m am (totAngMom_a-1) ) + +. (expo_inv_p *. (Coordinate.coord center_pq.(k) xyz)) + *.(vrr0 (m+1) am (totAngMom_a-1) ) + +. (if amm.(xyz) < 0 then 0. else + ((float_of_int am.(xyz)) *. expo_inv_p *. 0.5) + *. (vrr0 m amm (totAngMom_a-2) + +. expo_inv_p *. (vrr0 (m+1) amm (totAngMom_a-2) ) ) ) +) coef_prod + |> getk +) in - if not found then Zmap.add map.(m) key result; result @@ -74,7 +73,8 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) and vrr m angMom_a angMom_c totAngMom_a totAngMom_c = match (totAngMom_a, totAngMom_c) with - | (0,0) -> Array.mapi (fun j c -> c *. zero_m_array.(j).(m)) coef_prod + | (0,0) -> Array.mapi (fun k c -> c *. zero_m_array.(k).(m)) coef_prod + |> getk | (_,0) -> vrr0 m angMom_a totAngMom_a | (_,_) -> @@ -99,26 +99,21 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) am.(xyz) <- am.(xyz) - 1; cm.(xyz) <- cm.(xyz) - 1; cmm.(xyz) <- cmm.(xyz) - 2; - if cm.(xyz) < 0 then Array.map (fun _ -> 0.) coef_prod else - let result = - Array.map2 (-.) - ( Array.mapi (fun j x -> -. x *. expo_d.(j) *. expo_inv_q.(j) *. (Coordinate.coord center_cd.(j) xyz) ) (vrr m angMom_a cm totAngMom_a (totAngMom_c-1) ) ) - ( Array.mapi (fun j x -> x *. expo_inv_q.(j) *. (Coordinate.coord center_pq.(j) xyz)) (vrr (m+1) angMom_a cm totAngMom_a (totAngMom_c-1) ) ) - in - let result = - if cmm.(xyz) < 0 then result else - Array.map2 (+.) result - ( Array.mapi (fun j x -> x *. (float_of_int cm.(xyz)) *. expo_inv_q.(j) *. 0.5 ) - (vrr m angMom_a cmm totAngMom_a (totAngMom_c-2)) ) - |> Array.map2 (+.) - ( Array.mapi (fun j x -> x *. expo_inv_q.(j)) - (vrr (m+1) angMom_a cmm totAngMom_a (totAngMom_c-2) ) ) - in - if am.(xyz) lor cm.(xyz) < 0 then result else - Array.map2 (-.) result - (Array.mapi (fun j x -> x *. (float_of_int angMom_a.(xyz)) *. expo_inv_p *. expo_inv_q.(j) *. 0.5 ) - (vrr (m+1) am cm (totAngMom_a-1) (totAngMom_c-1) ) ) - ) +if cm.(xyz) < 0 then 0. else + Array.mapi (fun k _ -> + (-. expo_d.(k) *. expo_inv_q.(k) *. (Coordinate.coord center_cd.(k) xyz) ) + *.(vrr m angMom_a cm totAngMom_a (totAngMom_c-1) ) + -. (expo_inv_q.(k) *. (Coordinate.coord center_pq.(k) xyz)) + *.(vrr (m+1) angMom_a cm totAngMom_a (totAngMom_c-1) ) + +. (if cmm.(xyz) < 0 then 0. else + ((float_of_int cm.(xyz)) *. expo_inv_q.(k) *. 0.5 ) + *.(vrr m angMom_a cmm totAngMom_a (totAngMom_c-2) + +. expo_inv_q.(k) *. (vrr (m+1) angMom_a cmm totAngMom_a (totAngMom_c-2) ) ) ) + -. (if am.(xyz) lor cm.(xyz) < 0 then 0. else + ((float_of_int angMom_a.(xyz)) *. expo_inv_p *. expo_inv_q.(k) *. 0.5 ) + *.(vrr (m+1) am cm (totAngMom_a-1) (totAngMom_c-1) ) )) coef_prod + |> getk + ) in if not found then Zmap.add map.(m) key result; @@ -136,10 +131,8 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) | 1 -> let xyz = if angMom_b.(0) = 1 then 0 else if angMom_b.(1) = 1 then 1 else 2 in let ap = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |] in ap.(xyz) <- ap.(xyz) + 1; - Array.map2 (+.) - (vrr m ap angMom_c (totAngMom_a+1) totAngMom_c) - (Array.map (fun x -> x *. (Coordinate.coord center_ab xyz)) - (vrr m angMom_a angMom_c totAngMom_a totAngMom_c) ) + vrr m ap angMom_c (totAngMom_a+1) totAngMom_c + +. (Coordinate.coord center_ab xyz) *. (vrr m angMom_a angMom_c totAngMom_a totAngMom_c) | _ -> let ap = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |] and bm = [| angMom_b.(0) ; angMom_b.(1) ; angMom_b.(2) |] @@ -151,12 +144,11 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) in ap.(xyz) <- ap.(xyz) + 1; bm.(xyz) <- bm.(xyz) - 1; - if (bm.(xyz) < 0) then Array.map (fun _ -> 0.) coef_prod else - Array.map2 (+.) - (hrr0 m ap bm angMom_c (totAngMom_a+1) (totAngMom_b-1) totAngMom_c ) - (Array.map (fun x -> x *. (Coordinate.coord center_ab xyz)) - (hrr0 m angMom_a bm angMom_c totAngMom_a (totAngMom_b-1) totAngMom_c ) - ) +if (bm.(xyz) < 0) then 0. else + hrr0 m ap bm angMom_c (totAngMom_a+1) (totAngMom_b-1) totAngMom_c + +. (Coordinate.coord center_ab xyz) + *.(hrr0 m angMom_a bm angMom_c totAngMom_a (totAngMom_b-1) + totAngMom_c ) and hrr m angMom_a angMom_b angMom_c angMom_d totAngMom_a totAngMom_b totAngMom_c totAngMom_d = @@ -176,12 +168,12 @@ let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) in cp.(xyz) <- cp.(xyz) + 1; dm.(xyz) <- dm.(xyz) - 1; - Array.map2 (+.) - (hrr m angMom_a angMom_b cp dm totAngMom_a totAngMom_b - (totAngMom_c+1) (totAngMom_d-1) ) - (Array.mapi (fun j x -> x *. (Coordinate.coord center_cd.(j) xyz)) - (hrr m angMom_a angMom_b angMom_c dm totAngMom_a totAngMom_b - totAngMom_c (totAngMom_d-1) ) ) + let h1, h2 = + hrr m angMom_a angMom_b cp dm totAngMom_a totAngMom_b (totAngMom_c+1) (totAngMom_d-1) , + hrr m angMom_a angMom_b angMom_c dm totAngMom_a totAngMom_b totAngMom_c (totAngMom_d-1) + in + Array.map (fun center_cd -> h1 +. h2 *. (Coordinate.coord center_cd xyz)) center_cd + |> getk in hrr m angMom_a angMom_b angMom_c angMom_d totAngMom_a totAngMom_b totAngMom_c totAngMom_d @@ -222,7 +214,7 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q | Angular_momentum.(S,S,S,S) -> contracted_class.(0) <- Array.fold_left (fun accu shell_ab -> accu +. - Array.fold_left (fun accu shell_cd -> + Array.fold_left (fun accu shell_cd -> let coef_prod = shell_ab.Shell_pair.coef *. shell_cd.Shell_pair.coef in @@ -234,19 +226,19 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q let expo_pq_inv = shell_ab.Shell_pair.expo_inv +. shell_cd.Shell_pair.expo_inv in - let center_pq = + let center_pq = Coordinate.(shell_ab.Shell_pair.center |- shell_cd.Shell_pair.center) in let norm_pq_sq = Coordinate.dot center_pq center_pq in - let zero_m_array = + let zero_m_array = zero_m ~maxm ~expo_pq_inv ~norm_pq_sq in let coef_prod = - shell_ab.Shell_pair.coef *. shell_cd.Shell_pair.coef + shell_ab.Shell_pair.coef *. shell_cd.Shell_pair.coef in let integral = zero_m_array.(0) @@ -257,46 +249,49 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q ) 0. shell_p | _ -> - begin - Array.iter (fun shell_ab -> - let b = shell_ab.Shell_pair.j in - let common = - Array.map (fun shell_cd -> - let coef_prod = - shell_ab.Shell_pair.coef *. shell_cd.Shell_pair.coef - in - let coef_prod = - if (abs_float coef_prod) < 1.e-4*.cutoff then - 0. else coef_prod - in - let expo_pq_inv = - shell_ab.Shell_pair.expo_inv +. shell_cd.Shell_pair.expo_inv - in - let center_pq = - Coordinate.(shell_ab.Shell_pair.center |- shell_cd.Shell_pair.center) - in - let norm_pq_sq = - Coordinate.dot center_pq center_pq - in - - let zero_m_array = - zero_m ~maxm ~expo_pq_inv ~norm_pq_sq - in - - let d = shell_cd.Shell_pair.j in - - (zero_m_array, shell_cd.Shell_pair.expo_inv, - Contracted_shell.expo shell_d d, shell_cd.Shell_pair.center_ab, - center_pq,coef_prod) - ) shell_q +Array.iter (fun shell_ab -> + let b = shell_ab.Shell_pair.j in + let common = + Array.map (fun shell_cd -> + let coef_prod = + shell_ab.Shell_pair.coef *. shell_cd.Shell_pair.coef in - let (zero_m_array, expo_inv, d, center_cd, center_pq, coef_prod) = - Array.map (fun (x,_,_,_,_,_) -> x) common, - Array.map (fun (_,x,_,_,_,_) -> x) common, - Array.map (fun (_,_,x,_,_,_) -> x) common, - Array.map (fun (_,_,_,x,_,_) -> x) common, - Array.map (fun (_,_,_,_,x,_) -> x) common, - Array.map (fun (_,_,_,_,_,x) -> x) common + let coef_prod = + if (abs_float coef_prod) < 1.e-4*.cutoff then + 0. else coef_prod + in + let expo_pq_inv = + shell_ab.Shell_pair.expo_inv +. shell_cd.Shell_pair.expo_inv + in + let center_pq = + Coordinate.(shell_ab.Shell_pair.center |- shell_cd.Shell_pair.center) + in + let norm_pq_sq = + Coordinate.dot center_pq center_pq + in + + let zero_m_array = + zero_m ~maxm ~expo_pq_inv ~norm_pq_sq + in + + let d = shell_cd.Shell_pair.j in + + (zero_m_array, shell_cd.Shell_pair.expo_inv, + Contracted_shell.expo shell_d d, shell_cd.Shell_pair.center_ab, + center_pq,coef_prod) + ) shell_q + in + Array.iteri (fun cd shell_cd -> + try + let (zero_m_array, expo_inv, d, center_cd, + center_pq,coef_prod) = common.(cd) + in + let zero_m_array = [| zero_m_array |] + and expo_inv = [| expo_inv |] + and d = [| d |] + and center_cd = [| center_cd |] + and center_pq = [| center_pq |] + and coef_prod = [| coef_prod |] in let map = Array.init maxm (fun _ -> Zmap.create (Array.length class_indices)) in (* Compute the integral class from the primitive shell quartet *) @@ -308,23 +303,28 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q [| a.(6) ; a.(7) ; a.(8) |], [| a.(9) ; a.(10) ; a.(11) |] ) in + try + let norm = - shell_ab.Shell_pair.norm_fun angMomA angMomB *. shell_q.(0).Shell_pair.norm_fun angMomC angMomD + shell_ab.Shell_pair.norm_fun angMomA angMomB *. shell_cd.Shell_pair.norm_fun angMomC angMomD in - let integral = Array.fold_left (fun accu x -> accu +. norm *. x) 0. ( - hvrr_two_e 0 (angMomA, angMomB, angMomC, angMomD) - (Contracted_shell.totAngMom shell_a, Contracted_shell.totAngMom shell_b, - Contracted_shell.totAngMom shell_c, Contracted_shell.totAngMom shell_d) - (maxm, zero_m_array) - (Contracted_shell.expo shell_b b, d) - (shell_ab.Shell_pair.expo_inv, expo_inv) - (shell_ab.Shell_pair.center_ab, center_cd, center_pq) - coef_prod map ) + let integral = chop norm (fun () -> + hvrr_two_e 0 (angMomA, angMomB, angMomC, angMomD) + (Contracted_shell.totAngMom shell_a, Contracted_shell.totAngMom shell_b, + Contracted_shell.totAngMom shell_c, Contracted_shell.totAngMom shell_d) + (maxm, zero_m_array) + (Contracted_shell.expo shell_b b, d) + (shell_ab.Shell_pair.expo_inv, expo_inv) + (shell_ab.Shell_pair.center_ab, center_cd, center_pq) + coef_prod map ) in contracted_class.(i) <- contracted_class.(i) +. integral + with NullQuartet -> () ) class_indices - ) shell_p - end + with NullQuartet -> () + ) shell_q + ) shell_p + end; let result =