add curves, modify saving functions (curve and km details)

examples/advanced/kaplan-meier-he/README.md

@@ -41,12 +41,14 @@ python baseline_kaplan_meier.py
 ```
 By default, this will generate a KM curve image `km_curve_baseline.png` under the current working directory.
 
-Then we run the federated job with simulator
+Then we run a 5-client federated job with simulator 
 ```commandline
-nvflare simulator -w workspace_km_he -n 2 -t 2 jobs/kaplan-meier-he
+nvflare simulator -w workspace_km_he -n 5 -t 5 jobs/kaplan-meier-he
 ```
 By default, this will generate a KM curve image `km_curve_fl.png` under each client's directory.
 
 ## Display Result
 
-By comparing the two curves, we can observe that the two are identical.
+By comparing the two curves, we can observe that the two are identical:
+![KM survival baseline](figs/km_curve_baseline.png)
+![KM survival fl](figs/km_curve_fl.png)

examples/advanced/kaplan-meier-he/jobs/kaplan-meier-he/app/config/config_fed_client.conf

@@ -3,7 +3,7 @@
   format_version = 2
 
   # This is the application script which will be invoked. Client can replace this script with user's own training script.
-  app_script = "km_train.py"
+  app_script = "kaplan_meier_train.py"
 
   # Additional arguments needed by the training code. For example, in lightning, these can be --trainer.batch_size=xxx.
   app_config = ""

examples/advanced/kaplan-meier-he/jobs/kaplan-meier-he/app/config/config_fed_server.conf

@@ -12,7 +12,7 @@
             task_name = "train"
             wf_class_path = "kaplan_meier_wf.KM",
             wf_args {
-                min_clients = 2
+                min_clients = 5
             }
         }
       }

examples/advanced/kaplan-meier-he/jobs/kaplan-meier-he/app/custom/km_train.py → examples/advanced/kaplan-meier-he/jobs/kaplan-meier-he/app/custom/kaplan_meier_train.py

@@ -31,7 +31,7 @@
 np.random.seed(77)
 
 
-def prepare_data(num_of_clients: int = 2, bin_days: int = 7):
+def prepare_data(num_of_clients: int = 5, bin_days: int = 7):
     # Load data
     data_x, data_y = load_veterans_lung_cancer()
     # Get total data count
@@ -57,11 +57,43 @@ def prepare_data(num_of_clients: int = 2, bin_days: int = 7):
     return event_clients, time_clients
 
 
-def save(result: dict):
+def details_save(kmf):
+    # Get the survival function at all observed time points
+    survival_function_at_all_times = kmf.survival_function_
+    # Get the timeline (time points)
+    timeline = survival_function_at_all_times.index.values
+    # Get the KM estimate
+    km_estimate = survival_function_at_all_times["KM_estimate"].values
+    # Get the event count at each time point
+    event_count = kmf.event_table.iloc[:, 0].values  # Assuming the first column is the observed events
+    # Get the survival rate at each time point (using the 1st column of the survival function)
+    survival_rate = 1 - survival_function_at_all_times.iloc[:, 0].values
+    # Return the results
+    results = {
+        "timeline": timeline.tolist(),
+        "km_estimate": km_estimate.tolist(),
+        "event_count": event_count.tolist(),
+        "survival_rate": survival_rate.tolist(),
+    }
     file_path = os.path.join(os.getcwd(), "km_global.json")
-    print(f"save the result to {file_path} \n")
+    print(f"save the details of KM analysis result to {file_path} \n")
     with open(file_path, "w") as json_file:
-        json.dump(result, json_file, indent=4)
+        json.dump(results, json_file, indent=4)
+
+
+def plot_and_save(kmf):
+    # Plot and save the Kaplan-Meier survival curve
+    plt.figure()
+    plt.title("Federated HE")
+    kmf.plot_survival_function()
+    plt.ylim(0, 1)
+    plt.ylabel("prob")
+    plt.xlabel("time")
+    plt.legend("", frameon=False)
+    plt.tight_layout()
+    file_path = os.path.join(os.getcwd(), "km_curve_fl.png")
+    print(f"save the curve plot to {file_path} \n")
+    plt.savefig(file_path)
 
 
 def main():
@@ -164,36 +196,12 @@ def main():
             # Fit the model
             kmf.fit(durations=time_unfold, event_observed=event_unfold)
 
-            # Plot and save the Kaplan-Meier survival curve
-            plt.figure()
-            plt.title("Federated HE")
-            kmf.plot_survival_function()
-            plt.ylim(0, 1)
-            plt.ylabel("prob")
-            plt.xlabel("time")
-            plt.legend("", frameon=False)
-            plt.tight_layout()
-            plt.savefig(os.path.join(os.getcwd(), "km_curve_fl.png"))
-
-            # Save global result to a json file
-            # Get the survival function at all observed time points
-            survival_function_at_all_times = kmf.survival_function_
-            # Get the timeline (time points)
-            timeline = survival_function_at_all_times.index.values
-            # Get the KM estimate
-            km_estimate = survival_function_at_all_times["KM_estimate"].values
-            # Get the event count at each time point
-            event_count = kmf.event_table.iloc[:, 0].values  # Assuming the first column is the observed events
-            # Get the survival rate at each time point (using the 1st column of the survival function)
-            survival_rate = 1 - survival_function_at_all_times.iloc[:, 0].values
-            # Return the results
-            results = {
-                "timeline": timeline.tolist(),
-                "km_estimate": km_estimate.tolist(),
-                "event_count": event_count.tolist(),
-                "survival_rate": survival_rate.tolist(),
-            }
-            save(results)
+            # Plot and save the KM curve
+            print("plot KM curve!!!!!!!!!!!!")
+            plot_and_save(kmf)
+
+            # Save details of the KM result to a json file
+            details_save(kmf)
 
             # Send a simple response to server
             response = FLModel(params={}, params_type=ParamsType.FULL)