validmind 2.5.8__py3-none-any.whl → 2.5.18__py3-none-any.whl

validmind/tests/data_validation/JarqueBera.py

@@ -0,0 +1,70 @@
+# Copyright © 2023-2024 ValidMind Inc. All rights reserved.
+# See the LICENSE file in the root of this repository for details.
+# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
+
+import pandas as pd
+from statsmodels.stats.stattools import jarque_bera
+
+from validmind import tags, tasks
+
+
+@tasks("classification", "regression")
+@tags("tabular_data", "data_distribution", "statistical_test", "statsmodels")
+def JarqueBera(dataset):
+    """
+    Assesses normality of dataset features in an ML model using the Jarque-Bera test.
+
+    ### Purpose
+
+    The purpose of the Jarque-Bera test as implemented in this metric is to determine if the features in the dataset of
+    a given Machine Learning model follow a normal distribution. This is crucial for understanding the distribution and
+    behavior of the model's features, as numerous statistical methods assume normal distribution of the data.
+
+    ### Test Mechanism
+
+    The test mechanism involves computing the Jarque-Bera statistic, p-value, skew, and kurtosis for each feature in
+    the dataset. It utilizes the 'jarque_bera' function from the 'statsmodels' library in Python, storing the results
+    in a dictionary. The test evaluates the skewness and kurtosis to ascertain whether the dataset follows a normal
+    distribution. A significant p-value (typically less than 0.05) implies that the data does not possess normal
+    distribution.
+
+    ### Signs of High Risk
+
+    - A high Jarque-Bera statistic and a low p-value (usually less than 0.05) indicate high-risk conditions.
+    - Such results suggest the data significantly deviates from a normal distribution. If a machine learning model
+    expects feature data to be normally distributed, these findings imply that it may not function as intended.
+
+    ### Strengths
+
+    - Provides insights into the shape of the data distribution, helping determine whether a given set of data follows
+    a normal distribution.
+    - Particularly useful for risk assessment for models that assume a normal distribution of data.
+    - By measuring skewness and kurtosis, it provides additional insights into the nature and magnitude of a
+    distribution's deviation.
+
+    ### Limitations
+
+    - Only checks for normality in the data distribution. It cannot provide insights into other types of distributions.
+    - Datasets that aren't normally distributed but follow some other distribution might lead to inaccurate risk
+    assessments.
+    - Highly sensitive to large sample sizes, often rejecting the null hypothesis (that data is normally distributed)
+    even for minor deviations in larger datasets.
+    """
+
+    df = dataset.df[dataset.feature_columns_numeric]
+
+    jb_values = {}
+    for col in df.columns:
+        jb_stat, jb_pvalue, jb_skew, jb_kurtosis = jarque_bera(df[col].values)
+        jb_values[col] = {
+            "stat": jb_stat,
+            "pvalue": jb_pvalue,
+            "skew": jb_skew,
+            "kurtosis": jb_kurtosis,
+        }
+
+    jb_df = pd.DataFrame.from_dict(jb_values, orient="index")
+    jb_df.reset_index(inplace=True)
+    jb_df.columns = ["column", "stat", "pvalue", "skew", "kurtosis"]
+
+    return jb_df

validmind/tests/data_validation/KPSS.py

