Merge pull request #11 from kevinsung/rust-modules

split rust code into modules

src/contract/diag_coulomb.rs

@@ -0,0 +1,168 @@
+// (C) Copyright IBM 2023
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use ndarray::Array;
+use ndarray::Zip;
+use numpy::Complex64;
+use numpy::PyReadonlyArray1;
+use numpy::PyReadonlyArray2;
+use numpy::PyReadwriteArray2;
+use pyo3::prelude::*;
+
+/// Contract a diagonal Coulomb operator into a buffer.
+#[pyfunction]
+pub fn contract_diag_coulomb_into_buffer_num_rep(
+    vec: PyReadonlyArray2<Complex64>,
+    mat: PyReadonlyArray2<f64>,
+    norb: usize,
+    mat_alpha_beta: PyReadonlyArray2<f64>,
+    occupations_a: PyReadonlyArray2<usize>,
+    occupations_b: PyReadonlyArray2<usize>,
+    mut out: PyReadwriteArray2<Complex64>,
+) {
+    let vec = vec.as_array();
+    let mat = mat.as_array();
+    let mat_alpha_beta = mat_alpha_beta.as_array();
+    let occupations_a = occupations_a.as_array();
+    let occupations_b = occupations_b.as_array();
+    let mut out = out.as_array_mut();
+
+    let shape = vec.shape();
+    let dim_a = shape[0];
+    let dim_b = shape[1];
+    let n_alpha = occupations_a.shape()[1];
+    let n_beta = occupations_b.shape()[1];
+
+    let mut alpha_coeffs = Array::zeros(dim_a);
+    let mut beta_coeffs = Array::zeros(dim_b);
+    let mut coeff_map = Array::zeros((dim_a, norb));
+
+    Zip::from(&mut beta_coeffs)
+        .and(occupations_b.rows())
+        .par_for_each(|val, orbs| {
+            let mut coeff = Complex64::new(0.0, 0.0);
+            for j in 0..n_beta {
+                let orb_1 = orbs[j];
+                for k in j..n_beta {
+                    let orb_2 = orbs[k];
+                    coeff += mat[(orb_1, orb_2)];
+                }
+            }
+            *val = coeff;
+        });
+
+    Zip::from(&mut alpha_coeffs)
+        .and(occupations_a.rows())
+        .and(coeff_map.rows_mut())
+        .par_for_each(|val, orbs, mut row| {
+            let mut coeff = Complex64::new(0.0, 0.0);
+            for j in 0..n_alpha {
+                let orb_1 = orbs[j];
+                row += &mat_alpha_beta.row(orb_1);
+                for k in j..n_alpha {
+                    let orb_2 = orbs[k];
+                    coeff += mat[(orb_1, orb_2)];
+                }
+            }
+            *val = coeff;
+        });
+
+    Zip::from(vec.rows())
+        .and(out.rows_mut())
+        .and(&alpha_coeffs)
+        .and(coeff_map.rows())
+        .par_for_each(|source, target, alpha_coeff, coeff_map| {
+            Zip::from(source)
+                .and(target)
+                .and(&beta_coeffs)
+                .and(occupations_b.rows())
+                .for_each(|source, target, beta_coeff, orbs| {
+                    let mut coeff = *alpha_coeff + *beta_coeff;
+                    orbs.for_each(|&orb| coeff += coeff_map[orb]);
+                    *target += coeff * source;
+                })
+        });
+}
+
+/// Contract a diagonal Coulomb operator into a buffer, Z representation.
+#[pyfunction]
+pub fn contract_diag_coulomb_into_buffer_z_rep(
+    vec: PyReadonlyArray2<Complex64>,
+    mat: PyReadonlyArray2<f64>,
+    norb: usize,
+    mat_alpha_beta: PyReadonlyArray2<f64>,
+    strings_a: PyReadonlyArray1<i64>,
+    strings_b: PyReadonlyArray1<i64>,
+    mut out: PyReadwriteArray2<Complex64>,
+) {
+    let vec = vec.as_array();
+    let mat = mat.as_array();
+    let mat_alpha_beta = mat_alpha_beta.as_array();
+    let strings_a = strings_a.as_array();
+    let strings_b = strings_b.as_array();
+    let mut out = out.as_array_mut();
+
+    let shape = vec.shape();
+    let dim_a = shape[0];
+    let dim_b = shape[1];
+
+    let mut alpha_coeffs = Array::zeros(dim_a);
+    let mut beta_coeffs = Array::zeros(dim_b);
+    let mut coeff_map = Array::zeros((dim_a, norb));
+
+    Zip::from(&mut beta_coeffs)
+        .and(strings_b)
+        .par_for_each(|val, str0| {
+            let mut coeff = Complex64::new(0.0, 0.0);
+            for j in 0..norb {
+                let sign_j = if str0 >> j & 1 == 1 { -1 } else { 1 } as f64;
+                for k in j + 1..norb {
+                    let sign_k = if str0 >> k & 1 == 1 { -1 } else { 1 } as f64;
+                    coeff += sign_j * sign_k * mat[(j, k)];
+                }
+            }
+            *val = coeff;
+        });
+
+    Zip::from(&mut alpha_coeffs)
+        .and(strings_a)
+        .and(coeff_map.rows_mut())
+        .par_for_each(|val, str0, mut row| {
+            let mut coeff = Complex64::new(0.0, 0.0);
+            for j in 0..norb {
+                let sign_j = if str0 >> j & 1 == 1 { -1 } else { 1 } as f64;
+                row += &(sign_j * &mat_alpha_beta.row(j));
+                for k in j + 1..norb {
+                    let sign_k = if str0 >> k & 1 == 1 { -1 } else { 1 } as f64;
+                    coeff += sign_j * sign_k * mat[(j, k)];
+                }
+            }
+            *val = coeff;
+        });
+
+    Zip::from(vec.rows())
+        .and(out.rows_mut())
+        .and(&alpha_coeffs)
+        .and(coeff_map.rows())
+        .par_for_each(|source, target, alpha_coeff, coeff_map| {
+            Zip::from(source)
+                .and(target)
+                .and(&beta_coeffs)
+                .and(strings_b)
+                .for_each(|source, target, beta_coeff, str0| {
+                    let mut coeff = *alpha_coeff + *beta_coeff;
+                    for j in 0..norb {
+                        let sign_j = if str0 >> j & 1 == 1 { -1 } else { 1 } as f64;
+                        coeff += sign_j * coeff_map[j];
+                    }
+                    *target += 0.25 * coeff * source;
+                })
+        });
+}

