PWR-HYBRID-3/code/scripts/reach_heatup_pj.jl

#!/usr/bin/env julia
#
# reach_heatup_pj.jl — nonlinear reach on heatup, prompt-jump model.
#
# Reduced from 10-state to 9-state (n is algebraic).  Removes the
# Λ⁻¹ stiffness that capped the full-state reach at ~10 s.  We push
# horizons up: 60 s, 300 s, 1800 s, 5400 s, full T_max = 18000 s (5 hr).
#
# State (10D with augmented time):
#   x[1..6] = C_1..C_6   (delayed-neutron precursors)
#   x[7]    = T_f
#   x[8]    = T_c
#   x[9]    = T_cold
#   x[10]   = t (augmented time, dt/dt = 1)
#
# n is algebraic:  n = Λ · Σ λ_i C_i / (β - ρ),  ρ = K_p · (T_ref - T_c).

using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))

using LinearAlgebra
using ReachabilityAnalysis, LazySets
using JSON
using MAT

# --- Inlined plant constants (must match pke_params) ---
const LAMBDA  = 1e-4
const BETA_1, BETA_2, BETA_3, BETA_4, BETA_5, BETA_6 =
    0.000215, 0.001424, 0.001274, 0.002568, 0.000748, 0.000273
const BETA    = BETA_1 + BETA_2 + BETA_3 + BETA_4 + BETA_5 + BETA_6
const LAM_1, LAM_2, LAM_3, LAM_4, LAM_5, LAM_6 =
    0.0124, 0.0305, 0.111, 0.301, 1.14, 3.01

const P0   = 1e9
const M_F, C_F, M_C, C_C, HA, W_M, M_SG =
    50000.0, 300.0, 20000.0, 5450.0, 5e7, 5000.0, 30000.0

# Note: feedback-linearization in ctrl_heatup_unsat cancels the alpha
# terms exactly, so under closed-loop the effective rho is just Kp·e.
# We don't need ALPHA_F, ALPHA_C in the reach RHS as a result.

const T_COLD0 = 290.0
const DT_CORE = P0 / (W_M * C_C)
const T_HOT0  = T_COLD0 + DT_CORE
const T_C0    = (T_HOT0 + T_COLD0) / 2
const T_F0    = T_C0 + P0 / HA

const T_STANDBY    = T_C0 - 33.333333
const RAMP_RATE_CS = 28.0 / 3600
const KP_HEATUP    = 1e-4

# --- Taylorized PJ heatup RHS, 10D with augmented time ---
@taylorize function rhs_heatup_pj_taylor!(dx, x, p, t)
    # rho_total under closed-loop feedback linearization = Kp · (T_ref - T_c).
    rho = KP_HEATUP * (T_STANDBY + RAMP_RATE_CS * x[10] - x[8])

    # Algebraic prompt-jump n.
    sum_lam_C = LAM_1*x[1] + LAM_2*x[2] + LAM_3*x[3] +
                LAM_4*x[4] + LAM_5*x[5] + LAM_6*x[6]
    denom = BETA - rho
    n = LAMBDA * sum_lam_C / denom
    inv_factor = sum_lam_C / denom

    # Precursor balance under PJ.
    dx[1] = BETA_1 * inv_factor - LAM_1 * x[1]
    dx[2] = BETA_2 * inv_factor - LAM_2 * x[2]
    dx[3] = BETA_3 * inv_factor - LAM_3 * x[3]
    dx[4] = BETA_4 * inv_factor - LAM_4 * x[4]
    dx[5] = BETA_5 * inv_factor - LAM_5 * x[5]
    dx[6] = BETA_6 * inv_factor - LAM_6 * x[6]

    # Thermal — n appears algebraically in fuel eq.
    dx[7] = (P0 * n - HA * (x[7] - x[8])) / (M_F * C_F)
    dx[8] = (HA * (x[7] - x[8]) - 2 * W_M * C_C * (x[8] - x[9])) / (M_C * C_C)
    dx[9] = (2 * W_M * C_C * (x[8] - x[9])) / (M_SG * C_C)

    # Augmented time.
    dx[10] = one(x[1])

    return nothing
end

# --- Build X_entry (PJ form: no n) from predicates.json ---
pred_path = joinpath(@__DIR__, "..", "..", "reachability", "predicates.json")
pred_raw  = JSON.parsefile(pred_path)
entry     = pred_raw["mode_boundaries"]["q_heatup"]["X_entry_polytope"]

n_lo, n_hi           = entry["n_range"]
T_f_lo, T_f_hi       = entry["T_f_range_C"]
T_c_lo, T_c_hi       = entry["T_c_range_C"]
T_cold_lo, T_cold_hi = entry["T_cold_range_C"]

