Update tree plotting

Author: giuliabaldini

README.md

@@ -109,7 +109,7 @@ Your `data-path` should be a file with a header containing only the comma-separa
 Assuming that you have run the program and stored the results in a JSON file, you can plot the tree produced by the algorithm.
 
 ```bash
-poetry run python exact_kmeans.plot_tree.py --output-json output.json --plot-folder plots
+poetry run python -m exact_kmeans.plot_tree --output-json output.json --plot-folder plots
 ```
 
 This will create a tree that looks like this:

exact_kmeans/plot_tree.py

@@ -83,7 +83,7 @@ def plot(
         jj = json.load(f)
     plot(
         nodes=jj["processed_cluster_sizes"],
-        optimal_objective=jj["optimal_objective"],
+        optimal_objective=jj["objective"] if "objective" in jj else None,
         filename=args.output_json.name,
         plot_folder=args.plot_folder,
     )