Fix TypeError: can only join an iterable with AECCAR in `Volumetric…

…Data.write_file` (#3343)

* bugfix for aeccars `write_file`

* bugfix: chargemol_output_path

* add aeccar write_file test

* correctly assert allclose in TestAeccars

* rename single-letter variables

---------

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

pymatgen/command_line/chargemol_caller.py

@@ -196,7 +196,7 @@ def _from_data_dir(self, chargemol_output_path=None):
                 Default: None (current working directory).
         """
         if chargemol_output_path is None:
-            chargemol_output_path = ""
+            chargemol_output_path = "."
 
         charge_path = f"{chargemol_output_path}/DDEC6_even_tempered_net_atomic_charges.xyz"
         self.ddec_charges = self._get_data_from_xyz(charge_path)

pymatgen/io/vasp/outputs.py

@@ -11,6 +11,7 @@
 import warnings
 import xml.etree.ElementTree as ET
 from collections import defaultdict
+from collections.abc import Iterable
 from dataclasses import dataclass
 from glob import glob
 from io import StringIO
@@ -3348,7 +3349,7 @@ def parse_file(filename):
                 data_aug = {"total": all_dataset_aug.get(0)}
             return poscar, data, data_aug
 
-    def write_file(self, file_name, vasp4_compatible=False):
+    def write_file(self, file_name: str | Path, vasp4_compatible: bool = False) -> None:
         """
         Write the VolumetricData object to a vasp compatible file.
 
@@ -3358,8 +3359,7 @@ def write_file(self, file_name, vasp4_compatible=False):
         """
 
         def _print_fortran_float(flt):
-            """
-            Fortran codes print floats with a leading zero in scientific
+            """Fortran codes print floats with a leading zero in scientific
             notation. When writing CHGCAR files, we adopt this convention
             to ensure written CHGCAR files are byte-to-byte identical to
             their input files as far as possible.
@@ -3375,42 +3375,45 @@ def _print_fortran_float(flt):
                 return f"0.{s[0]}{s[2:12]}E{int(s[13:]) + 1:+03}"
             return f"-.{s[1]}{s[3:13]}E{int(s[14:]) + 1:+03}"
 
-        with zopen(file_name, "wt") as f:
-            p = Poscar(self.structure)
+        with zopen(file_name, "wt") as file:
+            poscar = Poscar(self.structure)
 
             # use original name if it's been set (e.g. from Chgcar)
-            comment = getattr(self, "name", p.comment)
+            comment = getattr(self, "name", poscar.comment)
 
             lines = comment + "\n"
             lines += "   1.00000000000000\n"
             for vec in self.structure.lattice.matrix:
                 lines += f" {vec[0]:12.6f}{vec[1]:12.6f}{vec[2]:12.6f}\n"
             if not vasp4_compatible:
-                lines += "".join(f"{s:5}" for s in p.site_symbols) + "\n"
-            lines += "".join(f"{x:6}" for x in p.natoms) + "\n"
+                lines += "".join(f"{s:5}" for s in poscar.site_symbols) + "\n"
+            lines += "".join(f"{x:6}" for x in poscar.natoms) + "\n"
             lines += "Direct\n"
             for site in self.structure:
-                a, b, c = site.frac_coords
-                lines += f"{a:10.6f}{b:10.6f}{c:10.6f}\n"
+                dim, b, c = site.frac_coords
+                lines += f"{dim:10.6f}{b:10.6f}{c:10.6f}\n"
             lines += " \n"
-            f.write(lines)
-            a = self.dim
+            file.write(lines)
+            dim = self.dim
 
             def write_spin(data_type):
                 lines = []
                 count = 0
-                f.write(f"   {a[0]}   {a[1]}   {a[2]}\n")
-                for k, j, i in itertools.product(list(range(a[2])), list(range(a[1])), list(range(a[0]))):
+                file.write(f"   {dim[0]}   {dim[1]}   {dim[2]}\n")
+                for k, j, i in itertools.product(list(range(dim[2])), list(range(dim[1])), list(range(dim[0]))):
                     lines.append(_print_fortran_float(self.data[data_type][i, j, k]))
                     count += 1
                     if count % 5 == 0:
-                        f.write(" " + "".join(lines) + "\n")
+                        file.write(" " + "".join(lines) + "\n")
                         lines = []
                     else:
                         lines.append(" ")
                 if count % 5 != 0:
-                    f.write(" " + "".join(lines) + " \n")
-                f.write("".join(self.data_aug.get(data_type, [])))
+                    file.write(" " + "".join(lines) + " \n")
+
+                data = self.data_aug.get(data_type, [])
+                if isinstance(data, Iterable):
+                    file.write("".join(data))
 
             write_spin("total")
             if self.is_spin_polarized and self.is_soc:
@@ -3453,21 +3456,21 @@ class Chgcar(VolumetricData):
     def __init__(self, poscar, data, data_aug=None):
         """
         Args:
-            poscar (Poscar or Structure): Object containing structure.
+            poscar (Poscar | Structure): Object containing structure.
             data: Actual data.
             data_aug: Augmentation charge data.
         """
         # allow for poscar or structure files to be passed
         if isinstance(poscar, Poscar):
-            tmp_struct = poscar.structure
+            struct = poscar.structure
             self.poscar = poscar
             self.name = poscar.comment
         elif isinstance(poscar, Structure):
-            tmp_struct = poscar
+            struct = poscar
             self.poscar = Poscar(poscar)
             self.name = None
 
-        super().__init__(tmp_struct, data, data_aug=data_aug)
+        super().__init__(struct, data, data_aug=data_aug)
         self._distance_matrix = {}
 
     @staticmethod

tests/io/vasp/test_outputs.py

@@ -1552,6 +1552,22 @@ def test_as_dict_and_from_dict(self):
         )
 
 
+class TestAeccars(PymatgenTest):
+    # https://github.com/materialsproject/pymatgen/pull/3343
+    def test_read_write_file(self):
+        aeccar0_test = Chgcar.from_file(f"{TEST_FILES_DIR}/bader/AECCAR0.gz")
+        aeccar0_outpath = f"{self.tmp_path}/AECCAR0_test"
+        aeccar0_test.write_file(aeccar0_outpath)
+        aeccar0_read = Chgcar.from_file(aeccar0_outpath)
+        assert_allclose(aeccar0_test.data["total"], aeccar0_read.data["total"])
+
+        aeccar2 = Chgcar.from_file(f"{TEST_FILES_DIR}/bader/AECCAR2.gz")
+        aeccar2_outpath = f"{self.tmp_path}/AECCAR2_test"
+        aeccar2.write_file(aeccar2_outpath)
+        aeccar2_read = Chgcar.from_file(aeccar2_outpath)
+        assert_allclose(aeccar2.data["total"], aeccar2_read.data["total"])
+
+
 class TestElfcar(PymatgenTest):
     def test_init(self):
         elfcar = Elfcar.from_file(f"{TEST_FILES_DIR}/ELFCAR.gz")