Get cell

structuretoolkit/__init__.py

@@ -51,6 +51,7 @@
     get_wrapped_coordinates,
     pymatgen_to_ase,
     select_index,
+    get_cell,
 )
 
 # Visualize

structuretoolkit/common/__init__.py

@@ -6,6 +6,7 @@
     get_vertical_length,
     get_wrapped_coordinates,
     select_index,
+    get_cell,
 )
 from structuretoolkit.common.pymatgen import (
     ase_to_pymatgen,

structuretoolkit/common/helper.py

@@ -2,7 +2,7 @@
 from ase.atoms import Atoms
 from ase.data import atomic_numbers
 from scipy.sparse import coo_matrix
-from typing import Optional
+from typing import Optional, Union
 
 
 def get_extended_positions(
@@ -226,3 +226,32 @@ def apply_strain(
     structure_copy.set_cell(cell, scale_atoms=True)
     if return_box:
         return structure_copy
+
+
+def get_cell(cell: Union[Atoms, list, np.ndarray, float]):
+    """
+    Get cell of an ase structure, or convert a float or a (3,)-array into a
+    orthogonal cell.
+
+    Args:
+        cell (Atoms|ndarray|list|float|tuple): Cell
+
+    Returns:
+        (3, 3)-array: Cell
+    """
+    if isinstance(cell, Atoms):
+        return cell.cell
+    # Convert float into (3,)-array. No effect if it is (3,3)-array or
+    # (3,)-array. Raises error if the shape is not correct
+    try:
+        cell = cell * np.ones(3)
+    except ValueError:
+        raise ValueError(
+            "cell must be a float, (3,)-ndarray/list/tuple or"
+            " (3,3)-ndarray/list/tuple"
+        )
+
+    if np.shape(cell) == (3, 3):
+        return cell
+    # Convert (3,)-array into (3,3)-array. Raises error if the shape is wrong
+    return cell * np.eye(3)

tests/test_helpers.py

@@ -0,0 +1,26 @@
+# coding: utf-8
+# Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
+# Distributed under the terms of "New BSD License", see the LICENSE file.
+
+import unittest
+import numpy as np
+from ase.build import bulk
+import structuretoolkit as stk
+
+
+class TestHelpers(unittest.TestCase):
+    def test_get_cell(self):
+        self.assertEqual((3 * np.eye(3)).tolist(), stk.get_cell(3).tolist())
+        self.assertEqual(
+            ([1, 2, 3] * np.eye(3)).tolist(), stk.get_cell([1, 2, 3]).tolist()
+        )
+        atoms = bulk("Fe")
+        self.assertEqual(
+            atoms.cell.tolist(), stk.get_cell(atoms).tolist()
+        )
+        with self.assertRaises(ValueError):
+            stk.get_cell(np.arange(4))
+
+
+if __name__ == "__main__":
+    unittest.main()