Bubble chart

lib/matplotex/figure/areal.ex

@@ -141,6 +141,7 @@ defmodule Matplotex.Figure.Areal do
       def materialized_by_region(figure) do
         figure
         |> Lead.set_regions_areal()
+        |> Lead.transform_sizes()
         |> Cast.cast_xticks_by_region()
         |> Cast.cast_yticks_by_region()
         |> Cast.cast_hgrids_by_region()
@@ -184,6 +185,7 @@ defmodule Matplotex.Figure.Areal do
 
         {x, y}
       end
+
       def flatten_for_data(_, _bottom) do
         raise InputError, bottom: "Wrong data provided for opts bottom"
       end
@@ -433,10 +435,36 @@ defmodule Matplotex.Figure.Areal do
         |> transformation(y, xlim, ylim, width, height, transition)
         |> Algebra.flip_y_coordinate()
       end)
+      |> maybe_wrap_with_sizes(dataset)
+      |> maybe_wrap_with_colors()
 
     %Dataset{dataset | transformed: transformed}
   end
 
+  defp maybe_wrap_with_sizes(transformed, %Dataset{sizes: sizes} = dataset)
+       when length(transformed) == length(sizes) do
+    {Enum.zip(transformed, sizes), dataset}
+  end
+
+  defp maybe_wrap_with_sizes(
+         transformed,
+         %Dataset{colors: colors, marker_size: marker_size} = dataset
+       )
+       when length(transformed) == length(colors) do
+    {Enum.zip(transformed, List.duplicate(marker_size, length(transformed))), dataset}
+  end
+
+  defp maybe_wrap_with_sizes(transformed, dataset) do
+    {transformed, dataset}
+  end
+
+  defp maybe_wrap_with_colors({transformed, %Dataset{colors: colors}})
+       when length(transformed) == length(colors) do
+    Enum.zip(transformed, colors)
+  end
+
+  defp maybe_wrap_with_colors({transformed, _dataset}), do: transformed
+
   defp transform_with_bottom(x, y, xlim, ylim, width, height, transition) when is_list(y) do
     y_top = Enum.sum(y)
     y_bottom = y |> tl() |> Enum.sum()

lib/matplotex/figure/areal/scatter.ex

@@ -1,5 +1,6 @@
 defmodule Matplotex.Figure.Areal.Scatter do
   @moduledoc false
+  alias Matplotex.Colorscheme.Garner
   alias Matplotex.Figure.Areal.PlotOptions
   alias Matplotex.Figure.Areal.Region
   alias Matplotex.Figure.Areal.Ticker
@@ -31,7 +32,7 @@ defmodule Matplotex.Figure.Areal.Scatter do
   def create(%Figure{axes: %__MODULE__{dataset: data} = axes} = figure, {x, y}, opts \\ []) do
     x = determine_numeric_value(x)
     y = determine_numeric_value(y)
-    dataset = Dataset.cast(%Dataset{x: x, y: y}, opts)
+    dataset = Dataset.cast(%Dataset{x: x, y: y}, opts) |> Dataset.update_cmap()
     datasets = data ++ [dataset]
     xydata = flatten_for_data(datasets)
 
@@ -97,6 +98,45 @@ defmodule Matplotex.Figure.Areal.Scatter do
     end
   end
 
