Skip to content

atomlib.io

FileType module-attribute 

FileType: TypeAlias = Literal[
    "cif", "xyz", "xsf", "cfg", "lmp", "mslice", "qe"
]

ReadFunc module-attribute 

ReadFunc = Callable[[FileOrPath], HasAtoms]

WriteFunc module-attribute 

WriteFunc = Callable[[HasAtoms, FileOrPath], None]

CIF dataclass

Source code in atomlib/io/cif.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
@dataclass
class CIF:
    data_blocks: t.Tuple[CIFDataBlock, ...]

    def __post_init__(self):
        # ensure that all data_blocks after the first have a name
        for data_block in self.data_blocks[1:]:
            if data_block.name is None:
                data_block.name = ""

    @staticmethod
    def from_file(file: FileOrPath) -> CIF:
        return CIF(tuple(CIFDataBlock.from_file(file)))

    def __len__(self) -> int:
        return self.data_blocks.__len__()

    def get_block(self, block: t.Union[int, str]) -> CIFDataBlock:
        try:
            if isinstance(block, int):
                return self.data_blocks[block]
            return next(b for b in self.data_blocks if b.name == block)
        except (IndexError, StopIteration):
            raise ValueError(f"Couldn't find block '{block}' in CIF file. File has {len(self)} blocks.")

    def write(self, file: FileOrPath):
        with open_file(file, 'w') as f:
            print("# generated by atomlib", file=f, end=None)
            for data_block in self.data_blocks:
                print(file=f)
                data_block._write(f)

data_blocks instance-attribute 

data_blocks: Tuple[CIFDataBlock, ...]

from_file staticmethod 

from_file(file: FileOrPath) -> CIF
Source code in atomlib/io/cif.py 
46
47
48
@staticmethod
def from_file(file: FileOrPath) -> CIF:
    return CIF(tuple(CIFDataBlock.from_file(file)))

get_block 

get_block(block: Union[int, str]) -> CIFDataBlock
Source code in atomlib/io/cif.py 
53
54
55
56
57
58
59
def get_block(self, block: t.Union[int, str]) -> CIFDataBlock:
    try:
        if isinstance(block, int):
            return self.data_blocks[block]
        return next(b for b in self.data_blocks if b.name == block)
    except (IndexError, StopIteration):
        raise ValueError(f"Couldn't find block '{block}' in CIF file. File has {len(self)} blocks.")

write 

write(file: FileOrPath)
Source code in atomlib/io/cif.py 
61
62
63
64
65
66
def write(self, file: FileOrPath):
    with open_file(file, 'w') as f:
        print("# generated by atomlib", file=f, end=None)
        for data_block in self.data_blocks:
            print(file=f)
            data_block._write(f)

XYZ dataclass

Source code in atomlib/io/xyz.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
@dataclass
class XYZ:
    atoms: polars.DataFrame
    comment: t.Optional[str] = None
    params: t.Dict[str, str] = field(default_factory=dict)

    @staticmethod
    def from_atoms(atoms: HasAtoms) -> XYZ:
        params = {}
        if isinstance(atoms, HasAtomCell):
            coords = atoms.get_cell().to_ortho().to_linear().inner.ravel()
            lattice_str = " ".join((f"{c:.8f}" for c in coords))
            params['Lattice'] = lattice_str

            pbc_str = " ".join(str(int(v)) for v in atoms.get_cell().pbc)
            params['pbc'] = pbc_str

        return XYZ(
            atoms.get_atoms('local')._get_frame(),
            params=params
        )

    @staticmethod
    def from_file(file: FileOrPath) -> XYZ:
        logging.info(f"Loading XYZ {file.name if hasattr(file, 'name') else file!r}...")  # type: ignore

        with open_file(file, 'r') as f:
            try:
                # TODO be more gracious about whitespace here
                length = int(f.readline())
            except ValueError:
                raise ValueError("Error parsing XYZ file: Invalid length") from None
            except IOError as e:
                raise IOError(f"Error parsing XYZ file: {e}") from None

            comment = f.readline().rstrip('\n')
            # TODO handle if there's not a gap here

            try:
                params = ExtXYZParser(comment).parse()
            except ValueError:
                params = None

            schema = _get_columns_from_params(params)

            df = parse_whitespace_separated(f, schema, start_line=1)

            # map atomic numbers -> symbols (on columns which are Int8)
            df = df.with_columns(
                get_sym(df.select(polars.col('symbol').cast(polars.Int8, strict=False)).to_series())
                    .fill_null(df['symbol']).alias('symbol')
            )
            # ensure all symbols are recognizable (this will raise ValueError if not)
            get_elem(df['symbol'])

            if length < len(df):
                warnings.warn(f"Warning: truncating structure of length {len(df)} "
                            f"to match declared length of {length}")
                df = df[:length]
            elif length > len(df):
                warnings.warn(f"Warning: structure length {len(df)} doesn't match "
                            f"declared length {length}.\nData could be corrupted.")

            try:
                params = ExtXYZParser(comment).parse()
                return XYZ(df, comment, params)
            except ValueError:
                pass

            return XYZ(df, comment)

    def write(self, file: FileOrPath, fmt: XYZFormat = 'exyz'):
        with open_file(file, 'w', newline='\r\n') as f:

            f.write(f"{len(self.atoms)}\n")
            if len(self.params) > 0 and fmt == 'exyz':
                f.write(" ".join(_param_strings(self.params)))
            else:
                f.write(self.comment or "")
            f.write("\n")

            # not my best work
            col_space = (3, 12, 12, 12)
            f.writelines(
                "".join(
                    f"{val:< {space}.8f}" if isinstance(val, float) else f"{val:<{space}}" for (val, space) in zip(_flatten(row), col_space)
                ).strip() + '\n' for row in self.atoms.select(('symbol', 'coords')).rows()
            )

    def cell_matrix(self) -> t.Optional[NDArray[numpy.float64]]:
        if (s := self.params.get('Lattice')) is None:
            return None

        try:
            items = list(map(float, s.split()))
            if not len(items) == 9:
                raise ValueError("Invalid length")
            return numpy.array(items, dtype=numpy.float64).reshape((3, 3)).T
        except ValueError:
            warnings.warn(f"Warning: Invalid format for key 'Lattice=\"{s}\"'")
        return None

    def pbc(self) -> t.Optional[NDArray[numpy.bool_]]:
        if (s := self.params.get('pbc')) is None:
            return None

        val_map = {'0': False, 'f': False, '1': True, 't': True}
        try:
            items = [val_map[v.lower()] for v in s.split()]
            if not len(items) == 3:
                raise ValueError("Invalid length")
            return numpy.array(items, dtype=numpy.bool_)
        except ValueError:
            warnings.warn(f"Warning: Invalid format for key 'pbc=\"{s}\"'")
        return None

atoms instance-attribute 

atoms: DataFrame

comment class-attribute instance-attribute 

comment: Optional[str] = None

params class-attribute instance-attribute 

params: Dict[str, str] = field(default_factory=dict)

from_atoms staticmethod 

from_atoms(atoms: HasAtoms) -> XYZ
Source code in atomlib/io/xyz.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@staticmethod
def from_atoms(atoms: HasAtoms) -> XYZ:
    params = {}
    if isinstance(atoms, HasAtomCell):
        coords = atoms.get_cell().to_ortho().to_linear().inner.ravel()
        lattice_str = " ".join((f"{c:.8f}" for c in coords))
        params['Lattice'] = lattice_str

        pbc_str = " ".join(str(int(v)) for v in atoms.get_cell().pbc)
        params['pbc'] = pbc_str

    return XYZ(
        atoms.get_atoms('local')._get_frame(),
        params=params
    )

from_file staticmethod 

