MAINT & STY: update based on PR #34 comments and style fixes

Author: xiaochendu

nff/analysis/loss_plot.py

@@ -4,38 +4,35 @@
 from . import mpl_settings
 
 
-def plot_loss(energy_history, forces_history, figname, train_key="train", val_key="val"):
+def plot_loss(
+    energy_history: dict,
+    forces_history: dict,
+    figname: str,
+    train_key: str = "train",
+    val_key: str = "val",
+) -> None:
     """Plot the loss history of the model.
-    Args:
-        energy_history (dict): energy loss history of the model for training and validation
-        forces_history (dict): forces loss history of the model for training and validation
-        figname (str): name of the figure
-
 
-    Returns:
-        None
+    Args:
+        energy_history: energy loss history of the model for training and validation
+        forces_history: forces loss history of the model for training and validation
+        figname: name of the figure
+        train_key: key for training data in the history dictionary
+        val_key: key for validation data in the history dictionary
     """
     epochs = np.arange(1, len(energy_history[train_key]) + 1)
     fig, ax_fig = plt.subplots(1, 2, figsize=(5, 2.5), dpi=mpl_settings.DPI)
-    ax_fig[0].semilogy(
-        epochs, energy_history[train_key], label="train", color=mpl_settings.colors[1]
-    )
+    ax_fig[0].semilogy(epochs, energy_history[train_key], label="train", color=mpl_settings.colors[1])
     ax_fig[0].semilogy(epochs, energy_history[val_key], label="val", color=mpl_settings.colors[2])
     ax_fig[0].legend()
     ax_fig[0].set_xlabel("Epoch")
     ax_fig[0].set_ylabel("Loss")
-    ax_fig[0].set_xlabel("Epoch")
-    ax_fig[0].set_ylabel("Loss")
 
-    ax_fig[1].semilogy(
-        epochs, forces_history[train_key], label="train", color=mpl_settings.colors[1]
-    )
+    ax_fig[1].semilogy(epochs, forces_history[train_key], label="train", color=mpl_settings.colors[1])
     ax_fig[1].semilogy(epochs, forces_history[val_key], label="val", color=mpl_settings.colors[2])
     ax_fig[1].legend()
     ax_fig[1].set_xlabel("Epoch")
     ax_fig[1].set_ylabel("Loss")
-    ax_fig[1].set_xlabel("Epoch")
-    ax_fig[1].set_ylabel("Loss")
 
     plt.tight_layout()
     plt.savefig(f"{figname}.png")

nff/analysis/mpl_settings.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import json
 from pathlib import Path
-from typing import List, Optional
+from typing import List
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
@@ -61,11 +63,11 @@
 plt.rcParams.update(custom_settings)
 
 
-def update_custom_settings(custom_settings: dict = custom_settings) -> None:
+def update_custom_settings(custom_settings: dict | None = custom_settings) -> None:
     """Update the custom settings for Matplotlib.
 
     Args:
-        custom_settings (dict, optional): Custom settings for Matplotlib. Defaults to
+        custom_settings: Custom settings for Matplotlib. Defaults to
             custom_settings.
     """
     current_settings = plt.rcParams.copy()
@@ -77,10 +79,7 @@ def hex_to_rgb(value: str) -> list[float]:
     """Converts hex to rgb colors.
 
     Args:
-        value (str): string of 6 characters representing a hex colour.
-
-    Returns:
-        list: length 3 of RGB values
+        value: string of 6 characters representing a hex color.
     """
     value = value.strip("#")  # removes hash symbol if present
     lv = len(value)
@@ -91,7 +90,7 @@ def rgb_to_dec(value: list[float]) -> list[float]:
     """Converts rgb to decimal colors (i.e. divides each value by 256).
 
     Args:
-        value (list[float]): string of 6 characters representing a hex colour.
+        value: string of 6 characters representing a hex color.
 
     Returns:
         list: length 3 of RGB values
@@ -107,12 +106,9 @@ def get_continuous_cmap(
     each color in hex_list is mapped to the respective location in float_list.
 
     Args:
-        hex_list (list[str]): list of hex code strings
-        float_list (list[float]): list of floats between 0 and 1, same length as hex_list. Must
-            start with 0 and end with 1.
-
-    Returns:
-        matplotlib.colors.LinearSegmentedColormap: continuous
+        hex_list: list of hex code strings
+        float_list: list of floats between 0 and 1, same length as hex_list.
+            Must start with 0 and end with 1.
     """
     rgb_list = [rgb_to_dec(hex_to_rgb(i)) for i in hex_list]
     if float_list:
