policyengine-core 2.4.0__tar.gz → 2.5.0__tar.gz

{policyengine-core-2.4.0/policyengine_core.egg-info → policyengine-core-2.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: policyengine-core
-Version: 2.4.0
+Version: 2.5.0
 Summary: Core microsimulation engine enabling country-specific policy models.
 Home-page: https://github.com/policyengine/policyengine-core
 Author: PolicyEngine

{policyengine-core-2.4.0 → policyengine-core-2.5.0}/policyengine_core/charts/__init__.py

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

policyengine-core-2.5.0/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-2.4.0 → policyengine-core-2.5.0}/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"

{policyengine-core-2.4.0 → policyengine-core-2.5.0/policyengine_core.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: policyengine-core
-Version: 2.4.0
+Version: 2.5.0
 Summary: Core microsimulation engine enabling country-specific policy models.
 Home-page: https://github.com/policyengine/policyengine-core
 Author: PolicyEngine

{policyengine-core-2.4.0 → policyengine-core-2.5.0}/policyengine_core.egg-info/SOURCES.txt

@@ -11,6 +11,7 @@ policyengine_core.egg-info/entry_points.txt
 policyengine_core.egg-info/requires.txt
 policyengine_core.egg-info/top_level.txt
 policyengine_core/charts/__init__.py
+policyengine_core/charts/bar.py
 policyengine_core/charts/formatting.py
 policyengine_core/commons/__init__.py
 policyengine_core/commons/formulas.py

{policyengine-core-2.4.0 → policyengine-core-2.5.0}/setup.py

@@ -48,7 +48,7 @@ dev_requirements = [
 
 setup(
     name="policyengine-core",
-    version="2.4.0",
+    version="2.5.0",
     author="PolicyEngine",
     author_email="hello@policyengine.org",
     classifiers=[