AeroViz 0.1.21__py3-none-any.whl

AeroViz/plot/radar.py

@@ -0,0 +1,184 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.patches import Circle, RegularPolygon
+from matplotlib.path import Path
+from matplotlib.projections import register_projection
+from matplotlib.projections.polar import PolarAxes
+from matplotlib.spines import Spine
+from matplotlib.transforms import Affine2D
+
+from AeroViz.plot.utils import *
+
+__all__ = ['radar']
+
+
+def radar_factory(num_vars, frame='circle'):
+    """
+    Create a radar chart with `num_vars` axes.
+
+    This function creates a RadarAxes projection and registers it.
+
+    Parameters
+    ----------
+    num_vars : int
+        Number of variables for radar chart.
+    frame : {'circle', 'polygon'}
+        Shape of frame surrounding axes.
+
+    """
+    # calculate evenly-spaced axis angles
+    theta = np.linspace(0, 2 * np.pi, num_vars, endpoint=False)
+
+    class RadarTransform(PolarAxes.PolarTransform):
+
+        def transform_path_non_affine(self, path):
+            # Paths with non-unit interpolation steps correspond to gridlines,
+            # in which case we force interpolation (to defeat PolarTransform's
+            # autoconversion to circular arcs).
+            if path._interpolation_steps > 1:
+                path = path.interpolated(num_vars)
+            return Path(self.transform(path.vertices), path.codes)
+
+    class RadarAxes(PolarAxes):
+
+        name = 'radar'
+        PolarTransform = RadarTransform
+
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            # rotate plot such that the first axis is at the top
+            self.set_theta_zero_location('N')
+
+        def fill(self, *args, closed=True, **kwargs):
+            """Override fill so that line is closed by default"""
+            return super().fill(closed=closed, *args, **kwargs)
+
+        def plot(self, *args, **kwargs):
+            """Override plot so that line is closed by default"""
+            lines = super().plot(*args, **kwargs)
+            for line in lines:
+                self._close_line(line)
+
+        def _close_line(self, line):
+            x, y = line.get_data()
+            if x[0] != x[-1]:
+                x = np.append(x, x[0])
+                y = np.append(y, y[0])
+                line.set_data(x, y)
+
+        def set_varlabels(self, labels):
+            self.set_thetagrids(np.degrees(theta), labels)
+
+        @staticmethod
+        def _gen_axes_patch():
+            # The Axes patch must be centered at (0.5, 0.5) and of radius 0.5
+            # in axes coordinates.
+            if frame == 'circle':
+                return Circle((0.5, 0.5), 0.5)
+            elif frame == 'polygon':
+                return RegularPolygon((0.5, 0.5), num_vars,
+                                      radius=.5, edgecolor="k")
+            else:
+                raise ValueError("Unknown value for 'frame': %s" % frame)
+
+        def _gen_axes_spines(self):
+            if frame == 'circle':
+                return super()._gen_axes_spines()
+            elif frame == 'polygon':
+                # spine_type must be 'left'/'right'/'top'/'bottom'/'circle'.
+                spine = Spine(axes=self,
+                              spine_type='circle',
+                              path=Path.unit_regular_polygon(num_vars))
+                # unit_regular_polygon gives a polygon of radius 1 centered at
+                # (0, 0) but we want a polygon of radius 0.5 centered at (0.5,
+                # 0.5) in axes coordinates.
+                spine.set_transform(Affine2D().scale(.5).translate(.5, .5)
+                                    + self.transAxes)
+                return {'polar': spine}
+            else:
+                raise ValueError("Unknown value for 'frame': %s" % frame)
+
+    register_projection(RadarAxes)
+    return theta
+
+
+@set_figure(figsize=(3, 3))
+def radar(data, labels=None, legend_labels=None, **kwargs) -> tuple[plt.Figure, plt.Axes]:
+    """
+    Creates a radar chart based on the provided data.
+
+    Parameters
+    ----------
+    data : list of list
+        A 2D list where each inner list represents a factor, and each element
+        within the inner lists represents a value for a species.
+        Shape: (n_factors, n_species)
+        Example: [[0.88, 0.01, 0.03, ...], [0.07, 0.95, 0.04, ...], ...]
+
+    labels : list, optional
+        A list of strings representing the names of species (variables).
+        If provided, it should have the same length as the number of elements
+        in each inner list of `data`.
+        Example: ['Sulfate', 'Nitrate', 'EC', 'OC1', 'OC2', 'OP', 'CO', 'O3']
+
+    legend_labels : list, optional
+        A list of strings for labeling each factor in the legend.
+        If provided, it should have the same length as the number of inner lists in `data`.
+
+    **kwargs : dict
+        Additional keyword arguments to be passed to the plotting function.
+        This may include 'title' for setting the chart title.
+
+    Returns
+    -------
+    tuple[plt.Figure, plt.Axes]
+        A tuple containing the Figure and Axes objects of the created plot.
+
+    Example
+    -------
+    >>> data = [[0.88, 0.01, 0.03, 0.03, 0.00, 0.06, 0.01, 0.00],
+    >>>         [0.07, 0.95, 0.04, 0.05, 0.00, 0.02, 0.01, 0.00],
+    >>>         [0.01, 0.02, 0.85, 0.19, 0.05, 0.10, 0.00, 0.00],
+    >>>         [0.02, 0.01, 0.07, 0.01, 0.21, 0.12, 0.98, 0.00],
+    >>>         [0.01, 0.01, 0.02, 0.71, 0.74, 0.70, 0.30, 0.20]]
+    >>> labels = ['Sulfate', 'Nitrate', 'EC', 'OC1', 'OC2', 'OP', 'CO', 'O3']
+    >>> fig, ax = radar(data, labels=labels, title='Basecase')
+
+    Note
+    ----
+    The first dimension of `data` represents each factor, while the second
+    dimension represents each species.
+    """
+    theta = radar_factory(np.array(data).shape[1], frame='polygon')
+
+    fig, ax = plt.subplots(subplot_kw=dict(projection='radar'))
+    fig.subplots_adjust(wspace=0.25, hspace=0.20, top=0.80, bottom=0.05, right=0.80)
+
+    colors = ['b', 'r', 'g', 'm', 'y']
+
+    # Plot the four cases from the example data on separate axes
+    for d, color in zip(data, colors):
+        ax.plot(theta, d, color=color)
+        ax.fill(theta, d, facecolor=color, alpha=0.25, label='_nolegend_')
+
+    ax.set_varlabels(labels)
+    ax.set_rgrids([0.2, 0.4, 0.6, 0.8])
+    ax.set(title=kwargs.get('title', ''))
+
+    # add legend relative to top-left plot
+    legend_labels = legend_labels or ('Factor 1', 'Factor 2', 'Factor 3', 'Factor 4', 'Factor 5')
+    legend = ax.legend(legend_labels, loc=(0.95, 0.95), labelspacing=0.1)
+
+    plt.show()
+
+    return fig, ax
+
+
+if __name__ == '__main__':
+    data = [[0.88, 0.01, 0.03, 0.03, 0.00, 0.06, 0.01, 0.00],
+            [0.07, 0.95, 0.04, 0.05, 0.00, 0.02, 0.01, 0.00],
+            [0.01, 0.02, 0.85, 0.19, 0.05, 0.10, 0.00, 0.00],
+            [0.02, 0.01, 0.07, 0.01, 0.21, 0.12, 0.98, 0.00],
+            [0.01, 0.01, 0.02, 0.71, 0.74, 0.70, 0.30, 0.20]]
+
+    fig, ax = radar(data=data, labels=['Sulfate', 'Nitrate', 'EC', 'OC1', 'OC2', 'OP', 'CO', 'O3'], title='Basecase')

