chemparseplot 1.6.1__tar.gz → 1.8.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (84) hide show
  1. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/PKG-INFO +16 -1
  2. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/_version.py +2 -2
  3. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/__init__.py +2 -0
  4. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/__init__.py +2 -0
  5. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/chemgp_hdf5.py +13 -12
  6. chemparseplot-1.8.0/chemparseplot/parse/chemgp_jsonl.py +441 -0
  7. chemparseplot-1.8.0/chemparseplot/parse/eon/_trajectory_common.py +160 -0
  8. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/eon/dimer_trajectory.py +46 -33
  9. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/eon/min_trajectory.py +38 -26
  10. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/eon/minimization.py +3 -3
  11. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/eon/neb.py +39 -5
  12. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/eon/saddle_search.py +53 -33
  13. chemparseplot-1.8.0/chemparseplot/parse/eon/stitch.py +251 -0
  14. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/neb_utils.py +9 -1
  15. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/orca/__init__.py +2 -0
  16. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/orca/neb/interp.py +2 -2
  17. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/orca/neb/opi_parser.py +96 -78
  18. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/plumed.py +28 -36
  19. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/sella/saddle_search.py +41 -22
  20. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/trajectory/__init__.py +2 -0
  21. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/trajectory/hdf5.py +89 -75
  22. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/trajectory/neb.py +17 -6
  23. chemparseplot-1.8.0/chemparseplot/parse/types.py +156 -0
  24. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/plot/__init__.py +9 -0
  25. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/plot/chemgp.py +1 -2
  26. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/plot/geomscan.py +10 -11
  27. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/plot/neb.py +620 -175
  28. chemparseplot-1.8.0/chemparseplot/plot/optimization.py +527 -0
  29. chemparseplot-1.8.0/chemparseplot/plot/structs.py +203 -0
  30. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/plot/theme.py +21 -0
  31. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/units.py +3 -1
  32. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/util.py +29 -32
  33. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/pyproject.toml +19 -1
  34. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/readme.md +3 -0
  35. chemparseplot-1.8.0/tests/_optional_imports.py +24 -0
  36. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/conftest.py +13 -2
  37. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/orca/test_interp.py +8 -0
  38. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_chemgp_hdf5.py +15 -5
  39. chemparseplot-1.8.0/tests/parse/test_chemgp_jsonl_types.py +142 -0
  40. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_dimer_trajectory.py +142 -0
  41. chemparseplot-1.8.0/tests/parse/test_eon_neb_types.py +80 -0
  42. chemparseplot-1.8.0/tests/parse/test_eon_trajectory_common.py +117 -0
  43. chemparseplot-1.8.0/tests/parse/test_min_trajectory.py +274 -0
  44. chemparseplot-1.8.0/tests/parse/test_minimization_result.py +22 -0
  45. chemparseplot-1.8.0/tests/parse/test_neb_stitch.py +70 -0
  46. chemparseplot-1.8.0/tests/parse/test_parser_types.py +86 -0
  47. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_plumed.py +10 -2
  48. chemparseplot-1.8.0/tests/parse/test_saddle_results.py +44 -0
  49. chemparseplot-1.8.0/tests/parse/test_sella_saddle.py +80 -0
  50. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_trajectory_hdf5.py +30 -0
  51. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_trajectory_neb.py +27 -0
  52. chemparseplot-1.8.0/tests/plot/test_neb_helpers.py +113 -0
  53. chemparseplot-1.8.0/tests/plot/test_optimization.py +505 -0
  54. chemparseplot-1.8.0/tests/plot/test_orca_neb_payload.py +43 -0
  55. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/plot/test_strip_rendering.py +47 -6
  56. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/test_coverage_batch.py +34 -4
  57. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/test_full_coverage.py +122 -35
  58. chemparseplot-1.8.0/tests/tutorials/test_chemparseplot.py +132 -0
  59. chemparseplot-1.6.1/chemparseplot/parse/chemgp_jsonl.py +0 -305
  60. chemparseplot-1.6.1/chemparseplot/plot/optimization.py +0 -251
  61. chemparseplot-1.6.1/chemparseplot/plot/structs.py +0 -111
  62. chemparseplot-1.6.1/tests/parse/test_min_trajectory.py +0 -146
  63. chemparseplot-1.6.1/tests/plot/test_optimization.py +0 -220
  64. chemparseplot-1.6.1/tests/tutorials/test_chemparseplot.py +0 -151
  65. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/.gitignore +0 -0
  66. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/LICENSE +0 -0
  67. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/converter.py +0 -0
  68. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/eon/__init__.py +0 -0
  69. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/eon/gprd.py +0 -0
  70. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/file_.py +0 -0
  71. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/orca/geomscan.py +0 -0
  72. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/orca/neb/__init__.py +0 -0
  73. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/patterns.py +0 -0
  74. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/parse/projection.py +0 -0
  75. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/chemparseplot/plot/plumed.py +0 -0
  76. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/orca/test_geomscan.py +0 -0
  77. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_converter.py +0 -0
  78. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_neb_utils.py +0 -0
  79. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_patterns.py +0 -0
  80. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/parse/test_projection.py +0 -0
  81. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/plot/__init__.py +0 -0
  82. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/plot/test_chemgp_utils.py +0 -0
  83. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/plot/test_neb_renderers.py +0 -0
  84. {chemparseplot-1.6.1 → chemparseplot-1.8.0}/tests/plot/test_projection_refactor.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: chemparseplot