from_file(file: FileOrPath) -> XYZ
Source code in atomlib/io/xyz.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
@staticmethod
def from_file(file: FileOrPath) -> XYZ:
    logging.info(f"Loading XYZ {file.name if hasattr(file, 'name') else file!r}...")  # type: ignore

    with open_file(file, 'r') as f:
        try:
            # TODO be more gracious about whitespace here
            length = int(f.readline())
        except ValueError:
            raise ValueError("Error parsing XYZ file: Invalid length") from None
        except IOError as e:
            raise IOError(f"Error parsing XYZ file: {e}") from None

        comment = f.readline().rstrip('\n')
        # TODO handle if there's not a gap here

        try:
            params = ExtXYZParser(comment).parse()
        except ValueError:
            params = None

        schema = _get_columns_from_params(params)

        df = parse_whitespace_separated(f, schema, start_line=1)

        # map atomic numbers -> symbols (on columns which are Int8)
        df = df.with_columns(
            get_sym(df.select(polars.col('symbol').cast(polars.Int8, strict=False)).to_series())
                .fill_null(df['symbol']).alias('symbol')
        )
        # ensure all symbols are recognizable (this will raise ValueError if not)
        get_elem(df['symbol'])

        if length < len(df):
            warnings.warn(f"Warning: truncating structure of length {len(df)} "
                        f"to match declared length of {length}")
            df = df[:length]
        elif length > len(df):
            warnings.warn(f"Warning: structure length {len(df)} doesn't match "
                        f"declared length {length}.\nData could be corrupted.")

        try:
            params = ExtXYZParser(comment).parse()
            return XYZ(df, comment, params)
        except ValueError:
            pass

        return XYZ(df, comment)

write 

write(file: FileOrPath, fmt: XYZFormat = 'exyz')
Source code in atomlib/io/xyz.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def write(self, file: FileOrPath, fmt: XYZFormat = 'exyz'):
    with open_file(file, 'w', newline='\r\n') as f:

        f.write(f"{len(self.atoms)}\n")
        if len(self.params) > 0 and fmt == 'exyz':
            f.write(" ".join(_param_strings(self.params)))
        else:
            f.write(self.comment or "")
        f.write("\n")

        # not my best work
        col_space = (3, 12, 12, 12)
        f.writelines(
            "".join(
                f"{val:< {space}.8f}" if isinstance(val, float) else f"{val:<{space}}" for (val, space) in zip(_flatten(row), col_space)
            ).strip() + '\n' for row in self.atoms.select(('symbol', 'coords')).rows()
        )

cell_matrix 

cell_matrix() -> Optional[NDArray[float64]]
Source code in atomlib/io/xyz.py 
152
153
154
155
156
157
158
159
160
161
162
163
def cell_matrix(self) -> t.Optional[NDArray[numpy.float64]]:
    if (s := self.params.get('Lattice')) is None:
        return None

    try:
        items = list(map(float, s.split()))
        if not len(items) == 9:
            raise ValueError("Invalid length")
        return numpy.array(items, dtype=numpy.float64).reshape((3, 3)).T
    except ValueError:
        warnings.warn(f"Warning: Invalid format for key 'Lattice=\"{s}\"'")
    return None

pbc 

pbc() -> Optional[NDArray[bool_]]
Source code in atomlib/io/xyz.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
def pbc(self) -> t.Optional[NDArray[numpy.bool_]]:
    if (s := self.params.get('pbc')) is None:
        return None

    val_map = {'0': False, 'f': False, '1': True, 't': True}
    try:
        items = [val_map[v.lower()] for v in s.split()]
        if not len(items) == 3:
            raise ValueError("Invalid length")
        return numpy.array(items, dtype=numpy.bool_)
    except ValueError:
        warnings.warn(f"Warning: Invalid format for key 'pbc=\"{s}\"'")
    return None

XSF dataclass

Source code in atomlib/io/xsf.py 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
@dataclass
class XSF:
    periodicity: Periodicity = 'crystal'
    primitive_cell: t.Optional[LinearTransform3D] = None
    conventional_cell: t.Optional[LinearTransform3D] = None

    prim_coords: t.Optional[polars.DataFrame] = None
    conv_coords: t.Optional[polars.DataFrame] = None
    atoms: t.Optional[polars.DataFrame] = None

    def get_atoms(self) -> polars.DataFrame:
        if self.prim_coords is not None:
            return self.prim_coords
        if self.atoms is not None:
            return self.atoms
        if self.conv_coords is not None:
            raise NotImplementedError()  # TODO untransform conv_coords by conventional_cell?
        raise ValueError("No coordinates specified in XSF file.")

    def get_pbc(self) -> NDArray[numpy.bool_]:
        return _periodicity_to_pbc(self.periodicity)

    @staticmethod
    def from_cell(cell: HasAtomCell) -> XSF:
        ortho = cell.get_transform('local', 'cell_box').to_linear()
        return XSF(
            primitive_cell=ortho,
            conventional_cell=ortho,
            prim_coords=cell.get_atoms('linear').inner,
            periodicity=_pbc_to_periodicity(cell.pbc)
        )

    @staticmethod
    def from_atoms(atoms: HasAtoms) -> XSF:
        return XSF(
            periodicity='molecule',
            atoms=atoms.get_atoms('local').inner
        )

    @staticmethod
    def from_file(file: FileOrPath) -> XSF:
        logging.info(f"Loading XSF {file.name if hasattr(file, 'name') else file!r}...")  # type: ignore
        with open_file(file) as f:
            return XSFParser(f).parse()

    def __post_init__(self):
        if self.prim_coords is None and self.conv_coords is None and self.atoms is None:
            raise ValueError("Error: No coordinates are specified (atoms, primitive, or conventional).")

        if self.prim_coords is not None and self.conv_coords is not None:
            logging.warning("Warning: Both 'primcoord' and 'convcoord' are specified. 'convcoord' will be ignored.")
        elif self.conv_coords is not None and self.conventional_cell is None:
            raise ValueError("If 'convcoord' is specified, 'convvec' must be specified as well.")

        if self.periodicity == 'molecule':
            if self.atoms is None:
                raise ValueError("'atoms' must be specified for molecules.")

    def write(self, path: FileOrPath):
        with open_file(path, 'w') as f:
            print(self.periodicity.upper(), file=f)
            if self.primitive_cell is not None:
                print('PRIMVEC', file=f)
                self._write_cell(f, self.primitive_cell)
            if self.conventional_cell is not None:
                print('CONVVEC', file=f)
                self._write_cell(f, self.conventional_cell)
            print(file=f)

            if self.prim_coords is not None:
                print("PRIMCOORD", file=f)
                print(f"{len(self.prim_coords)} 1", file=f)
                self._write_coords(f, self.prim_coords)
            if self.conv_coords is not None:
                print("CONVCOORD", file=f)
                print(f"{len(self.conv_coords)} 1", file=f)
                self._write_coords(f, self.conv_coords)
            if self.atoms is not None:
                print("ATOMS", file=f)
                self._write_coords(f, self.atoms)

    def _write_cell(self, f: TextIOBase, cell: LinearTransform3D):
        for row in cell.inner.T:
            for val in row:
                f.write(f"{val:12.7f}")
            f.write('\n')

    def _write_coords(self, f: TextIOBase, coords: polars.DataFrame):
        for (elem, [x, y, z]) in coords.select(['elem', 'coords']).rows():
            print(f"{elem:2d} {x:11.6f} {y:11.6f} {z:11.6f}", file=f)
        print(file=f)

periodicity class-attribute instance-attribute 

periodicity: Periodicity = 'crystal'

primitive_cell class-attribute instance-attribute 

primitive_cell: Optional[LinearTransform3D] = None

conventional_cell class-attribute instance-attribute 

conventional_cell: Optional[LinearTransform3D] = None

prim_coords class-attribute instance-attribute 

prim_coords: Optional[DataFrame] = None

conv_coords class-attribute instance-attribute 

conv_coords: Optional[DataFrame] = None

atoms class-attribute instance-attribute 

atoms: Optional[DataFrame] = None

get_atoms 

get_atoms() -> DataFrame
Source code in atomlib/io/xsf.py 
67
68
69
70
71
72
73
74
def get_atoms(self) -> polars.DataFrame:
    if self.prim_coords is not None:
        return self.prim_coords
    if self.atoms is not None:
        return self.atoms
    if self.conv_coords is not None:
        raise NotImplementedError()  # TODO untransform conv_coords by conventional_cell?
    raise ValueError("No coordinates specified in XSF file.")

