diff --git a/Basis/ERI.ml b/Basis/ERI.ml index a063491..fddcc96 100644 --- a/Basis/ERI.ml +++ b/Basis/ERI.ml @@ -23,11 +23,11 @@ let zero_m ~maxm ~expo_pq_inv ~norm_pq_sq = (** Compute all the integrals of a contracted class *) let contracted_class shell_a shell_b shell_c shell_d : float Zmap.t = - TwoElectronRR.contracted_class ~zero_m shell_a shell_b shell_c shell_d + TwoElectronRRVectorized.contracted_class ~zero_m shell_a shell_b shell_c shell_d (** Compute all the integrals of a contracted class *) let contracted_class_shell_pairs ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t = - TwoElectronRR.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q + TwoElectronRRVectorized.contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q let cutoff2 = cutoff *. cutoff diff --git a/Basis/TwoElectronRRVectorized.ml b/Basis/TwoElectronRRVectorized.ml new file mode 100644 index 0000000..efe23f3 --- /dev/null +++ b/Basis/TwoElectronRRVectorized.ml @@ -0,0 +1,346 @@ +open Util + +let cutoff2 = cutoff *. cutoff + +exception NullQuartet + +(** Horizontal and Vertical Recurrence Relations (HVRR) *) +let hvrr_two_e m (angMom_a, angMom_b, angMom_c, angMom_d) + (totAngMom_a, totAngMom_b, totAngMom_c, totAngMom_d) + (maxm, zero_m_array) + (expo_b, expo_d) + (expo_inv_p, expo_inv_q) + (center_ab, center_cd, center_pq) coef_prod + map + = + + 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 + and totAngMom_d = Angular_momentum.to_int totAngMom_d + in + + (** Vertical recurrence relations *) + let rec vrr0 m angMom_a = function + | 0 -> Array.mapi (fun j c -> c *. zero_m_array.(j).(m)) coef_prod + | 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 + + | totAngMom_a -> + let key = Zkey.of_int_tuple (Zkey.Three + (angMom_a.(0)+1, angMom_a.(1)+1, angMom_a.(2)+1) ) + in + + let (found, result) = + try (true, Zmap.find map.(m) key) with + | Not_found -> (false, + let am = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |] + and amm = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |] + and xyz = + match angMom_a with + | [|0;0;_|] -> 2 + | [|0;_;_|] -> 1 + | _ -> 0 + 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 + ) + + in + + if not found then + Zmap.add map.(m) key result; + result + + 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) -> vrr0 m angMom_a totAngMom_a + | (_,_) -> + + let key = Zkey.of_int_tuple (Zkey.Six + ((angMom_a.(0)+1, angMom_a.(1)+1, angMom_a.(2)+1), + (angMom_c.(0)+1, angMom_c.(1)+1, angMom_c.(2)+1)) ) + + in + + let (found, result) = + try (true, Zmap.find map.(m) key) with + | Not_found -> (false, + let am = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |] + and cm = [| angMom_c.(0) ; angMom_c.(1) ; angMom_c.(2) |] + and cmm = [| angMom_c.(0) ; angMom_c.(1) ; angMom_c.(2) |] + and xyz = + match angMom_c with + | [|0;0;_|] -> 2 + | [|0;_;_|] -> 1 + | _ -> 0 + in + 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) ) ) + ) + in + if not found then + Zmap.add map.(m) key result; + result + + + + + (** Horizontal recurrence relations *) + and hrr0 m angMom_a angMom_b angMom_c + totAngMom_a totAngMom_b totAngMom_c = + + match totAngMom_b with + | 0 -> vrr m angMom_a angMom_c totAngMom_a totAngMom_c + | 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) ) + | _ -> + let ap = [| angMom_a.(0) ; angMom_a.(1) ; angMom_a.(2) |] + and bm = [| angMom_b.(0) ; angMom_b.(1) ; angMom_b.(2) |] + and xyz = + match angMom_b with + | [|0;0;_|] -> 2 + | [|0;_;_|] -> 1 + | _ -> 0 + 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 ) + ) + + and hrr m angMom_a angMom_b angMom_c angMom_d + totAngMom_a totAngMom_b totAngMom_c totAngMom_d = + + match (totAngMom_b, totAngMom_d) with + | (0,0) -> + vrr m angMom_a angMom_c totAngMom_a totAngMom_c + | (_,0) -> hrr0 m angMom_a angMom_b angMom_c totAngMom_a totAngMom_b totAngMom_c + | (_,_) -> + let cp = [| angMom_c.(0) ; angMom_c.(1) ; angMom_c.(2) |] + and dm = [| angMom_d.(0) ; angMom_d.(1) ; angMom_d.(2) |] + and xyz = + match angMom_d with + | [|0;0;_|] -> 2 + | [|0;_;_|] -> 1 + | _ -> 0 + 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) ) ) + in + hrr m angMom_a angMom_b angMom_c angMom_d totAngMom_a totAngMom_b + totAngMom_c totAngMom_d + + + + +let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q : float Zmap.t = + + let shell_a = shell_p.(0).Shell_pair.shell_a + and shell_b = shell_p.(0).Shell_pair.shell_b + and shell_c = shell_q.(0).Shell_pair.shell_a + and shell_d = shell_q.(0).Shell_pair.shell_b + in + let maxm = + let open Angular_momentum in + (to_int @@ Contracted_shell.totAngMom shell_a) + (to_int @@ Contracted_shell.totAngMom shell_b) + + (to_int @@ Contracted_shell.totAngMom shell_c) + (to_int @@ Contracted_shell.totAngMom shell_d) + in + + (* Pre-computation of integral class indices *) + let class_indices = + Angular_momentum.zkey_array + (Angular_momentum.Quartet + Contracted_shell.(totAngMom shell_a, totAngMom shell_b, + totAngMom shell_c, totAngMom shell_d)) + in + + let contracted_class = + Array.make (Array.length class_indices) 0.; + in + + (* Compute all integrals in the shell for each pair of significant shell pairs *) + + begin + match Contracted_shell.(totAngMom shell_a, totAngMom shell_b, + totAngMom shell_c, totAngMom shell_d) with + | Angular_momentum.(S,S,S,S) -> + contracted_class.(0) <- + Array.fold_left (fun accu shell_ab -> accu +. + Array.fold_left (fun accu shell_cd -> + let coef_prod = + shell_ab.Shell_pair.coef *. shell_cd.Shell_pair.coef + in + (** Screening on the product of coefficients *) + try + if (abs_float coef_prod) < 1.e-4*.cutoff then + raise NullQuartet; + + 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 coef_prod = + shell_ab.Shell_pair.coef *. shell_cd.Shell_pair.coef + in + let integral = + zero_m_array.(0) + in + accu +. coef_prod *. integral + with NullQuartet -> accu + ) 0. 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 + 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 + in + let map = Array.init maxm (fun _ -> Zmap.create (Array.length class_indices)) in + (* Compute the integral class from the primitive shell quartet *) + Array.iteri (fun i key -> + let a = Zkey.to_int_array Zkey.Kind_12 key in + let (angMomA,angMomB,angMomC,angMomD) = + ( [| a.(0) ; a.(1) ; a.(2) |], + [| a.(3) ; a.(4) ; a.(5) |], + [| a.(6) ; a.(7) ; a.(8) |], + [| a.(9) ; a.(10) ; a.(11) |] ) + in + let norm = + shell_ab.Shell_pair.norm_fun angMomA angMomB *. shell_q.(0).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 ) + in + contracted_class.(i) <- contracted_class.(i) +. integral + ) class_indices + ) shell_p + end + end; + + let result = + Zmap.create (Array.length contracted_class) + in + Array.iteri (fun i key -> Zmap.add result key contracted_class.(i)) class_indices; + result + + + +(** Computes all the two-electron integrals of the contracted shell quartet *) +let contracted_class ~zero_m shell_a shell_b shell_c shell_d : float Zmap.t = + + let shell_p = Shell_pair.create_array shell_a shell_b + and shell_q = Shell_pair.create_array shell_c shell_d + in + contracted_class_shell_pairs ~zero_m shell_p shell_q + +