From 856fcf5e1f8f77feac221ebdd066b6cf3e2d7dc7 Mon Sep 17 00:00:00 2001
From: James R Swift <james.rt.swift@gmail.com>
Date: Tue, 30 Jul 2024 12:35:58 +0100
Subject: [PATCH] Violin plots (draft)

---
 src/plot/violin.typ        | 103 +++++++++++++++----------------------
 tests/plot/violin/test.typ |  19 +++++--
 2 files changed, 57 insertions(+), 65 deletions(-)

diff --git a/src/plot/violin.typ b/src/plot/violin.typ
index a36f059..3d42990 100644
--- a/src/plot/violin.typ
+++ b/src/plot/violin.typ
@@ -1,74 +1,55 @@
-#import "/src/cetz.typ": draw, process, util, matrix
+#import "line.typ"
 
-#import "util.typ"
-#import "sample.typ"
-
-#let _prepare(self, ctx) = {
-  let (x, y) = (ctx.x, ctx.y)
-
-  // bin the data
-  let (min, max) = (calc.min(..self.data),calc.max(..self.data))
-  let range = max - min
-  
-  let binned = self.data.sorted().fold( (0,)*self.bins, (acc, it) => {
-    let bin =  int(self.bins * (it - min) / (max - min))
-    acc.at(bin - 1) += 1
-    return acc
-  })
-
-  self.line-data = binned.enumerate().map( ((x, y)) => {
-    (0 + y / 10, min + x)
-  })
-
-  // Generate stroke paths
-  self.stroke-paths = util.compute-stroke-paths(self.line-data,
-    (x.min, y.min), (x.max, y.max))
-
-  // Compute fill paths if filling is requested
-  self.fill = self.at("fill", default: false)
-  if self.fill {
-    self.fill-paths = util.compute-fill-paths(
-      self.line-data,
-      (x.min, y.min), 
-      (x.max, y.max)
-    )
-  }
-
-  return self
-}
-
-#let _stroke(self, ctx) = {
-  let (x, y) = (ctx.x, ctx.y)
-
-  for p in self.stroke-paths {
-    draw.line(..p, fill: none)
-  }
-}
-
-#let _fill(self, ctx) = {
-  // fill-segments-to(self.fill-paths, y.min)
+#let kernal-normal(x, stdev: 1.5) = {
+  (1/calc.sqrt(2*calc.pi*calc.pow(stdev,2))) * calc.exp( - (x*x)/(2*calc.pow(stdev,2)))
 }
 
 #let violin( 
   data,
+  x-key: 0,
+  y-key: 1,
+  side: "left", // "left", "right", "both"
   style: (:),
+  kernel: kernal-normal.with(stdev: 1.5),
+  bandwidth: 1,
+  extend: 0.5,
   axes: ("x", "y"),
-  bins: 7,
 ) = {
 
-  ((
-    type: "violin",
-    axes: axes,
-    data: data,
-    bins: bins,
-    style: style,
-    plot-prepare: _prepare,
-    plot-stroke: _stroke,
-    plot-fill: _fill,
-    plot-legend-preview: self => {
-      draw.rect((0,0), (1,1), ..self.style)
+  for category in data {
+    let (x, ys) = (category.at(x-key), category.at(y-key))
+    let n = ys.len()
+    let (min, max) = (calc.min(..ys), calc.max(..ys))
+    let domain = (min - (max - min)*extend, max + (max - min)*extend)
+
+    let convolve = (t, side: side)=>{
+      let val = ys.map((y)=>kernel((y - t)/bandwidth)).sum()
+      (x + (1/(n*bandwidth) * val) * if side == "left" {-1} else {1}, t)
+    }
+
+    if (side in ("left", "both")) {
+      line.add(
+        convolve.with(side: "left"),
+        domain: domain,
+        line: "raw",
+        fill-type: "shape",
+        fill: true,
+        style: style
+      )
+    }
+
+    if (side in ("right", "both")){
+      line.add(
+        convolve.with(side: "right"),
+        domain: domain,
+        line: "raw",
+        fill-type: "shape",
+        fill: true,
+        style: style
+      )
     }
-  ),)
 
 
+  }
+
 }
\ No newline at end of file
diff --git a/tests/plot/violin/test.typ b/tests/plot/violin/test.typ
index 80e70f1..4abb751 100644
--- a/tests/plot/violin/test.typ
+++ b/tests/plot/violin/test.typ
@@ -10,10 +10,21 @@
   plot.plot(size: (9, 6),
     // x-tick-step: none,
     // y-tick-step: none,
-    y-max: 10,
+    y-min: -10, y-max: 20,
+    x-min: -1, x-max: 2,
   {
-    let vals = (5,4,6,8,5,4,1,5,5,5,4,2,5,4,6,5,4,5,8,4,5,)
-    cetz-plot.plot.add(vals.map(it=>(0,it)), mark: "x", style: (stroke: none))
-    cetz-plot.plot.violin(vals)
+    let vals = (
+      (0,(5,4,6,8,5.1,4.1,1,5.2,5.3,5.4,4.2,2,5.5,4.3,6,5,4,5,8,4,5,)),
+      (1,(5,4,6,8,5.1,4.1,1,5.2,5.3,5.4,4.2,2,5.5,4.3,6,5,4,5,8,4,5,)),
+    )
+    // for (x, ys) in vals {
+    //   cetz-plot.plot.add(ys.map(y=>(x,y)), mark: "x", style: (stroke: none))
+    // }
+    cetz-plot.plot.violin(
+      vals,
+      extend: 0.35,
+      side: "right",
+      bandwidth: 0.5
+    )
   })
 })