Add flexible filename and extension to phonon new_plotter (#792)

* added a way to set a flexible filename and extension for the phonon band structure and DOS new_plotter object.

* replacing Si strcuture with si_strcuture fixture

* implementing requested changes

* change default filename_(bs|dos) extension to PDF

* use frozenset for SUPPORTED_CODES

* save filename_bs, filename_dos to tmp dir and simplify assert file exists

* tweak comment

---------

Co-authored-by: Janosh Riebesell <[email protected]>

src/atomate2/common/flows/phonons.py

@@ -27,7 +27,7 @@
     from atomate2.forcefields.jobs import ForceFieldRelaxMaker, ForceFieldStaticMaker
     from atomate2.vasp.jobs.base import BaseVaspMaker
 
-SUPPORTED_CODES = ["vasp", "aims", "forcefields"]
+SUPPORTED_CODES = frozenset(("vasp", "aims", "forcefields"))
 
 
 @dataclass

src/atomate2/common/schemas/phonons.py

@@ -359,9 +359,9 @@ def from_forces_born(
             filename_band_yaml, labels_dict=kpath_dict, has_nac=born is not None
         )
         new_plotter = PhononBSPlotter(bs=bs_symm_line)
-
         new_plotter.save_plot(
-            "phonon_band_structure.eps", units=kwargs.get("units", "THz")
+            filename=kwargs.get("filename_bs", "phonon_band_structure.pdf"),
+            units=kwargs.get("units", "THz"),
         )
 
         # will determine if imaginary modes are present in the structure
@@ -389,7 +389,8 @@ def from_forces_born(
         new_plotter_dos = PhononDosPlotter()
         new_plotter_dos.add_dos(label="total", dos=dos)
         new_plotter_dos.save_plot(
-            filename="phonon_dos.eps", units=kwargs.get("units", "THz")
+            filename=kwargs.get("filename_dos", "phonon_dos.pdf"),
+            units=kwargs.get("units", "THz"),
         )
 
         # compute vibrational part of free energies per formula unit

tests/forcefields/flows/test_phonon.py

@@ -1,3 +1,6 @@
+import os
+from pathlib import Path
+
 import torch
 from jobflow import run_locally
 from numpy.testing import assert_allclose
@@ -14,7 +17,7 @@
 from atomate2.forcefields.flows.phonons import PhononMaker
 
 
-def test_phonon_wf(clean_dir):
+def test_phonon_wf(clean_dir, tmp_path: Path):
     # TODO brittle due to inability to adjust dtypes in CHGNetRelaxMaker
     torch.set_default_dtype(torch.float32)
 
@@ -29,7 +32,11 @@ def test_phonon_wf(clean_dir):
         create_thermal_displacements=False,
         store_force_constants=False,
         prefer_90_degrees=False,
-        generate_frequencies_eigenvectors_kwargs={"tstep": 100},
+        generate_frequencies_eigenvectors_kwargs={
+            "tstep": 100,
+            "filename_bs": (filename_bs := f"{tmp_path}/phonon_bs_test.png"),
+            "filename_dos": (filename_dos := f"{tmp_path}/phonon_dos_test.pdf"),
+        },
     ).make(structure)
 
     # run the flow or job and ensure that it finished running successfully
@@ -100,3 +107,7 @@ def test_phonon_wf(clean_dir):
         [5058.44158791, 5385.88058579, 6765.19854165, 8723.78588089, 10919.0199409],
         atol=1000,
     )
+
+    # check phonon plots exist
+    assert os.path.isfile(filename_bs)
+    assert os.path.isfile(filename_dos)