@@ -16,35 +16,40 @@ logger = get_logger(__name__)
 @dataclass
 class KPSS(Metric):
     """
-    Executes KPSS unit root test to validate stationarity of time-series data in machine learning model.
-
-    **Purpose**: The Kwiatkowski-Phillips-Schmidt-Shin (KPSS) unit root test is utilized to ensure the stationarity of
-    data within the machine learning model. It specifically works on time-series data to establish the order of
-    integration, which is a prime requirement for accurate forecasting, given the fundamental condition for any time
-    series model is that the series should be stationary.
-
-    **Test Mechanism**: This metric evaluates the KPSS score for every feature present in the dataset. Within the KPSS
-    score, there are various components, namely: a statistic, a p-value, a used lag, and critical values. The core
-    scheme behind the KPSS score is to test the hypothesis that an observable time series is stationary around a
-    deterministic trend. If the computed statistic surpasses the critical value, the null hypothesis is dismissed,
-    inferring the series is non-stationary.
-
-    **Signs of High Risk**:
-    - High KPSS score represents a considerable risk, particularly if the calculated statistic is higher than the
-    critical value.
-    - If the null hypothesis is rejected and the series is recognized as non-stationary, it heavily influences the
-    model's forecasting capability rendering it less effective.
-
-    **Strengths**:
-    - The KPSS test directly measures the stationarity of a series, allowing it to fulfill a key prerequisite for many
-    time-series models, making it a valuable tool for model validation.
-    - The logics underpinning the test are intuitive and simple, making it understandable and accessible for developers
-    and risk management teams.
-
-    **Limitations**:
-    - The KPSS test presumes the absence of a unit root in the series and does not differentiate between series that
-    are stationary and those border-lining stationarity.
-    - The test might show restricted power against specific alternatives.
+    Assesses the stationarity of time-series data in a machine learning model using the KPSS unit root test.
+
+    ### Purpose
+
+    The KPSS (Kwiatkowski-Phillips-Schmidt-Shin) unit root test is utilized to ensure the stationarity of data within a
+    machine learning model. It specifically works on time-series data to establish the order of integration, which is
+    essential for accurate forecasting. A fundamental requirement for any time series model is that the series should
+    be stationary.
+
+    ### Test Mechanism
+
+    This test calculates the KPSS score for each feature in the dataset. The KPSS score includes a statistic, a
+    p-value, a used lag, and critical values. The core principle behind the KPSS test is to evaluate the hypothesis
+    that an observable time series is stationary around a deterministic trend. If the computed statistic exceeds the
+    critical value, the null hypothesis (that the series is stationary) is rejected, indicating that the series is
+    non-stationary.
+
+    ### Signs of High Risk
+
+    - High KPSS score, particularly if the calculated statistic is higher than the critical value.
+    - Rejection of the null hypothesis, indicating that the series is recognized as non-stationary, can severely affect
+    the model's forecasting capability.
+
+    ### Strengths
+
+    - Directly measures the stationarity of a series, fulfilling a key prerequisite for many time-series models.
+    - The underlying logic of the test is intuitive and simple, making it easy to understand and accessible for both
+    developers and risk management teams.
+
+    ### Limitations
+
+    - Assumes the absence of a unit root in the series and doesn't differentiate between series that are stationary and
+    those border-lining stationarity.
+    - The test may have restricted power against certain alternatives.
     - The reliability of the test is contingent on the number of lags selected, which introduces potential bias in the
     measurement.
     """

validmind/tests/data_validation/LJungBox.py

@@ -0,0 +1,66 @@
+# Copyright © 2023-2024 ValidMind Inc. All rights reserved.
+# See the LICENSE file in the root of this repository for details.
+# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
+
+import pandas as pd
+from statsmodels.stats.diagnostic import acorr_ljungbox
+
+from validmind import tags, tasks
+
+
+@tasks("regression")
+@tags("time_series_data", "forecasting", "statistical_test", "statsmodels")
+def LJungBox(dataset):
+    """
+    Assesses autocorrelations in dataset features by performing a Ljung-Box test on each feature.
+
+    ### Purpose
+
+    The Ljung-Box test is a type of statistical test utilized to ascertain whether there are autocorrelations within a
+    given dataset that differ significantly from zero. In the context of a machine learning model, this test is
+    primarily used to evaluate data utilized in regression tasks, especially those involving time series and
+    forecasting.
+
+    ### Test Mechanism
+
+    The test operates by iterating over each feature within the dataset and applying the `acorr_ljungbox`
+    function from the `statsmodels.stats.diagnostic` library. This function calculates the Ljung-Box statistic and
+    p-value for each feature. These results are then stored in a pandas DataFrame where the columns are the feature names,
+    statistic, and p-value respectively. Generally, a lower p-value indicates a higher likelihood of significant
+    autocorrelations within the feature.
+
+    ### Signs of High Risk
+
+    - High Ljung-Box statistic values or low p-values.
+    - Presence of significant autocorrelations in the respective features.
+    - Potential for negative impact on model performance or bias if autocorrelations are not properly handled.
+
+    ### Strengths
+
+    - Powerful tool for detecting autocorrelations within datasets, especially in time series data.
+    - Provides quantitative measures (statistic and p-value) for precise evaluation.
+    - Helps avoid issues related to autoregressive residuals and other challenges in regression models.
+
+    ### Limitations
+
+    - Cannot detect all types of non-linearity or complex interrelationships among variables.
+    - Testing individual features may not fully encapsulate the dynamics of the data if features interact with each other.
+    - Designed more for traditional statistical models and may not be fully compatible with certain types of complex
+      machine learning models.
+    """
+
+    df = dataset.df
+
+    ljung_box_values = {}
+    for col in df.columns:
+        lb_results = acorr_ljungbox(df[col].values, return_df=True)
+        ljung_box_values[col] = {
+            "stat": lb_results.iloc[0]["lb_stat"],
+            "pvalue": lb_results.iloc[0]["lb_pvalue"],
+        }
+
+    ljung_box_df = pd.DataFrame.from_dict(ljung_box_values, orient="index")
+    ljung_box_df.reset_index(inplace=True)
+    ljung_box_df.columns = ["column", "stat", "pvalue"]
+
+    return ljung_box_df

