From 07cd14407b9c7e1acbf91d6405291be5b3c0d757 Mon Sep 17 00:00:00 2001
From: NickEdwards7502 <nickedwards7502@gmail.com>
Date: Fri, 20 Sep 2024 14:18:34 +1000
Subject: [PATCH] REFACTOR: Make ExportModel compatible with JsonRFAnalyser
 (#237)

---
 .../au/csiro/variantspark/api/ExportModel.scala      | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/main/scala/au/csiro/variantspark/api/ExportModel.scala b/src/main/scala/au/csiro/variantspark/api/ExportModel.scala
index 7a09a73f..42076a20 100644
--- a/src/main/scala/au/csiro/variantspark/api/ExportModel.scala
+++ b/src/main/scala/au/csiro/variantspark/api/ExportModel.scala
@@ -47,7 +47,7 @@ class ExportModel(rfModel: RandomForestModel, featureSource: FeatureSource) {
     val writer = new BufferedWriter(new OutputStreamWriter(outputStream))
 
     try {
-      writer.write("{\n\"forest\": ")
+      writer.write("{\n\"params\" : ")
       writer.write(writePretty(rfModel.params))
       writer.write(",\n\"trees\": [\n")
 
@@ -83,12 +83,16 @@ class ExportModel(rfModel: RandomForestModel, featureSource: FeatureSource) {
         writer.flush()
       }
 
-      writer.write("\n]}\n")
+      writer.write("\n]")
+      if (rfModel.params.oob == true) {
+        writer.write(", \n\"oobErrors\" : ")
+        writer.write(writePretty(rfModel.oobErrors))
+      }
+      writer.write("\n}\n")
     } finally {
       writer.close()
       outputStream.close()
+      println(s"Model saved successfully to: ${jsonFilename}")
     }
-
-    println(s"Model saved successfully to: ${jsonFilename}")
   }
 }