Merge pull request #108 from PolicyEngine/charts

Add chart utils

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    added:
+    - Chart formatting utilities.

docs/_config.yml

@@ -1,9 +1,10 @@
-title: PolicyEngine Core documentation
+title: PolicyEngine Core
 author: PolicyEngine
 copyright: "2022"
+logo: logo.png
 
 execute:
-  execute_notebooks: "force"
+  execute_notebooks: off
 
 repository:
   url: https://github.com/policyengine/policyengine-core
@@ -13,9 +14,11 @@ repository:
 sphinx:
   config:
     html_js_files:
-      - https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js
+      - https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.js
     html_theme: furo
     pygments_style: default
+    html_css_files:
+      - style.css
   extra_extensions:
     - "sphinx.ext.autodoc"
     - "sphinxarg.ext"

docs/_static/style.css

@@ -0,0 +1,9 @@
+@import url('https://fonts.googleapis.com/css2?family=Roboto+Serif:[email protected]&family=Roboto:wght@300&display=swap');
+
+h1, h2, h3, h4, h5, h6 {
+    font-family: "Roboto";
+}
+
+body {
+    font-family: "Roboto Serif";
+}

docs/_toc.yml

@@ -12,6 +12,7 @@ parts:
     - file: usage/parameters
     - file: usage/datasets
     - file: usage/reforms
+    - file: usage/charts
   - caption: Python API
     chapters:
     - file: python_api/commons

docs/contributing/intro.md

@@ -5,7 +5,7 @@ Any and all contributions are welcome to this project. You can help by:
 * Filing issues. Tell us about bugs you've found, or features you'd like to see.
 * Fixing issues. File a pull request to fix an issue you or someone else has filed.
 
-If you file an issue or a pull request, one of the maintainers (primarily @nikhilwoodruff) will respond to it within at least a week. If you don't hear back, feel free to ping us on the issue or pull request.
+If you file an issue or a pull request, one of the maintainers (primarily [@nikhilwoodruff](https://github.com/nikhilwoodruff)) will respond to it within at least a week. If you don't hear back, feel free to ping us on the issue or pull request.
 
 ## Changelog Entries
 

policyengine_core/charts/__init__.py

@@ -10,3 +10,5 @@
     format_fig,
     display_fig,
 )
+
+from .bar import *

policyengine_core/charts/bar.py

@@ -0,0 +1,133 @@
+import pandas as pd
+from .formatting import *
+import plotly.express as px
+from microdf import MicroSeries
+from typing import Callable
+import numpy as np
+
+
+def bar_chart(
+    data: pd.Series,
+    showlegend: bool = False,
+    remove_zero_index: bool = True,
+    text_format: str = "+.1%",
+    positive_colour: str = BLUE,
+    negative_colour: str = DARK_GRAY,
+    hover_text_function: Callable = None,
+    **kwargs,
+):
+    """Create a PolicyEngine bar chart.
+
+    Args:
+        data (pd.Series): A pandas series.
+        showlegend (bool, optional): Whether to show the legend. Defaults to False.
+        remove_zero_index (bool, optional): Whether to remove the zero index. Defaults to True.
+        text_format (str, optional): The format of the text. Defaults to "+.1%".
+        positive_colour (str, optional): The colour of positive values. Defaults to BLUE.
+        negative_colour (str, optional): The colour of negative values. Defaults to DARK_GRAY.
+        hover_text_labels (list, optional): The hover text labels. Defaults to None.
+
+    Returns:
+        go.Figure: A plotly figure.
+    """
+
+    hover_text_labels = [
+        hover_text_function(index, value)
+        if hover_text_function is not None
+        else None
+        for index, value in data.items()
+    ]
+
+    fig = (
+        px.bar(
+            data,
+            text=data.apply(lambda x: f"{x:{text_format}}"),
+            custom_data=[hover_text_labels] if hover_text_labels else None,
+        )
+        .update_layout(
+            showlegend=showlegend,
+            uniformtext=dict(
+                mode="hide",
+                minsize=12,
+            ),
+            **kwargs,
+        )
+        .update_traces(
+            marker_color=[
+                positive_colour if v > 0 else negative_colour
+                for v in data.values
+            ],
+            hovertemplate="%{customdata[0]}<extra></extra>"
+            if hover_text_labels is not None
+            else None,
+        )
+    )
+    return format_fig(fig)
+
+
+def cross_section_bar_chart(
+    data: MicroSeries,
+    cross_section: MicroSeries,
+    slices: list = [-0.05, -0.01, 0.01, 0.05],
+    add_infinities: bool = True,
+    text_format: str = ".1%",
+    hover_text_function: Callable = None,
+    category_names=None,
+    color_discrete_map: dict = None,
+    **kwargs,
+):
+    df = pd.DataFrame()
+    slices = [-np.inf, *slices, np.inf] if add_infinities else slices
+    for i, lower, upper in zip(range(len(slices)), slices[:-1], slices[1:]):
+        for cross_section_value in cross_section.unique():
+            in_slice = (data >= lower) * (data < upper)
+            value = (
+                data[cross_section == cross_section_value][in_slice].count()
+                / data[cross_section == cross_section_value].count()
+            )
+            category = (
+                category_names[i]
+                if category_names is not None
+                else f"{lower:.0%} to {upper:.0%}"
+            )
+            row = {
+                "Category": category,
+                "Cross section": cross_section_value,
+                "Value": value,
+                "Hover text": hover_text_function(
+                    cross_section_value, category, value
+                )
+                if hover_text_function is not None
+                else None,
+            }
+            df = df.append(row, ignore_index=True)
+
+    fig = (
+        px.bar(
+            df,
+            x="Value",
+            y="Cross section",
+            color="Category",
+            barmode="stack",
+            orientation="h",
+            text=df["Value"].apply(lambda x: f"{x:{text_format}}"),
+            color_discrete_map=color_discrete_map,
+            custom_data=["Hover text"],
+        )
+        .update_layout(
+            xaxis=dict(
+                tickformat="%",
+                title="Fraction of observations",
+            ),
+            yaxis=dict(tickvals=list(range(1, 11))),
+            uniformtext=dict(
+                mode="hide",
+                minsize=12,
+            ),
+            **kwargs,
+        )
+        .update_traces(
+            hovertemplate="%{customdata[0]}<extra></extra>",
+        )
+    )
+    return format_fig(fig)

policyengine_core/charts/formatting.py

@@ -50,12 +50,6 @@ def format_fig(fig: go.Figure) -> go.Figure:
         template="plotly_white",
         height=600,
         width=800,
-        margin=dict(
-            t=100,
-            b=100,
-            l=100,
-            r=100,
-        ),
     )
     # don't show modebar
     fig.update_layout(
@@ -71,3 +65,16 @@ def display_fig(fig: go.Figure) -> HTML:
     return HTML(
         format_fig(fig).to_html(full_html=False, include_plotlyjs="cdn")
     )
+
+
+def cardinal(n: int) -> int:
+    """Convert an integer to a cardinal string."""
+    ending_number = n % 10
+    if ending_number == 1:
+        return f"{n}st"
+    elif ending_number == 2:
+        return f"{n}nd"
+    elif ending_number == 3:
+        return f"{n}rd"
+    else:
+        return f"{n}th"