Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

split rust code into modules #11

Merged
merged 1 commit into from
Sep 18, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
168 changes: 168 additions & 0 deletions src/contract/diag_coulomb.rs
Original file line number Diff line number Diff line change
@@ -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;
})
});
}
12 changes: 12 additions & 0 deletions src/contract/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
39 changes: 39 additions & 0 deletions src/contract/num_op_sum.rs
Original file line number Diff line number Diff line change
@@ -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);
});
}
Loading