Merge pull request #388 from phonopy/refactor-read-write-fc3-fc2

Refactor reading and writing fc2 and fc3

phono3py/api_phono3py.py

@@ -353,15 +353,18 @@ class Phono3py:
 
     @property
     def fc3_nonzero_indices(self) -> NDArray | None:
-        """Return non-zero indices of fc3.
+        """Setter and getter of non-zero indices of fc3.
 
-        ndarray
-            Non-zero indices of fc3. This is available only when
-            SymfcFCSolver is used.
+        ndarray, optional
+            Non-zero indices of fc3.
 
         """
         return self._fc3_nonzero_indices
 
+    @fc3_nonzero_indices.setter
+    def fc3_nonzero_indices(self, fc3_nonzero_indices):
+        self._fc3_nonzero_indices = fc3_nonzero_indices
+
     @property
     def fc2(self) -> NDArray | None:
         """Setter and getter of second order force constants (fc2).

phono3py/cui/load.py

@@ -42,6 +42,7 @@ from typing import Optional, Union
 
 import numpy as np
 import phonopy.cui.load_helper as load_helper
+from numpy.typing import NDArray
 from phonopy.harmonic.force_constants import show_drift_force_constants
 from phonopy.interface.calculator import get_calculator_physical_units
 from phonopy.physical_units import get_physical_units