n_mid = 0.5 * (n_lo + n_hi)
C_mid_1 = (BETA_1 / (LAM_1 * LAMBDA)) * n_mid
C_mid_2 = (BETA_2 / (LAM_2 * LAMBDA)) * n_mid
C_mid_3 = (BETA_3 / (LAM_3 * LAMBDA)) * n_mid
C_mid_4 = (BETA_4 / (LAM_4 * LAMBDA)) * n_mid
C_mid_5 = (BETA_5 / (LAM_5 * LAMBDA)) * n_mid
C_mid_6 = (BETA_6 / (LAM_6 * LAMBDA)) * n_mid

# C_i scale linearly with n; sweep across the n_lo..n_hi band.
x_lo = [C_mid_1 * (n_lo / n_mid); C_mid_2 * (n_lo / n_mid);
        C_mid_3 * (n_lo / n_mid); C_mid_4 * (n_lo / n_mid);
        C_mid_5 * (n_lo / n_mid); C_mid_6 * (n_lo / n_mid);
        T_f_lo; T_c_lo; T_cold_lo;
        0.0]
x_hi = [C_mid_1 * (n_hi / n_mid); C_mid_2 * (n_hi / n_mid);
        C_mid_3 * (n_hi / n_mid); C_mid_4 * (n_hi / n_mid);
        C_mid_5 * (n_hi / n_mid); C_mid_6 * (n_hi / n_mid);
        T_f_hi; T_c_hi; T_cold_hi;
        0.0]

X0 = Hyperrectangle(low=x_lo, high=x_hi)

println("\n=== Nonlinear heatup reach, prompt-jump model ===")
println("  X_entry (n-implied range): n ∈ [$(n_lo), $(n_hi)]")
println("  T_c ∈ [$(T_c_lo), $(T_c_hi)] °C")

T_RAMP_END = (T_C0 - T_STANDBY) / RAMP_RATE_CS
println("  T_ramp_end = $(round(T_RAMP_END; digits=0)) s ($(round(T_RAMP_END/60; digits=1)) min)")
println("  Probing horizons up to T_max(heatup) = 18000 s (5 hr).")

# Probe at increasing horizons.  Stop early if any probe fails.
probe_horizons = (60.0, 300.0, 1800.0, 5400.0)

results = Dict{Float64, Any}()
for T_probe in probe_horizons
    println("\n--- Probe T = $T_probe s ($(round(T_probe/60; digits=1)) min) ---")
    sys  = BlackBoxContinuousSystem(rhs_heatup_pj_taylor!, 10)
    prob = InitialValueProblem(sys, X0)
    try
        alg = TMJets(orderT=4, orderQ=2, abstol=1e-9, maxsteps=100000)
        t_start = time()
        sol = solve(prob; T=T_probe, alg=alg)
        elapsed = time() - t_start
        flow = flowpipe(sol)
        n_sets = length(flow)
        println("  TMJets: $n_sets reach-sets, wall-time $(round(elapsed; digits=1)) s")

        flow_hr = overapproximate(flow, Hyperrectangle)
        n_steps = length(flow_hr)
        # Per-timestep envelopes for plotting (saved below).
        t_arr     = zeros(n_steps)
        Tc_lo_ts  = zeros(n_steps); Tc_hi_ts  = zeros(n_steps)
        Tf_lo_ts  = zeros(n_steps); Tf_hi_ts  = zeros(n_steps)
        Tco_lo_ts = zeros(n_steps); Tco_hi_ts = zeros(n_steps)
        n_lo_ts   = zeros(n_steps); n_hi_ts   = zeros(n_steps)
        rho_lo_ts = zeros(n_steps); rho_hi_ts = zeros(n_steps)
        for (k, R) in enumerate(flow_hr)
            s = set(R)
            t_arr[k]    = high(s, 10)  # time-state upper bound ≈ current time
            Tc_lo_ts[k] = low(s, 8);  Tc_hi_ts[k] = high(s, 8)
            Tf_lo_ts[k] = low(s, 7);  Tf_hi_ts[k] = high(s, 7)
            Tco_lo_ts[k] = low(s, 9); Tco_hi_ts[k] = high(s, 9)
            # Algebraic n reconstruction at this reach-set.
            sumLC_lo = LAM_1*low(s,1) + LAM_2*low(s,2) + LAM_3*low(s,3) +
                       LAM_4*low(s,4) + LAM_5*low(s,5) + LAM_6*low(s,6)
            sumLC_hi = LAM_1*high(s,1) + LAM_2*high(s,2) + LAM_3*high(s,3) +
                       LAM_4*high(s,4) + LAM_5*high(s,5) + LAM_6*high(s,6)
            rho_lo_here = KP_HEATUP * (T_STANDBY - high(s, 8))
            rho_hi_here = KP_HEATUP * (T_C0 - low(s, 8))
            rho_lo_ts[k] = rho_lo_here
            rho_hi_ts[k] = rho_hi_here
            denom_lo = BETA - rho_hi_here
            denom_hi = BETA - rho_lo_here
            if denom_lo > 0
                n_lo_ts[k] = LAMBDA * sumLC_lo / denom_hi
                n_hi_ts[k] = LAMBDA * sumLC_hi / denom_lo
            end
        end
        # Also global-envelope values (for backward compat).
        Tc_lo_env = minimum(Tc_lo_ts); Tc_hi_env = maximum(Tc_hi_ts)
        Tf_lo_env = minimum(Tf_lo_ts); Tf_hi_env = maximum(Tf_hi_ts)
        Tcold_lo  = minimum(Tco_lo_ts); Tcold_hi = maximum(Tco_hi_ts)
        n_env_lo = minimum(n_lo_ts); n_env_hi = maximum(n_hi_ts)

        println("  n envelope (reconstructed): [$(round(n_env_lo; sigdigits=4)), $(round(n_env_hi; sigdigits=4))]")
        println("  T_f envelope:               [$(round(Tf_lo_env; digits=2)), $(round(Tf_hi_env; digits=2))] °C")
        println("  T_c envelope:               [$(round(Tc_lo_env; digits=2)), $(round(Tc_hi_env; digits=2))] °C")
        println("  T_cold envelope:            [$(round(Tcold_lo; digits=2)), $(round(Tcold_hi; digits=2))] °C")
        results[T_probe] = (status="OK", n_sets=n_sets, elapsed=elapsed,
                            Tc=(Tc_lo_env, Tc_hi_env),
                            Tf=(Tf_lo_env, Tf_hi_env),
                            Tcold=(Tcold_lo, Tcold_hi),
                            n=(n_env_lo, n_env_hi),
                            t_arr=t_arr,
                            Tc_lo_ts=Tc_lo_ts, Tc_hi_ts=Tc_hi_ts,
                            Tf_lo_ts=Tf_lo_ts, Tf_hi_ts=Tf_hi_ts,
                            Tco_lo_ts=Tco_lo_ts, Tco_hi_ts=Tco_hi_ts,
                            n_lo_ts=n_lo_ts, n_hi_ts=n_hi_ts,
                            rho_lo_ts=rho_lo_ts, rho_hi_ts=rho_hi_ts)
    catch err
        msg = sprint(showerror, err)
        println("  FAILED: ", first(msg, 300))
        results[T_probe] = (status="FAILED", err=first(msg, 300))
        break
    end