validmind/tests/data_validation/LaggedCorrelationHeatmap.py

@@ -17,30 +17,37 @@ class LaggedCorrelationHeatmap(Metric):
     Assesses and visualizes correlation between target variable and lagged independent variables in a time-series
     dataset.
 
-    **Purpose**: The LaggedCorrelationHeatmap metric is utilized to appraise and illustrate the correlation between the
-    target variable and delayed copies (lags) of independent variables in a time-series dataset. It assists in
-    revealing relationships in time-series data where the influence of an independent variable on the dependent
-    variable is not immediate but occurs after a period (lags).
-
-    **Test Mechanism**: To execute this test, Python's Pandas library pairs with Plotly to perform computations and
-    present the visualization in the form of a heatmap. The test begins by extracting the target variable and
-    corresponding independent variables from the dataset. Then, generation of lags of independent variables takes
-    place, followed by the calculation of correlation between these lagged variables and the target variable. The
-    outcome is a correlation matrix that gets recorded and illustrated as a heatmap, where different color intensities
-    represent the strength of the correlation, making patterns easier to identify.
-
-    **Signs of High Risk**:
+    ### Purpose
+
+    The LaggedCorrelationHeatmap metric is utilized to appraise and illustrate the correlation between the target
+    variable and delayed copies (lags) of independent variables in a time-series dataset. It assists in revealing
+    relationships in time-series data where the influence of an independent variable on the dependent variable is not
+    immediate but occurs after a period (lags).
+
+    ### Test Mechanism
+
+    To execute this test, Python's Pandas library pairs with Plotly to perform computations and present the
+    visualization in the form of a heatmap. The test begins by extracting the target variable and corresponding
+    independent variables from the dataset. Then, generation of lags of independent variables takes place, followed by
+    the calculation of correlation between these lagged variables and the target variable. The outcome is a correlation
+    matrix that gets recorded and illustrated as a heatmap, where different color intensities represent the strength of
+    the correlation, making patterns easier to identify.
+
+    ### Signs of High Risk
+
     - Insignificant correlations across the heatmap, indicating a lack of noteworthy relationships between variables.
     - Correlations that break intuition or previous understanding, suggesting potential issues with the dataset or the
     model.
 
-    **Strengths**:
+    ### Strengths
+
     - This metric serves as an exceptional tool for exploring and visualizing time-dependent relationships between
     features and the target variable in a time-series dataset.
     - It aids in identifying delayed effects that might go unnoticed with other correlation measures.
     - The heatmap offers an intuitive visual representation of time-dependent correlations and influences.
 
-    **Limitations**:
+    ### Limitations
+
     - The metric presumes linear relationships between variables, potentially ignoring non-linear relationships.
     - The correlation considered is linear; therefore, intricate non-linear interactions might be overlooked.
     - The metric is only applicable for time-series data, limiting its utility outside of this context.

validmind/tests/data_validation/MissingValues.py

