Support both U_iso and B_iso in CIF files

README.md

@@ -22,7 +22,7 @@ This means you can write your own code to utilize advanced format features even
 
 | Format                 | Ext.    | Read               | Write              | Notes |
 | :--------------------- | :------ | :----------------: | :----------------: | :---- |
-| [CIF][cif]             | .cif    | :white_check_mark: | :x:                | CIF1 & CIF2. Isotropic B-factor only |
+| [CIF][cif]             | .cif    | :white_check_mark: | :white_check_mark: | CIF1 & CIF2. Isotropic B-factor only |
 | [XCrysDen][xsf]        | .xsf    | :white_check_mark: | :white_check_mark: |       |
 | [AtomEye CFG][cfg]     | .cfg    | :white_check_mark: | :white_check_mark: | Currently basic format only |
 | [Basic XYZ][xyz]       | .xyz    | :white_check_mark: | :white_check_mark: |       |
@@ -49,4 +49,4 @@ This means you can write your own code to utilize advanced format features even
 [docs-stable-badge]: https://img.shields.io/badge/docs-stable-blue
 [docs-stable-url]: https://hexane360.github.io/atomlib/latest/
 [commit-badge]: https://img.shields.io/github/last-commit/hexane360/atomlib
-[commit-url]: https://github.com/hexane360/atomlib/commits
+[commit-url]: https://github.com/hexane360/atomlib/commits

atomlib/io/__init__.py

@@ -6,6 +6,7 @@
 import typing as t
 
 import numpy
+import polars
 
 from .cif import CIF, CIFDataBlock
 from .xyz import XYZ, XYZFormat
@@ -50,18 +51,21 @@ def read_cif(f: t.Union[FileOrPath, CIF, CIFDataBlock], block: t.Union[int, str,
     logging.debug("cif data: %r", cif.data_dict)
 
     # TODO: support atom_site_Cartn_[xyz]
-    # TODO: support atom_site_B_iso_or_equiv
     df = cif.stack_tags('atom_site_fract_x', 'atom_site_fract_y', 'atom_site_fract_z',
-                        '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))
+                        'atom_site_type_symbol', 'atom_site_label', 'atom_site_occupancy',
+                        'atom_site_U_iso_or_equiv', 'atom_site_B_iso_or_equiv',
+                        rename=('x', 'y', 'z', 'symbol', 'label', 'frac_occupancy', 'wobble', 'wobble_B'),
+                        required=(True, True, True, False, False, False, False, False))
+    if 'wobble_B' in df.columns:
+        if 'wobble' in df.columns:
+            raise ValueError("CIF file specifies both 'atom_site_U_iso_or_equiv' and 'atom_site_B_iso_or_equiv'")
+        df = df.rename({'wobble_B': 'wobble'}) \
+            .with_columns(polars.col('wobble') * (3./8. / numpy.pi**2))
     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]])
+        # infer symbol from label, insert at beginning
+        df = df.insert_column(0, get_sym(get_elem(df['label'])))
     atoms = Atoms(df)
 
     # parse and apply symmetry

tests/input_files/AlN_Bfactor.cif

@@ -0,0 +1,30 @@
+_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_type_symbol
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_B_iso_or_equiv
+ Al  0.666592  0.333183  0.498984  0.1
+ Al  0.333183  0.666366  0.998681  0.2
+  N  0.666592  0.333183  0.880179  0.3
+  N  0.333183  0.666366  0.380482  0.4

tests/input_files/AlN_Ufactor.cif

@@ -0,0 +1,30 @@
+_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_type_symbol
+ _atom_site_fract_x
+ _atom_site_fract_y
+ _atom_site_fract_z
+ _atom_site_U_iso_or_equiv
+ Al  0.666592  0.333183  0.498984  0.1
+ Al  0.333183  0.666366  0.998681  0.2
+  N  0.666592  0.333183  0.880179  0.3
+  N  0.333183  0.666366  0.380482  0.4

tests/test_io.py

@@ -164,6 +164,16 @@ def test_cif_aln(aln):
     return aln
 
 
+@check_parse_structure('AlN_Ufactor.cif')
+def test_cif_aln_u_factor(aln: AtomCell):
+    return aln.with_wobble(numpy.array([0.1, 0.2, 0.3, 0.4]))
+
+
+@check_parse_structure('AlN_Bfactor.cif')
+def test_cif_aln_b_factor(aln: AtomCell):
+    return aln.with_wobble(numpy.array([0.1, 0.2, 0.3, 0.4]) * 3. / (8. * numpy.pi**2))
+
+
 @check_parse_structure('label_only.cif')
 def test_cif_aln_labelonly(aln):
     return AtomCell(