Improved mslice output. Added tests

hexane360 · May 20, 2023 · 7c1eed8 · 7c1eed8
1 parent 2e0681c
commit 7c1eed8
Show file tree

Hide file tree

Showing 9 changed files with 381 additions and 21 deletions.
diff --git a/structlib/cell.py b/structlib/cell.py
@@ -26,7 +26,7 @@
 CoordinateFrame = t.Union[
     t.Literal['cell'], t.Literal['cell_frac'], t.Literal['cell_box'],
     t.Literal['ortho'], t.Literal['ortho_frac'], t.Literal['ortho_box'],
-    t.Literal['local'], t.Literal['global']
+    t.Literal['linear'], t.Literal['local'], t.Literal['global'],
 ]
 """
 A coordinate frame to use.
@@ -36,6 +36,7 @@
  - 'ortho': Angstroms along orthogonal cell
  - 'ortho_frac': Fraction of orthogonal cell
  - 'ortho_box': Fraction of orthogonal box
+ - 'linear': Angstroms in local coordinate system (without affine transformation)
  - 'local': Angstroms in local coordinate system (with affine transformation)
  - 'global': Angstroms in global coordinate system (with all transformations)
 """
@@ -132,6 +133,9 @@ def _get_transform_to_local(self, frame: CoordinateFrame) -> AffineTransform3D:
         if frame == 'local' or frame == 'global':
             return LinearTransform3D()
 
+        if frame == 'linear':
+            return self.affine.to_translation()
+
         if frame.startswith('cell'):
             transform = self.affine @ self.ortho
             cell_size = self.cell_size
@@ -187,6 +191,14 @@ def is_orthogonal_in_local(self, tol: float = 1e-8) -> bool:
             for row in normed
         )
 
+    def _cell_size_in_local(self) -> Vec3:
+        """Calculate cell_size in the local coordinate system. Assumes ``self.is_orthogonal_in_local()``."""
+        return numpy.abs(self.get_transform('local', 'ortho').transform_vec(self.cell_size))
+
+    def _n_cells_in_local(self) -> NDArray[numpy.int_]:
+        """Calculate n_cells after any local rotation. Assumes ``self.is_orthogonal_in_local()``."""
+        return numpy.abs(numpy.round(self.get_transform('local', 'ortho').transform_vec(self.n_cells)).astype(int))
+
     def to_ortho(self) -> AffineTransform3D:
         return self.get_transform('local', 'cell_box')
 

diff --git a/structlib/io/mslice.py b/structlib/io/mslice.py
@@ -10,6 +10,7 @@
 import builtins
 from copy import deepcopy
 from pathlib import Path
+from warnings import warn
 import typing as t
 
 import numpy
@@ -109,6 +110,7 @@ def read_mslice(path: MSliceFile) -> AtomCell:
             'wobble': polars.Float64, 'fracoccupancy': polars.Float64,
         })
         .rename({'atomicnumber': 'elem', 'fracoccupancy': 'frac_occupancy'}) \
+        # 1d sigma -> <u^2>
         .with_columns((3. * polars.col('wobble')**2).alias('wobble'))
     )
     cell = Cell.from_ortho(LinearTransform3D.scale(cell_size), n_cells, [True, True, False])
@@ -119,19 +121,23 @@ def read_mslice(path: MSliceFile) -> AtomCell:
 def write_mslice(cell: HasAtomCell, f: FileOrPath, template: t.Optional[MSliceFile] = None, *,
                  slice_thickness: t.Optional[float] = None,
                  scan_points: t.Optional[ArrayLike] = None,
-                 scan_extent: t.Optional[ArrayLike] = None):
+                 scan_extent: t.Optional[ArrayLike] = None,
+                 n_cells: t.Optional[ArrayLike] = None):
     if not isinstance(cell, HasAtomCell):
         raise TypeError("mslice format requires an AtomCell.")
 
     if not cell.is_orthogonal_in_local():
         raise ValueError("mslice requires an orthogonal AtomCell.")
 
