Added docstrings to functions

Author: m9o8

CLAUDE.md

@@ -67,6 +67,11 @@ This is an academic research project investigating algorithmic collusion among L
 - Focus on final 50 periods (251-300) for equilibrium analysis
 - Bootstrap validation and robustness testing
 
+**`/src/plotting/`** - Visualization Tools
+- `final_figures.py` generates publication-ready statistical plots
+- `plotting.py` contains core visualization utilities
+- `run_animation.py` creates animated visualizations of price dynamics
+
 ### Data Flow
 
 1. **Configuration** → Market parameters (α, β, μ, costs) and agent setup

src/README.md

@@ -0,0 +1,75 @@
+# Source Code Architecture
+
+This directory contains the core implementation for algorithmic collusion experiments testing Folk Theorem predictions in multi-agent pricing games.
+
+## Module Structure
+
+### `agents/` - LLM Agent Implementation
+
+- `LLMAgent` - Mistral API integration for strategic pricing decisions
+- `Agent` - Abstract base class defining agent interface
+- `FakeAgent` - Deterministic agent for testing and baselines
+
+### `environment/` - Market Simulation
+
+- `CalvanoDemandEnvironment` - Implements Calvano et al. (2020) demand specification
+- Calculates quantities, profits, and theoretical benchmarks
+- Supports both static and time-varying cost structures
+
+### `experiment/` - Experiment Orchestration
+
+- `Experiment` - Coordinates multi-agent pricing games
+- `StorageManager` - Handles data persistence in Parquet format
+- `RateLimiter` - Manages API rate limiting for cost control
+
+### `prompts/` - Prompt Engineering
+
+- `PromptManager` - Dynamic context injection and memory management
+- Systematic P1/P2 prompt specifications for coordination testing
+- Structured response validation using Pydantic models
+
+### `analysis/` - Statistical Analysis
+
+- `CollusionAnalysis` - Core regression models for Folk Theorem testing
+- Bootstrap validation and robustness testing
+- Focus on run-level equilibrium analysis (periods 251-300)
+
+### `plotting/` - Visualization Tools
+
+- Publication-ready figures with theoretical benchmark overlays
+- Time series analysis and convergence visualization
+- Animated price dynamics for presentation
+
+### `utils/` - Utility Functions
+
+- Logging configuration for experiment tracking
+- Price convergence analysis and statistical helpers
+- Data processing utilities
+
+## Key Design Principles
+
+**Rate Limiting**: Built-in API management prevents cost overruns
+**Memory Architecture**: 100-period rolling windows for strategic learning
+**Data Integrity**: Parquet storage with comprehensive metadata
+**Reproducibility**: Structured logging and configuration management
+**Modularity**: Clean separation between agents, environment, and analysis
+
+## Usage Patterns
+
+```python
+from src.agents import LLMAgent
+from src.environment import CalvanoDemandEnvironment
+from src.experiment import Experiment
+
+# Create environment
+env = CalvanoDemandEnvironment("oligopoly", "2-agent market")
+
+# Initialize agents
+agents = [LLMAgent(f"Agent_{i}", api_key=api_key, ...) for i in range(2)]
+
+# Run experiment
+experiment = Experiment(agents, env, n_rounds=300)
+results = experiment.run()
+```
+
+This architecture supports the systematic testing of Folk Theorem predictions through controlled LLM agent interactions in oligopoly markets.

src/agents/LLM_agent.py

@@ -1,4 +1,10 @@
-# src/agents/LLM_agent.py
+"""
+LLM Agent Implementation for Market Simulation
+
+This module provides the LLMAgent class that interfaces with Mistral API
+for strategic pricing decisions in oligopoly market experiments.
+"""
+
 import json
 import logging
 import asyncio
@@ -14,6 +20,25 @@
 
 
 class LLMAgent(Agent):
