✨ Correctly transform the U matrix when applying transformations

picometer/atom.py

@@ -119,6 +119,20 @@ def fract_xyz(self) -> np.ndarray:
     def cart_xyz(self) -> np.ndarray:
         return self.orthogonalise(self.fract_xyz)
 
+    @property
+    def fract_uij(self) -> np.ndarray:
+        """Return a 3D array i.e. stack of 3x3 fract. displacement tensors."""
+        t = self.table
+        default = pd.Series([np.nan] * len(t), index=t.index)
+        uij = np.zeros((len(t), 3, 3), dtype=np.float64)
+        uij[:, 0, 0] = t.get('U11', default).to_numpy(dtype=np.float64)
+        uij[:, 1, 1] = t.get('U22', default).to_numpy(dtype=np.float64)
+        uij[:, 2, 2] = t.get('U33', default).to_numpy(dtype=np.float64)
+        uij[:, 0, 1] = uij[:, 1, 0] = t.get('U12', default).to_numpy(dtype=np.float64)
+        uij[:, 0, 2] = uij[:, 2, 0] = t.get('U13', default).to_numpy(dtype=np.float64)
+        uij[:, 1, 2] = uij[:, 2, 1] = t.get('U23', default).to_numpy(dtype=np.float64)
+        return uij
+
     def fractionalise(self, cart_xyz: np.ndarray) -> np.ndarray:
         """Multiply 3xN vector by crystallographic matrix to get fract coord"""
         return np.linalg.inv(self.base.A_d.T) @ cart_xyz
@@ -158,6 +172,17 @@ def transform(self, symm_op_code: str) -> 'AtomSet':
         data['fract_x'] = fract_xyz[:, 0]
         data['fract_y'] = fract_xyz[:, 1]
         data['fract_z'] = fract_xyz[:, 2]
+        if {'U11', 'U22', 'U33', 'U12', 'U13', 'U23'}.issubset(data.columns):
+            uij = self.fract_uij  # shape: (n_atoms, 3, 3)
+            mask = ~np.isnan(uij).all(axis=(1, 2))  # atoms with defined Uij
+            if np.any(mask):
+                uij_rot = (s := symm_op.tf) @ uij[mask] @ s.T
+                data.loc[mask, 'U11'] = uij_rot[:, 0, 0]
+                data.loc[mask, 'U22'] = uij_rot[:, 1, 1]
+                data.loc[mask, 'U33'] = uij_rot[:, 2, 2]
+                data.loc[mask, 'U12'] = uij_rot[:, 0, 1]
+                data.loc[mask, 'U13'] = uij_rot[:, 0, 2]
+                data.loc[mask, 'U23'] = uij_rot[:, 1, 2]
         return self.__class__(self.base, data)
 
     @property

picometer/instructions.py

@@ -15,6 +15,7 @@
 from typing import Any, Union, Protocol
 
 from numpy import rad2deg
+import numpy as np
 import pandas as pd
 import yaml
 
@@ -218,6 +219,31 @@ def _load_model_state(self, cif_path, block_name):
         label = cif_path + (':' + block_name if block_name else '')
         self.processor.model_states[label] = ModelState(atoms=atoms)
         logger.info(f'Loaded model state {label}')
+
+        if self.processor.settings['complete_uiso_from_umatrix']:
+            if 'U11' in atoms.table.columns:
+                if 'Uiso' not in atoms.table.columns:
+                    atoms.table['Uiso'] = pd.NA
+                u_equiv = atoms.table[['U11', 'U22', 'U33']].mean(axis=1)
+                atoms.table['Uiso'].fillna(u_equiv, inplace=True)
+
+        if self.processor.settings['complete_umatrix_from_uiso']:
+            u_columns = ['U11', 'U12', 'U13', 'U22', 'U23', 'U33']
+            if 'Uiso' in atoms.table.columns:
+                for col in u_columns:
+                    if col not in atoms.table.columns:
+                        atoms.table[col] = pd.NA
+                mask1 = atoms.table['Uiso'].notna()
+                mask2 = atoms.table[['U11', 'U22', 'U33']].isna().all(axis=1)
+                # based on http://dx.doi.org/10.1107/S0021889802008580
+                n_mat = np.diag([atoms.base.a_r, atoms.base.b_r, atoms.base.c_r])
+                n_inv = np.linalg.inv(n_mat)
+                u_star = (m := np.linalg.inv(atoms.base.A_d.T)) @ m.T
+                u_cif = n_inv @ u_star @ n_inv.T
+                for atom_label in atoms.table.index[mask1 & mask2]:
+                    u_atom = atoms.table.at[atom_label, 'Uiso'] * u_cif
+                    atoms.table.loc[atom_label, u_columns] = u_atom[np.triu_indices(3)]
+
         if not self.processor.settings['auto_write_unit_cell']:
             return
         et = self.processor.evaluation_table
@@ -326,6 +352,7 @@ class DisplacementInstructionHandler(SerialInstructionHandler):
 
     def handle_one(self, instruction: Instruction, ms_key: str, ms: ModelState) -> None:
         focus = ms.nodes.locate(self.processor.selection)
+        assert len(focus) > 0
         for label, displacements in focus.table.iterrows():
             for suffix in 'Uiso U11 U22 U33 U23 U13 U12'.split():
                 label_ = label + '_' + suffix

picometer/settings.py

@@ -19,6 +19,8 @@ class DefaultSettings:
     """Store default values of all settings. Use `AnyValue` if no default."""
     auto_write_unit_cell: bool = True
     clear_selection_after_use: bool = True
+    complete_uiso_from_umatrix: bool = False
+    complete_umatrix_from_uiso: bool = False
 
     @classmethod
     def get_field(cls, key: str) -> Field:

picometer/settings.yaml

@@ -1,3 +1,5 @@
 settings:
   auto_write_unit_cell: True
   clear_selection_after_use: True
+  complete_uiso_from_umatrix: False
+  complete_umatrix_from_uiso: False