AeroViz/plot/regression.py

@@ -0,0 +1,200 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from matplotlib.pyplot import Figure, Axes
+
+from AeroViz.plot.utils import *
+
+__all__ = [
+    'linear_regression',
+    'multiple_linear_regression',
+]
+
+
+@set_figure
+def linear_regression(df: pd.DataFrame,
+                      x: str | list[str],
+                      y: str | list[str],
+                      labels: str | list[str] = None,
+                      ax: Axes | None = None,
+                      diagonal=False,
+                      positive: bool = True,
+                      fit_intercept: bool = True,
+                      **kwargs
+                      ) -> tuple[Figure, Axes]:
+    """
+    Create a scatter plot with regression lines for the given data.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        Input DataFrame containing the data.
+    x : str or list of str
+        Column name(s) for the x-axis variable(s). If a list, only the first element is used.
+    y : str or list of str
+        Column name(s) for the y-axis variable(s).
+    labels : str or list of str, optional
+        Labels for the y-axis variable(s). If None, column names are used as labels. Default is None.
+    ax : Axes, optional
+        Matplotlib Axes object to use for the plot. If None, a new subplot is created. Default is None.
+    diagonal : bool, optional
+        If True, a diagonal line (1:1 line) is added to the plot. Default is False.
+    positive : bool, optional
+        Whether to constrain the regression coefficients to be positive. Default is True.
+    fit_intercept: bool, optional
+        Whether to calculate the intercept for this model. Default is True.
+    **kwargs
+        Additional keyword arguments for plot customization.
+
+    Returns
+    -------
+    fig : Figure
+        The matplotlib Figure object.
+    ax : Axes
+        The matplotlib Axes object with the scatter plot.
+
+    Notes
+    -----
+    - The function creates a scatter plot with optional regression lines.
+    - The regression line is fitted for each y variable.
+    - Customization options are provided via **kwargs.
+
+    Examples
+    --------
+    >>> from AeroViz import plot
+    >>>
+    >>> plot.linear_regression(df, x='X', y=['Y1', 'Y2'], labels=['Label1', 'Label2'],
+    ...                        diagonal=True, xlim=(0, 10), ylim=(0, 20),
+    ...                        xlabel="X-axis", ylabel="Y-axis", title="Scatter Plot with Regressions")
+    """
+    fig, ax = plt.subplots(**kwargs.get('fig_kws', {})) if ax is None else (ax.get_figure(), ax)
+
+    if not isinstance(x, str):
+        x = x[0]
+
+    if not isinstance(y, list):
+        y = [y]
+
+    if labels is None:
+        labels = y
+
+    df = df.dropna(subset=[x, *y])
+    x_array = df[[x]].to_numpy()
+
+    color_cycle = Color.linecolor
+
+    handles, text_list = [], []
+
+    for i, y_var in enumerate(y):
+        y_array = df[[y_var]].to_numpy()
+
+        color = color_cycle[i % len(color_cycle)]
+
+        scatter = ax.scatter(x_array, y_array, s=25, color=color['face'], edgecolors=color['edge'], alpha=0.8,
+                             label=labels[i])
+        handles.append(scatter)
+
+        text, y_predict, slope = linear_regression_base(x_array, y_array,
+                                                        columns=labels[i],
+                                                        positive=positive,
+                                                        fit_intercept=fit_intercept)
+
+        text_list.append(f'{labels[i]}:\n{text}')
+        plt.plot(x_array, y_predict, linewidth=3, color=color['line'], alpha=1, zorder=3)
+
+    ax.set(xlim=kwargs.get('xlim'), ylim=kwargs.get('ylim'), xlabel=Unit(x), ylabel=Unit(y[0]),
+           title=kwargs.get('title'))
+
+    # Add regression info to the legend
+    leg = plt.legend(handles=handles, labels=text_list, loc='upper left', prop={'weight': 'bold'})
+
+    for text, color in zip(leg.get_texts(), [color['line'] for color in color_cycle]):
+        text.set_color(color)
+
+    if diagonal:
+        ax.axline((0, 0), slope=1., color='k', lw=2, ls='--', alpha=0.5, label='1:1')
+        plt.text(0.97, 0.97, r'$\bf 1:1\ Line$', color='k', ha='right', va='top', transform=ax.transAxes)
+
+    plt.show()
+
+    return fig, ax
+
+
+@set_figure
+def multiple_linear_regression(df: pd.DataFrame,
+                               x: str | list[str],
+                               y: str | list[str],
+                               labels: str | list[str] = None,
+                               ax: Axes | None = None,
+                               diagonal=False,
+                               positive: bool = True,
+                               fit_intercept: bool = True,
+                               **kwargs
+                               ) -> tuple[Figure, Axes]:
+    """
+    Perform multiple linear regression analysis and plot the results.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        Input DataFrame containing the data.
+    x : str or list of str
+        Column name(s) for the independent variable(s). Can be a single string or a list of strings.
+    y : str or list of str
+        Column name(s) for the dependent variable(s). Can be a single string or a list of strings.
+    labels : str or list of str, optional
+        Labels for the dependent variable(s). If None, column names are used as labels. Default is None.
+    ax : Axes, optional
+        Matplotlib Axes object to use for the plot. If None, a new subplot is created. Default is None.
+    diagonal : bool, optional
+        Whether to include a diagonal line (1:1 line) in the plot. Default is False.
+    positive : bool, optional
+        Whether to constrain the regression coefficients to be positive. Default is True.
+    fit_intercept: bool, optional
+        Whether to calculate the intercept for this model. Default is True.
+    **kwargs
+        Additional keyword arguments for plot customization.
+
+    Returns
+    -------
+    tuple[Figure, Axes]
+        The Figure and Axes containing the regression plot.
+
+    Notes
+    -----
+    This function performs multiple linear regression analysis using the input DataFrame.
+    It supports multiple independent variables and can plot the regression results.
+
+    Examples
+    --------
+    >>> from AeroViz import plot
+    >>>
+    >>> plot.multiple_linear_regression(df, x=['X1', 'X2'], y='Y', labels=['Y1', 'Y2'],
+    ...                                 diagonal=True, fit_intercept=True,
+    ...                                 xlabel="X-axis", ylabel="Y-axis", title="Multiple Linear Regression Plot")
+    """
+    fig, ax = plt.subplots(**kwargs.get('fig_kws', {})) if ax is None else (ax.get_figure(), ax)
+
+    if not isinstance(x, list):
+        x = [x]
+
+    if not isinstance(y, str):
+        y = y[0]
+
+    if labels is None:
+        labels = x
+
+    df = df[[*x, y]].dropna()
+    x_array = df[[*x]].to_numpy()
+    y_array = df[[y]].to_numpy()
+
+    text, y_predict, coefficients = linear_regression_base(x_array, y_array,
+                                                           columns=labels,
+                                                           positive=positive,
+                                                           fit_intercept=fit_intercept)
+
+    df = pd.DataFrame(np.concatenate([y_array, y_predict], axis=1), columns=['y_actual', 'y_predict'])
+
+    linear_regression(df, x='y_actual', y='y_predict', ax=ax, regression=True, diagonal=diagonal)
+
+    return fig, ax

