Merge pull request #34 from learningmatter-mit/vssr_pourbaix

Merge `vssr_pourbaix`

Author: HojeChun

models/foundation_models/chgnet/0.2.0/README.md

@@ -4,18 +4,24 @@
 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=(12, 6), dpi=mpl_settings.DPI)
+    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[val_key], label="val", color=mpl_settings.colors[2])
     ax_fig[0].legend()

models/foundation_models/chgnet/0.2.0/chgnet_0.2.0_e30f77s348m32.pth.tar

@@ -1,34 +1,44 @@
+from __future__ import annotations
+
 import json
 from pathlib import Path
-from typing import List, Optional
+from typing import List
 
-import matplotlib
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 import numpy as np
 
 plt.style.use("default")
 
-DPI = 100
-LINEWIDTH = 2
-FONTSIZE = 20
-LABELSIZE = 18
+dir_name = Path(__file__).parent
+
+DPI = 300
+LINEWIDTH = 1.25
+FONTSIZE = 8
+LABELSIZE = 8
 ALPHA = 0.8
-LINE_MARKERSIZE = 15 * 25
-MARKERSIZE = 15
-GRIDSIZE = 40
-MAJOR_TICKLEN = 6
-MINOR_TICKLEN = 3
-TICKPADDING = 5
+MARKERSIZE = 25
+GRIDSIZE = 20
+MAJOR_TICKLEN = 4
+MINOR_TICKLEN = 2
+TICKPADDING = 3
 SECONDARY_CMAP = "inferno"
 
-params = {
+custom_settings = {
     "mathtext.default": "regular",
     "font.family": "Arial",
     "font.size": FONTSIZE,
     "axes.labelsize": LABELSIZE,
     "axes.titlesize": FONTSIZE,
+    "axes.linewidth": LINEWIDTH,
     "grid.linewidth": LINEWIDTH,
     "lines.linewidth": LINEWIDTH,
+    "lines.color": "black",
+    "axes.labelcolor": "black",
+    "axes.edgecolor": "black",
+    "axes.titlecolor": "black",
+    "axes.titleweight": "bold",
+    "axes.grid": False,
     "lines.markersize": MARKERSIZE,
     "xtick.major.size": MAJOR_TICKLEN,
     "xtick.minor.size": MINOR_TICKLEN,
@@ -38,66 +48,67 @@
     "ytick.minor.size": MINOR_TICKLEN,
     "ytick.major.pad": TICKPADDING,
     "ytick.minor.pad": TICKPADDING,
-    "axes.linewidth": LINEWIDTH,
-    "legend.fontsize": LABELSIZE,
-    "figure.dpi": DPI,
-    "savefig.dpi": DPI,
     "ytick.major.width": LINEWIDTH,
     "xtick.major.width": LINEWIDTH,
     "ytick.minor.width": LINEWIDTH,
     "xtick.minor.width": LINEWIDTH,
+    "legend.fontsize": LABELSIZE,
+    "figure.dpi": DPI,
+    "savefig.dpi": DPI,
+    "savefig.format": "png",
+    "savefig.bbox": "tight",
+    "savefig.pad_inches": 0.1,
+    "figure.facecolor": "white",
 }
-plt.rcParams.update(params)
+plt.rcParams.update(custom_settings)
+
 
+def update_custom_settings(custom_settings: dict | None = custom_settings) -> None:
+    """Update the custom settings for Matplotlib.
 
-def hex_to_rgb(value: str) -> tuple:
+    Args:
+        custom_settings: Custom settings for Matplotlib. Defaults to
+            custom_settings.
     """
-    Converts hex to rgb colours
+    current_settings = plt.rcParams.copy()
+    new_settings = current_settings | custom_settings
+    plt.rcParams.update(new_settings)
 
-    Parameters
-    ----------
-    value: string of 6 characters representing a hex colour
 
-    Returns
-    ----------
-    tuple of 3 integers representing the RGB values
-    """
+def hex_to_rgb(value: str) -> list[float]:
+    """Converts hex to rgb colors.
 
+    Args:
+        value: string of 6 characters representing a hex color.
+    """
     value = value.strip("#")  # removes hash symbol if present
     lv = len(value)
     return tuple(int(value[i : i + lv // 3], 16) for i in range(0, lv, lv // 3))
 
 
-def rgb_to_dec(value: list):
-    """
-    Converts rgb to decimal colours (i.e. divides each value by 256)
+def rgb_to_dec(value: list[float]) -> list[float]:
+    """Converts rgb to decimal colors (i.e. divides each value by 256).
 
-    Parameters
-    ----------
-    value: list of 3 integers representing the RGB values
+    Args:
+        value: string of 6 characters representing a hex color.
 
-    Returns
-    ----------
-    list of 3 floats representing the RGB values
+    Returns:
+        list: length 3 of RGB values
     """
-
     return [v / 256 for v in value]
 
 
-def get_continuous_cmap(hex_list: List[str], float_list: Optional[List[float]] = None) -> matplotlib.colors.Colormap:
-    """
-    Creates and returns a color map that can be used in heat map figures.
-    If float_list is not provided, colour map graduates linearly between each color in hex_list.
-    If float_list is provided, each color in hex_list is mapped to the respective location in float_list.
-
-    Parameters
-    ----------
-    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.
-
-    Returns
-    ----------
-    Colormap
+def get_continuous_cmap(
+    hex_list: list[str], float_list: list[float] | None = None
+) -> mpl.colors.LinearSegmentedColormap:
+    """Creates a color map that can be used in heat map figures. If float_list is not provided,
+    color map graduates linearly between each color in hex_list. If float_list is provided,
+    each color in hex_list is mapped to the respective location in float_list.
+
+    Args:
+        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:
@@ -109,15 +120,12 @@ def get_continuous_cmap(hex_list: List[str], float_list: Optional[List[float]] =
     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))]
         cdict[col] = col_list
-    cmp = matplotlib.colors.LinearSegmentedColormap("j_cmap", segmentdata=cdict, N=256)
-    return cmp
+    return mpl.colors.LinearSegmentedColormap("j_cmap", segmentdata=cdict, N=256)
 
 
-# colors taken from Johannes Dietschreit's script and interpolated with correct lightness and Bezier
+# Colors taken from Johannes Dietschreit's script and interpolated with correct lightness and Bezier
 # http://www.vis4.net/palettes/#/100|s|fce1a4,fabf7b,f08f6e,d12959,6e005f|ffffe0,ff005e,93003a|1|1
 hex_list: List[str]
-dir_name = Path(__file__).parent
-
 with open(dir_name / "config/mpl_settings.json", "r") as f:
     hex_list = json.load(f)["plot_colors"]