About the difference between the csv file I produced using file "test_chgnet.py" and the csv file under the chgnet folder

[Claudia-Hello]

Hi,
I tried to reproduce the predictions of the CHGNet model on the WBM test set and got the csv file as follows:
The three columns include"material_id", "chgnet_energy", and "chgnet_structure"`.

However, the contents of the 2023-12-21-chgnet-0.3.0-wbm-IS2RE.csv file in the chgnet folder are as follows:
The three columns include "material_id", "chgnet_energy", and "e_form_per_atom_chgnet".

As you can see, the last column of my predicted csv file is "chgnet_structure", while the last column of 2023-12-21-chgnet-0.3.0-wbm-IS2RE.csv is "e_form_per_atom_chgnet". Also, my predicted csv file results in only 5319, while the original 2023-12-21-chgnet-0.3.0-wbm-IS2RE.csv file results in 256,963.

To summarize, I would like to ask, how did you get such predicted results as 2023-12-21-chgnet-0.3.0-wbm-IS2RE.csv file? Can you tell me the details?
Thank you!

[janosh]

test_chgnet.py is designed to be run as a slurm job array where each sub task runs on a different split of the WBM test set.

matbench-discovery/models/chgnet/test_chgnet.py

Lines 35 to 49 in 71a5edc

	slurm_array_task_count = 50
	device = "cuda" if torch.cuda.is_available() else "cpu"
	chgnet = StructOptimizer(use_device=device) # load default pre-trained CHGNnet model
	job_name = f"chgnet-{chgnet.version}-wbm-{task_type}"
	out_dir = os.getenv("SBATCH_OUTPUT", f"{module_dir}/{today}-{job_name}")
	slurm_vars = slurm_submit(
	job_name=job_name,
	out_dir=out_dir,
	account="matgen",
	time="11:55:0",
	array=f"1-{slurm_array_task_count}",
	slurm_flags="--qos shared --constraint cpu --nodes 1 --mem 10G",
	# slurm_flags="--qos regular --constraint gpu --gpus 1",
	)

once all jobs are finished you run join_chgnet_results.py to collect all output files into 1 file and apply the MP formation energy corrections:

matbench-discovery/models/chgnet/join_chgnet_results.py

Lines 30 to 50 in 71a5edc

	task_type = Task.IS2RE
	date = "2023-12-21"
	glob_pattern = f"{date}-chgnet-*-wbm-{task_type}*/*.json.gz"
	file_paths = sorted(glob(f"{module_dir}/{glob_pattern}"))
	print(f"Found {len(file_paths):,} files for {glob_pattern = }")
	dfs: dict[str, pd.DataFrame] = {}
	# %%
	failed = {}
	for file_path in tqdm(file_paths):
	if file_path in dfs:
	continue
	try:
	df = pd.read_json(file_path).set_index(Key.mat_id)
	except Exception as exc:
	failed[file_path] = str(exc)
	continue
	# drop trajectory to save memory
	dfs[file_path] = df.drop(columns="chgnet_trajectory", errors="ignore")

join_chgnet_results.py saves the final output as yyy-mm-dd-chgnet-0.3.0-wbm-IS2RE.csv.gz to disk