add cluster support to gen-data.rb, show centroids on generated png

1 files changed, 9 insertions, 4 deletions

diff --git a/gen-data.rb b/gen-data.rb
index e010ed1..88fe9c4 100755
--- a/gen-data.rb
+++ b/gen-data.rb
@@ -1,18 +1,23 @@
 #!/usr/bin/env ruby
 
-# get number of rows 
-NUM_ROWS = (ARGV.shift || 100).to_i
+# get number of rows and number of clusters
+NUM_ROWS = (ARGV.shift || 1000).to_i
+NUM_CLUSTERS = (ARGV.shift || 4).to_i
 
 # data shape (num_floats, num_ints)
 SHAPE = [2, 0]
 
 # cluster centers
-Q = [[-1, -1], [-1, +1], [+1, -1], [+1, +1]].map { |c|
+Q = [
+  [-1, -1], [-1, +1], [+1, -1], [+1, +1],
+  [0, 0],
+  [-1, 0], [0, +1], [0, -1], [+1, 0],
+].map { |c|
   c.map { |v| 0.5 + 0.25 * v }
 }
 
 def noise(y, x, scale)
-  Q[y % Q.size][x] + scale * rand * Math.cos(Math::PI * rand)
+  Q[y % NUM_CLUSTERS][x] + scale * rand * Math.cos(Math::PI * rand)
 end
 
 # print output