Spaces:
Running
Running
| # Copyright (C) 2026 Hengzhe Zhao. All rights reserved. | |
| # Licensed under dual license: AGPL-3.0 (open-source) or commercial. See LICENSE. | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from .config import ModelSpec | |
| def get_best_device() -> torch.device: | |
| """Auto-detect the best available compute device.""" | |
| if torch.cuda.is_available(): | |
| return torch.device("cuda") | |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| return torch.device("mps") | |
| return torch.device("cpu") | |
| def get_device_info() -> str: | |
| """Return a human-readable string describing the active compute device.""" | |
| if torch.cuda.is_available(): | |
| name = torch.cuda.get_device_name(0) | |
| return f"{name} (CUDA)" | |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| import platform | |
| chip = platform.processor() or "Apple Silicon" | |
| return f"{chip} (MPS)" | |
| import platform | |
| proc = platform.processor() or "unknown" | |
| return f"CPU ({proc})" | |
| class ChoiceTensors: | |
| X: torch.Tensor | |
| y: torch.Tensor | |
| panel_idx: torch.Tensor | |
| n_individuals: int | |
| n_obs: int | |
| n_alts: int | |
| feature_names: list[str] | |
| id_values: np.ndarray | |
| def load_long_csv(path: str | Path) -> pd.DataFrame: | |
| """Read a long-format CSV file.""" | |
| return pd.read_csv(path) | |
| def validate_long_format(df: pd.DataFrame, spec: ModelSpec) -> None: | |
| """Validate core long-format assumptions.""" | |
| required_cols = { | |
| spec.id_col, | |
| spec.task_col, | |
| spec.alt_col, | |
| spec.choice_col, | |
| *[v.column for v in spec.variables], | |
| } | |
| missing = [c for c in required_cols if c not in df.columns] | |
| if missing: | |
| raise ValueError(f"Missing required columns: {missing}") | |
| key_cols = [spec.id_col, spec.task_col, spec.alt_col] | |
| if df.duplicated(subset=key_cols).any(): | |
| dup_rows = int(df.duplicated(subset=key_cols).sum()) | |
| raise ValueError( | |
| f"Found {dup_rows} duplicated (id, task, alt) rows. " | |
| "Each alternative in each task should appear once." | |
| ) | |
| group_sizes = df.groupby([spec.id_col, spec.task_col]).size() | |
| if group_sizes.empty: | |
| raise ValueError("Input dataframe is empty after grouping by id and task.") | |
| if (group_sizes < 2).any(): | |
| raise ValueError("Each (id, task) must have at least two alternatives.") | |
| if group_sizes.nunique() != 1: | |
| raise ValueError( | |
| "Each (id, task) must have the same number of alternatives. " | |
| "Variable-size choice sets are not supported in this baseline." | |
| ) | |
| def _choice_indices(choice_matrix: np.ndarray, alt_matrix: np.ndarray) -> np.ndarray: | |
| """Convert either one-hot choices or chosen-alt labels to index targets.""" | |
| unique_vals = np.unique(choice_matrix) | |
| # One-hot / binary indicator format. | |
| if np.isin(unique_vals, [0, 1]).all(): | |
| row_sums = choice_matrix.sum(axis=1) | |
| if not np.allclose(row_sums, 1.0): | |
| bad = int(np.where(~np.isclose(row_sums, 1.0))[0][0]) | |
| raise ValueError( | |
| f"Choice indicator rows must sum to 1. Row {bad} sums to {row_sums[bad]}." | |
| ) | |
| return np.argmax(choice_matrix, axis=1).astype(np.int64) | |
| # Label format: each row in a task repeats the same chosen alternative code. | |
| row_constant = np.all(choice_matrix == choice_matrix[:, [0]], axis=1) | |
| if not row_constant.all(): | |
| raise ValueError( | |
| "Choice column is neither one-hot nor a repeated chosen-alt label per task." | |
| ) | |
| chosen_codes = choice_matrix[:, 0] | |
| matches = alt_matrix == chosen_codes[:, None] | |
| valid = matches.sum(axis=1) == 1 | |
| if not valid.all(): | |
| bad = int(np.where(~valid)[0][0]) | |
| raise ValueError( | |
| "Could not map choice code to exactly one alternative in each task. " | |
| f"First invalid task index: {bad}." | |
| ) | |
| return np.argmax(matches, axis=1).astype(np.int64) | |
| def prepare_choice_tensors( | |
| df: pd.DataFrame, | |
| spec: ModelSpec, | |
| device: torch.device | None = None, | |
| ) -> ChoiceTensors: | |
| """ | |
| Convert long-format dataframe into tensors used by estimators. | |
| Expected format: one row per (id, task, alternative), with choice as either: | |
| - one-hot indicator (0/1), or | |
| - chosen alternative label repeated across alternatives in the task. | |
| """ | |
| validate_long_format(df, spec) | |
| if device is None: | |
| device = get_best_device() | |
| sort_cols = [spec.id_col, spec.task_col, spec.alt_col] | |
| work = df.sort_values(sort_cols).reset_index(drop=True) | |
| group_cols = [spec.id_col, spec.task_col] | |
| n_obs = int(work.groupby(group_cols).ngroups) | |
| n_alts = int(work.groupby(group_cols).size().iloc[0]) | |
| n_vars = len(spec.variables) | |
| feature_cols = [v.column for v in spec.variables] | |
| X_flat = work.loc[:, feature_cols].astype(float).to_numpy(dtype=np.float32) | |
| X = X_flat.reshape(n_obs, n_alts, n_vars) | |
| choice_mat = ( | |
| work.loc[:, spec.choice_col] | |
| .to_numpy(dtype=work.loc[:, spec.choice_col].dtype) | |
| .reshape(n_obs, n_alts) | |
| ) | |
| alt_mat = work.loc[:, spec.alt_col].to_numpy().reshape(n_obs, n_alts) | |
| y = _choice_indices(choice_mat, alt_mat) | |
| task_table = work.loc[:, group_cols].drop_duplicates() | |
| obs_ids = task_table.loc[:, spec.id_col].to_numpy() | |
| unique_ids, panel_idx = np.unique(obs_ids, return_inverse=True) | |
| return ChoiceTensors( | |
| X=torch.tensor(X, dtype=torch.float32, device=device), | |
| y=torch.tensor(y, dtype=torch.long, device=device), | |
| panel_idx=torch.tensor(panel_idx, dtype=torch.long, device=device), | |
| n_individuals=len(unique_ids), | |
| n_obs=n_obs, | |
| n_alts=n_alts, | |
| feature_names=[v.name for v in spec.variables], | |
| id_values=unique_ids, | |
| ) | |