Merge pull request #730 from JuliaPlots/pv/violin

support datalimits and respect area in violin

Author: SimonDanisch

docs/src/plotting_functions/violin.md

@@ -17,22 +17,53 @@ ys = randn(1000)
 violin(xs, ys)
 ```
 
+
+```@example
+using CairoMakie
+CairoMakie.activate!() # hide
+AbstractPlotting.inline!(true) # hide
+
+xs = rand(1:3, 1000)
+ys = map(xs) do x
+    return x == 1 ? randn() : x == 2 ? 0.5 * randn() : 5 * rand()
+end
+
+violin(xs, ys, datalimits = extrema)
+```
+
+
 ```@example
 using CairoMakie
 CairoMakie.activate!() # hide
 AbstractPlotting.inline!(true) # hide
 
-xs1 = rand(1:3, 1000)
-ys1 = randn(1000)
-dodge1 = rand(1:2, 1000)
+N = 1000
+xs = rand(1:3, N)
+dodge = rand(1:2, N)
+side = rand([:left, :right], N)
+color = @. ifelse(side == :left, :orange, :teal)
+ys = map(side) do s
+    return s == :left ? randn() : rand()
+end
+
+violin(xs, ys, dodge = dodge, side = side, color = color)
+```
+
+```@example
+using CairoMakie
+CairoMakie.activate!() # hide
+AbstractPlotting.inline!(true) # hide
 
-xs2 = rand(1:3, 1000)
-ys2 = randn(1000)
-dodge2 = rand(1:2, 1000)
+N = 1000
+xs = rand(1:3, N)
+side = rand([:left, :right], N)
+color = map(xs, side) do x, s
+    colors = s == :left ? [:red, :orange, :yellow] : [:blue, :teal, :cyan]
+    return colors[x]
+end
+ys = map(side) do s
+    return s == :left ? randn() : rand()
+end
 