@@ -19,34 +19,39 @@ class MissingValues(ThresholdTest):
     """
     Evaluates dataset quality by ensuring missing value ratio across all features does not exceed a set threshold.
 
-    **Purpose**: This test is designed to evaluate the quality of a dataset by measuring the number of missing values
+    ### Purpose
+
+    The Missing Values test is designed to evaluate the quality of a dataset by measuring the number of missing values
     across all features. The objective is to ensure that the ratio of missing data to total data is less than a
-    predefined threshold, defaulting to 1, to maintain the data quality necessary for reliable predictive strength in a
-    machine learning model.
-
-    **Test Mechanism**: The mechanism for this test involves iterating through each column of the dataset, counting
-    missing values (represented as NaNs), and calculating the percentage they represent against the total number of
-    rows. The test then checks if these missing value counts are less than the predefined `min_threshold`. The results
-    are shown in a table summarizing each column, the number of missing values, the percentage of missing values in
-    each column, and a Pass/Fail status based on the threshold comparison.
-
-    **Signs of High Risk**:
-    - When the number of missing values in any column exceeds the `min_threshold` value, it indicates a high risk.
-    - A high risk is also flagged when missing values are present across many columns. In both instances, the test
-    would return a "Fail" mark.
-
-    **Strengths**:
-    - The test offers a quick and granular identification of missing data across each feature in the dataset.
-    - It provides an effective, straightforward means of maintaining data quality, which is vital for constructing
-    efficient machine learning models.
-
-    **Limitations**:
-    - Even though the test can efficiently identify missing values, it does not suggest the root causes of these
-    missing values or recommend ways to impute or handle them.
-    - The test might overlook features with a significant amount of missing data, but still less than the
-    `min_threshold`. This could impact the model, especially if `min_threshold` is set too high.
-    - The test does not account for data encoded as values (like "-999" or "None"), which might not technically
-    classify as missing but could bear similar implications.
+    predefined threshold, defaulting to 1, in order to maintain the data quality necessary for reliable predictive
+    strength in a machine learning model.
+
+    ### Test Mechanism
+
+    The mechanism for this test involves iterating through each column of the dataset, counting missing values
+    (represented as NaNs), and calculating the percentage they represent against the total number of rows. The test
+    then checks if these missing value counts are less than the predefined `min_threshold`. The results are shown in a
+    table summarizing each column, the number of missing values, the percentage of missing values in each column, and a
+    Pass/Fail status based on the threshold comparison.
+
+    ### Signs of High Risk
+
+    - When the number of missing values in any column exceeds the `min_threshold` value.
+    - Presence of missing values across many columns, leading to multiple instances of failing the threshold.
+
+    ### Strengths
+
+    - Quick and granular identification of missing data across each feature in the dataset.
+    - Provides an effective and straightforward means of maintaining data quality, essential for constructing efficient
+    machine learning models.
+
+    ### Limitations
+
+    - Does not suggest the root causes of the missing values or recommend ways to impute or handle them.
+    - May overlook features with significant missing data but still less than the `min_threshold`, potentially
+    impacting the model.
+    - Does not account for data encoded as values like "-999" or "None," which might not technically classify as
+    missing but could bear similar implications.
     """
 
     name = "missing"

validmind/tests/data_validation/MissingValuesBarPlot.py

