fn_trimat_ind.hpp 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. // Copyright 2008-2016 Conrad Sanderson (http://conradsanderson.id.au)
  2. // Copyright 2008-2016 National ICT Australia (NICTA)
  3. //
  4. // Licensed under the Apache License, Version 2.0 (the "License");
  5. // you may not use this file except in compliance with the License.
  6. // You may obtain a copy of the License at
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. // ------------------------------------------------------------------------
  15. //! \addtogroup fn_trimat_ind
  16. //! @{
  17. arma_warn_unused
  18. inline
  19. uvec
  20. trimatu_ind(const SizeMat& s, const sword k = 0)
  21. {
  22. arma_extra_debug_sigprint();
  23. const uword n_rows = s.n_rows;
  24. const uword n_cols = s.n_cols;
  25. const uword row_offset = (k < 0) ? uword(-k) : uword(0);
  26. const uword col_offset = (k > 0) ? uword( k) : uword(0);
  27. arma_debug_check( ((row_offset > 0) && (row_offset >= n_rows)) || ((col_offset > 0) && (col_offset >= n_cols)), "trimatu_ind(): requested diagonal is out of bounds" );
  28. const uword N = (std::min)(n_rows - row_offset, n_cols - col_offset);
  29. uvec tmp(n_rows * n_cols); // worst case scenario
  30. uword* tmp_mem = tmp.memptr();
  31. uword count = 0;
  32. for(uword i=0; i < n_cols; ++i)
  33. {
  34. const uword col = i + col_offset;
  35. if(i < N)
  36. {
  37. const uword end_row = i + row_offset;
  38. const uword index_offset = (n_rows * col);
  39. for(uword row=0; row <= end_row; ++row)
  40. {
  41. tmp_mem[count] = index_offset + row;
  42. ++count;
  43. }
  44. }
  45. else
  46. {
  47. if(col < n_cols)
  48. {
  49. const uword index_offset = (n_rows * col);
  50. for(uword row=0; row < n_rows; ++row)
  51. {
  52. tmp_mem[count] = index_offset + row;
  53. ++count;
  54. }
  55. }
  56. }
  57. }
  58. uvec out;
  59. out.steal_mem_col(tmp, count);
  60. return out;
  61. }
  62. arma_warn_unused
  63. inline
  64. uvec
  65. trimatl_ind(const SizeMat& s, const sword k = 0)
  66. {
  67. arma_extra_debug_sigprint();
  68. const uword n_rows = s.n_rows;
  69. const uword n_cols = s.n_cols;
  70. const uword row_offset = (k < 0) ? uword(-k) : uword(0);
  71. const uword col_offset = (k > 0) ? uword( k) : uword(0);
  72. arma_debug_check( ((row_offset > 0) && (row_offset >= n_rows)) || ((col_offset > 0) && (col_offset >= n_cols)), "trimatl_ind(): requested diagonal is out of bounds" );
  73. const uword N = (std::min)(n_rows - row_offset, n_cols - col_offset);
  74. uvec tmp(n_rows * n_cols); // worst case scenario
  75. uword* tmp_mem = tmp.memptr();
  76. uword count = 0;
  77. for(uword col=0; col < col_offset; ++col)
  78. {
  79. const uword index_offset = (n_rows * col);
  80. for(uword row=0; row < n_rows; ++row)
  81. {
  82. tmp_mem[count] = index_offset + row;
  83. ++count;
  84. }
  85. }
  86. for(uword i=0; i<N; ++i)
  87. {
  88. const uword start_row = i + row_offset;
  89. const uword col = i + col_offset;
  90. const uword index_offset = (n_rows * col);
  91. for(uword row=start_row; row < n_rows; ++row)
  92. {
  93. tmp_mem[count] = index_offset + row;
  94. ++count;
  95. }
  96. }
  97. uvec out;
  98. out.steal_mem_col(tmp, count);
  99. return out;
  100. }
  101. //! @}