Introduce matrix and color-vector drawing foundations

Author: kojix2

spec/matrix_spec.cr

@@ -0,0 +1,74 @@
+require "./spec_helper"
+
+describe UnicodePlot::MatrixView do
+  it "stores rectangular matrix dimensions" do
+    m = UnicodePlot::MatrixView.new([
+      [1.0, 2.0, 3.0],
+      [4.0, 5.0, 6.0],
+    ])
+
+    m.nrows.should eq(2)
+    m.ncols.should eq(3)
+    m[0, 1].should eq(2.0)
+    m[1, 2].should eq(6.0)
+  end
+
+  it "rejects jagged matrices" do
+    expect_raises(ArgumentError, /same length/) do
+      UnicodePlot::MatrixView.new([
+        [1.0, 2.0],
+        [3.0],
+      ])
+    end
+  end
+
+  it "allows empty matrices at the view level" do
+    m = UnicodePlot::MatrixView.new([] of Array(Float64))
+
+    m.nrows.should eq(0)
+    m.ncols.should eq(0)
+    m.empty?.should be_true
+  end
+end
+
+describe ".matrix_axis_coords" do
+  it "uses one-based coordinates when fact is nil" do
+    UnicodePlot.matrix_axis_coords(3, nil, 0.0).should eq([1.0, 2.0, 3.0])
+  end
+
+  it "applies offset to one-based coordinates when fact is nil" do
+    UnicodePlot.matrix_axis_coords(3, nil, 10.0).should eq([11.0, 12.0, 13.0])
+  end
+
+  it "uses zero-based scaled coordinates when fact is given" do
+    UnicodePlot.matrix_axis_coords(3, 0.5, 10.0).should eq([10.0, 10.5, 11.0])
+  end
+end
+
+describe ".matrix_extrema_finite" do
+  it "returns min and max of finite values" do
+    m = UnicodePlot::MatrixView.new([
+      [1.0, Float64::NAN],
+      [3.0, -2.0],
+    ])
+
+    UnicodePlot.matrix_extrema_finite(m).should eq({-2.0, 3.0})
+  end
+
+  it "ignores infinities" do
+    m = UnicodePlot::MatrixView.new([
+      [Float64::INFINITY, 2.0],
+      [-Float64::INFINITY, 5.0],
+    ])
+
+    UnicodePlot.matrix_extrema_finite(m).should eq({2.0, 5.0})
+  end
+
+  it "returns nil when there are no finite values" do
+    m = UnicodePlot::MatrixView.new([
+      [Float64::NAN, Float64::INFINITY],
+    ])
+
+    UnicodePlot.matrix_extrema_finite(m).should be_nil
+  end
+end

spec/unicode_plot_spec.cr

@@ -420,6 +420,13 @@ describe UnicodePlot do
       z = [[Float64::NAN, Float64::NAN]]
       p = UnicodePlot.heatmap(z)
       p.to_s.should be_a(String)
+      p.colormap.lim.should eq({0.0, 1.0})
+    end
+
+    it "uses finite extrema for colormap limits" do
+      z = [[Float64::INFINITY, 1.0], [2.0, -Float64::INFINITY]]
+      p = UnicodePlot.heatmap(z)
+      p.colormap.lim.should eq({1.0, 2.0})
     end
 
     it "accepts Int matrix via generic numeric overload" do
@@ -437,6 +444,12 @@ describe UnicodePlot do
         UnicodePlot.heatmap([[1.0, 2.0], [3.0, 4.0]], colormap: :mystery)
       end
     end
+
+    it "rejects jagged matrix input" do
+      expect_raises(ArgumentError, /same length/) do
+        UnicodePlot.heatmap([[1.0, 2.0], [3.0]])
+      end
+    end
   end
 
   describe "stairs" do
@@ -492,6 +505,32 @@ describe UnicodePlot do
 
       c.to_s.should be_a(String)
     end