AeroViz/plot/scatter.py

@@ -0,0 +1,174 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from matplotlib.colors import Normalize
+from matplotlib.pyplot import Figure, Axes
+from matplotlib.ticker import ScalarFormatter
+
+from AeroViz.plot.utils import *
+
+__all__ = ['scatter']
+
+
+def check_empty(*arrays):
+    for i, arr in enumerate(arrays):
+        if arr.size == 0:
+            raise ValueError(f"Array is empty!")
+
+
+@set_figure
+def scatter(df: pd.DataFrame,
+            x: str,
+            y: str,
+            c: str | None = None,
+            color: str | None = '#7a97c9',
+            s: str | None = None,
+            cmap='jet',
+            regression=False,
+            regression_line_color: str | None = sns.xkcd_rgb["denim blue"],
+            diagonal=False,
+            ax: Axes | None = None,
+            **kwargs
+            ) -> tuple[Figure, Axes]:
+    """
+    Creates a scatter plot with optional color and size encoding.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        The DataFrame containing the data to plot.
+    x : str
+        The column name for the x-axis values.
+    y : str
+        The column name for the y-axis values.
+    c : str, optional
+        The column name for c encoding. Default is None.
+    color : str, optional
+        The column name for color encoding. Default is None.
+    s : str, optional
+        The column name for size encoding. Default is None.
+    cmap : str, optional
+        The colormap to use for the color encoding. Default is 'jet'.
+    regression : bool, optional
+        If True, fits and plots a linear regression line. Default is False.
+    regression_line_color : str, optional
+        The color of the regression line. Default is 'sns.xkcd_rgb["denim blue"]'.
+    diagonal : bool, optional
+        If True, plots a 1:1 diagonal line. Default is False.
+    ax : Axes, optional
+        The matplotlib Axes to plot on. If not provided, a new figure and axes are created.
+    **kwargs : Any
+        Additional keyword arguments passed to customize the plot, such as `fig_kws` for figure creation and `xlabel`,
+        `ylabel`, `xlim`, `ylim`, `title` for axis labeling and limits.
+
+    Returns
+    -------
+    fig : Figure
+        The matplotlib Figure object.
+    ax : Axes
+        The matplotlib Axes object with the scatter plot.
+
+    Notes
+    -----
+    - If both `c` and `s` are provided, the scatter plot will encode data points using both color and size.
+    - If only `c` is provided, data points will be color-coded according to the values in the `c` column.
+    - If only `s` is provided, data points will be sized according to the values in the `s` column.
+    - If neither `c` nor `s` is provided, a basic scatter plot is created.
+    - The `regression` option will add a linear regression line and display the equation on the plot.
+    - The `diagonal` option will add a 1:1 reference line to the plot.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> from AeroViz.plot import scatter
+    >>> df = pd.DataFrame({
+    >>>     'x': [1, 2, 3, 4],
+    >>>     'y': [1.1, 2.0, 2.9, 4.1],
+    >>>     'color': [10, 20, 30, 40],
+    >>>     'size': [100, 200, 300, 400]
+    >>> })
+    >>> fig, ax = scatter(df, x='x', y='y', c='color', s='size', regression=True, diagonal=True)
+    """
+    fig, ax = plt.subplots(**kwargs.get('fig_kws', {})) if ax is None else (ax.get_figure(), ax)
+
+    if c is not None and s is not None:
+        df_ = df.dropna(subset=[x, y, c, s]).copy()
+        x_data, y_data, c_data, s_data = df_[x].to_numpy(), df_[y].to_numpy(), df_[c].to_numpy(), df_[s].to_numpy()
+        check_empty(x_data, y_data, c_data, s_data)
+
+        scatter = ax.scatter(x_data, y_data, c=c_data,
+                             norm=Normalize(vmin=np.percentile(c_data, 10), vmax=np.percentile(c_data, 90)),
+                             cmap=cmap, s=50 * (s_data / s_data.max()) ** 1.5, alpha=0.7, edgecolors=None)
+        colorbar = True
+
+        dot = np.linspace(s_data.min(), s_data.max(), 6).round(-1)
+
+        for dott in dot[1:-1]:
+            plt.scatter([], [], c='k', alpha=0.8, s=50 * (dott / s_data.max()) ** 1.5, label='{:.0f}'.format(dott))
+
+        plt.legend(title=Unit(s))
+
+    elif c is not None:
+        df_ = df.dropna(subset=[x, y, c]).copy()
+        x_data, y_data, c_data = df_[x].to_numpy(), df_[y].to_numpy(), df_[c].to_numpy()
+        check_empty(x_data, y_data, c_data)
+
+        scatter = ax.scatter(x_data, y_data, c=c_data, vmin=c_data.min(), vmax=np.percentile(c_data, 90), cmap=cmap,
+                             alpha=0.7,
+                             edgecolors=None)
+        colorbar = True
+
+    elif s is not None:
+        df_ = df.dropna(subset=[x, y, s]).copy()
+        x_data, y_data, s_data = df_[x].to_numpy(), df_[y].to_numpy(), df_[s].to_numpy()
+        check_empty(x_data, y_data, s_data)
+
+        scatter = ax.scatter(x_data, y_data, s=50 * (s_data / s_data.max()) ** 1.5, color=color, alpha=0.5,
+                             edgecolors='white')
+        colorbar = False
+
+        # dealing
+        dot = np.linspace(s_data.min(), s_data.max(), 6).round(-1)
+
+        for dott in dot[1:-1]:
+            plt.scatter([], [], c='k', alpha=0.8, s=50 * (dott / s_data.max()) ** 1.5, label='{:.0f}'.format(dott))
+
+        plt.legend(title=Unit(s))
+
+    else:
+        df_ = df.dropna(subset=[x, y]).copy()
+        x_data, y_data = df_[x].to_numpy(), df_[y].to_numpy()
+        check_empty(x_data, y_data)
+
+        scatter = ax.scatter(x_data, y_data, s=30, color=color, alpha=0.5, edgecolors='white')
+        colorbar = False
+
+    ax.set(xlim=kwargs.get('xlim', (x_data.min(), x_data.max())),
+           ylim=kwargs.get('ylim', (y_data.min(), y_data.max())),
+           xlabel=kwargs.get('xlabel', Unit(x)),
+           ylabel=kwargs.get('ylabel', Unit(y)),
+           title=kwargs.get('title', ''))
+
+    ax.xaxis.set_major_formatter(ScalarFormatter())
+    ax.yaxis.set_major_formatter(ScalarFormatter())
+
+    if colorbar:
+        plt.colorbar(scatter, extend='both', label=Unit(c))
+
+    if regression:
+        text, y_predict, slope = linear_regression_base(x_data, y_data)
+        ax.plot(x_data, y_predict, linewidth=3, color=regression_line_color, alpha=1, zorder=3)
+        plt.text(0.05, 0.95, text, fontdict={'weight': 'bold'}, color=regression_line_color,
+                 ha='left', va='top', transform=ax.transAxes)
+
+    if diagonal:
+        ax.axline((0, 0), slope=1., color='k', lw=2, ls='--', alpha=0.5, label='1:1')
+
+        data_range = min(ax.get_xlim()[1] - ax.get_xlim()[0], ax.get_ylim()[1] - ax.get_ylim()[0])
+        plt.text(0.9 * data_range, 0.9 * data_range, r'$\bf 1:1\ Line$', color='k', ha='left', va='bottom',
+                 bbox=dict(facecolor='white', edgecolor='none', alpha=0.1, pad=3))
+
+    plt.show()
+
+    return fig, ax

