Added plot method to results

src/caf/distribute/gravity_model/core.py

@@ -14,6 +14,7 @@
 import pandas as pd
 from caf.toolkit import cost_utils, io, timing
 from scipy import optimize
+from matplotlib import pyplot as plt
 
 # Local Imports
 # pylint: disable=import-error,wrong-import-position
@@ -60,6 +61,7 @@ class GravityModelResults:
     value_distribution: np.ndarray
 
 
+
 @dataclasses.dataclass
 class GravityModelCalibrateResults(GravityModelResults):
     """A collection of results from a run of the Gravity Model.
@@ -94,6 +96,25 @@ class GravityModelCalibrateResults(GravityModelResults):
     cost_function: cost_functions.CostFunction
     cost_params: dict[str, Any]
 
+    def plot_distributions(self):
+        fig, ax = plt.subplots(figsize=(10, 6))
+        df_1 = self.cost_distribution.df
+        df_1['normalised'] = df_1[self.target_cost_distribution.trips_col] / df_1[self.target_cost_distribution.trips_col].sum()
+        df_2 = self.target_cost_distribution.df
+        df_2['normalised'] = df_2[self.cost_distribution.trips_col] / df_2[
+            self.cost_distribution.trips_col].sum()
+        ax.bar(df_1[self.cost_distribution.avg_col], df_1['normalised'], width=df_1[self.cost_distribution.max_col] - df_1[self.cost_distribution.min_col], label='Achieved Distribution',
+                color='blue', alpha=0.7)
+        ax.bar(df_1[self.cost_distribution.avg_col], df_2['normalised'], width=df_2[self.target_cost_distribution.max_col] - df_2[self.target_cost_distribution.min_col], label='Target Distribution',
+                color='orange', alpha=0.7)
+
+        ax.set_xlabel('Cost')
+        ax.set_ylabel('Trips')
+        ax.set_title('Distribution Comparison')
+        ax.legend()
+
+        return fig
+
 
 @dataclasses.dataclass
 class GravityModelRunResults(GravityModelResults):