From 96f8f46e8f1393b35661e918018f624920fc5939 Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Thu, 21 Nov 2024 22:38:57 +0000
Subject: [PATCH] Demonstrate testing the LDK scorer

Note that this requires LDK PR #3420
---
 src/main.rs | 74 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 71 insertions(+), 3 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 0d82fc1..abf059b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,12 +6,22 @@ use lightning::routing::gossip::{NetworkGraph, NodeId, ReadOnlyNetworkGraph};
 /// The simulation state.
 ///
 /// You're free to put whatever you want here.
-pub struct State {
+pub struct State<'a> {
+	/// As a demonstration, the default run calculates the probability the LDK historical model
+	/// assigns to results.
+	scorer: lightning::routing::scoring::ProbabilisticScorer<&'a NetworkGraph<DevNullLogger>, DevNullLogger>,
+	// We demonstrate calculating log-loss of the LDK historical model
+	loss_sum: f64,
+	result_count: u64,
 }
 
 /// Creates a new [`State`] before any probe results are processed.
 pub fn do_setup<'a>(graph: &'a NetworkGraph<DevNullLogger>) -> State {
-	State {}
+	State {
+		scorer: lightning::routing::scoring::ProbabilisticScorer::new(Default::default(), graph, internal::DevNullLogger),
+		loss_sum: 0.0,
+		result_count: 0,
+	}
 }
 
 /// Processes one probe result.
@@ -19,13 +29,71 @@ pub fn do_setup<'a>(graph: &'a NetworkGraph<DevNullLogger>) -> State {
 /// The network graph as it existed at `result.timestamp` is provided, as well as a reference to
 /// the current state.
 pub fn process_probe_result(network_graph: ReadOnlyNetworkGraph, result: ProbeResult, state: &mut State) {
+	// Update the model's time
+	use lightning::routing::scoring::ScoreUpdate;
+	let cur_time = std::time::Duration::from_secs(result.timestamp);
+	state.scorer.time_passed(cur_time);
 
+	// Evaluate the model
+	for hop in result.channels_with_sufficient_liquidity.iter() {
+		// You can get additional information about the channel from the network_graph:
+		let _chan = network_graph.channels().get(&hop.short_channel_id).unwrap();
+		let mut model_probability =
+			state.scorer.historical_estimated_payment_success_probability(hop.short_channel_id, &hop.dst_node_id, hop.amount_msat, &Default::default())
+			.unwrap_or_else(||
+				state.scorer
+					.live_estimated_payment_success_probability(hop.short_channel_id, &hop.dst_node_id, hop.amount_msat, &Default::default())
+					.expect("We should have some estimated probability, even without history data")
+			);
+		if model_probability < 0.01 { model_probability = 0.01; }
+		state.loss_sum -= model_probability.log2();
+		state.result_count += 1;
+	}
+	if let Some(hop) = &result.channel_that_rejected_payment {
+		// You can get additional information about the channel from the network_graph:
+		let _chan = network_graph.channels().get(&hop.short_channel_id).unwrap();
+		let mut model_probability =
+			state.scorer.historical_estimated_payment_success_probability(hop.short_channel_id, &hop.dst_node_id, hop.amount_msat, &Default::default())
+			.unwrap_or_else(||
+				state.scorer
+					.live_estimated_payment_success_probability(hop.short_channel_id, &hop.dst_node_id, hop.amount_msat, &Default::default())
+					.expect("We should have some estimated probability, even without history data")
+			);
+		if model_probability > 0.99 { model_probability = 0.99; }
+		state.loss_sum -= (1.0 - model_probability).log2();
+		state.result_count += 1;
+	}
+
+	// Update the model with the information we learned
+	let mut path = lightning::routing::router::Path {
+		hops: Vec::new(),
+		blinded_tail: None,
+	};
+	let mut hops = result.channels_with_sufficient_liquidity.iter().chain(result.channel_that_rejected_payment.iter()).peekable();
+	while hops.peek().is_some() {
+		// Sadly, we have to munge the `DirectedChannel` into an LDK `RouteHop`.
+		let hop = hops.next().unwrap();
+		path.hops.push(lightning::routing::router::RouteHop {
+			pubkey: hop.dst_node_id.try_into().unwrap(),
+			node_features: lightning::types::features::NodeFeatures::empty(),
+			short_channel_id: hop.short_channel_id,
+			channel_features: lightning::types::features::ChannelFeatures::empty(),
+			fee_msat: hop.amount_msat - hops.peek().map(|hop| hop.amount_msat).unwrap_or(0),
+			cltv_expiry_delta: 42,
+			maybe_announced_channel: true,
+		});
+	}
+	if let Some(hop) = result.channel_that_rejected_payment {
+		state.scorer.payment_path_failed(&path, hop.short_channel_id, cur_time);
+	} else {
+		state.scorer.payment_path_successful(&path, cur_time);
+	}
 }
 
 /// This is run after all probe results have been processed, and should be used for printing
 /// results or any required teardown.
 pub fn results_complete(state: State) {
-
+	println!("Avg log-loss {}", state.loss_sum / (state.result_count as f64));
 }
 
 /// A hop in a route, consisting of a channel and the source public key, as well as the amount
-- 
2.39.5