minor refactors to codebase to implement DRO

engine/lib/__init__.py

@@ -1,4 +1,4 @@
-from .demand import estimate_demand, generate_demand_for_actor
+from .demand import estimate_demand, estimate_weighted_demand, generate_demand_for_actor
 from .behavior import sample_behavior, get_transition_models, trajectory_to_events
 from .render import DashboardRenderer, style_axis
 from .wrappers import EconomicMetricsWrapper

engine/lib/demand.py

@@ -1,9 +1,23 @@
-import logging
 import numpy as np
-from logging import getLogger
-logger = getLogger(__name__)
 
-def generate_demand_for_actor(prices: np.ndarray, params: tuple, noise_std: float = 1.0, distribution_method=np.random.normal) -> np.ndarray:
+CATEGORY_WEIGHTS = {"cart": 4.0, "dwell": 2.0, "nav": 1.0, "filter": 0.5}
+ACTION_CATEGORIES = {
+    "cart": {"add_item", "add_to_cart", "remove", "checkout", "purchase"},
+    "dwell": {"hover_title", "hover_paragraph", "hover_link"},
+    "nav": {"page_view", "view_item", "view", "learn_more"},
+    "filter": {"search", "filter_date", "filter_price", "sort"},
+}
+DEFAULT_ACTION_WEIGHTS = {
+    a: CATEGORY_WEIGHTS[c] for c, actions in ACTION_CATEGORIES.items() for a in actions
+}
+
+
+def generate_demand_for_actor(
+    prices: np.ndarray,
+    params: tuple,
+    noise_std: float = 1.0,
+    distribution_method=np.random.normal,
+) -> np.ndarray:
     """d(p;0) = max(0, valuation - price) + epsi for single actor type
     params: (mean, std) for valuation distribution D_H or D_A"""
     val = distribution_method(*params, size=len(prices))
@@ -13,17 +27,50 @@ def generate_demand_for_actor(prices: np.ndarray, params: tuple, noise_std: floa
     return demand / total * 100 if total > 0 else demand
 
 
-def estimate_demand(trajectories):
-    demand_estimate = {}
+def estimate_demand(trajectories, action_weights=None):
+    return estimate_weighted_demand(trajectories, action_weights)
+
+
+def _parse_event_state(state: str):
+    if "_product" not in state:
+        return state, None
+    action, raw_pid = state.rsplit("_product", 1)
+    return action, int(raw_pid) if raw_pid.isdigit() else None
+
+
+def _weight_for_action(action: str, action_weights: dict) -> float:
+    if action in action_weights:
+        return action_weights[action]
+    if action.startswith("hover"):
+        return CATEGORY_WEIGHTS["dwell"]
+    if action.startswith("filter") or action in {"search", "sort"}:
+        return CATEGORY_WEIGHTS["filter"]
+    if action.startswith("add") or action in {"checkout", "purchase", "remove"}:
+        return CATEGORY_WEIGHTS["cart"]
+    return CATEGORY_WEIGHTS["nav"]
+
+
+def estimate_weighted_demand(trajectories, action_weights=None):
+    action_weights = (
+        DEFAULT_ACTION_WEIGHTS if action_weights is None else action_weights
+    )
+    scores = {}
     for traj in trajectories:
-        for event in traj:
-            if 'view_product' in event:
-                product_id = int(event.split('_')[-1].replace('product', ''))
-                demand_estimate[product_id] = demand_estimate.get(product_id, 0) + 1
-    total_views = sum(demand_estimate.values())
-    for product_id in demand_estimate:
-        demand_estimate[product_id] = (demand_estimate[product_id] / total_views) * 100  # normalize to percentage
-    return demand_estimate
+        for state in traj:
+            action, product_id = _parse_event_state(state)
+            if product_id is None:
+                continue
+            w = _weight_for_action(action, action_weights)
+            if w <= 0:
+                continue
+            scores[product_id] = scores.get(product_id, 0.0) + w
+    total = sum(scores.values())
+    return (
+        {pid: (score / total) * 100 for pid, score in scores.items()}
+        if total > 0
+        else {}
+    )
+
 
 # Example usage
 if __name__ == "__main__":
@@ -36,6 +83,7 @@ if __name__ == "__main__":
     print("Human Demand:", demand_h)
     print("Agent Demand:", demand_a)
     from .behavior import sample_behavior
+
     N, alpha = 200, 0.3
     n_h, n_a = int(N * (1 - alpha)), int(N * alpha)
     human_t = [sample_behavior(demand_h, human=True) for _ in range(n_h)]