@@ -491,7 +492,7 @@ def compute_force_constants_from_datasets(
 
 def _load_fc3(
     ph3py: Phono3py,
-    fc3_filename: Optional[os.PathLike] = None,
+    fc3_filename: str | os.PathLike | None = None,
     log_level: int = 0,
 ):
     p2s_map = ph3py.primitive.p2s_map
@@ -500,21 +501,31 @@ def _load_fc3(
     else:
         _fc3_filename = fc3_filename
     fc3 = read_fc3_from_hdf5(filename=_fc3_filename, p2s_map=p2s_map)
-    _check_fc3_shape(ph3py, fc3, filename=_fc3_filename)
-    if log_level:
-        print(f'fc3 was read from "{_fc3_filename}".')
-    ph3py.fc3 = fc3
+    if isinstance(fc3, dict):
+        # fc3 is read from a file with type-1 format.
+        assert "fc3" in fc3
+        _check_fc3_shape(ph3py, fc3["fc3"], filename=_fc3_filename)
+        ph3py.fc3 = fc3["fc3"]
+        assert "fc3_nonzero_indices" in fc3
+        ph3py.fc3_nonzero_indices = fc3["fc3_nonzero_indices"]
+        if log_level:
+            print(f'fc3 and nonzero indices were read from "{_fc3_filename}".')
+    else:
+        _check_fc3_shape(ph3py, fc3, filename=_fc3_filename)
+        ph3py.fc3 = fc3
+        if log_level:
+            print(f'fc3 was read from "{_fc3_filename}".')
 
 
 def _select_and_load_dataset(
     ph3py: Phono3py,
-    ph3py_yaml: Optional[Phono3pyYaml] = None,
-    forces_fc3_filename: Optional[Union[os.PathLike, Sequence]] = None,
-    phono3py_yaml_filename: Optional[os.PathLike] = None,
-    cutoff_pair_distance: Optional[float] = None,
-    calculator: Optional[str] = None,
+    ph3py_yaml: Phono3pyYaml | None = None,
+    forces_fc3_filename: os.PathLike | Sequence | None = None,
+    phono3py_yaml_filename: os.PathLike | None = None,
+    cutoff_pair_distance: float | None = None,
+    calculator: str | None = None,
     log_level: int = 0,
-) -> Optional[dict]:
+) -> dict | None:
     dataset = None
     if (
         ph3py_yaml is not None
@@ -561,7 +572,7 @@ def _select_and_load_dataset(
 
 
 def _load_fc2(
-    ph3py: Phono3py, fc2_filename: Optional[os.PathLike] = None, log_level: int = 0
+    ph3py: Phono3py, fc2_filename: os.PathLike | None = None, log_level: int = 0
 ):
     phonon_p2s_map = ph3py.phonon_primitive.p2s_map
     if fc2_filename is None:
@@ -667,7 +678,7 @@ def _get_dataset_for_fc2(
     return dataset
 
 
-def _check_fc2_shape(ph3py: Phono3py, fc2, filename="fc2.hdf5"):
+def _check_fc2_shape(ph3py: Phono3py, fc2, filename: str | os.PathLike = "fc2.hdf5"):
     if ph3py.phonon_supercell_matrix is None:
         smat = ph3py.supercell_matrix
     else:
@@ -675,7 +686,9 @@ def _check_fc2_shape(ph3py: Phono3py, fc2, filename="fc2.hdf5"):
     _check_fc_shape(ph3py, fc2, smat, filename)
 
 
-def _check_fc3_shape(ph3py: Phono3py, fc3, filename="fc3.hdf5"):
+def _check_fc3_shape(
+    ph3py: Phono3py, fc3: NDArray, filename: str | os.PathLike = "fc3.hdf5"
+):
     smat = ph3py.supercell_matrix
     _check_fc_shape(ph3py, fc3, smat, filename)
 

phono3py/cui/phono3py_script.py

@@ -688,7 +688,7 @@ def _store_force_constants(ph3py: Phono3py, settings: Phono3pySettings, log_leve
     if not read_fc2:
         write_fc2_to_hdf5(
             ph3py.fc2,
-            p2s_map=ph3py.primitive.p2s_map,
+            p2s_map=ph3py.phonon_primitive.p2s_map,
             physical_unit="eV/angstrom^2",
             compression=settings.hdf5_compression,
         )

phono3py/file_IO.py

@@ -35,6 +35,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 from __future__ import annotations
 
+import os
 import pathlib
 import warnings
 from collections.abc import Sequence
@@ -341,7 +342,9 @@ def write_fc3_to_hdf5(
             w.create_dataset("p2s_map", data=p2s_map)
 
 
-def read_fc3_from_hdf5(filename="fc3.hdf5", p2s_map=None):
+def read_fc3_from_hdf5(
+    filename: str | os.PathLike = "fc3.hdf5", p2s_map: NDArray | None = None
+) -> NDArray | dict:
     """Read fc3 from fc3.hdf5.
 
     fc3 can be in full or compact format. They are distinguished by
@@ -353,20 +356,44 @@ def read_fc3_from_hdf5(filename="fc3.hdf5", p2s_map=None):
 
     """
     with h5py.File(filename, "r") as f:
-        fc3 = f["fc3"][:]
+        if "fc3" not in f:
+            raise KeyError(
+                f"{filename} does not have 'fc3' dataset. "
+                "This file is not a valid fc3.hdf5."
+            )
+        fc3: NDArray = f["fc3"][:]  # type: ignore
+        if fc3.dtype == np.dtype("double") and fc3.flags.c_contiguous:
+            pass
+        else:
+            raise TypeError(
+                f"{filename} has to be read by h5py as numpy ndarray of "
+                "dtype='double' and c_contiguous."
+            )
+
         if "p2s_map" in f:
             p2s_map_in_file = f["p2s_map"][:]
             check_force_constants_indices(
                 fc3.shape[:2], p2s_map_in_file, p2s_map, filename
             )
-        if fc3.dtype == np.dtype("double") and fc3.flags.c_contiguous:
+
+        fc3_nonzero_indices = None  # type: ignore
+        if "fc3_nonzero_indices" in f:
+            fc3_nonzero_indices: NDArray = f["fc3_nonzero_indices"][:]  # type: ignore
+            if (
+                fc3_nonzero_indices.dtype == np.dtype("byte")
+                and fc3_nonzero_indices.flags.c_contiguous
+            ):
+                pass
+            else:
+                raise TypeError(
+                    f"{filename} has to be read by h5py as numpy ndarray of "
+                    "dtype='byte' and c_contiguous."
+                )
+
+        if fc3_nonzero_indices is None:
             return fc3
         else:
-            msg = (
-                "%s has to be read by h5py as numpy ndarray of "
-                "dtype='double' and c_contiguous." % filename
-            )
-            raise TypeError(msg)
+            return {"fc3": fc3, "fc3_nonzero_indices": fc3_nonzero_indices}
 
 
 def write_fc2_to_hdf5(
@@ -429,7 +456,7 @@ def read_fc2_from_hdf5(filename="fc2.hdf5", p2s_map=None):
 
 def write_datasets_to_hdf5(
     dataset: dict,
-    phonon_dataset: dict = None,
+    phonon_dataset: dict | None = None,
     filename: str = "datasets.hdf5",
     compression: str = "gzip",
 ):

test/cui/test_phono3py_load_script.py

@@ -58,15 +58,32 @@ def test_phono3py_load():
     """Test phono3py-load script."""
     # Check sys.exit(0)
     argparse_control = _get_phono3py_load_args(
-        cwd / ".." / "phono3py_params_Si-111-222.yaml"
+        cwd / ".." / "phono3py_params_Si-111-222.yaml",
+    )
+    with pytest.raises(SystemExit) as excinfo:
+        main(**argparse_control)
+    assert excinfo.value.code == 0
+
+    argparse_control = _get_phono3py_load_args(
+        cwd_called / "phono3py.yaml",
+        is_bterta=True,
+        temperatures=[
+            "300",
+        ],
+        mesh_numbers=["5", "5", "5"],
     )
     with pytest.raises(SystemExit) as excinfo:
         main(**argparse_control)
     assert excinfo.value.code == 0
 
     # Clean files created by phono3py-load script.
-    for created_filename in ("phono3py.yaml", "fc2.hdf5", "fc3.hdf5"):
-        file_path = pathlib.Path(cwd_called / created_filename)
+    for created_filename in (
+        "phono3py.yaml",
+        "fc2.hdf5",
+        "fc3.hdf5",
+        "kappa-m555.hdf5",
+    ):
+        file_path = cwd_called / created_filename
         if file_path.exists():
             file_path.unlink()