@@ -12,37 +12,41 @@ from validmind.vm_models import Figure, Metric
 @dataclass
 class MissingValuesBarPlot(Metric):
     """
-    Creates a bar plot showcasing the percentage of missing values in each column of the dataset with risk
-    categorization based on a user-defined threshold.
-
-    **Purpose:** The 'MissingValuesBarPlot' metric provides a color-coded visual representation of the percentage of
-    missing values for each column in an ML model's dataset. The primary purpose of this metric is to easily identify
-    and quantify missing data, which are essential steps in data preprocessing. The presence of missing data can
-    potentially skew the model's predictions and decrease its accuracy. Additionally, this metric uses a pre-set
-    threshold to categorize various columns into ones that contain missing data above the threshold (high risk) and
-    below the threshold (less risky).
-
-    **Test Mechanism:** The test mechanism involves scanning each column in the input dataset and calculating the
-    percentage of missing values. It then compares each column's missing data percentage with the predefined threshold,
-    categorizing columns with missing data above the threshold as high-risk. The test generates a bar plot in which
-    columns with missing data are represented on the y-axis and their corresponding missing data percentages are
-    displayed on the x-axis. The color of each bar reflects the missing data percentage in relation to the threshold:
-    grey for values below the threshold and light coral for those exceeding it. The user-defined threshold is
-    represented by a red dashed line on the plot.
-
-    **Signs of High Risk:**
+    Assesses the percentage and distribution of missing values in the dataset via a bar plot, with emphasis on
+    identifying high-risk columns based on a user-defined threshold.
+
+    ### Purpose
+
+    The 'MissingValuesBarPlot' metric provides a color-coded visual representation of the percentage of missing values
+    for each column in an ML model's dataset. The primary purpose of this metric is to easily identify and quantify
+    missing data, which are essential steps in data preprocessing. The presence of missing data can potentially skew
+    the model's predictions and decrease its accuracy. Additionally, this metric uses a pre-set threshold to categorize
+    various columns into ones that contain missing data above the threshold (high risk) and below the threshold (less
+    risky).
+
+    ### Test Mechanism
+
+    The test mechanism involves scanning each column in the input dataset and calculating the percentage of missing
+    values. It then compares each column's missing data percentage with the predefined threshold, categorizing columns
+    with missing data above the threshold as high-risk. The test generates a bar plot in which columns with missing
+    data are represented on the y-axis and their corresponding missing data percentages are displayed on the x-axis.
+    The color of each bar reflects the missing data percentage in relation to the threshold: grey for values below the
+    threshold and light coral for those exceeding it. The user-defined threshold is represented by a red dashed line on
+    the plot.
+
+    ### Signs of High Risk
 
     - Columns with higher percentages of missing values beyond the threshold are high-risk. These are visually
     represented by light coral bars on the bar plot.
 
-    **Strengths:**
+    ### Strengths
 
     - Helps in quickly identifying and quantifying missing data across all columns of the dataset.
     - Facilitates pattern recognition through visual representation.
     - Enables customization of the level of risk tolerance via a user-defined threshold.
     - Supports both classification and regression tasks, sharing its versatility.
 
-    **Limitations:**
+    ### Limitations
 
     - It only considers the quantity of missing values, not differentiating between different types of missingness
     (Missing completely at random - MCAR, Missing at random - MAR, Not Missing at random - NMAR).

validmind/tests/data_validation/PearsonCorrelationMatrix.py

@@ -2,103 +2,90 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-from dataclasses import dataclass
 
 import plotly.graph_objects as go
 
-from validmind.vm_models import Figure, Metric
+from validmind import tags, tasks
 
 
-@dataclass
-class PearsonCorrelationMatrix(Metric):
+@tags("tabular_data", "numerical_data", "correlation")
+@tasks("classification", "regression")
+def PearsonCorrelationMatrix(dataset):
     """
     Evaluates linear dependency between numerical variables in a dataset via a Pearson Correlation coefficient heat map.
 
-    **Purpose**: This test is intended to evaluate the extent of linear dependency between all pairs of numerical
-    variables in the given dataset. It provides the Pearson Correlation coefficient, which reveals any high
-    correlations present. The purpose of doing this is to identify potential redundancy, as variables that are highly
-    correlated can often be removed to reduce the dimensionality of the dataset without significantly impacting the
-    model's performance.
+    ### Purpose
 
-    **Test Mechanism**: This metric test generates a correlation matrix for all numerical variables in the dataset
-    using the Pearson correlation formula. A heat map is subsequently created to visualize this matrix effectively. The
-    color of each point on the heat map corresponds to the magnitude and direction (positive or negative) of the
-    correlation, with a range from -1 (perfect negative correlation) to 1 (perfect positive correlation). Any
-    correlation coefficients higher than 0.7 (in absolute terms) are indicated in white in the heat map, suggesting a
-    high degree of correlation.
+    This test is intended to evaluate the extent of linear dependency between all pairs of numerical variables in the
+    given dataset. It provides the Pearson Correlation coefficient, which reveals any high correlations present. The
+    purpose of doing this is to identify potential redundancy, as variables that are highly correlated can often be
+    removed to reduce the dimensionality of the dataset without significantly impacting the model's performance.
+
+    ### Test Mechanism
+
+    This metric test generates a correlation matrix for all numerical variables in the dataset using the Pearson
+    correlation formula. A heat map is subsequently created to visualize this matrix effectively. The color of each
+    point on the heat map corresponds to the magnitude and direction (positive or negative) of the correlation, with a
+    range from -1 (perfect negative correlation) to 1 (perfect positive correlation). Any correlation coefficients
+    higher than 0.7 (in absolute terms) are indicated in white in the heat map, suggesting a high degree of correlation.
+
+    ### Signs of High Risk
 
-    **Signs of High Risk**:
     - A large number of variables in the dataset showing a high degree of correlation (coefficients approaching ±1).
     This indicates redundancy within the dataset, suggesting that some variables may not be contributing new
     information to the model.
