Skip to content

atomlib.io.lmp

IO for LAAMPS data files, as described here.

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)

LMPSection dataclass

Source code in atomlib/io/lmp.py 
231
232
233
234
235
236
@dataclass
class LMPSection:
    name: str
    body: t.Tuple[str, ...]
    style: t.Optional[str] = None
    start_line: int = 0

name instance-attribute 

name: str

body instance-attribute 

body: Tuple[str, ...]

style class-attribute instance-attribute 

style: Optional[str] = None

start_line class-attribute instance-attribute 

start_line: int = 0

LMPReader

Source code in atomlib/io/lmp.py 
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
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
class LMPReader:
    def __init__(self, f: TextIOBase):
        self.line = 0
        self._file: TextIOBase = f
        self._buf: t.Optional[str] = None

    def _split_comment(self, line: str) -> t.Tuple[str, t.Optional[str]]:
        split = _COMMENT_RE.split(line, maxsplit=1)
        return (split[0], split[1] if len(split) > 1 else None)

    def parse(self) -> LMP:
        # parse comment
        comment = self.next_line(skip_blank=False)
        if comment is None:
            raise ValueError("Unexpected EOF (file is blank)")
        if comment.isspace():
            comment = None
        else:
            comment = comment[:-1]

        headers = self.parse_headers()
        sections = self.parse_sections(headers)

        return LMP(comment, headers, sections)

    def parse_headers(self) -> t.Dict[str, t.Any]:
        headers: t.Dict[str, t.Any] = {}
        while True:
            line = self.peek_line()
            if line is None:
                break
            body = self._split_comment(line)[0]

            if (match := _HEADER_KW_RE.search(body)) is None:
                # probably a body
                break
            self.next_line()

            name = match[0]
            value = body[:match.start(0)].strip()

            try:
                if name in _HEADER_PARSE:
                    value = _HEADER_PARSE[name](value)
            except Exception as e:
                raise ValueError(f"While parsing header '{name}' at line {self.line}: Failed to parse value '{value}") from e

            #print(f"header {name} => {value} (type {type(value)})")
            headers[name] = value

        return headers

    def parse_sections(self, headers: t.Dict[str, t.Any]) -> t.Tuple[LMPSection, ...]:
        first = True

        sections: t.List[LMPSection] = []

        while True:
            start_line = self.line
            line = self.next_line()
            if line is None:
                break
            name, comment = self._split_comment(line)
            name = name.strip()

            try:
                n_lines_header = _SECTION_KWS[name]
            except KeyError:
                if first:
                    raise ValueError(f"While parsing line {self.line}: Unknown header or section keyword '{line}'") from None
                else:
                    raise ValueError(f"While parsing line {self.line}: Unknown section keyword '{line}'") from None

            try:
                if n_lines_header is None:
                    # special case for PairIJ Coeffs:
                    n = int(headers['atom types'])
                    n_lines = (n * (n + 1)) // 2
                else:
                    n_lines = int(headers[n_lines_header])
            except KeyError:
                raise ValueError(f"While parsing body section '{name}' at line {self.line}: "
                                 f"Missing required header '{n_lines_header or 'atom types'}'") from None

            style = comment if name in _SECTION_STYLE_KWS else None
            if style is not None:
                style = style.strip()

            #print(f"section '{name}' @ {self.line}, {n_lines} lines, style {style}")

            lines = self.collect_lines(n_lines)
            if lines is None:
                raise ValueError(f"While parsing body section '{name}' starting at line {self.line}: "
                                 f"Unexpected EOF before {n_lines} lines were read")

            sections.append(LMPSection(
                name, tuple(lines), style, start_line
            ))
            first = False

        return tuple(sections)

    def _try_fill_buf(self, skip_blank: bool = True) -> t.Optional[str]:
        if self._buf is None:
            try:
                # skip blank lines
                while True:
                    self._buf = next(self._file)
                    self.line += 1
                    if not (skip_blank and self._buf.isspace()):
                        break
            except StopIteration:
                pass
        return self._buf

    def peek_line(self, skip_blank: bool = True) -> t.Optional[str]:
        return self._try_fill_buf(skip_blank)

    def next_line(self, skip_blank: bool = True) -> t.Optional[str]:
        line = self._try_fill_buf(skip_blank)
        self._buf = None
        return line

    def collect_lines(self, n: int) -> t.Optional[t.List[str]]:
        assert self._buf is None
        lines = []
        try:
            for _ in range(n):
                while True:
                    line = next(self._file)
                    if not line.isspace():
                        lines.append(line)
                        break
        except StopIteration:
            return None
        self.line += n
        return lines