get_pbc 

get_pbc() -> NDArray[bool_]
Source code in atomlib/io/xsf.py 
76
77
def get_pbc(self) -> NDArray[numpy.bool_]:
    return _periodicity_to_pbc(self.periodicity)

from_cell staticmethod 

from_cell(cell: HasAtomCell) -> XSF
Source code in atomlib/io/xsf.py 
79
80
81
82
83
84
85
86
87
@staticmethod
def from_cell(cell: HasAtomCell) -> XSF:
    ortho = cell.get_transform('local', 'cell_box').to_linear()
    return XSF(
        primitive_cell=ortho,
        conventional_cell=ortho,
        prim_coords=cell.get_atoms('linear').inner,
        periodicity=_pbc_to_periodicity(cell.pbc)
    )

from_atoms staticmethod 

from_atoms(atoms: HasAtoms) -> XSF
Source code in atomlib/io/xsf.py 
89
90
91
92
93
94
@staticmethod
def from_atoms(atoms: HasAtoms) -> XSF:
    return XSF(
        periodicity='molecule',
        atoms=atoms.get_atoms('local').inner
    )

from_file staticmethod 

from_file(file: FileOrPath) -> XSF
Source code in atomlib/io/xsf.py 
 96
 97
 98
 99
100
@staticmethod
def from_file(file: FileOrPath) -> XSF:
    logging.info(f"Loading XSF {file.name if hasattr(file, 'name') else file!r}...")  # type: ignore
    with open_file(file) as f:
        return XSFParser(f).parse()

write 

write(path: FileOrPath)
Source code in atomlib/io/xsf.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def write(self, path: FileOrPath):
    with open_file(path, 'w') as f:
        print(self.periodicity.upper(), file=f)
        if self.primitive_cell is not None:
            print('PRIMVEC', file=f)
            self._write_cell(f, self.primitive_cell)
        if self.conventional_cell is not None:
            print('CONVVEC', file=f)
            self._write_cell(f, self.conventional_cell)
        print(file=f)

        if self.prim_coords is not None:
            print("PRIMCOORD", file=f)
            print(f"{len(self.prim_coords)} 1", file=f)
            self._write_coords(f, self.prim_coords)
        if self.conv_coords is not None:
            print("CONVCOORD", file=f)
            print(f"{len(self.conv_coords)} 1", file=f)
            self._write_coords(f, self.conv_coords)
        if self.atoms is not None:
            print("ATOMS", file=f)
            self._write_coords(f, self.atoms)

CFG dataclass

Source code in atomlib/io/cfg.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@dataclass
class CFG:
    atoms: polars.DataFrame

    cell: LinearTransform3D
    transform: t.Optional[LinearTransform3D] = None
    eta: t.Optional[LinearTransform3D] = None

    length_scale: t.Optional[float] = None
    length_unit: t.Optional[str] = None
    rate_scale: t.Optional[float] = None
    rate_unit: t.Optional[str] = None

    @staticmethod
    def from_file(file: FileOrPath) -> CFG:
        with open_file(file, 'r') as f:
            return CFGParser(f).parse()

    @staticmethod
    def from_atoms(atoms: HasAtoms) -> CFG:
        if isinstance(atoms, HasAtomCell):
            cell = atoms.get_transform('cell_box').inverse().to_linear()
            atoms = atoms.get_atoms('cell_box')
        else:
            cell = LinearTransform3D.identity()

        # ensure we have masses and velocities
        atoms = atoms.with_mass().with_velocity()
        return CFG(atoms._get_frame(), cell, length_scale=1.0, length_unit="Angstrom")

    def write(self, file: FileOrPath):
        with open_file(file, 'w', newline='\r\n') as f:
            f.write(f"Number of particles = {len(self.atoms)}\n")

            if self.length_scale is not None:
                unit = f" [{self.length_unit}]" if self.length_unit is not None else ""
                f.write(f"\nA = {self.length_scale:.8}{unit}\n\n")

            cell = self.cell.inner
            for (i, j) in product(range(3), repeat=2):
                f.write(f"H0({i+1},{j+1}) = {cell[j,i]:.8} A\n")

            if self.transform is not None:
                f.write("\n")
                transform = self.transform.inner
                for (i, j) in product(range(3), repeat=2):
                    f.write(f"Transform({i+1},{j+1}) = {transform[j,i]:.8}\n")

            if self.eta is not None:
                f.write("\n")
                eta = self.eta.inner
                for i in range(3):
                    for j in range(i, 3):
                        f.write(f"eta({i+1},{j+1}) = {eta[j,i]:.8}\n")

            if self.rate_scale is not None:
                unit = f" [{self.rate_unit}]" if self.rate_unit is not None else ""
                f.write(f"\nR = {self.rate_scale:.8}{unit}\n")

            f.write("\n")
            for row in self.atoms.select(('mass', 'symbol', 'coords', 'velocity')).rows():
                (mass, sym, (x, y, z), (v_x, v_y, v_z)) = row
                f.write(f"{mass:.4f} {sym:>2} {x:.8} {y:.8} {z:.8} {v_x:.8} {v_y:.8} {v_z:.8}\n")

atoms instance-attribute 

atoms: DataFrame

cell instance-attribute 

cell: LinearTransform3D

transform class-attribute instance-attribute 

transform: Optional[LinearTransform3D] = None

eta class-attribute instance-attribute 

eta: Optional[LinearTransform3D] = None

length_scale class-attribute instance-attribute 

length_scale: Optional[float] = None

length_unit class-attribute instance-attribute 

length_unit: Optional[str] = None

rate_scale class-attribute instance-attribute 

rate_scale: Optional[float] = None

rate_unit class-attribute instance-attribute 

rate_unit: Optional[str] = None

from_file staticmethod 

from_file(file: FileOrPath) -> CFG
Source code in atomlib/io/cfg.py 
35
36
37
38
@staticmethod
def from_file(file: FileOrPath) -> CFG:
    with open_file(file, 'r') as f:
        return CFGParser(f).parse()

from_atoms staticmethod 

from_atoms(atoms: HasAtoms) -> CFG
Source code in atomlib/io/cfg.py 
40
41
42
43
44
45
46
47
48
49
50
@staticmethod
def from_atoms(atoms: HasAtoms) -> CFG:
    if isinstance(atoms, HasAtomCell):
        cell = atoms.get_transform('cell_box').inverse().to_linear()
        atoms = atoms.get_atoms('cell_box')
    else:
        cell = LinearTransform3D.identity()

    # ensure we have masses and velocities
    atoms = atoms.with_mass().with_velocity()
    return CFG(atoms._get_frame(), cell, length_scale=1.0, length_unit="Angstrom")

write 

write(file: FileOrPath)
Source code in atomlib/io/cfg.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def write(self, file: FileOrPath):
    with open_file(file, 'w', newline='\r\n') as f:
        f.write(f"Number of particles = {len(self.atoms)}\n")

        if self.length_scale is not None:
            unit = f" [{self.length_unit}]" if self.length_unit is not None else ""
            f.write(f"\nA = {self.length_scale:.8}{unit}\n\n")

        cell = self.cell.inner
        for (i, j) in product(range(3), repeat=2):
            f.write(f"H0({i+1},{j+1}) = {cell[j,i]:.8} A\n")

        if self.transform is not None:
            f.write("\n")
            transform = self.transform.inner
            for (i, j) in product(range(3), repeat=2):
                f.write(f"Transform({i+1},{j+1}) = {transform[j,i]:.8}\n")

        if self.eta is not None:
            f.write("\n")
            eta = self.eta.inner
            for i in range(3):
                for j in range(i, 3):
                    f.write(f"eta({i+1},{j+1}) = {eta[j,i]:.8}\n")

        if self.rate_scale is not None:
            unit = f" [{self.rate_unit}]" if self.rate_unit is not None else ""
            f.write(f"\nR = {self.rate_scale:.8}{unit}\n")

        f.write("\n")
        for row in self.atoms.select(('mass', 'symbol', 'coords', 'velocity')).rows():
            (mass, sym, (x, y, z), (v_x, v_y, v_z)) = row
            f.write(f"{mass:.4f} {sym:>2} {x:.8} {y:.8} {z:.8} {v_x:.8} {v_y:.8} {v_z:.8}\n")

