challenge_ai/tests/conftest.py

"""Shared fixtures: a spec matching Attachment 1 and helpers for tiny markets."""

from __future__ import annotations

import math

import numpy as np
import pandas as pd
import pytest

from battery_dispatch.config import HALF_HOUR, BatterySpec, RunConfig
from battery_dispatch.markets import MARKET_1, MARKET_2, Market

# Equal to Attachment 1, but stated literally so the tests do not depend on the
# workbook being present or unchanged.
SPEC = BatterySpec(
    max_charge_mw=2.0,
    max_discharge_mw=2.0,
    capacity_mwh=4.0,
    charge_efficiency=0.95,
    discharge_efficiency=0.95,
    lifetime_years=10.0,
    lifetime_cycles=5000.0,
    degradation_pct_per_cycle=0.001,
    capex_gbp=500_000.0,
    fixed_opex_gbp_per_year=5_000.0,
)


@pytest.fixture
def spec() -> BatterySpec:
    return SPEC


@pytest.fixture
def config() -> RunConfig:
    return RunConfig()


def markets_from(
    market_1_prices: list[float],
    market_2_prices: list[float] | None = None,
) -> list[Market]:
    """Build both markets from a half-hourly Market 1 price list.

    ``market_2_prices`` is given per *hour*; when omitted Market 2 is priced so
    low that it never trades, isolating Market 1 behaviour.
    """
    n = len(market_1_prices)
    if market_2_prices is None:
        broadcast = np.zeros(n)
    else:
        assert n == 2 * len(market_2_prices), "one Market 2 price per hour"
        broadcast = np.repeat(np.asarray(market_2_prices, dtype=float), 2)

    return [
        Market(MARKET_1, 1, np.asarray(market_1_prices, dtype=float)),
        Market(MARKET_2, 2, broadcast),
    ]


def prices_frame(
    market_1_prices: list[float],
    market_2_prices: list[float],
    start: str = "2018-01-01",
) -> pd.DataFrame:
    """A price frame shaped like the one ``load_prices`` returns."""
    n = len(market_1_prices)
    index = pd.date_range(start, periods=n, freq="30min")
    return pd.DataFrame(
        {
            "market_1_price": np.asarray(market_1_prices, dtype=float),
            "market_2_price": np.repeat(np.asarray(market_2_prices, dtype=float), 2),
            "hour_index": np.arange(n) // 2,
        },
        index=index,
    )


def make_valid_schedule(spec: BatterySpec, n: int = 48) -> pd.DataFrame:
    """A schedule (shaped like the optimiser's output) that satisfies every
    invariant ``validate_schedule`` checks -- built directly from the model's
    own energy-balance and revenue formulas, not by running the optimiser.

    ``validate_schedule`` is tested in isolation: it only needs *some*
    structurally valid schedule to mutate one violation into. Deriving that
    from a real rolling-horizon solve would mean spinning up CBC once per
    test for no reason connected to what is under test, so this constructs
    one by hand instead -- charge for two hours, sit idle, discharge for two
    hours, comfortably within every power and capacity bound.
    """
    charge1 = np.zeros(n)
    charge2 = np.zeros(n)
    discharge1 = np.zeros(n)
    discharge2 = np.zeros(n)
    charge1[0:4] = 1.0
    charge2[0:4] = 0.5
    discharge1[36:40] = 0.5
    discharge2[36:40] = 0.3

    total_charge = charge1 + charge2
    total_discharge = discharge1 + discharge2
    delta = HALF_HOUR * (
        spec.charge_efficiency * total_charge - total_discharge / spec.discharge_efficiency
    )
    soc = np.cumsum(delta)
    assert soc.min() >= -1e-9 and soc.max() <= spec.capacity_mwh - 1e-9, (
        "fixture parameters must stay clear of the battery's bounds"
    )

    price1 = np.array([40.0 + 10.0 * math.sin(i / 5) for i in range(n)])
    price2_hourly = np.array([35.0 + 8.0 * math.sin(h / 4) for h in range(n // 2)])
    price2 = np.repeat(price2_hourly, 2)

    frame = pd.DataFrame(
        {
            "market_1_price": price1,
            "market_2_price": price2,
            "charge_market_1_mw": charge1,
            "discharge_market_1_mw": discharge1,
            "charge_market_2_mw": charge2,
            "discharge_market_2_mw": discharge2,
            "soc_mwh": soc,
            "capacity_mwh": np.full(n, spec.capacity_mwh),
        },
        index=pd.date_range("2018-01-01", periods=n, freq="30min"),
    )
    frame.index.name = "timestamp"
    for name in (MARKET_1, MARKET_2):
        net_export = frame[f"discharge_{name}_mw"] - frame[f"charge_{name}_mw"]
        frame[f"revenue_{name}_gbp"] = HALF_HOUR * frame[f"{name}_price"] * net_export
    return frame