AeroViz/plot/templates/__init__.py

@@ -0,0 +1,6 @@
+from .ammonium_rich import ammonium_rich
+from .contour import *
+from .corr_matrix import corr_matrix, cross_corr_matrix
+from .diurnal_pattern import *
+from .koschmieder import *
+from .metal_heatmap import metal_heatmaps, process_data_with_two_df

AeroViz/plot/templates/ammonium_rich.py

@@ -0,0 +1,34 @@
+import matplotlib.pyplot as plt
+from matplotlib.pyplot import Figure, Axes
+from pandas import DataFrame
+
+from AeroViz.plot.utils import set_figure, Unit
+
+
+@set_figure(figsize=(5, 4))
+def ammonium_rich(df: DataFrame,
+                  **kwargs
+                  ) -> tuple[Figure, Axes]:
+    df = df[['NH4+', 'SO42-', 'NO3-', 'PM2.5']].dropna().copy().div([18, 96, 62, 1])
+    df['required_ammonium'] = df['NO3-'] + 2 * df['SO42-']
+
+    fig, ax = plt.subplots()
+
+    scatter = ax.scatter(df['required_ammonium'].to_numpy(), df['NH4+'].to_numpy(), c=df['PM2.5'].to_numpy(),
+                         vmin=0, vmax=70, cmap='jet', marker='o', s=10, alpha=1)
+
+    ax.axline((0, 0), slope=1., color='k', lw=2, ls='--', alpha=0.5, label='1:1')
+    plt.text(0.97, 0.97, r'$\bf 1:1\ Line$', color='k', ha='right', va='top', transform=ax.transAxes)
+
+    ax.set(xlim=(0, 1.2),
+           ylim=(0, 1.2),
+           xlabel=r'$\bf NO_{3}^{-}\ +\ 2\ \times\ SO_{4}^{2-}\ (mole\ m^{-3})$',
+           ylabel=r'$\bf NH_{4}^{+}\ (mole\ m^{-3})$',
+           title=kwargs.get('title', ''))
+
+    color_bar = plt.colorbar(scatter, label=Unit('PM2.5'), extend='both')
+
+    # fig.savefig(f'Ammonium_rich_{title}')
+    plt.show()
+
+    return fig, ax