+    """
+    LLM-based agent for strategic pricing decisions in oligopoly markets.
+
+    This agent uses Mistral API to make pricing decisions based on market history
+    and strategic context. It maintains a rolling memory window and validates
+    responses using Pydantic models.
+
+    Args:
+        name: Unique identifier for the agent
+        prefix: Prompt prefix defining agent's strategic context
+        api_key: Mistral API key for authentication
+        model_name: Mistral model to use (e.g., 'mistral-large-2411')
+        response_model: Pydantic model for response validation
+        prompt_template: Optional template for prompt formatting
+        memory_length: Number of periods to maintain in memory (default: 100)
+        env_params: Market environment parameters
+        logger: Logger instance for experiment tracking
+    """
+
     def __init__(
         self,
         name: str,
@@ -54,6 +79,22 @@ def __init__(
         )
 
     async def act(self, prompt: str) -> Dict:
+        """
+        Execute pricing decision using LLM API.
+
+        Makes API call to Mistral with retry logic and validates response
+        using the configured Pydantic model.
+
+        Args:
+            prompt: Market context and decision prompt
+
+        Returns:
+            Dict containing agent name and validated response content
+
+        Raises:
+            ValidationError: If response validation fails after all retries
+            Exception: If API call fails after all retries
+        """
         async with Mistral(api_key=self.api_key) as client:
             for attempt in range(1, MAX_RETRIES + 1):
                 self.logger.info(f"🔄 Attempt {attempt} for agent {self.name}")
@@ -90,6 +131,18 @@ async def act(self, prompt: str) -> Dict:
                 await asyncio.sleep(RETRY_DELAY_SECONDS * attempt)
 
     def __filter_in_answer_fields(self, model: Type[BaseModel]) -> Type[BaseModel]:
+        """
+        Filter Pydantic model fields marked for inclusion in responses.
+
+        Creates a new model containing only fields with 'in_answer=True'
+        in their json_schema_extra metadata.
+
+        Args:
+            model: Source Pydantic model to filter
+
+        Returns:
+            New Pydantic model with only filtered fields
+        """
         # Filter the fields that have 'in_answer=True'
         in_answer_fields = {
             name: (field.annotation, field.default)
@@ -100,5 +153,6 @@ def __filter_in_answer_fields(self, model: Type[BaseModel]) -> Type[BaseModel]:
         return create_model(model.__name__ + "Filtered", **in_answer_fields)
 
     @property
-    def type(self) -> bool:
+    def type(self) -> str:
+        """Return agent type identifier."""
         return "LLM_agent"

src/agents/base_agent.py

@@ -1,9 +1,30 @@
-# src/agents/base_agent.py
+"""
+Base Agent Interface for Market Simulation
+
+This module defines the abstract base class for all market agents,
+providing common functionality for pricing decisions and market interaction.
+"""
+
 from abc import ABC, abstractmethod
 from typing import Dict, Any
 
 
 class Agent(ABC):
+    """
+    Abstract base class for market agents in oligopoly experiments.
+
+    Defines the interface for agents that make pricing decisions based on
+    market context and strategic parameters.
+
+    Args:
+        name: Unique identifier for the agent
+        prefix: Strategic context prefix for decision-making
+        prompt_template: Template for formatting decision prompts
+        env_index: Agent's index in the market environment
+        env_params: Market environment parameters (a, alpha, c)
+        logger: Logger instance for experiment tracking
+    """
+
     def __init__(
         self,
         name: str,
@@ -35,14 +56,25 @@ def __init__(
 
     @property
     def requires_prompt(self) -> bool:
+        """Whether agent requires prompt for decision-making."""
         return True  # default for LLM-based agents
 
     @property
     def type(self) -> str:
+        """Return agent type identifier."""
         return None
 
     @abstractmethod
     async def act(self, prompt: str) -> Dict[str, Any]:
+        """
+        Execute pricing decision based on market context.
+
+        Args:
+            prompt: Market context and decision prompt
+
+        Returns:
+            Dict containing agent response and metadata
+        """
         pass
 
     def get_marginal_cost(self, round_num: int) -> float:

src/agents/fake_agent.py

@@ -1,9 +1,28 @@
-# src/agents/fake_agent.py
+"""
+Fake Agent for Testing and Simulation
+
+This module provides a deterministic agent that follows pre-defined price
+sequences, useful for testing and baseline comparisons.
+"""
+
 import numpy as np
 from src.agents.base_agent import Agent
 
 
 class FakeAgent(Agent):
+    """
+    Deterministic agent that follows pre-defined price sequences.
+
+    Used for testing market mechanisms and providing baseline comparisons
+    with deterministic pricing behavior.
+
+    Args:
+        name: Unique identifier for the agent
+        time_series_data: Array of predetermined prices to follow
+        nbr_rounds: Number of rounds the agent will participate in
+        **kwargs: Additional arguments passed to parent Agent class
+    """
+
     def __init__(
         self, name: str, time_series_data: np.ndarray, nbr_rounds: int, **kwargs
     ):
@@ -19,6 +38,15 @@ def __init__(
         self.current_index = 0
 
     async def act(self, prompt: str) -> dict:
+        """
+        Return the next predetermined price in the sequence.
+
+        Args:
+            prompt: Ignored, as fake agent doesn't use prompts
+
+        Returns:
+            Dict containing agent name and next price in sequence
+        """
         chosen_price = self.time_series_data[self.current_index]
         self.current_index += 1
         return {
@@ -28,8 +56,10 @@ async def act(self, prompt: str) -> dict:
 
     @property
     def requires_prompt(self) -> bool:
+        """Fake agents don't require prompts."""
         return False
 
     @property
     def type(self) -> str:
+        """Return agent type identifier."""
         return "fake_agent"

src/analysis/group_size.py

@@ -1,3 +1,10 @@
+"""
+Statistical Analysis for Folk Theorem Testing
+
+This module provides core regression models and statistical analysis
+for testing Folk Theorem predictions in multi-agent pricing experiments.
+"""
+
 import warnings
 
 import matplotlib.pyplot as plt
@@ -12,6 +19,18 @@
 
 
 class CollusionAnalysis:
+    """
+    Statistical analysis framework for testing Folk Theorem predictions.
+
+    Provides regression models, robustness tests, and visualization tools
+    for analyzing pricing coordination breakdown as group size increases.
+
+    Args:
+        df: Polars DataFrame with experimental data
+            Required columns: run_id, period, agent_id, group_size,
+            prompt_type, price, alpha, monopoly_prices, nash_prices
+    """
+
     def __init__(self, df: pl.DataFrame):
         """
         Initialize with a Polars DataFrame containing experimental data