Fixed CIF parsing with label only

hexane360 · Apr 14, 2023 · 518b78d · 518b78d
1 parent a63206d
commit 518b78d
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 12 deletions.
diff --git a/structlib/io/__init__.py b/structlib/io/__init__.py
@@ -18,7 +18,7 @@
 from ..core import AtomCollection, Atoms, SimpleAtoms, AtomCell, Cell
 from ..types import to_vec3
 from ..transform import LinearTransform3D
-from ..elem import get_sym
+from ..elem import get_sym, get_elem
 from ..util import FileOrPath
 
 FileType = t.Literal['cif', 'xyz', 'xsf', 'cfg', 'lmp', 'mslice', 'qe']
@@ -39,9 +39,16 @@ def read_cif(f: t.Union[FileOrPath, CIF]) -> AtomCollection:
     logging.debug("cif data: %r", cif.data)
 
     df = cif.stack_tags('atom_site_fract_x', 'atom_site_fract_y', 'atom_site_fract_z',
-                        'atom_site_type_symbol', 'atom_site_occupancy', 'atom_site_U_iso_or_equiv',
-                        rename=('x', 'y', 'z', 'symbol', 'frac_occupancy', 'wobble'),
-                        required=(True, True, True, True, False))
+                        'atom_site_type_symbol', 'atom_site_label', 'atom_site_occupancy', 'atom_site_U_iso_or_equiv',
+                        rename=('x', 'y', 'z', 'symbol', 'label', 'frac_occupancy', 'wobble'),
+                        required=(True, True, True, False, False, False, False))
+    if 'symbol' not in df.columns:
+        if 'label' not in df.columns:
+            raise ValueError("Tag 'atom_site_type_symbol' or 'atom_site_label' missing from CIF file")
+        # infer symbol from label
+        df = df.with_columns(get_sym(get_elem(df['label'])))
+        # reorder symbol column
+        df = df.select([*df.columns[:3], df.columns[-1], *df.columns[3:-1]])
     atoms = Atoms(df)
 
     # parse and apply symmetry

diff --git a/tests/input_files/AlN.cif b/tests/input_files/AlN.cif
@@ -19,12 +19,11 @@ _cell_angle_gamma                 120.0000
 _cell_volume                      36.8853
 
 loop_
- _atom_site_label
  _atom_site_type_symbol
  _atom_site_fract_x
  _atom_site_fract_y
  _atom_site_fract_z
-Al1  Al  0.666592  0.333183  0.498984
-Al2  Al  0.333183  0.666366  0.998681
-N1  N  0.666592  0.333183  0.880179
-N2  N  0.333183  0.666366  0.380482
+ Al  0.666592  0.333183  0.498984
+ Al  0.333183  0.666366  0.998681
+  N  0.666592  0.333183  0.880179
+  N  0.333183  0.666366  0.380482
diff --git a/tests/input_files/label_only.cif b/tests/input_files/label_only.cif
@@ -0,0 +1,29 @@
+_audit_creation_date      'Jan. 28, 2023'
+_audit_creation_method    'Draft CIF file generated with Atomsk'
+
+_chemical_formula_iupac           'Al2 N2'
+_chemical_formula_moiety          'Al2 N2'
+_chemical_formula_sum             'Al2 N2'
+_chemical_formula_weight          81.977
+
+loop_
+_space_group_symop_operation_xyz
+'+x,+y,+z'
+
+_cell_length_a                    3.1300
+_cell_length_b                    3.1300
+_cell_length_c                    5.0200
+_cell_angle_alpha                 90.0000
+_cell_angle_beta                  90.0000
+_cell_angle_gamma                 120.0000
+_cell_volume                      36.8853
+
+loop_
+ _atom_site_label
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ Al(0)  0.666592  0.333183  0.498984
+ Al(1)  0.333183  0.666366  0.998681
+  N(0)  0.666592  0.333183  0.880179
+  N(1)  0.333183  0.666366  0.380482
diff --git a/tests/test_io.py b/tests/test_io.py
@@ -1,5 +1,3 @@
-from pathlib import Path
-
 import numpy
 import polars
 import pytest
@@ -166,4 +164,14 @@ def test_cfg_hex(aln_ortho):
         'v_y': [0.0, 0.0, 0.0, 0.0],
         'v_z': [0.0, 0.0, 0.0, 0.0],
         'mass': [26.9815, 14.0067, 26.9815, 14.0067],
-    }), aln_ortho, frame='local')
+    }), aln_ortho, frame='local')
+
+
+@check_parse_structure('label_only.cif')
+def test_cif_aln_labelonly(aln):
+    return AtomCell(
+        aln.get_atoms('local')
+            .with_column(polars.Series('label', ['Al(0)', 'Al(1)', 'N(0)', 'N(1)']))
+            .select(['x', 'y', 'z', 'symbol', 'label', 'elem']),
+        aln.cell, frame='local'
+    )