line instance-attribute 

line = 0

parse 

parse() -> LMP
Source code in atomlib/io/lmp.py 
249
250
251
252
253
254
255
256
257
258
259
260
261
262
def parse(self) -> LMP:
    # parse comment
    comment = self.next_line(skip_blank=False)
    if comment is None:
        raise ValueError("Unexpected EOF (file is blank)")
    if comment.isspace():
        comment = None
    else:
        comment = comment[:-1]

    headers = self.parse_headers()
    sections = self.parse_sections(headers)

    return LMP(comment, headers, sections)

parse_headers 

parse_headers() -> Dict[str, Any]
Source code in atomlib/io/lmp.py 
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
def parse_headers(self) -> t.Dict[str, t.Any]:
    headers: t.Dict[str, t.Any] = {}
    while True:
        line = self.peek_line()
        if line is None:
            break
        body = self._split_comment(line)[0]

        if (match := _HEADER_KW_RE.search(body)) is None:
            # probably a body
            break
        self.next_line()

        name = match[0]
        value = body[:match.start(0)].strip()

        try:
            if name in _HEADER_PARSE:
                value = _HEADER_PARSE[name](value)
        except Exception as e:
            raise ValueError(f"While parsing header '{name}' at line {self.line}: Failed to parse value '{value}") from e

        #print(f"header {name} => {value} (type {type(value)})")
        headers[name] = value

    return headers

parse_sections 

parse_sections(
    headers: Dict[str, Any]
) -> Tuple[LMPSection, ...]
Source code in atomlib/io/lmp.py 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def parse_sections(self, headers: t.Dict[str, t.Any]) -> t.Tuple[LMPSection, ...]:
    first = True

    sections: t.List[LMPSection] = []

    while True:
        start_line = self.line
        line = self.next_line()
        if line is None:
            break
        name, comment = self._split_comment(line)
        name = name.strip()

        try:
            n_lines_header = _SECTION_KWS[name]
        except KeyError:
            if first:
                raise ValueError(f"While parsing line {self.line}: Unknown header or section keyword '{line}'") from None
            else:
                raise ValueError(f"While parsing line {self.line}: Unknown section keyword '{line}'") from None

        try:
            if n_lines_header is None:
                # special case for PairIJ Coeffs:
                n = int(headers['atom types'])
                n_lines = (n * (n + 1)) // 2
            else:
                n_lines = int(headers[n_lines_header])
        except KeyError:
            raise ValueError(f"While parsing body section '{name}' at line {self.line}: "
                             f"Missing required header '{n_lines_header or 'atom types'}'") from None

        style = comment if name in _SECTION_STYLE_KWS else None
        if style is not None:
            style = style.strip()

        #print(f"section '{name}' @ {self.line}, {n_lines} lines, style {style}")

        lines = self.collect_lines(n_lines)
        if lines is None:
            raise ValueError(f"While parsing body section '{name}' starting at line {self.line}: "
                             f"Unexpected EOF before {n_lines} lines were read")

        sections.append(LMPSection(
            name, tuple(lines), style, start_line
        ))
        first = False

    return tuple(sections)

peek_line 

peek_line(skip_blank: bool = True) -> Optional[str]
Source code in atomlib/io/lmp.py 
354
355
def peek_line(self, skip_blank: bool = True) -> t.Optional[str]:
    return self._try_fill_buf(skip_blank)

next_line 

next_line(skip_blank: bool = True) -> Optional[str]
Source code in atomlib/io/lmp.py 
357
358
359
360
def next_line(self, skip_blank: bool = True) -> t.Optional[str]:
    line = self._try_fill_buf(skip_blank)
    self._buf = None
    return line

collect_lines 

collect_lines(n: int) -> Optional[List[str]]
Source code in atomlib/io/lmp.py 
362
363
364
365
366
367
368
369
370
371
372
373
374
375
def collect_lines(self, n: int) -> t.Optional[t.List[str]]:
    assert self._buf is None
    lines = []
    try:
        for _ in range(n):
            while True:
                line = next(self._file)
                if not line.isspace():
                    lines.append(line)
                    break
    except StopIteration:
        return None
    self.line += n
    return lines

write_lmp 

write_lmp(atoms: HasAtoms, f: FileOrPath)
Source code in atomlib/io/lmp.py 
378
379
380
def write_lmp(atoms: HasAtoms, f: FileOrPath):
    LMP.from_atoms(atoms).write(f)
    return