source/gconvex_authors/scripts/solve_mechanisms.py

"""Mechanism-design sweep over dimensions, kernels, and costs.

This script instantiates the generic :class:`mech_design.mechanism.Mechanism`
model for a grid of problem dimensions and experiment families
(`kernel`, `cost_fn`, and type distribution), trains each configuration,
and stores checkpoints/snapshots under ``WRITING_ROOT/mech/``. Types ``X``
are screened against outcome/bundle vectors ``Y``, where each ``Y`` is a
finite convex candidate in the underlying grid model.

Two main families are swept:
- Dot-product kernel with no production cost on bounded types in ``[0, 1]^d``.
- Hinge kernel with quadratic production costs on lognormal types, corresponding
  to screening with one-sided surplus above types and quadratic production costs.

Typical usage
-------------
Run a full sweep (using the default output root from ``config.py``)::

    python scripts/solve_mechanisms.py --niters 20000 --batch-size 512

Use ``--clear`` to remove previous runs in that directory before starting a
fresh sweep.
"""
import argparse
import os
import shutil
from math import exp
import numpy as np
import torch
from config import WRITING_ROOT
import mech_design.mechanism as mechanism_module

def dot_kernel(X: torch.Tensor, Y: torch.Tensor) -> torch.Tensor:
    """Dot-product kernel :math:`k(x, y) = \\langle x, y \\rangle`.

    Here ``X`` is a vector of types/valuations and ``Y`` is a bundle/outcome
    vector; each ``Y`` corresponds to a finite convex candidate in the
    underlying model.

    Parameters
    ----------
    X, Y:
        Batches of type and outcome vectors with matching last dimension. In
        the common case, both have shape ``(batch_size, dim)``.

    Returns
    -------
    torch.Tensor
        Kernel values with the leading batch shape, typically
        ``(batch_size,)``.
    """
    return (X * Y).sum(dim=-1)

def hinge_kernel(X: torch.Tensor, Y: torch.Tensor, κ: float = 1.0) -> torch.Tensor:
    """One-sided linear kernel that only rewards allocations above the type.

    Computes ``κ * sum(max(Y - X, 0))`` along the last dimension, i.e. shortfall
    below ``X`` is ignored while surplus above ``X`` contributes linearly. This
    is the kernel used in the screening-with-one-sided-surplus experiments,
    combined with :func:`quadratic_cost`.

    Parameters
    ----------
    X, Y:
        Type and outcome vectors with matching last dimension; each ``Y`` is a
        finite convex candidate bundle.
    κ:
        Overall scale applied to the hinge payoff.

    Returns
    -------
    torch.Tensor
        Kernel values with the leading batch shape.
    """
    return κ * (Y - X).clamp_min(0).sum(dim=-1)