3
- Version: 1.6.1
3
+ Version: 1.8.0
4
4
  Summary: Parsers and plotting tools for computational chemistry
5
5
  Project-URL: Documentation, https://chemparseplot.rgoswami.me
6
6
  Project-URL: Issues, https://github.com/HaoZeke/chemparseplot/issues
@@ -49,9 +49,21 @@ Requires-Dist: readcon>=0.7.0; extra == 'neb'
49
49
  Provides-Extra: plot
50
50
  Requires-Dist: cmcrameri>=1.7; extra == 'plot'
51
51
  Requires-Dist: matplotlib>=3.8.2; extra == 'plot'
52
+ Provides-Extra: release
53
+ Requires-Dist: build>=0.10; extra == 'release'
54
+ Requires-Dist: towncrier>=24.8.0; extra == 'release'
55
+ Requires-Dist: twine>=4.0; extra == 'release'
56
+ Requires-Dist: wheel>=0.40; extra == 'release'
52
57
  Provides-Extra: test
58
+ Requires-Dist: ase>=3.22; extra == 'test'
59
+ Requires-Dist: h5py>=3.0; extra == 'test'
60
+ Requires-Dist: matplotlib>=3.8.2; extra == 'test'
61
+ Requires-Dist: polars>=0.20; extra == 'test'
53
62
  Requires-Dist: pytest-cov>=4.1.0; extra == 'test'
54
63
  Requires-Dist: pytest>=7.4.3; extra == 'test'
64
+ Requires-Dist: readcon>=0.7.0; extra == 'test'
65
+ Requires-Dist: rgpycrumbs; extra == 'test'
66
+ Requires-Dist: scipy>=1.11; extra == 'test'
55
67
  Provides-Extra: xyzrender
56
68
  Requires-Dist: xyzrender>=0.1.2; extra == 'xyzrender'
57
69
  Description-Content-Type: text/markdown
