foxhunt/crates/ml/tests/sp12_reward_math_tests.rs

#![allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory.

//! SP12 v3 reward-math GPU oracle tests.
//!
//! Validates the three reward composition device functions added to
//! `trade_physics.cuh` for SP12 v3 (commit 17cfbb250):
//!   1. `compute_asymmetric_capped_pnl`  — prospect-theory asymmetric cap
//!   2. `compute_min_hold_penalty`        — soft-saturation patience penalty
//!   3. `compute_lump_sum_opp_cost`       — lump-sum carrying-cost on exit
//!
//! Each test drives a synthetic input through the production device
//! function (via the standalone wrapper `sp12_reward_math_test_kernel.cu`)
//! and verifies the GPU-computed result against an analytically-known
//! expected value. Per `feedback_no_cpu_test_fallbacks`: GPU oracle only,
//! no CPU reference impl. Per `feedback_no_htod_htoh_only_mapped_pinned`:
//! every CPU↔GPU buffer is a `MappedF32Buffer` (cuMemHostAlloc with
//! DEVICEMAP|PORTABLE); zero `htod_copy`, zero `dtoh_sync_copy`.
//!
//! All tests are `#[ignore = "requires GPU"]`-gated to match every other
//! GPU oracle test in this crate (sp4/sp5/sp11). Run on a GPU host:
//!
//!     SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \
//!       cargo test -p ml --test sp12_reward_math_tests --features cuda \
//!         -- --ignored --nocapture
//!
//! The wrapper kernel is single-thread / single-block; each test launch
//! drives one (input → output) triple in parallel (the kernel writes all
//! three outputs in one invocation), so the test helper returns a tuple
//! `(capped_pnl, min_hold_penalty, opp_cost)` and individual assertions
//! pick the relevant slot.

#![cfg(feature = "cuda")]

use std::sync::Arc;

use cudarc::driver::{CudaContext, CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer;

/// Test-only cubin built by `crates/ml/build.rs`.
const SP12_REWARD_MATH_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/sp12_reward_math_test_kernel.cubin"));

/// Resolve a CUDA stream against device 0. Mirrors `make_test_stream` in
/// every other oracle test in this crate (`sp4_producer_unit_tests.rs`,
/// `sp11_producer_unit_tests.rs`, `distributional_q_tests.rs`).
fn make_test_stream() -> Arc<CudaStream> {
    let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?");
    ctx.default_stream()
}

/// Load the test kernel function from the embedded cubin.
fn load_test_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
    let module = stream
        .context()
        .load_cubin(SP12_REWARD_MATH_CUBIN.to_vec())
        .expect("load sp12_reward_math_test_kernel cubin");
    module
        .load_function("sp12_reward_math_test_kernel")
        .expect("load sp12_reward_math_test_kernel function")
}

/// Inputs for one launch of the SP12 reward-math wrapper kernel.
///
/// Carries the inputs each device function consumes. The wrapper kernel
/// computes all three results in one invocation; tests project the
/// returned tuple's relevant slot.
#[derive(Clone, Copy, Debug)]
struct Sp12RewardInputs {
    /* asymmetric-cap inputs */
    base_reward: f32,
    neg_cap:     f32,
    pos_cap:     f32,
    /* min-hold-penalty inputs */
    hold_time_mh:         f32,
    min_hold_target:      f32,
    min_hold_temperature: f32,
    min_hold_penalty_max: f32,
    /* lump-sum opp-cost inputs */
    shaping_scale:     f32,
    holding_cost_rate: f32,
    position:          f32,
    hold_time_oc:      f32,
}

impl Sp12RewardInputs {
    /// Production constants (from `state_layout.cuh`) for the asymmetric cap.
    /// `min_hold_*` and opp-cost params get neutral defaults so a test that
    /// only cares about the asymmetric cap doesn't pay attention to the
    /// other two outputs.
    fn defaults() -> Self {
        Self {
            base_reward: 0.0,
            neg_cap:    -10.0,
            pos_cap:     5.0,
            hold_time_mh:         30.0, // == target → penalty 0
            min_hold_target:      30.0,
            min_hold_temperature: 10.0,
            min_hold_penalty_max:  3.0,
            shaping_scale:     1.0,
            holding_cost_rate: 0.001,
            position:          0.0,    // → opp_cost 0
            hold_time_oc:      0.0,
        }
    }
}