LMP dataclass

Source code in atomlib/io/lmp.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
@dataclass
class LMP:
    comment: t.Optional[str]
    headers: t.Dict[str, t.Any]
    sections: t.Tuple[LMPSection, ...]

    def get_cell(self) -> Cell:
        dims = numpy.array([
            self.headers.get(f"{c}lo {c}hi", (-0.5, 0.5))
            for c in "xyz"
        ])
        origin = dims[:, 0]
        tilts = self.headers.get("xy xz yz", (0., 0., 0.))

        ortho = numpy.diag(dims[:, 1] - dims[:, 0])
        (ortho[0, 1], ortho[0, 2], ortho[1, 2]) = tilts

        return Cell.from_ortho(LinearTransform3D(ortho).translate(origin))

    def get_atoms(self, type_map: t.Optional[t.Dict[int, t.Union[str, int]]] = None) -> AtomCell:
        if type_map is not None:
            try:
                type_map_df = polars.DataFrame({
                    'type': polars.Series(type_map.keys(), dtype=polars.Int32),
                    'elem': polars.Series(list(map(get_elem, type_map.values())), dtype=polars.UInt8),
                    'symbol': polars.Series([get_sym(v) if isinstance(v, int) else v for v in type_map.values()], dtype=polars.Utf8),
                })
            except ValueError as e:
                raise ValueError("Invalid type map") from e
        else:
            type_map_df = None

        cell = self.get_cell()

        def _apply_type_labels(df: polars.DataFrame, section_name: str, labels: t.Optional[polars.DataFrame] = None) -> polars.DataFrame:
            if labels is not None:
                #df = df.with_columns(polars.col('type').replace(d, default=polars.col('type').cast(polars.Int32, strict=False), return_dtype=polars.Int32))
                df = df.with_columns(polars.col('type').replace_strict(labels['symbol'], labels['type'], default=polars.col('type').cast(polars.Int32, strict=False), return_dtype=polars.Int32))
                if df['type'].is_null().any():
                    raise ValueError(f"While parsing section {section_name}: Unknown atom label or invalid atom type")
            try:
                return df.with_columns(polars.col('type').cast(polars.Int32))
            except polars.ComputeError:
                raise ValueError(f"While parsing section {section_name}: Invalid atom type(s)")

        atoms: t.Optional[polars.DataFrame] = None
        labels: t.Optional[polars.DataFrame] = None
        masses: t.Optional[polars.DataFrame] = None
        velocities = None

        for section in self.sections:
            start_line = section.start_line + 1

            if section.name == 'Atoms':
                if section.style not in (None, 'atomic'):
                    # TODO support other styles
                    raise ValueError(f"Only 'atomic' atom_style is supported, instead got '{section.style}'")

                atoms = parse_whitespace_separated(section.body, {
                    'i': polars.Int64, 'type': polars.Utf8,
                    'coords': polars.Array(polars.Float64, 3),
                }, start_line=start_line)
                atoms = _apply_type_labels(atoms, 'Atoms', labels)
            elif section.name == 'Atom Type Labels':
                labels = parse_whitespace_separated(section.body, {'type': polars.Int32, 'symbol': polars.Utf8}, start_line=start_line)
            elif section.name == 'Masses':
                masses = parse_whitespace_separated(section.body, {'type': polars.Utf8, 'mass': polars.Float64}, start_line=start_line)
                masses = _apply_type_labels(masses, 'Masses', labels)
            elif section.name == 'Velocities':
                velocities = parse_whitespace_separated(section.body, {
                    'i': polars.Int64, 'velocity': polars.Array(polars.Float64, 3),
                }, start_line=start_line)

        # now all 'type's should be in Int32

        if atoms is None:
            if self.headers['atoms'] > 0:
                raise ValueError("Missing required section 'Atoms'")
            return AtomCell(Atoms.empty(), cell=cell, frame='local')

        # next we need to assign element symbols
        # first, if type_map is specified, use that:
        #if type_map_elem is not None and type_map_sym is not None:
        if type_map_df is not None:
            try:
                atoms = checked_left_join(atoms, type_map_df, on='type')
            except CheckedJoinError as e:
                raise ValueError(f"Missing type_map specification for atom type(s): {', '.join(map(repr, e.missing_keys))}")
        elif labels is not None:
            try:
                labels = labels.with_columns(get_elem(labels['symbol']))
            except ValueError as e:
                raise ValueError("Failed to auto-detect elements from type labels. Please pass 'type_map' explicitly") from e
            try:
                atoms = checked_left_join(atoms, labels, 'type')
            except CheckedJoinError as e:
                raise ValueError(f"Missing labels for atom type(s): {', '.join(map(repr, e.missing_keys))}")
        # otherwise we have no way
        else:
            raise ValueError("Failed to auto-detect elements from type labels. Please pass 'type_map' explicitly")

        if velocities is not None:
            # join velocities
            try:
                # TODO use join_asof here?
                atoms = checked_left_join(atoms, velocities, 'i')
            except CheckedJoinError as e:
                raise ValueError(f"Missing velocities for {len(e.missing_keys)}/{len(atoms)} atoms")

        if masses is not None:
            # join masses
            try:
                atoms = checked_left_join(atoms, masses, 'type')
            except CheckedJoinError as e:
                raise ValueError(f"Missing masses for atom type(s): {', '.join(map(repr, e.missing_keys))}")

        return AtomCell(atoms, cell=cell, frame='local')

    @staticmethod
    def from_atoms(atoms: HasAtoms) -> LMP:
        if isinstance(atoms, HasAtomCell):
            # we're basically converting everything to the ortho frame, but including the affine shift

            # transform affine shift into ortho frame
            origin = atoms.get_transform('ortho', 'local').to_linear().round_near_zero() \
                .transform(atoms.get_cell().affine.translation())

            # get the orthogonalization transform only, without affine
            ortho = atoms.get_transform('ortho', 'cell_box').to_linear().round_near_zero().inner

            # get atoms in ortho frame, and then add the affine shift
            frame = atoms.get_atoms('ortho').transform_atoms(AffineTransform3D.translate(origin)) \
                .round_near_zero().with_type()
        else:
            bbox = atoms.bbox_atoms()
            ortho = numpy.diag(bbox.size)
            origin = bbox.min

            frame = atoms.get_atoms('local').with_type()

        types = frame.unique(subset='type')
        types = types.with_mass().sort('type')

        now = localtime()
        comment = f"# Generated by atomlib on {now.isoformat(' ', 'seconds')}"

        headers = {}
        sections = []

        headers['atoms'] = len(frame)
        headers['atom types'] = len(types)

        for (s, low, diff) in zip(('x', 'y', 'z'), origin, ortho.diagonal()):
            headers[f"{s}lo {s}hi"] = (low, low + diff)

        headers['xy xz yz'] = (ortho[0, 1], ortho[0, 2], ortho[1, 2])

        body = [
            f" {ty:8} {sym:>4}\n"
            for (ty, sym) in types.select('type', 'symbol').rows()
        ]
        sections.append(LMPSection("Atom Type Labels", tuple(body)))

        if 'mass' in types:
            body = [
                f" {ty:8} {mass:14.7f}  # {sym}\n"
                for (ty, sym, mass) in types.select(('type', 'symbol', 'mass')).rows()
            ]
            sections.append(LMPSection("Masses", tuple(body)))

        body = [
            f" {i+1:8} {ty:4} {x:14.7f} {y:14.7f} {z:14.7f}\n"
            for (i, (ty, (x, y, z))) in enumerate(frame.select(('type', 'coords')).rows())
        ]
        sections.append(LMPSection("Atoms", tuple(body), 'atomic'))

        if (velocities := frame.velocities()) is not None:
            body = [
                f" {i+1:8} {v_x:14.7f} {v_y:14.7f} {v_z:14.7f}\n"
                for (i, (v_x, v_y, v_z)) in enumerate(velocities)
            ]
            sections.append(LMPSection("Velocities", tuple(body)))

        return LMP(comment, headers, tuple(sections))

    @staticmethod
    def from_file(file: FileOrPath) -> LMP:
        with open_file(file, 'r') as f:
            return LMPReader(f).parse()

    def write(self, file: FileOrPath):
        with open_file(file, 'w') as f:
            print((self.comment or "") + '\n', file=f)

            # print headers
            for (name, val) in self.headers.items():
                val = _HEADER_FMT.get(name, lambda s: f"{s:8}")(val)
                print(f" {val} {name}", file=f)

            # print sections
            for section in self.sections:
                line = section.name
                if section.style is not None:
                    line += f'  # {section.style}'
                print(f"\n{line}\n", file=f)

                f.writelines(section.body)

