Add support for non-diagonal supercells

docs/source/user_guide/command_line.rst

@@ -314,7 +314,7 @@ Calculate phonons with a 2x2x2 supercell, after geometry optimization (using the
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2 2 2 --minimize --arch mace_mp --model-path small
+    janus phonons --struct tests/data/NaCl.cif --supercell "2 2 2" --minimize --arch mace_mp --model-path small
 
 
 This will save the Phonopy parameters, including displacements and force constants, to ``NaCl-phonopy.yml`` and ``NaCl-force_constants.hdf5``,
@@ -324,7 +324,7 @@ Additionally, the ``--bands`` option can be added to calculate the band structur
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2 2 2 --minimize --arch mace_mp --model-path small --bands
+    janus phonons --struct tests/data/NaCl.cif --supercell "2 2 2" --minimize --arch mace_mp --model-path small --bands
 
 
 If you need eigenvectors and group velocities written, add the ``--write-full`` option. This will generate a much larger file, but can be used to visualise phonon modes.
@@ -333,12 +333,19 @@ Further calculations, including thermal properties, DOS, and PDOS, can also be c
 
 .. code-block:: bash
 
-    janus phonons --struct tests/data/NaCl.cif --supercell 2 3 4 --dos --pdos --thermal --temp-start 0 --temp-end 300 --temp-step 50
+    janus phonons --struct tests/data/NaCl.cif --supercell "2 3 4" --dos --pdos --thermal --temp-start 0 --temp-end 300 --temp-step 50
 
 
 This will create additional output files: ``NaCl-thermal.dat`` for the thermal properties (heat capacity, entropy, and free energy)
 between 0K and 300K, ``NaCl-dos.dat`` for the DOS, and ``NaCl-pdos.dat`` for the PDOS.
 
+To define the supercell, the ``--supercell`` option can be used, which *must* be passed in as a space-separated string.
+Similar to Phonopy, the supercell matrix can be defined in three ways:
+
+1. One integer (``--supercell "2"``) specifying all diagonal elements.
+2. Three integers (``--supercell "2 2 2"``) specifying each individual diagonal element.
+3. Nine integers (``--supercell "2 0 0 0 2 0 0 0 2"``) specifying all elements, filling the matrix row-wise.
+
 For all options, run ``janus phonons --help``.
 
 

janus_core/calculations/phonons.py

@@ -59,7 +59,14 @@ class Phonons(BaseCalculation):
     calcs : Optional[MaybeSequence[PhononCalcs]]
         Phonon calculations to run. Default calculates force constants only.
     supercell : MaybeList[int]
-        Size of supercell for calculation. Default is 2.
+        The size of a supercell for calculation, or the supercell itself.
+        If a single number is provided, it is interpreted as the size, so a
+        diagonal supercell of that size in all dimensions is constructed.
+        If three values are provided, they are interpreted as the diagonal
+        values of a diagonal supercell. If nine values are provided, they
+        are assumed to be the full supercell matrix in the style of Phonopy,
+        so the first three values will be used as the first row, the second
+        three as the second row, etc. Default is 2.
     displacement : float
         Displacement for force constants calculation, in A. Default is 0.01.
     mesh : tuple[int, int, int]
@@ -194,7 +201,14 @@ def __init__(
         calcs : Optional[MaybeSequence[PhononCalcs]]
             Phonon calculations to run. Default calculates force constants only.
         supercell : MaybeList[int]
-            Size of supercell for calculation. Default is 2.
+            The size of a supercell for calculation, or the supercell itself.
+            If a single number is provided, it is interpreted as the size, so a
+            diagonal supercell of that size in all dimensions is constructed.
+            If three values are provided, they are interpreted as the diagonal
+            values of a diagonal supercell. If nine values are provided, they
+            are assumed to be the full supercell matrix in the style of Phonopy,
+            so the first three values will be used as the first row, the second
+            three as the second row, etc. Default is 2.
         displacement : float
             Displacement for force constants calculation, in A. Default is 0.01.
         mesh : tuple[int, int, int]
@@ -249,8 +263,14 @@ def __init__(
 
         # Ensure supercell is a valid list
         self.supercell = [supercell] * 3 if isinstance(supercell, int) else supercell
-        if len(self.supercell) != 3:
-            raise ValueError("`supercell` must be an integer, or list of length 3")
+        if len(self.supercell) not in [3, 9]:
+            raise ValueError(
+                "`supercell` must be an integer, or list of length 3 or 9. A list of "
+                "length 3 must specify the values of a diagonal supercell, and a list "
+                "of length 9 must specify all values of a full supercell matrix, where "
+                "the first three values are the first row, the second three are the "
+                "second row, etc."
+            )
 
         # Read last image by default
         read_kwargs.setdefault("index", -1)
@@ -370,11 +390,19 @@ def calc_force_constants(
 
         cell = self._ASE_to_PhonopyAtoms(self.struct)
 
-        supercell_matrix = (
-            (self.supercell[0], 0, 0),
-            (0, self.supercell[1], 0),
-            (0, 0, self.supercell[2]),
-        )
+        if len(self.supercell) == 3:
+            supercell_matrix = (
+                (self.supercell[0], 0, 0),
+                (0, self.supercell[1], 0),
+                (0, 0, self.supercell[2]),
+            )
+        else:
+            supercell_matrix = (
+                tuple(self.supercell[:3]),
+                tuple(self.supercell[3:6]),
+                tuple(self.supercell[6:]),
+            )
+
         phonon = phonopy.Phonopy(cell, supercell_matrix)
         phonon.generate_displacements(distance=self.displacement)
         disp_supercells = phonon.supercells_with_displacements

janus_core/cli/phonons.py

@@ -29,8 +29,15 @@ def phonons(
     ctx: Context,
     struct: StructPath,
     supercell: Annotated[
-        tuple[int, int, int], Option(help="Supercell lattice vectors.")
-    ] = (2, 2, 2),
+        str,
+        Option(
+            help="Supercell matrix, in the Phonopy style. Must be passed as a string "
+            "in one of three forms: single integer ('2'), which specifies all "
+            "diagonal elements; three integers ('1 2 3'), which specifies each "
+            "individual diagonal element; or nine values ('1 2 3 4 5 6 7 8 9'), "
+            "which specifies all elements, filling the matrix row-wise."
+        ),
+    ] = "2 2 2",
     displacement: Annotated[
         float, Option(help="Displacement for force constants calculation, in A.")
     ] = 0.01,
@@ -114,8 +121,12 @@ def phonons(
         Typer (Click) Context. Automatically set.
     struct : Path
         Path of structure to simulate.
-    supercell : tuple[int, int, int]
-        Supercell lattice vectors. Default is (2, 2, 2).
+    supercell : str
+        Supercell matrix, in the Phonopy style. Must be passed as a string in one of
+        three forms: single integer ('2'), which specifies all diagonal elements;
+        three integers ('1 2 3'), which specifies each individual diagonal element;
+        or nine values ('1 2 3 4 5 6 7 8 9'), which specifies all elements, filling the
+        matrix row-wise.
     displacement : float
         Displacement for force constants calculation, in A. Default is 0.01.
     mesh : tuple[int, int, int]
@@ -206,6 +217,22 @@ def phonons(
         raise ValueError("'fmax' must be passed through the --fmax option")
     minimize_kwargs["fmax"] = fmax
 
+    try:
+        supercell = [int(x) for x in supercell.split()]
+    except ValueError as exc:
+        raise ValueError(
+            "Please pass lattice vectors as integers in the form '1 2 3'"
+        ) from exc
+
+    supercell_length = len(supercell)
+    if supercell_length == 1:
+        supercell = supercell[0]
+    elif supercell_length not in [3, 9]:
+        raise ValueError(
+            "Please pass lattice vectors as space-separated integers in quotes. "
+            "For example, '1 2 3'."
+        )
+
     calcs = []
     if bands:
         calcs.append("bands")

tests/test_phonons_cli.py

@@ -228,9 +228,7 @@ def test_plot(tmp_path):
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--supercell",
-            1,
-            1,
-            1,
+            "1 1 1",
             "--pdos",
             "--dos",
             "--bands",
@@ -258,7 +256,15 @@ def test_plot(tmp_path):
     assert phonon_summary["inputs"]["calcs"][2] == "pdos"
 
 
-def test_supercell(tmp_path):
+test_data = [
+    ("1", [[1, 0, 0], [0, 1, 0], [0, 0, 1]]),
+    ("1 2 3", [[1, 0, 0], [0, 2, 0], [0, 0, 3]]),
+    ("1 1 0 -1 1 0 0 0 2", [[1, 1, 0], [-1, 1, 0], [0, 0, 2]]),
+]
+
+
+@pytest.mark.parametrize("supercell,supercell_matrix", test_data)
+def test_supercell(supercell, supercell_matrix, tmp_path):
     """Test setting the supercell."""
     file_prefix = tmp_path / "NaCl"
     param_file = tmp_path / "NaCl-phonopy.yml"
@@ -270,9 +276,7 @@ def test_supercell(tmp_path):
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--supercell",
-            1,
-            2,
-            3,
+            supercell,
             "--no-hdf5",
             "--file-prefix",
             file_prefix,
@@ -286,10 +290,13 @@ def test_supercell(tmp_path):
 
     assert "supercell_matrix" in params
     assert len(params["supercell_matrix"]) == 3
-    assert params["supercell_matrix"] == [[1, 0, 0], [0, 2, 0], [0, 0, 3]]
+    assert params["supercell_matrix"] == supercell_matrix
+
+
+test_data = ["2x2x2", "2.1 2.1 2.1", "2 2 a", "2 2", "2 2 2 2 2 2"]
 
 
-@pytest.mark.parametrize("supercell", [(2,), (2, 2), (2, 2, "a"), ("2x2x2",)])
+@pytest.mark.parametrize("supercell", test_data)
 def test_invalid_supercell(supercell, tmp_path):
     """Test errors are raise for invalid supercells."""
     file_prefix = tmp_path / "test"
@@ -301,7 +308,7 @@ def test_invalid_supercell(supercell, tmp_path):
             "--struct",
             DATA_PATH / "NaCl.cif",
             "--supercell",
-            *supercell,
+            supercell,
             "--file-prefix",
             file_prefix,
         ],
@@ -379,9 +386,7 @@ def test_valid_traj_input(read_kwargs, tmp_path):
             "--struct",
             DATA_PATH / "NaCl-traj.xyz",
             "--supercell",
-            1,
-            1,
-            1,
+            "1 1 1",
             "--read-kwargs",
             read_kwargs,
             "--no-hdf5",