fix plot function cumulative_metrics() when 'RMSE' in metrics

simplify test_glob_to_df
upload_to_figshare.py remove CHUNK_SIZE global

Author: janosh

.pre-commit-config.yaml

@@ -7,7 +7,7 @@ default_install_hook_types: [pre-commit, commit-msg]
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.0.289
+    rev: v0.0.290
     hooks:
       - id: ruff
         args: [--fix]

matbench_discovery/plots.py

@@ -667,8 +667,9 @@ def cumulative_metrics(
         df_preds (pd.DataFrame): Distance to convex hull predicted by models, one column
             per model (in eV / atom). Same as true energy to convex hull plus predicted
             minus true formation energy.
-        metrics (Sequence[str], optional): Which metrics to plot. Defaults to
-            ('Precision', 'Recall'). Also accepts 'F1'.
+        metrics (Sequence[str], optional): Which metrics to plot. Any subset of
+            ("Precision", "Recall", "F1", "MAE", "RMSE").
+            Defaults to ('Precision', 'Recall').
         stability_threshold (float, optional): Max distance above convex hull before
             material is considered unstable. Defaults to 0.
         project_end_point ('x' | 'y' | 'xy' | '', optional): Whether to project end
@@ -735,9 +736,9 @@ def cumulative_metrics(
             f1_interp = cubic_interpolate(model_range, f1_cum[:n_pred_stable])
             dfs["F1"][model_name] = dict(zip(xs_model, f1_interp(xs_model)))
 
+        cum_counts = np.arange(1, len(each_true) + 1)
         if "MAE" in metrics:
             cum_errors = (each_true - each_pred).abs().cumsum()
-            cum_counts = np.arange(1, len(each_true) + 1)
             mae_cum = cum_errors / cum_counts
             mae_interp = cubic_interpolate(model_range, mae_cum[:n_pred_stable])
             dfs["MAE"][model_name] = dict(zip(xs_model, mae_interp(xs_model)))
@@ -848,6 +849,7 @@ def cumulative_metrics(
                     text=optimal_recall,
                     showarrow=False,
                     # rotate text parallel to line
+                    # angle not quite right, could be improved
                     textangle=math.degrees(math.cos(n_stable)),
                     **grid_pos,
                 )

scripts/upload_to_figshare.py

@@ -28,7 +28,6 @@
     TOKEN = file.read().split("figshare_token=")[1].split("\n")[0]
 
 BASE_URL = "https://api.figshare.com/v2"
-CHUNK_SIZE = 10_000_000  # ~10MB
 
 with open(f"{ROOT}/pyproject.toml", "rb") as file:
     pyproject = tomllib.load(file)["project"]
@@ -96,12 +95,16 @@ def create_article(metadata: dict[str, str | int | float]) -> int:
     return result["id"]
 
 
-def get_file_hash_and_size(file_name: str) -> tuple[str, int]:
-    """Get the md5 hash and size of a file."""
+def get_file_hash_and_size(
+    file_name: str, chunk_size: int = 10_000_000
+) -> tuple[str, int]:
+    """Get the md5 hash and size of a file. File is read in chunks of chunk_size bytes.
+    Default chunk size is 10_000_000 ~= 10MB.
+    """
     md5 = hashlib.md5()
     size = 0
     with open(file_name, "rb") as file:
-        while data := file.read(CHUNK_SIZE):
+        while data := file.read(chunk_size):
             size += len(data)
             md5.update(data)
     return md5.hexdigest(), size

tests/test_data.py

@@ -52,12 +52,12 @@ def test_load(
 ) -> None:
     filepath = DATA_FILES[data_key]
     # intercept HTTP requests and write dummy df to disk instead
-    with patch("urllib.request.urlretrieve") as urlretrieve:
+    with patch("urllib.request.urlretrieve") as url_retrieve:
         # dummy df with random floats and material_id column
         df_csv = pd._testing.makeDataFrame().reset_index(names="material_id")
 
         writer = dummy_df_serialized.to_json if ".json" in filepath else df_csv.to_csv
-        urlretrieve.side_effect = lambda url, path: writer(path)
+        url_retrieve.side_effect = lambda _url, path: writer(path)
         out = load(
             data_key,
             hydrate=hydrate,
@@ -70,7 +70,7 @@ def test_load(
     assert f"Downloading {data_key!r} from {figshare_urls[data_key][0]}" in stdout
 
     # check we called read_csv/read_json once for each data_name
-    assert urlretrieve.call_count == 1
+    assert url_retrieve.call_count == 1
 
     assert isinstance(out, pd.DataFrame), f"{data_key} not a DataFrame"
 
@@ -201,21 +201,17 @@ def test_df_wbm() -> None:
     assert set(df_wbm) > {"bandgap_pbe", "formula", "material_id"}
 
 
-@pytest.mark.parametrize("pattern", ["tmp/*df.csv", "tmp/*df.json"])
-def test_glob_to_df(pattern: str) -> None:
-    try:
-        df = pd._testing.makeMixedDataFrame()
+@pytest.mark.parametrize("pattern", ["*df.csv", "*df.json"])
+def test_glob_to_df(pattern: str, tmp_path: Path) -> None:
+    df = pd._testing.makeMixedDataFrame()
 
-        os.makedirs(f"{ROOT}/tmp", exist_ok=True)
-        df.to_csv(f"{ROOT}/tmp/dummy_df.csv", index=False)
-        df.to_json(f"{ROOT}/tmp/dummy_df.json")
+    os.makedirs(f"{tmp_path}", exist_ok=True)
+    df.to_csv(f"{tmp_path}/dummy_df.csv", index=False)
+    df.to_json(f"{tmp_path}/dummy_df.json")
 
-        df_out = glob_to_df(pattern)
-        assert df_out.shape == df.shape
-        assert list(df_out) == list(df)
+    df_out = glob_to_df(f"{tmp_path}/{pattern}")
+    assert df_out.shape == df.shape
+    assert list(df_out) == list(df)
 
-        with pytest.raises(FileNotFoundError):
-            glob_to_df("foo")
-    finally:
-        os.remove(f"{ROOT}/tmp/dummy_df.csv")
-        os.remove(f"{ROOT}/tmp/dummy_df.json")
+    with pytest.raises(FileNotFoundError):
+        glob_to_df("foo")

tests/test_models.py

@@ -34,9 +34,10 @@ def test_model_dirs_have_metadata() -> None:
             err_msg = f"Invalid {key=}, expected {required[key]} in {model_dir}"
             assert isinstance(metadata[key], required[key]), err_msg  # type: ignore
 
-        authors, date_added, mbd_version, model_name, model_version, repo = (
+        authors, date_added, mbd_version, yml_model_name, model_version, repo = (
             metadata[key] for key in list(required)[:-1]
         )
+        assert model_name == yml_model_name, f"{model_name=} != {yml_model_name=}"
 
         # make sure all keys are valid
         for name in model_name if isinstance(model_name, list) else [model_name]:

tests/test_plots.py

@@ -38,11 +38,7 @@
 @pytest.mark.parametrize("backend", ["matplotlib", "plotly"])
 @pytest.mark.parametrize(
     "metrics",
-    [
-        ("Recall",),
-        ("Recall", "MAE"),
-        ("Recall", "Precision", "F1"),
-    ],
+    [("Recall",), ("Recall", "MAE"), ("Recall", "Precision", "RMSE")],
 )
 def test_cumulative_metrics(
     project_end_point: AxLine,
@@ -68,9 +64,8 @@ def test_cumulative_metrics(
         assert {ax.get_ylabel() for ax in fig.axes} >= {*metrics}
     elif backend == "plotly":
         assert isinstance(fig, go.Figure)
-        # TODO fix AssertionError {'Recall', 'metric=F1'} == {'F1', 'Recall'}
-        # subplot_titles = [anno.text for anno in fig.layout.annotations][:len(metrics)]
-        # assert set(subplot_titles) == set(metrics)
+        subplot_titles = {anno.text.split("=")[-1] for anno in fig.layout.annotations}
+        assert subplot_titles >= set(metrics)
 
 
 def test_cumulative_metrics_raises() -> None: