GSoC 2025: My Journey with ArviZ
Introduction
This is a biweekly blog series chronicling my 12-week journey as a GSoC contributor this summer - a space where I’ll be documenting key milestones, reflections, and learnings from the project.
What is GSoC?
Google Summer of Code is a global program that pairs contributors with open source organizations to work on impactful coding projects over the summer. For many, it’s a first dive into real-world, collaborative software development - beyond classroom assignments or personal repos. GSoC offers mentorship, a stipend, and most importantly, the experience of writing code that matters. You don’t just fix bugs or add features; you contribute to projects used by people around the world. It’s a rewarding blend of challenge, learning, and purpose - and a chance to give back to the open source community that powers so much of the tech we rely on.
ArviZ and My Project
For my GSoC project, I am contributing to ArviZ, a Python library for the exploratory analysis and visualization of Bayesian inference models. These models are built on the principles of Bayesian statistics, where prior beliefs are combined with observed data to update our understanding of unknown parameters. Unlike traditional models that provide single-point estimates, Bayesian models yield full probability distributions, capturing uncertainty in a more nuanced way. ArviZ helps visualize these complex outputs and provides essential tools for model diagnostics, comparison, and criticism. The project I’m working on focuses on the ongoing refactor of ArviZ into a modular structure-specifically contributing to arviz-plots. My goal is to migrate and enhance core visualization tools from the original ArviZ, such as pair_plot, parallel_plot etc, while introducing new features like dark mode and improved support for circular variables in trace plot. The mentors of my project are Osvaldo A Martin, Oriol Abril Pla and Rohan Babbar
For further knowledge about ArviZ, you could refer to the documentation
Bonding Period (May 8 - June 2)
The GSoC bonding period began on May 8th (the day the results were announced) and continued until June 2nd. Prior to this, I had already been contributing to ArviZ since February, although there was a short pause in my activity due to my end-semester exams. During the bonding period, we had our onboarding meeting where we outlined the initial goals for the project. The first task was to investigate the differences between the plotting functionalities in the new arviz-plots package and the original ArviZ (now referred to as arviz-legacy).
While many plotting functions had already been ported to arviz-plots, several features from the legacy implementation were still missing-primarily because they hadn’t been an immediate priority or say they were not required in immediate future. For instance, the groupby functionality in plot_ppc had not yet been implemented.
My focus during the bonding period was to carefully identify and document these gaps. After conducting a detailed comparison of the plotting functions between the two versions and engaging in several discussions with my mentors, I compiled a prioritized list of features to work on during the coding phase of GSoC.
Week 1 & 2 (June 2 - June 16)
The GSoC coding period officially began on June 2nd, 2025, and my first task was to implement the pair_focus_plot function in arviz-plots. Although I had been contributing to ArviZ since February-mainly in the backend folder, where I worked on features like step_histogram and support for square root scaling across Matplotlib, Plotly, and Bokeh-this was my first major contribution involving a new batteries-included plot.
Unlike the isolated features I had worked on earlier, adding a full plotting function in arviz-plots required a deeper understanding of the library’s internal design, especially the faceting and aesthetic mapping managers: PlotCollection and PlotMatrix. These managers abstract much of the complexity of arranging plots, handling aesthetics, and coordinating behavior across subplots.
To get up to speed, I spent considerable time studying the ArviZ-Plots documentation, experimenting with different manager methods such as .wrap()
, .grid()
, and .map()
. Each serves a distinct purpose:
.wrap()
is used to instantiate a PlotCollection and iterate over subsets, wrapping plots in a grid layout..grid()
arranges plots based on row and column iterators..map()
applies the plotting function across all plots while respecting aesthetic mappings.
I also dove into the DataTree
structure, which is used internally to represent hierarchical data in ArviZ. Understanding how to navigate and manipulate this data was essential for writing a reliable and flexible plotting function.
Following is a glimpse of posterior group of the DataTree
I worked with:
What is pair_focus_plot?
The pair_focus_plot function creates visualizations to explore relationships between a focus variable and several other variables. It includes:
- A scatter plot for the focus variable vs. each other variable
- A divergence scatter plot highlighting divergent samples
- X and Y axis labels
To calculate the divergence mask, I referred to the implementation in the trace_plot function. After thoroughly testing different parameter combinations and edge cases with the PlotCollection manager, I was able to implement pair_focus_plot successfully.
Challenges Faced
Axis Sharing :
One of the design challenges involved axis and label sharing. Depending on the context:
- If different variables are plotted on the x-axis (with the same domain), it makes sense to share only the axis, not the labels.
- If the same variable is plotted repeatedly, both the axis and label should be shared. However, in our current design choice, Plotly always forces shared labels along with shared axes, otherwise it leads to overlapping labels and titles of plots in upper and lower rows. After discussions with the mentors, we decided to keep the current behavior for now and revisit this issue in a future refinement.
Removing Focus Variable to Avoid Diagonal Plot:
Another challenge was preventing a scatter plot of the focus variable against itself, which results in an uninformative diagonal line. This required removing the focus variable from the data. However, that introduced complications with maintaining coordinate dimensions in the DataTree. As a temporary workaround, we decided to allow plotting the focus variable against itself for now and revisit a cleaner solution later.
Additional Work
During this period, I also:
- Contributed to documentation using Sphinx.
- Wrote tests using the Hypothesis library (which I used for the first time).
- Gained a deeper appreciation for the faceting manager-it makes complex visualizations much more manageable for both users and contributors.
Final result
result of the pair_focus_plot function, which can be generated using the following code snippet:
from arviz_base import load_arviz_data
import arviz_plots as azp
azp.style.use("arviz-variat")
idata = load_arviz_data("centered_eight")
visuals = {"divergence": True }
print(idata)
pc = azp.plot_pair_focus(
idata,
focus_var=idata.posterior["theta"].sel(school="Choate"),
var_names=["tau","theta"],
backend="bokeh",
visuals = visuals,
)
pc.show()
Result
plotted theta-choate
against tau
and theta
variables, with divergence highlighted.
After wrapping up the implementation, documentation, and testing for pair_focus_plot, I moved on to my next task: implementing the pair_plot function.
Continuing with pair_plot: All Variables, All At Once
After completing pair_focus_plot, my next task was to implement pair_plot in arviz-plots. Conceptually, it’s similar to pair_focus_plot, but instead of plotting one focus variable against others, pair_plot visualizes all variables against each other. Each cell in the resulting grid displays a scatter plot of a variable pair, with additional support for highlighting divergent samples-just like in pair_focus_plot. Labels for rows and columns are added along the diagonal to help identify the plotted variables.
Getting Comfortable with PlotMatrix
Since this plot required creating a full matrix of plots, I shifted focus from PlotCollection to the PlotMatrix faceting manager. I spent some time getting comfortable with its methods and behavior-particularly how it manages grid construction and handles aesthetics across the matrix. Once I was confident with the API and internal flow, I began the implementation
A Key Challenge: Handling Divergence Masks using PlotMatrix
One interesting challenge I faced during this implementation was related to how PlotMatrix.map() works. Internally, it passes kwargs argument to the visual function by directly merging it into fun_kwargs. However, this becomes tricky when dealing with the divergence mask.
We only want to pass the relevant subset of the divergence mask corresponding to the current pair of variables being plotted-not the full mask. Initially, this wasn’t working as expected, and divergent data points weren’t being visualized.
To fix this, I modified the way divergence masks were passed so that only the appropriate slice of the mask reached the visual function. Once this adjustment was in place, the divergence scatter plots rendered correctly across the matrix.
Final Result
Result of pair_plot function, which can be generated using the following code snippet:
from arviz_base import load_arviz_data
import arviz_plots as azp
azp.style.use("arviz-variat")
data = load_arviz_data("centered_eight")
visuals = {"divergence": True}
pc = azp.plot_pair(
data,
backend="bokeh" ,
coords={"school": ["Choate", "Deerfield"]},
visuals=visuals,
)
pc.show()
Result
plotted mu
,tau
, theta-choate
and theta-deerfield
against each other, with divergence highlighted.
I will be adding documentation and tests for pair_plot in the coming weeks.
Week 3 & 4 (June 17 - July 1)
Deep Dive into pair_plot
and Beyond
After kicking off pair_plot
at the end of Week 2, I continued working on it through Weeks 3 and 4. While the concept initially seemed straightforward, implementing it turned out to be far more intricate. The challenge lay not just in plotting variable pairs, but in ensuring the visual output aligned with user expectations and supported features like flexible triangle types, axis sharing and diagonal marginals.
Getting Comfortable with PlotMatrix
Having just begun exploring PlotMatrix
in the previous weeks, I spent more time getting familiar with its internals. Unlike PlotCollection
, PlotMatrix
offers specialized methods like .map_triangle()
and .map()
for targeting specific regions of the grid, including the diagonal. However, we also needed functionality to map a function across an entire row or column, which wasn’t directly supported out of the box.
Initially, I attempted to create a new method called map_row_col
, borrowing logic from map_triangle
, to address this. While it worked, the approach introduced a lot of code duplication. After discussions with my mentors, we brainstormed a cleaner solution using the existing .map()
method in combination with additional metadata. We ended up extending PlotMatrix
with three new attributes - orientation
, fixed_var_name
, and fixed_selection
- and adjusted the get_target()
helper to support the new behavior. This allowed us to reuse existing logic efficiently and achieve the same result with cleaner code.
Enhancing the Diagonal: Adding Marginals
The diagonal plots in pair_plot
initially only displayed variable names as text labels, but we wanted to support marginal distributions by default. At first, I started implementing the marginal plots from scratch, but that led to unnecessary duplication. Eventually, we decided to reuse the existing plot_dist
function, which made the implementation more consistent and modular.
This refactor taught me a lot about good design practices—especially the importance of code reusability and maintaining clean abstraction layers. I found it incredibly rewarding to work through these challenges and get a deeper understanding of the design decisions behind ArviZ’s plotting architecture.
Results
Result of pair_plot
function, which can be generated using the following code snippet:
Code (triangle="lower"
by default)
from arviz_base import load_arviz_data
import arviz_plots as azp
azp.style.use("arviz-variat")
data = load_arviz_data("centered_eight")
pc = azp.plot_pair(
data,
backend="bokeh" ,
coords={"school": ["Choate", "Deerfield"]},
visuals={"divergence": True,},
)
pc.show()
Result (triangle="lower"
by default)
Code (triangle="upper"
)
from arviz_base import load_arviz_data
import arviz_plots as azp
azp.style.use("arviz-variat")
data = load_arviz_data("centered_eight")
pc = azp.plot_pair(
data,
backend="bokeh" ,
coords={"school": ["Choate", "Deerfield"]},
triangle="upper",
visuals={"divergence": True,},
)
pc.show()
Result (triangle="upper"
)
Code (triangle="both"
)
from arviz_base import load_arviz_data
import arviz_plots as azp
azp.style.use("arviz-variat")
data = load_arviz_data("centered_eight")
pc = azp.plot_pair(
data,
backend="bokeh" ,
coords={"school": ["Choate", "Deerfield"]},
triangle="both",
visuals={"divergence": True,},
)
pc.show()
Result (triangle="both"
)
Current and future plan (From Violin Plot to Parallel Plot)
Once pair_plot
was in good shape, I began working on a violin plot
. However, after receiving feedback from several contributors, we decided to drop the violin plot
in favor of focusing on more impactful visualizations.
I then shifted to implementing the parallel plot
, which I’m currently working on. Additionally, a new task was added to my to-do list: implementing filled faces in plot_dist
, which I’ll be tackling after the parallel plot
is complete.
These past two weeks were packed with problem-solving, learning, and refining-giving me a much clearer picture of ArviZ’s plotting system and how to contribute effectively to its evolving modular structure.
Week 5 & 6 (June 2 - July 16)
After we decided to drop the violin_plot
task in the previous weeks, I shifted my focus to implementing the parallel_plot
. While this plot didn’t require as many visual components as pair_plot
, it came with its own unique challenges particularly around data structure and aesthetic management.
Implementing parallel_plot
: A Structural Challenge
Unlike previous plots, parallel_plot
required all variables to be plotted on the same coordinate system, meaning we needed to restructure the data into a format that could accommodate this layout. The solution - we arrived at after a lot of brainstorming with my mentor - was to stack the sampling dimensions into a single multi-indexed dimension called sample
. This allowed each variable to be aligned along a common axis for the parallel coordinates plot
.
Once the structure was in place, we worked on several optimizations:
- Divergence highlighting: By adding divergence information as a coordinate to the data, we were able to plot both
divergent
andnon-divergent
samples in a single call tomultiple_lines
, making the implementation cleaner and more efficient. - Tick label management: With many variables plotted along the x-axis, label collisions became an issue. We introduced label rotation to improve readability. However, long variable names still cause crowding and plot compression-something we’ve noted for future refinement.
With these issues addressed, parallel_plot
was complete and functioning as intended.
Result
Enhancing dist_plot
: Supporting Filled Distributions
After wrapping up parallel_plot
, I moved on to enhancing dist_plot
to support filled faces (shaded areas under the curve), in addition to the existing line-based visualizations like KDE, ECDF, and histogram lines.
This enhancement, while seemingly simple, had its own complexities. Since dist_plot
is used internally by many other plotting functions-each passing its own set of aesthetic mappings-we had to ensure that filled faces respected all existing aesthetic variations. After carefully handling these edge cases, dist_plot
now supports filled visuals by default (with the option to disable them), making it more expressive and flexible.
This change also opens the door to generalizing ridge_plot
by making it use the updated dist_plot
.
Result
Behind the Scenes: Tests and Docs
In parallel with these implementations, I spent a significant amount of time refining documentation and improving tests for pair_plot
and parallel_plot
, ensuring better clarity, examples, and robust coverage.
Future Plan
My next focus will be implementing a dark theme for ArviZ plots and I will also keep refining the works done till now wherever required.
Week 7 & 8 (July 17 - July 31)
Over the past two weeks, I focused on improving the usability and flexibility parallel_plot
and dist_plot
, while also laying important groundwork for dark_theme
support.
Tackling the xtick label compression in parallel plot
One of the most pressing issues was the xtick
labels in the parallel_plot
. When these labels were rotated to prevent overlap-especially when visualizing multiple variables-they began consuming nearly 40–50% of the plot height, leading to severe compression of the actual visualization.
Initially, I experimented with creating a separate row for the labels, similar to how it’s handled in forest
and ridge
plots. However, this approach introduced unnecessary complexity. Following my mentor’s advice, I pivoted to a simpler yet effective solution: increasing the figure’s height dynamically during figsize calculation. This involved calculating the maximum label length after setting the font size, and then manually adjusting the height accordingly. This approach worked smoothly and resolved the plot compression issue.
Refactoring histogram variations: step vs filled
Another major focus was cleaning up the implementation of histogram variants. ArviZ-plots
supports both step histograms
and filled histograms
, but previously, both were handled by the same function. This led to API inconsistencies-for example, step histograms needed line-specific styling arguments (linestyle, linewidth), which were irrelevant for filled ones.
To address this, we:
- Separated the backend functions for
step
andfilled
histograms. - Added a dedicated
visual
function to call the appropriate backend method forstep
histograms.
We modified the dist_plot behavior by defaulting face_kwargs
to False
, as filled plots aren’t typically required in most use cases. Also, added corresponding tests to test behaviour of dist_plot
with face_kwargs=True
and face_kwargs=False
.
With this, the refactoring of dist_plot
is nearly complete.
Initiating dark theme support
I also began work on dark_theme
support, which brought a new set of challenges. A core part of this task was selecting a palette of colors that:
- Offer strong contrast against dark backgrounds.
- Remain color-blind friendly.
After consulting multiple accessibility and contrast resources, I finalized a suitable set of colors and implemented initial dark theme support.
However, a number of enhancements remain:
- Many tick labels, titles, and other text elements have manually hardcoded colors that don’t adapt to the theme.
- Bokeh (used in some visualizations) currently lacks support for styling the grid lines that separate subplots. On a dark background, this results in visible white grid lines that previously blended in with light themes.
Result
Future Plans:
In the coming weeks, my focus will shift toward:
- Dynamically adapting all label and title colors to the selected theme.
- Resolving Bokeh’s subplot grid line issue.
- Implementing support for circular variables in dist_plot, expanding its flexibility even further.
Week 9 & 10 (August 1 - August 15)
Over the past two weeks, we continued working on the dark theme
implementation. One of the most challenging parts was designing a color cycle that would remain both color-blind friendly
and have good contrast against a dark background
. Initially, we aimed for a cycle of 10 unique colors, but after several iterations and testing with my mentor, we decided to reduce the cycle to 8 colors. Interestingly, this set of colors worked well on both light and dark backgrounds, so we unified the color cycles of the arviz-variat
and arviz-tumma
themes.
Once the colors were finalized, the next hurdle was handling the hardcoded colors in the codebase. To make the visuals adapt dynamically based on the theme, we implemented two utility functions:
- Backend function – returns the background color corresponding to the active theme.
- Utility function – takes the background color and computes a contrasting color for text and other elements. This function can also return a contrasting gray variant if a
gray
flag is set. Internally, it calculates theYIQ brightness value
of the background color to determine whether a light or dark contrasting color should be chosen.
With these functions in place, we replaced all instances of hardcoded white, black, or gray colors with dynamically computed ones. This completed the dark theme implementation for the Tumma theme.
Results
Dark themed parallel plot
Dark themed bf plot
Dark themed forest plot
Next Steps
Looking ahead, our focus will shift to implementing the plot_lm function. This function is designed to plot the credible intervals and mean of posterior predictive, alongside a scatter plot of observed data, specifically for regression-like datasets.
We also decided to pause work on dist_plot enhancements (such as circular variable support
and dot plots
), since there is already ongoing work in that file.
Week 11 & 12 (August 16 - August 30)
During these weeks, we began implementing the plot_lm function. We initially decided that the plot should contain three main elements:
- Confidence interval (CI) band
- Mean line
- Observation scatter plot
To create these elements, the function needs access to four types of data:
- Observations for the target variable
- Observations for the independent variable
- Predictions or posterior predictives
- Independent variables corresponding to the predictions/posterior predictives (used to generate the CI band and mean line)
The most challenging part for me was understanding the structure of the data tree being used and the default behavior of the function. It took some time to fully grasp this, but once I did, I implemented the initial version of plot_lm.
This initial implementation included:
- A
CI band
andmean line
, calculated from eitherpredictions
orposterior predictives
and their corresponding independent variables.- If
predictions
were used, the independent variable came from theprediction_constant_data
group. - If
posterior predictives
were used, the independent variable came from theconstant data
group ( independent variable corresponding to observed data).
- If
- A
scatter plot
ofobserved data
, using the target variable from theobserved_data
group and the independent variable from theconstant_data
group.
After this, we extended the functionality to support:
- Multiple independent variables
- Multiple CI interval bands
- Median line as an alternative to the mean line
This required some changes to the way we handled data structures, but since the foundation was well set up earlier, the extension was not too difficult.We added support for multiple independent variables ( now it could handle multiple independent variables like x1, x2, x3 … which are used to predict single target variable y) and multiple CI interval bands ( now we can plot 90% interval band along with a 50% interval band together ). Additionally, we extended functionality to support a median line (in addition to the mean line) for predictions or posterior predictives.
We also discussed adding a smoothing function to smooth CI interval lines and mean lines. Once this function is implemented in arviz-stats, it can be easily integrated into plot_lm.
Results
lm_plot with multiple ci_bands and multiple independent variables
lm_plot with single ci_band and multiple independent variables
lm_plot with single ci_band and single independent variable
Ongoing Work
I am currently focusing on writing documentation and tests for the plot_lm function.