@@ -115,7 +115,8 @@ def test_ml_ff_md_maker(
115115 traj_file = "md_traj.json.gz" ,
116116 traj_file_fmt = "pmg" ,
117117 store_trajectory = "partial" ,
118- ionic_step_data = ("energy" , "forces" , "stress" , "mol_or_struct" ),
118+ # check that `structure` alias to `mol_or_struct` works:
119+ ionic_step_data = ("energy" , "forces" , "stress" , "structure" ),
119120 calculator_kwargs = calculator_kwargs ,
120121 use_emmet_models = use_emmet_models ,
121122 ).make (structure )
@@ -135,10 +136,15 @@ def test_ml_ff_md_maker(
135136 # Check that the ionic steps have the expected physical properties
136137 assert all (
137138 key in step .model_dump ()
138- for key in ("energy" , "forces" , "stress" , "mol_or_struct" , "structure" )
139+ for key in ("energy" , "forces" , "stress" , "mol_or_struct" )
139140 for step in task_doc .output .ionic_steps
140141 )
141142
143+ # `structure` aliases `mol_or_struct`
144+ assert all (
145+ step .structure == step .mol_or_struct for step in task_doc .output .ionic_steps
146+ )
147+
142148 # Check that the trajectory has expected physical properties
143149 assert task_doc .included_objects == ["trajectory" ]
144150 assert len (task_doc .objects ["trajectory" ]) == n_steps + 1
0 commit comments