* pd_frames.do
*
* Table 2 of "Which H2O frame should you use to explain your model?"
*
* RMSE of the centered partial dependence against the analytic truth, for
* five frames of increasing size, over 15 replications. The model is refit
* in each replication but held fixed across frames within a replication, so
* differences across frames are estimation error and not model differences.
*
* Every curve is interpolated onto one common grid with -ipolate- before any
* error is computed. This matters. -h2omlgraph pdp- takes its grid from the
* range of the predictor in the frame it is given, so scoring each frame on
* its own grid would score small frames on a narrower and easier range.

clear all
set more off
// Number of simulations
local nsim 15
// Number of grid points
local G 20

capture program drop makedata
program define makedata
    args n seed
    clear
    set seed `seed'
    set obs `n'
    gen double x1 = runiform(-3,3)
    gen double x2 = rnormal()
    gen double x3 = 0.9*x2 + sqrt(1-0.9^2)*rnormal()
    gen double x4 = runiform(-2,2)
    gen byte   g  = runiform()<0.5
    gen double x5 = runiform()
    gen byte   x6 = runiform()<0.5
    gen double y  = 4*sin(1.5*x1) + 2*x2 + 3*x4*(2*g-1) + rnormal(0,1.5)
end

capture program drop pdgrid
program define pdgrid, rclass
    args feat lo hi G
    preserve
        quietly use _pd, clear
        quietly keep `feat' mean_response
        quietly rename `feat' gx
        quietly gen byte isgrid = 0
        local n0 = _N
        local ntot = `n0' + `G'
        quietly set obs `ntot'
        forvalues j = 1/`G' {
            local row = `n0' + `j'
            local gv  = `lo' + (`j'-1)*(`hi'-`lo')/(`G'-1)
            quietly replace gx     = `gv' in `row'
            quietly replace isgrid = 1    in `row'
        }
        sort gx
        quietly ipolate mean_response gx, gen(pdi)
        quietly keep if isgrid==1
        sort gx
        matrix C = J(1,`G',.)
        forvalues j = 1/`G' {
            matrix C[1,`j'] = pdi[`j']
        }
        return matrix curve = C
    restore
end

capture program drop rmsevs
program define rmsevs, rclass
    args cname lo hi G feat
    matrix CC = `cname'
    local sc = 0
    local st = 0
    forvalues j = 1/`G' {
        local gv = `lo' + (`j'-1)*(`hi'-`lo')/(`G'-1)
        if "`feat'"=="x1"  local tv = 4*sin(1.5*`gv')
        else               local tv = 0
        local sc = `sc' + CC[1,`j']
        local st = `st' + `tv'
    }
    local mc = `sc'/`G'
    local mt = `st'/`G'
    local ss = 0
    forvalues j = 1/`G' {
        local gv = `lo' + (`j'-1)*(`hi'-`lo')/(`G'-1)
        if "`feat'"=="x1"  local tv = 4*sin(1.5*`gv')
        else               local tv = 0
        local d  = (CC[1,`j'] - `mc') - (`tv' - `mt')
        local ss = `ss' + `d'*`d'
    }
    return scalar rmse = sqrt(`ss'/`G')
end

capture program drop frameN
program define frameN, rclass
    args frm
    quietly h2omlpostestframe `frm'
    quietly h2omlestat metrics
    return scalar N = r(N)
end

h2o init
tempfile res
postfile pf int rep str12 frame str2 feat int nobs double rmse using `res'

forvalues r = 1/`nsim' {
    makedata 4000 `=20260729 + `r''
    quietly _h2oframe put, into(all`r')
    * ---- the idiomatic route: split, then rbind for the unions ----
    quietly _h2oframe split all`r', into(tr`r' va`r' te`r') split(0.6 0.2 0.2) rseed(19)
    quietly _h2oframe rbind tr`r' va`r', into(tv`r')
    quietly _h2oframe split te`r', into(q1_`r' q2_`r' q3_`r' q4_`r') ///
        split(0.25 0.25 0.25 0.25) rseed(19)
    quietly _h2oframe rbind q1_`r' q2_`r', into(h4`r')

    quietly _h2oframe change tr`r'
    quietly h2oml gbregress y x1 x2 x3 x4 g x5 x6, ///
        ntrees(500) maxdepth(12) minobsleaf(1) lrate(0.1) h2orseed(19)

    * h200 = q1, h400 = q1+q2, h800 = te, train = tr, trainvalid = tv
    foreach spec in "q1_ h200" "h4 h400" "te h800" "tr train" "tv trainvalid" {
        local fr  : word 1 of `spec'
        local lab : word 2 of `spec'
        frameN `fr'`r'
        local nn = r(N)
        quietly h2omlgraph pdp x1, frame(`fr'`r') savedata(_pd, replace) nodraw
        pdgrid x1 -2.8 2.8 `G'
        matrix A = r(curve)
        rmsevs A -2.8 2.8 `G' x1
        post pf (`r') ("`lab'") ("x1") (`nn') (r(rmse))
        quietly h2omlgraph pdp x4, frame(`fr'`r') savedata(_pd, replace) nodraw
        pdgrid x4 -1.8 1.8 `G'
        matrix A = r(curve)
        rmsevs A -1.8 1.8 `G' x4
        post pf (`r') ("`lab'") ("x4") (`nn') (r(rmse))
    }
    di "rep `r' done"
}
postclose pf

use `res', clear
save pdres_rbind, replace
di _n "===== realized frame sizes ====="
table frame, statistic(mean nobs) statistic(min nobs) statistic(max nobs) nformat(%6.0f)
di _n "===== PD RMSE on the common grid, `nsim' reps ====="
table frame feat, statistic(mean rmse) nformat(%9.4f)

* --------------------------------------------------------------------------
* Paired comparison against the training frame, within replication.
* The mean excess errors are differences of the column means above, but the
* win counts and the standard deviation of the paired difference are not, and
* the post quotes both. The standard deviation reported for trainvalid is the
* one the post compares the +0.0024 and -0.0009 gaps against.
* --------------------------------------------------------------------------
preserve
    keep rep frame feat rmse
    quietly reshape wide rmse, i(rep feat) j(frame) string
    di _n "===== paired excess error over the TRAINING frame ====="
    foreach f in x1 x4 {
        di "  `f':"
        foreach fr in h200 h400 h800 trainvalid {
            quietly gen double d = rmse`fr' - rmsetrain
            quietly summarize d if feat=="`f'"
            local m = r(mean)
            local s = r(sd)
            quietly count if feat=="`f'" & d>0
            local w = r(N)
            quietly count if feat=="`f'"
            di "     `fr'  excess=" %8.4f `m' "  sd=" %8.4f `s' ///
                "   train better in `w'/" r(N)
            quietly drop d
        }
    }
restore

di "----pd_frames.do successful----"