add aggregate output: count/total/mean/median gross pay per category

uc_salary_pipeline.py

@@ -32,9 +32,14 @@ expects, and writes the result.
 
 OUTPUT
 ------
-One row per employee per year:
+Per-employee CSV (--output), one row per employee per year:
   year, campus, title, category, category_label,
   gross_pay, regular_pay, overtime_pay, other_pay, first_name, last_name
+
+Optionally, an aggregate CSV (--aggregate-output) with one row per
+category (broken down further by --group-by year/campus):
+  [year] [campus] category, category_label,
+  employee_count, total_gross_pay, mean_gross_pay, median_gross_pay
 """
 
 import argparse
@@ -42,6 +47,7 @@ import csv
 import json
 import logging
 import re
+import statistics
 import sys
 from pathlib import Path
 
@@ -98,12 +104,21 @@ def parse_uc_wages_record(record: dict) -> dict | None:
 
 
 def process_file(path: Path, campus: str, year: int, writer: csv.writer,
-                 senior_threshold: float, support_threshold: float) -> int:
+                 senior_threshold: float, support_threshold: float,
+                 agg: dict | None = None, group_dims: tuple = ()) -> int:
     """Classify every record in one UC-wages JSON file, writing a CSV row
-    per employee. Returns the number of rows written."""
+    per employee. Returns the number of rows written.
+
+    If `agg` is given, also accumulates each record's gross_pay under an
+    (group_dims..., category) key for later aggregate reporting.
+    """
     with open(path, encoding="utf-8") as fh:
         data = json.load(fh)
 
+    group_prefix = tuple(
+        year if dim == "year" else campus for dim in group_dims
+    )
+
     count = 0
     for record in data.get("records", []):
         parsed = parse_uc_wages_record(record)
@@ -123,9 +138,38 @@ def process_file(path: Path, campus: str, year: int, writer: csv.writer,
             parsed["first_name"], parsed["last_name"],
         ])
         count += 1
+        if agg is not None:
+            agg.setdefault(group_prefix + (category,), []).append(parsed["gross_pay"])
     return count
 
 
+def write_aggregate_csv(agg: dict, group_dims: tuple, outpath: Path) -> None:
+    """Write one row per (group_dims..., category) with employee count,
+    total/mean/median gross pay."""
+    header = list(group_dims) + [
+        "category", "category_label",
+        "employee_count", "total_gross_pay", "mean_gross_pay", "median_gross_pay",
+    ]
+
+    def sort_key(key):
+        *prefix, category = key
+        return (*prefix, classifier.CATEGORIES.index(category))
+
+    with open(outpath, "w", newline="", encoding="utf-8") as fh:
+        writer = csv.writer(fh)
+        writer.writerow(header)
+        for key in sorted(agg, key=sort_key):
+            *prefix, category = key
+            values = agg[key]
+            total = sum(values)
+            writer.writerow([
+                *prefix, category, classifier.CATEGORY_LABELS[category],
+                len(values), round(total, 2), round(total / len(values), 2),
+                round(statistics.median(values), 2),
+            ])
+    log.info("Aggregate CSV -> %s", outpath)
+
+
 # ---------------------------------------------------------------------------
 # CLI
 # ---------------------------------------------------------------------------
@@ -151,6 +195,17 @@ def main():
                              "(default: uc_wages_data/).")
     parser.add_argument("--output", default="uc_salary_classified.csv",
                         help="Output CSV path (default: uc_salary_classified.csv).")
+    parser.add_argument("--aggregate-output",
+                        help="If given, also write an aggregate CSV here with "
+                             "employee count, total/mean/median gross pay per "
+                             "category (see --group-by).")
+    parser.add_argument("--group-by", nargs="*", choices=["year", "campus"],
+                        default=["year"],
+                        help="Additional breakdown dimensions for "
+                             "--aggregate-output, alongside category "
+                             "(default: year). Pass both for a year x campus "
+                             "x category breakdown, or nothing to collapse "
+                             "everything to one row per category.")
     parser.add_argument("--skip-download", action="store_true",
                         help="Classify whatever is already cached in "
                              "--data-dir without hitting the network.")
@@ -171,6 +226,9 @@ def main():
     else:
         log.info("--skip-download set, classifying cached data only.")
 
+    group_dims = tuple(dim for dim in ("year", "campus") if dim in args.group_by)
+    agg = {} if args.aggregate_output else None
+
     output_path = Path(args.output)
     total_rows = 0
     with open(output_path, "w", newline="", encoding="utf-8") as out_fh:
@@ -184,7 +242,8 @@ def main():
                     log.info("No data for %s %s, skipping.", campus, year)
                     continue
                 count = process_file(path, campus, year, writer,
-                                    args.senior_threshold, args.support_threshold)
+                                    args.senior_threshold, args.support_threshold,
+                                    agg=agg, group_dims=group_dims)
                 total_rows += count
                 log.info("%s %s: classified %d records", campus, year, count)
 
@@ -194,6 +253,9 @@ def main():
 
     log.info("Total rows written: %d -> %s", total_rows, output_path)
 
+    if agg is not None:
+        write_aggregate_csv(agg, group_dims, Path(args.aggregate_output))
+
 
 if __name__ == "__main__":
     main()