+
+    it "rejects point color vectors with wrong length" do
+      p = UnicodePlot.scatterplot([1.0, 2.0], [1.0, 2.0])
+
+      expect_raises(ArgumentError, /colors/) do
+        p.points!(
+          [1.0, 2.0],
+          [1.0, 2.0],
+          [UnicodePlot.plot_color(:red)],
+          p.canvas.blend?
+        )
+      end
+    end
+
+    it "rejects segment color vectors with wrong length" do
+      p = UnicodePlot.lineplot([1.0, 2.0, 3.0], [1.0, 2.0, 3.0])
+
+      expect_raises(ArgumentError, /segment_colors/) do
+        p.lines!(
+          [1.0, 2.0, 3.0],
+          [1.0, 2.0, 3.0],
+          [UnicodePlot.plot_color(:red)],
+          p.canvas.blend?
+        )
+      end
+    end
   end
 
   describe "BlockCanvas" do

src/unicode_plot.cr

@@ -1,6 +1,7 @@
 require "./unicode_plot/terminal_size"
 require "./unicode_plot/common"
 require "./unicode_plot/color"
+require "./unicode_plot/matrix"
 require "./unicode_plot/graphics"
 require "./unicode_plot/canvas"
 require "./unicode_plot/canvas/braille_canvas"

src/unicode_plot/canvas.cr

@@ -190,6 +190,24 @@ module UnicodePlot
       self
     end
 
+    # Plot a vector of points with per-point colors
+    def points!(xs : Array(Float64), ys : Array(Float64), colors : Array(UInt32), blend : Bool) : self
+      unless xs.size == ys.size
+        raise ArgumentError.new("xs and ys must have the same length")
+      end
+
+      unless xs.size == colors.size
+        raise ArgumentError.new("colors must have the same length as xs and ys")
+      end
+
+      xs.each_with_index do |x, i|
+        y = ys[i]
+        next unless x.finite? && y.finite?
+        points!(x, y, colors[i], blend)
+      end
+      self
+    end
+
     # Digital differential analyser (DDA) line drawing
     def lines!(
       x1 : Float64, y1 : Float64,
@@ -246,6 +264,26 @@ module UnicodePlot
       self
     end
 
+    # Vector line drawing (polyline) with per-segment colors
+    def lines!(xs : Array(Float64), ys : Array(Float64), segment_colors : Array(UInt32), blend : Bool) : self
+      unless xs.size == ys.size
+        raise ArgumentError.new("xs and ys must have the same length")
+      end
+
+      expected = xs.size > 0 ? xs.size - 1 : 0
+      unless segment_colors.size == expected
+        raise ArgumentError.new("segment_colors must have xs.size - 1 elements")
+      end
+
+      (1...xs.size).each do |i|
+        x, y = xs[i], ys[i]
+        xm1, ym1 = xs[i - 1], ys[i - 1]
+        next unless x.finite? && y.finite? && xm1.finite? && ym1.finite?
+        lines!(xm1, ym1, x, y, segment_colors[i - 1], blend)
+      end
+      self
+    end
+
     # Annotate a character at canvas coords (x, y)
     def annotate!(x : Float64, y : Float64, char : Char, color : UInt32, blend : Bool) : self
       return self unless valid_x?(x) && valid_y?(y)

src/unicode_plot/interface/heatmap.cr

@@ -133,25 +133,15 @@ z : Array(Array(Float64)),
               width : Int32? = nil,
               fix_ar : Bool = false,
               out_stream : IO? = nil,) : Plot
-    raise ArgumentError.new("z must not be empty") if z.empty?
+    matrix = MatrixView.new(z)
+    raise ArgumentError.new("z must not be empty") if matrix.empty?
 
-    data_nrows = z.size
-    data_ncols = z.empty? ? 0 : z[0].size
+    data_nrows = matrix.nrows
+    data_ncols = matrix.ncols
 
     # Build Y and X coordinate arrays (mirrors Julia's Y/X computation).
-    y_arr = if yfact.nil?
-              Array.new(data_nrows) { |i| (i + 1).to_f + yoffset }
-            else
-              f = yfact || 1.0
-              Array.new(data_nrows) { |i| (i.to_f * f + yoffset).as(Float64) }
-            end
-
-    x_arr = if xfact.nil?
-              Array.new(data_ncols) { |i| (i + 1).to_f + xoffset }
-            else
-              f = xfact || 1.0
-              Array.new(data_ncols) { |i| (i.to_f * f + xoffset).as(Float64) }
-            end
+    y_arr = matrix_axis_coords(data_nrows, yfact, yoffset)
+    x_arr = matrix_axis_coords(data_ncols, xfact, xoffset)
 
     # autolims: {0,0} means auto → use full data range.
     eff_ylim = if ylim == {0.0, 0.0}