comment instance-attribute 

comment: Optional[str]

headers instance-attribute 

headers: Dict[str, Any]

sections instance-attribute 

sections: Tuple[LMPSection, ...]

get_cell 

get_cell() -> Cell
Source code in atomlib/io/lmp.py 
28
29
30
31
32
33
34
35
36
37
38
39
def get_cell(self) -> Cell:
    dims = numpy.array([
        self.headers.get(f"{c}lo {c}hi", (-0.5, 0.5))
        for c in "xyz"
    ])
    origin = dims[:, 0]
    tilts = self.headers.get("xy xz yz", (0., 0., 0.))

    ortho = numpy.diag(dims[:, 1] - dims[:, 0])
    (ortho[0, 1], ortho[0, 2], ortho[1, 2]) = tilts

    return Cell.from_ortho(LinearTransform3D(ortho).translate(origin))

get_atoms 

get_atoms(
    type_map: Optional[Dict[int, Union[str, int]]] = None
) -> AtomCell
Source code in atomlib/io/lmp.py 
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def get_atoms(self, type_map: t.Optional[t.Dict[int, t.Union[str, int]]] = None) -> AtomCell:
    if type_map is not None:
        try:
            type_map_df = polars.DataFrame({
                'type': polars.Series(type_map.keys(), dtype=polars.Int32),
                'elem': polars.Series(list(map(get_elem, type_map.values())), dtype=polars.UInt8),
                'symbol': polars.Series([get_sym(v) if isinstance(v, int) else v for v in type_map.values()], dtype=polars.Utf8),
            })
        except ValueError as e:
            raise ValueError("Invalid type map") from e
    else:
        type_map_df = None

    cell = self.get_cell()

    def _apply_type_labels(df: polars.DataFrame, section_name: str, labels: t.Optional[polars.DataFrame] = None) -> polars.DataFrame:
        if labels is not None:
            #df = df.with_columns(polars.col('type').replace(d, default=polars.col('type').cast(polars.Int32, strict=False), return_dtype=polars.Int32))
            df = df.with_columns(polars.col('type').replace_strict(labels['symbol'], labels['type'], default=polars.col('type').cast(polars.Int32, strict=False), return_dtype=polars.Int32))
            if df['type'].is_null().any():
                raise ValueError(f"While parsing section {section_name}: Unknown atom label or invalid atom type")
        try:
            return df.with_columns(polars.col('type').cast(polars.Int32))
        except polars.ComputeError:
            raise ValueError(f"While parsing section {section_name}: Invalid atom type(s)")

    atoms: t.Optional[polars.DataFrame] = None
    labels: t.Optional[polars.DataFrame] = None
    masses: t.Optional[polars.DataFrame] = None
    velocities = None

    for section in self.sections:
        start_line = section.start_line + 1

        if section.name == 'Atoms':
            if section.style not in (None, 'atomic'):
                # TODO support other styles
                raise ValueError(f"Only 'atomic' atom_style is supported, instead got '{section.style}'")

            atoms = parse_whitespace_separated(section.body, {
                'i': polars.Int64, 'type': polars.Utf8,
                'coords': polars.Array(polars.Float64, 3),
            }, start_line=start_line)
            atoms = _apply_type_labels(atoms, 'Atoms', labels)
        elif section.name == 'Atom Type Labels':
            labels = parse_whitespace_separated(section.body, {'type': polars.Int32, 'symbol': polars.Utf8}, start_line=start_line)
        elif section.name == 'Masses':
            masses = parse_whitespace_separated(section.body, {'type': polars.Utf8, 'mass': polars.Float64}, start_line=start_line)
            masses = _apply_type_labels(masses, 'Masses', labels)
        elif section.name == 'Velocities':
            velocities = parse_whitespace_separated(section.body, {
                'i': polars.Int64, 'velocity': polars.Array(polars.Float64, 3),
            }, start_line=start_line)

    # now all 'type's should be in Int32

    if atoms is None:
        if self.headers['atoms'] > 0:
            raise ValueError("Missing required section 'Atoms'")
        return AtomCell(Atoms.empty(), cell=cell, frame='local')

    # next we need to assign element symbols
    # first, if type_map is specified, use that:
    #if type_map_elem is not None and type_map_sym is not None:
    if type_map_df is not None:
        try:
            atoms = checked_left_join(atoms, type_map_df, on='type')
        except CheckedJoinError as e:
            raise ValueError(f"Missing type_map specification for atom type(s): {', '.join(map(repr, e.missing_keys))}")
    elif labels is not None:
        try:
            labels = labels.with_columns(get_elem(labels['symbol']))
        except ValueError as e:
            raise ValueError("Failed to auto-detect elements from type labels. Please pass 'type_map' explicitly") from e
        try:
            atoms = checked_left_join(atoms, labels, 'type')
        except CheckedJoinError as e:
            raise ValueError(f"Missing labels for atom type(s): {', '.join(map(repr, e.missing_keys))}")
    # otherwise we have no way
    else:
        raise ValueError("Failed to auto-detect elements from type labels. Please pass 'type_map' explicitly")

    if velocities is not None:
        # join velocities
        try:
            # TODO use join_asof here?
            atoms = checked_left_join(atoms, velocities, 'i')
        except CheckedJoinError as e:
            raise ValueError(f"Missing velocities for {len(e.missing_keys)}/{len(atoms)} atoms")

    if masses is not None:
        # join masses
        try:
            atoms = checked_left_join(atoms, masses, 'type')
        except CheckedJoinError as e:
            raise ValueError(f"Missing masses for atom type(s): {', '.join(map(repr, e.missing_keys))}")

    return AtomCell(atoms, cell=cell, frame='local')

from_atoms staticmethod 