AeroViz/plot/templates/contour.py

@@ -0,0 +1,47 @@
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.pyplot import Figure, Axes
+from scipy.optimize import curve_fit
+
+from AeroViz.plot.utils import *
+
+__all__ = ['contour']
+
+
+@set_figure
+def contour(df, ax: Axes | None = None, **kwargs) -> tuple[Figure, Axes]:
+    fig, ax = plt.subplots(**kwargs.get('fig_kws', {})) if ax is None else (ax.get_figure(), ax)
+
+    npoints = 1000
+    xreg = np.linspace(df.PM25.min(), df.PM25.max(), 83)
+    yreg = np.linspace(df.gRH.min(), df.gRH.max(), 34)
+    X, Y = np.meshgrid(xreg, yreg)
+
+    d_f = df.copy()
+    df['gRH'] = d_f['gRH'].round(2)
+    df['PM25'] = d_f['PM25'].round(2)
+
+    def func(data, *params):
+        return params[0] * data ** (params[1])
+
+    initial_guess = [1.0, 1.0]
+
+    fit_df = df[['PM25', 'gRH', 'Extinction']].dropna()
+    popt, pcov = curve_fit(func, xdata=(fit_df['PM25'] * fit_df['gRH']), ydata=fit_df['Extinction'], p0=initial_guess,
+                           maxfev=2000000, method='trf')
+
+    x, y = df.PM25, df.gRH
+
+    # pcolor = ax.pcolormesh(X, Y, (X * 4.5 * Y ** (1 / 3)), cmap='jet', shading='auto', vmin=0, vmax=843, alpha=0.8)
+    Z = func(X * Y, *popt)
+    cont = ax.contour(X, Y, Z, colors='black', levels=5, vmin=0, vmax=Z.max())
+    conf = ax.contourf(X, Y, Z, cmap='YlGnBu', levels=100, vmin=0, vmax=Z.max())
+    ax.clabel(cont, colors=['black'], fmt=lambda s: f"{s:.0f} 1/Mm")
+    ax.set(xlabel=Unit('PM25'), ylabel=Unit('gRH'), xlim=(x.min(), x.max()), ylim=(y.min(), y.max()))
+
+    color_bar = plt.colorbar(conf, pad=0.02, fraction=0.05, label='Extinction (1/Mm)')
+    color_bar.ax.set_xticklabels(color_bar.ax.get_xticks().astype(int))
+
+    plt.show()
+
+    return fig, ax