Coverage for cosmolayer/parser/chaos.py: 100%
49 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""CHAOS dataset JSON parser.
3This module parses per-molecule JSON records from the CHAOS dataset
4(arXiv:2511.19002) into the same shape produced by the DMol-3 and TURBOMOLE
5text parsers, so they can be consumed interchangeably by
6:func:`cosmolayer.parser.parser.parse_cosmo_file`.
8Unlike the DMol-3/TURBOMOLE modules, CHAOS records are JSON rather than
9fixed-column text, so this module reads them with :func:`json.loads` and
10dict/list indexing instead of the regex-based machinery in
11:mod:`cosmolayer.parser.common`.
13Two CHAOS-specific unit quirks matter here and are handled below:
15- ``solvation.SegmentList`` segment positions are reported in Bohr (atomic
16 units), unlike ``structural.Coordinates``, which is already in Å.
17- ``solvation.CavVolume`` (and ``solvation.CavArea``, unused here) are also
18 reported in atomic units (Bohr\\ :sup:`3` and Bohr\\ :sup:`2`
19 respectively), even though the per-atom/per-segment ``area``/``charge``
20 fields elsewhere in the ``solvation`` block are already in Ų/e.
21"""
23import json
25import pandas as pd
27from .common import BOHR_TO_ANGSTROM
29FORMAT_NAME = "CHAOS"
31REQUIRED_TOP_LEVEL_KEYS = frozenset({"general", "structural", "solvation"})
33ATOM_POSITION_CONVERSION_FACTOR = 1.0
34SEGMENT_POSITION_CONVERSION_FACTOR = BOHR_TO_ANGSTROM
35VOLUME_CONVERSION_FACTOR = BOHR_TO_ANGSTROM**3
38def _has_required_fields(data: object) -> bool:
39 """Check that a parsed JSON value has the shape a CHAOS record needs.
41 Validates both the top-level keys (``general``, ``structural``,
42 ``solvation``) and the nested fields (``general.AtomList``,
43 ``structural.Coordinates``, ``solvation.SegmentList``,
44 ``solvation.CavVolume``) that :func:`get_atom_dataframe`,
45 :func:`get_segment_dataframe`, and :func:`get_volume` require. A value
46 missing any of these is not recognized as CHAOS JSON, so callers fall
47 through to raising ``ValueError`` instead of a raw
48 ``KeyError``/``TypeError``.
49 """
50 if not (isinstance(data, dict) and REQUIRED_TOP_LEVEL_KEYS <= data.keys()):
51 return False
52 try:
53 data["general"]["AtomList"]
54 data["structural"]["Coordinates"]
55 data["solvation"]["SegmentList"]
56 data["solvation"]["CavVolume"]
57 except (KeyError, TypeError):
58 return False
59 return True
62def is_chaos_json(contents: str) -> bool:
63 """Detect whether ``contents`` is a CHAOS dataset JSON record.
65 Parameters
66 ----------
67 contents : str
68 Candidate file contents.
70 Returns
71 -------
72 bool
73 True if ``contents`` parses as JSON and has the top-level and
74 nested fields a CHAOS record needs (see :func:`_has_required_fields`).
75 """
76 try:
77 data = json.loads(contents)
78 except json.JSONDecodeError:
79 return False
80 return _has_required_fields(data)
83def parse_record(contents: str) -> dict | None:
84 """Parse ``contents`` as a CHAOS JSON record, if it looks like one.
86 Parses ``contents`` exactly once and validates its shape in the same
87 pass, so callers that need both the format-detection answer and the
88 parsed record (e.g. :func:`cosmolayer.parser.parser.parse_cosmo_file`)
89 avoid re-parsing the same JSON text once per field they read.
91 Parameters
92 ----------
93 contents : str
94 Candidate file contents.
96 Returns
97 -------
98 dict or None
99 The parsed record if ``contents`` is valid CHAOS JSON (see
100 :func:`_has_required_fields`), else ``None``.
101 """
102 try:
103 data = json.loads(contents)
104 except json.JSONDecodeError:
105 return None
106 return data if _has_required_fields(data) else None
109def get_atom_dataframe(data: dict) -> pd.DataFrame:
110 """Parse per-atom data from a CHAOS JSON record.
112 Combines ``general.AtomList`` (element symbols) with
113 ``structural.Coordinates`` (Cartesian coordinates, already in Å) in
114 ``general.AtomList`` order, which is the atom-numbering convention used
115 throughout the rest of the record (including
116 ``solvation.SegmentList``'s parent-atom index).
118 Parameters
119 ----------
120 data : dict
121 Contents of a CHAOS JSON file.
123 Returns
124 -------
125 pd.DataFrame
126 Columns: ``id`` (synthesized as ``f"{element}{index}"``), ``x``,
127 ``y``, ``z`` (Å), ``element``.
129 Raises
130 ------
131 ValueError
132 If ``structural.Coordinates`` is JSON ``null``, or any per-atom
133 entry (or component) is ``null``.
134 """
135 atom_list = data["general"]["AtomList"]
136 coordinates = data["structural"]["Coordinates"]
137 if coordinates is None or any(
138 xyz is None or any(c is None for c in xyz) for xyz in coordinates
139 ):
140 raise ValueError(
141 "CHAOS record has null structural.Coordinates; cannot build an atom table."
142 )
143 rows = [
144 {
145 "id": f"{atom['element']}{atom['index']}",
146 "x": xyz[0] * ATOM_POSITION_CONVERSION_FACTOR,
147 "y": xyz[1] * ATOM_POSITION_CONVERSION_FACTOR,
148 "z": xyz[2] * ATOM_POSITION_CONVERSION_FACTOR,
149 "element": atom["element"],
150 }
151 for atom, xyz in zip(atom_list, coordinates, strict=True)
152 ]
153 return pd.DataFrame(rows, columns=["id", "x", "y", "z", "element"])
156def get_segment_dataframe(data: dict) -> pd.DataFrame:
157 """Parse per-segment cavity data from a CHAOS JSON record.
159 Each entry of ``solvation.SegmentList`` is
160 ``[segment_index, parent_atom_index, x, y, z, charge, area, sigma,
161 potential]``, both indices 1-based. ``x, y, z`` are in Bohr and are
162 converted to Šhere; ``charge`` (e) and ``area`` (Ų) need no
163 conversion. ``sigma`` and ``potential`` are dropped, matching the
164 columns produced by the DMol-3/TURBOMOLE parsers.
166 Parameters
167 ----------
168 data : dict
169 Contents of a CHAOS JSON file.
171 Returns
172 -------
173 pd.DataFrame
174 Columns: ``atom`` (0-based index into the atom dataframe from
175 :func:`get_atom_dataframe`), ``x``, ``y``, ``z`` (Å), ``charge``
176 (e), ``area`` (Ų).
178 Raises
179 ------
180 ValueError
181 If ``solvation.SegmentList`` is JSON ``null``, or any entry (or
182 component) is ``null``.
183 """
184 segment_list = data["solvation"]["SegmentList"]
185 if segment_list is None or any(
186 entry is None or any(c is None for c in entry) for entry in segment_list
187 ):
188 raise ValueError(
189 "CHAOS record has null solvation.SegmentList; cannot build a segment table."
190 )
191 rows = [
192 {
193 "atom": parent_atom_index - 1,
194 "x": x * SEGMENT_POSITION_CONVERSION_FACTOR,
195 "y": y * SEGMENT_POSITION_CONVERSION_FACTOR,
196 "z": z * SEGMENT_POSITION_CONVERSION_FACTOR,
197 "charge": charge,
198 "area": area,
199 }
200 for (
201 _segment_index,
202 parent_atom_index,
203 x,
204 y,
205 z,
206 charge,
207 area,
208 _sigma,
209 _potential,
210 ) in segment_list
211 ]
212 return pd.DataFrame(rows, columns=["atom", "x", "y", "z", "charge", "area"])
215def get_volume(data: dict) -> float:
216 """Parse the cavity volume from a CHAOS JSON record.
218 ``solvation.CavVolume`` is reported in Bohr³ and is converted to ų
219 here.
221 Parameters
222 ----------
223 data : dict
224 Contents of a CHAOS JSON file.
226 Returns
227 -------
228 float
229 Cavity volume in ų.
231 Raises
232 ------
233 ValueError
234 If ``solvation.CavVolume`` is JSON ``null``.
235 """
236 cav_volume = data["solvation"]["CavVolume"]
237 if cav_volume is None:
238 raise ValueError(
239 "CHAOS record has null solvation.CavVolume; cannot compute a cavity volume."
240 )
241 return float(cav_volume) * VOLUME_CONVERSION_FACTOR