+    if not numpy.all(cell.pbc[:2]):
+        warn("AtomCell may not be periodic", UserWarning, stacklevel=2)
+
+    cell_size = cell._cell_size_in_local()
+
     # get atoms in local frame (which we verified aligns with the cell's axes)
     # then scale into fractional coordinates
-    bbox = cell.bbox()
-    cell_size = bbox.size
-    atoms = cell.get_atoms('local') \
-        .transform(AffineTransform3D.translate(bbox.min).scale(cell_size).inverse()) \
+    atoms = cell.get_atoms('linear') \
+        .transform(AffineTransform3D.scale(1/cell_size)) \
         .with_wobble().with_occupancy()
 
     if template is None:
@@ -165,8 +171,7 @@ def set_attr(struct: et.Element, name: str, type: str, val: str):
         node.text = val
 
     # TODO how to store atoms in unexploded form
-    #(n_a, n_b, n_c) = map(str, atoms.n_cells)
-    (n_a, n_b, n_c) = map(str, (1, 1, 1))
+    (n_a, n_b, n_c) = map(str, (1, 1, 1) if n_cells is None else numpy.asarray(n_cells).astype(int))
     set_attr(struct, 'repeata', 'int16', n_a)
     set_attr(struct, 'repeatb', 'int16', n_b)
     set_attr(struct, 'repeatc', 'int16', n_c)
@@ -196,11 +201,11 @@ def set_attr(struct: et.Element, name: str, type: str, val: str):
     for elem in db.findall("./object[@type='STRUCTUREATOM']"):
         db.remove(elem)
 
-    atoms = atoms.with_wobble((polars.col('wobble') / 3.).sqrt())  # pyMultislicer wants wobble in one dimension
+    # <u^2> -> 1d sigma
+    atoms = atoms.with_wobble((polars.col('wobble') / 3.).sqrt())
     rows = atoms.select(('elem', 'x', 'y', 'z', 'wobble', 'frac_occupancy')).rows()
     for (i, (elem, x, y, z, wobble, frac_occupancy)) in enumerate(rows):
         e = _atom_elem(i, elem, x, y, z, wobble, frac_occupancy)
-        #et.indent(e, space="    ", level=1)
         db.append(e)
 
     et.indent(db, space="    ", level=0)

diff --git a/structlib/io/test_mslice.py b/structlib/io/test_mslice.py
@@ -1,13 +1,36 @@
 
-
 from io import StringIO
 
+import pytest
+import numpy
+
 from .mslice import write_mslice
-from ..testing import check_equals_file
-from ..make import fcc
+from ..testing import check_equals_file, INPUT_PATH, OUTPUT_PATH
+from ..make import fcc, fluorite
+from ..io import read
+from ..atomcell import HasAtomCell
+from ..transform import AffineTransform3D
+
+
+@pytest.fixture(scope='module')
+def ceo2_ortho_cell():
+    return fluorite('CeO2', 5.47, cell='ortho') \
+        .transform(AffineTransform3D.rotate_euler(x=numpy.pi/2.).translate(y=10.))
 
 
 @check_equals_file('Al_from_template.mslice')
 def test_mslice_default_template(buf: StringIO):
     cell = fcc('Al', 4.05, cell='conv').with_wobble(0.030)
     write_mslice(cell, buf, slice_thickness=2.025)
+
+
+@check_equals_file('CeO2_ortho_rotated.mslice')
+def test_mslice_custom_template(buf: StringIO, ceo2_ortho_cell):
+    write_mslice(ceo2_ortho_cell, buf, template=INPUT_PATH / 'bare_template.mslice')
+
+
+@check_equals_file('Al_roundtrip.mslice')
+def test_mslice_roundtrip(buf: StringIO):
+    cell = read(OUTPUT_PATH / 'Al_roundtrip.mslice')
+    assert isinstance(cell, HasAtomCell)
+    write_mslice(cell, buf, slice_thickness=2.025)
diff --git a/structlib/test_cell.py b/structlib/test_cell.py
@@ -30,6 +30,15 @@ def mono_cell() -> Cell:
     return Cell.from_ortho(mono_ortho, n_cells=[2, 3, 5])
 
 