@@ -229,17 +219,17 @@ z : Array(Array(Float64)),
     )
 
     # zmin/zmax from data or explicit zlim.
+    submatrix = MatrixView.new(z_sub)
     has_extrema = false
     zmin, zmax = if zlim != {0.0, 0.0}
                    has_extrema = true
                    {zlim[0], zlim[1]}
                  else
-                   finite_vals = z_sub.each.flat_map(&.each).select(&.finite?).to_a
-                   if finite_vals.empty?
-                     {0.0, 1.0}
-                   else
+                   if extrema = matrix_extrema_finite(submatrix)
                      has_extrema = true
-                     {finite_vals.min, finite_vals.max}
+                     extrema
+                   else
+                     {0.0, 1.0}
                    end
                  end
     # Colorbar auto-detection (height >= 7 matches Julia's `height < 7` threshold).

src/unicode_plot/matrix.cr

@@ -0,0 +1,76 @@
+module UnicodePlot
+  struct MatrixView(T)
+    getter rows : Array(Array(T))
+    getter nrows : Int32
+    getter ncols : Int32
+
+    def initialize(@rows : Array(Array(T)))
+      @nrows = @rows.size
+      @ncols = @rows.empty? ? 0 : @rows[0].size
+
+      @rows.each_with_index do |row, i|
+        unless row.size == @ncols
+          raise ArgumentError.new(
+            "matrix rows must all have the same length; row #{i} has #{row.size}, expected #{@ncols}"
+          )
+        end
+      end
+    end
+
+    def empty? : Bool
+      @nrows == 0 || @ncols == 0
+    end
+
+    def [](row : Int32, col : Int32) : T
+      @rows[row][col]
+    end
+
+    def each_cell(& : Int32, Int32, T ->)
+      @rows.each_with_index do |row_values, row|
+        row_values.each_with_index do |value, col|
+          yield row, col, value
+        end
+      end
+    end
+  end
+
+  def matrix_axis_coords(
+    n : Int32,
+    fact : Float64?,
+    offset : Float64 = 0.0,
+  ) : Array(Float64)
+    if fact.nil?
+      Array.new(n) { |i| (i + 1).to_f64 + offset }
+    else
+      f = fact || 1.0
+      Array.new(n) { |i| i.to_f64 * f + offset }
+    end
+  end
+
+  def matrix_extrema_finite(matrix : MatrixView(T)) : {Float64, Float64}? forall T
+    {% unless T <= Number %}
+      {% raise "matrix_extrema_finite requires numeric matrix elements" %}
+    {% end %}
+
+    min_value = nil.as(Float64?)
+    max_value = nil.as(Float64?)
+
+    matrix.each_cell do |_row, _col, value|
+      v = value.to_f64
+      next unless v.finite?
+
+      if min_value.nil?
+        min_value = v
+        max_value = v
+      else
+        current_min = min_value || v
+        current_max = max_value || v
+        min_value = v if v < current_min
+        max_value = v if v > current_max
+      end
+    end
+
+    return if min_value.nil?
+    {min_value || 0.0, max_value || 0.0}
+  end
+end

src/unicode_plot/plot.cr

@@ -173,11 +173,21 @@ module UnicodePlot
       self
     end
 
+    def lines!(xs : Array(Float64), ys : Array(Float64), segment_colors : Array(UInt32), blend : Bool) : self
+      canvas.lines!(xs, ys, segment_colors, blend)
+      self
+    end
+
     def points!(xs : Array(Float64), ys : Array(Float64), color : UInt32, blend : Bool) : self
       canvas.points!(xs, ys, color, blend)
       self
     end
 
+    def points!(xs : Array(Float64), ys : Array(Float64), colors : Array(UInt32), blend : Bool) : self
+      canvas.points!(xs, ys, colors, blend)
+      self
+    end
+
     def annotate!(x : Float64, y : Float64, char : Char, color : UInt32, blend : Bool) : self
       canvas.annotate!(x, y, char, color, blend)
       self