@@ -122,9 +118,7 @@ def get_continuous_cmap(
 
     cdict = dict()
     for num, col in enumerate(["red", "green", "blue"]):
-        col_list = [
-            [float_list[i], rgb_list[i][num], rgb_list[i][num]] for i in range(len(float_list))
-        ]
+        col_list = [[float_list[i], rgb_list[i][num], rgb_list[i][num]] for i in range(len(float_list))]
         cdict[col] = col_list
     return mpl.colors.LinearSegmentedColormap("j_cmap", segmentdata=cdict, N=256)
 

nff/analysis/parity_plot.py

@@ -1,13 +1,16 @@
-from typing import Dict, Literal, Union
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Literal
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
-import torch
 from matplotlib.lines import Line2D
 from scipy import stats
 from scipy.stats import gaussian_kde
 
+if TYPE_CHECKING:
+    from torch import Tensor
 from nff.data import to_tensor
 from nff.utils import cuda
 
@@ -18,8 +21,8 @@
 
 
 def plot_parity(
-    results: Dict[str, Union[list, torch.Tensor]],
-    targets: Dict[str, Union[list, torch.Tensor]],
+    results: Dict[str, list | Tensor],
+    targets: Dict[str, list | Tensor],
     figname: str,
     plot_type: Literal["hexbin", "scatter"] = "hexbin",
     energy_key: str = "energy",
@@ -29,13 +32,13 @@ def plot_parity(
     """Perform a parity plot between the results and the targets.
 
     Args:
-        results (dict): dictionary containing the results
-        targets (dict): dictionary containing the targets
-        figname (str): name of the figure
-        plot_type (str): type of plot to use, either "hexbin" or "scatter"
-        energy_key (str): key for the energy
-        force_key (str): key for the forces
-        units (dict): dictionary containing the units of the keys
+        results: dictionary containing the results
+        targets: dictionary containing the targets
+        figname: name of the figure
+        plot_type: type of plot to use, either "hexbin" or "scatter"
+        energy_key: key for the energy
+        force_key: key for the forces
+        units: dictionary containing the units of the keys
 
     Returns:
         float: MAE of the energy
@@ -113,8 +116,8 @@ def plot_parity(
 
 
 def plot_err_var(
-    err: Union[torch.Tensor, np.ndarray],
-    var: Union[torch.Tensor, np.ndarray],
+    err: Tensor | np.ndarray,
+    var: Tensor | np.ndarray,
     figname: str,
     units: str = "eV/Å",
     x_min: float = 0.0,
@@ -128,20 +131,17 @@ def plot_err_var(
     """Plot the error vs variance of the forces.
 
     Args:
-        err (torch.Tensor): error of the forces
-        var (torch.Tensor): variance of the forces
-        figname (str): name of the figure
-        units (str): units of the error and variance
-        x_min (float): minimum value of the x-axis
-        x_max (float): maximum value of the x-axis
-        y_min (float): minimum value of the y-axis
-        y_max (float): maximum value of the y-axis
-        sample_frac (float): fraction of the data to sample for the plot
-        num_bins (int): number of bins to use for binning
-        cb_format (str): format of the colorbar
-
-    Returns:
-        None
+        err: error of the forces
+        var: variance of the forces
+        figname: name of the figure
+        units: units of the error and variance
+        x_min: minimum value of the x-axis
+        x_max: maximum value of the x-axis
+        y_min: minimum value of the y-axis
+        y_max: maximum value of the y-axis
+        sample_frac: fraction of the data to sample for the plot
+        num_bins: number of bins to use for binning
+        cb_format: format of the colorbar
     """
     fig, ax = plt.subplots(1, 1, figsize=(6, 6), dpi=mpl_settings.DPI)
 
@@ -166,12 +166,7 @@ def plot_err_var(
     x = pd.Series(var)
     y = pd.Series(err)
 
-    kernel = gaussian_kde(
-        np.vstack([x.sample(n=len(x), random_state=2), y.sample(n=len(y), random_state=2)])
-    )
-    kernel = gaussian_kde(
-        np.vstack([x.sample(n=len(x), random_state=2), y.sample(n=len(y), random_state=2)])
-    )
+    kernel = gaussian_kde(np.vstack([x.sample(n=len(x), random_state=2), y.sample(n=len(y), random_state=2)]))
     c = kernel(np.vstack([x, y]))
     hb = ax.scatter(
         var,