{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1051c8cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com\n",
      "Requirement already satisfied: numpy in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (2.4.4)\n",
      "Requirement already satisfied: cvxpy in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (1.8.2)\n",
      "Requirement already satisfied: torch in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (2.8.0+cu126)\n",
      "Requirement already satisfied: cvxpylayers in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (1.1.0)\n",
      "Requirement already satisfied: osqp>=1.0.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from cvxpy) (1.0.4)\n",
      "Requirement already satisfied: clarabel>=0.5.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from cvxpy) (0.11.1)\n",
      "Requirement already satisfied: scs>=3.2.4.post1 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from cvxpy) (3.2.9)\n",
      "Requirement already satisfied: scipy>=1.13.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from cvxpy) (1.15.3)\n",
      "Requirement already satisfied: highspy>=1.11.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from cvxpy) (1.12.0)\n",
      "Requirement already satisfied: filelock in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (3.20.0)\n",
      "Requirement already satisfied: typing-extensions>=4.10.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (4.15.0)\n",
      "Requirement already satisfied: setuptools in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (82.0.1)\n",
      "Requirement already satisfied: sympy>=1.13.3 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (1.14.0)\n",
      "Requirement already satisfied: networkx in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (3.4.2)\n",
      "Requirement already satisfied: jinja2 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (3.1.6)\n",
      "Requirement already satisfied: fsspec in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (2025.7.0)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.6.77 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (12.6.77)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.6.77 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (12.6.77)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.6.80 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (12.6.80)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.10.2.21 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (9.10.2.21)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.6.4.1 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (12.6.4.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.3.0.4 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (11.3.0.4)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.7.77 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (10.3.7.77)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.7.1.2 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (11.7.1.2)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.5.4.2 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (12.5.4.2)\n",
      "Requirement already satisfied: nvidia-cusparselt-cu12==0.7.1 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (0.7.1)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.27.3 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (2.27.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.6.77 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (12.6.77)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.6.85 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (12.6.85)\n",
      "Requirement already satisfied: nvidia-cufile-cu12==1.11.1.6 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (1.11.1.6)\n",
      "Requirement already satisfied: triton==3.4.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from torch) (3.4.0)\n",
      "Requirement already satisfied: diffcp>=1.1.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from cvxpylayers) (1.1.8)\n",
      "Requirement already satisfied: cffi in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from clarabel>=0.5.0->cvxpy) (2.0.0)\n",
      "Requirement already satisfied: threadpoolctl>=1.1 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from diffcp>=1.1.0->cvxpylayers) (3.6.0)\n",
      "Requirement already satisfied: joblib in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from osqp>=1.0.0->cvxpy) (1.5.1)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from sympy>=1.13.3->torch) (1.3.0)\n",
      "Requirement already satisfied: pycparser in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from cffi->clarabel>=0.5.0->cvxpy) (3.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /localscratch/zzhao628/anaconda3/envs/bilevel/lib/python3.12/site-packages (from jinja2->torch) (3.0.3)\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install numpy cvxpy torch cvxpylayers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "md-intro",
   "metadata": {},
   "source": [
    "# DFL-Bench"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c-imports",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:35.492363Z",
     "iopub.status.busy": "2026-05-16T16:31:35.492203Z",
     "iopub.status.idle": "2026-05-16T16:31:37.208897Z",
     "shell.execute_reply": "2026-05-16T16:31:37.207976Z"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import cvxpy as cp\n",
    "import torch\n",
    "from cvxpylayers.torch import CvxpyLayer\n",
    "\n",
    "BENCH = Path('DFL-Bench/tasks')\n",
    "if not BENCH.exists():\n",
    "    BENCH = Path('..') / BENCH\n",
    "# print('using BENCH =', BENCH.resolve())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "md-power-data",
   "metadata": {},
   "source": [
    "## 1. Power scheduling — data loading\n",
    "\n",
    "You may want to design your own data loader for different features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c-power-load",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:37.210971Z",
     "iopub.status.busy": "2026-05-16T16:31:37.210717Z",
     "iopub.status.idle": "2026-05-16T16:31:37.231751Z",
     "shell.execute_reply": "2026-05-16T16:31:37.231121Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_load (1000, 24)\n",
      "test_ctx_load (219, 4, 24)\n",
      "test_day4_temp (219, 24)   n_test = 219\n"
     ]
    }
   ],
   "source": [
    "with open(BENCH / 'power_scheduling' / 'train.json') as f:\n",
    "    tr = json.load(f)\n",
    "with open(BENCH / 'power_scheduling' / 'test.json') as f:\n",
    "    te = json.load(f)\n",
    "\n",
    "train_ids = sorted(tr.keys(), key=int)\n",
    "train_temp = np.array([tr[k]['temp'] for k in train_ids], dtype=np.float64)\n",
    "train_load = np.array([tr[k]['load'] for k in train_ids], dtype=np.float64)\n",
    "\n",
    "test_ids = sorted(te.keys(), key=int)\n",
    "n_test = len(test_ids)\n",
    "test_ctx_load = np.stack([\n",
    "    [te[k][f'day_{d}']['load'] for d in range(4)] for k in test_ids\n",
    "]).astype(np.float64)\n",
    "test_day4_temp = np.array([te[k]['day_4']['temp'] for k in test_ids])\n",
    "\n",
    "print('train_load', train_load.shape)\n",
    "print('test_ctx_load', test_ctx_load.shape)\n",
    "print('test_day4_temp', test_day4_temp.shape, '  n_test =', n_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "md-nv-data",
   "metadata": {},
   "source": "## 2. Newsvendor — data loading\nDownload the data from the [dropbox link](https://www.dropbox.com/scl/fo/sv0xn8vkeejic05g675he/AANKfJb0eyCBC9xeu7FeFUQ?rlkey=2yxjuzadk36pvj9smu8vk0vlm&e=1&st=xx4ur2gs&dl=0); unzip the file in `DFL-Bench/tasks/newsvendor/`.\nYou may want to design your own data loader for different features."
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c-nv-load",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:37.234082Z",
     "iopub.status.busy": "2026-05-16T16:31:37.233959Z",
     "iopub.status.idle": "2026-05-16T16:31:46.037253Z",
     "shell.execute_reply": "2026-05-16T16:31:46.035951Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "items: 8230 (of 8230)\n",
      "sale (1358, 8230) price (1358, 8230)   dates: 2011-01-29 .. 2014-10-17\n"
     ]
    }
   ],
   "source": [
    "import ijson\n",
    "\n",
    "# MAX_ITEMS = 200  # set to None for full 8230 items\n",
    "MAX_ITEMS = None\n",
    "tr_path = BENCH / 'newsvendor' / 'train.json'\n",
    "te_path = BENCH / 'newsvendor' / 'test.json'\n",
    "\n",
    "with open(tr_path, 'rb') as f:\n",
    "    static = next(ijson.items(f, 'static'))\n",
    "all_item_ids = sorted(static.keys())\n",
    "item_ids = all_item_ids if MAX_ITEMS is None else all_item_ids[:MAX_ITEMS]\n",
    "item_idx = {k: i for i, k in enumerate(item_ids)}\n",
    "n_items = len(item_ids)\n",
    "print(f'items: {n_items} (of {len(all_item_ids)})')\n",
    "\n",
    "train_dates, sale_rows, price_rows = [], [], []\n",
    "with open(tr_path, 'rb') as f:\n",
    "    for date, items in ijson.kvitems(f, 'time'):\n",
    "        s = np.full(n_items, np.nan, dtype=np.float32)\n",
    "        p = np.full(n_items, np.nan, dtype=np.float32)\n",
    "        for iid, rec in items.items():\n",
    "            j = item_idx.get(iid)\n",
    "            if j is None: continue\n",
    "            if rec.get('sale')  is not None: s[j] = float(rec['sale'])\n",
    "            if rec.get('price') is not None: p[j] = float(rec['price'])\n",
    "        train_dates.append(date); sale_rows.append(s); price_rows.append(p)\n",
    "\n",
    "order = np.argsort(train_dates)\n",
    "train_dates = [train_dates[i] for i in order]\n",
    "sale  = np.stack([sale_rows[i]  for i in order])\n",
    "price = np.stack([price_rows[i] for i in order])\n",
    "print('sale', sale.shape, 'price', price.shape,\n",
    "      '  dates:', train_dates[0], '..', train_dates[-1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c-nv-test",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:46.039765Z",
     "iopub.status.busy": "2026-05-16T16:31:46.039550Z",
     "iopub.status.idle": "2026-05-16T16:31:47.859815Z",
     "shell.execute_reply": "2026-05-16T16:31:47.858833Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test_ctx_sale (26, 13, 8230)   n_test = 26\n"
     ]
    }
   ],
   "source": [
    "# Test: 26 instances of 14 days each (last day, price only).\n",
    "with open(te_path) as f:\n",
    "    te_raw = json.load(f)\n",
    "test_ids_nv = sorted(te_raw.keys(), key=int)\n",
    "n_test_nv = len(test_ids_nv)\n",
    "\n",
    "test_ctx_sale = np.full((n_test_nv, 13, n_items), np.nan, dtype=np.float32)\n",
    "test_price13  = np.full((n_test_nv, n_items),     np.nan, dtype=np.float32)\n",
    "for k, tid in enumerate(test_ids_nv):\n",
    "    inst = te_raw[tid]\n",
    "    dates_k = sorted(inst.keys())\n",
    "    for d in range(13):\n",
    "        for iid, rec in inst[dates_k[d]].items():\n",
    "            j = item_idx.get(iid)\n",
    "            if j is None: continue\n",
    "            if rec.get('sale') is not None:\n",
    "                test_ctx_sale[k, d, j] = float(rec['sale'])\n",
    "    for iid, rec in inst[dates_k[13]].items():\n",
    "        j = item_idx.get(iid)\n",
    "        if j is None: continue\n",
    "        if rec.get('price') is not None:\n",
    "            test_price13[k, j] = float(rec['price'])\n",
    "print('test_ctx_sale', test_ctx_sale.shape, '  n_test =', n_test_nv)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "md-baselines",
   "metadata": {},
   "source": [
    "## 3. Trivial baselines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "870fe647",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(219, 4, 24)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_ctx_load.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c-baseline-power",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:47.862276Z",
     "iopub.status.busy": "2026-05-16T16:31:47.862134Z",
     "iopub.status.idle": "2026-05-16T16:31:47.865837Z",
     "shell.execute_reply": "2026-05-16T16:31:47.865271Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y_pred_power (219, 24) mean 1.67\n"
     ]
    }
   ],
   "source": [
    "y_pred_power = test_ctx_load.mean(axis=1)          # (n_test, 24)\n",
    "print('y_pred_power', y_pred_power.shape, 'mean', y_pred_power.mean().round(3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d254f816",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(26, 13, 8230)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_ctx_sale.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c-baseline-nv",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:47.868006Z",
     "iopub.status.busy": "2026-05-16T16:31:47.867666Z",
     "iopub.status.idle": "2026-05-16T16:31:47.873249Z",
     "shell.execute_reply": "2026-05-16T16:31:47.872469Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y_pred_nv (26, 8230) mean 2.09 total/day 17199.7\n"
     ]
    }
   ],
   "source": [
    "y_pred_nv = np.nan_to_num(np.nanmean(test_ctx_sale, axis=1), nan=0.0)  # (n_test_nv, n_items)\n",
    "print('y_pred_nv', y_pred_nv.shape,\n",
    "      'mean', y_pred_nv.mean().round(3),\n",
    "      'total/day', y_pred_nv.sum(1).mean().round(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "md-power-solve",
   "metadata": {},
   "source": [
    "## 4. Power — solve + score + submit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c-power-solve",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:47.875090Z",
     "iopub.status.busy": "2026-05-16T16:31:47.874924Z",
     "iopub.status.idle": "2026-05-16T16:31:48.107981Z",
     "shell.execute_reply": "2026-05-16T16:31:48.106607Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "z: (219, 24)\n"
     ]
    }
   ],
   "source": [
    "POWER_PARAMS = {'n': 24, 'c_ramp': 0.4, 'gamma_under': 50.0, 'gamma_over': 0.5}\n",
    "n = POWER_PARAMS['n']\n",
    "z_var = cp.Variable(n)\n",
    "y_par = cp.Parameter(n)\n",
    "diff = z_var - y_par\n",
    "obj = (POWER_PARAMS['gamma_under'] * cp.sum(cp.pos(-diff))\n",
    "       + POWER_PARAMS['gamma_over'] * cp.sum(cp.pos(diff)))\n",
    "cons = [cp.abs(z_var[1:] - z_var[:-1]) <= POWER_PARAMS['c_ramp']]\n",
    "power_layer = CvxpyLayer(cp.Problem(cp.Minimize(obj), cons),\n",
    "                         parameters=[y_par], variables=[z_var])\n",
    "\n",
    "with torch.no_grad():\n",
    "    (z_t,) = power_layer(torch.tensor(y_pred_power, dtype=torch.float32),\n",
    "                         solver_args={'solve_method': 'SCS'})\n",
    "z_np = z_t.cpu().numpy()\n",
    "max_ramp = np.abs(z_np[:, 1:] - z_np[:, :-1]).max()\n",
    "print('z:', z_np.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c-power-eval",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:48.110488Z",
     "iopub.status.busy": "2026-05-16T16:31:48.110319Z",
     "iopub.status.idle": "2026-05-16T16:31:48.118086Z",
     "shell.execute_reply": "2026-05-16T16:31:48.117520Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'violation_mean': 0.0, 'violation_std': 0.0}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def eval_power(z, y_actual, p=POWER_PARAMS):\n",
    "    res = (p['gamma_under'] * np.clip(y_actual - z, 0, None)\n",
    "           + p['gamma_over'] * np.clip(z - y_actual, 0, None)).sum(1)\n",
    "    viol = np.clip(np.abs(z[:, 1:] - z[:, :-1]) - p['c_ramp'], 0, None).sum(1)\n",
    "    return {'violation_mean': float(viol.mean()), 'violation_std': float(viol.std())}\n",
    "\n",
    "eval_power(z_np, y_pred_power) # Note that this is not the true score in the leaderboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c-power-write",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:48.119393Z",
     "iopub.status.busy": "2026-05-16T16:31:48.119257Z",
     "iopub.status.idle": "2026-05-16T16:31:48.127889Z",
     "shell.execute_reply": "2026-05-16T16:31:48.127354Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wrote submissions/submission_power_scheduling.json\n"
     ]
    }
   ],
   "source": [
    "out = Path('submissions') / 'submission_power_scheduling.json'\n",
    "out.parent.mkdir(parents=True, exist_ok=True)\n",
    "with open(out, 'w') as f:\n",
    "    json.dump({tid: z_np[i].tolist() for i, tid in enumerate(test_ids)}, f)\n",
    "print('wrote', out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "md-nv-solve",
   "metadata": {},
   "source": [
    "## 5. Newsvendor — solve + score + submit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c-nv-solve",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:48.129441Z",
     "iopub.status.busy": "2026-05-16T16:31:48.129317Z",
     "iopub.status.idle": "2026-05-16T16:31:48.186199Z",
     "shell.execute_reply": "2026-05-16T16:31:48.185425Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "z (26, 8230) per-day total 17199.7  budget 25000.0\n"
     ]
    }
   ],
   "source": [
    "NEWSVENDOR_PARAMS = {'h': 1.0, 'b': 9.0, 'B': 25000.0}\n",
    "B = NEWSVENDOR_PARAMS['B']\n",
    "\n",
    "z_var = cp.Variable(n_items, nonneg=True)\n",
    "y_par = cp.Parameter(n_items)\n",
    "diff = z_var - y_par\n",
    "obj = (NEWSVENDOR_PARAMS['h'] * cp.sum(cp.pos(diff))\n",
    "       + NEWSVENDOR_PARAMS['b'] * cp.sum(cp.pos(-diff)))\n",
    "nv_layer = CvxpyLayer(cp.Problem(cp.Minimize(obj), [cp.sum(z_var) <= B]),\n",
    "                      parameters=[y_par], variables=[z_var])\n",
    "\n",
    "with torch.no_grad():\n",
    "    (z_t,) = nv_layer(torch.tensor(y_pred_nv, dtype=torch.float32),\n",
    "                      solver_args={'solve_method': 'SCS'})\n",
    "z_n_np = z_t.cpu().numpy()\n",
    "print('z', z_n_np.shape,\n",
    "      'per-day total', z_n_np.sum(1).mean().round(1),\n",
    "      ' budget', round(B, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c-nv-eval",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:48.187748Z",
     "iopub.status.busy": "2026-05-16T16:31:48.187598Z",
     "iopub.status.idle": "2026-05-16T16:31:48.191514Z",
     "shell.execute_reply": "2026-05-16T16:31:48.191047Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'violation_mean': 0.0, 'violation_std': 0.0}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def eval_newsvendor(z, y_actual, h=1.0, b=9.0, B=B):\n",
    "    res = (b * np.clip(y_actual - z, 0, None) + h * np.clip(z - y_actual, 0, None)).sum(1)\n",
    "    viol = np.clip(z.sum(1) - B, 0, None)\n",
    "    return {'violation_mean': float(viol.mean()), 'violation_std': float(viol.std())}\n",
    "\n",
    "eval_newsvendor(z_n_np, y_pred_nv) # Note that this is not the true score in the leaderboard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c-nv-write",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-05-16T16:31:48.193057Z",
     "iopub.status.busy": "2026-05-16T16:31:48.192933Z",
     "iopub.status.idle": "2026-05-16T16:31:48.202650Z",
     "shell.execute_reply": "2026-05-16T16:31:48.202127Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "wrote submissions/submission_newsvendor.json\n"
     ]
    }
   ],
   "source": [
    "out = Path('submissions') / 'submission_newsvendor.json'\n",
    "out.parent.mkdir(parents=True, exist_ok=True)\n",
    "with open(out, 'w') as f:\n",
    "    json.dump({\n",
    "        tid: {iid: float(z_n_np[i, j]) for j, iid in enumerate(item_ids)}\n",
    "        for i, tid in enumerate(test_ids_nv)\n",
    "    }, f)\n",
    "print('wrote', out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ad3637b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5240ea4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bilevel",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}