From 032f1a0913968bc6773fb35913e18247eac2564c Mon Sep 17 00:00:00 2001 From: Anthony Scemama Date: Fri, 9 Feb 2018 20:32:39 +0100 Subject: [PATCH] Working on contraction --- Basis/TwoElectronRRVectorized.ml | 349 ++++++++++++++++++++----------- 1 file changed, 226 insertions(+), 123 deletions(-) diff --git a/Basis/TwoElectronRRVectorized.ml b/Basis/TwoElectronRRVectorized.ml index 33021b3..7cefe8e 100644 --- a/Basis/TwoElectronRRVectorized.ml +++ b/Basis/TwoElectronRRVectorized.ml @@ -8,14 +8,6 @@ let cutoff2 = cutoff *. cutoff exception NullQuartet exception Found -let at_least_one_valid arr = - try - Vec.iter (fun x -> if (abs_float x > cutoff) then raise Found) arr ; false - with Found -> true - -(*TODO : REMOVE *) -let sum integral = - Array.fold_left (+.) 0. integral (** Horizontal and Vertical Recurrence Relations (HVRR) *) let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) @@ -34,6 +26,7 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) Array.make nq 0. 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 @@ -41,7 +34,8 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) in (** Vertical recurrence relations *) - let rec vrr0_v l m angMom_a = function + let rec vrr0_v m angMom_a = function + (* | 1 -> let xyz = match angMom_a with @@ -49,16 +43,36 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) | (_,1,_) -> 1 | _ -> 2 in - let f = expo_b.{l} *. (Coordinate.coord center_ab xyz) in - Array.init nq (fun k -> coef_prod.{k+1,l} *. expo_inv_p.{l} *. - (center_pq.{xyz+1,k+1,l} *. zero_m_array.(m+1).{k+1,l} - -. f *. zero_m_array.(m).{k+1,l} ) ) - | 0 -> Array.init nq (fun k -> zero_m_array.(m).{k+1,l} *. coef_prod.{k+1,l}) + let a = Mat.mul center_pq.(xyz) zero_m_array.(m+1) in + let b = copy expo_b in + scal (Coordinate.coord center_ab xyz) b; + let c = Mat.map (fun x -> x) zero_m_array.(m) in + Mat.scal_cols c b; + let d = Mat.sub a c in + Mat.scal_cols d expo_inv_p; + Some (Mat.mul coef_prod d) + *) + | 1 -> + let xyz = + match angMom_a with + | (1,_,_) -> 0 + | (_,1,_) -> 1 + | _ -> 2 + in + Some ( + Array.init np (fun ab -> let l=ab+1 in + Array.init nq (fun k -> + let f = expo_b.{l} *. (Coordinate.coord center_ab xyz) in + coef_prod.{k+1,l} *. expo_inv_p.{l} *. + (center_pq.(xyz).{k+1,l} *. zero_m_array.(m+1).{k+1,l} + -. f *. zero_m_array.(m).{k+1,l} ) )) + ) + | 0 -> Some (Mat.mul zero_m_array.(m) coef_prod |> Mat.transpose_copy |> Mat.to_array) | totAngMom_a -> let key = Zkey.of_int_tuple (Zkey.Three angMom_a) in - try Zmap.find map_1d.(m).(l-1) key with + try Zmap.find map_1d.(m) key with | Not_found -> let result = let am, amm, amxyz, xyz = @@ -67,46 +81,84 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) | (x,y,0) -> (x,y-1,0),(x,y-2,0), y-1, 1 | (x,y,z) -> (x,y,z-1),(x,y,z-2), z-1, 2 in - if amxyz < 0 then empty else - let v1 = - let f = - -. expo_b.{l} *. expo_inv_p.{l} *. (Coordinate.coord center_ab xyz) - in - if (abs_float f < cutoff) then empty else - Array.map (fun v1k -> f *. v1k) (vrr0_v l m am (totAngMom_a-1) ) + if amxyz < 0 then + None + else + let v1_top, p1_top = + if abs_float (Coordinate.coord center_ab xyz) < cutoff then + None, + vrr0_v (m+1) am (totAngMom_a-1) + else + vrr0_v m am (totAngMom_a-1), + vrr0_v (m+1) am (totAngMom_a-1) + in + let v1_top2, p1_top2 = + if amxyz < 1 then (None,None) else + vrr0_v m amm (totAngMom_a-2), + vrr0_v (m+1) amm (totAngMom_a-2) in - let p1 = - Array.mapi (fun k v2k -> v1.(k) +. expo_inv_p.{l} *. center_pq.{xyz+1,k+1,l} *. v2k) (vrr0_v l (m+1) am (totAngMom_a-1)) - in - if amxyz < 1 then p1 else - let f = (float_of_int amxyz) *. expo_inv_p.{l} *. 0.5 - in - if (abs_float f < cutoff) then empty else - let v1 = vrr0_v l m amm (totAngMom_a-2) - in - let v2 = - if (abs_float (f *. expo_inv_p.{l})) < cutoff then empty else - vrr0_v l (m+1) amm (totAngMom_a-2) - in - Array.init nq (fun k -> p1.(k) +. - f *. (v1.(k) +. v2.(k) *. expo_inv_p.{l} ) ) - in Zmap.add map_1d.(m).(l-1) key result; + Some ( + Array.init np (fun ab -> let l = ab+1 in + let v1 = + let f = + -. expo_b.{l} *. expo_inv_p.{l} *. (Coordinate.coord center_ab xyz) + in + match v1_top with + | Some v1_top -> + v1_top.(l-1) + |> Array.map (fun x -> f *. x) + | None -> empty + in + let p1 = + match p1_top with + | Some p1_top -> + p1_top.(l-1) + | _ -> assert false + in + let p1 = + Array.init nq (fun k -> + v1.(k) +. expo_inv_p.{l} *. center_pq.(xyz).{k+1,l} *. p1.(k)) + in + if amxyz < 1 then p1 else + let f = (float_of_int amxyz) *. expo_inv_p.{l} *. 0.5 + in + let v1 = + match v1_top2 with + | Some v1_top2 -> v1_top2.(l-1) + | None -> assert false + in + let v2 = + match p1_top2 with + | Some p1_top2 -> p1_top2.(l-1) + | None -> assert false + in + Array.init nq (fun k -> + p1.(k) +. f *. (v1.(k) +. v2.(k) *. expo_inv_p.{l} ) ) + ) + ) + in Zmap.add map_1d.(m) key result; result - and vrr_v l m angMom_a angMom_c totAngMom_a totAngMom_c = + + and vrr_v m angMom_a angMom_c totAngMom_a totAngMom_c = match (totAngMom_a, totAngMom_c) with - | (i,0) -> if (i>0) then - vrr0_v l m angMom_a totAngMom_a + | (i,0) -> + if (i>0) then + begin + match vrr0_v m angMom_a totAngMom_a with + | Some x -> Some (Mat.of_array x |> Mat.transpose_copy) + | None -> None + end else - Array.init nq (fun k -> zero_m_array.(m).{k+1,l} *. coef_prod.{k+1,l}) + Some (Mat.mul zero_m_array.(m) coef_prod ) | (_,_) -> let key = Zkey.of_int_tuple (Zkey.Six (angMom_a, angMom_c)) in - try Zmap.find map_2d.(m).(l-1) key with + try Zmap.find map_2d.(m) key with | Not_found -> let result = begin @@ -129,32 +181,39 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) (acx-2,acy,acz), aax,acx, 0 in - (* - if cxyz < 1 then empty else - *) let f1 = let f = (Coordinate.coord center_cd xyz) in - Array.init nq (fun k -> - expo_d.{k+1} *. expo_inv_q.{k+1} *. f) - |> Vec.of_array - in - let f2 = - Array.init nq (fun k -> - expo_inv_q.{k+1} *. center_pq.{xyz+1,k+1,l} ) - |> Vec.of_array + Vec.init nq (fun k -> + expo_d.{k} *. expo_inv_q.{k} *. f) in let v1 = - if (at_least_one_valid f1) then - vrr_v l m angMom_a cm totAngMom_a (totAngMom_c-1) - else empty - and v2 = - if (at_least_one_valid f2) then - vrr_v l (m+1) angMom_a cm totAngMom_a (totAngMom_c-1) - else empty + if (abs_float @@ amax f1 > cutoff) then + vrr_v m angMom_a cm totAngMom_a (totAngMom_c-1) + else None in + let f2 = + Mat.init_cols nq np (fun k l -> + expo_inv_q.{k} *. center_pq.(xyz).{k,l} ) + in + let v2 = + if (Mat.as_vec f2 |> amax |> abs_float) < cutoff then + None + else + vrr_v (m+1) angMom_a cm totAngMom_a (totAngMom_c-1) + in + let p1 = - Array.init nq (fun k -> -. v1.(k) *. f1.{k+1} -. v2.(k) *. f2.{k+1}) + match v1, v2 with + | Some v1, Some v2 -> + Some (Mat.init_cols nq np (fun k l -> + -. v1.{k,l} *. f1.{k} -. v2.{k,l} *. f2.{k,l}) ) + | None, Some v2 -> + Some (Mat.init_cols nq np (fun k l -> -. v2.{k,l} *. f2.{k,l}) ) + | Some v1, None -> + Some (Mat.init_cols nq np (fun k l -> -. v1.{k,l} *. f1.{k} ) ) + | None, None -> None in + let p2 = if cxyz < 2 then p1 else let fcm = @@ -167,50 +226,88 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) Vec.mul f1 expo_inv_q in let v1 = - if (at_least_one_valid f1) then - vrr_v l m angMom_a cmm totAngMom_a (totAngMom_c-2) - else empty + if (abs_float @@ amax f1 > cutoff) then + vrr_v m angMom_a cmm totAngMom_a (totAngMom_c-2) + else None in let v2 = - if (at_least_one_valid f2) then - vrr_v l (m+1) angMom_a cmm totAngMom_a (totAngMom_c-2) - else empty + if (abs_float @@ amax f2 > cutoff) then + vrr_v (m+1) angMom_a cmm totAngMom_a (totAngMom_c-2) + else None in - Array.init nq (fun k -> p1.(k) +. f1.{k+1} *. v1.(k) +. f2.{k+1} *. v2.(k)) + match p1, v1, v2 with + | Some p1, Some v1, Some v2 -> + Some (Mat.init_cols nq np (fun k l -> p1.{k,l} +. f1.{k} *. v1.{k,l} +. f2.{k} *. v2.{k,l}) ) + | Some p1, Some v1, None -> + Some (Mat.init_cols nq np (fun k l -> p1.{k,l} +. f1.{k} *. v1.{k,l} )) + | Some p1, None, Some v2 -> + Some (Mat.init_cols nq np (fun k l -> p1.{k,l} +. f2.{k} *. v2.{k,l}) ) + | None , Some v1, Some v2 -> + Some (Mat.init_cols nq np (fun k l -> f1.{k} *. v1.{k,l} +. f2.{k} *. v2.{k,l}) ) + | Some p1, None, None -> Some p1 + | None , Some v1, None -> + Some (Mat.init_cols nq np (fun k l -> f1.{k} *. v1.{k,l})) + | None, None, Some v2 -> + Some (Mat.init_cols nq np (fun k l -> f2.{k} *. v2.{k,l}) ) + | None, None, None -> None in if (axyz < 1) || (cxyz < 1) then p2 else - let fa = - (float_of_int axyz) *. expo_inv_p.{l} *. 0.5 + let v = + vrr_v (m+1) am cm (totAngMom_a-1) (totAngMom_c-1) in - let f1 = - Vec.map (fun e -> fa *. e ) expo_inv_q - in - if (at_least_one_valid f1) then - let v = - vrr_v l (m+1) am cm (totAngMom_a-1) (totAngMom_c-1) - in - Array.init nq (fun k -> p2.(k) -. f1.{k+1} *. v.(k)) - else p2 + begin + match (p2, v) with + | Some p2, Some v -> Some ( + Array.init np (fun ab -> let l = ab+1 in + let fa = + (float_of_int axyz) *. expo_inv_p.{l} *. 0.5 + in + let f1 = + Vec.map (fun e -> fa *. e ) expo_inv_q + in + Vec.init nq (fun k -> p2.{k,l} -. f1.{k} *. v.{k,l}) + ) + |> Mat.of_col_vecs ) + | Some p2, None -> Some p2 + | None, Some v -> Some ( + Array.init np (fun ab -> let l = ab+1 in + let fa = + (float_of_int axyz) *. expo_inv_p.{l} *. 0.5 + in + let f1 = + Vec.map (fun e -> fa *. e ) expo_inv_q + in + Vec.init nq (fun k -> -. f1.{k} *. v.{k,l}) + ) + |> Mat.of_col_vecs ) + | None, None -> None + end end - in Zmap.add map_2d.(m).(l-1) key result; + in Zmap.add map_2d.(m) key result; result (** Horizontal recurrence relations *) - and hrr0_v l angMom_a angMom_b angMom_c + and hrr0_v angMom_a angMom_b angMom_c totAngMom_a totAngMom_b totAngMom_c = match totAngMom_b with | 0 -> begin match (totAngMom_a, totAngMom_c) with - | (0,0) -> - Array.init nq (fun k -> zero_m_array.(0).{k+1,l} *. coef_prod.{k+1,l}) - |> sum - | (_,0) -> vrr0_v l 0 angMom_a totAngMom_a |> sum - | (_,_) -> vrr_v l 0 angMom_a angMom_c totAngMom_a totAngMom_c |> sum + | (0,0) -> Mat.gemm_trace zero_m_array.(0) coef_prod + | (_,0) -> + begin + match vrr0_v 0 angMom_a totAngMom_a with + | Some matrix -> Mat.sum (Mat.of_array matrix) + | None -> 0. + end + | (_,_) -> + match vrr_v 0 angMom_a angMom_c totAngMom_a totAngMom_c with + | Some matrix -> Mat.sum matrix + | None -> 0. end | 1 -> let (aax, aay, aaz) = angMom_a in @@ -222,13 +319,17 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) in let f = Coordinate.coord center_ab xyz in let v1 = - vrr_v l 0 ap angMom_c (totAngMom_a+1) totAngMom_c + match vrr_v 0 ap angMom_c (totAngMom_a+1) totAngMom_c with + | Some matrix -> Mat.sum matrix + | None -> 0. in - if (abs_float f < cutoff) then sum v1 else - let v2 = - vrr_v l 0 angMom_a angMom_c totAngMom_a totAngMom_c + if (abs_float f < cutoff) then v1 else + let v2 = + match vrr_v 0 angMom_a angMom_c totAngMom_a totAngMom_c with + | Some matrix -> Mat.sum matrix + | None -> 0. in - Array.map2 (fun v1 v2 -> v1 +. v2 *. f) v1 v2 |> sum + v1 +. v2 *. f | _ -> let (aax, aay, aaz) = angMom_a and (abx, aby, abz) = angMom_b in @@ -247,24 +348,27 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) in let h1 = - hrr0_v l ap bm angMom_c (totAngMom_a+1) (totAngMom_b-1) totAngMom_c + hrr0_v ap bm angMom_c (totAngMom_a+1) (totAngMom_b-1) totAngMom_c in let f = (Coordinate.coord center_ab xyz) in if (abs_float f < cutoff) then h1 else let h2 = - hrr0_v l angMom_a bm angMom_c totAngMom_a (totAngMom_b-1) totAngMom_c + hrr0_v angMom_a bm angMom_c totAngMom_a (totAngMom_b-1) totAngMom_c in h1 +. h2 *. f - and hrr_v l angMom_a angMom_b angMom_c angMom_d + and hrr_v angMom_a angMom_b angMom_c angMom_d totAngMom_a totAngMom_b totAngMom_c totAngMom_d = match (totAngMom_b, totAngMom_d) with | (_,0) -> if (totAngMom_b = 0) then - vrr_v l 0 angMom_a angMom_c totAngMom_a totAngMom_c - |> sum + begin + match vrr_v 0 angMom_a angMom_c totAngMom_a totAngMom_c with + | Some matrix -> Mat.sum matrix + | None -> 0. + end else - hrr0_v l angMom_a angMom_b angMom_c totAngMom_a totAngMom_b totAngMom_c + hrr0_v angMom_a angMom_b angMom_c totAngMom_a totAngMom_b totAngMom_c | (_,_) -> let (acx, acy, acz) = angMom_c and (adx, ady, adz) = angMom_d in @@ -275,21 +379,19 @@ let hvrr_two_e_vector (angMom_a, angMom_b, angMom_c, angMom_d) | _ -> (acx, acy, acz+1), (adx, ady, adz-1), 2 in let h1 = - hrr_v l angMom_a angMom_b cp dm totAngMom_a totAngMom_b (totAngMom_c+1) (totAngMom_d-1) + hrr_v angMom_a angMom_b cp dm totAngMom_a totAngMom_b (totAngMom_c+1) (totAngMom_d-1) and h2 = - hrr_v l angMom_a angMom_b angMom_c dm totAngMom_a totAngMom_b totAngMom_c (totAngMom_d-1) + hrr_v angMom_a angMom_b angMom_c dm totAngMom_a totAngMom_b totAngMom_c (totAngMom_d-1) in let f = (Coordinate.coord center_cd xyz) in h1 +. f *. h2 in - Array.init np (fun ab -> - hrr_v (ab+1) + hrr_v (angMom_a.(0),angMom_a.(1),angMom_a.(2)) (angMom_b.(0),angMom_b.(1),angMom_b.(2)) (angMom_c.(0),angMom_c.(1),angMom_c.(2)) (angMom_d.(0),angMom_d.(1),angMom_d.(2)) totAngMom_a totAngMom_b totAngMom_c totAngMom_d - ) |> sum @@ -350,7 +452,7 @@ 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 zm_array = Mat.init_rows np nq (fun i j -> + let zm_array = Mat.init_cols np nq (fun i j -> (** Screening on the product of coefficients *) try if (abs_float coef.{j,i} ) < 1.e-3*.cutoff then @@ -385,20 +487,21 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q let center_pq = - let result = - Array3.create Float64 fortran_layout 3 nq np - in - Array.iteri (fun ab shell_ab -> - Array.iteri (fun cd shell_cd -> - let cpq = - Coordinate.(shell_ab.ShellPair.center |- shell_cd.ShellPair.center) - in - result.{1,cd+1,ab+1} <- Coordinate.x cpq; - result.{2,cd+1,ab+1} <- Coordinate.y cpq; - result.{3,cd+1,ab+1} <- Coordinate.z cpq; - ) sq - ) sp; - result + Array.init 3 (fun xyz -> + Mat.init_cols nq np (fun cd ab -> + let shell_ab = sp.(ab-1) + and shell_cd = sq.(cd-1) + in + let cpq = + Coordinate.(shell_ab.ShellPair.center |- shell_cd.ShellPair.center) + in + match xyz with + | 0 -> Coordinate.x cpq; + | 1 -> Coordinate.y cpq; + | 2 -> Coordinate.z cpq; + | _ -> assert false + ) + ) in let zero_m_array = let result = @@ -411,9 +514,9 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q expo_inv_p.{ab+1} +. expo_inv_q.{cd+1} in let norm_pq_sq = - center_pq.{1,cd+1,ab+1} *. center_pq.{1,cd+1,ab+1} +. - center_pq.{2,cd+1,ab+1} *. center_pq.{2,cd+1,ab+1} +. - center_pq.{3,cd+1,ab+1} *. center_pq.{3,cd+1,ab+1} + center_pq.(0).{cd+1,ab+1} *. center_pq.(0).{cd+1,ab+1} +. + center_pq.(1).{cd+1,ab+1} *. center_pq.(1).{cd+1,ab+1} +. + center_pq.(2).{cd+1,ab+1} *. center_pq.(2).{cd+1,ab+1} in zero_m ~maxm ~expo_pq_inv ~norm_pq_sq @@ -444,8 +547,8 @@ let contracted_class_shell_pairs ~zero_m ?schwartz_p ?schwartz_q shell_p shell_q |> Array.concat in - let map_1d = Array.init maxm (fun _ -> Array.init np (fun _ -> Zmap.create (4*maxm))) - and map_2d = Array.init maxm (fun _ -> Array.init np (fun _ -> Zmap.create (Array.length class_indices))) + let map_1d = Array.init maxm (fun _ -> Zmap.create (4*maxm)) + and map_2d = 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 ->