add iteration for sparse types

ext/coo/coo_def.c

@@ -227,4 +227,5 @@ VALUE coo_to_csc(VALUE self) {
 }
 
 
-#include "elementwise.c"
+#include "elementwise.c"
+#include "iteration.c"

ext/coo/iteration.c

@@ -0,0 +1,27 @@
+VALUE coo_each(VALUE self) {
+  coo_matrix* input;
+
+  TypedData_Get_Struct(self, coo_matrix, &coo_data_type, input);
+
+  for (size_t index = 0; index < input->count; index++) {
+    rb_yield(DBL2NUM(input->elements[index]));
+  }
+
+  return self;
+}
+
+VALUE coo_each_with_indices(VALUE self) {
+  coo_matrix* input;
+
+  TypedData_Get_Struct(self, coo_matrix, &coo_data_type, input);
+
+  VALUE* result_array = ALLOC_N(VALUE, input->ndims + 1);
+  for (size_t index = 0; index < input->count; index++) {
+    result_array[0] = DBL2NUM(input->elements[index]);
+    result_array[1] = INT2NUM(input->ia[index]);
+    result_array[2] = INT2NUM(input->ja[index]);
+    rb_yield(rb_ary_new4(input->ndims + 1, result_array));
+  }
+
+  return self;
+}

ext/csc/csc_def.c

@@ -185,4 +185,5 @@ VALUE csc_to_coo(VALUE self) {
 }
 
 
-#include "elementwise.c"
+#include "elementwise.c"
+#include "iteration.c"

ext/csc/iteration.c

@@ -0,0 +1,30 @@
+VALUE csc_each(VALUE self) {
+  csc_matrix* input;
+
+  TypedData_Get_Struct(self, csc_matrix, &csc_data_type, input);
+
+  for (size_t index = 0; index < input->count; index++) {
+    rb_yield(DBL2NUM(input->elements[index]));
+  }
+
+  return self;
+}
+
+VALUE csc_each_column(VALUE self) {
+  csc_matrix* input;
+
+  TypedData_Get_Struct(self, csc_matrix, &csc_data_type, input);
+
+  for (size_t j = 0; j < input->shape[1]; j++) {
+    size_t col_len = input->jp[j + 1] - input->jp[j];
+    VALUE* col_array = ALLOC_N(VALUE, col_len);
+
+    for(size_t i = input->jp[j]; i < input->jp[j + 1]; i++) {
+      col_array[i - input->jp[j]] = DBL2NUM(input->elements[i]);
+    }
+
+    rb_yield(rb_ary_new4(col_len, col_array));
+  }
+
+  return self;
+}

ext/csr/csr_def.c

@@ -185,4 +185,5 @@ VALUE csr_to_coo(VALUE self) {
 }
 
 
-#include "elementwise.c"
+#include "elementwise.c"
+#include "iteration.c"

ext/csr/iteration.c

@@ -0,0 +1,30 @@
+VALUE csr_each(VALUE self) {
+  csr_matrix* input;
+
+  TypedData_Get_Struct(self, csr_matrix, &csr_data_type, input);
+
+  for (size_t index = 0; index < input->count; index++) {
+    rb_yield(DBL2NUM(input->elements[index]));
+  }
+
+  return self;
+}
+
+VALUE csr_each_row(VALUE self) {
+  csr_matrix* input;
+
+  TypedData_Get_Struct(self, csr_matrix, &csr_data_type, input);
+
+  for (size_t i = 0; i < input->shape[0]; i++) {
+    size_t row_len = input->ip[i + 1] - input->ip[i];
+    VALUE* row_array = ALLOC_N(VALUE, row_len);
+
+    for(size_t j = input->ip[i]; j < input->ip[i + 1]; j++) {
+      row_array[j - input->ip[i]] = DBL2NUM(input->elements[j]);
+    }
+
+    rb_yield(rb_ary_new4(row_len, row_array));
+  }
+
+  return self;
+}