src/contract/mod.rs

@@ -0,0 +1,12 @@
+// (C) Copyright IBM 2023
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+pub mod diag_coulomb;
+pub mod num_op_sum;

src/contract/num_op_sum.rs

@@ -0,0 +1,39 @@
+// (C) Copyright IBM 2023
+//
+// This code is licensed under the Apache License, Version 2.0. You may
+// obtain a copy of this license in the LICENSE.txt file in the root directory
+// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+//
+// Any modifications or derivative works of this code must retain this
+// copyright notice, and modified files need to carry a notice indicating
+// that they have been altered from the originals.
+
+use ndarray::Zip;
+use numpy::Complex64;
+use numpy::PyReadonlyArray1;
+use numpy::PyReadonlyArray2;
+use numpy::PyReadwriteArray2;
+use pyo3::prelude::*;
+
+/// Contract a sum of number operators into a buffer.
+#[pyfunction]
+pub fn contract_num_op_sum_spin_into_buffer(
+    vec: PyReadonlyArray2<Complex64>,
+    coeffs: PyReadonlyArray1<f64>,
+    occupations: PyReadonlyArray2<usize>,
+    mut out: PyReadwriteArray2<Complex64>,
+) {
+    let vec = vec.as_array();
+    let coeffs = coeffs.as_array();
+    let occupations = occupations.as_array();
+    let mut out = out.as_array_mut();
+
+    Zip::from(vec.rows())
+        .and(out.rows_mut())
+        .and(occupations.rows())
+        .par_for_each(|source, mut target, orbs| {
+            let mut coeff = Complex64::new(0.0, 0.0);
+            orbs.for_each(|&orb| coeff += coeffs[orb]);
+            target += &(coeff * &source);
+        });
+}