feat: add persistence roundtrip examples for all model types

derrickburns · claude · derrickburns · commit 04a9ffc57895 · 2025-10-18T11:35:26.000-07:00
Added executable persistence roundtrip examples with comprehensive assertions: - PersistenceRoundTrip.scala (GeneralizedKMeansModel) - enhanced with assertions - PersistenceRoundTripKMedoids.scala - tests medoid preservation and indices - PersistenceRoundTripSoftKMeans.scala - tests beta, minMembership, probability col - PersistenceRoundTripStreamingKMeans.scala - tests weight preservation and streaming updates All examples follow save/load pattern and verify: - Model parameters roundtrip correctly - Centers preserve correct values - Model-specific state (medoids, weights, etc.) is restored - Predictions work after loading - Special behavior (streaming updates) continues after load Usage for cross-version testing: sbt -Dspark.version=3.4.3 "runMain examples.PersistenceRoundTrip save ./model" sbt -Dspark.version=3.5.1 "runMain examples.PersistenceRoundTrip load ./model" 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
diff --git a/src/main/scala/examples/PersistenceRoundTrip.scala b/src/main/scala/examples/PersistenceRoundTrip.scala
@@ -37,10 +37,23 @@ object PersistenceRoundTrip {
 
       case "load" =>
         val loaded = com.massivedatascience.clusterer.ml.GeneralizedKMeansModel.load(path)
+
+        // Assertions to verify roundtrip correctness
+        assert(loaded.numClusters == 2, s"Expected k=2, got ${loaded.numClusters}")
+        assert(loaded.clusterCenters.length == 2, s"Expected 2 centers, got ${loaded.clusterCenters.length}")
+        assert(loaded.numFeatures == 2, s"Expected dim=2, got ${loaded.numFeatures}")
+
+        // Verify predictions work
         val preds = loaded.transform(df)
         val n = preds.count()
         assert(n == 4, s"expected 4 rows after load, got $n")
-        println(s"Loaded model from $path; predictions=$n")
+
+        // Verify center values are reasonable (should be near (0.5, 0.5) and (9.5, 9.5))
+        val centers = loaded.clusterCenters.sortBy(_.apply(0))
+        assert(math.abs(centers(0)(0) - 0.5) < 1.0, s"Center 0 x-coord should be near 0.5, got ${centers(0)(0)}")
+        assert(math.abs(centers(1)(0) - 9.5) < 1.0, s"Center 1 x-coord should be near 9.5, got ${centers(1)(0)}")
+
+        println(s"✅ Loaded model from $path; predictions=$n; all assertions passed")
       case other =>
         sys.error(s"Unknown mode: $other")
     }
