TRIQS / nda

C++ library for multi-dimensional arrays
https://triqs.github.io/nda
Other
15 stars 12 forks source link

Issue with reshaping + matrix products in fortran layout #34

Closed jasonkaye closed 1 year ago

jasonkaye commented 1 year ago

I am trying to write a code that: (1) Takes an mxm matrix A and an mxnxp array B, both in Fortran layout, (2) Reshapes B to an mx(np) matrix, (3) Computes A B, (4) Reshapes the result to an mxnxp array and returns it.

I have copied below a minimal working example showing that this fails to produce the correct result, when compared with a simple code that simply performs the matrix-vector product of A with the columns B(:,i,j).

The code prints the difference of the results of the two calculations, as well as the two results themselves.

I am using commit d020afe692f3c2263707168a907e6926e57fb21f on a Linux machine.

#include "nda/nda.hpp"

using namespace nda;

static constexpr auto _ = nda::range::all;

// Test Fortran layout matrix-array products
// Array is reshaped to a matrix, matrix-matrix product is taken, and result is
// reshape back

// Simple matvec function for checking

nda::vector<dcomplex> mymatvec(nda::matrix_const_view<dcomplex,F_layout> a,
nda::vector_const_view<dcomplex> x) {

  int m = a.shape(0);
  int n = a.shape(1);

  auto b = nda::vector<dcomplex>(m);
  b = 0;

  for (int j = 0; j < n; ++j) {
    b += a(_,j)*x(j);
  }

  return b;
}

// Multiply mxm matrix A by mxnxp matrix B, combining last two indices of B
nda::array<dcomplex,3> arraymult(nda::matrix<dcomplex,F_layout> a, nda::array<dcomplex,3,F_layout> b) {

  int m = b.shape(0);
  int n = b.shape(1);
  int p = b.shape(2);

  auto brs = nda::reshaped_view(b,std::array<int,2>({m,n*p}));
  auto bmat = nda::matrix_const_view<dcomplex,F_layout>(brs);
  return nda::reshape(a * bmat,std::array<int,3>({m,n,p}));

}

int main() {

  int m = 2;
  int n = 2;
  int p = 3;

  auto a = nda::matrix<dcomplex,F_layout>(m,m);
  auto b = nda::array<dcomplex,3,F_layout>(m,n,p);

  // Fill in A and B with random entries

  for (int i = 0; i<m; ++i) {
    for (int j = 0; j<m; ++j) {
      a(i,j) = std::sin(100.0*(i+2*j)) + 1i*std::cos(100.0*(i+2*j));
    }
  }

  for (int k = 0; k<p; ++k) {
    for (int j = 0; j<n; ++j) {
      for (int i = 0; i<m; ++i) {
        b(i,j,k) = std::sin(100.0*(i+j+k)) + 1i*std::cos(100.0*(i+j+k));
      }
    }
  }

  // Compute products the old-fashioned way

  auto ctrue = nda::array<dcomplex,3,F_layout>(m,n,p);
  for (int k = 0; k<p; ++k) {
    for (int j = 0; j<n; ++j) {
      ctrue(_,j,k) = mymatvec(a,b(_,j,k));
    }
  }

  // Compute products with my function

  auto c = arraymult(a,b);

  // Check agreement

  PRINT(max_element(abs(c-ctrue)));

  PRINT(c);
  PRINT(ctrue);
}
Wentzell commented 1 year ago

Thank you @jasonkaye for pointing this out!

I have addressed this issue on the DEV_CUDA_LINALG branch in commit 3f9efb7.