end

println("\n=== Summary ===")
for T_probe in probe_horizons
    haskey(results, T_probe) || continue
    r = results[T_probe]
    if r.status == "OK"
        println("  T = $(T_probe) s: OK, $(r.n_sets) reach-sets, $(round(r.elapsed; digits=1))s wall")
    else
        println("  T = $(T_probe) s: FAILED")
    end
end

# Save the longest-successful probe's envelope arrays for the app.
mat_out = joinpath(@__DIR__, "..", "..", "reachability", "reach_heatup_pj_result.mat")
saved = Dict{String, Any}()
saved["probe_horizons"] = collect(probe_horizons)
for T_probe in probe_horizons
    haskey(results, T_probe) || continue
    r = results[T_probe]
    if r.status == "OK"
        pre = "T_$(Int(T_probe))_"
        saved[pre * "Tc_lo"]    = r.Tc[1]
        saved[pre * "Tc_hi"]    = r.Tc[2]
        saved[pre * "Tf_lo"]    = r.Tf[1]
        saved[pre * "Tf_hi"]    = r.Tf[2]
        saved[pre * "Tcold_lo"] = r.Tcold[1]
        saved[pre * "Tcold_hi"] = r.Tcold[2]
        saved[pre * "n_lo"]     = r.n[1]
        saved[pre * "n_hi"]     = r.n[2]
        # Per-timestep envelopes
        saved[pre * "t_arr"]     = r.t_arr
        saved[pre * "Tc_lo_ts"]  = r.Tc_lo_ts
        saved[pre * "Tc_hi_ts"]  = r.Tc_hi_ts
        saved[pre * "Tf_lo_ts"]  = r.Tf_lo_ts
        saved[pre * "Tf_hi_ts"]  = r.Tf_hi_ts
        saved[pre * "Tco_lo_ts"] = r.Tco_lo_ts
        saved[pre * "Tco_hi_ts"] = r.Tco_hi_ts
        saved[pre * "n_lo_ts"]   = r.n_lo_ts
        saved[pre * "n_hi_ts"]   = r.n_hi_ts
        saved[pre * "rho_lo_ts"] = r.rho_lo_ts
        saved[pre * "rho_hi_ts"] = r.rho_hi_ts
    end
end
matwrite(mat_out, saved)
println("\nSaved envelope summaries to $mat_out")