Skip to content

Commit

Permalink
Density heatmap (#56)
Browse files Browse the repository at this point in the history
  • Loading branch information
@cristineguadelupe
cristineguadelupe authored Aug 2, 2023
1 parent 373594c commit e5c2496
Show file tree
Hide file tree
Showing 2 changed files with 247 additions and 18 deletions.
87 changes: 73 additions & 14 deletions lib/vega_lite/data.ex
View file
Original file line number Diff line number Diff line change
Expand Up @@ -133,24 +133,54 @@ defmodule VegaLite.Data do
@spec heatmap(VegaLite.t(), Table.Reader.t(), keyword()) :: VegaLite.t()
def heatmap(vl, data, fields) do
for key <- [:x, :y], is_nil(fields[key]) do
raise ArgumentError, "the #{key} axis is required to plot a heatmap"
raise ArgumentError, "the #{key} field is required to plot a heatmap"
end

{cols, fields, used_fields} = build_options(data, fields, &heatmap_defaults/2)
text_fields = Keyword.take(fields, [:text, :x, :y])
rect_fields = Keyword.delete(fields, :text)
opts = build_options(data, fields, &heatmap_defaults/2)
build_heatmap_layers(vl, data, opts)
end

layers =
[encode_layer(cols, :rect, rect_fields)] ++
if fields[:text] do
[encode_layer(cols, :text, text_fields)]
else
[]
end
@doc """
Returns the specification of a density heat map for a given data and a list of fields to be encoded.
vl
|> Vl.data_from_values(data, only: used_fields)
|> Vl.layers(layers)
As a specialized chart, the density heatmap expects the `:x` and `:y` axes, a `:color` field and
optionally a `:text` field. All data must be `:quantitative` and the default aggregation
function is `:count`.
## Examples
data = [
%{"total_bill" => 16.99, "tip" => 1.0},
%{"total_bill" => 10.34, "tip" => 1.66}
]
Data.density_heatmap(data, x: "total_bill", y: "tip", color: "total_bill", text: "tip")
"""
@spec density_heatmap(Table.Reader.t(), keyword()) :: VegaLite.t()
def density_heatmap(data, fields), do: density_heatmap(Vl.new(), data, fields)

@doc """
Same as density_heatmap/2, but takes a valid `VegaLite` specification as the first argument.
## Examples
data = [
%{"total_bill" => 16.99, "tip" => 1.0},
%{"total_bill" => 10.34, "tip" => 1.66}
]
Vl.new(title: "Density Heatmap", width: 500)
|> Data.heatmap(data, x: "total_bill", y: "tip", color: "total_bill", text: "tip")
"""
@spec density_heatmap(VegaLite.t(), Table.Reader.t(), keyword()) :: VegaLite.t()
def density_heatmap(vl, data, fields) do
for key <- [:x, :y, :color], is_nil(fields[key]) do
raise ArgumentError, "the #{key} field is required to plot a density heatmap"
end

opts = build_options(data, fields, &density_heatmap_defaults/2)
build_heatmap_layers(vl, data, opts)
end

defp heatmap_defaults(field, opts) when field in [:x, :y] do
Expand All @@ -165,6 +195,18 @@ defmodule VegaLite.Data do
opts
end

defp density_heatmap_defaults(field, opts) when field in [:x, :y] do
opts |> Keyword.put_new(:type, :quantitative) |> Keyword.put_new(:bin, true)
end

defp density_heatmap_defaults(field, opts) when field in [:color, :text] do
opts |> Keyword.put_new(:type, :quantitative) |> Keyword.put_new(:aggregate, :count)
end

defp density_heatmap_defaults(_field, opts) do
opts
end

## Shared helpers

defp encode_mark(vl, opts) when is_list(opts) do
Expand Down Expand Up @@ -195,6 +237,23 @@ defmodule VegaLite.Data do
{columns_for(data), standardize_fields(fields, fun), used_fields}
end

defp build_heatmap_layers(vl, data, {cols, fields, used_fields}) do
text_fields = Keyword.take(fields, [:text, :x, :y])
rect_fields = Keyword.delete(fields, :text)

layers =
[encode_layer(cols, :rect, rect_fields)] ++
if fields[:text] do
[encode_layer(cols, :text, text_fields)]
else
[]
end

vl
|> Vl.data_from_values(data, only: used_fields)
|> Vl.layers(layers)
end

defp used_fields(fields) do
for {_key, field} <- fields do
if is_list(field), do: field[:field], else: field
Expand Down
178 changes: 174 additions & 4 deletions test/vega_lite/data_test.exs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -339,16 +339,186 @@ defmodule VegaLite.DataTest do
assert vl == Data.heatmap(@data, x: "height", y: "weight", color: "height", text: "width")
end

test "raises an error when the x axis is not given" do
assert_raise ArgumentError, "the x axis is required to plot a heatmap", fn ->
test "raises an error when the x field is not given" do
assert_raise ArgumentError, "the x field is required to plot a heatmap", fn ->
Data.heatmap(@data, y: "y")
end
end

test "raises an error when the y axis is not given" do
assert_raise ArgumentError, "the y axis is required to plot a heatmap", fn ->
test "raises an error when the y field is not given" do
assert_raise ArgumentError, "the y field is required to plot a heatmap", fn ->
Data.heatmap(@data, x: "x", text: "text")
end
end
end

describe "density heatmap" do
test "simple density heatmap" do
vl =
Vl.new()
|> Vl.data_from_values(@data, only: ["height", "weight"])
|> Vl.layers([
Vl.new()
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:color, "height", type: :quantitative, aggregate: :count)
])

assert vl == Data.density_heatmap(@data, x: "height", y: "weight", color: "height")
end

test "with title" do
vl =
Vl.new(title: "Density heatmap")
|> Vl.data_from_values(@data, only: ["height", "weight"])
|> Vl.layers([
Vl.new()
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:color, "height", type: :quantitative, aggregate: :count)
])

assert vl ==
Vl.new(title: "Density heatmap")
|> Data.density_heatmap(@data, x: "height", y: "weight", color: "height")
end

test "with specified bins" do
vl =
Vl.new()
|> Vl.data_from_values(@data, only: ["height", "weight"])
|> Vl.layers([
Vl.new()
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: [maxbins: 10])
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: [maxbins: 10])
|> Vl.encode_field(:color, "height", type: :quantitative, aggregate: :count)
])

assert vl ==
Data.density_heatmap(@data,
x: [field: "height", bin: [maxbins: 10]],
y: [field: "weight", bin: [maxbins: 10]],
color: "height"
)
end

test "with specified aggregate for color" do
vl =
Vl.new()
|> Vl.data_from_values(@data, only: ["height", "weight"])
|> Vl.layers([
Vl.new()
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:color, "height", type: :quantitative, aggregate: :mean)
])

assert vl ==
Data.density_heatmap(@data,
x: "height",
y: "weight",
color: [field: "height", aggregate: :mean]
)
end

test "with text" do
vl =
Vl.new()
|> Vl.data_from_values(@data, only: ["height", "weight"])
|> Vl.layers([
Vl.new()
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:color, "height", type: :quantitative, aggregate: :count),
Vl.new()
|> Vl.mark(:text)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:text, "height", type: :quantitative, aggregate: :count)
])

assert vl ==
Data.density_heatmap(@data,
x: "height",
y: "weight",
color: "height",
text: "height"
)
end

test "with specified aggregate for text" do
vl =
Vl.new()
|> Vl.data_from_values(@data, only: ["height", "weight"])
|> Vl.layers([
Vl.new()
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:color, "height", type: :quantitative, aggregate: :count),
Vl.new()
|> Vl.mark(:text)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:text, "height", type: :quantitative, aggregate: :mean)
])

assert vl ==
Data.density_heatmap(@data,
x: "height",
y: "weight",
color: "height",
text: [field: "height", aggregate: :mean]
)
end

test "with text different from the axes" do
vl =
Vl.new()
|> Vl.data_from_values(@data, only: ["height", "weight", "width"])
|> Vl.layers([
Vl.new()
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:color, "height", type: :quantitative, aggregate: :count),
Vl.new()
|> Vl.mark(:text)
|> Vl.encode_field(:x, "height", type: :quantitative, bin: true)
|> Vl.encode_field(:y, "weight", type: :quantitative, bin: true)
|> Vl.encode_field(:text, "width", type: :quantitative, aggregate: :count)
])

assert vl ==
Data.density_heatmap(@data,
x: "height",
y: "weight",
color: "height",
text: "width"
)
end

test "raises an error when the x field is not given" do
assert_raise ArgumentError, "the x field is required to plot a density heatmap", fn ->
Data.density_heatmap(@data, y: "y")
end
end

test "raises an error when the y field is not given" do
assert_raise ArgumentError, "the y field is required to plot a density heatmap", fn ->
Data.density_heatmap(@data, x: "x", text: "text")
end
end

test "raises an error when the color field is not given" do
assert_raise ArgumentError, "the color field is required to plot a density heatmap", fn ->
Data.density_heatmap(@data, x: "x", y: "y")
end
end
end
end

0 comments on commit e5c2496

Please sign in to comment.