Files

264 lines
9.1 KiB
Python

from __future__ import annotations
import re
import struct
import tarfile
from dataclasses import dataclass
from pathlib import Path
class AmesimResultsError(ValueError):
"""Raised when AMESim result files cannot be parsed consistently."""
@dataclass(frozen=True)
class AmesimVariable:
index: int
label: str
data_path: str | None
param_id: int | None
hidden: bool
@dataclass(frozen=True)
class AmesimResults:
times: tuple[float, ...]
variables: tuple[AmesimVariable, ...]
saved_variable_indices: tuple[int, ...]
series_by_data_path: dict[str, tuple[float, ...]]
final_values_by_data_path: dict[str, float]
@property
def point_count(self) -> int:
return len(self.times)
@property
def saved_variable_count(self) -> int:
return len(self.saved_variable_indices)
def series(self, data_path: str) -> tuple[float, ...]:
return self.series_by_data_path[data_path]
def final_value(self, data_path: str) -> float:
return self.final_values_by_data_path[data_path]
_DATA_PATH_RE = re.compile(r"Data_Path=(\S+)")
_PARAM_ID_RE = re.compile(r"Param_Id=(\d+)")
def load_test_mql_amesim_results(
archive_path: str | Path,
*,
time_stop_s: float | None = None,
) -> AmesimResults:
return load_amesim_results_from_archive(
archive_path=archive_path,
var_member=None,
results_member=None,
time_stop_s=time_stop_s,
)
def load_amesim_results_from_archive(
*,
archive_path: str | Path,
var_member: str | None,
results_member: str | None,
time_stop_s: float | None = None,
) -> AmesimResults:
with tarfile.open(archive_path) as archive:
var_member, results_member = _resolve_result_members(
archive,
var_member=var_member,
results_member=results_member,
)
var_file = archive.extractfile(var_member)
results_file = archive.extractfile(results_member)
if var_file is None:
raise AmesimResultsError(f"Missing AMESim variable member: {var_member}")
if results_file is None:
raise AmesimResultsError(f"Missing AMESim results member: {results_member}")
var_lines = var_file.read().decode("latin1").splitlines()
variables = tuple(
_parse_variable_line(index, line) for index, line in enumerate(var_lines)
)
if time_stop_s is not None:
return _parse_amesim_results_window(
results_file,
variables,
time_stop_s=time_stop_s,
)
results_data = results_file.read()
return parse_amesim_results_bytes(results_data, variables)
def _resolve_result_members(
archive: tarfile.TarFile,
*,
var_member: str | None,
results_member: str | None,
) -> tuple[str, str]:
member_names = set(archive.getnames())
if var_member is not None or results_member is not None:
if var_member is None or results_member is None:
raise AmesimResultsError(
"var_member and results_member must either both be set or both be omitted."
)
return var_member, results_member
preferred = ("test_mql_.var", "test_mql_.results")
if preferred[0] in member_names and preferred[1] in member_names:
return preferred
pairs = sorted(
(name, f"{name[:-4]}.results")
for name in member_names
if name.endswith(".var") and f"{name[:-4]}.results" in member_names
)
if len(pairs) != 1:
raise AmesimResultsError(
"Unable to identify a unique AMESim .var/.results member pair."
)
return pairs[0]
def _parse_amesim_results_window(
results_file,
variables: tuple[AmesimVariable, ...],
*,
time_stop_s: float,
) -> AmesimResults:
header = results_file.read(8)
if len(header) < 8:
raise AmesimResultsError("AMESim results data is too small.")
point_count, encoded_saved_variable_count = struct.unpack("<2i", header)
saved_variable_count = abs(encoded_saved_variable_count)
if point_count <= 0 or saved_variable_count <= 0:
raise AmesimResultsError("Invalid AMESim results header.")
mapping_data = results_file.read(saved_variable_count * 4)
if len(mapping_data) != saved_variable_count * 4:
raise AmesimResultsError("AMESim results variable mapping is truncated.")
saved_variable_indices = struct.unpack(
f"<{saved_variable_count}i",
mapping_data,
)
if any(index < 0 or index >= len(variables) for index in saved_variable_indices):
raise AmesimResultsError(
"AMESim results variable mapping references unknown .var rows."
)
row_length = 1 + saved_variable_count
row_byte_count = row_length * 8
times: list[float] = []
series_lists: dict[str, list[float]] = {}
saved_paths: list[tuple[int, str]] = []
for column, variable_index in enumerate(saved_variable_indices, start=1):
data_path = variables[variable_index].data_path
if data_path is None:
continue
series_lists[data_path] = []
saved_paths.append((column, data_path))
for _row_index in range(point_count):
row = results_file.read(row_byte_count)
if len(row) != row_byte_count:
raise AmesimResultsError("AMESim results matrix is truncated.")
time_s = struct.unpack_from("<d", row, 0)[0]
times.append(time_s)
for column, data_path in saved_paths:
series_lists[data_path].append(
struct.unpack_from("<d", row, column * 8)[0]
)
# Keep one real sample after the requested stop so endpoint finite
# differences do not silently fall back to a backward-only slope.
if time_s > time_stop_s + 1.0e-12:
break
return AmesimResults(
times=tuple(times),
variables=variables,
saved_variable_indices=tuple(saved_variable_indices),
series_by_data_path={
data_path: tuple(values) for data_path, values in series_lists.items()
},
final_values_by_data_path={},
)
def parse_amesim_results_bytes(
results_data: bytes,
variables: tuple[AmesimVariable, ...],
) -> AmesimResults:
if len(results_data) < 8:
raise AmesimResultsError("AMESim results data is too small.")
point_count, encoded_saved_variable_count = struct.unpack_from("<2i", results_data, 0)
saved_variable_count = abs(encoded_saved_variable_count)
if point_count <= 0 or saved_variable_count <= 0:
raise AmesimResultsError("Invalid AMESim results header.")
mapping_offset = 8
mapping_size = saved_variable_count * 4
data_offset = mapping_offset + mapping_size
saved_variable_indices = struct.unpack_from(
f"<{saved_variable_count}i",
results_data,
mapping_offset,
)
if any(index < 0 or index >= len(variables) for index in saved_variable_indices):
raise AmesimResultsError(
"AMESim results variable mapping references unknown .var rows."
)
row_length = 1 + saved_variable_count
main_value_count = point_count * row_length
main_byte_count = main_value_count * 8
main_end = data_offset + main_byte_count
if main_end > len(results_data):
raise AmesimResultsError("AMESim results matrix is truncated.")
main_values = struct.unpack_from(f"<{main_value_count}d", results_data, data_offset)
times = tuple(main_values[row * row_length] for row in range(point_count))
series_by_data_path: dict[str, tuple[float, ...]] = {}
for column, variable_index in enumerate(saved_variable_indices, start=1):
variable = variables[variable_index]
if variable.data_path is None:
continue
series_by_data_path[variable.data_path] = tuple(
main_values[row * row_length + column]
for row in range(point_count)
)
final_values_by_data_path: dict[str, float] = {}
trailing_bytes = len(results_data) - main_end
expected_final_bytes = (1 + len(variables)) * 8
if trailing_bytes >= expected_final_bytes:
final_values = struct.unpack_from(f"<{1 + len(variables)}d", results_data, main_end)
for variable, value in zip(variables, final_values[1:]):
if variable.data_path is not None:
final_values_by_data_path[variable.data_path] = value
return AmesimResults(
times=times,
variables=variables,
saved_variable_indices=tuple(saved_variable_indices),
series_by_data_path=series_by_data_path,
final_values_by_data_path=final_values_by_data_path,
)
def _parse_variable_line(index: int, line: str) -> AmesimVariable:
data_path_match = _DATA_PATH_RE.search(line)
param_id_match = _PARAM_ID_RE.search(line)
label = line
if data_path_match is not None:
label = line[: data_path_match.start()].strip()
return AmesimVariable(
index=index,
label=label,
data_path=data_path_match.group(1) if data_path_match else None,
param_id=int(param_id_match.group(1)) if param_id_match else None,
hidden="HIDDEN" in line,
)