vse_sim.dataframe

  1"""Pandas-first helpers for VSE simulation results."""
  2
  3from __future__ import annotations
  4
  5from collections.abc import Iterable
  6from dataclasses import dataclass
  7from inspect import signature
  8from pathlib import Path
  9
 10import pandas as pd
 11
 12DEFAULT_GROUP_BY = ("method", "chooser")
 13
 14
 15def _group_columns(group_by) -> list[str]:
 16    return [group_by] if isinstance(group_by, str) else list(group_by)
 17
 18
 19def _unique_columns(*groups) -> tuple[str, ...]:
 20    columns = []
 21    for group in groups:
 22        for column in _group_columns(group):
 23            if column not in columns:
 24                columns.append(column)
 25    return tuple(columns)
 26
 27
 28def _voter_metadata(voter) -> dict:
 29    metadata = {}
 30    for attribute in ("cluster", "personality"):
 31        if hasattr(voter, attribute):
 32            metadata[attribute] = getattr(voter, attribute)
 33    if hasattr(voter, "dims"):
 34        for dimension, value in enumerate(voter.dims):
 35            metadata[f"dimension_{dimension}"] = value
 36    return metadata
 37
 38
 39def _call_dataframe_method(method, copy=True, **kwargs) -> pd.DataFrame:
 40    if "copy" in signature(method).parameters:
 41        return method(copy=copy, **kwargs)
 42    frame = method(**kwargs)
 43    return frame.copy() if copy else frame
 44
 45
 46def to_dataframe(data, copy=True, **kwargs) -> pd.DataFrame:
 47    """Convert VSE objects, result rows, or records to a pandas DataFrame."""
 48    if isinstance(data, VseResults):
 49        return data.to_dataframe(copy=copy)
 50    if isinstance(data, pd.DataFrame):
 51        return data.copy() if copy else data
 52    dataframe_method = getattr(data, "to_dataframe", None)
 53    if callable(dataframe_method):
 54        return _call_dataframe_method(dataframe_method, copy=copy, **kwargs)
 55    return pd.DataFrame(data, **kwargs)
 56
 57
 58def rows_to_dataframe(rows: Iterable[dict] | pd.DataFrame | VseResults, copy=True) -> pd.DataFrame:
 59    """Convert simulation rows, a DataFrame, or ``VseResults`` to a DataFrame."""
 60    return to_dataframe(rows, copy=copy)
 61
 62
 63def voter_to_dataframe(
 64    voter,
 65    voter_id=None,
 66    voter_column="voter",
 67    candidate_column="candidate",
 68    value_column="utility",
 69) -> pd.DataFrame:
 70    """Return one voter's candidate utilities as a tidy DataFrame."""
 71    metadata = _voter_metadata(voter)
 72    rows = []
 73    for candidate, value in enumerate(voter):
 74        row = {candidate_column: candidate, value_column: value, **metadata}
 75        if voter_id is not None:
 76            row[voter_column] = voter_id
 77        rows.append(row)
 78    return pd.DataFrame(rows)
 79
 80
 81def voters_to_dataframe(
 82    voters,
 83    wide=False,
 84    voter_column="voter",
 85    candidate_column="candidate",
 86    value_column="utility",
 87    candidate_prefix="candidate_",
 88) -> pd.DataFrame:
 89    """Return voter utilities as a tidy or wide DataFrame."""
 90    if wide:
 91        rows = []
 92        for voter_id, voter in enumerate(voters):
 93            row = {
 94                voter_column: voter_id,
 95                **{
 96                    f"{candidate_prefix}{candidate}": value for candidate, value in enumerate(voter)
 97                },
 98                **_voter_metadata(voter),
 99            }
