workforce-plotting 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- workforce_plotting-0.2.0/LICENSE +21 -0
- workforce_plotting-0.2.0/MANIFEST.in +1 -0
- workforce_plotting-0.2.0/PKG-INFO +135 -0
- workforce_plotting-0.2.0/README.md +111 -0
- workforce_plotting-0.2.0/examples/paper_examples.py +107 -0
- workforce_plotting-0.2.0/pyproject.toml +24 -0
- workforce_plotting-0.2.0/setup.cfg +4 -0
- workforce_plotting-0.2.0/src/workforce_plotting/__init__.py +9 -0
- workforce_plotting-0.2.0/src/workforce_plotting/figures.py +158 -0
- workforce_plotting-0.2.0/src/workforce_plotting.egg-info/PKG-INFO +135 -0
- workforce_plotting-0.2.0/src/workforce_plotting.egg-info/SOURCES.txt +13 -0
- workforce_plotting-0.2.0/src/workforce_plotting.egg-info/dependency_links.txt +1 -0
- workforce_plotting-0.2.0/src/workforce_plotting.egg-info/requires.txt +14 -0
- workforce_plotting-0.2.0/src/workforce_plotting.egg-info/top_level.txt +1 -0
- workforce_plotting-0.2.0/tests/test_plots.py +101 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Corey Gelb-Bicknell
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
recursive-include examples *.py
|
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: workforce-plotting
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Simple plotnine figures with Opportunity Insights styling.
|
|
5
|
+
Author: Corey Gelb-Bicknell
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Repository, https://github.com/coreygb1/workforce-plotting
|
|
8
|
+
Requires-Python: >=3.12
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Dist: plotnine<0.16,>=0.15
|
|
12
|
+
Requires-Dist: pandas>=2.2
|
|
13
|
+
Requires-Dist: oi-tools>=0.11.7
|
|
14
|
+
Provides-Extra: polars
|
|
15
|
+
Requires-Dist: polars>=1; extra == "polars"
|
|
16
|
+
Requires-Dist: pyarrow>=14; extra == "polars"
|
|
17
|
+
Provides-Extra: dev
|
|
18
|
+
Requires-Dist: pytest>=8; extra == "dev"
|
|
19
|
+
Requires-Dist: build; extra == "dev"
|
|
20
|
+
Requires-Dist: twine; extra == "dev"
|
|
21
|
+
Requires-Dist: polars>=1; extra == "dev"
|
|
22
|
+
Requires-Dist: pyarrow>=14; extra == "dev"
|
|
23
|
+
Dynamic: license-file
|
|
24
|
+
|
|
25
|
+
# workforce-plotting
|
|
26
|
+
|
|
27
|
+
Five small plotnine functions with a shared Opportunity Insights style. Supply
|
|
28
|
+
prepared estimates and confidence bounds; the functions only draw them. Every
|
|
29
|
+
function returns a `plotnine.ggplot` that you can edit with ordinary layers.
|
|
30
|
+
|
|
31
|
+
## Install
|
|
32
|
+
|
|
33
|
+
Python 3.12 or newer is required (matching `oi-tools`). For a Polars workflow:
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
uv add 'workforce-plotting[polars]'
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
Or use pip:
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
pip install 'workforce-plotting[polars]'
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
For local development from this checkout:
|
|
46
|
+
|
|
47
|
+
```bash
|
|
48
|
+
pip install -e '.[dev]'
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
For Polars input, install the `polars` extra, which also supplies the Arrow
|
|
52
|
+
conversion dependency. Pandas input needs no extra.
|
|
53
|
+
|
|
54
|
+
## Functions
|
|
55
|
+
|
|
56
|
+
| Function | Use |
|
|
57
|
+
|---|---|
|
|
58
|
+
| `plot_lines(data, x=..., y=..., series=...)` | Event studies, trajectories, cohort trends, model complexity |
|
|
59
|
+
| `plot_validation(data, x=..., y=...)` | RCT versus digital-twin estimates with an identity line |
|
|
60
|
+
| `plot_bars(data, x=..., y=..., series=...)` | Already-calculated values by category |
|
|
61
|
+
| `plot_scatter(data, x=..., y=..., fitted=...)` | Points with an optional precomputed fitted line |
|
|
62
|
+
| `plot_intervals(data, x=..., y=..., lower=..., upper=...)` | Estimates and supplied confidence intervals |
|
|
63
|
+
| `theme_workforce()` | Shared white background, gray axes, sans-serif type, no grid |
|
|
64
|
+
|
|
65
|
+
`plot_lines` accepts `lower` and `upper` columns and `points=False`.
|
|
66
|
+
`plot_validation` accepts `x_lower`, `x_upper`, `y_lower`, `y_upper`,
|
|
67
|
+
`gaps="horizontal"` or `"vertical"`, and a common `(min, max)` pair of `limits`.
|
|
68
|
+
The identity line has slope one in data units; the panel follows the figure's
|
|
69
|
+
aspect ratio. Other functions accept `series` for discrete groups. Use ordered
|
|
70
|
+
categorical columns to control color/legend order. Colors are teal, orange, navy,
|
|
71
|
+
purple, blue, red, and dark green from the public `oi_tools.figures.OIColors` API.
|
|
72
|
+
For a different mapping, add a plotnine manual color or fill scale.
|
|
73
|
+
|
|
74
|
+
```python
|
|
75
|
+
import plotnine as pn
|
|
76
|
+
from workforce_plotting import plot_lines
|
|
77
|
+
|
|
78
|
+
p = plot_lines(df, x="year", y="earnings", series="group")
|
|
79
|
+
p += pn.labs(x="Years since enrollment", y="Mean annual earnings")
|
|
80
|
+
p.save("earnings.svg")
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
Dataframes should be tidy: one row per plotted point, with separate columns for
|
|
84
|
+
x, y, and an optional series. Lines connect points in x order within each series.
|
|
85
|
+
Include a row with a missing y where a line should break; simply omitting a row
|
|
86
|
+
connects the remaining observations. Fitted lines use the same column names as
|
|
87
|
+
the scatter data. No smoothing, RMSE calculation, confidence-interval estimation,
|
|
88
|
+
or automatic saving happens inside the plotting functions.
|
|
89
|
+
|
|
90
|
+
Add `pn.coord_flip()` for horizontal bars/intervals; use ordinary plotnine scales
|
|
91
|
+
for dollars, percentages, limits, or log axes. Fonts, dimensions, and text size
|
|
92
|
+
use Lucida Sans when installed and DejaVu Sans otherwise. They can be overridden with `theme_workforce(base_family=..., base_size=..., width=..., height=...)`.
|
|
93
|
+
Titles, logos, notes, and paper-specific annotations belong in the example/caller.
|
|
94
|
+
|
|
95
|
+
## Styling your own plots
|
|
96
|
+
|
|
97
|
+
Use the theme and colors without using the plot constructors. `WFColors` is an
|
|
98
|
+
alias of `oi_tools.figures.OIColors`, and `theme_workforce` builds on `theme_oi`.
|
|
99
|
+
The `oi-tools` dependency remains; you only need workforce imports in your code.
|
|
100
|
+
All plotting functions and styling live together in `workforce_plotting/figures.py`.
|
|
101
|
+
|
|
102
|
+
```python
|
|
103
|
+
import polars as pl
|
|
104
|
+
import plotnine as pn
|
|
105
|
+
from workforce_plotting import theme_workforce, WFColors
|
|
106
|
+
|
|
107
|
+
df = pl.DataFrame({"year": [2020, 2021, 2022], "earnings": [20000, 23000, 25000]})
|
|
108
|
+
p = (
|
|
109
|
+
pn.ggplot(df.to_pandas(), pn.aes(x="year", y="earnings"))
|
|
110
|
+
+ pn.geom_line(color=WFColors.TEAL)
|
|
111
|
+
+ pn.geom_point(color=WFColors.TEAL)
|
|
112
|
+
+ theme_workforce()
|
|
113
|
+
+ pn.labs(x="Year", y="Annual earnings")
|
|
114
|
+
)
|
|
115
|
+
p.save("earnings.svg")
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
Plot constructors accept Polars directly. When building a plotnine plot yourself,
|
|
119
|
+
the explicit conversion above works with the supported plotnine version.
|
|
120
|
+
|
|
121
|
+
## Editable examples
|
|
122
|
+
|
|
123
|
+
`examples/paper_examples.py` contains synthetic recipes for the paper/deck plot
|
|
124
|
+
families. These are invented values, not research results.
|
|
125
|
+
|
|
126
|
+
```bash
|
|
127
|
+
python examples/paper_examples.py outputs/examples
|
|
128
|
+
pytest
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
## Changes from 0.1
|
|
132
|
+
|
|
133
|
+
This is a fresh API. The old Matplotlib functions, plotting dataclasses, curve
|
|
134
|
+
fitting machinery, and `aipw_analysis` compatibility package have been removed.
|
|
135
|
+
Plot creation and saving are now separate. The public API consists of the five plot functions, `theme_workforce`, and `WFColors`.
|
|
@@ -0,0 +1,111 @@
|
|
|
1
|
+
# workforce-plotting
|
|
2
|
+
|
|
3
|
+
Five small plotnine functions with a shared Opportunity Insights style. Supply
|
|
4
|
+
prepared estimates and confidence bounds; the functions only draw them. Every
|
|
5
|
+
function returns a `plotnine.ggplot` that you can edit with ordinary layers.
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
Python 3.12 or newer is required (matching `oi-tools`). For a Polars workflow:
|
|
10
|
+
|
|
11
|
+
```bash
|
|
12
|
+
uv add 'workforce-plotting[polars]'
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
Or use pip:
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
pip install 'workforce-plotting[polars]'
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
For local development from this checkout:
|
|
22
|
+
|
|
23
|
+
```bash
|
|
24
|
+
pip install -e '.[dev]'
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
For Polars input, install the `polars` extra, which also supplies the Arrow
|
|
28
|
+
conversion dependency. Pandas input needs no extra.
|
|
29
|
+
|
|
30
|
+
## Functions
|
|
31
|
+
|
|
32
|
+
| Function | Use |
|
|
33
|
+
|---|---|
|
|
34
|
+
| `plot_lines(data, x=..., y=..., series=...)` | Event studies, trajectories, cohort trends, model complexity |
|
|
35
|
+
| `plot_validation(data, x=..., y=...)` | RCT versus digital-twin estimates with an identity line |
|
|
36
|
+
| `plot_bars(data, x=..., y=..., series=...)` | Already-calculated values by category |
|
|
37
|
+
| `plot_scatter(data, x=..., y=..., fitted=...)` | Points with an optional precomputed fitted line |
|
|
38
|
+
| `plot_intervals(data, x=..., y=..., lower=..., upper=...)` | Estimates and supplied confidence intervals |
|
|
39
|
+
| `theme_workforce()` | Shared white background, gray axes, sans-serif type, no grid |
|
|
40
|
+
|
|
41
|
+
`plot_lines` accepts `lower` and `upper` columns and `points=False`.
|
|
42
|
+
`plot_validation` accepts `x_lower`, `x_upper`, `y_lower`, `y_upper`,
|
|
43
|
+
`gaps="horizontal"` or `"vertical"`, and a common `(min, max)` pair of `limits`.
|
|
44
|
+
The identity line has slope one in data units; the panel follows the figure's
|
|
45
|
+
aspect ratio. Other functions accept `series` for discrete groups. Use ordered
|
|
46
|
+
categorical columns to control color/legend order. Colors are teal, orange, navy,
|
|
47
|
+
purple, blue, red, and dark green from the public `oi_tools.figures.OIColors` API.
|
|
48
|
+
For a different mapping, add a plotnine manual color or fill scale.
|
|
49
|
+
|
|
50
|
+
```python
|
|
51
|
+
import plotnine as pn
|
|
52
|
+
from workforce_plotting import plot_lines
|
|
53
|
+
|
|
54
|
+
p = plot_lines(df, x="year", y="earnings", series="group")
|
|
55
|
+
p += pn.labs(x="Years since enrollment", y="Mean annual earnings")
|
|
56
|
+
p.save("earnings.svg")
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
Dataframes should be tidy: one row per plotted point, with separate columns for
|
|
60
|
+
x, y, and an optional series. Lines connect points in x order within each series.
|
|
61
|
+
Include a row with a missing y where a line should break; simply omitting a row
|
|
62
|
+
connects the remaining observations. Fitted lines use the same column names as
|
|
63
|
+
the scatter data. No smoothing, RMSE calculation, confidence-interval estimation,
|
|
64
|
+
or automatic saving happens inside the plotting functions.
|
|
65
|
+
|
|
66
|
+
Add `pn.coord_flip()` for horizontal bars/intervals; use ordinary plotnine scales
|
|
67
|
+
for dollars, percentages, limits, or log axes. Fonts, dimensions, and text size
|
|
68
|
+
use Lucida Sans when installed and DejaVu Sans otherwise. They can be overridden with `theme_workforce(base_family=..., base_size=..., width=..., height=...)`.
|
|
69
|
+
Titles, logos, notes, and paper-specific annotations belong in the example/caller.
|
|
70
|
+
|
|
71
|
+
## Styling your own plots
|
|
72
|
+
|
|
73
|
+
Use the theme and colors without using the plot constructors. `WFColors` is an
|
|
74
|
+
alias of `oi_tools.figures.OIColors`, and `theme_workforce` builds on `theme_oi`.
|
|
75
|
+
The `oi-tools` dependency remains; you only need workforce imports in your code.
|
|
76
|
+
All plotting functions and styling live together in `workforce_plotting/figures.py`.
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
import polars as pl
|
|
80
|
+
import plotnine as pn
|
|
81
|
+
from workforce_plotting import theme_workforce, WFColors
|
|
82
|
+
|
|
83
|
+
df = pl.DataFrame({"year": [2020, 2021, 2022], "earnings": [20000, 23000, 25000]})
|
|
84
|
+
p = (
|
|
85
|
+
pn.ggplot(df.to_pandas(), pn.aes(x="year", y="earnings"))
|
|
86
|
+
+ pn.geom_line(color=WFColors.TEAL)
|
|
87
|
+
+ pn.geom_point(color=WFColors.TEAL)
|
|
88
|
+
+ theme_workforce()
|
|
89
|
+
+ pn.labs(x="Year", y="Annual earnings")
|
|
90
|
+
)
|
|
91
|
+
p.save("earnings.svg")
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
Plot constructors accept Polars directly. When building a plotnine plot yourself,
|
|
95
|
+
the explicit conversion above works with the supported plotnine version.
|
|
96
|
+
|
|
97
|
+
## Editable examples
|
|
98
|
+
|
|
99
|
+
`examples/paper_examples.py` contains synthetic recipes for the paper/deck plot
|
|
100
|
+
families. These are invented values, not research results.
|
|
101
|
+
|
|
102
|
+
```bash
|
|
103
|
+
python examples/paper_examples.py outputs/examples
|
|
104
|
+
pytest
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
## Changes from 0.1
|
|
108
|
+
|
|
109
|
+
This is a fresh API. The old Matplotlib functions, plotting dataclasses, curve
|
|
110
|
+
fitting machinery, and `aipw_analysis` compatibility package have been removed.
|
|
111
|
+
Plot creation and saving are now separate. The public API consists of the five plot functions, `theme_workforce`, and `WFColors`.
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
"""Editable paper/deck recipes using invented values, not research results.
|
|
2
|
+
|
|
3
|
+
Run: python examples/paper_examples.py /tmp/workforce-plotting-examples
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
import sys
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
|
|
9
|
+
import pandas as pd
|
|
10
|
+
import plotnine as pn
|
|
11
|
+
from mizani.labels import label_dollar
|
|
12
|
+
|
|
13
|
+
from workforce_plotting import (
|
|
14
|
+
plot_bars, plot_intervals, plot_lines, plot_scatter, plot_validation,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def examples():
|
|
19
|
+
# Event study: prepared means and bounds, connected separately by group.
|
|
20
|
+
event = pd.DataFrame({
|
|
21
|
+
"year": [-2, -1, 0, 1, 2, 3] * 3,
|
|
22
|
+
"group": ["Treated"] * 6 + ["RCT control"] * 6 + ["Digital twin"] * 6,
|
|
23
|
+
"earnings": [12, 14, 15, 22, 27, 30, 12, 14, 15, 18, 21, 23,
|
|
24
|
+
12, 14, 15, 19, 21, 24],
|
|
25
|
+
})
|
|
26
|
+
event["group"] = pd.Categorical(event["group"],
|
|
27
|
+
categories=["Treated", "RCT control", "Digital twin"], ordered=True)
|
|
28
|
+
event["earnings"] *= 1000
|
|
29
|
+
event["lower"] = event["earnings"] - 1200
|
|
30
|
+
event["upper"] = event["earnings"] + 1200
|
|
31
|
+
earnings = (plot_lines(event, x="year", y="earnings", series="group",
|
|
32
|
+
lower="lower", upper="upper")
|
|
33
|
+
+ pn.scale_y_continuous(labels=label_dollar(scale=0.001, suffix="K", accuracy=1))
|
|
34
|
+
+ pn.labs(x="Years since enrollment", y="Mean annual earnings")
|
|
35
|
+
+ pn.theme(legend_position=(0.2, 0.82), legend_direction="vertical")
|
|
36
|
+
+ pn.guides(color=pn.guide_legend(override_aes={"linetype": ""})))
|
|
37
|
+
|
|
38
|
+
# Complexity: supplied RMSE at each sample share, on a log x-axis.
|
|
39
|
+
complexity = pd.DataFrame({
|
|
40
|
+
"share": [1, 3.33, 10, 33.3, 100] * 4,
|
|
41
|
+
"model": ["Full population"] * 5 + ["1%"] * 5 + ["0.04% (CPS-sized)"] * 5 + ["0.02%"] * 5,
|
|
42
|
+
"rmse": [7400, 3900, 2200, 1900, 1600, 6200, 3900, 3600, 3300, 3000,
|
|
43
|
+
6100, 3400, 2800, 2400, 2300, 6300, 4800, 4400, 4200, 4200],
|
|
44
|
+
})
|
|
45
|
+
complexity["model"] = pd.Categorical(complexity["model"],
|
|
46
|
+
categories=["Full population", "1%", "0.04% (CPS-sized)", "0.02%"], ordered=True)
|
|
47
|
+
complexity_plot = (plot_lines(complexity, x="share", y="rmse", series="model")
|
|
48
|
+
+ pn.scale_x_log10(breaks=[1, 3.33, 10, 33.3, 100],
|
|
49
|
+
labels=["1%", "3.33%", "10%", "33.3%", "100%"])
|
|
50
|
+
+ pn.scale_y_continuous(breaks=[0, 2000, 4000, 6000, 8000],
|
|
51
|
+
labels=["0", "$2K", "$4K", "$6K", "$8K"])
|
|
52
|
+
+ pn.coord_cartesian(ylim=(0, 8000))
|
|
53
|
+
+ pn.labs(x="Model complexity, share of tuned model", y="Noise-adjusted RMSE")
|
|
54
|
+
+ pn.theme(legend_position=(0.80, 0.78), legend_direction="vertical")
|
|
55
|
+
+ pn.guides(color=pn.guide_legend(override_aes={"linetype": ""})))
|
|
56
|
+
|
|
57
|
+
validation = pd.DataFrame({
|
|
58
|
+
"rct": [-2500, -1000, 200, 1500, 2500, 4000, 6000, 8000, 10000],
|
|
59
|
+
"twin": [-1500, -1800, 900, 2500, 1800, 5200, 6500, 7200, 11000],
|
|
60
|
+
})
|
|
61
|
+
validation["lower"] = validation["twin"] - 900
|
|
62
|
+
validation["upper"] = validation["twin"] + 900
|
|
63
|
+
validation_plot = (plot_validation(validation, x="twin", y="rct",
|
|
64
|
+
gaps="horizontal", limits=(-5000, 12000))
|
|
65
|
+
+ pn.scale_x_continuous(labels=label_dollar(scale=0.001, suffix="K", accuracy=1))
|
|
66
|
+
+ pn.scale_y_continuous(labels=label_dollar(scale=0.001, suffix="K", accuracy=1))
|
|
67
|
+
+ pn.labs(x="Digital-twin estimate", y="RCT estimate"))
|
|
68
|
+
|
|
69
|
+
validation_ci = (plot_validation(validation, x="twin", y="rct",
|
|
70
|
+
x_lower="lower", x_upper="upper", limits=(-5000, 12000))
|
|
71
|
+
+ pn.scale_x_continuous(labels=label_dollar(scale=0.001, suffix="K", accuracy=1))
|
|
72
|
+
+ pn.scale_y_continuous(labels=label_dollar(scale=0.001, suffix="K", accuracy=1))
|
|
73
|
+
+ pn.labs(x="Digital-twin estimate", y="RCT estimate"))
|
|
74
|
+
|
|
75
|
+
comparisons = pd.DataFrame({"specification": ["Demographics", "+ Earnings", "+ Parents"],
|
|
76
|
+
"rmse": [3500, 2200, 1800]})
|
|
77
|
+
bars = (plot_bars(comparisons, x="specification", y="rmse")
|
|
78
|
+
+ pn.scale_x_discrete(limits=comparisons["specification"].tolist())
|
|
79
|
+
+ pn.scale_y_continuous(
|
|
80
|
+
labels=label_dollar(scale=0.001, suffix="K", accuracy=1),
|
|
81
|
+
expand=(0, 0, 0.05, 0))
|
|
82
|
+
+ pn.labs(x="", y="Noise-adjusted RMSE"))
|
|
83
|
+
|
|
84
|
+
placement = pd.DataFrame({"placement": [20, 30, 40, 50, 60],
|
|
85
|
+
"effect": [1500, 2600, 2100, 4200, 4800]})
|
|
86
|
+
# This line is supplied by the analysis; plotting does not fit it.
|
|
87
|
+
fitted = pd.DataFrame({"placement": [20, 60], "effect": [1300, 4700]})
|
|
88
|
+
scatter = (plot_scatter(placement, x="placement", y="effect", fitted=fitted)
|
|
89
|
+
+ pn.labs(x="Placement rate (%)", y="Earnings effect"))
|
|
90
|
+
|
|
91
|
+
intervals = pd.DataFrame({"program": ["Program A", "Program B", "Program C"],
|
|
92
|
+
"effect": [1200, 2500, -500], "lower": [200, 1400, -1500],
|
|
93
|
+
"upper": [2200, 3600, 500]})
|
|
94
|
+
interval_plot = (plot_intervals(intervals, x="program", y="effect",
|
|
95
|
+
lower="lower", upper="upper") + pn.coord_flip()
|
|
96
|
+
+ pn.labs(x="", y="Earnings effect"))
|
|
97
|
+
return {"event_study": earnings, "complexity": complexity_plot,
|
|
98
|
+
"validation": validation_plot, "validation_ci": validation_ci, "bars": bars,
|
|
99
|
+
"scatter": scatter, "intervals": interval_plot}
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
if __name__ == "__main__":
|
|
103
|
+
output = Path(sys.argv[1] if len(sys.argv) > 1 else "outputs/examples")
|
|
104
|
+
output.mkdir(parents=True, exist_ok=True)
|
|
105
|
+
for name, plot in examples().items():
|
|
106
|
+
plot.save(output / f"{name}.png", dpi=150, verbose=False)
|
|
107
|
+
plot.save(output / f"{name}.svg", verbose=False)
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=77", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "workforce-plotting"
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "Simple plotnine figures with Opportunity Insights styling."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.12"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
authors = [{ name = "Corey Gelb-Bicknell" }]
|
|
13
|
+
dependencies = ["plotnine>=0.15,<0.16", "pandas>=2.2", "oi-tools>=0.11.7"]
|
|
14
|
+
|
|
15
|
+
[project.optional-dependencies]
|
|
16
|
+
polars = ["polars>=1", "pyarrow>=14"]
|
|
17
|
+
dev = ["pytest>=8", "build", "twine", "polars>=1", "pyarrow>=14"]
|
|
18
|
+
|
|
19
|
+
[project.urls]
|
|
20
|
+
Repository = "https://github.com/coreygb1/workforce-plotting"
|
|
21
|
+
|
|
22
|
+
[tool.setuptools.packages.find]
|
|
23
|
+
where = ["src"]
|
|
24
|
+
include = ["workforce_plotting*"]
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
"""Plotnine figures and styling for workforce research."""
|
|
2
|
+
|
|
3
|
+
from .figures import (
|
|
4
|
+
WFColors, plot_bars, plot_intervals, plot_lines, plot_scatter,
|
|
5
|
+
plot_validation, theme_workforce,
|
|
6
|
+
)
|
|
7
|
+
|
|
8
|
+
__all__ = ["plot_lines", "plot_validation", "plot_bars", "plot_scatter",
|
|
9
|
+
"plot_intervals", "theme_workforce", "WFColors"]
|
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
"""Workforce plots, shared styling, and the Opportunity Insights color palette."""
|
|
2
|
+
|
|
3
|
+
import pandas as pd
|
|
4
|
+
import plotnine as pn
|
|
5
|
+
from matplotlib import font_manager
|
|
6
|
+
from oi_tools.figures import OIColors as WFColors, theme_oi
|
|
7
|
+
|
|
8
|
+
PALETTE = [WFColors.TEAL, WFColors.ORANGE, WFColors.NAVY, WFColors.PURPLE,
|
|
9
|
+
WFColors.BLUE, WFColors.RED, WFColors.DARK_GREEN]
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def theme_workforce(base_size=12, base_family=None, width=8, height=5):
|
|
13
|
+
"""White background, gray axes, no grid, and a bottom legend.
|
|
14
|
+
|
|
15
|
+
Dimensions are in inches. Install Lucida Sans to match the reference deck.
|
|
16
|
+
If unavailable, Matplotlib falls back to its default sans-serif font.
|
|
17
|
+
"""
|
|
18
|
+
if base_family is None:
|
|
19
|
+
base_family = ("Lucida Sans" if "Lucida Sans" in font_manager.get_font_names()
|
|
20
|
+
else "DejaVu Sans")
|
|
21
|
+
return theme_oi(base_size, base_family, width, height) + pn.theme(
|
|
22
|
+
axis_line=pn.element_line(color="#A0A0A0", size=0.54),
|
|
23
|
+
axis_ticks=pn.element_line(color="#A0A0A0", size=0.54),
|
|
24
|
+
axis_text=pn.element_text(color="black", size=base_size),
|
|
25
|
+
axis_title=pn.element_text(size=base_size, lineheight=1.8),
|
|
26
|
+
axis_title_y=pn.element_text(margin={"r": 16}),
|
|
27
|
+
axis_title_x=pn.element_text(margin={"t": 14}),
|
|
28
|
+
legend_text=pn.element_text(size=base_size),
|
|
29
|
+
legend_title=pn.element_blank(),
|
|
30
|
+
plot_margin_top=0.044,
|
|
31
|
+
plot_margin_right=0.015,
|
|
32
|
+
plot_margin_bottom=0.025,
|
|
33
|
+
panel_grid=pn.element_blank(),
|
|
34
|
+
axis_ticks_minor=pn.element_blank(),
|
|
35
|
+
legend_key=pn.element_blank(),
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _data(data):
|
|
40
|
+
"""Keep pandas inputs; convert Polars using its public conversion method."""
|
|
41
|
+
if isinstance(data, pd.DataFrame):
|
|
42
|
+
return data
|
|
43
|
+
return data.to_pandas()
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def _interval(lower, upper):
|
|
47
|
+
if (lower is None) != (upper is None):
|
|
48
|
+
raise ValueError("Supply both lower and upper interval columns, or neither.")
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def plot_lines(data, *, x, y, series=None, lower=None, upper=None, points=True):
|
|
52
|
+
"""Lines with optional points and supplied confidence intervals.
|
|
53
|
+
|
|
54
|
+
Use an ordered categorical series to control legend/color order. Add scales,
|
|
55
|
+
reference lines, facets, and labels using ordinary plotnine layers.
|
|
56
|
+
"""
|
|
57
|
+
_interval(lower, upper)
|
|
58
|
+
mapping = pn.aes(x=x, y=y, group=1)
|
|
59
|
+
if series is not None:
|
|
60
|
+
mapping.update(color=series, group=series)
|
|
61
|
+
p = pn.ggplot(_data(data), mapping)
|
|
62
|
+
if lower is not None:
|
|
63
|
+
p += pn.geom_errorbar(pn.aes(ymin=lower, ymax=upper), width=0.08,
|
|
64
|
+
size=0.4, **({} if series else {"color": WFColors.TEAL}))
|
|
65
|
+
p += pn.geom_line(size=0.914, **({} if series else {"color": WFColors.TEAL}))
|
|
66
|
+
if points:
|
|
67
|
+
p += pn.geom_point(size=3.45, stroke=0, **({} if series else {"color": WFColors.TEAL}))
|
|
68
|
+
if series is not None:
|
|
69
|
+
p += pn.scale_color_manual(values=PALETTE)
|
|
70
|
+
return p + theme_workforce()
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def plot_validation(data, *, x, y, x_lower=None, x_upper=None,
|
|
74
|
+
y_lower=None, y_upper=None, gaps=None, limits=None):
|
|
75
|
+
"""Compare estimates against a 45-degree identity line.
|
|
76
|
+
|
|
77
|
+
x/y confidence bounds are optional. ``gaps`` may be 'vertical' (from y to x)
|
|
78
|
+
or 'horizontal' (from x to y). Both axes use the same numeric limits.
|
|
79
|
+
Default limits include all supplied estimates and interval endpoints.
|
|
80
|
+
"""
|
|
81
|
+
_interval(x_lower, x_upper)
|
|
82
|
+
_interval(y_lower, y_upper)
|
|
83
|
+
if gaps not in (None, "vertical", "horizontal"):
|
|
84
|
+
raise ValueError("gaps must be None, 'vertical', or 'horizontal'.")
|
|
85
|
+
data = _data(data)
|
|
86
|
+
if limits is None:
|
|
87
|
+
columns = [v for v in (x, y, x_lower, x_upper, y_lower, y_upper)
|
|
88
|
+
if v is not None]
|
|
89
|
+
values = pd.concat([data[v] for v in columns]).dropna()
|
|
90
|
+
if values.empty:
|
|
91
|
+
raise ValueError("Supply at least one nonmissing estimate or explicit limits.")
|
|
92
|
+
limits = (values.min(), values.max())
|
|
93
|
+
p = pn.ggplot(data, pn.aes(x=x, y=y)) + pn.geom_abline(
|
|
94
|
+
intercept=0, slope=1, color="black", size=0.6)
|
|
95
|
+
if gaps == "vertical":
|
|
96
|
+
p += pn.geom_segment(pn.aes(xend=x, yend=x), color=WFColors.ORANGE, size=0.5, alpha=0.8)
|
|
97
|
+
elif gaps == "horizontal":
|
|
98
|
+
p += pn.geom_segment(pn.aes(xend=y, yend=y), color=WFColors.ORANGE, size=0.5, alpha=0.8)
|
|
99
|
+
if x_lower is not None:
|
|
100
|
+
p += pn.geom_errorbarh(pn.aes(xmin=x_lower, xmax=x_upper), height=0,
|
|
101
|
+
color=WFColors.TEAL, size=0.5, alpha=0.5)
|
|
102
|
+
if y_lower is not None:
|
|
103
|
+
p += pn.geom_errorbar(pn.aes(ymin=y_lower, ymax=y_upper), width=0,
|
|
104
|
+
color=WFColors.TEAL, size=0.5, alpha=0.5)
|
|
105
|
+
return (p + pn.geom_point(color=WFColors.TEAL, size=3.45, stroke=0)
|
|
106
|
+
+ pn.coord_cartesian(xlim=limits, ylim=limits) + theme_workforce())
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def plot_bars(data, *, x, y, series=None):
|
|
110
|
+
"""Bars for already-calculated values; optional series are side by side.
|
|
111
|
+
|
|
112
|
+
Set x to an ordered categorical column, or add scale_x_discrete(limits=...),
|
|
113
|
+
to choose the category order. Add coord_flip() for horizontal bars.
|
|
114
|
+
"""
|
|
115
|
+
mapping = pn.aes(x=x, y=y, group=1)
|
|
116
|
+
if series is not None:
|
|
117
|
+
mapping.update(fill=series, group=series)
|
|
118
|
+
p = pn.ggplot(_data(data), mapping) + pn.geom_col(
|
|
119
|
+
width=0.7, position="dodge", **({} if series else {"fill": WFColors.TEAL}))
|
|
120
|
+
if series is not None:
|
|
121
|
+
p += pn.scale_fill_manual(values=PALETTE)
|
|
122
|
+
return p + theme_workforce()
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def plot_scatter(data, *, x, y, series=None, fitted=None):
|
|
126
|
+
"""Points, optionally with a separate dataframe of precomputed fitted values.
|
|
127
|
+
|
|
128
|
+
``fitted`` uses the same x, y, and (when specified) series column names.
|
|
129
|
+
The function never fits or smooths a curve.
|
|
130
|
+
"""
|
|
131
|
+
mapping = pn.aes(x=x, y=y, group=1)
|
|
132
|
+
if series is not None:
|
|
133
|
+
mapping.update(color=series, group=series)
|
|
134
|
+
p = pn.ggplot(_data(data), mapping)
|
|
135
|
+
if fitted is not None:
|
|
136
|
+
p += pn.geom_line(data=_data(fitted), size=0.914,
|
|
137
|
+
**({} if series else {"color": WFColors.NAVY}))
|
|
138
|
+
p += pn.geom_point(size=3.45, stroke=0, **({} if series else {"color": WFColors.TEAL}))
|
|
139
|
+
if series is not None:
|
|
140
|
+
p += pn.scale_color_manual(values=PALETTE)
|
|
141
|
+
return p + theme_workforce()
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
def plot_intervals(data, *, x, y, lower, upper, series=None):
|
|
145
|
+
"""Points with supplied interval bounds, dodged when series are supplied.
|
|
146
|
+
|
|
147
|
+
x is the category and y the estimate. Add coord_flip() for a horizontal
|
|
148
|
+
coefficient plot. Interval bounds are used as supplied, without adjustment.
|
|
149
|
+
"""
|
|
150
|
+
mapping = pn.aes(x=x, y=y, ymin=lower, ymax=upper)
|
|
151
|
+
if series is not None:
|
|
152
|
+
mapping.update(color=series, group=series)
|
|
153
|
+
p = pn.ggplot(_data(data), mapping) + pn.geom_pointrange(
|
|
154
|
+
position=pn.position_dodge(width=0.5), size=0.5,
|
|
155
|
+
**({} if series else {"color": WFColors.TEAL}))
|
|
156
|
+
if series is not None:
|
|
157
|
+
p += pn.scale_color_manual(values=PALETTE)
|
|
158
|
+
return p + theme_workforce()
|
|
@@ -0,0 +1,135 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: workforce-plotting
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Simple plotnine figures with Opportunity Insights styling.
|
|
5
|
+
Author: Corey Gelb-Bicknell
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Repository, https://github.com/coreygb1/workforce-plotting
|
|
8
|
+
Requires-Python: >=3.12
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Dist: plotnine<0.16,>=0.15
|
|
12
|
+
Requires-Dist: pandas>=2.2
|
|
13
|
+
Requires-Dist: oi-tools>=0.11.7
|
|
14
|
+
Provides-Extra: polars
|
|
15
|
+
Requires-Dist: polars>=1; extra == "polars"
|
|
16
|
+
Requires-Dist: pyarrow>=14; extra == "polars"
|
|
17
|
+
Provides-Extra: dev
|
|
18
|
+
Requires-Dist: pytest>=8; extra == "dev"
|
|
19
|
+
Requires-Dist: build; extra == "dev"
|
|
20
|
+
Requires-Dist: twine; extra == "dev"
|
|
21
|
+
Requires-Dist: polars>=1; extra == "dev"
|
|
22
|
+
Requires-Dist: pyarrow>=14; extra == "dev"
|
|
23
|
+
Dynamic: license-file
|
|
24
|
+
|
|
25
|
+
# workforce-plotting
|
|
26
|
+
|
|
27
|
+
Five small plotnine functions with a shared Opportunity Insights style. Supply
|
|
28
|
+
prepared estimates and confidence bounds; the functions only draw them. Every
|
|
29
|
+
function returns a `plotnine.ggplot` that you can edit with ordinary layers.
|
|
30
|
+
|
|
31
|
+
## Install
|
|
32
|
+
|
|
33
|
+
Python 3.12 or newer is required (matching `oi-tools`). For a Polars workflow:
|
|
34
|
+
|
|
35
|
+
```bash
|
|
36
|
+
uv add 'workforce-plotting[polars]'
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
Or use pip:
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
pip install 'workforce-plotting[polars]'
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
For local development from this checkout:
|
|
46
|
+
|
|
47
|
+
```bash
|
|
48
|
+
pip install -e '.[dev]'
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
For Polars input, install the `polars` extra, which also supplies the Arrow
|
|
52
|
+
conversion dependency. Pandas input needs no extra.
|
|
53
|
+
|
|
54
|
+
## Functions
|
|
55
|
+
|
|
56
|
+
| Function | Use |
|
|
57
|
+
|---|---|
|
|
58
|
+
| `plot_lines(data, x=..., y=..., series=...)` | Event studies, trajectories, cohort trends, model complexity |
|
|
59
|
+
| `plot_validation(data, x=..., y=...)` | RCT versus digital-twin estimates with an identity line |
|
|
60
|
+
| `plot_bars(data, x=..., y=..., series=...)` | Already-calculated values by category |
|
|
61
|
+
| `plot_scatter(data, x=..., y=..., fitted=...)` | Points with an optional precomputed fitted line |
|
|
62
|
+
| `plot_intervals(data, x=..., y=..., lower=..., upper=...)` | Estimates and supplied confidence intervals |
|
|
63
|
+
| `theme_workforce()` | Shared white background, gray axes, sans-serif type, no grid |
|
|
64
|
+
|
|
65
|
+
`plot_lines` accepts `lower` and `upper` columns and `points=False`.
|
|
66
|
+
`plot_validation` accepts `x_lower`, `x_upper`, `y_lower`, `y_upper`,
|
|
67
|
+
`gaps="horizontal"` or `"vertical"`, and a common `(min, max)` pair of `limits`.
|
|
68
|
+
The identity line has slope one in data units; the panel follows the figure's
|
|
69
|
+
aspect ratio. Other functions accept `series` for discrete groups. Use ordered
|
|
70
|
+
categorical columns to control color/legend order. Colors are teal, orange, navy,
|
|
71
|
+
purple, blue, red, and dark green from the public `oi_tools.figures.OIColors` API.
|
|
72
|
+
For a different mapping, add a plotnine manual color or fill scale.
|
|
73
|
+
|
|
74
|
+
```python
|
|
75
|
+
import plotnine as pn
|
|
76
|
+
from workforce_plotting import plot_lines
|
|
77
|
+
|
|
78
|
+
p = plot_lines(df, x="year", y="earnings", series="group")
|
|
79
|
+
p += pn.labs(x="Years since enrollment", y="Mean annual earnings")
|
|
80
|
+
p.save("earnings.svg")
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
Dataframes should be tidy: one row per plotted point, with separate columns for
|
|
84
|
+
x, y, and an optional series. Lines connect points in x order within each series.
|
|
85
|
+
Include a row with a missing y where a line should break; simply omitting a row
|
|
86
|
+
connects the remaining observations. Fitted lines use the same column names as
|
|
87
|
+
the scatter data. No smoothing, RMSE calculation, confidence-interval estimation,
|
|
88
|
+
or automatic saving happens inside the plotting functions.
|
|
89
|
+
|
|
90
|
+
Add `pn.coord_flip()` for horizontal bars/intervals; use ordinary plotnine scales
|
|
91
|
+
for dollars, percentages, limits, or log axes. Fonts, dimensions, and text size
|
|
92
|
+
use Lucida Sans when installed and DejaVu Sans otherwise. They can be overridden with `theme_workforce(base_family=..., base_size=..., width=..., height=...)`.
|
|
93
|
+
Titles, logos, notes, and paper-specific annotations belong in the example/caller.
|
|
94
|
+
|
|
95
|
+
## Styling your own plots
|
|
96
|
+
|
|
97
|
+
Use the theme and colors without using the plot constructors. `WFColors` is an
|
|
98
|
+
alias of `oi_tools.figures.OIColors`, and `theme_workforce` builds on `theme_oi`.
|
|
99
|
+
The `oi-tools` dependency remains; you only need workforce imports in your code.
|
|
100
|
+
All plotting functions and styling live together in `workforce_plotting/figures.py`.
|
|
101
|
+
|
|
102
|
+
```python
|
|
103
|
+
import polars as pl
|
|
104
|
+
import plotnine as pn
|
|
105
|
+
from workforce_plotting import theme_workforce, WFColors
|
|
106
|
+
|
|
107
|
+
df = pl.DataFrame({"year": [2020, 2021, 2022], "earnings": [20000, 23000, 25000]})
|
|
108
|
+
p = (
|
|
109
|
+
pn.ggplot(df.to_pandas(), pn.aes(x="year", y="earnings"))
|
|
110
|
+
+ pn.geom_line(color=WFColors.TEAL)
|
|
111
|
+
+ pn.geom_point(color=WFColors.TEAL)
|
|
112
|
+
+ theme_workforce()
|
|
113
|
+
+ pn.labs(x="Year", y="Annual earnings")
|
|
114
|
+
)
|
|
115
|
+
p.save("earnings.svg")
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
Plot constructors accept Polars directly. When building a plotnine plot yourself,
|
|
119
|
+
the explicit conversion above works with the supported plotnine version.
|
|
120
|
+
|
|
121
|
+
## Editable examples
|
|
122
|
+
|
|
123
|
+
`examples/paper_examples.py` contains synthetic recipes for the paper/deck plot
|
|
124
|
+
families. These are invented values, not research results.
|
|
125
|
+
|
|
126
|
+
```bash
|
|
127
|
+
python examples/paper_examples.py outputs/examples
|
|
128
|
+
pytest
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
## Changes from 0.1
|
|
132
|
+
|
|
133
|
+
This is a fresh API. The old Matplotlib functions, plotting dataclasses, curve
|
|
134
|
+
fitting machinery, and `aipw_analysis` compatibility package have been removed.
|
|
135
|
+
Plot creation and saving are now separate. The public API consists of the five plot functions, `theme_workforce`, and `WFColors`.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
MANIFEST.in
|
|
3
|
+
README.md
|
|
4
|
+
pyproject.toml
|
|
5
|
+
examples/paper_examples.py
|
|
6
|
+
src/workforce_plotting/__init__.py
|
|
7
|
+
src/workforce_plotting/figures.py
|
|
8
|
+
src/workforce_plotting.egg-info/PKG-INFO
|
|
9
|
+
src/workforce_plotting.egg-info/SOURCES.txt
|
|
10
|
+
src/workforce_plotting.egg-info/dependency_links.txt
|
|
11
|
+
src/workforce_plotting.egg-info/requires.txt
|
|
12
|
+
src/workforce_plotting.egg-info/top_level.txt
|
|
13
|
+
tests/test_plots.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
workforce_plotting
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
import pandas as pd
|
|
2
|
+
import plotnine as pn
|
|
3
|
+
import pytest
|
|
4
|
+
from oi_tools.figures import OIColors
|
|
5
|
+
|
|
6
|
+
from workforce_plotting import plot_bars, plot_intervals, plot_lines, plot_scatter, plot_validation
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
@pytest.fixture
|
|
10
|
+
def data():
|
|
11
|
+
return pd.DataFrame({"x": [2, 0, 1, 2, 0, 1], "y": [4., 1, 2, 8, 3, 5],
|
|
12
|
+
"group": ["A"] * 3 + ["B"] * 3,
|
|
13
|
+
"lo": [3., 0, 1, 7, 2, 4], "hi": [5., 2, 3, 9, 4, 6]})
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def test_all_plots_render_without_changing_values(data):
|
|
17
|
+
before = data.copy(deep=True)
|
|
18
|
+
plots = [plot_lines(data, x="x", y="y", series="group", lower="lo", upper="hi"),
|
|
19
|
+
plot_bars(data, x="x", y="y", series="group"),
|
|
20
|
+
plot_scatter(data, x="x", y="y", series="group", fitted=data),
|
|
21
|
+
plot_intervals(data, x="x", y="y", lower="lo", upper="hi", series="group"),
|
|
22
|
+
plot_validation(data, x="x", y="y", y_lower="lo", y_upper="hi")]
|
|
23
|
+
import matplotlib.pyplot as plt
|
|
24
|
+
for p in plots:
|
|
25
|
+
assert isinstance(p, pn.ggplot)
|
|
26
|
+
fig = p.draw()
|
|
27
|
+
plt.close(fig)
|
|
28
|
+
pd.testing.assert_frame_equal(before, data)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def test_lines_sort_and_group_and_keep_missing_gaps(data):
|
|
32
|
+
data.loc[2, "y"] = float("nan")
|
|
33
|
+
p = plot_lines(data, x="x", y="y", series="group")
|
|
34
|
+
p._build()
|
|
35
|
+
lines = p._build_objs.layers[0].data
|
|
36
|
+
assert lines.groupby("group")["x"].apply(list).tolist() == [[0, 1, 2], [0, 1, 2]]
|
|
37
|
+
assert lines["y"].isna().sum() == 1
|
|
38
|
+
assert lines.groupby("group")["color"].first().tolist() == [OIColors.TEAL, OIColors.ORANGE]
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def test_validation_gaps_and_bounds(data):
|
|
42
|
+
for direction in ["horizontal", "vertical"]:
|
|
43
|
+
p = plot_validation(data, x="x", y="y", gaps=direction,
|
|
44
|
+
y_lower="lo", y_upper="hi")
|
|
45
|
+
p._build()
|
|
46
|
+
segments = p._build_objs.layers[1].data
|
|
47
|
+
if direction == "horizontal":
|
|
48
|
+
assert list(segments.xend) == list(data.y)
|
|
49
|
+
assert list(segments.yend) == list(data.y)
|
|
50
|
+
else:
|
|
51
|
+
assert list(segments.xend) == list(data.x)
|
|
52
|
+
assert list(segments.yend) == list(data.x)
|
|
53
|
+
intervals = p._build_objs.layers[2].data
|
|
54
|
+
assert list(intervals.ymin) == list(data.lo)
|
|
55
|
+
assert list(intervals.ymax) == list(data.hi)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def test_polars_conversion(data):
|
|
59
|
+
pl = pytest.importorskip("polars")
|
|
60
|
+
p = plot_lines(pl.from_pandas(data), x="x", y="y")
|
|
61
|
+
pd.testing.assert_frame_equal(p.data, data)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def test_invalid_partial_interval(data):
|
|
65
|
+
with pytest.raises(ValueError, match="both"):
|
|
66
|
+
plot_lines(data, x="x", y="y", lower="lo")
|
|
67
|
+
with pytest.raises(ValueError, match="gaps"):
|
|
68
|
+
plot_validation(data, x="x", y="y", gaps="unknown")
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def test_categorical_x_connects_single_series():
|
|
72
|
+
data = pd.DataFrame({"x": pd.Categorical(["late", "early", "middle"],
|
|
73
|
+
categories=["early", "middle", "late"], ordered=True), "y": [3, 1, 2]})
|
|
74
|
+
p = plot_lines(data, x="x", y="y")
|
|
75
|
+
p._build()
|
|
76
|
+
line = p._build_objs.layers[0].data
|
|
77
|
+
assert line.group.nunique() == 1
|
|
78
|
+
assert line.y.tolist() == [1, 2, 3]
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def test_validation_preserves_bounds_outside_view(data):
|
|
82
|
+
p = plot_validation(data, x="x", y="y", x_lower="lo", x_upper="hi",
|
|
83
|
+
y_lower="lo", y_upper="hi", limits=(1, 3))
|
|
84
|
+
p._build()
|
|
85
|
+
horizontal, vertical = p._build_objs.layers[1:3]
|
|
86
|
+
assert horizontal.data.xmin.tolist() == data.lo.tolist()
|
|
87
|
+
assert horizontal.data.xmax.tolist() == data.hi.tolist()
|
|
88
|
+
assert vertical.data.ymin.tolist() == data.lo.tolist()
|
|
89
|
+
assert vertical.data.ymax.tolist() == data.hi.tolist()
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def test_standalone_theme_and_colors():
|
|
93
|
+
import matplotlib.pyplot as plt
|
|
94
|
+
from workforce_plotting import WFColors, theme_workforce
|
|
95
|
+
from workforce_plotting.figures import theme_workforce as module_theme
|
|
96
|
+
|
|
97
|
+
assert WFColors is OIColors
|
|
98
|
+
assert module_theme is theme_workforce
|
|
99
|
+
p = (pn.ggplot(pd.DataFrame({"x": [1, 2], "y": [3, 4]}), pn.aes("x", "y"))
|
|
100
|
+
+ pn.geom_line(color=WFColors.TEAL) + theme_workforce())
|
|
101
|
+
plt.close(p.draw())
|