diff --git a/src/main/scala/examples/PersistenceRoundTripKMedoids.scala b/src/main/scala/examples/PersistenceRoundTripKMedoids.scala
@@ -0,0 +1,74 @@
+package examples
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.ml.linalg.Vectors
+import com.massivedatascience.clusterer.ml.KMedoids
+
+/**
+ * Usage:
+ *   sbt -Dspark.version=3.4.3 "runMain examples.PersistenceRoundTripKMedoids save ./tmp_kmedoids_34"
+ *   sbt -Dspark.version=3.5.1 "runMain examples.PersistenceRoundTripKMedoids load ./tmp_kmedoids_34"
+ */
+object PersistenceRoundTripKMedoids {
+  def main(args: Array[String]): Unit = {
+    require(args.length == 2, "args: save|load <path>")
+    val mode = args(0)
+    val path = args(1)
+
+    val spark = SparkSession.builder().appName("PersistenceRoundTripKMedoids").master("local[*]").getOrCreate()
+    import spark.implicits._
+
+    val df = Seq(
+      Tuple1(Vectors.dense(0.0, 0.0)),
+      Tuple1(Vectors.dense(0.1, 0.1)),
+      Tuple1(Vectors.dense(1.0, 1.0)),
+      Tuple1(Vectors.dense(9.0, 9.0)),
+      Tuple1(Vectors.dense(9.1, 9.1)),
+      Tuple1(Vectors.dense(10.0, 10.0))
+    ).toDF("features")
+
+    mode match {
+      case "save" =>
+        val kmedoids = new KMedoids()
+          .setK(2)
+          .setMaxIter(10)
+          .setSeed(456)
+        val model = kmedoids.fit(df)
+        model.write.overwrite().save(path)
+        println(s"Saved KMedoids model to $path")
+        println(s"  Medoids: ${model.medoids.mkString(", ")}")
+        println(s"  Medoid indices: ${model.medoidIndices.mkString(", ")}")
+
+      case "load" =>
+        val loaded = com.massivedatascience.clusterer.ml.KMedoidsModel.load(path)
+
+        // Assertions to verify roundtrip correctness
+        assert(loaded.numClusters == 2, s"Expected k=2, got ${loaded.numClusters}")
+        assert(loaded.medoids.length == 2, s"Expected 2 medoids, got ${loaded.medoids.length}")
+        assert(loaded.medoidIndices.length == 2, s"Expected 2 medoid indices, got ${loaded.medoidIndices.length}")
+        assert(loaded.numFeatures == 2, s"Expected dim=2, got ${loaded.numFeatures}")
+
+        // Verify predictions work
+        val preds = loaded.transform(df)
+        val n = preds.count()
+        assert(n == 6, s"expected 6 rows after load, got $n")
+
+        // Verify medoids are actual data points (one near 0, one near 9-10)
+        val medoids = loaded.medoids.sortBy(_.apply(0))
+        assert(medoids(0)(0) < 2.0, s"Medoid 0 should be near cluster at (0,0), got ${medoids(0)}")
+        assert(medoids(1)(0) > 8.0, s"Medoid 1 should be near cluster at (9-10,9-10), got ${medoids(1)}")
+
+        // Verify medoid indices are valid
+        assert(loaded.medoidIndices.forall(i => i >= 0 && i < 6), s"Medoid indices should be in [0,5], got ${loaded.medoidIndices.mkString(", ")}")
+
+        println(s"✅ Loaded KMedoids model from $path; predictions=$n; all assertions passed")
+        println(s"  Medoids: ${loaded.medoids.mkString(", ")}")
+        println(s"  Medoid indices: ${loaded.medoidIndices.mkString(", ")}")
+
+      case other =>
+        sys.error(s"Unknown mode: $other")
+    }
+
+    spark.stop()
+  }
+}
diff --git a/src/main/scala/examples/PersistenceRoundTripSoftKMeans.scala b/src/main/scala/examples/PersistenceRoundTripSoftKMeans.scala
@@ -0,0 +1,79 @@
+package examples
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.ml.linalg.Vectors
+import com.massivedatascience.clusterer.ml.SoftKMeans
+
+/**
+ * Usage:
+ *   sbt -Dspark.version=3.4.3 "runMain examples.PersistenceRoundTripSoftKMeans save ./tmp_soft_34"
+ *   sbt -Dspark.version=3.5.1 "runMain examples.PersistenceRoundTripSoftKMeans load ./tmp_soft_34"
+ */
+object PersistenceRoundTripSoftKMeans {
+  def main(args: Array[String]): Unit = {
+    require(args.length == 2, "args: save|load <path>")
+    val mode = args(0)
+    val path = args(1)
+
+    val spark = SparkSession.builder().appName("PersistenceRoundTripSoftKMeans").master("local[*]").getOrCreate()
+    import spark.implicits._
+
+    val df = Seq(
+      Tuple1(Vectors.dense(0.0, 0.0)),
+      Tuple1(Vectors.dense(0.1, 0.1)),
+      Tuple1(Vectors.dense(1.0, 1.0)),
+      Tuple1(Vectors.dense(9.0, 9.0)),
+      Tuple1(Vectors.dense(9.1, 9.1)),
+      Tuple1(Vectors.dense(10.0, 10.0))
+    ).toDF("features")
+
+    mode match {
+      case "save" =>
+        val softKMeans = new SoftKMeans()
+          .setK(2)
+          .setDivergence("kullbackLeibler")
+          .setBeta(2.0)
+          .setMinMembership(0.01)
+          .setSeed(789)
+        val model = softKMeans.fit(df)
+        model.write.overwrite().save(path)
+        println(s"Saved SoftKMeans model to $path")
+        println(s"  Centers: ${model.clusterCenters.mkString(", ")}")
+        println(s"  Beta: ${model.betaValue}")
+
+      case "load" =>
+        val loaded = com.massivedatascience.clusterer.ml.SoftKMeansModel.load(path)
+
+        // Assertions to verify roundtrip correctness
+        assert(loaded.numClusters == 2, s"Expected k=2, got ${loaded.numClusters}")
+        assert(loaded.clusterCenters.length == 2, s"Expected 2 centers, got ${loaded.clusterCenters.length}")
+        assert(loaded.clusterCenters(0).size == 2, s"Expected dim=2, got ${loaded.clusterCenters(0).size}")
+
+        // Verify soft clustering parameters
+        assert(math.abs(loaded.betaValue - 2.0) < 0.001, s"Expected beta=2.0, got ${loaded.betaValue}")
+        assert(math.abs(loaded.minMembershipValue - 0.01) < 0.001, s"Expected minMembership=0.01, got ${loaded.minMembershipValue}")
+
+        // Verify predictions work and include probability column
+        val preds = loaded.transform(df)
+        val n = preds.count()
+        assert(n == 6, s"expected 6 rows after load, got $n")
+
+        // Verify probability column exists
+        assert(preds.columns.contains("probability"), "Expected 'probability' column in predictions")
+
+        // Verify centers are reasonable (one near 0, one near 9-10)
+        val centers = loaded.clusterCenters.sortBy(_.apply(0))
+        assert(centers(0)(0) < 2.0, s"Center 0 should be near cluster at (0,0), got ${centers(0)}")
+        assert(centers(1)(0) > 8.0, s"Center 1 should be near cluster at (9-10,9-10), got ${centers(1)}")
+
+        println(s"✅ Loaded SoftKMeans model from $path; predictions=$n; all assertions passed")
+        println(s"  Centers: ${loaded.clusterCenters.mkString(", ")}")
+        println(s"  Beta: ${loaded.betaValue}, MinMembership: ${loaded.minMembershipValue}")
+
+      case other =>
+        sys.error(s"Unknown mode: $other")
+    }
+
+    spark.stop()
+  }
+}
diff --git a/src/main/scala/examples/PersistenceRoundTripStreamingKMeans.scala b/src/main/scala/examples/PersistenceRoundTripStreamingKMeans.scala
@@ -0,0 +1,95 @@
+package examples
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.ml.linalg.Vectors
+import com.massivedatascience.clusterer.ml.StreamingKMeans
+
+/**
+ * Usage:
+ *   sbt -Dspark.version=3.4.3 "runMain examples.PersistenceRoundTripStreamingKMeans save ./tmp_streaming_34"
+ *   sbt -Dspark.version=3.5.1 "runMain examples.PersistenceRoundTripStreamingKMeans load ./tmp_streaming_34"
+ */
+object PersistenceRoundTripStreamingKMeans {
+  def main(args: Array[String]): Unit = {
+    require(args.length == 2, "args: save|load <path>")
+    val mode = args(0)
+    val path = args(1)
+
+    val spark = SparkSession.builder().appName("PersistenceRoundTripStreamingKMeans").master("local[*]").getOrCreate()
+    import spark.implicits._
+
+    val df1 = Seq(
+      Tuple1(Vectors.dense(0.0, 0.0)),
+      Tuple1(Vectors.dense(0.1, 0.1)),
+      Tuple1(Vectors.dense(1.0, 1.0))
+    ).toDF("features")
+
+    val df2 = Seq(
+      Tuple1(Vectors.dense(9.0, 9.0)),
+      Tuple1(Vectors.dense(9.1, 9.1)),
+      Tuple1(Vectors.dense(10.0, 10.0))
+    ).toDF("features")
+
+    mode match {
+      case "save" =>
+        val streamingKMeans = new StreamingKMeans()
+          .setK(2)
+          .setDivergence("squaredEuclidean")
+          .setDecayFactor(0.9)
+          .setSmoothing(1e-9)
+          .setSeed(42)
+
+        // Initialize model with first batch
+        val model1 = streamingKMeans.fit(df1)
+        println(s"After batch 1 - Centers: ${model1.clusterCenters.mkString(", ")}")
+        println(s"After batch 1 - Weights: ${model1.currentWeights.mkString(", ")}")
+
+        // Simulate streaming update with second batch
+        val model2 = model1.update(df2)
+        println(s"After batch 2 - Centers: ${model2.clusterCenters.mkString(", ")}")
+        println(s"After batch 2 - Weights: ${model2.currentWeights.mkString(", ")}")
+
+        model2.write.overwrite().save(path)
+        println(s"Saved StreamingKMeans model to $path")
+
+      case "load" =>
+        val loaded = com.massivedatascience.clusterer.ml.StreamingKMeansModel.load(path)
+
+        // Assertions to verify roundtrip correctness
+        assert(loaded.numClusters == 2, s"Expected k=2, got ${loaded.numClusters}")
+        assert(loaded.clusterCenters.length == 2, s"Expected 2 centers, got ${loaded.clusterCenters.length}")
+        assert(loaded.numFeatures == 2, s"Expected dim=2, got ${loaded.numFeatures}")
+
+        // Verify streaming-specific parameters
+        assert(loaded.divergenceName == "squaredEuclidean", s"Expected squaredEuclidean divergence, got ${loaded.divergenceName}")
+        assert(math.abs(loaded.decayFactorValue - 0.9) < 0.001, s"Expected decayFactor=0.9, got ${loaded.decayFactorValue}")
+        assert(math.abs(loaded.smoothingValue - 1e-9) < 1e-10, s"Expected smoothing=1e-9, got ${loaded.smoothingValue}")
+
+        // CRITICAL: Verify cluster weights were restored (essential for streaming!)
+        val currentWeights = loaded.currentWeights
+        assert(currentWeights.length == 2, s"Expected 2 cluster weights, got ${currentWeights.length}")
+        assert(currentWeights.forall(_ > 0), s"Cluster weights should be positive, got ${currentWeights.mkString(", ")}")
+        println(s"Cluster weights restored: ${currentWeights.mkString(", ")}")
+
+        // Verify predictions work
+        val preds = loaded.transform(df1)
+        val n = preds.count()
+        assert(n == 3, s"expected 3 rows after load, got $n")
+
+        // Verify we can continue streaming after load
+        val continued = loaded.update(df2)
+        assert(continued.clusterCenters.length == 2, "Should be able to continue streaming after load")
+        println(s"After continued update - Centers: ${continued.clusterCenters.mkString(", ")}")
+        println(s"After continued update - Weights: ${continued.currentWeights.mkString(", ")}")
+
+        println(s"✅ Loaded StreamingKMeans model from $path; predictions=$n; all assertions passed")
+        println(s"  Centers: ${loaded.clusterCenters.mkString(", ")}")
+        println(s"  Weights: ${loaded.currentWeights.mkString(", ")}")
+
+      case other =>
+        sys.error(s"Unknown mode: $other")
+    }
+
+    spark.stop()
+  }
+}