/// Returns `(capped_pnl, min_hold_penalty, opp_cost)` from one wrapper
/// kernel launch. Single-thread single-block — the device functions are
/// pure register arithmetic, no parallelism to exploit.
fn launch_sp12_reward_math(
    stream: &Arc<CudaStream>,
    f: &CudaFunction,
    inputs: Sp12RewardInputs,
) -> (f32, f32, f32) {
    // Safety: CUDA context active on this thread (resolved through the
    // stream's context above). All CPU↔GPU buffers are mapped-pinned.
    let in_base_reward = unsafe { MappedF32Buffer::new(1) }
        .expect("alloc in_base_reward (1 f32)");
    in_base_reward.write_from_slice(&[inputs.base_reward]);

    let in_hold_time_mh = unsafe { MappedF32Buffer::new(1) }
        .expect("alloc in_hold_time_mh (1 f32)");
    in_hold_time_mh.write_from_slice(&[inputs.hold_time_mh]);

    let in_position = unsafe { MappedF32Buffer::new(1) }
        .expect("alloc in_position (1 f32)");
    in_position.write_from_slice(&[inputs.position]);

    let in_hold_time_oc = unsafe { MappedF32Buffer::new(1) }
        .expect("alloc in_hold_time_oc (1 f32)");
    in_hold_time_oc.write_from_slice(&[inputs.hold_time_oc]);

    let out_capped_pnl = unsafe { MappedF32Buffer::new(1) }
        .expect("alloc out_capped_pnl (1 f32)");
    let out_min_hold_penalty = unsafe { MappedF32Buffer::new(1) }
        .expect("alloc out_min_hold_penalty (1 f32)");
    let out_opp_cost = unsafe { MappedF32Buffer::new(1) }
        .expect("alloc out_opp_cost (1 f32)");

    let in_base_reward_dev = in_base_reward.dev_ptr;
    let in_hold_time_mh_dev = in_hold_time_mh.dev_ptr;
    let in_position_dev = in_position.dev_ptr;
    let in_hold_time_oc_dev = in_hold_time_oc.dev_ptr;
    let out_capped_pnl_dev = out_capped_pnl.dev_ptr;
    let out_min_hold_penalty_dev = out_min_hold_penalty.dev_ptr;
    let out_opp_cost_dev = out_opp_cost.dev_ptr;

    unsafe {
        stream
            .launch_builder(f)
            // asymmetric cap
            .arg(&in_base_reward_dev)
            .arg(&inputs.neg_cap)
            .arg(&inputs.pos_cap)
            .arg(&out_capped_pnl_dev)
            // min-hold penalty
            .arg(&in_hold_time_mh_dev)
            .arg(&inputs.min_hold_target)
            .arg(&inputs.min_hold_temperature)
            .arg(&inputs.min_hold_penalty_max)
            .arg(&out_min_hold_penalty_dev)
            // lump-sum opp_cost
            .arg(&inputs.shaping_scale)
            .arg(&inputs.holding_cost_rate)
            .arg(&in_position_dev)
            .arg(&in_hold_time_oc_dev)
            .arg(&out_opp_cost_dev)
            .launch(LaunchConfig {
                grid_dim:         (1, 1, 1),
                block_dim:        (1, 1, 1),
                shared_mem_bytes: 0,
            })
            .expect("launch sp12_reward_math_test_kernel");
    }
    stream.synchronize().expect("sync after sp12_reward_math launch");

    (
        out_capped_pnl.read_all()[0],
        out_min_hold_penalty.read_all()[0],
        out_opp_cost.read_all()[0],
    )
}

// ─── Asymmetric capped P&L (Change 1 — prospect-theory loss aversion) ──────

/// Above the positive cap → clipped to `pos_cap`.
#[test]
#[ignore = "requires GPU"]
fn sp12_asymmetric_cap_clips_above_pos_cap() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.base_reward = 20.0;
    let (capped, _, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (capped - 5.0).abs() < 1e-5,
        "+20 should clip to pos_cap=+5, got {capped}"
    );
}