from_atoms(atoms: HasAtoms) -> LMP
Source code in atomlib/io/lmp.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
@staticmethod
def from_atoms(atoms: HasAtoms) -> LMP:
    if isinstance(atoms, HasAtomCell):
        # we're basically converting everything to the ortho frame, but including the affine shift

        # transform affine shift into ortho frame
        origin = atoms.get_transform('ortho', 'local').to_linear().round_near_zero() \
            .transform(atoms.get_cell().affine.translation())

        # get the orthogonalization transform only, without affine
        ortho = atoms.get_transform('ortho', 'cell_box').to_linear().round_near_zero().inner

        # get atoms in ortho frame, and then add the affine shift
        frame = atoms.get_atoms('ortho').transform_atoms(AffineTransform3D.translate(origin)) \
            .round_near_zero().with_type()
    else:
        bbox = atoms.bbox_atoms()
        ortho = numpy.diag(bbox.size)
        origin = bbox.min

        frame = atoms.get_atoms('local').with_type()

    types = frame.unique(subset='type')
    types = types.with_mass().sort('type')

    now = localtime()
    comment = f"# Generated by atomlib on {now.isoformat(' ', 'seconds')}"

    headers = {}
    sections = []

    headers['atoms'] = len(frame)
    headers['atom types'] = len(types)

    for (s, low, diff) in zip(('x', 'y', 'z'), origin, ortho.diagonal()):
        headers[f"{s}lo {s}hi"] = (low, low + diff)

    headers['xy xz yz'] = (ortho[0, 1], ortho[0, 2], ortho[1, 2])

    body = [
        f" {ty:8} {sym:>4}\n"
        for (ty, sym) in types.select('type', 'symbol').rows()
    ]
    sections.append(LMPSection("Atom Type Labels", tuple(body)))

    if 'mass' in types:
        body = [
            f" {ty:8} {mass:14.7f}  # {sym}\n"
            for (ty, sym, mass) in types.select(('type', 'symbol', 'mass')).rows()
        ]
        sections.append(LMPSection("Masses", tuple(body)))

    body = [
        f" {i+1:8} {ty:4} {x:14.7f} {y:14.7f} {z:14.7f}\n"
        for (i, (ty, (x, y, z))) in enumerate(frame.select(('type', 'coords')).rows())
    ]
    sections.append(LMPSection("Atoms", tuple(body), 'atomic'))

    if (velocities := frame.velocities()) is not None:
        body = [
            f" {i+1:8} {v_x:14.7f} {v_y:14.7f} {v_z:14.7f}\n"
            for (i, (v_x, v_y, v_z)) in enumerate(velocities)
        ]
        sections.append(LMPSection("Velocities", tuple(body)))

    return LMP(comment, headers, tuple(sections))

from_file staticmethod 

from_file(file: FileOrPath) -> LMP
Source code in atomlib/io/lmp.py 
207
208
209
210
@staticmethod
def from_file(file: FileOrPath) -> LMP:
    with open_file(file, 'r') as f:
        return LMPReader(f).parse()

write 

write(file: FileOrPath)
Source code in atomlib/io/lmp.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def write(self, file: FileOrPath):
    with open_file(file, 'w') as f:
        print((self.comment or "") + '\n', file=f)

        # print headers
        for (name, val) in self.headers.items():
            val = _HEADER_FMT.get(name, lambda s: f"{s:8}")(val)
            print(f" {val} {name}", file=f)

        # print sections
        for section in self.sections:
            line = section.name
            if section.style is not None:
                line += f'  # {section.style}'
            print(f"\n{line}\n", file=f)

            f.writelines(section.body)

write_mslice 

write_mslice(
    cell: HasAtomCell,
    f: BinaryFileOrPath,
    template: Optional[MSliceFile] = None,
    *,
    slice_thickness: Optional[float] = None,
    scan_points: Optional[ArrayLike] = None,
    scan_extent: Optional[ArrayLike] = None,
    noise_sigma: Optional[float] = None,
    conv_angle: Optional[float] = None,
    energy: Optional[float] = None,
    defocus: Optional[float] = None,
    tilt: Optional[Tuple[float, float]] = None,
    tds: Optional[bool] = None,
    n_cells: Optional[ArrayLike] = None
)

Write a structure to an mslice file. The structure must be orthogonal and aligned with the local coordinate system. It should be periodic in X and Y.

template may be a file, path, or ElementTree containing an existing mslice file. Its structure will be modified to make the final output. If not specified, a default template will be used.

Additional options modify simulation properties. If an option is not specified, the template's properties are used.