-    - This could potentially lead to overfitting.
-
-    **Strengths**:
-    - The primary strength of this metric test is its ability to detect and quantify the linearity of relationships
-    between variables. This allows for the identification of redundant variables, which in turn can help in simplifying
-    models and potentially improving their performance.
-    - The visualization aspect (heatmap) is another strength as it offers an easy-to-understand overview of the
-    correlations, beneficial for those not comfortable navigating numerical matrices.
-
-    **Limitations**:
-    - The primary limitation of Pearson Correlation is its inability to detect non-linear relationships between
-    variables, which can lead to missed opportunities for dimensionality reduction.
-    - It only measures the degree of linear relationship and not the strength of effect of one variable on the other.
-    - The cutoff value of 0.7 for high correlation is a somewhat arbitrary choice and some valid dependencies might be
-    missed if they have a correlation coefficient less than this value.
-    """
+    - Potential risk of overfitting.
 
-    name = "pearson_correlation_matrix"
-    required_inputs = ["dataset"]
-    tasks = ["classification", "regression"]
-    tags = ["tabular_data", "numerical_data", "correlation"]
-
-    def run(self):
-        columns = self.params.get("columns", list(self.inputs.dataset.df.columns))
-
-        corr_matrix = self.inputs.dataset.df[columns].corr(numeric_only=True)
-        heatmap = go.Heatmap(
-            z=corr_matrix.values,
-            x=list(corr_matrix.columns),
-            y=list(corr_matrix.index),
-            colorscale="rdbu",
-            zmin=-1,
-            zmax=1,
-        )
-
-        annotations = []
-        for i, row in enumerate(corr_matrix.values):
-            for j, value in enumerate(row):
-                color = "#ffffff" if abs(value) > 0.7 else "#000000"
-                annotations.append(
-                    go.layout.Annotation(
-                        text=str(round(value, 2)),
-                        x=corr_matrix.columns[j],
-                        y=corr_matrix.index[i],
-                        showarrow=False,
-                        font=dict(color=color),
-                    )
-                )
+    ### Strengths
+
+    - Detects and quantifies the linearity of relationships between variables, aiding in identifying redundant
+    variables to simplify models and potentially improve performance.
+    - The heatmap visualization provides an easy-to-understand overview of correlations, beneficial for users not
+    comfortable with numerical matrices.
+
+    ### Limitations
+
+    - Limited to detecting linear relationships, potentially missing non-linear relationships which impede
+    opportunities for dimensionality reduction.
+    - Measures only the degree of linear relationship, not the strength of one variable's effect on another.
+    - The 0.7 correlation threshold is arbitrary and might exclude valid dependencies with lower coefficients.
+    """
 
-        layout = go.Layout(
-            annotations=annotations,
-            xaxis=dict(side="top"),
-            yaxis=dict(scaleanchor="x", scaleratio=1),
-            width=800,
-            height=800,
-            autosize=True,
-            paper_bgcolor="rgba(0,0,0,0)",
-            plot_bgcolor="rgba(0,0,0,0)",
-        )
-
-        fig = go.Figure(data=[heatmap], layout=layout)
-
-        return self.cache_results(
-            figures=[
-                Figure(
-                    for_object=self,
-                    key=self.key,
-                    figure=fig,
+    corr_matrix = dataset.df.corr(numeric_only=True)
+    heatmap = go.Heatmap(
+        z=corr_matrix.values,
+        x=list(corr_matrix.columns),
+        y=list(corr_matrix.index),
+        colorscale="rdbu",
+        zmin=-1,
+        zmax=1,
+    )
+
+    annotations = []
+    for i, row in enumerate(corr_matrix.values):
+        for j, value in enumerate(row):
+            color = "#ffffff" if abs(value) > 0.7 else "#000000"
+            annotations.append(
+                go.layout.Annotation(
+                    text=str(round(value, 2)),
+                    x=corr_matrix.columns[j],
+                    y=corr_matrix.index[i],
+                    showarrow=False,
+                    font=dict(color=color),
                 )
-            ]
-        )
+            )
+
+    layout = go.Layout(
+        annotations=annotations,
+        xaxis=dict(side="top"),
+        yaxis=dict(scaleanchor="x", scaleratio=1),
+        width=800,
+        height=800,
+        autosize=True,
+        paper_bgcolor="rgba(0,0,0,0)",
+        plot_bgcolor="rgba(0,0,0,0)",
+    )
+
+    fig = go.Figure(data=[heatmap], layout=layout)
+
+    return fig