1package chart23import (4 "image"5 "image/color"6 "math"7 "strconv"89 "github.com/charmbracelet/x/ansi"1011 "src.rest/cat/012/internal/sheet"12)1314// scatterPlan lays out a scatter chart as Sheets does: the first series15// is the X values and every other one a series of Y values plotted16// against them, each point a dot, with an optional linear trend line per17// series. The value axis on the left is Y, with the chart's axis options;18// X runs along the bottom.19//20// 60 ┤ ⠛ ⠛⠛21// 40 ┤ ⠛⠛ ⠛22// 20 ┤⠛23// 0 ┼──┬──┬──┬24// 0 2 4 625type scatterPlan struct {26 xs []float6427 ys []sheet.ChartSeries28 trend bool29 grid bool30 msg string31 xsc scale32 ysc scale33 labelW int34 axisX int35 axisRow int36 tickRow int37 plot image.Rectangle38}3940// scatterSeries splits d into X values and Y series: the first series is41// X when there are two or more, and otherwise the categories are, read as42// numbers or else counted from 1.43func scatterSeries(d sheet.ChartData) ([]float64, []sheet.ChartSeries) {44 if len(d.Series) >= 2 {45 return d.Series[0].Values, d.Series[1:]46 }47 xs := make([]float64, len(d.Categories))48 for i, c := range d.Categories {49 x, err := strconv.ParseFloat(c, 64)50 if err != nil {51 x = float64(i + 1)52 }53 xs[i] = x54 }55 return xs, d.Series56}5758func newScatterPlan(d sheet.ChartData, w, h int, o sheet.ChartOptions) *scatterPlan {59 d = prepare(d, o, false)60 xs, ys := scatterSeries(d)61 p := &scatterPlan{xs: xs, ys: ys, trend: o.Trend, grid: !o.NoGrid}62 xlo, xhi, okX := finite([]sheet.ChartSeries{{Values: xs}})63 ylo, yhi, okY := finite(ys)64 if !okX || !okY {65 p.msg = "No numbers to chart"66 return p67 }68 p.tickRow = h - 169 p.axisRow = h - 270 if p.axisRow < 2 {71 p.msg = "Too small to chart"72 return p73 }74 p.ysc = newAxis(ylo, yhi, p.axisRow, 2, 6, d.Format, o)75 for j := 0; j <= p.ysc.n; j++ {76 p.labelW = max(p.labelW, ansi.StringWidth(p.ysc.label(p.ysc.tick(j))))77 }78 p.axisX = p.labelW + 179 x0 := p.axisX + 180 labelW := 181 for range 2 {82 plotW := w - x0 - (labelW+1)/283 if plotW < 2 {84 p.msg = "Too small to chart"85 return p86 }87 p.xsc = newScale(xlo, xhi, plotW, labelW+2, 6, d.Format)88 for j := 0; j <= p.xsc.n; j++ {89 labelW = max(labelW, ansi.StringWidth(p.xsc.label(p.xsc.tick(j))))90 }91 }92 p.plot = image.Rect(x0, 0, x0+p.xsc.cells(), p.axisRow)93 return p94}9596// points calls fn with each point of every series, in cells from the97// plot's bottom left.98func (p *scatterPlan) points(fn func(j int, x, y float64)) {99 for j, s := range p.ys {100 for i, x := range p.xs {101 y := value(s, i)102 if math.IsNaN(y) || math.IsNaN(x) || math.IsInf(x, 0) {103 continue104 }105 fn(j, p.xsc.at(x), p.ysc.at(y))106 }107 }108}109110// fit returns the least squares line y = a + b*x through series j's111// points, and false with fewer than two distinct X values.112func (p *scatterPlan) fit(j int) (a, b float64, ok bool) {113 var n, sx, sy, sxx, sxy float64114 for i, x := range p.xs {115 y := value(p.ys[j], i)116 if math.IsNaN(y) || math.IsNaN(x) || math.IsInf(x, 0) {117 continue118 }119 n++120 sx, sy, sxx, sxy = sx+x, sy+y, sxx+x*x, sxy+x*y121 }122 d := n*sxx - sx*sx123 if n < 2 || math.Abs(d) < 1e-12*max(1, sxx*n) {124 return 0, 0, false125 }126 b = (n*sxy - sx*sy) / d127 return (sy - b*sx) / n, b, true128}129130// trendPoints calls fn with points along series j's trend line, in cells131// from the plot's bottom left, steps per cell across the plot; ok is132// false where the line leaves the value axis.133func (p *scatterPlan) trendPoints(j int, steps float64, fn func(x, y float64, ok bool)) {134 a, b, ok := p.fit(j)135 if !ok {136 return137 }138 cells := float64(p.xsc.cells())139 n := int(cells * steps)140 for k := 0; k <= n; k++ {141 x := cells * float64(k) / float64(max(n, 1))142 v := a + b*(p.xsc.lo+x/cells*(p.xsc.hi-p.xsc.lo))143 y := p.ysc.pos(v)144 fn(x, y, !math.IsNaN(y) && y >= 0 && y <= float64(p.ysc.cells()))145 }146}147148func (p *scatterPlan) draw(g *Grid, o Options) {149 drawValueAxis(g, p.ysc, p.labelW, p.axisX, p.axisRow, p.plot.Min.X, p.plot.Max.X)150 drawXTicks(g, p.xsc, p.axisX, p.plot.Min.X, p.axisRow, p.tickRow)151 if o.Image {152 return153 }154 if p.grid {155 hGridlines(g, p.ysc, p.axisRow, p.plot)156 }157 w, h := p.plot.Dx(), p.axisRow158 b := newBraille(w, h)159 // Dots: two per cell across, four down, from the top left.160 dot := func(x, y float64) (int, int) {161 dx := int(math.Round(x / float64(max(p.xsc.cells(), 1)) * float64(w*2-1)))162 dy := h*4 - 1 - int(math.Round(y/float64(max(p.ysc.cells(), 1))*float64(p.ysc.cells()*4-1)))163 return min(max(dx, 0), w*2-1), min(max(dy, 0), h*4-1)164 }165 if p.trend {166 for j := range p.ys {167 px, py, have := 0, 0, false168 p.trendPoints(j, 2, func(x, y float64, ok bool) {169 if !ok {170 have = false171 return172 }173 dx, dy := dot(x, y)174 if have {175 b.line(px, py, dx, dy, j)176 }177 px, py, have = dx, dy, true178 })179 }180 }181 p.points(func(j int, x, y float64) {182 dx, dy := dot(x, y)183 // A point is two dots by two, kept inside the plot.184 dx, dy = min(dx, w*2-2), max(dy-1, 0)185 for k := range 4 {186 b.dot(dx+k%2, dy+k/2, j)187 }188 })189 b.draw(g, p.plot.Min.X, 0)190}191192// image draws the points as discs and the trend lines through them.193func (p *scatterPlan) image(cv *canvas) {194 bottom := float64(p.plot.Dy()) * cv.ch195 px := func(x float64) float64 { return x * cv.cw }196 py := func(y float64) float64 { return bottom - y*cv.ch }197 if p.grid {198 for j := 1; j <= p.ysc.n; j++ {199 cv.hline(int(math.Round(py(float64(j*p.ysc.k)))), cv.pal.Grid)200 }201 for j := 1; j <= p.xsc.n; j++ {202 cv.vline(int(math.Round(px(float64(j*p.xsc.k))))-1, cv.pal.Grid)203 }204 }205 if p.trend {206 width := max(cv.ch/12, 1.5)207 for j := range p.ys {208 var run [][2]float64209 flush := func() {210 if len(run) > 1 {211 cv.polyline(run, width, cv.pal.Series[j%Colors])212 }213 run = run[:0]214 }215 p.trendPoints(j, 1, func(x, y float64, ok bool) {216 if !ok {217 flush()218 return219 }220 run = append(run, [2]float64{px(x), py(y)})221 })222 flush()223 }224 }225 r := max(min(cv.cw, cv.ch)*0.3, 2.5)226 p.points(func(j int, x, y float64) {227 cv.disc(px(x), py(y), r, cv.pal.Series[j%Colors])228 })229}230231// disc fills a circle of radius r centered at cx, cy, antialiased.232func (cv *canvas) disc(cx, cy, r float64, c color.RGBA) {233 if cv.vec != nil {234 cv.vec.disc(cx, cy, r, c)235 return236 }237 for y := int(math.Floor(cy - r - 1)); y <= int(math.Ceil(cy+r+1)); y++ {238 for x := int(math.Floor(cx - r - 1)); x <= int(math.Ceil(cx+r+1)); x++ {239 d := math.Hypot(float64(x)+0.5-cx, float64(y)+0.5-cy)240 cv.blend(x, y, c, r+0.5-d)241 }242 }243}