Dataset size is the number of structures

metatensor · Sep 23, 2024 · a6f4da6 · a6f4da6
1 parent 507316a
commit a6f4da6
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/src/metatrain/utils/data/dataset.py b/src/metatrain/utils/data/dataset.py
@@ -246,7 +246,7 @@ def get_stats(dataset: Union[Dataset, Subset], dataset_info: DatasetInfo) -> str
     """Returns the statistics of a dataset or subset as a string."""
 
     dataset_len = len(dataset)
-    stats = f"Dataset of size {dataset_len}"
+    stats = f"Dataset containing {dataset_len} structures"
     if dataset_len == 0:
         return stats
 

diff --git a/tests/cli/test_train_model.py b/tests/cli/test_train_model.py
@@ -77,7 +77,7 @@ def test_train(capfd, monkeypatch, tmp_path, output):
     assert "Training dataset:" in stdout_log
     assert "Validation dataset:" in stdout_log
     assert "Test dataset:" in stdout_log
-    assert "size 50" in stdout_log
+    assert "50 structures" in stdout_log
     assert "mean " in stdout_log
     assert "std " in stdout_log
     assert "[INFO]" in stdout_log

diff --git a/tests/utils/data/test_dataset.py b/tests/utils/data/test_dataset.py
@@ -652,7 +652,7 @@ def test_get_stats():
     stats = get_stats(dataset, dataset_info)
     stats_2 = get_stats(dataset_2, dataset_info)
 
-    assert "size 100" in stats
+    assert "100 structures" in stats
     assert "mtt::U0" in stats
     assert "energy" in stats_2
     assert "mean " in stats