+  def capture(
+        [{{{x, y}, s}, color} | to_capture],
+        captured,
+        %Dataset{
+          marker: marker,
+          colors: colors,
+          cmap: cmap
+        } = dataset
+      ) do
+    color = colors |> Enum.min_max() |> Garner.garn_color(color, cmap)
+
+    capture(
+      to_capture,
+      captured ++
+        [
+          Marker.generate_marker(marker, x, y, color, s)
+        ],
+      dataset
+    )
+  end
+
+  def capture(
+        [{{x, y}, s} | to_capture],
+        captured,
+        %Dataset{
+          color: color,
+          marker: marker
+        } = dataset
+      ) do
+    capture(
+      to_capture,
+      captured ++
+        [
+          Marker.generate_marker(marker, x, y, color, s)
+        ],
+      dataset
+    )
+  end
+
   def capture(
         [{x, y} | to_capture],
         captured,

lib/matplotex/figure/dataset.ex

@@ -1,5 +1,7 @@
 defmodule Matplotex.Figure.Dataset do
   @moduledoc false
+  alias Matplotex.Colorscheme.Colormap
+
   @default_color "blue"
   @default_marker "o"
   @default_linestyle "_"
@@ -31,4 +33,8 @@ defmodule Matplotex.Figure.Dataset do
   def cast(dataset, values) do
     struct(dataset, values)
   end
+
+  def update_cmap(%__MODULE__{cmap: cmap} = dataset) do
+    %__MODULE__{dataset | cmap: Colormap.fetch_cmap(cmap)}
+  end
 end

lib/matplotex/figure/lead.ex

@@ -1,5 +1,7 @@
 defmodule Matplotex.Figure.Lead do
   @moduledoc false
+  alias Matplotex.InputError
+  alias Matplotex.Figure.Dataset
   alias Matplotex.Utils.Algebra
   alias Matplotex.Figure.Areal.XyRegion.Coords, as: XyCoords
   alias Matplotex.Figure.Font
@@ -45,6 +47,48 @@ defmodule Matplotex.Figure.Lead do
     %Figure{figure | axes: %{axes | border: {lx, by, rx, ty}}}
   end
 
+  def transform_sizes(%Figure{axes: %{dataset: datasets} = axes} = figure) do
+    datasets = transform_dataset_sizes(figure, datasets, [])
+    %Figure{figure | axes: %{axes | dataset: datasets}}
+  end
+
+  defp transform_dataset_sizes(_figure, [%Dataset{sizes: nil} | _to_transorm] = datasets, _),
+    do: datasets
+
+  defp transform_dataset_sizes(
+         %Figure{
+           axes: %{
+             region_content: %Region{width: width_region_content, height: height_region_content}
+           }
+         } = figure,
+         [%Dataset{sizes: sizes} = dataset | to_transorm],
+         transformed
+       )
+       when length(sizes) > 0 do
+    content_area = width_region_content * height_region_content
+    total_size = Enum.sum(sizes)
+
+    area_size_ratio =
+      if total_size > 0 do
+        content_area / total_size * 2
+      else
+        raise InputError, message: "Invalid sizes for fractionizing area"
+      end
+
+    sizes =
+      sizes
+      |> Nx.tensor()
+      |> Nx.multiply(area_size_ratio)
+      |> Nx.to_list()
+
+    transform_dataset_sizes(figure, to_transorm, [%Dataset{dataset | sizes: sizes} | transformed])
+  end
+
+  defp transform_dataset_sizes(_figure, [%Dataset{sizes: []} | _to_transorm] = datasets, _),
+    do: datasets
+
+  defp transform_dataset_sizes(_, [], transformed), do: transformed
+
   defp set_frame_size(%Figure{margin: margin, figsize: {fwidth, fheight}, axes: axes} = figure) do
     frame_size = {fwidth - fwidth * 2 * margin, fheight - fheight * 2 * margin}
     lx = fwidth * margin

test/matplotex/figure/areal/scatter_test.exs

@@ -1,4 +1,6 @@
 defmodule Matplotex.Figure.Areal.ScatterTest do
+  alias Matplotex.Figure.Areal
+  alias Matplotex.Figure.Dataset
   alias Matplotex.Figure.Areal.Scatter
   alias Matplotex.Figure
   use Matplotex.PlotCase
@@ -12,6 +14,51 @@ defmodule Matplotex.Figure.Areal.ScatterTest do
       assert %Figure{axes: %{data: {x, _y}, element: elements}} = Scatter.materialize(figure)
       assert Enum.count(elements, fn elem -> elem.type == "plot.marker" end) == length(x)
     end
+
+    test "generates elements with various saizes if it passed a size attrbute" do
+      x = [1, 2, 3, 4, 5]
+      y = [10, 20, 30, 40, 50]
+      sizes = [1, 2, 3, 4, 5]
+
+      assert %Figure{axes: %{element: elements}} =
+               x |> Matplotex.scatter(y, sizes: sizes) |> Figure.materialize()
+
+      [h | tail] =
+        elements
+        |> Enum.filter(fn x -> x.type == "plot.marker" end)
+        |> Enum.map(fn x ->
+          x.r
+        end)
+
+      refute Enum.all?(tail, fn x -> x == h end)
+    end
+
+    test "generates elements with various saizes and colors if it passed a size and color attrbute" do
+      x = [1, 2, 3, 4, 5]
+      y = [10, 20, 30, 40, 50]
+      sizes = [1, 2, 3, 4, 5]
+
+      assert %Figure{axes: %{element: elements}} =
+               x |> Matplotex.scatter(y, sizes: sizes, colors: sizes) |> Figure.materialize()
+
+      [h | tail] =
+        elements
+        |> Enum.filter(fn x -> x.type == "plot.marker" end)
+        |> Enum.map(fn x ->
+          x.r
+        end)
+
+      refute Enum.all?(tail, fn x -> x == h end)
+
+      [h | tail] =
+        elements
+        |> Enum.filter(fn x -> x.type == "plot.marker" end)
+        |> Enum.map(fn x ->
+          x.fill
+        end)
+
+      refute Enum.all?(tail, fn x -> x == h end)
+    end
   end
 
   describe "generate_ticks/2" do
@@ -24,4 +71,99 @@ defmodule Matplotex.Figure.Areal.ScatterTest do
       assert length(ticks) == 6
     end
   end
+
+  describe "do_transform" do
+    test "zips transformed values with sizes if the dataset contains sizes in eaqual size" do
+      x = [1, 2, 3, 4, 5]
+      y = [10, 20, 30, 40, 50]
+      sizes = [1, 2, 3, 4, 5]
+      width = 2
+      height = 2
+
+      assert %Figure{axes: %{dataset: [dataset]}} =
+               x |> Matplotex.scatter(y, figsize: {width, height}, sizes: sizes)
+
+      assert %Dataset{transformed: transformed} =
+               Areal.do_transform(
+                 dataset,
+                 Enum.min_max(x),
+                 Enum.min_max(y),
+                 width,
+                 height,
+                 {0, 0}
+               )
+
+      assert Enum.all?(transformed, &match?({{_, _}, _}, &1))
+    end
+
+    test "zips transformed values with marker size if colors exist without sizes" do
+      x = [1, 2, 3, 4, 5]
+      y = [10, 20, 30, 40, 50]
+      colors = [1, 2, 3, 4, 5]
+      width = 2
+      height = 2
+
+      assert %Figure{axes: %{dataset: [%Dataset{marker_size: _marker_size} = dataset]}} =
+               x |> Matplotex.scatter(y, figsize: {width, height}, colors: colors)
+
+      assert %Dataset{transformed: transformed} =
+               Areal.do_transform(
+                 dataset,
+                 Enum.min_max(x),
+                 Enum.min_max(y),
+                 width,
+                 height,
+                 {0, 0}
+               )
+
+      assert Enum.all?(transformed, &match?({{{_, _}, _marker_size}, _}, &1))
+    end
+
+    test "zips transformed values with colors if the dataset contanis colors" do
+      x = [1, 2, 3, 4, 5]
+      y = [10, 20, 30, 40, 50]
+      colors = [1, 2, 3, 4, 5]
+      width = 2
+      height = 2
+
+      assert %Figure{axes: %{dataset: [dataset]}} =
+               x |> Matplotex.scatter(y, figsize: {width, height}, colors: colors)
+
+      assert %Dataset{transformed: transformed} =
+               Areal.do_transform(
+                 dataset,
+                 Enum.min_max(x),
+                 Enum.min_max(y),
+                 width,
+                 height,
+                 {0, 0}
+               )
+
+      assert Enum.all?(transformed, &match?({{{_, _}, _}, _}, &1))
+    end
+
+    test "zips both size and colors if the dataset contains size and color" do
+      x = [1, 2, 3, 4, 5]
+      y = [10, 20, 30, 40, 50]
+      sizes = [1, 2, 3, 4, 5]
+      colors = [1, 2, 3, 4, 5]
+      width = 2
+      height = 2
+
+      assert %Figure{axes: %{dataset: [dataset]}} =
+               x |> Matplotex.scatter(y, figsize: {width, height}, sizes: sizes, colors: colors)
+
+      assert %Dataset{transformed: transformed} =
+               Areal.do_transform(
+                 dataset,
+                 Enum.min_max(x),
+                 Enum.min_max(y),
+                 width,
+                 height,
+                 {0, 0}
+               )
+
+      assert Enum.all?(transformed, &match?({{{_, _}, _}, _}, &1))
+    end
+  end
 end

test/matplotex/figure/lead_test.exs

@@ -1,4 +1,5 @@
 defmodule Matplotex.Figure.LeadTest do
+  alias Matplotex.Figure.Dataset
   alias Matplotex.Figure.TwoD
   alias Matplotex.Figure.Areal.Region
   alias Matplotex.Figure
@@ -259,4 +260,25 @@ defmodule Matplotex.Figure.LeadTest do
       assert cx != 0 && cy != 0
     end
   end
+
+  describe "transform_sizes/1" do
+    test "converts sizes to equivalent radius to the buble" do
+      x = [1, 2, 3, 4, 5]
+      y = [10, 20, 30, 40, 50]
+      sizes = [1, 2, 3, 4, 5]
+      width = 2
+      height = 2
+
+      figure =
+        x
+        |> Matplotex.scatter(y, figsize: {width, height}, sizes: sizes)
+        |> Lead.set_regions_areal()
+
+      assert %Figure{axes: %{dataset: [%Dataset{sizes: transformed_sizes}]}} =
+               Lead.transform_sizes(figure)
+
+      assert length(transformed_sizes) == length(sizes)
+      assert Enum.sum(transformed_sizes) < width * height * 1.5
+    end
+  end
 end