+@pytest.fixture(scope='module')
+def affine_cell() -> Cell:
+    ortho = cell_to_ortho([3., 4., 5.]) \
+        .rotate([0., 0., 1.], numpy.pi/2.) \
+        .translate(-1., -1., -1.)
+
+    return Cell.from_ortho(ortho, n_cells=[2, 3, 5])
+
+
 def test_cell_from_ortho(mono_cell: Cell):
     assert mono_cell.affine.to_linear().is_orthogonal()
     #assert mono_cell.affine.inner == pytest.approx(AffineTransform3D.rotate([0., 0., 1.], numpy.pi/2.).inner)
@@ -62,6 +71,23 @@ def test_mono_cell(mono_cell: Cell, frame_in, frame_out, pts_in, pts_out):
     assert_array_almost_equal(mono_cell.get_transform(frame_out, frame_in).transform(pts_in), pts_out)
 
 
+@pytest.mark.parametrize(('frame_in', 'frame_out', 'pts_in', 'pts_out'), [
+    ('linear', 'local', [[0., 0., 0.]], [[-1., -1., -1.]]),
+    ('local', 'linear', [[0., 0., 0.]], [[1., 1., 1.]]),
+])
+def test_affine_cell(affine_cell: Cell, frame_in, frame_out, pts_in, pts_out):
+    pts_in = numpy.eye(3) if pts_in is None else pts_in
+    assert_array_almost_equal(affine_cell.get_transform(frame_in, frame_out).transform(pts_out), pts_in)
+    assert_array_almost_equal(affine_cell.get_transform(frame_out, frame_in).transform(pts_in), pts_out)
+
+
+def test_cell_in_local(affine_cell: Cell):
+    assert affine_cell.is_orthogonal_in_local()
+
+    assert_array_equal(affine_cell._n_cells_in_local(), [3, 2, 5])
+    assert_array_almost_equal(affine_cell._cell_size_in_local(), [4., 3., 5.])
+
+
 def test_transform_affine_cell():
     cell = Cell.from_unit_cell([2., 3., 5.]) \
         .repeat((9, 6, 3)) \

diff --git a/structlib/transform.py b/structlib/transform.py
@@ -161,6 +161,9 @@ def to_linear(self) -> LinearTransform3D:
         """Return the linear part of an affine transformation."""
         return LinearTransform3D(self.inner[:3, :3])
 
+    def to_translation(self) -> AffineTransform3D:
+        return AffineTransform3D.translate(self.translation())
+
     def det(self) -> float:
         """Return the determinant of an affine transformation."""
         return numpy.linalg.det(self.inner[:3, :3])

