[SuperBarChart] A stacked barchart object using Pandas dataframes. #manim #barchart #pandas #dataframe

20250210_superbarchart.py

# https://discord.com/channels/581738731934056449/1338585711648968775/1338585711648968775
from manim import *
from pandas import DataFrame

class SuperBarChart(VGroup):
    def __init__(self,
            x_length,
            y_range,
            y_length,
            dataframe,
            colors=[[RED,YELLOW,GREEN,BLUE,MAROON]],
            buff = 0.2,
            label_buff = 0.2,
            bar_stroke_width = 1,
            bar_fill_opacity = 0.6,
            **kwargs
        ):
        self.buff        = buff
        self.label_buff  = label_buff
        self.colors = colors
        self.bar_stroke_width = bar_stroke_width
        self.bar_fill_opacity =  bar_fill_opacity
        self.barwidth    = (1-buff)/len(dataframe.columns)
        self.axes = Axes(
            x_range=[0,len(dataframe.index)+1,1],
            x_length=x_length,
            y_range=y_range,
            y_length=y_length,
            tips=False,
            x_axis_config={"include_numbers":False},
            y_axis_config={"include_numbers":True},
        )
        self.unit_x = (self.axes.c2p(1,0)-self.axes.c2p(0,0))[0]
        self.unit_y = (self.axes.c2p(0,1)-self.axes.c2p(0,0))[1]
        VGroup.__init__(self, **kwargs)
        self.add(self.axes)

        self.labels = VGroup()
        self.series = VGroup()
        self.update_chartdata(dataframe=dataframe)

        self.add(self.labels)
        self.add(self.series)
        self.center()

    def update_chartdata(self, dataframe):
        self.dataframe   = dataframe
        self.indices     = self.dataframe.index
        self.seriesnames = self.dataframe.columns
        labels = Tex(*[label for label in self.indices])
        newlabels = VGroup(
            labels[i].shift(self.axes.c2p(i+1,0)-[labels[i].get_x(),labels.get_top()[1]+self.label_buff,0])
            for i in range(len(labels))
        )
        self.labels.become(newlabels)

        series = VGroup()
        for j,seriename in enumerate(self.seriesnames):
            serie = VGroup()
            datas = self.dataframe[seriename]
            colorscale = self.colors[j % len(self.colors)]
            for x,data in enumerate(datas):
                if hasattr(data,"__len__"):
                    bar = VGroup(
                        Rectangle(
                            width=self.unit_x*self.barwidth,
                            height=datapoint*self.unit_y,
                            color=colorscale[i % len(colorscale)],
                            stroke_width=self.bar_stroke_width,
                            fill_opacity = self.bar_fill_opacity,
                        )
                        for i,datapoint in enumerate(data)
                    ).arrange(UP, buff=0)
                else:
                    bar = VGroup(
                        Rectangle(
                            width=self.unit_x*self.barwidth,
                            height=data*self.unit_y,
                            color=colorscale[0],
                            stroke_width=self.bar_stroke_width,
                            fill_opacity = self.bar_fill_opacity,
                        )
                    )
                bar.shift(self.axes.c2p(x+0.5*(1+self.buff)+(0.5+j)*self.barwidth)-bar.get_bottom())
                serie += bar
            series += serie
        self.series.become(series)

        
class superbarcharttest(Scene):
    def construct(self):
        dataframe = DataFrame(
            {
                "one": [[1.0, 1.2, 1.5], 2.0, 3.0, 4.0],
                "two": [4.0, 3.0, 2.0, 1.0],
                "three": [3.0, 2.0, 1.0, 4.0 ],
            },
            index = ["a", "b", "c", "d"]
        )
        chart = SuperBarChart(
            x_length=12,
            y_range=[0,20,5],
            y_length=6,
            dataframe = dataframe,
            colors=[[RED,YELLOW,GREEN,BLUE,MAROON],[TEAL],[ORANGE]],
        )
        self.add(chart)
        
class superbarcharttest_dynamic(Scene):
    def construct(self):
        dataframe = DataFrame(
            {
                "one": [[1.0, 1.2, 1.5], 2.0, 3.0, 4.0],
                "two": [4.0, 3.0, 2.0, 1.0],
                "three": [3.0, 2.0, 1.0, [3,3.0] ],
            },
            index = ["a", "b", "c", "d"]
        )
        chart = SuperBarChart(
            x_length=12,
            y_range=[0,7,1],
            y_length=6,
            dataframe = dataframe,
            colors=[[RED,YELLOW,GREEN,BLUE,MAROON],[TEAL],[ORANGE,BLUE]],
        )
        self.play(Create(chart.axes))
        self.play(Write(chart.labels))
        for serie in chart.series:
            self.play(Create(serie))
        self.wait()
        dataframe2 = DataFrame(
            {
                "one": [[1.5, 1.2, 1.0], 4.0, 2.0, 1.0],
                "two": [1.0, 3.0, 4.0, 1.0],
                "three": [5.0, 0.5, 1.0, [1,5] ],
            },
            index = ["a", "b", "c", "d"]
        )
        self.play(chart.animate.update_chartdata(dataframe2))
        self.wait()

readme.md