/// Below the negative cap → clipped to `neg_cap`.
#[test]
#[ignore = "requires GPU"]
fn sp12_asymmetric_cap_clips_below_neg_cap() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.base_reward = -20.0;
    let (capped, _, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (capped - (-10.0)).abs() < 1e-5,
        "-20 should clip to neg_cap=-10, got {capped}"
    );
}

/// Inside the band (zero) → identity passthrough.
#[test]
#[ignore = "requires GPU"]
fn sp12_asymmetric_cap_passes_through_zero() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.base_reward = 0.0;
    let (capped, _, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(capped.abs() < 1e-5, "0 should pass through, got {capped}");
}

/// Exactly at `pos_cap` → boundary identity.
#[test]
#[ignore = "requires GPU"]
fn sp12_asymmetric_cap_at_pos_cap() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.base_reward = 5.0;
    let (capped, _, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (capped - 5.0).abs() < 1e-5,
        "+5 (== pos_cap) should pass through, got {capped}"
    );
}

/// Exactly at `neg_cap` → boundary identity.
#[test]
#[ignore = "requires GPU"]
fn sp12_asymmetric_cap_at_neg_cap() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.base_reward = -10.0;
    let (capped, _, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (capped - (-10.0)).abs() < 1e-5,
        "-10 (== neg_cap) should pass through, got {capped}"
    );
}

// ─── Min-hold soft penalty (Change 2 — patience curriculum) ────────────────

/// `hold_time == target` → deficit 0, soft factor 0, penalty 0.
#[test]
#[ignore = "requires GPU"]
fn sp12_min_hold_penalty_at_target_is_zero() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.hold_time_mh         = 30.0;
    inputs.min_hold_target      = 30.0;
    inputs.min_hold_temperature = 10.0;
    inputs.min_hold_penalty_max = 3.0;
    let (_, penalty, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        penalty.abs() < 1e-5,
        "hold_time == target should give 0 penalty, got {penalty}"
    );
}

/// `hold_time = 0`, target = 30, T = 10 → deficit 30, factor 30/40 = 0.75,
/// penalty = 3 × 0.75 = 2.25.
#[test]
#[ignore = "requires GPU"]
fn sp12_min_hold_penalty_zero_hold_time_max_target() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.hold_time_mh         = 0.0;
    inputs.min_hold_target      = 30.0;
    inputs.min_hold_temperature = 10.0;
    inputs.min_hold_penalty_max = 3.0;
    let (_, penalty, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (penalty - 2.25).abs() < 1e-5,
        "hold=0,target=30,T=10,max=3 should give penalty=2.25, got {penalty}"
    );
}

/// `hold_time = 15`, target = 30, T = 10 → deficit 15, factor 15/25 = 0.6,
/// penalty = 3 × 0.6 = 1.8.
#[test]
#[ignore = "requires GPU"]
fn sp12_min_hold_penalty_half_way_deficit() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.hold_time_mh         = 15.0;
    inputs.min_hold_target      = 30.0;
    inputs.min_hold_temperature = 10.0;
    inputs.min_hold_penalty_max = 3.0;
    let (_, penalty, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (penalty - 1.8).abs() < 1e-5,
        "hold=15,target=30,T=10,max=3 should give penalty=1.8, got {penalty}"
    );
}

/// `hold_time > target` → no penalty (early-exit branch in device fn).
#[test]
#[ignore = "requires GPU"]
fn sp12_min_hold_penalty_past_target_no_penalty() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.hold_time_mh         = 50.0;
    inputs.min_hold_target      = 30.0;
    inputs.min_hold_temperature = 10.0;
    inputs.min_hold_penalty_max = 3.0;
    let (_, penalty, _) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        penalty.abs() < 1e-5,
        "hold > target should give 0 penalty, got {penalty}"
    );
}

