From 333a00fc07ed929dda081f78b591667cd18a3031 Mon Sep 17 00:00:00 2001
From: Joseph Lizier <joseph.lizier@gmail.com>
Date: Tue, 1 Apr 2025 22:55:28 +1100
Subject: [PATCH] Adding method to compute kNN distances for KSG AIS estimator

---
 .../ActiveInfoStorageCalculatorKraskov.java   | 45 +++++++++++++++++++
 .../TransferEntropyCalculatorKraskov.java     |  2 +-
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/java/source/infodynamics/measures/continuous/kraskov/ActiveInfoStorageCalculatorKraskov.java b/java/source/infodynamics/measures/continuous/kraskov/ActiveInfoStorageCalculatorKraskov.java
index 4aa232c..27318df 100755
--- a/java/source/infodynamics/measures/continuous/kraskov/ActiveInfoStorageCalculatorKraskov.java
+++ b/java/source/infodynamics/measures/continuous/kraskov/ActiveInfoStorageCalculatorKraskov.java
@@ -23,6 +23,7 @@ import java.util.Hashtable;
 import infodynamics.measures.continuous.ActiveInfoStorageCalculator;
 import infodynamics.measures.continuous.ActiveInfoStorageCalculatorViaMutualInfo;
 import infodynamics.measures.continuous.MutualInfoCalculatorMultiVariate;
+import infodynamics.utils.MatrixUtils;
 
 /**
  * An Active Information Storage (AIS) calculator (implementing {@link ActiveInfoStorageCalculator})
@@ -230,4 +231,48 @@ public class ActiveInfoStorageCalculatorKraskov
 		}
 	}
 
+	/**
+     * Debug method to return the k nearest neighbour distances that 
+     *  would be utilised for each sample point here.
+     * Note that this is specifically the max-norm across the target-targetPast variables, which is used for each
+     *  range search in algorithm 1 (although algorithm 2 would use the max distance
+     *  for each variable within the kNNs in their separate range searches). 
+     *  
+     * @param startTimePoint
+     * @param numTimePoints
+     * @return
+     * @throws Exception
+     */
+	public double[] kNNDistances(int startTimePoint, int numTimePoints) throws Exception {
+		// Defer the call to the underlying KSG CMI estimator
+		return ((MutualInfoCalculatorMultiVariateKraskov) miCalc).kNNDistances(startTimePoint, numTimePoints);
+	}
+
+	/**
+	 * Debug method to return the k nearest neighbour distances that
+	 *  would be utilised in {@link #computeLocalUsingPreviousObservations(double[])}
+	 *  for a cross AIS.
+     * Note that this is specifically the max-norm across the target-targetPast variables, which is used for each
+     *  range search in algorithm 1 (although algorithm 2 would use the max distance
+     *  for each variable within the kNNs in their separate range searches). 
+	 *  
+	 * @param startTimePoint
+	 * @param numTimePoints
+	 * @param newObservations
+	 * @return
+	 * @throws Exception
+	 */
+	public double[] kNNDistancesForNewSamples(int startTimePoint, int numTimePoints,
+			double[] newObservations) throws Exception {
+		if (newObservations.length - (k-1)*tau - 1 <= 0) {
+			// There are no observations to compute for here
+			throw new Exception("Not enough samples to embed");
+		}
+		double[][] newDestPastVectors = 
+				MatrixUtils.makeDelayEmbeddingVector(newObservations, k, tau, (k-1)*tau, newObservations.length - (k-1)*tau - 1);
+		double[][] newDestNextVectors =
+				MatrixUtils.makeDelayEmbeddingVector(newObservations, 1, (k-1)*tau + 1, newObservations.length - (k-1)*tau - 1);
+		return ((MutualInfoCalculatorMultiVariateKraskov) miCalc).
+				kNNDistancesForNewSamples(startTimePoint, numTimePoints, newDestPastVectors, newDestNextVectors);
+	}
 }
diff --git a/java/source/infodynamics/measures/continuous/kraskov/TransferEntropyCalculatorKraskov.java b/java/source/infodynamics/measures/continuous/kraskov/TransferEntropyCalculatorKraskov.java
index 71be4fe..5a78be7 100755
--- a/java/source/infodynamics/measures/continuous/kraskov/TransferEntropyCalculatorKraskov.java
+++ b/java/source/infodynamics/measures/continuous/kraskov/TransferEntropyCalculatorKraskov.java
@@ -292,7 +292,7 @@ public class TransferEntropyCalculatorKraskov
 		}
 		if (newDestObservations.length < startTimeForFirstDestEmbedding + 2) {
 			// There are no observations to compute for here
-			return new double[newDestObservations.length];
+			throw new Exception("Not enough samples to embed");
 		}
 		// Now embed as per computeLocalUsingPreviousObservations() in super:
 		double[][] newDestPastVectors =