Coverage for calorine/gpumd/_header.py: 98%
132 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-10-08 12:52 +0000
1"""Parsing of the comment header GPUMD writes on its output files.
3GPUMD opens these files in append mode and writes one header per ``run``, so
4a single file can hold several blocks, each with its own ``dt_output`` and
5``num_atoms``. The readers in :mod:`calorine.gpumd.io` either lay the blocks
6end to end or refuse to combine them, and the helpers here all take the blocks
7that :func:`_read_gpumd_header_blocks` returns.
8"""
10from collections.abc import Iterable
11from dataclasses import dataclass, field
13import numpy as np
14from pandas import DataFrame
17_DT_OUTPUT_UNIT_TO_FS = {'fs': 1.0, 'ps': 1e3}
19# The line names a GPUMD output header uses. A leading `#` block that
20# carries none of them is a comment of some other origin, and the file is
21# read through the historical column-count parsing.
22_GPUMD_HEADER_KEYS = frozenset((
23 'format_version', 'num_atoms', 'cell', 'dt_output', 'columns',
24 'columns_correlation', 'columns_shc',
25 'num_correlation_rows', 'num_frequency_rows', 'num_output_groups',
26))
29@dataclass
30class _HeaderBlock:
31 """One header block and the data rows that follow it."""
33 lines: dict
34 n_rows: int = 0
35 tags: list[str] | None = None
36 num_atoms: int | None = None
37 dt_output_fs: float | None = None
40@dataclass
41class _GpumdHeader:
42 """The header of a GPUMD output file, with ``n_rows`` the number of data
43 rows that follow it. A file written without one has ``tags`` of ``None``
44 and no blocks, and its rows are still counted."""
46 tags: list[str] | None = None
47 blocks: list[_HeaderBlock] = field(default_factory=list)
48 n_rows: int = 0
51def _header_block(lines: dict, n_rows: int, filename: str) -> _HeaderBlock:
52 """Builds a block from its raw tokens, reading ``num_atoms`` as an ``int``
53 and ``dt_output`` in femtoseconds, its unit in the file being ``fs`` for
54 some outputs and ``ps`` for others.
55 """
56 block = _HeaderBlock(lines=lines, n_rows=n_rows)
57 if 'num_atoms' in lines:
58 block.num_atoms = int(lines['num_atoms'][0])
59 if 'dt_output' in lines:
60 tokens = lines['dt_output']
61 if len(tokens) != 2: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true
62 raise ValueError(
63 f'Malformed `dt_output` line in the header of `{filename}`:'
64 f' expected a value and a unit, got {" ".join(tokens) or "nothing"}.'
65 )
66 value, unit = tokens
67 if unit not in _DT_OUTPUT_UNIT_TO_FS:
68 raise ValueError(
69 f'Unrecognized `dt_output` unit `{unit}` in the header of'
70 f' `{filename}`. Known units are:'
71 f' {", ".join(sorted(_DT_OUTPUT_UNIT_TO_FS))}.'
72 )
73 block.dt_output_fs = float(value) * _DT_OUTPUT_UNIT_TO_FS[unit]
74 return block
77def _read_gpumd_header_blocks(filename: str) -> tuple[list[_HeaderBlock], int]:
78 """Returns one block per header of ``filename``, in order, together with
79 the number of data rows the file holds. The block list is empty for a
80 file with no GPUMD header, while the row count is reported either way.
81 """
82 blocks = []
83 with open(filename, 'r') as f:
84 for line in f:
85 line = line.strip()
86 if not line:
87 continue
88 if line.startswith('#'):
89 tokens = line[1:].split()
90 if not tokens:
91 continue
92 # A repeated line name is the only marker of the next header
93 # when the block before it carries no data rows, which is what
94 # a `run` that produced no output leaves behind.
95 if not blocks or blocks[-1][1] > 0 or tokens[0] in blocks[-1][0]:
96 blocks.append([{}, 0])
97 blocks[-1][0][tokens[0]] = tokens[1:]
98 else:
99 if not blocks:
100 blocks.append([{}, 0])
101 blocks[-1][1] += 1
102 n_rows = sum(rows for _, rows in blocks)
103 if not blocks or not _GPUMD_HEADER_KEYS.intersection(blocks[0][0]):
104 return [], n_rows
105 return ([_header_block(lines, rows, filename) for lines, rows in blocks], n_rows)
108def _require_single_block(blocks: list[_HeaderBlock], filename: str, keyword: str):
109 """Raises when ``filename`` holds more than one header block, the
110 output of ``keyword`` being one correlation function per ``run``.
111 """
112 if len(blocks) > 1:
113 raise ValueError(
114 f'`{filename}` holds {len(blocks)} header blocks, one per `run`'
115 f' keyword. The output of `{keyword}` is a separate correlation'
116 ' function per run, and the blocks cannot be combined into one'
117 ' data frame. Write one file per run and read them one at a time.'
118 )
121def _require_header_key(lines: dict, key: str, filename: str):
122 """Returns ``lines[key]``, raising when the header of ``filename`` does
123 not carry that line.
124 """
125 if key not in lines:
126 raise ValueError(
127 f'Header in `{filename}` is missing the required `{key}` line.'
128 )
129 return lines[key]
132def _translate_header_columns(tokens: list[str], tag_map: dict,
133 group_suffixes: bool = False) -> list[str]:
134 """Translates each column token via ``tag_map``, carrying an absent token
135 through unchanged. With ``group_suffixes``, for output that GPUMD writes
136 per group, a trailing ``_<group index>`` is split off before the lookup
137 and reattached after it.
138 """
139 tags = []
140 for token in tokens:
141 if token in tag_map:
142 tags.append(tag_map[token])
143 continue
144 base, separator, index = token.rpartition('_')
145 if group_suffixes and separator and index.isdigit() and base in tag_map:
146 tags.append(f'{tag_map[base]}_{index}')
147 continue
148 tags.append(token)
149 return tags
152def _read_gpumd_header(filename: str, tag_map: dict, required: Iterable[str] = (),
153 group_suffixes: bool = False) -> _GpumdHeader:
154 """Returns the header blocks of ``filename`` with the column names of
155 each translated. ``required`` names the lines every block has to carry
156 besides ``columns``, and ``group_suffixes`` is passed on to
157 :func:`_translate_header_columns`.
158 """
159 blocks, n_rows = _read_gpumd_header_blocks(filename)
160 if not blocks:
161 return _GpumdHeader(n_rows=n_rows)
162 for block in blocks:
163 for key in (*required, 'columns'):
164 _require_header_key(block.lines, key, filename)
165 block.tags = _translate_header_columns(block.lines['columns'], tag_map, group_suffixes)
166 # a run stopped before its first output leaves a block without rows behind
167 tags = next((block.tags for block in blocks if block.n_rows), blocks[0].tags)
168 return _GpumdHeader(tags=tags, blocks=blocks, n_rows=n_rows)
171def _read_gpumd_table(filename: str, header: _GpumdHeader,
172 legacy_tags=None) -> DataFrame:
173 """Reads the data rows of a GPUMD output file into a data frame whose
174 columns come from the header, or from ``legacy_tags(ncols)`` for
175 header-less output. The blocks of the header that are followed by rows
176 have to name the same columns.
177 """
178 numbered_blocks = [(number, block) for number, block in enumerate(header.blocks, start=1)
179 if block.n_rows]
180 for number, block in numbered_blocks[1:]:
181 if block.tags != header.tags:
182 first_number, first_block = numbered_blocks[0]
183 raise ValueError(
184 f'The header blocks of `{filename}` name different columns.'
185 f' Block {first_number} names {" ".join(first_block.lines["columns"])},'
186 f' and block {number} names {" ".join(block.lines["columns"])}.'
187 )
188 if header.n_rows == 0:
189 # `np.loadtxt` warns on a file that holds no data rows, which is what
190 # a `run` interrupted before its first dump leaves behind.
191 return DataFrame(columns=header.tags or [], dtype=float)
192 data = np.loadtxt(filename, ndmin=2)
193 if header.blocks and header.n_rows != data.shape[0]: 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true
194 raise ValueError(
195 f'The headers of `{filename}` are followed by {header.n_rows} data'
196 f' rows, of which {data.shape[0]} could be parsed.'
197 )
198 ncols = data.shape[1]
199 if header.tags is not None:
200 tags = header.tags
201 if ncols != len(tags):
202 raise ValueError(
203 f'Input file contains {ncols} data columns, but its'
204 f' header declares {len(tags)} columns ({filename}).'
205 )
206 else:
207 tags = legacy_tags(ncols)
208 return DataFrame(data=data, columns=tags)
211def _resolve_normalization(normalize: bool, value, blocks: list[_HeaderBlock],
212 n_rows: int, filename: str, name: str):
213 """Returns the divisor of each row, or ``None`` for none. An explicit
214 ``value`` turns normalization on, an unstated ``normalize`` is ``None``
215 and lets ``value`` speak for itself, and ``False`` contradicts it.
216 """
217 if value is not None:
218 if normalize is False:
219 raise ValueError(
220 f'`normalize=False` and `{name}={value}` ask for opposite'
221 f' things. Pass `{name}` on its own to divide by it, and'
222 ' neither of them to leave the data as written.'
223 )
224 if value <= 0:
225 raise ValueError(f'{name} must be positive, got {value}.')
226 return np.full(n_rows, value, dtype=float)
227 if not normalize:
228 return None
229 if not blocks:
230 raise ValueError(
231 f'`{filename}` has no header to take the atom count from.'
232 f' Pass `{name}` to normalize it.'
233 )
234 return np.concatenate(
235 [np.full(block.n_rows, block.num_atoms, dtype=float) for block in blocks])
238def _block_spans(blocks: list[_HeaderBlock]):
239 """Yields each block with the row it starts at and the time in
240 femtoseconds elapsed before it.
241 """
242 start = 0
243 elapsed = 0.0
244 for block in blocks:
245 yield block, start, elapsed
246 start += block.n_rows
247 elapsed += block.n_rows * block.dt_output_fs
250def _insert_time_column(df: DataFrame, blocks: list[_HeaderBlock]) -> DataFrame:
251 """Inserts a ``time`` (ps) column, in place, synthesized from the
252 ``dt_output`` of each block, the first row of a block sitting one
253 ``dt_output`` after the last row of the block before it.
254 """
255 if not blocks:
256 return df
257 times = [elapsed + (np.arange(block.n_rows) + 1) * block.dt_output_fs
258 for block, _, elapsed in _block_spans(blocks)]
259 df.insert(0, 'time', np.concatenate(times) * 1e-3) # fs -> ps
260 return df
263def _offset_time_column(df: DataFrame, blocks: list[_HeaderBlock]) -> DataFrame:
264 """Lays the ``time`` column GPUMD restarts at every ``run`` end to end,
265 in place.
266 """
267 if len(blocks) < 2:
268 return df
269 offsets = np.zeros(len(df))
270 for block, start, elapsed in _block_spans(blocks):
271 offsets[start:start + block.n_rows] = elapsed
272 df['time'] = df['time'].to_numpy() + offsets * 1e-3 # fs -> ps
273 return df