/// Higher temperature smooths the curve: same deficit gives a SMALLER
/// penalty under T=20 than under T=10. With deficit=30 (hold=0,target=30):
///   T=10 → factor=30/40=0.75 → penalty=2.25
///   T=20 → factor=30/50=0.60 → penalty=1.80
/// Verifies `compute_min_hold_penalty` is monotone-decreasing in T at
/// fixed deficit (the curriculum-anneal direction).
#[test]
#[ignore = "requires GPU"]
fn sp12_min_hold_penalty_temperature_smooths_transition() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);

    let mut inputs_t10 = Sp12RewardInputs::defaults();
    inputs_t10.hold_time_mh         = 0.0;
    inputs_t10.min_hold_target      = 30.0;
    inputs_t10.min_hold_temperature = 10.0;
    inputs_t10.min_hold_penalty_max = 3.0;
    let (_, penalty_t10, _) = launch_sp12_reward_math(&stream, &f, inputs_t10);

    let mut inputs_t20 = Sp12RewardInputs::defaults();
    inputs_t20.hold_time_mh         = 0.0;
    inputs_t20.min_hold_target      = 30.0;
    inputs_t20.min_hold_temperature = 20.0;
    inputs_t20.min_hold_penalty_max = 3.0;
    let (_, penalty_t20, _) = launch_sp12_reward_math(&stream, &f, inputs_t20);

    assert!(
        penalty_t10 > penalty_t20,
        "higher T should give lower penalty for same deficit; T=10 → {penalty_t10}, T=20 → {penalty_t20}"
    );
    assert!(
        (penalty_t20 - 1.8).abs() < 1e-5,
        "hold=0,target=30,T=20,max=3 should give penalty=1.8, got {penalty_t20}"
    );
}

// ─── Lump-sum opportunity cost (Change 3 — event-driven density) ───────────

/// Standard exit: `scale=1, rate=0.001, position=0.5, hold_time=20` →
/// `r_opp_cost = -0.001 × 0.5 × 20 = -0.01`.
#[test]
#[ignore = "requires GPU"]
fn sp12_lump_sum_opp_cost_at_exit() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.shaping_scale     = 1.0;
    inputs.holding_cost_rate = 0.001;
    inputs.position          = 0.5;
    inputs.hold_time_oc      = 20.0;
    let (_, _, opp_cost) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (opp_cost - (-0.01)).abs() < 1e-6,
        "expected -0.01 (= -1 * 0.001 * 0.5 * 20), got {opp_cost}"
    );
}

/// Zero position → zero cost regardless of `hold_time`.
#[test]
#[ignore = "requires GPU"]
fn sp12_lump_sum_opp_cost_zero_position() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.shaping_scale     = 1.0;
    inputs.holding_cost_rate = 0.001;
    inputs.position          = 0.0;
    inputs.hold_time_oc      = 20.0;
    let (_, _, opp_cost) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        opp_cost.abs() < 1e-6,
        "zero position should give zero cost, got {opp_cost}"
    );
}

/// Zero hold_time → zero cost regardless of `position`.
#[test]
#[ignore = "requires GPU"]
fn sp12_lump_sum_opp_cost_zero_hold_time() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.shaping_scale     = 1.0;
    inputs.holding_cost_rate = 0.001;
    inputs.position          = 0.5;
    inputs.hold_time_oc      = 0.0;
    let (_, _, opp_cost) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        opp_cost.abs() < 1e-6,
        "zero hold_time should give zero cost, got {opp_cost}"
    );
}

/// Negative position uses `|position|`: same magnitude as positive case.
#[test]
#[ignore = "requires GPU"]
fn sp12_lump_sum_opp_cost_negative_position_uses_abs() {
    let stream = make_test_stream();
    let f = load_test_kernel(&stream);
    let mut inputs = Sp12RewardInputs::defaults();
    inputs.shaping_scale     = 1.0;
    inputs.holding_cost_rate = 0.001;
    inputs.position          = -0.5;
    inputs.hold_time_oc      = 20.0;
    let (_, _, opp_cost) = launch_sp12_reward_math(&stream, &f, inputs);
    assert!(
        (opp_cost - (-0.01)).abs() < 1e-6,
        "|-0.5| should match +0.5 magnitude (-0.01), got {opp_cost}"
    );
}