100            rows.append(row)
101        return pd.DataFrame(rows)
102
103    rows = []
104    for voter_id, voter in enumerate(voters):
105        rows.extend(
106            voter_to_dataframe(
107                voter,
108                voter_id=voter_id,
109                voter_column=voter_column,
110                candidate_column=candidate_column,
111                value_column=value_column,
112            ).to_dict("records")
113        )
114    return pd.DataFrame(rows)
115
116
117def ballots_to_dataframe(
118    ballots,
119    wide=False,
120    method=None,
121    voter_column="voter",
122    candidate_column="candidate",
123    value_column="ballot",
124    candidate_prefix="candidate_",
125) -> pd.DataFrame:
126    """Return ballots as a tidy or wide DataFrame."""
127    if wide:
128        rows = []
129        for voter_id, ballot in enumerate(ballots):
130            row = {
131                voter_column: voter_id,
132                **{
133                    f"{candidate_prefix}{candidate}": value
134                    for candidate, value in enumerate(ballot)
135                },
136            }
137            if method is not None:
138                row["method"] = str(method)
139            rows.append(row)
140        return pd.DataFrame(rows)
141
142    rows = []
143    for voter_id, ballot in enumerate(ballots):
144        for candidate, value in enumerate(ballot):
145            row = {
146                voter_column: voter_id,
147                candidate_column: candidate,
148                value_column: value,
149            }
150            if method is not None:
151                row["method"] = str(method)
152            rows.append(row)
153    return pd.DataFrame(rows)
154
155
156def ballots_from_dataframe(
157    ballots,
158    voter_column="voter",
159    candidate_column="candidate",
160    value_column="ballot",
161    candidate_prefix="candidate_",
162):
163    """Convert tidy or wide ballot DataFrames back to method-ready ballots."""
164    if not isinstance(ballots, pd.DataFrame):
165        return ballots if type(ballots) is list else list(ballots)
166
167    if {voter_column, candidate_column, value_column} <= set(ballots.columns):
168        return (
169            ballots.pivot(index=voter_column, columns=candidate_column, values=value_column)
170            .sort_index()
171            .sort_index(axis=1)
172            .to_numpy()
173            .tolist()
174        )
175
176    candidate_columns = [
177        column for column in ballots.columns if str(column).startswith(candidate_prefix)
178    ]
179    if candidate_columns:
180        candidate_columns = sorted(
181            candidate_columns, key=lambda column: int(str(column).split("_")[-1])
182        )
183        return ballots[candidate_columns].to_numpy().tolist()
184
185    return ballots.to_numpy().tolist()
186
187
188def scores_to_dataframe(
189    scores,
190    method=None,
191    candidate_column="candidate",
192    value_column="score",
193) -> pd.DataFrame:
194    """Return candidate-level method scores as a DataFrame."""
195    rows = [
196        {
197            candidate_column: candidate,
198            value_column: score,
199        }
200        for candidate, score in enumerate(scores)
201    ]
202    frame = pd.DataFrame(rows)
203    if method is not None:
204        frame.insert(0, "method", str(method))
205    return frame
206
207
208def summarize_vse(
209    rows: Iterable[dict] | pd.DataFrame | VseResults,
210    group_by=DEFAULT_GROUP_BY,
211    sort_by="mean_vse",
212    ascending=False,
213) -> pd.DataFrame:
214    """Summarize VSE scores by method, chooser, or another grouping."""
215    return VseResults(rows_to_dataframe(rows, copy=False)).summarize(
216        group_by=group_by,
217        sort_by=sort_by,
218        ascending=ascending,
219    )
220
221
222def read_results_csv(path) -> "VseResults":
223    """Load a VSE result CSV written by ``CsvBatch.saveFile``."""
224    return VseResults.from_csv(path)
225
226
227@dataclass(frozen=True)
228class VseResults:
229    """Pandas-backed simulation result set.
230
231    ``frame`` is the canonical tabular representation. Convenience methods
232    return DataFrames or matplotlib axes so notebook workflows can keep chaining.
233    """
234
235    frame: pd.DataFrame
236
237    @classmethod
238    def from_rows(cls, rows: Iterable[dict] | pd.DataFrame | VseResults) -> "VseResults":
239        return cls(to_dataframe(rows))
240
241    @classmethod
242    def from_csv(cls, path) -> "VseResults":
243        return cls(pd.read_csv(Path(path), comment="#"))
244
245    @classmethod
246    def concat(cls, results: Iterable[VseResults | pd.DataFrame | Iterable[dict]]) -> "VseResults":
247        frames = [rows_to_dataframe(result) for result in results]
248        return cls(pd.concat(frames, ignore_index=True))
249
250    def __len__(self) -> int:
251        return len(self.frame)
252
253    @property
254    def dataframe(self) -> pd.DataFrame:
255        """Return the backing DataFrame for fluent notebook work."""
256        return self.frame
257
258    @property
259    def df(self) -> pd.DataFrame:
260        """Alias for ``dataframe``."""
261        return self.frame
262
263    def to_dataframe(self, copy=True) -> pd.DataFrame:
264        return self.frame.copy() if copy else self.frame
265
266    def to_csv(self, path, index=False, **kwargs):
267        """Write the result DataFrame to CSV and return the path."""
268        self.frame.to_csv(path, index=index, **kwargs)
269        return path
270
271    def summarize(
272        self,
273        group_by=DEFAULT_GROUP_BY,
274        sort_by="mean_vse",
275        ascending=False,
276    ) -> pd.DataFrame:
277        """Return aggregate VSE metrics grouped by one or more columns."""
278        group_columns = _group_columns(group_by)
279        summary = (
280            self.frame.groupby(group_columns, as_index=False)
281            .agg(
282                rows=("vse", "size"),
283                elections=("eid", "nunique"),
284                mean_vse=("vse", "mean"),
285                median_vse=("vse", "median"),
286                min_vse=("vse", "min"),
287                max_vse=("vse", "max"),
288                std_vse=("vse", "std"),
289            )
290            .fillna({"std_vse": 0})
291        )
292        return summary.sort_values(
293            [sort_by, *group_columns],
294            ascending=[ascending, *[True] * len(group_columns)],
295        ).reset_index(drop=True)
296
297    def leaderboard(self, n=10, group_by="method", by="mean_vse") -> pd.DataFrame:
298        """Return the top groups by a summary metric."""
299        return self.summarize(group_by=group_by, sort_by=by).head(n)
300
301    def pivot(
302        self,
303        index="method",
304        columns="chooser",
305        values="mean_vse",
306        group_by=None,
307    ) -> pd.DataFrame:
308        """Return a comparison matrix from summarized result data."""
309        index_columns = _group_columns(index)
310        column_columns = _group_columns(columns)
311        if group_by is None:
312            group_by = _unique_columns(index, columns)
313        summary = self.summarize(group_by=group_by).copy()
314        pivot_columns = []
315        for column in column_columns:
316            if column in index_columns:
317                column_alias = f"__vse_sim_pivot_{column}"
318                summary[column_alias] = summary[column]
319                pivot_columns.append(column_alias)
320            else:
321                pivot_columns.append(column)
322        pivoted = summary.pivot(
323            index=index,
324            columns=pivot_columns[0] if isinstance(columns, str) else pivot_columns,
325            values=values,
326        )
327        pivoted.columns = pivoted.columns.set_names(column_columns)
328        return pivoted
329
330    def report(self, group_by=DEFAULT_GROUP_BY) -> dict[str, pd.DataFrame]:
331        """Build common report tables from a result set."""
332        tables = {
333            "results": self.to_dataframe(),
334            "summary": self.summarize(group_by=group_by),
335            "method_summary": self.summarize(group_by="method"),
336        }
337        if "chooser" in self.frame:
338            tables["chooser_summary"] = self.summarize(group_by="chooser")
339            tables["method_by_chooser"] = self.pivot()
340        return tables
341
342    def plot_vse(
343        self,
344        group_by="method",
345        value="mean_vse",
346        kind="bar",
347        ax=None,
348        title=None,
349        **kwargs,
350    ):
351        """Plot summarized VSE scores and return the matplotlib axes."""
352        summary = self.summarize(group_by=group_by)
353        group_columns = _group_columns(group_by)
354        labels = summary[group_columns].astype(str).agg(" | ".join, axis=1)
355        plot_frame = summary.assign(label=labels).set_index("label")
356        axes = plot_frame[value].plot(kind=kind, ax=ax, **kwargs)
357        axes.set_xlabel("VSE" if kind == "barh" else "")
358        axes.set_ylabel("" if kind == "barh" else "VSE")
359        axes.set_title(title or f"{value} by {' / '.join(group_columns)}")
360        return axes
361
362
363__all__ = [
364    "DEFAULT_GROUP_BY",
365    "VseResults",
366    "ballots_from_dataframe",
367    "ballots_to_dataframe",
368    "read_results_csv",
369    "rows_to_dataframe",
370    "scores_to_dataframe",
371    "summarize_vse",
372    "to_dataframe",
373    "voter_to_dataframe",
374    "voters_to_dataframe",
375]