-fig = Figure()
-ax = Axis(fig[1, 1])
-violin!(ax, xs1, ys1, dodge = dodge1, side = :left, color = :orange)
-violin!(ax, xs2, ys2, dodge = dodge2, side = :right, color = :teal)
-fig
+violin(xs, ys, side = side, color = color)
 ```

src/stats/violin.jl

@@ -8,6 +8,8 @@ Draw a violin plot.
 - `orientation=:vertical`: orientation of the violins (`:vertical` or `:horizontal`)
 - `width=0.8`: width of the violin
 - `show_median=true`: show median as midline
+- `side=:both`: specify `:left` or `:right` to only plot the violin on one side
+- `datalimits`: specify values to trim the `violin`. Can be a `Tuple` or a `Function` (e.g. `datalimits=extrema`)
 """
 @recipe(Violin, x, y) do scene
     Theme(;
@@ -21,7 +23,8 @@ Draw a violin plot.
         n_dodge = automatic,
         x_gap = 0.2,
         dodge_gap = 0.03,
-        trim = false,
+        datalimits = (-Inf, Inf),
+        max_density = automatic,
         strokecolor = :white,
         show_median = false,
         mediancolor = automatic,
@@ -31,33 +34,67 @@ end
 
 conversion_trait(x::Type{<:Violin}) = SampleBased()
 
-function plot!(plot::Violin)
-    x, y, width, side, show_median = plot[1], plot[2], plot[:width], plot[:side], plot[:show_median]
-    npoints, boundary, bandwidth = plot[:npoints], plot[:boundary], plot[:bandwidth]
-    dodge, n_dodge, x_gap, dodge_gap = plot[:dodge], plot[:n_dodge], plot[:x_gap], plot[:dodge_gap]
+getuniquevalue(v, idxs) = v
+
+function getuniquevalue(v::AbstractVector, idxs)
+    u = view(v, idxs)
+    f = first(u)
+    msg = "Collection must have the same value across all indices"
+    all(isequal(f), u) || throw(ArgumentError(msg))
+    return f
+end
 
-    signals = lift(x, y, width, dodge, n_dodge, x_gap, dodge_gap, side, show_median, npoints, boundary, bandwidth) do x, y, width, dodge, n_dodge, x_gap, dodge_gap, vside, show_median, n, bound, bw
+function plot!(plot::Violin)
+    x, y = plot[1], plot[2]
+    args = @extract plot (width, side, color, show_median, npoints, boundary, bandwidth,
+        datalimits, max_density, dodge, n_dodge, x_gap, dodge_gap)
+    signals = lift(x, y, args...) do x, y, width, vside, color, show_median, n, bound, bw, limits, max_density, dodge, n_dodge, x_gap, dodge_gap
         x̂, violinwidth = xw_from_dodge(x, width, 1, x_gap, dodge, n_dodge, dodge_gap)
-        vertices = Vector{Point2f0}[]
-        lines = Pair{Point2f0, Point2f0}[]
-        for (key, idxs) in StructArrays.finduniquesorted(x̂)
+
+        # Allow `side` to be either scalar or vector
+        sides = broadcast(x̂, vside) do _, s
+            return s == :left ? - 1 : s == :right ? 1 : 0
+        end
+
+        sa = StructArray((x = x̂, side = sides))
+
+        specs = map(StructArrays.finduniquesorted(sa)) do (key, idxs)
             v = view(y, idxs)
             k = KernelDensity.kde(v;
                 npoints = n,
                 (bound === automatic ? NamedTuple() : (boundary = bound,))...,
                 (bw === automatic ? NamedTuple() : (bandwidth = bw,))...,
             )
-            spec = (x = key, kde = k, median = median(v))
-            min, max = extrema_nan(spec.kde.density)
+            l1, l2 = limits isa Function ? limits(v) : limits
+            i1, i2 = searchsortedfirst(k.x, l1), searchsortedlast(k.x, l2)
+            kde = (x = view(k.x, i1:i2), density = view(k.density, i1:i2))
+            c = getuniquevalue(color, idxs)
+            return (x = key.x, side = key.side, color = to_color(c), kde = kde, median = median(v))
+        end
+
+        max = if max_density === automatic
+            maximum(specs) do spec
+                _, max = extrema_nan(spec.kde.density)
+                return max
+            end
+        else
+            max_density
+        end
+
+        vertices = Vector{Point2f0}[]
+        lines = Pair{Point2f0, Point2f0}[]
+        colors = RGBA{Float32}[]
+
+        for spec in specs
             scale = 0.5*violinwidth/max
             xl = reverse(spec.x .- spec.kde.density .* scale)
             xr = spec.x .+ spec.kde.density .* scale
             yl = reverse(spec.kde.x)
             yr = spec.kde.x
 
-            x_coord, y_coord = if vside == :left
+            x_coord, y_coord = if spec.side == -1 # left violin
                 [spec.x; xl; spec.x], [yl[1]; yl; yl[end]]
-            elseif vside == :right
+            elseif spec.side == 1 # right violin
                 [spec.x; xr; spec.x], [yr[1]; yr; yr[end]]
             else
                 [spec.x; xr; spec.x; xl], [yr[1]; yr; yl[1]; yl]
@@ -72,22 +109,30 @@ function plot!(plot::Violin)
                 ym₋, ym₊ = spec.kde.density[ip-1], spec.kde.density[ip]
                 xm₋, xm₊ = spec.kde.x[ip-1], spec.kde.x[ip]
                 ym = (xm * (ym₊ - ym₋) + xm₊ * ym₋ - xm₋ * ym₊) / (xm₊ - xm₋)
-                median_left = Point2f0(vside == :right ? spec.x : spec.x - ym * scale, xm)
-                median_right = Point2f0(vside == :left ? spec.x : spec.x + ym * scale, xm)
+                median_left = Point2f0(spec.side == 1 ? spec.x : spec.x - ym * scale, xm)
+                median_right = Point2f0(spec.side == -1 ? spec.x : spec.x + ym * scale, xm)
                 push!(lines, median_left => median_right)
             end
+
+            push!(colors, spec.color)
         end
-        return vertices, lines
+
+        return (vertices = vertices, lines = lines, colors = colors)
     end
-    t = copy(Theme(plot))
-    mediancolor = pop!(t, :mediancolor)
-    poly!(plot, t, lift(first, signals))
+
+    poly!(
+        plot,
+        lift(s -> s.vertices, signals),
+        color = lift(s -> s.colors, signals),
+        strokecolor = plot[:strokecolor],
+        strokewidth = plot[:strokewidth],
+    )
     linesegments!(
         plot,
-        lift(last, signals),
+        lift(s -> s.lines, signals),
         color = lift(
             (mc, sc) -> mc === automatic ? sc : mc,
-            mediancolor,
+            plot[:mediancolor],
             plot[:strokecolor],
         ),
         linewidth = plot[:medianlinewidth],