Implements sqrt(2) x sqrt(2) supercell construction script (#27)

* orthorhombic cells

* Ruff fix

* Ruff fix

* Ruff fix

* Fixing mypy error

* Fixing mypy error

* Fixing mypy error

* Fixing mypy error

* Added sqrt(2) x sqrt(2) supercell generator script

* Ignoring pymatgen type hints until setuptools is fixed

* Type fix

vasppy/doscar.py

@@ -274,7 +274,7 @@ def plot_pdos(
             fig, ax = plt.subplots(1, 1, figsize=(8.0, 3.0))
         else:
             fig = None
-        assert(isinstance(ax, Axes))
+        assert isinstance(ax, Axes)
         if not colors:
             colors = mcd.TABLEAU_COLORS
         assert isinstance(colors, Iterable)

vasppy/poscar.py

@@ -5,7 +5,7 @@
 from vasppy import cell
 from vasppy.units import angstrom_to_bohr
 from pymatgen.core import Lattice as pmg_Lattice # type: ignore
-from pymatgen.core import Structure as pmg_Structure
+from pymatgen.core import Structure as pmg_Structure # type: ignore
 from pymatgen.io.cif import CifWriter  # type: ignore
 from collections import Counter
 
@@ -175,15 +175,19 @@ def output(self, coordinate_type="Direct", opts=None):
         if opts is None:
             opts = {}
         if not opts.get("coordinates_only"):
-            self.output_header(coordinate_type=coordinate_type)
+            self.output_header(coordinate_type=coordinate_type, opts=opts)
         self.output_coordinates_only(coordinate_type=coordinate_type, opts=opts)
 
     def output_header(self, coordinate_type="Direct", opts=None):
         if opts is None:
             opts = {}
         print(self.title)
         print(self.scaling)
-        for row in self.cell.matrix:
+        if opts.get("orthorhombic"):
+            matrix = self.cell.matrix * np.eye(3)
+        else:
+            matrix = self.cell.matrix
+        for row in matrix:
             print("".join(["   {: .10f}".format(element) for element in row]))
         print(" ".join(self.atoms))
         print(" ".join([str(n) for n in self.atom_numbers]))

vasppy/scripts/proc_poscar.py

@@ -3,7 +3,6 @@
 from vasppy.poscar import Poscar
 import argparse
 
-
 def parse_command_line_arguments():
     # command line arguments
     parser = argparse.ArgumentParser(description="Manipulates VASP POSCAR files")
@@ -49,6 +48,12 @@ def parse_command_line_arguments():
         action="store_true",
         help="label coordinates with atom number",
     )
+    parser.add_argument(
+        "-o",
+        "--orthorhombic",
+        action="store_true",
+        help="force orthorhombic cell matrix (set off-diagonal elements to zero)",
+    )
     parser.add_argument(
         "--scale",
         action="store_true",
@@ -96,6 +101,7 @@ def main():
         "numbered": args.number_atoms,
         "coordinates_only": args.coordinates_only,
         "selective": args.selective,
+        "orthorhombic": args.orthorhombic,
     }
     poscar.output(coordinate_type=coordinate_type, opts=output_opts)
 

vasppy/scripts/r2r2_expansion.py

@@ -0,0 +1,48 @@
+#! /usr/bin/env python3
+
+from vasppy.poscar import Poscar
+import argparse
+import numpy as np
+from typing import Literal
+
+def parse_command_line_arguments():
+    # command line arguments
+    parser = argparse.ArgumentParser(description="Generate a sqrt(2) x sqrt(2) supercell from a VASP POSCAR")
+    parser.add_argument("poscar", help="filename of the VASP POSCAR to be processed")
+    parser.add_argument(
+        "-a",
+        "--axis",
+        choices=['x', 'y', 'z'],
+        type=str,
+        help="normal vector for sqrt(2) x sqrt(2) expansion",
+        required=True,
+    )
+    args = parser.parse_args()
+    return args
+
+def sqrt2_by_sqrt2_expansion(
+        poscar: Poscar,
+        axis: Literal['x', 'y', 'z']
+        ) -> Poscar:
+        axis_vectors = {'x': [1, 0, 0],
+                        'y': [0, 1, 0],
+                        'z': [0, 0, 1]}
+        poscar.cell.rotate(axis=axis_vectors[axis], theta=np.pi/4)
+        poscar = poscar.replicate(2, 2, 1)
+        cc = poscar.cartesian_coordinates()
+        poscar.cell.matrix = np.diag(poscar.cell.matrix.diagonal())
+        poscar.coordinates = cc.dot(np.linalg.inv(poscar.cell.matrix))
+        return poscar
+
+def main():
+    args = parse_command_line_arguments()
+    poscar = Poscar.from_file(args.poscar)
+    sqrt2_by_sqrt2_expansion(
+        poscar = poscar,
+        axis = args.axis
+    ).output()
+
+if __name__ == '__main__':
+    main()
+
+

vasppy/xdatcar.py

@@ -5,6 +5,8 @@
 
 
 class Xdatcar:
+    """Class for parsing and working with VASP XDATCAR files."""
+
     lines_offset = 9
 
     def __init__(self):
@@ -20,7 +22,16 @@ def __init__(self):
         self.poscar = []
         self.poscar.append(Poscar())
 
-    def read_from(self, filename):
+    def read_from(self, filename: str) -> None:
+        """Read XDATCAR data from a VASP XDATCAR file.
+
+        Args:
+            filename (str): The XDATCAR file to read.
+
+        Returns:
+            None
+
+        """
         self.poscar[0].read_from(filename)
         with open(filename) as f:
             lines = f.read()