explode-collapse.cc 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. // -*- mode: c++; coding: utf-8 -*-
  2. // ra-ra/examples - Demo collapse and explode operations.
  3. // (c) Daniel Llorens - 2016
  4. // This library is free software; you can redistribute it and/or modify it under
  5. // the terms of the GNU Lesser General Public License as published by the Free
  6. // Software Foundation; either version 3 of the License, or (at your option) any
  7. // later version.
  8. #include <iostream>
  9. #include "ra/ra.hh"
  10. using std::cout, std::endl, std::flush;
  11. template <class T, int rank> using Array = ra::Big<T, rank>;
  12. template <class T, int ... sizes> using Small = ra::Small<T, sizes ...>;
  13. // FIXME remove when stdlib has it
  14. template <class T> struct std::formatter<std::complex<T>>: ra::ostream_formatter {};
  15. int main()
  16. {
  17. // These operations let you view an array of rank n as an array of rank (n-k) of
  18. // subarrays of rank k (explode) or the opposite (collapse). However, there is a
  19. // caveat: the subarrays must be compact (stored contiguously) and the sizes of
  20. // the subarrays must be known statically. In effect, the subarray must be of
  21. // type ra::Small. For example:
  22. cout << "\nexplode\n" << endl;
  23. {
  24. Array<int, 2> A({2, 3}, ra::_0 - ra::_1);
  25. auto B = ra::explode<Small<int, 3>>(A);
  26. cout << "B(0): " << B(0) << endl; // [0 -1 -2], note the static size
  27. cout << "B(1): " << B(1) << endl; // [1, 0, -1], note the static size
  28. B(1) = 9;
  29. cout << "A after: " << A << endl;
  30. }
  31. // The collapsed/exploded arrays are views into the source. This is similar to a
  32. // Blitz++ feature called 'multicomponents'. For instance, suppose we have an
  33. // array of (static-sized) [x, y, z] vectors and wish to initialize the x, y, z
  34. // components separately:
  35. cout << "\ncollapse\n" << endl;
  36. {
  37. Array<Small<int, 3>, 1> A = {{0, -1, -2}, {1, 0, -1}, {2, 1, -1}};
  38. auto B = ra::collapse<int>(A);
  39. cout << "B before: " << B << endl;
  40. B(ra::all, 0) = 9; // initialize all x components
  41. B(ra::all, 1) = 7; // all y
  42. B(ra::all, 2) = 3; // all z
  43. cout << "B after: " << B << endl;
  44. cout << "A after: " << A << endl;
  45. }
  46. // collapse/explode can also be used to make views into the real and imag parts
  47. // of a complex array. To use explode in this way, the last axis must be
  48. // contiguous and of size 2, just as when using Small<>.
  49. cout << "\nreal/imag views\n" << endl;
  50. {
  51. using complex = std::complex<double>;
  52. Array<complex, 2> A({2, 2}, 0.);
  53. std::print(cout, "{}", complex {1, 0});
  54. // construct array views into real and imag parts of A
  55. auto B = ra::collapse<double>(A);
  56. auto Bre = B(ra::all, ra::all, 0);
  57. auto Bim = B(ra::all, ra::all, 1);
  58. // set real and imag parts of A
  59. Bre = ra::_0-ra::_1;
  60. Bim = -3;
  61. cout << "Aa: " << A << endl;
  62. // of course, you could also do
  63. imag_part(A) = 7.;
  64. cout << "Ab: " << A << endl;
  65. // which doesn't construct an array view as collapse() does, but relies on the
  66. // expression template mechanism instead.
  67. // Constructing an array view is useful for example when you need to pass a
  68. // pointer to an external library. But be aware of the steps!
  69. double * re_ptr = Bre.data();
  70. double * im_ptr = Bim.data();
  71. re_ptr[0] = 77;
  72. im_ptr[0] = 99;
  73. cout << "Ac: " << A << endl;
  74. }
  75. return 0;
  76. }