Combinatorial Bandits for Load Tripping

load_trip.jl

using LinearAlgebra
using Random
using CairoMakie

"""
    ucb_attack(demand::Matrix, k; α = 1.2)

Run top‐k UCB on a demand matrix of size (n, T).
Returns a binary matrix of chosen attacks at each time step.

Params:
- `demand`: A matrix of size (n, T) representing the demand at each load and time step.
- `k`: The number of loads to trip at each time step.
- `α`: Exploration parameter for UCB (default is 1.2).
"""
function ucb_attack(demand::Matrix{<:Real}, k; α = 1.2)
    n, T = size(demand) # n loads, T time steps
    N  = zeros(Int, n) # number of times each load is attacked
    μ  = zeros(Float64, n) # empirical mean of each load
    A  = spzeros(n, T)  # attack record-- A[i,t] = 1 means load i is attacked at time t, meaning that a_i is set to 0
    for t in 1:T
        ucb = μ .+ α .* sqrt.(2log(t) ./ max.(1, N))
        idx = partialsortperm(ucb, 1:k; rev = true)   # top-k indices
        A[idx, t] .= 1.0   # trip loads
        # empirical‐mean update for chosen loads
        for i in idx
            N[i] += 1
            μ[i] += (demand[i, t] - μ[i]) / N[i]
        end
    end
    return A
end

"""
Calculate the total flow under the tripping attacks.
Returns a vector of total flow at each time step.
The flow is the sum of generation minus demand at each time step.
If a load is tripped, its demand is set to 0.
This is a simplified model and does not consider the actual flow in the network.
"""
function calc_tripped_flow(attacks, demand::Matrix{<:Real},generation::Matrix{<:Real})
    n, T = size(demand)
    tripped_flow = zeros(T)
    for t in 1:T
        a_t = ones(n) .- attacks[:, t] #WARNING: note the negation here
        p_t = generation[:, t] - a_t .* demand[:, t]
        tripped_flow[t] = sum(p_t)
    end
    return tripped_flow
end

"""
Simulate a covert tripping experiment.
This function runs the UCB attack and calculates the total flow under the tripping attacks.
Returns the attack matrix, total flow under attacks, and best case total flow.
"""
function run_covert_tripping_experiment(k;demand=demand,generation=generation)
    #----- find the maximum mean demands for computing the regret
    μ_d = mean(demand; dims = 2)[:,1] # mean demand for each load
    idx_TopK_μ = partialsortperm(μ_d,1:k; rev = true)  # top-k loads by mean demand
    demand_best = deepcopy(demand)  # copy of demand for best case 
    # note: "best" means from the perspective of the attacker, not the system operator
    for i=1:n
        if i ∈ idx_TopK_μ
            demand_best[i, :] .= 0.0  # set top-k loads to 0 in best case
        end
    end
    P_best = generation .- demand_best
    S_best = sum(P_best, dims = 1)'[:,1]  # best case total flow

    #----- run the UCB attack
    @info "Running UCB attack with k = $k"
    @info "Best case total flow: $(sum(S_best))"
    attacks  = ucb_attack(demand, k)
    tripped_flow = calc_tripped_flow(attacks, demand, generation)
    return attacks, tripped_flow, S_best
end

#-----------------------------------------------------------
#----- Begin main script
#-----------------------------------------------------------

#################################################
# REPLACE THIS WITH YOUR OWN DATA
# For now, we will use synthetic data.
####################################################
Random.seed!(42) # for reproducibility
T = 1000 # number of time steps
n = 100 # number of loads and top-k to trip
K_values = [5, 10, 15, 20] # different k values to test
demand   = randn(n, T) .+ 2.5  # synthetic positive demand
generation = randn(n, T) .+ 5.0  # synthetic positive generation
######################################
#----- if using synthetic, make a few demands extremely high to simulate a real‐world scenario
vulnerable_loads = randperm(n)[1:10]  # randomly select 10 loads to be vulnerable
for i in vulnerable_loads
    demand[i, :] .+= 5.0 * rand(T)  # increase demand
end
######################################

with_theme(theme_latexfonts()) do

    # PLOT 1: showing the total flows and attack strategies over time
    fig1 = Figure(size=(800,500), fontsize = 18)
    ax1a = Axis(fig1[1, 1], title = "Cumulative Regret vs. Time",
        xlabel = "Time", ylabel = "Cumulative Regret")

    ax1b = Axis(fig1[2, 1],
        title = "Attack strategies (1 = Tripped, 0 = Not Tripped)",
        xlabel = "Time", ylabel = "Load Index",
        # xticks = 0:50:T
    )

    # PLOT 2: showing the regret over time
    fig2 = Figure()
    ax2a = Axis(fig2[1, 1], title = "Total Flow Over Time",
        xlabel = "Time", ylabel = "Total Flow",
        # xticks = 0:50:T
    )

    # Calculate a SINGLE benchmark S_best with the maximum k value
    max_k = maximum(K_values)
    _, _, S_best_benchmark = run_covert_tripping_experiment(max_k; demand=demand, generation=generation)
    lines!(ax2a, 1:T, S_best_benchmark, label = "Best Case Flow (k=$max_k)", color = :black, linewidth = 2)

    # Set consistent colors for each k value
    colors = Makie.wong_colors()
    #----- run the covert tripping experiment for multiple values of k
    for (i, k) in enumerate(K_values)
        @info "Running experiment for k = $k"
        attacks, tripped_flow, _ = run_covert_tripping_experiment(k; demand=demand, generation=generation)
        
        # Plot regret against the consistent benchmark
        lines!(ax1a, 1:T, cumsum(S_best_benchmark - tripped_flow), 
            label = "(k=$k)", color = colors[i])


        # Plot spy with appropriate colormap
        spy!(ax1b, attacks', colormap = cmaps[i], color = (colors[i],0.3), 
            markersize = 2) 
        
        # Plot flow with consistent color
        lines!(ax2a, 1:T, tripped_flow, label = "Tripped Flow (k=$k)", color = colors[i])
        
        
    end

    # Finalize and display the figures
    axislegend(ax1a, position = :lt,nbanks=2)
    axislegend(ax2a, position = :rt)


    display(fig2)
    display(fig1)
    # png and pdfs
    save("covert_tripping_regret.png", fig1)
    save("covert_tripping_flows.png", fig2)
    save("covert_tripping_regret.pdf", fig1)
    save("covert_tripping_flows.pdf", fig2)
end