charty 0.2.3 → 0.2.4

data/lib/charty/backends/pyplot.rb

@@ -8,11 +8,14 @@ module Charty
       class << self
         def prepare
           require 'matplotlib/pyplot'
+          require 'numpy'
         end
       end
 
       def initialize
         @pyplot = ::Matplotlib::Pyplot
+        @default_line_width = ::Matplotlib.rcParams["lines.linewidth"]
+        @default_marker_size = ::Matplotlib.rcParams["lines.markersize"]
       end
 
       def self.activate_iruby_integration
@@ -90,7 +93,16 @@ module Charty
           min_l = palette.colors.map {|c| c.to_rgb.to_hsl.l }.min
           lum = min_l*0.6
           gray = Colors::RGB.new(lum, lum, lum).to_hex_string
-          draw_box_plot(context, subplot, colors, gray)
+          Array(context.data).each_with_index do |group_data, i|
+            next if group_data.empty?
+
+            box_data = group_data.compact
+            next if box_data.empty?
+
+            color = colors.next
+            draw_box_plot(box_data, vert: "v", position: i, color: color,
+                          gray: gray, width: 0.8, whisker: 1.5, flier_size: 5)
+          end
         when :bubble
           context.series.each do |data|
             ax.scatter(data.xs.to_a, data.ys.to_a, s: data.zs.to_a, alpha: 0.5,
@@ -125,93 +137,409 @@ module Charty
         end
       end
 
-      private def draw_box_plot(context, subplot, colors, gray)
-        Array(context.data).each_with_index do |group_data, i|
-          next if group_data.empty?
-
-          box_data = group_data.compact
-          next if box_data.empty?
-
-          artist_dict = @pyplot.boxplot(box_data, vert: "v", patch_artist: true,
-                                        positions: [i], widths: 0.8)
-
-          color = colors.next
-          artist_dict["boxes"].each do |box|
-            box.update({facecolor: color, zorder: 0.9, edgecolor: gray}, {})
-          end
-          artist_dict["whiskers"].each do |whisker|
-            whisker.update({color: gray, linestyle: "-"}, {})
-          end
-          artist_dict["caps"].each do |cap|
-            cap.update({color: gray}, {})
-          end
-          artist_dict["medians"].each do |median|
-            median.update({color: gray}, {})
-          end
-          artist_dict["fliers"].each do |flier|
-            flier.update({
-              markerfacecolor: gray,
-              marker: "d",
-              markeredgecolor: gray,
-              markersize: 5
-            }, {})
-          end
-        end
-      end
-
       # ==== NEW PLOTTING API ====
 
       def begin_figure
-        # do nothing
+        @legend_keys = []
+        @legend_labels = []
       end
 
-      def bar(bar_pos, values, color: nil, width: 0.8r, align: :center, orient: :v)
+      def bar(bar_pos, _group_names, values, colors, orient, label: nil, width: 0.8r,
+              align: :center, conf_int: nil, error_colors: nil, error_width: nil, cap_size: nil)
         bar_pos = Array(bar_pos)
         values = Array(values)
-        color = Array(color).map(&:to_hex_string)
+        colors = Array(colors).map(&:to_hex_string)
         width = Float(width)
+
+        ax = @pyplot.gca
+        kw = {color: colors, align: align}
+        kw[:label] = label unless label.nil?
+
         if orient == :v
-          @pyplot.bar(bar_pos, values, width: width, color: color, align: align)
+          ax.bar(bar_pos, values, width, **kw)
         else
-          @pyplot.barh(bar_pos, values, width: width, color: color, align: align)
+          ax.barh(bar_pos, values, width, **kw)
+        end
+
+        if conf_int
+          error_colors = Array(error_colors).map(&:to_hex_string)
+          confidence_intervals(ax, bar_pos, conf_int, orient, error_colors, error_width, cap_size)
+        end
+      end
+
+      private def confidence_intervals(ax, at_group, conf_int, orient, colors, error_width=nil, cap_size=nil, **options)
+        options[:lw] = error_width || @default_line_width * 1.8
+
+        at_group.each_index do |i|
+          at = at_group[i]
+          ci_low, ci_high = conf_int[i]
+          color = colors[i]
+
+          if orient == :v
+            ax.plot([at, at], [ci_low, ci_high], color: color, **options)
+            unless cap_size.nil?
+              ax.plot([at - cap_size / 2.0, at + cap_size / 2.0], [ci_low,  ci_low],  color: color, **options)
+              ax.plot([at - cap_size / 2.0, at + cap_size / 2.0], [ci_high, ci_high], color: color, **options)
+            end
+          else
+            ax.plot([ci_low, ci_high], [at, at], color: color, **options)
+            unless cap_size.nil?
+              ax.plot([ci_low,  ci_low],  [at - cap_size / 2.0, at + cap_size / 2.0], color: color, **options)
+              ax.plot([ci_high, ci_high], [at - cap_size / 2.0, at + cap_size / 2.0], color: color, **options)
+            end
+          end
         end
       end
 
-      def box_plot(plot_data, positions, color:, gray:,
-                   width: 0.8r, flier_size: 5, whisker: 1.5, notch: false)
-        color = Array(color).map(&:to_hex_string)
+      def box_plot(plot_data, group_names,
+                   orient:, colors:, gray:, dodge:, width: 0.8r,
+                   flier_size: 5, whisker: 1.5, notch: false)
+        colors = Array(colors).map(&:to_hex_string)
         gray = gray.to_hex_string
         width = Float(width)
         flier_size = Float(flier_size)
         whisker = Float(whisker)
+
         plot_data.each_with_index do |group_data, i|
-          next if group_data.nil? || group_data.empty?
+          unless group_data.nil?
+            draw_box_plot(group_data,
+                          vert: (orient == :v),
+                          position: i,
+                          color: colors[i],
+                          gray: gray,
+                          width: width,
+                          whisker: whisker,
+                          flier_size: flier_size)
+          end
+        end
+      end
 
-          artist_dict = @pyplot.boxplot(group_data, vert: :v,
-                                        patch_artist: true,
-                                        positions: [i],
-                                        widths: width,
-                                        whis: whisker, )
+      def grouped_box_plot(plot_data, group_names, color_names,
+                           orient:, colors:, gray:, dodge:, width: 0.8r,
+                           flier_size: 5, whisker: 1.5, notch: false)
+        colors = Array(colors).map(&:to_hex_string)
+        gray = gray.to_hex_string
+        width = Float(width)
+        flier_size = Float(flier_size)
+        whisker = Float(whisker)
 
-          artist_dict["boxes"].each do |box|
-            box.update({facecolor: color[i], zorder: 0.9, edgecolor: gray}, {})
+        offsets = color_offsets(color_names, dodge, width)
+        orig_width = width
+        width = Float(nested_width(color_names, dodge, width))
+
+        color_names.each_with_index do |color_name, i|
+          add_box_plot_legend(gray, colors[i], color_names[i])
+
+          plot_data[i].each_with_index do |group_data, j|
+            next if group_data.empty?
+
+            position = j + offsets[i]
+            draw_box_plot(group_data,
+                          vert: (orient == :v),
+                          position: position,
+                          color: colors[i],
+                          gray: gray,
+                          width: width,
+                          whisker: whisker,
+                          flier_size: flier_size)
           end
-          artist_dict["whiskers"].each do |whisker|
-            whisker.update({color: gray, linestyle: "-"}, {})
+        end
+      end
+
+      private def add_box_plot_legend(gray, color, name)
+        patch = @pyplot.Rectangle.new([0, 0], 0, 0, edgecolor: gray, facecolor: color, label: name)
+        @pyplot.gca.add_patch(patch)
+      end
+
+      private def draw_box_plot(group_data, vert:, position:, color:, gray:, width:, whisker:, flier_size:)
+        # TODO: Do not convert to Array when group_data is Pandas::Series or Numpy::NDArray,
+        # and use MemoryView if available when group_data is Numo::NArray
+        artist_dict = @pyplot.boxplot(Array(group_data),
+                                      vert: vert,
+                                      patch_artist: true,
+                                      positions: [position],
+                                      widths: width,
+                                      whis: whisker)
+
+        artist_dict["boxes"].each do |box|
+          box.update({facecolor: color, zorder: 0.9, edgecolor: gray}, {})
+        end
+        artist_dict["whiskers"].each do |whisker|
+          whisker.update({color: gray, linestyle: "-"}, {})
+        end
+        artist_dict["caps"].each do |cap|
+          cap.update({color: gray}, {})
+        end
+        artist_dict["medians"].each do |median|
+          median.update({color: gray}, {})
+        end
+        artist_dict["fliers"].each do |flier|
+          flier.update({
+            markerfacecolor: gray,
+            marker: "d",
+            markeredgecolor: gray,
+            markersize: flier_size
+          }, {})
+        end
+      end
+
+      private def color_offsets(color_names, dodge, width)
+        n_names = color_names.length
+        if dodge
+          each_width = width / n_names
+          offsets = Charty::Linspace.new(0 .. (width - each_width), n_names).to_a
+          mean = Statistics.mean(offsets)
+          offsets.map {|x| x - mean }
+        else
+          Array.new(n_names, 0)
+        end
+      end
+
+      private def nested_width(color_names, dodge, width)
+        if dodge
+          width.to_r / color_names.length * 0.98r
+        else
+          width
+        end
+      end
+
+      def scatter(x, y, variables, legend:, color:, color_mapper:,
+                  style:, style_mapper:, size:, size_mapper:)
+        kwd = {}
+        kwd[:edgecolor] = "w"
+
+        ax = @pyplot.gca
+        points = ax.scatter(x.to_a, y.to_a, **kwd)
+
+        unless color.nil?
+          color = color_mapper[color].map(&:to_hex_string)
+          points.set_facecolors(color)
+        end
+
+        unless size.nil?
+          size = size_mapper[size].map(&method(:scale_scatter_point_size))
+          points.set_sizes(size)
+        end
+
+        unless style.nil?
+          paths = style_mapper[style, :marker].map(&method(:marker_to_path))
+          points.set_paths(paths)
+        end
+
+        sizes = points.get_sizes
+        points.set_linewidths(0.08 * Numpy.sqrt(Numpy.percentile(sizes, 10)))
+
+        if legend
+          add_relational_plot_legend(
+            ax, legend, variables, color_mapper, size_mapper, style_mapper,
+            [:color, :s, :marker]
+          ) do |label, kwargs|
+            ax.scatter([], [], label: label, **kwargs)
           end
-          artist_dict["caps"].each do |cap|
-            cap.update({color: gray}, {})
+        end
+      end
+
+      PYPLOT_MARKERS = {
+               circle: "o",
+                    x: "X",
+                cross: "P",
+          triangle_up: "^",
+        triangle_down: "v",
+               square: [4, 0, 45].freeze,
+              diamond: [4, 0, 0].freeze,
+                 star: [5, 1, 0].freeze,
+         star_diamond: [4, 1, 0].freeze,
+          star_square: [4, 1, 45].freeze,
+             pentagon: [5, 0, 0].freeze,
+              hexagon: [6, 0, 0].freeze,
+      }.freeze
+
+      private def marker_to_path(marker)
+        @path_cache ||= {}
+        if @path_cache.key?(marker)
+          @path_cache[marker]
+        elsif PYPLOT_MARKERS.key?(marker)
+          val = PYPLOT_MARKERS[marker]
+          ms = Matplotlib.markers.MarkerStyle.new(val)
+          @path_cache[marker] = ms.get_path().transformed(ms.get_transform())
+        else
+          raise ArgumentError, "Unknown marker name: %p" % marker
+        end
+      end
+
+      RELATIONAL_PLOT_LEGEND_BRIEF_TICKS = 6
+
+      private def add_relational_plot_legend(ax, verbosity, variables, color_mapper, size_mapper, style_mapper,
+                                             legend_attributes, &func)
+        brief_ticks = RELATIONAL_PLOT_LEGEND_BRIEF_TICKS
+        verbosity = :auto if verbosity == true
+
+        legend_titles = [:color, :size, :style].filter_map {|v| variables[v] }
+        legend_title = legend_titles.pop if legend_titles.length == 1
+
+        legend_kwargs = {}
+        update_legend = ->(var_name, val_name, **kw) do
+          key = [var_name, val_name]
+          if legend_kwargs.key?(key)
+            legend_kwargs[key].update(kw)
+          else
+            legend_kwargs[key] = kw
           end
-          artist_dict["medians"].each do |median|
-            median.update({color: gray}, {})
+        end
+
+        title_kwargs = {visible: false, color: "w", s: 0, linewidth: 0, marker: "", dashes: ""}
+
+        # color legend
+
+        brief_color = case verbosity
+                      when :brief
+                        color_mapper.map_type == :numeric
+                      when :auto
+                        if color_mapper.levels.nil?
+                          false
+                        else
+                          color_mapper.levels.length > brief_ticks
+                        end
+                      else
+                        false
+                      end
+        case
+        when brief_color
+          # TODO: Also support LogLocator
+          # locator = Matplotlib.ticker.LogLocator.new(numticks: brief_ticks)
+          locator = Matplotlib.ticker.MaxNLocator.new(nbins: brief_ticks)
+          limits = color_map.levels.minmax
+          color_levels, color_formatted_levels = locator_to_legend_entries(locator, limits)
+        when color_mapper.levels.nil?
+          color_levels = color_formatted_levels = []
+        else
+          color_levels = color_formatted_levels = color_mapper.levels
+        end
+
+        if legend_title.nil? && variables.key?(:color)
+          update_legend.([variables[:color], :title], variables[:color], **title_kwargs)
+        end
+
+        color_levels.length.times do |i|
+          next if color_levels[i].nil?
+          color_value = color_mapper[color_levels[i]].to_hex_string
+          update_legend.(variables[:color], color_formatted_levels[i], color: color_value)
+        end
+
+        brief_size = case verbosity
+                     when :brief
+                       size_mapper.map_type == :numeric
+                     when :auto
+                       if size_mapper.levels.nil?
+                         false
+                       else
+                         size_mapper.levels.length > brief_ticks
+                       end
+                     else
+                       false
+                     end
+        case
+        when brief_size
+          # TODO: Also support LogLocator
+          # locator = Matplotlib.ticker.LogLocator(numticks: brief_ticks)
+          locator = Matplotlib.ticker.MaxNLocator.new(nbins: brief_ticks)
+          limits = size_mapper.levels.minmax
+          size_levels, size_formatted_levels = locator_to_legend_entries(locator, limits)
+        when size_mapper.levels.nil?
+          size_levels = size_formatted_levels = []
+        else
+          size_levels = size_formatted_levels = size_mapper.levels
+        end
+
+        if legend_title.nil? && variables.key?(:size)
+          update_legend.([variables[:size], :title], variables[:size], **title_kwargs)
+        end
+
+        size_levels.length.times do |i|
+          next if size_levels[i].nil?
+          size_value = scale_scatter_point_size(size_mapper[size_levels[i]])
+          update_legend.(variables[:size], size_formatted_levels[i], linewidth: size_value, s: size_value)
+        end
+
+        if legend_title.nil? && variables.key?(:style)
+          update_legend.([variables[:style], :title], variables[:style], **title_kwargs)
+        end
+
+        unless style_mapper.levels.nil?
+          style_mapper.levels.each do |level|
+            next if level.nil?
+            attrs = style_mapper[level]
+            marker = if attrs.key?(:marker)
+                       PYPLOT_MARKERS[attrs[:marker]]
+                     else
+                       ""
+                     end
+            # TODO: support dashes
+            update_legend.(variables[:style], level,
+                           marker: marker,
+                           dashes: attrs.fetch(:dashes, ""))
           end
-          artist_dict["fliers"].each do |flier|
-            flier.update({
-              markerfacecolor: gray,
-              marker: "d",
-              markeredgecolor: gray,
-              markersize: flier_size
-            }, {})
+        end
+
+        legend_kwargs.each do |key, kw|
+          _, label = key
+          kw[:color] ||= ".2"
+          use_kw = legend_attributes.filter_map {|attr|
+            [attr, kw[attr]] if kw.key?(attr)
+          }.to_h
+          use_kw[:visible] = kw[:visible] if kw.key?(:visible)
+          func.(label, use_kw)
+        end
+
+        handles = ax.get_legend_handles_labels()[0].to_a
+        unless handles.empty?
+          legend = ax.legend(title: legend_title || "")
+          adjust_legend_subtitles(legend)
+        end
+      end
+
+      private def scale_scatter_point_size(x)
+        min = 0.5 * @default_marker_size**2
+        max = 2.0 * @default_marker_size**2
+
+        min + x * (max - min)
+      end
+
+      private def locator_to_legend_entries(locator, limits)
+        vmin, vmax = limits
+        raw_levels = locator.tick_values(vmin, vmax).to_a
+        raw_levels.reject! {|v| v < limits[0] || limits[1] < v }
+
+        formatter = case locator
+                    when Matplotlib.ticker.LogLocator
+                      Matplotlib.ticker.LogFormatter.new
+                    else
+                      Matplotlib.ticker.ScalarFormatter.new
+                    end
+
+        dummy_axis = Object.new
+        dummy_axis.define_singleton_method(:get_view_interval) { limits }
+        formatter.axis =  dummy_axis
+
+        formatter.set_locs(raw_levels)
+        formatted_levels = raw_levels.map {|x| formatter.(x) }
+
+        return raw_levels, formatted_levels
+      end
+
+      private def adjust_legend_subtitles(legend)
+        font_size = Matplotlib.rcParams.get("legend.title_fontsize", nil)
+        hpackers = legend.findobj(Matplotlib.offsetbox.VPacker)[0].get_children()
+        hpackers.each do |hpack|
+          draw_area, text_area = hpack.get_children()
+          handles = draw_area.get_children()
+          unless handles.all? {|a| a.get_visible() }
+            draw_area.set_width(0)
+            unless font_size.nil?
+              text_area.get_children().each do |text|
+                text.set_size(font_size)
+              end
+            end
           end
         end
       end
@@ -228,18 +556,42 @@ module Charty
         @pyplot.gca.set_xticks(Array(values))
       end
 
+      def set_yticks(values)
+        @pyplot.gca.set_yticks(Array(values))
+      end
+
       def set_xtick_labels(labels)
         @pyplot.gca.set_xticklabels(Array(labels).map(&method(:String)))
       end
 
+      def set_ytick_labels(labels)
+        @pyplot.gca.set_yticklabels(Array(labels).map(&method(:String)))
+      end
+
       def set_xlim(min, max)
         @pyplot.gca.set_xlim(Float(min), Float(max))
       end
 
+      def set_ylim(min, max)
+        @pyplot.gca.set_ylim(Float(min), Float(max))
+      end
+
       def disable_xaxis_grid
         @pyplot.gca.xaxis.grid(false)
       end
 
+      def disable_yaxis_grid
+        @pyplot.gca.xaxis.grid(false)
+      end
+
+      def invert_yaxis
+        @pyplot.gca.invert_yaxis
+      end
+
+      def legend(loc:, title:)
+        @pyplot.gca.legend(loc: loc, title: title)
+      end
+
       def show
         @pyplot.show
       end