From 564f1b07a8ae8d2a58de56ca358631d2a83ec18d Mon Sep 17 00:00:00 2001
From: Atsushi Togo <atz.togo@gmail.com>
Date: Mon, 22 Jul 2024 15:18:01 +0900
Subject: [PATCH] Add --rd option to velph-phono3py-init

---
 src/phelel/velph/cli/phono3py/__init__.py     | 13 ++++-
 .../cli/phono3py/init/test_phono3py_init.py   | 58 +++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 test/velph/cli/phono3py/init/test_phono3py_init.py

diff --git a/src/phelel/velph/cli/phono3py/__init__.py b/src/phelel/velph/cli/phono3py/__init__.py
index 02664ae..4898b4b 100644
--- a/src/phelel/velph/cli/phono3py/__init__.py
+++ b/src/phelel/velph/cli/phono3py/__init__.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import pathlib
+from typing import Optional
 
 import click
 import tomli
@@ -29,13 +30,21 @@ def cmd_phono3py():
     type=click.Path(),
     default="velph.toml",
 )
+@click.option(
+    "--rd",
+    "random_displacements",
+    nargs=1,
+    default=None,
+    type=int,
+    help="Number of snapshots of supercells with random directional displacement.",
+)
 @click.help_option("-h", "--help")
-def cmd_init(toml_filename: str):
+def cmd_init(toml_filename: str, random_displacements: Optional[int]):
     """Generate displacements and write phelel_disp.yaml."""
     with open(toml_filename, "rb") as f:
         toml_dict = tomli.load(f)
 
-    ph3py = run_init(toml_dict)
+    ph3py = run_init(toml_dict, number_of_snapshots=random_displacements)
 
     phono3py_yaml_filename = pathlib.Path("phono3py/phono3py_disp.yaml")
     phono3py_yaml_filename.parent.mkdir(parents=True, exist_ok=True)
diff --git a/test/velph/cli/phono3py/init/test_phono3py_init.py b/test/velph/cli/phono3py/init/test_phono3py_init.py
new file mode 100644
index 0000000..c512da9
--- /dev/null
+++ b/test/velph/cli/phono3py/init/test_phono3py_init.py
@@ -0,0 +1,58 @@
+"""Tests velph-phono3py-init."""
+
+import numpy as np
+import pytest
+import tomli
+from phelel.velph.cli.phono3py.init import run_init
+
+
+@pytest.mark.parametrize("distance", [0.03, 0.05])
+def test_phono3py_init_random_displacements(distance: float):
+    """Test of plusminus and diagonal with Ti."""
+    toml_str = f"""title = "VASP el-ph settings"
+
+[phelel]
+supercell_dimension = [4, 4, 2]
+amplitude = {distance}
+fft_mesh = [18, 18, 28]
+[vasp.supercell.incar]
+lwap = true
+isym = 0
+kpar = 2
+ncore = 24
+ismear = 0
+sigma = 0.2
+ediff = 1e-08
+encut = 329.532
+prec = "accurate"
+lreal = false
+lwave = false
+lcharg = false
+addgrid = true
+lsorbit = true
+[vasp.supercell.kpoints]
+mesh = [6, 6, 7]
+
+[unitcell]
+lattice = [
+  [     2.930720886111760,     0.000000000000000,     0.000000000000000 ], # a
+  [    -1.465360443055880,     2.538078738774425,     0.000000000000000 ], # b
+  [     0.000000000000000,     0.000000000000000,     4.646120482318025 ], # c
+]
+[[unitcell.points]]  # 1
+symbol = "Ti"
+coordinates = [  0.333333333333336,  0.666666666666664,  0.250000000000000 ]
+magnetic_moment = [ 0.00000000, 0.00000000, 0.00000000 ]
+[[unitcell.points]]  # 2
+symbol = "Ti"
+coordinates = [  0.666666666666664,  0.333333333333336,  0.750000000000000 ]
+magnetic_moment = [ 0.00000000, 0.00000000, 0.00000000 ]
+"""
+
+    toml_dict = tomli.loads(toml_str)
+    ph3 = run_init(toml_dict, number_of_snapshots=10)
+    np.testing.assert_array_equal(ph3.supercell_matrix, np.diag([4, 4, 2]))
+    natom = len(ph3.supercell)
+    assert natom == 64
+    assert ph3.displacements.shape == (10, natom, 3)
+    np.testing.assert_allclose(np.linalg.norm(ph3.displacements, axis=2), distance)