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.

Sign up for GitHub

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

Add kron(): kronecker product for 2D matrices #1039

Closed
wants to merge 1 commit into from
Closed

Add kron(): kronecker product for 2D matrices #1039

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
24 changes: 24 additions & 0 deletions src/linalg/impl_linalg.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -691,6 +691,30 @@ unsafe fn general_mat_vec_mul_impl<A, S1, S2>(
}
}

/// Kronecker product of 2D matrices.
///
/// The kronecker product of a LxN matrix A and a MxR matrix B is a (L*M)x(N*R)
/// matrix K formed by the block multiplication A_ij * B.
pub fn kron<T>(a: &Array2<T>, b: &Array2<T>) -> Array2<T>
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To include this we'd need to use the same kind of argument type we usually use for reading an array - an &ArrayBase like other places that take any kind of array.

Minor style issues - I'd suggest it can have the slightly longer name kronecker. Also note that ndarray uses the element type parameter A by convention, so we need to keep that to keep it consistent. Thanks

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think kron would align with numpy, as a datapoint, just wanted to mention. It also is an operation usually done on 2 dimensional matrices, though I believe numpy supports higher dimensional kronecker products.

where
T: LinalgScalar,
{
let dimar = a.shape()[0];
let dimac = a.shape()[1];
let dimbr = b.shape()[0];
let dimbc = b.shape()[1];
let mut out = Array2::zeros((dimar * dimbr, dimac * dimbc));
for (mut chunk, elem) in out
.exact_chunks_mut((dimbr, dimbc))
.into_iter()
.zip(a.iter())
{
let v: Array2<T> = Array2::from_elem((dimbr, dimbc), *(elem)) * b;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here we'd like to avoid creating an intermediate array. We'd just use Zip instead of from_elem and assign. Other improvements are less important (alloc uninit array instead of zeroed, other ways to loop?). We would however like to replace the other zip with Zip because the iteration order of exact_chunks_mut is not well specified.

chunk.assign(&v);
}
out
}

#[inline(always)]
/// Return `true` if `A` and `B` are the same type
fn same_type<A: 'static, B: 'static>() -> bool {
Expand Down
1 change: 1 addition & 0 deletions src/linalg/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

pub use self::impl_linalg::general_mat_mul;
pub use self::impl_linalg::general_mat_vec_mul;
pub use self::impl_linalg::kron;
pub use self::impl_linalg::Dot;

mod impl_linalg;
63 changes: 63 additions & 0 deletions tests/oper.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
)]
#![cfg(feature = "std")]
use ndarray::linalg::general_mat_mul;
use ndarray::linalg::kron;
use ndarray::prelude::*;
use ndarray::{rcarr1, rcarr2};
use ndarray::{Data, LinalgScalar};
Expand Down Expand Up @@ -820,3 +821,65 @@ fn vec_mat_mul() {
}
}
}

#[test]
fn kron_square_f64() {
let a = arr2(&[[1.0, 0.0], [0.0, 1.0]]);
let b = arr2(&[[0.0, 1.0], [1.0, 0.0]]);

assert_eq!(
kron(&a, &b),
arr2(&[
[0.0, 1.0, 0.0, 0.0],
[1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
[0.0, 0.0, 1.0, 0.0]
]),
);

assert_eq!(
kron(&b, &a),
arr2(&[
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0]
]),
)
}

#[test]
fn kron_square_i64() {
let a = arr2(&[[1, 0], [0, 1]]);
let b = arr2(&[[0, 1], [1, 0]]);

assert_eq!(
kron(&a, &b),
arr2(&[[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]),
);

assert_eq!(
kron(&b, &a),
arr2(&[[0, 0, 1, 0], [0, 0, 0, 1], [1, 0, 0, 0], [0, 1, 0, 0]]),
)
}

#[test]
fn kron_i64() {
let a = arr2(&[[1, 0]]);
let b = arr2(&[[0, 1], [1, 0]]);
let r = arr2(&[[0, 1, 0, 0], [1, 0, 0, 0]]);
assert_eq!(kron(&a, &b), r);

let a = arr2(&[[1, 0], [0, 0], [0, 1]]);
let b = arr2(&[[0, 1], [1, 0]]);
let r = arr2(&[
[0, 1, 0, 0],
[1, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 1],
[0, 0, 1, 0],
]);
assert_eq!(kron(&a, &b), r);
}