@@ -135,6 +147,9 @@ See the [installation guide](https://chemparseplot.rgoswami.me/installation.html
135
147
  - **Unit preserving** throughout via `pint`
136
148
  - **Computation** delegated to [`rgpycrumbs`](https://github.com/HaoZeke/rgpycrumbs) for surface fitting, interpolation,
137
149
  and structure analysis
150
+ - **NEB stitching** (`stitch_neb_segments`) for continuous multi-segment bands (v1.8+)
151
+ - **Metadata-native eOn** CON frames and typed parser results for NEB and single-ended tools
152
+ - **Unit-aware plot helpers** (`convert_neb_values`, shared strip rendering) used by `rgpycrumbs` CLIs
138
153
 
139
154
 
140
155
  <a id="org3c6cfc2"></a>
@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
18
18
  commit_id: str | None
19
19
  __commit_id__: str | None
20
20
 
21
- __version__ = version = '1.6.1'
22
- __version_tuple__ = version_tuple = (1, 6, 1)
21
+ __version__ = version = '1.8.0'
22
+ __version_tuple__ = version_tuple = (1, 8, 0)
23
23
 
24
24
  __commit_id__ = commit_id = None
@@ -3,3 +3,5 @@
3
3
  # SPDX-License-Identifier: MIT
4
4
 
5
5
  from chemparseplot import parse, units
6
+
7
+ __all__ = ["parse", "units"]
@@ -4,6 +4,8 @@
4
4
 
5
5
  from chemparseplot.parse import eon, orca, patterns
6
6
 
7
+ __all__ = ["eon", "orca", "patterns"]
8
+
7
9
  # Lazy imports for modules with optional heavy deps (h5py, pandas)
8
10
  # Import directly: from chemparseplot.parse.chemgp_hdf5 import read_h5_table
9
11
  # Or: from chemparseplot.parse import plumed, projection
@@ -22,6 +22,7 @@ HDF5 Layout
22
22
  from typing import Any
23
23
 
24
24
  import numpy as np
25
+ from chemparseplot.parse.types import ArrayGroup, ParserAttrs
25
26
 
26
27
 
27
28
  def read_h5_table(f: Any, name: str = "table") -> Any:
@@ -88,7 +89,7 @@ def read_h5_grid(
88
89
  return data, x_coords, y_coords
89
90
 
90
91
 
91
- def read_h5_path(f: Any, name: str) -> dict[str, np.ndarray]:
92
+ def read_h5_path(f: Any, name: str) -> ArrayGroup:
92
93
  """Read a path (ordered point sequence).
93
94
 
94
95
  Parameters
@@ -100,14 +101,14 @@ def read_h5_path(f: Any, name: str) -> dict[str, np.ndarray]:
100
101
 
101
102
  Returns
102
103
  -------
103
- dict
104
- Dictionary mapping coordinate names to arrays
104
+ ArrayGroup
105
+ Named mapping of coordinate names to arrays
105
106
  """
106
107
  g = f[f"paths/{name}"]
107
- return {k: g[k][()] for k in g.keys()}
108
+ return ArrayGroup(data={k: g[k][()] for k in g.keys()})
108
109
 
109
110
 
110
- def read_h5_points(f: Any, name: str) -> dict[str, np.ndarray]:
111
+ def read_h5_points(f: Any, name: str) -> ArrayGroup:
111
112
  """Read a point set.
112
113
 
113
114
  Parameters
@@ -119,14 +120,14 @@ def read_h5_points(f: Any, name: str) -> dict[str, np.ndarray]:
119
120
 
120
121
  Returns
121
122
  -------
122
- dict
123
- Dictionary mapping coordinate names to arrays
123
+ ArrayGroup
124
+ Named mapping of coordinate names to arrays
124
125
  """
125
126
  g = f[f"points/{name}"]
126
- return {k: g[k][()] for k in g.keys()}
127
+ return ArrayGroup(data={k: g[k][()] for k in g.keys()})
127
128
 
128
129
 
129
- def read_h5_metadata(f: Any) -> dict[str, Any]:
130
+ def read_h5_metadata(f: Any) -> ParserAttrs:
130
131
  """Read root-level metadata attributes.
131
132
 
132
133
  Parameters
@@ -136,10 +137,10 @@ def read_h5_metadata(f: Any) -> dict[str, Any]:
136
137
 
137
138
  Returns
138
139
  -------
139
- dict
140
- Dictionary of metadata attributes
140
+ ParserAttrs
141
+ Named mapping of metadata attributes
141
142
  """
142
- return {k: f.attrs[k] for k in f.attrs.keys()}
143
+ return ParserAttrs(data={k: f.attrs[k] for k in f.attrs.keys()})
143
144
 
144
145
 
145
146
  def validate_hdf5_structure(
@@ -0,0 +1,441 @@
1
+ # SPDX-FileCopyrightText: 2023-present Rohit Goswami <rog32@hi.is>
2
+ #
3
+ # SPDX-License-Identifier: MIT
4
+
5
+ """Parsers for ChemGP JSONL output formats.
6
+
7
+ ChemGP Rust examples produce JSONL files with method comparison data,
8
+ GP quality grids, and RFF approximation benchmarks. This module provides
9
+ structured parsing into typed containers for downstream plotting.
10
+
11
+ .. versionadded:: 1.5.0
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import json
17
+ from collections import defaultdict
18
+ from dataclasses import dataclass, field
19
+ from pathlib import Path
20
+
21
+ from chemparseplot.parse.types import ParserAttrs
22
+
23
+
24
+ @dataclass(frozen=True, slots=True)
25
+ class ComparisonRecord:
26
+ """One optimizer-step record from a comparison JSONL."""
27
+
28
+ method: str
29
+ oracle_calls: int
30
+ step: int | None = None
31
+ energy: float | None = None
32
+ force: float | None = None
33
+ max_force: float | None = None
34
+
35
+ @classmethod
36
+ def from_mapping(cls, rec: ParserAttrs) -> ComparisonRecord:
37
+ return cls(
38
+ method=str(rec["method"]),
39
+ oracle_calls=int(rec["oracle_calls"]),
40
+ step=int(rec["step"]) if "step" in rec else None,
41
+ energy=float(rec["energy"]) if "energy" in rec else None,
42
+ force=float(rec["force"]) if "force" in rec else None,
43
+ max_force=float(rec["max_force"]) if "max_force" in rec else None,
44
+ )
45
+
46
+
47
+ @dataclass(slots=True)
48
+ class OptimizerTrace:
49
+ """Single optimizer trace from a comparison JSONL.
50
+
51
+ Attributes
52
+ ----------
53
+ method : str
54
+ Optimizer name (e.g. ``"gp_minimize"``, ``"neb"``, ``"otgpd"``).
55
+ steps : list[int]
56
+ Step indices.
57
+ oracle_calls : list[int]
58
+ Cumulative oracle call counts.
59
+ energies : list[float] | None
60
+ Energy at each step (minimize, dimer).
61
+ forces : list[float] | None
62
+ Force norm at each step (dimer: ``force``, NEB: ``max_force``).
63
+ """
64
+
65
+ method: str
66
+ steps: list[int] = field(default_factory=list)
67
+ oracle_calls: list[int] = field(default_factory=list)
68
+ energies: list[float] | None = None
69
+ forces: list[float] | None = None
70
+
71
+ def add_record(self, rec: ComparisonRecord) -> None:
72
+ """Accumulate one typed optimizer record into the trace."""
73
+
74
+ self.steps.append(rec.step if rec.step is not None else len(self.steps))
75
+ self.oracle_calls.append(rec.oracle_calls)
76
+ if rec.energy is not None:
77
+ if self.energies is None:
78
+ self.energies = []
79
+ self.energies.append(rec.energy)
80
+ force_value = rec.force if rec.force is not None else rec.max_force
81
+ if force_value is not None:
82
+ if self.forces is None:
83
+ self.forces = []
84
+ self.forces.append(force_value)
85
+
86
+
87
+ @dataclass(slots=True)
88
+ class ComparisonData:
89
+ """Parsed optimizer comparison from a single JSONL file.
90
+
91
+ Attributes
92
+ ----------
93
+ traces : dict[str, OptimizerTrace]
94
+ Keyed by method name.
95
+ summary : ParserAttrs | None
96
+ Summary record if present.
97
+ """
98
+
99
+ traces: dict[str, OptimizerTrace] = field(default_factory=dict)
100
+ summary: ParserAttrs | None = None
101
+
102
+ def ensure_trace(self, method: str) -> OptimizerTrace:
103
+ """Return the named optimizer trace, creating it if needed."""
104
+
105
+ if method not in self.traces:
106
+ self.traces[method] = OptimizerTrace(method=method)
107
+ return self.traces[method]
108
+
109
+
110
+ def parse_comparison_jsonl(path: str | Path) -> ComparisonData:
111
+ """Parse a ChemGP optimizer comparison JSONL file.
112
+
113
+ Handles minimize, dimer, and NEB comparison formats. Each line is a
114
+ JSON object with a ``method`` field (or ``summary: true``).
115
+
116
+ Parameters
117
+ ----------
118
+ path
119
+ Path to the JSONL file.
120
+
121
+ Returns
122
+ -------
123
+ ComparisonData
124
+ Parsed traces keyed by method name.
125
+ """
126
+ data = ComparisonData()
127
+ with open(path) as f:
128
+ for line in f:
129
+ rec = ParserAttrs(data=json.loads(line.strip()))
130
+ if rec.get("summary"):
131
+ data.summary = rec
132
+ continue
133
+ record = ComparisonRecord.from_mapping(rec)
134
+ trace = data.ensure_trace(record.method)
135
+ trace.add_record(record)
136
+ return data
137
+
138
+
139
+ @dataclass(frozen=True, slots=True)
140
+ class RFFExactRecord:
141
+ """Exact-GP benchmark reference record."""
142
+
143
+ energy_mae: float
144
+ gradient_mae: float
145
+
146
+ @classmethod
147
+ def from_mapping(cls, rec: ParserAttrs) -> RFFExactRecord:
148
+ return cls(
149
+ energy_mae=float(rec["energy_mae"]),
150
+ gradient_mae=float(rec["gradient_mae"]),
151
+ )
152
+
153
+
154
+ @dataclass(frozen=True, slots=True)
155
+ class RFFApproxRecord:
156
+ """One random-feature approximation benchmark record."""
157
+
158
+ d_rff: int
159
+ energy_mae_vs_true: float
160
+ gradient_mae_vs_true: float
161
+ energy_mae_vs_gp: float
162
+ gradient_mae_vs_gp: float
163
+
164
+ @classmethod
165
+ def from_mapping(cls, rec: ParserAttrs) -> RFFApproxRecord:
166
+ return cls(
167
+ d_rff=int(rec["d_rff"]),
168
+ energy_mae_vs_true=float(rec["energy_mae_vs_true"]),
169
+ gradient_mae_vs_true=float(rec["gradient_mae_vs_true"]),
170
+ energy_mae_vs_gp=float(rec["energy_mae_vs_gp"]),
171
+ gradient_mae_vs_gp=float(rec["gradient_mae_vs_gp"]),
172
+ )
173
+
174
+
175
+ @dataclass(slots=True)
176
+ class RFFQualityData:
177
+ """Parsed RFF approximation quality data.
178
+
179
+ Attributes
180
+ ----------
181
+ exact_energy_mae : float
182
+ Exact GP energy MAE vs true surface.
183
+ exact_gradient_mae : float
184
+ Exact GP gradient MAE vs true surface.
185
+ d_rff_values : list[int]
186
+ RFF feature counts tested.
187
+ energy_mae_vs_true : list[float]
188
+ RFF energy MAE vs true surface.
189
+ gradient_mae_vs_true : list[float]
190
+ RFF gradient MAE vs true surface.
191
+ energy_mae_vs_gp : list[float]
192
+ RFF energy MAE vs exact GP.
193
+ gradient_mae_vs_gp : list[float]
194
+ RFF gradient MAE vs exact GP.
195
+ """
196
+
197
+ exact_energy_mae: float = 0.0
198
+ exact_gradient_mae: float = 0.0
199
+ d_rff_values: list[int] = field(default_factory=list)
200
+ energy_mae_vs_true: list[float] = field(default_factory=list)
201
+ gradient_mae_vs_true: list[float] = field(default_factory=list)
202
+ energy_mae_vs_gp: list[float] = field(default_factory=list)
203
+ gradient_mae_vs_gp: list[float] = field(default_factory=list)
204
+
205
+ def add_exact_gp(self, rec: RFFExactRecord) -> None:
206
+ """Store the exact-GP reference metrics."""
207
+
208
+ self.exact_energy_mae = rec.energy_mae
209
+ self.exact_gradient_mae = rec.gradient_mae
210
+
211
+ def add_rff(self, rec: RFFApproxRecord) -> None:
212
+ """Store one RFF approximation record."""
213
+
214
+ self.d_rff_values.append(rec.d_rff)
215
+ self.energy_mae_vs_true.append(rec.energy_mae_vs_true)
216
+ self.gradient_mae_vs_true.append(rec.gradient_mae_vs_true)
217
+ self.energy_mae_vs_gp.append(rec.energy_mae_vs_gp)
218
+ self.gradient_mae_vs_gp.append(rec.gradient_mae_vs_gp)
219
+
220
+
221
+ def parse_rff_quality_jsonl(path: str | Path) -> RFFQualityData:
222
+ """Parse a ChemGP RFF quality JSONL file.
223
+
224
+ Parameters
225
+ ----------
226
+ path
227
+ Path to the JSONL file.
228
+
229
+ Returns
230
+ -------
231
+ RFFQualityData
232
+ Parsed exact GP and RFF metrics.
233
+ """
234
+ data = RFFQualityData()
235
+ with open(path) as f:
236
+ for line in f:
237
+ rec = ParserAttrs(data=json.loads(line.strip()))
238
+ if rec["type"] == "exact_gp":
239
+ data.add_exact_gp(RFFExactRecord.from_mapping(rec))
240
+ elif rec["type"] == "rff":
241
+ data.add_rff(RFFApproxRecord.from_mapping(rec))
242
+ return data
243
+
244
+
245
+ @dataclass(slots=True)
246
+ class GPQualityGrid:
247
+ """GP quality grid data for a single training set size.
248
+
249
+ Attributes
250
+ ----------
251
+ n_train : int
252
+ Number of training points.
253
+ nx : int
254
+ Grid x resolution.
255
+ ny : int
256
+ Grid y resolution.
257
+ x : list[list[float]]
258
+ Grid x coordinates (ny x nx).
259
+ y : list[list[float]]
260
+ Grid y coordinates (ny x nx).
261
+ true_e : list[list[float]]
262
+ True energy on grid.
263
+ gp_e : list[list[float]]
264
+ GP predicted energy on grid.
265
+ gp_var : list[list[float]]
266
+ GP variance on grid.
267
+ train_x : list[float]
268
+ Training point x coordinates.
269
+ train_y : list[float]
270
+ Training point y coordinates.
271
+ train_e : list[float]
272
+ Training point energies.
273
+ """
274
+
275
+ n_train: int = 0
276
+ nx: int = 0
277
+ ny: int = 0
278
+ x: list[list[float]] = field(default_factory=list)
279
+ y: list[list[float]] = field(default_factory=list)
280
+ true_e: list[list[float]] = field(default_factory=list)
281
+ gp_e: list[list[float]] = field(default_factory=list)
282
+ gp_var: list[list[float]] = field(default_factory=list)
283
+ train_x: list[float] = field(default_factory=list)
284
+ train_y: list[float] = field(default_factory=list)
285
+ train_e: list[float] = field(default_factory=list)
286
+
287
+ @classmethod
288
+ def from_records(
289
+ cls,
290
+ *,
291
+ n_train: int,
292
+ meta: ParserAttrs,
293
+ records: list[GPGridRecord],
294
+ train_points: TrainingPointSet | None = None,
295
+ ) -> GPQualityGrid:
296
+ """Build a typed grid from JSONL records and parsed metadata."""
297
+
298
+ nx = int(meta["nx"]) if "nx" in meta else 0
299
+ ny = int(meta["ny"]) if "ny" in meta else 0
300
+ grid = cls(n_train=n_train, nx=nx, ny=ny)
301
+ grid.x = [[0.0] * nx for _ in range(ny)]
302
+ grid.y = [[0.0] * nx for _ in range(ny)]
303
+ grid.true_e = [[0.0] * nx for _ in range(ny)]
304
+ grid.gp_e = [[0.0] * nx for _ in range(ny)]
305
+ grid.gp_var = [[0.0] * nx for _ in range(ny)]
306
+
307
+ for rec in records:
308
+ grid.x[rec.iy][rec.ix] = rec.x
309
+ grid.y[rec.iy][rec.ix] = rec.y
310
+ grid.true_e[rec.iy][rec.ix] = rec.true_e
311
+ grid.gp_e[rec.iy][rec.ix] = rec.gp_e
312
+ grid.gp_var[rec.iy][rec.ix] = rec.gp_var
313
+
314
+ if train_points is not None:
315
+ grid.train_x = list(train_points.x)
316
+ grid.train_y = list(train_points.y)
317
+ grid.train_e = list(train_points.e)
318
+ return grid
319
+
320
+
321
+ @dataclass(frozen=True, slots=True)
322
+ class GPGridRecord:
323
+ """One grid-sample record from the GP-quality JSONL."""
324
+
325
+ ix: int
326
+ iy: int
327
+ x: float
328
+ y: float
329
+ true_e: float
330
+ gp_e: float
331
+ gp_var: float
332
+
333
+ @classmethod
334
+ def from_mapping(cls, rec: ParserAttrs) -> GPGridRecord:
335
+ return cls(
336
+ ix=int(rec["ix"]),
337
+ iy=int(rec["iy"]),
338
+ x=float(rec["x"]),
339
+ y=float(rec["y"]),
340
+ true_e=float(rec["true_e"]),
341
+ gp_e=float(rec["gp_e"]),
342
+ gp_var=float(rec["gp_var"]),
343
+ )
344
+
345
+
346
+ @dataclass(slots=True)
347
+ class TrainingPointSet:
348
+ """Accumulated training points for a single ``n_train`` value."""
349
+
350
+ x: list[float] = field(default_factory=list)
351
+ y: list[float] = field(default_factory=list)
352
+ e: list[float] = field(default_factory=list)
353
+
354
+ def append(self, *, x: float, y: float, energy: float) -> None:
355
+ self.x.append(x)
356
+ self.y.append(y)
357
+ self.e.append(energy)
358
+
359
+
360
+ @dataclass(frozen=True, slots=True)
361
+ class StationaryPoint:
362
+ """A stationary point (minimum or saddle) on the PES."""
363
+
364
+ kind: str # "minimum" or "saddle"
365
+ id: int
366
+ x: float
367
+ y: float
368
+ energy: float
369
+
370
+
371
+ @dataclass(slots=True)
372
+ class GPQualityData:
373
+ """Complete GP quality data from mb_gp_quality.jsonl.
374
+
375
+ Attributes
376
+ ----------
377
+ meta : ParserAttrs
378
+ Grid metadata (nx, ny, x_min, x_max, y_min, y_max).
379
+ stationary : list[StationaryPoint]
380
+ Minima and saddle points.
381
+ grids : dict[int, GPQualityGrid]
382
+ Grid data keyed by n_train.
383
+ """
384
+
385
+ meta: ParserAttrs = field(default_factory=ParserAttrs)
386
+ stationary: list[StationaryPoint] = field(default_factory=list)
387
+ grids: dict[int, GPQualityGrid] = field(default_factory=dict)
388
+
389
+
390
+ def parse_gp_quality_jsonl(path: str | Path) -> GPQualityData:
391
+ """Parse a ChemGP GP quality JSONL file.
392
+
393
+ Parameters
394
+ ----------
395
+ path
396
+ Path to the JSONL file (e.g. ``mb_gp_quality.jsonl``).
397
+
398
+ Returns
399
+ -------
400
+ GPQualityData
401
+ Structured grid data with metadata and stationary points.
402
+ """
403
+ data = GPQualityData()
404
+ train_points = defaultdict(TrainingPointSet)
405
+ grid_records = defaultdict(list)
406
+
407
+ with open(path) as f:
408
+ for line in f:
409
+ rec = ParserAttrs(data=json.loads(line.strip()))
410
+ t = rec["type"]
411
+ if t == "grid_meta":
412
+ data.meta = rec
413
+ elif t in ("minimum", "saddle"):
414
+ data.stationary.append(
415
+ StationaryPoint(
416
+ kind=t,
417
+ id=int(rec["id"]),
418
+ x=float(rec["x"]),
419
+ y=float(rec["y"]),
420
+ energy=float(rec["energy"]),
421
+ )
422
+ )
423
+ elif t == "train_point":
424
+ n = int(rec["n_train"])
425
+ train_points[n].append(
426
+ x=float(rec["x"]),
427
+ y=float(rec["y"]),
428
+ energy=float(rec["energy"]),
429
+ )
430
+ elif t == "grid":
431
+ grid_records[int(rec["n_train"])].append(GPGridRecord.from_mapping(rec))
432
+
433
+ for n_train, records in grid_records.items():
434
+ data.grids[n_train] = GPQualityGrid.from_records(
435
+ n_train=n_train,
436
+ meta=data.meta,
437
+ records=records,
438
+ train_points=train_points.get(n_train),
439
+ )
440
+
441
+ return data