[VCF] Add plot to FitResult and small changes in gev tests

Author: viccollado

bluemath_tk/distributions/_base_distributions.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from scipy.optimize import minimize
+import matplotlib.pyplot as plt
 
 from ..core.models import BlueMathModel
 
@@ -53,6 +54,10 @@ def __init__(self, dist, data, res):
         self.nll = res.fun
         self.res = res
 
+        # Auxiliar for diagnostics plots
+        self.n = self.data.shape[0]
+        self.ecdf = np.arange(1, self.n + 1) / (self.n + 1)
+
     def summary(self):
         """
         Print a summary of the fitting results
@@ -66,11 +71,145 @@ def summary(self):
         print(f"Negative Log-Likelihood value: {self.nll:.4f}")
         print(f"{self.message}")
 
-    def plot(self, ax=None, plot_type="hist"):
+    def plot(self, ax=None, plot_type="all"):
         """
-        Plots of fitting results
+        Plots of fitting results: PP-plot, QQ-plot, histogram with fitted distribution, and return period plot.
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes, optional
+            Axes to plot on. If None, a new figure and axes will be created.
+        plot_type : str, optional
+            Type of plot to create. Options are "hist" for histogram, "pp" for P-P plot,
+            "qq" for Q-Q plot, "return_period" for return period plot, or "all" for all plots.
+            Default is "all".
+
+        Returns
+        -------
+        fig : matplotlib.figure.Figure
+            Figure object containing the plots. If `ax` is provided, returns None.
+        
+        Raises
+        -------
+        ValueError
+            If `plot_type` is not one of the valid options ("hist", "pp", "qq", "return_period", "all").
         """
-        pass
+        if plot_type == "all":
+            fig, axs = plt.subplots(2, 2, figsize=(12, 10))
+            self.hist(ax=axs[0, 0])
+            self.pp(ax=axs[0, 1])
+            self.qq(ax=axs[1, 0])
+            self.return_period(ax=axs[1, 1])
+            plt.tight_layout()
+            return fig
+        elif plot_type == "hist":
+            return self.hist()
+        elif plot_type == "pp":
+            return self.pp()
+        elif plot_type == "qq":
+            return self.qq()
+        elif plot_type == "return_period":
+            return self.return_period()
+        else:
+            raise ValueError("Invalid plot type. Use 'hist', 'pp', 'qq', 'return_period', or 'all'.")
+
+    def pp(self, ax=None):
+        """
+        Probability plot of the fitted distribution.
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes, optional
+            Axes to plot on. If None, a new figure and axes will be created.
+        """
+        if ax is None:
+            fig, ax = plt.subplots(figsize=(8, 6))
+        else:
+            fig = None
+
+        probabilities = self.dist.cdf(np.sort(self.data), *self.params)
+        ax.plot([0, 1], [0, 1], color="tab:red", linestyle="--")
+        ax.plot(probabilities, self.ecdf, color="tab:blue", marker="o", linestyle="", alpha=0.7)
+        ax.set_xlabel("Fitted Probability")
+        ax.set_ylabel("Empirical Probability")
+        ax.set_title(f"PP Plot of {self.dist().name}")
+        ax.grid()
+
+        return fig
+
+    def qq(self, ax=None):
+        """
+        Quantile-Quantile plot of the fitted distribution.
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes, optional
+            Axes to plot on. If None, a new figure and axes will be created.
+        """
+        if ax is None:
+            fig, ax = plt.subplots(figsize=(8, 6))
+        else:
+            fig = None
+
+        quantiles = self.dist.qf(self.ecdf, *self.params)
+        ax.plot([np.min(self.data), np.max(self.data)], [np.min(self.data), np.max(self.data)], color="tab:red", linestyle="--")
+        ax.plot(quantiles, np.sort(self.data), color="tab:blue", marker="o", linestyle="", alpha=0.7)
+        ax.set_xlabel("Theoretical Quantiles")
+        ax.set_ylabel("Sample Quantiles")
+        ax.set_title(f"QQ Plot of {self.dist().name}")
+        ax.grid()
+
+        return fig
+
+    def hist(self, ax=None):
+        """
+        Histogram of the data with the fitted distribution overlayed.
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes, optional
+            Axes to plot on. If None, a new figure and axes will be created.
+        """
+        if ax is None:
+            fig, ax = plt.subplots(figsize=(8, 6))
+        else:
+            fig = None
+
+        ax.hist(self.data, bins=30, density=True, alpha=0.7, color='tab:blue', label='Data Histogram')
+        x = np.linspace(np.min(self.data), np.max(self.data), 1000)
+        ax.plot(x, self.dist.pdf(x, *self.params), color='tab:red', label='Fitted PDF')
+        ax.set_xlabel("Data Values")
+        ax.set_ylabel("Density")
+        ax.set_title(f"Histogram and Fitted PDF of {self.dist().name}")
+        ax.legend()
+        ax.grid()
+
+        return fig
+
+    def return_period(self, ax=None):
+        """
+        Return period plot of the fitted distribution.
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes, optional
+            Axes to plot on. If None, a new figure and axes will be created.
+        """
+        if ax is None:
+            fig, ax = plt.subplots(figsize=(8, 6))
+        else:
+            fig = None
+
+
+        sorted_data = np.sort(self.data)
+        exceedance_prob = 1 - self.ecdf
+        return_period = 1 / exceedance_prob
+
+        ax.plot(return_period, self.dist.qf(self.ecdf, *self.params), color='tab:red', label='Fitted Distribution')
+        ax.plot(return_period, sorted_data, marker="o", linestyle="", color="tab:blue", alpha=0.7, label='Empirical Data')
+        ax.set_xscale("log")
+        ax.set_xlabel("Return Period")
+        ax.set_ylabel("Data Values")
+        ax.set_title(f"Return Period Plot of {self.dist().name}")
+        ax.legend()
+        ax.grid()
+
+        return fig
 
 
 def fit_dist(dist, data: np.ndarray, **kwargs) -> FitResult:

tests/distributions/test_gev.py

@@ -105,7 +105,7 @@ def test_invalid_scale(self):
     def test_fit(self):
         # Generate data using specific parameters
         loc, scale, shape = 0.5, 1.5, 0.2
-        data = gev.random(1000, loc, scale, shape)
+        data = gev.random(1000, loc, scale, shape, random_state=42)
 
         # Fit the GEV distribution to the data
         fit_result = gev.fit(data)
@@ -118,7 +118,9 @@ def test_fit(self):
         self.assertIsInstance(fit_result.nll, float)
 
         # Verify that the fitted parameters are close to the original ones
-        np.testing.assert_allclose(fit_result.params, [loc, scale, shape], rtol=0.25)
+        self.assertAlmostEqual(fit_result.params[0], loc, delta=0.1)
+        self.assertAlmostEqual(fit_result.params[1], scale, delta=0.1)
+        self.assertAlmostEqual(fit_result.params[2], shape, delta=0.1)
 
 
 if __name__ == "__main__":