def quadratic_cost(Y: torch.Tensor, β: float = 0.1) -> torch.Tensor:
    """Additive quadratic production cost :math:`β \\lVert Y \\rVert_2^2`.

    Parameters
    ----------
    Y:
        Candidate bundle/allocation vectors.
    β:
        Scalar weight on the squared Euclidean norm of ``Y``.

    Returns
    -------
    torch.Tensor
        Per-sample cost with shape matching the leading batch dimensions of
        ``Y`` (typically ``(batch_size,)``).
    """
    return β * (Y**2).sum(dim=-1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run mechanism design sweep.")
    parser.add_argument(
        "-c",
        "--clear",
        action="store_true",
        help="Remove existing snapshots in the writing_dir before training.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=5_000,
        help="Chunk size used when computing the kernel to limit memory.",
    )
    parser.add_argument(
        "--niters",
        type=int,
        default=60_000,
        help="Number of training iterations (passed as nsteps).",
    )
    parser.add_argument(
        "--epsilon",
        type=float,
        default=2e-3,
        help="Barrier epsilon threshold for the constraints.",
    )
    args = parser.parse_args()
    writing_dir_base = f"{WRITING_ROOT}mech/"
    if args.clear:
        shutil.rmtree(writing_dir_base, ignore_errors=True)
        print(f"Cleared {writing_dir_base}")
    dims = [1, 2, 4, 6, 12, 50, 100, 250]
    initial_penalty_factors = [100.0/d for d in dims]
    npoints = [d + 10 for d in dims]  # number of candidate bundles per dimension
    epsilon = args.epsilon
    kernel_batch_size=args.batch_size        
    patience = 600
    cooldown = 300
    factor = 0.85
    nsteps = args.niters
    max_samples = 1_000
    temp = 60.
    temp_warmup_steps = 1000
    temp_schedule_initial = 5.0
    lrs = [0.05 for d in dims]  # potentially dimension-scaled learning rate
    convergence_tolerance = 1e-8
    sorted_model = True
    is_Y_parameter = True
    window = 1000
    # Utility level of the outside option when present. We shift it into the
    # model via a negative intercept; if this were exactly zero, optimization
    # may get stuck at the beginning because the default and non-default
    # options start too symmetrically.
    default_utility = 1.5*epsilon
    scheduler_threshold = 1e-2
    max_clipping_norm = 5

    # Pre-sampled type draws reused across dimensions so that different
    # kernels/costs see comparable inputs.
    unif_sample = torch.rand(max_samples, max(dims))
    lognormal_sample = torch.exp(torch.randn(max_samples, max(dims)) * 0.25)

    y_maxes = [1.0, None]
    costs =   [None, quadratic_cost]
    kernels = [dot_kernel, hinge_kernel]
    samples = [unif_sample, lognormal_sample]
    are_there_defaults = [True, False]


    # Each tuple (y_max, cost, kernel, sample, is_there_default) defines one
    # sweep family:
    #   - dot_kernel + no cost on bounded types in [0, 1]^d
    #   - hinge_kernel + quadratic_cost on unbounded lognormal types
    for y_max, cost, kernel, big_sample, is_there_default in zip(y_maxes, costs, kernels, samples, are_there_defaults):
        for dim, npoint, lr, initial_penalty_factor in zip(dims, npoints, lrs, initial_penalty_factors):
            print(f"Running mechanism design for dim={dim}, npoints={npoint}")
            sample = big_sample[:, :dim]
            if sorted_model:
                sample, _ = torch.sort(sample, dim=-1)
            model_kwargs = {
                "npoints": npoint,
                "kernel": kernel,
                "cost_fn": cost,
                "y_dim": dim,
                "temp": temp,
                "is_Y_parameter": is_Y_parameter,
                "is_there_default": is_there_default,
                "default_intercept": -default_utility,
                "y_min": 0.0,
                "y_max": y_max,
                "sorted_model": sorted_model,
            }
            mechanism_instance = mechanism_module.Mechanism(**model_kwargs, kernel_batch_size=kernel_batch_size)
            with torch.no_grad():
                q = torch.linspace(0.3, 0.9, mechanism_instance.num_candidates, device=sample.device)
                Y_target = torch.quantile(sample, q, dim=0).unsqueeze(0)  # sample is sorted already
                # keep this to satisfy the sorted parametrization
                Y_target, _ = Y_target.sort(dim=-1)
                diffs = torch.cat([Y_target[..., :1], Y_target[..., 1:] - Y_target[..., :-1]], dim=-1)
                raw = torch.log(torch.expm1(diffs.clamp_min(1e-8)))
                mechanism_instance.Y_rest_raw.copy_(raw)
                if cost is not None:
                    costs_per_candidate = cost(Y_target.squeeze(0)).view(1, -1)  # shape (1, npoints)
                    mechanism_instance.intercept_rest.copy_(costs_per_candidate)
                else:
                    mechanism_instance.intercept_rest.zero_()

            writing_dir_dim = os.path.join(writing_dir_base, f"dim{dim}/")
            mechanism, mechanism_data = mechanism_instance.fit(
                sample,
                already_sorted=True,
                modes=["soft"],
                compile=True,
                optimizers_kwargs_dict={"soft": {"lr": lr}},
                schedulers_kwargs_dict={
                    "soft": {
                        "patience": patience,
                        "threshold": scheduler_threshold,
                        "factor": factor,
                        "cooldown": cooldown,
                        "eps": 1e-8,
                    },
                },
                train_kwargs={
                    "nsteps": nsteps,
                    "max_clipping_norm": max_clipping_norm,
                    "initial_penalty_factor": initial_penalty_factor,
                    "steps_per_snapshot": 200,
                    "steps_per_update": 5,
                    "window": window,
                    "constraint_fns": [],
                    "use_wandb": False,
                    "writing_dir": writing_dir_dim,
                    "convergence_tolerance": convergence_tolerance,
                    "epsilon": epsilon,
                    "temp_warmup_steps": temp_warmup_steps,
                    "temp_schedule_initial": temp_schedule_initial,
                    "switch_threshold": 0.995,

            },
        )