diff --git a/tests/baseline_files/Al_roundtrip.mslice b/tests/baseline_files/Al_roundtrip.mslice
@@ -0,0 +1,172 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes" ?>
+<!DOCTYPE database SYSTEM "file:///System/Library/DTDs/CoreData.dtd">
+
+<database>
+    <databaseInfo>
+        <version>134481920</version>
+        <UUID>D1E2667A-3268-45DB-8AA0-DF69A046A0F1</UUID>
+        <nextObjectID>118</nextObjectID>
+        <metadata>
+            <plist version="1.0">
+                <dict>
+                    <key>NSPersistenceFrameworkVersion</key>
+                    <integer>977</integer>
+                    <key>NSStoreModelVersionHashes</key>
+                    <dict>
+                        <key>Aberration</key>
+                        <data>
+                        UXNYvhGfAZGC49DHY75VKihj5GPq2R8M+2L7QhekMmk=
+                        </data>
+                        <key>Detector</key>
+                        <data>
+                        uJuVKALoqyt/oG6mH49teZTkSE+mkbjp/rY1fplbw1I=
+                        </data>
+                        <key>Microscope</key>
+                        <data>
+                        UlR07Rk0b6X1qzhXcmo5lVucE24cBD4ON5bTEMF6nwo=
+                        </data>
+                        <key>SimParameters</key>
+                        <data>
+                        Dxbr+dxiyWuFqo9wbE7miwyXTDQWyxJHhHIaXjfifCY=
+                        </data>
+                        <key>Structure</key>
+                        <data>
+                        e95nHUW6ZbQm5tswhzS+MMc6na7juAQzhA9+z4fqLD0=
+                        </data>
+                        <key>StructureAtom</key>
+                        <data>
+                        boqYojFisRQc7Ui/nA4KYfOt58vOXLXUe8TNOJEQC+k=
+                        </data>
+                        <key>Unit</key>
+                        <data>
+                        91H2vpCfT2XscqlHbquFh62aqBOrVY0ozbfIR0V8fqM=
+                        </data>
+                    </dict>
+                    <key>NSStoreModelVersionHashesDigest</key>
+                    <string>9SQG0FT96IWNxz9CKG5H+mlch0ltlf1ugFqbZ55R8KvYfQFtoAfOXu8ggpr++FRYZdYwhJ2hvbLH6zOKyGB+uw==</string>
+                    <key>NSStoreModelVersionHashesVersion</key>
+                    <integer>3</integer>
+                    <key>NSStoreModelVersionIdentifiers</key>
+                    <array>
+                        <string>0.9</string>
+                    </array>
+                </dict>
+            </plist>
+        </metadata>
+    </databaseInfo>
+    <object type="UNIT" id="z102">
+        <attribute name="value" type="int32">10000</attribute>
+        <attribute name="name" type="string">µm</attribute>
+    </object>
+    <object type="MICROSCOPE" id="z103">
+        <attribute name="kv" type="float">200</attribute>
+        <attribute name="aperture" type="float">19.6</attribute>
+    </object>
+    <object type="ABERRATION" id="z104">
+        <attribute name="name" type="string">Defocus</attribute>
+        <attribute name="n" type="int16">1</attribute>
+        <attribute name="m" type="int16">0</attribute>
+        <attribute name="cnmb" type="float">0</attribute>
+        <attribute name="cnma" type="float">0</attribute>
+        <relationship name="unit" type="0/1" destination="UNIT" idrefs="z106"></relationship>
+    </object>
+    <object type="ABERRATION" id="z105">
+        <attribute name="name" type="string">Two-fold Stig</attribute>
+        <attribute name="n" type="int16">1</attribute>
+        <attribute name="m" type="int16">2</attribute>
+        <attribute name="cnmb" type="float">0</attribute>
+        <attribute name="cnma" type="float">0</attribute>
+        <relationship name="unit" type="0/1" destination="UNIT" idrefs="z106"></relationship>
+    </object>
+    <object type="UNIT" id="z106">
+        <attribute name="value" type="int32">1</attribute>
+        <attribute name="name" type="string">Å</attribute>
+    </object>
+    <object type="SIMPARAMETERS" id="z107">
+        <attribute name="realspace" type="bool">0</attribute>
+        <attribute name="wavesamples" type="int16">1024</attribute>
+        <attribute name="thermalconfigs" type="int16">30</attribute>
+        <attribute name="temperature" type="float">300</attribute>
+        <attribute name="slicethickness" type="float">2.02500000</attribute>
+        <attribute name="pacbed" type="bool">1</attribute>
+        <attribute name="overrideslices" type="bool">1</attribute>
+        <attribute name="numscany" type="int16">64</attribute>
+        <attribute name="numscanx" type="int16">64</attribute>
+        <attribute name="numsaves" type="int16">2</attribute>
+        <attribute name="inty" type="float">0</attribute>
+        <attribute name="intx" type="float">0</attribute>
+        <attribute name="includetds" type="bool">0</attribute>
+        <attribute name="finy" type="float">1</attribute>
+        <attribute name="finx" type="float">1</attribute>
+        <attribute name="cbed" type="bool">1</attribute>
+        <attribute name="saveSlices" type="bool">1</attribute>
+    </object>
+    <object type="UNIT" id="z108">
+        <attribute name="value" type="int32">10000000</attribute>
+        <attribute name="name" type="string">mm</attribute>
+    </object>
+    <object type="ABERRATION" id="z109">
+        <attribute name="name" type="string">Spherical</attribute>
+        <attribute name="n" type="int16">3</attribute>
+        <attribute name="m" type="int16">0</attribute>
+        <attribute name="cnmb" type="float">0</attribute>
+        <attribute name="cnma" type="float">0</attribute>
+        <relationship name="unit" type="0/1" destination="UNIT" idrefs="z102"></relationship>
+    </object>
+    <object type="STRUCTURE" id="z110">
+        <attribute name="tilty" type="float">0</attribute>
+        <attribute name="tiltx" type="float">0</attribute>
+        <attribute name="repeatc" type="int16">1</attribute>
+        <attribute name="repeatb" type="int16">1</attribute>
+        <attribute name="repeata" type="int16">1</attribute>
+        <attribute name="aparam" type="float">4.05000000</attribute>
+        <attribute name="bparam" type="float">4.05000000</attribute>
+        <attribute name="cparam" type="float">4.05000000</attribute>
+        <relationship name="atoms" type="0/1" destination="STRUCTUREATOM"></relationship>
+    </object>
+    <object type="DETECTOR" id="ABF">
+        <attribute name="outerangle" type="float">30</attribute>
+        <attribute name="innerangle" type="float">10</attribute>
+        <attribute name="dpc" type="bool">1</attribute>
+    </object>
+    <object type="DETECTOR" id="HAADF">
+        <attribute name="outerangle" type="float">70</attribute>
+        <attribute name="innerangle" type="float">50</attribute>
+    </object>
+    <object type="DETECTOR" id="LAADF">
+        <attribute name="outerangle" type="float">50</attribute>
+        <attribute name="innerangle" type="float">20</attribute>
+    </object>
+    <object type="STRUCTUREATOM" id="atom0">
+        <attribute name="x" type="float">0.00000000</attribute>
+        <attribute name="y" type="float">0.00000000</attribute>
+        <attribute name="z" type="float">0.00000000</attribute>
+        <attribute name="wobble" type="float">0.1000</attribute>
+        <attribute name="fracoccupancy" type="float">1.0000</attribute>
+        <attribute name="atomicnumber" type="int16">13</attribute>
+    </object>
+    <object type="STRUCTUREATOM" id="atom1">
+        <attribute name="x" type="float">0.00000000</attribute>
+        <attribute name="y" type="float">0.50000000</attribute>
+        <attribute name="z" type="float">0.50000000</attribute>
+        <attribute name="wobble" type="float">0.1000</attribute>
+        <attribute name="fracoccupancy" type="float">1.0000</attribute>
+        <attribute name="atomicnumber" type="int16">13</attribute>
+    </object>
+    <object type="STRUCTUREATOM" id="atom2">
+        <attribute name="x" type="float">0.50000000</attribute>
+        <attribute name="y" type="float">0.00000000</attribute>
+        <attribute name="z" type="float">0.50000000</attribute>
+        <attribute name="wobble" type="float">0.1000</attribute>
+        <attribute name="fracoccupancy" type="float">1.0000</attribute>
+        <attribute name="atomicnumber" type="int16">13</attribute>
+    </object>
+    <object type="STRUCTUREATOM" id="atom3">
+        <attribute name="x" type="float">0.50000000</attribute>
+        <attribute name="y" type="float">0.50000000</attribute>
+        <attribute name="z" type="float">0.00000000</attribute>
+        <attribute name="wobble" type="float">0.1000</attribute>
+        <attribute name="fracoccupancy" type="float">1.0000</attribute>
+        <attribute name="atomicnumber" type="int16">13</attribute>
+    </object>
+</database>