ext/dia/dia_def.c

@@ -98,4 +98,5 @@ VALUE dia_get_shape(VALUE self) {
 }
 
 
-#include "elementwise.c"
+#include "elementwise.c"
+#include "iteration.c"

ext/dia/iteration.c

@@ -0,0 +1,11 @@
+VALUE dia_each(VALUE self) {
+  dia_matrix* input;
+
+  TypedData_Get_Struct(self, dia_matrix, &dia_data_type, input);
+
+  for (size_t index = 0; index < input->count; index++) {
+    rb_yield(DBL2NUM(input->elements[index]));
+  }
+
+  return self;
+}

ext/ruby_sparse.c

@@ -166,6 +166,9 @@ VALUE coo_sin(VALUE self);
 VALUE coo_cos(VALUE self);
 VALUE coo_tan(VALUE self);
 
+VALUE coo_each(VALUE self);
+VALUE coo_each_with_indices(VALUE self);
+
 VALUE coo_from_nmatrix(VALUE self, VALUE nmat);
 
 // csr methods declaration
@@ -189,6 +192,9 @@ VALUE csr_sin(VALUE self);
 VALUE csr_cos(VALUE self);
 VALUE csr_tan(VALUE self);
 
+VALUE csr_each(VALUE self);
+VALUE csr_each_row(VALUE self);
+
 VALUE csr_from_nmatrix(VALUE self, VALUE nmat);
 
 // csc methods declaration
@@ -212,6 +218,9 @@ VALUE csc_sin(VALUE self);
 VALUE csc_cos(VALUE self);
 VALUE csc_tan(VALUE self);
 
+VALUE csc_each(VALUE self);
+VALUE csc_each_column(VALUE self);
+
 VALUE csc_from_nmatrix(VALUE self, VALUE nmat);
 
 // dia methods declaration
@@ -231,6 +240,8 @@ VALUE dia_sin(VALUE self);
 VALUE dia_cos(VALUE self);
 VALUE dia_tan(VALUE self);
 
+VALUE dia_each(VALUE self);
+
 VALUE dia_from_nmatrix(VALUE self, VALUE nmat);
 
 
@@ -264,6 +275,9 @@ void Init_ruby_sparse() {
   rb_define_method(COO, "cos", coo_cos, 0);
   rb_define_method(COO, "tan", coo_tan, 0);
 
+  rb_define_method(COO, "each", coo_each, 0);
+  rb_define_method(COO, "each_with_indices", coo_each_with_indices, 0);
+
   //rb_define_singleton_method(COO, "from_nmatrix", coo_from_nmatrix, 1);
 
 
@@ -287,6 +301,9 @@ void Init_ruby_sparse() {
   rb_define_method(CSR, "cos", csr_cos, 0);
   rb_define_method(CSR, "tan", csr_tan, 0);
 
+  rb_define_method(CSR, "each", csr_each, 0);
+  rb_define_method(CSR, "each_row", csr_each_row, 0);
+
   //rb_define_singleton_method(CSR, "from_nmatrix", csr_from_nmatrix, 1);
 
 
@@ -310,6 +327,9 @@ void Init_ruby_sparse() {
   rb_define_method(CSC, "cos", csc_cos, 0);
   rb_define_method(CSC, "tan", csc_tan, 0);
 
+  rb_define_method(CSC, "each", csc_each, 0);
+  rb_define_method(CSC, "each_column", csc_each_column, 0);
+
   //rb_define_singleton_method(CSC, "from_nmatrix", csc_from_nmatrix, 1);
 
 
@@ -328,6 +348,8 @@ void Init_ruby_sparse() {
   rb_define_method(DIA, "cos", dia_cos, 0);
   rb_define_method(DIA, "tan", dia_tan, 0);
 
+  rb_define_method(DIA, "each", dia_each, 0);
+
   //rb_define_singleton_method(DIA, "from_nmatrix", dia_from_nmatrix, 1);
 }