added legend, custom columns

src/SplomPlots.jl

@@ -4,18 +4,23 @@ using DataFrames
 using Plots
 using RDatasets
 
-function splom(df::DataFrame)
+function splom(df::DataFrame; group="", columns=[])
-    columns = Dict(names(df) .=> eltype.(eachcol(df)))
+    if columns == []
-
+        columns = Dict(names(df) .=> eltype.(eachcol(df)))
-    for (key, value) in columns
+    else
-        if value <: Number
+        columns = Dict(names(df[!,columns]) .=> eltype.(eachcol(df[!,columns])))
-            continue
-
-        else
-            pop!(columns, key)
-        end
     end
 
+        for (key, value) in columns
+            if value <: Number
+                continue
+
+            else
+                pop!(columns, key)
+            end
+        end
+
+    println(columns)
     cols = collect(keys(columns))
     col_pairs = [(x, y) for x in cols, y in cols]
     col_len = length(cols)
@@ -25,12 +30,38 @@ function splom(df::DataFrame)
 
     for (i, (x, y)) in enumerate(col_pairs)
         scatter_plots[i] = plot()
+
         if x == y
-            plot!(;
+            if i == col_len^2 && group != ""
-                xaxis=nothing, yaxis=nothing, showaxis=false
+                fakedat = (1:length(df[!, x])) .* 0
-            )
+                scatter!(
+                    fakedat,
+                    fakedat;
+                    xaxis=nothing,
+                    yaxis=nothing,
+                    showaxis=false,
+                    grid=false,
+                    xlims=(1, 0),
+                    ylims=(1, 0),
+                    group=df[!, group],
+                    legend=:topleft,
+                )
+            else
+                plot!(;
+                    xaxis=nothing,
+                    yaxis=nothing,
+                    showaxis=false,
+                )
+            end
         else
-            scatter!(df[!, x], df[!, y])
+            scatter!(
+                df[!, x],
+                df[!, y];
+                group=df.Species,
+                label="",
+            )
+
+            # end
         end
 
         if i % col_len == 1
@@ -42,11 +73,11 @@ function splom(df::DataFrame)
         end
     end
 
-    return plot(scatter_plots...; label="")
+    return plot(scatter_plots...;)
 end
 
 df = dataset("datasets", "iris")
 
-splom(df)
+splom(df; group=:Species,columns=[:SepalLength, :SepalWidth, :PetalWidth])
 
 end # module