Source code in atomlib/io/mslice.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def write_mslice(cell: HasAtomCell, f: BinaryFileOrPath, template: t.Optional[MSliceFile] = None, *,
                 slice_thickness: t.Optional[float] = None,  # angstrom
                 scan_points: t.Optional[ArrayLike] = None,
                 scan_extent: t.Optional[ArrayLike] = None,
                 noise_sigma: t.Optional[float] = None,  # angstrom
                 conv_angle: t.Optional[float] = None,  # mrad
                 energy: t.Optional[float] = None,  # keV
                 defocus: t.Optional[float] = None,  # angstrom
                 tilt: t.Optional[t.Tuple[float, float]] = None,  # (mrad, mrad)
                 tds: t.Optional[bool] = None,
                 n_cells: t.Optional[ArrayLike] = None):
    """
    Write a structure to an mslice file. The structure must be orthogonal and aligned
    with the local coordinate system. It should be periodic in X and Y.

    ``template`` may be a file, path, or ElementTree containing an existing mslice file.
    Its structure will be modified to make the final output. If not specified, a default
    template will be used.

    Additional options modify simulation properties. If an option is not specified, the
    template's properties are used.
    """
    #if not issubclass(type(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)

    box_size = cell._box_size_in_local()

    # get atoms in local frame (which we verified aligns with the cell's axes)
    # then scale into fractional coordinates
    atoms = cell.get_atoms('linear') \
        .transform(AffineTransform3D.scale(1/box_size)) \
        .with_wobble().with_occupancy()

    out: ElementTree
    if template is None:
        out = default_template()
    elif not isinstance(template, ElementTree):
        with open_file(template, 'r') as temp:
            out = et.parse(temp, None)
    else:
        out = deepcopy(template)

    # TODO clean up this code
    db: t.Optional[Element] = out.getroot() if out.getroot().tag == 'database' else out.find("./database", None)
    if db is None:
        raise ValueError("Couldn't find 'database' tag in template.")

    struct = db.find(".//object[@type='STRUCTURE']", None)
    if struct is None:
        raise ValueError("Couldn't find STRUCTURE object in template.")

    params = db.find(".//object[@type='SIMPARAMETERS']", None)
    if params is None:
        raise ValueError("Couldn't find SIMPARAMETERS object in template.")

    microscope = db.find(".//object[@type='MICROSCOPE']", None)
    if microscope is None:
        raise ValueError("Couldn't find MICROSCOPE object in template.")

    scan = db.find(".//object[@type='SCAN']", None)
    aberrations = db.findall(".//object[@type='ABERRATION']", None)

    def set_attr(struct: Element, name: str, type: str, val: str):
        node = t.cast(t.Optional[Element], struct.find(f".//attribute[@name='{name}']", None))
        if node is None:
            node = t.cast(Element, et.Element('attribute', dict(name=name, type=type), None))
            struct.append(node)
        else:
            node.attrib['type'] = type
        node.text = val  # type: ignore

    def parse_xml_object(obj: Element) -> t.Dict[str, t.Any]:
        """Parse the attributes of a passed XML object."""
        params = {}
        for attr in obj.iterchildren(None):
            if attr.tag == 'attribute':
                params[attr.attrib['name']] = convert_xml_value(attr.text, attr.attrib['type'])
            elif attr.tag == 'relationship':
                # todo give this a better API
                if 'idrefs' in attr.attrib:
                    params[f"{attr.attrib['name']}ID"] = attr.attrib['idrefs']
        return params

    # TODO how to store atoms in unexploded form
    (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)

    (a, b, c) = map(lambda v: f"{v:.8g}", box_size)
    set_attr(struct, 'aparam', 'float', a)
    set_attr(struct, 'bparam', 'float', b)
    set_attr(struct, 'cparam', 'float', c)

    if tilt is not None:
        (tiltx, tilty) = tilt
        set_attr(struct, 'tiltx', 'float', f"{tiltx:.4g}")
        set_attr(struct, 'tilty', 'float', f"{tilty:.4g}")

    if slice_thickness is not None:
        set_attr(params, 'slicethickness', 'float', f"{float(slice_thickness):.8g}")
    if tds is not None:
        set_attr(params, 'includetds', 'bool', str(int(bool(tds))))
    if conv_angle is not None:
        set_attr(microscope, 'aperture', 'float', f"{float(conv_angle):.8g}")
    if energy is not None:
        set_attr(microscope, 'kv', 'float', f"{float(energy):.8g}")
    if noise_sigma is not None:
        if scan is None:
            raise ValueError("New scan specification required for 'noise_sigma'.")
        set_attr(scan, 'noise_sigma', 'float', f"{float(noise_sigma):.8g}")

    if defocus is not None:
        for aberration in aberrations:
            obj = parse_xml_object(aberration)
            if obj['n'] == 1 and obj['m'] == 0:
                set_attr(aberration, 'cnma', 'float', f"{float(defocus):.8g}")  # A, + is over
                set_attr(aberration, 'cnmb', 'float', "0.0")
                break
        else:
            raise ValueError("Couldn't find defocus aberration to modify.")

    if scan_points is not None:
        (nx, ny) = numpy.broadcast_to(scan_points, 2,).astype(int)
        if scan is not None:
            set_attr(scan, 'nx', 'int16', str(nx))
            set_attr(scan, 'ny', 'int16', str(ny))
        else:
            set_attr(params, 'numscanx', 'int16', str(nx))
            set_attr(params, 'numscany', 'int16', str(ny))

    if scan_extent is not None:
        scan_extent = numpy.asarray(scan_extent, dtype=float)
        try:
            if scan_extent.ndim < 2:
                if not scan_extent.shape == (4,):
                    scan_extent = numpy.broadcast_to(scan_extent, (2,))
                    scan_extent = numpy.stack(((0., 0.), scan_extent), axis=-1)
            else:
                scan_extent = numpy.broadcast_to(scan_extent, (2, 2))
        except ValueError as e:
            raise ValueError(f"Invalid scan_extent '{scan_extent}'. Expected an array of shape (2,), (4,), or (2, 2).") from e

        if scan is not None:
            names = ('x_i', 'x_f', 'y_i', 'y_f')
            elem = scan
        else:
            names = ('intx', 'finx', 'inty', 'finy')
            elem = params

        for (name, val) in zip(names, scan_extent.ravel()):
            set_attr(elem, name, 'float', f"{float(val):.8g}")

    # remove existing atoms
    for elem in db.findall("./object[@type='STRUCTUREATOM']", None):
        db.remove(elem)

    # <u^2> -> 1d sigma
    atoms = atoms.with_wobble((polars.col('wobble') / 3.).sqrt())
    rows = atoms.select(('elem', 'coords', '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)
        db.append(e)

    et.indent(db, space="    ", level=0)  # type: ignore

    with open_file_binary(f, 'w') as f:
        doctype = b"""<!DOCTYPE database SYSTEM "file:///System/Library/DTDs/CoreData.dtd">\n"""
        out.write(f, encoding='UTF-8', xml_declaration=True, standalone=True, doctype=doctype)  # type: ignore
        f.write(b'\n')

write_qe 

write_qe(
    atomcell: HasAtomCell,
    f: FileOrPath,
    pseudo: Optional[Mapping[str, str]] = None,
)

Write a structure to a Quantum Espresso pw.x file.

PARAMETER DESCRIPTION
atomcell

Structure to write

TYPE: HasAtomCell

f

File or path to write to

TYPE: FileOrPath

pseudo

Mapping from atom symbol

TYPE: Optional[Mapping[str, str]] DEFAULT: None

Source code in atomlib/io/qe.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def write_qe(atomcell: HasAtomCell, f: FileOrPath, pseudo: t.Optional[t.Mapping[str, str]] = None):
    """
    Write a structure to a Quantum Espresso pw.x file.

    Args:
      atomcell: Structure to write
      f: File or path to write to
      pseudo: Mapping from atom symbol
    """
    if not isinstance(atomcell, HasAtomCell):
        raise TypeError("'qe' format requires an AtomCell.")

    atoms = atomcell.wrap().get_atoms('cell_box').with_mass()

    types = atoms.select(('symbol', 'mass')).unique(subset='symbol').sort('mass')
    if pseudo is not None:
        types = types.with_columns(_get_symbol_mapping(types, pseudo, ty=polars.Utf8).alias('pot'))
    else:
        types = types.with_columns((polars.col('symbol') + polars.lit('.UPF')).alias('pot'))
        #types = types.with_columns(polars.col('symbol').apply(lambda sym: f"{sym}.UPF").alias('pot'))

    with open_file(f, 'w') as f:
        print(f"""\
&SYSTEM
  ibrav=0,
  nat={len(atoms)},
  ntyp={len(types)}
/""", file=f)

        ortho = atomcell.get_transform('local', 'cell_box').to_linear().inner
        print("\nCELL_PARAMETERS angstrom", file=f)
        for row in ortho.T:
            print(f"  {row[0]:12.8f} {row[1]:12.8f} {row[2]:12.8f}", file=f)

        print("\nATOMIC_SPECIES", file=f)
        for (symbol, mass, pot) in types.select(('symbol', 'mass', 'pot')).rows():
            print(f"{symbol:>4} {mass:10.3f}  {pot}", file=f)

        print("\nATOMIC_POSITIONS crystal", file=f)
        for (symbol, (x, y, z)) in atoms.select(('symbol', 'coords')).rows():
            print(f"{symbol:>4} {x:.8f} {y:.8f} {z:.8f}", file=f)

        print(file=f)  # allows for easy concatenation

read_cif 

read_cif(
    f: Union[FileOrPath, CIF, CIFDataBlock],
    block: Union[int, str, None] = None,
) -> HasAtoms

Read a structure from a CIF file.

If block is specified, read data from the given block of the CIF file (index or name).

Source code in atomlib/io/__init__.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def read_cif(f: t.Union[FileOrPath, CIF, CIFDataBlock], block: t.Union[int, str, None] = None) -> HasAtoms:
    """
    Read a structure from a CIF file.

    If `block` is specified, read data from the given block of the CIF file (index or name).
    """

    if isinstance(f, (CIF, CIFDataBlock)):
        cif = f
    else:
        cif = CIF.from_file(f)

    if isinstance(cif, CIF):
        if len(cif) == 0:
            raise ValueError("No data present in CIF file.")
        if block is None:
            if len(cif) > 1:
                logging.warning("Multiple blocks present in CIF file. Defaulting to reading first block.")
            cif = cif.data_blocks[0]
        else:
            cif = cif.get_block(block)

    logging.debug("cif data: %r", cif.data_dict)

    # TODO: support atom_site_Cartn_[xyz]
    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', '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, insert at beginning
        df = df.insert_column(0, get_sym(get_elem(df['label'])))
    atoms = Atoms(df)

    # parse and apply symmetry
    sym_atoms = []
    for sym in cif.get_symmetry():
        sym_atoms.append(atoms.transform(sym))

    #s = '\n'.join(map(str, sym_atoms))
    #print(f"sym_atoms:\n{s}")
    #print(f"atoms: {atoms!s}")

    if len(sym_atoms) > 0:
        atoms = Atoms.concat(sym_atoms)._wrap().deduplicate()

    if (cell_size := cif.cell_size()) is not None:
        cell_size = to_vec3(cell_size)
        if (cell_angle := cif.cell_angle()) is not None:
            # degrees to radians
            cell_angle = to_vec3(cell_angle) * numpy.pi/180.
        return AtomCell.from_unit_cell(atoms, cell_size, cell_angle, frame='cell_frac')
    return Atoms(atoms)

write_cif 

write_cif(
    atoms: Union[HasAtoms, CIF, CIFDataBlock], f: FileOrPath
)

Write a structure to an XSF file.

Source code in atomlib/io/__init__.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
def write_cif(atoms: t.Union[HasAtoms, CIF, CIFDataBlock], f: FileOrPath):
    """Write a structure to an XSF file."""
    if isinstance(atoms, (CIF, CIFDataBlock)):
        cif = atoms
    elif isinstance(atoms, AtomCell):
        cif = CIF((CIFDataBlock.from_atomcell(atoms),))
    else:
        cif = CIF((CIFDataBlock.from_atoms(atoms),))

    cif.write(f)

read_xyz 

read_xyz(f: Union[FileOrPath, XYZ]) -> HasAtoms

Read a structure from an XYZ file.

Source code in atomlib/io/__init__.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
def read_xyz(f: t.Union[FileOrPath, XYZ]) -> HasAtoms:
    """Read a structure from an XYZ file."""
    if isinstance(f, XYZ):
        xyz = f
    else:
        xyz = XYZ.from_file(f)

    atoms = Atoms(xyz.atoms)

    if (cell_matrix := xyz.cell_matrix()) is not None:
        cell = Cell.from_ortho(LinearTransform3D(cell_matrix), pbc=xyz.pbc())
        return AtomCell(atoms, cell, frame='local')
    return Atoms(atoms)

write_xyz 

write_xyz(
    atoms: Union[HasAtoms, XYZ],
    f: FileOrPath,
    fmt: XYZFormat = "exyz",
)
Source code in atomlib/io/__init__.py 
120
121
122
123
def write_xyz(atoms: t.Union[HasAtoms, XYZ], f: FileOrPath, fmt: XYZFormat = 'exyz'):
    if not isinstance(atoms, XYZ):
        atoms = XYZ.from_atoms(atoms)
    atoms.write(f, fmt)

read_xsf 

read_xsf(f: Union[FileOrPath, XSF]) -> HasAtoms

Read a structure from a XSF file.

Source code in atomlib/io/__init__.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def read_xsf(f: t.Union[FileOrPath, XSF]) -> HasAtoms:
    """Read a structure from a XSF file."""
    if isinstance(f, XSF):
        xsf = f
    else:
        xsf = XSF.from_file(f)

    atoms = xsf.get_atoms()
    atoms = atoms.with_columns(get_sym(atoms['elem']))

    if (primitive_cell := xsf.primitive_cell) is not None:
        cell = Cell.from_ortho(primitive_cell, pbc=xsf.get_pbc())
        return AtomCell(atoms, cell, frame='local')
    return Atoms(atoms)

write_xsf 

write_xsf(atoms: Union[HasAtoms, XSF], f: FileOrPath)

Write a structure to an XSF file.

Source code in atomlib/io/__init__.py 
142
143
144
145
146
147
148
149
150
151
def write_xsf(atoms: t.Union[HasAtoms, XSF], f: FileOrPath):
    """Write a structure to an XSF file."""
    if isinstance(atoms, XSF):
        xsf = atoms
    elif isinstance(atoms, AtomCell):
        xsf = XSF.from_cell(atoms)
    else:
        xsf = XSF.from_atoms(atoms)

    xsf.write(f)

read_cfg 

read_cfg(f: Union[FileOrPath, CFG]) -> AtomCell

Read a structure from an AtomEye CFG file.

Source code in atomlib/io/__init__.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def read_cfg(f: t.Union[FileOrPath, CFG]) -> AtomCell:
    """Read a structure from an AtomEye CFG file."""
    if isinstance(f, CFG):
        cfg = f
    else:
        cfg = CFG.from_file(f)

    ortho = cfg.cell
    if cfg.transform is not None:
        ortho = cfg.transform @ ortho

    if cfg.length_scale is not None:
        ortho = ortho.scale(all=cfg.length_scale)

    if cfg.eta is not None:
        m = numpy.eye(3) + 2. * cfg.eta.inner
        # matrix sqrt using eigenvals, eigenvecs
        eigenvals, eigenvecs = numpy.linalg.eigh(m)
        sqrtm = (eigenvecs * numpy.sqrt(eigenvals)) @ eigenvecs.T
        ortho = LinearTransform3D(sqrtm) @ ortho

    frame = Atoms(cfg.atoms).transform(ortho, transform_velocities=True)
    return AtomCell.from_ortho(frame, ortho)

write_cfg 

write_cfg(atoms: Union[HasAtoms, CFG], f: FileOrPath)

Write a structure to an AtomEye CFG file.

Source code in atomlib/io/__init__.py 
179
180
181
182
183
def write_cfg(atoms: t.Union[HasAtoms, CFG], f: FileOrPath):
    """Write a structure to an AtomEye CFG file."""
    if not isinstance(atoms, CFG):
        atoms = CFG.from_atoms(atoms)
    atoms.write(f)

read_lmp 

read_lmp(
    f: Union[FileOrPath, LMP],
    type_map: Optional[Dict[int, Union[str, int]]] = None,
) -> AtomCell

Read a structure from a LAAMPS data file.

Source code in atomlib/io/__init__.py 
186
187
188
189
190
191
192
193
def read_lmp(f: t.Union[FileOrPath, LMP], type_map: t.Optional[t.Dict[int, t.Union[str, int]]] = None) -> AtomCell:
    """Read a structure from a LAAMPS data file."""
    if isinstance(f, LMP):
        lmp = f
    else:
        lmp = LMP.from_file(f)

    return lmp.get_atoms(type_map=type_map)

write_lmp 

write_lmp(atoms: Union[HasAtoms, LMP], f: FileOrPath)

Write a structure to a LAAMPS data file.

Source code in atomlib/io/__init__.py 
196
197
198
199
200
def write_lmp(atoms: t.Union[HasAtoms, LMP], f: FileOrPath):
    """Write a structure to a LAAMPS data file."""
    if not isinstance(atoms, LMP):
        atoms = LMP.from_atoms(atoms)
    atoms.write(f)

read 

read(path: FileOrPath, ty: FileType) -> HasAtoms

read(
    path: Union[str, Path, TextIO], ty: Literal[None] = None
) -> HasAtoms

read(
    path: FileOrPath, ty: Optional[FileType] = None
) -> HasAtoms

Read a structure from a file.

Currently, supported file types are 'cif', 'xyz', and 'xsf'. If no ty is specified, it is inferred from the file's extension.

Source code in atomlib/io/__init__.py 
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
def read(path: FileOrPath, ty: t.Optional[FileType] = None) -> HasAtoms:
    """
    Read a structure from a file.

    Currently, supported file types are 'cif', 'xyz', and 'xsf'.
    If no `ty` is specified, it is inferred from the file's extension.
    """
    if ty is None:
        if isinstance(path, (t.IO, IOBase)):
            try:
                name = path.name  # type: ignore
                if name is None:
                    raise AttributeError()
                ext = Path(name).suffix
            except AttributeError:
                raise TypeError("read() must be passed a file-type when reading an already-open file.") from None
        else:
            name = Path(path).name
            ext = Path(path).suffix

        if len(ext) == 0:
            raise ValueError(f"Can't infer extension for file '{name}'")

        return read(path, t.cast(FileType, ext))

    ty_strip = str(ty).lstrip('.').lower()
    try:
        read_fn = _READ_TABLE[t.cast(FileType, ty_strip)]
    except KeyError:
        raise ValueError(f"Unknown file type '{ty}'") from None
    if read_fn is None:
        raise ValueError(f"Reading is not supported for file type '{ty_strip}'")
    return read_fn(path)

write 

write(atoms: HasAtoms, path: FileOrPath, ty: FileType)

write(
    atoms: HasAtoms,
    path: Union[str, Path, TextIO],
    ty: Literal[None] = None,
)

write(
    atoms: HasAtoms,
    path: FileOrPath,
    ty: Optional[FileType] = None,
)

Write this structure to a file.

A file type may be specified using ty. If no ty is specified, it is inferred from the path's extension.

Source code in atomlib/io/__init__.py 
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
def write(atoms: HasAtoms, path: FileOrPath, ty: t.Optional[FileType] = None):
    """
    Write this structure to a file.

    A file type may be specified using `ty`.
    If no `ty` is specified, it is inferred from the path's extension.
    """

    if ty is None:
        if isinstance(path, (t.IO, IOBase)):
            try:
                name = path.name  # type: ignore
                if name is None:
                    raise AttributeError()
                ext = Path(name).suffix
            except AttributeError:
                raise TypeError("write() must be passed a file-type when reading an already-open file.") from None
        else:
            name = Path(path).name
            ext = Path(path).suffix

        if len(ext) == 0:
            raise ValueError(f"Can't infer extension for file '{name}'")

        return write(atoms, path, t.cast(FileType, ext))

    ty_strip = str(ty).lstrip('.').lower()
    try:
        write_fn = _WRITE_TABLE[t.cast(FileType, ty_strip)]
    except KeyError:
        raise ValueError(f"Unknown file type '{ty}'") from None
    if write_fn is None:
        raise ValueError(f"Writing is not supported for file type '{ty_strip}'")

    return write_fn(atoms, path)