Heavy calculations almost done. Added script for automatization "fill_table_from_nb.py"

This commit is contained in:
Artyom Titov 2026-07-29 17:48:33 +03:00
commit 3585c75b2a
39 changed files with 125384 additions and 88606 deletions

519
automation/fill_table_from_nb.py Executable file
View file

@ -0,0 +1,519 @@
#!/usr/bin/env python3
"""Заполнение Excel-шаблона результатами из Mathematica Notebook (.nb).
Пример:
python fill_table_from_nb.py \
--template table_template.xlsx \
--notebooks h_174-66_Pu240_2-3_osc_Qw1.nb \
--output table_filled.xlsx
Можно передать сразу несколько notebook-файлов. Имя каждого файла должно содержать:
_2-3_ или _4-5_ блок строк в таблице;
_osc_, _pr_ или _gid_ группу столбцов;
_Qw1 или _Qw2 левую или правую половину таблицы.
"""
from __future__ import annotations
import argparse
import glob
import math
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable
# Границы блоков в конкретном шаблоне.
PAIR_START_ROW = {
"2-3": 5,
"4-5": 19,
}
# Смещение группы относительно первого столбца данных своей половины.
MODEL_COLUMN_OFFSET = {
"osc": 0, # Jосц/J
"pr": 3, # Jпр/J
"gid": 6, # Jгид/J
}
# Порядок столбцов внутри каждой группы шаблона.
METHOD_COLUMN_OFFSET = {
"heavy": 0,
"new": 1,
"light": 2,
}
# Смещение строк относительно строки Qb.
VALUE_ROW_OFFSET = {
"Qb": 0,
"Qw": 1,
"hb": 2,
"hw": 3,
"L": 4,
"Q": 5,
"int1": 6,
"int2": 7,
}
QB_TO_METHOD = {
"Qb3": "heavy",
"Qb2": "new",
"Qb1": "light",
}
GROUP_RE = re.compile(r"Cell\[CellGroupData\[\{(.*?)\},\s*Open\s*\]\]", re.S)
INPUT_LABEL_RE = re.compile(r'CellLabel->\s*"In\[(\d+)\]:="')
ASSIGNMENT_RE = re.compile(r'RowBox\[\{"([A-Za-z][A-Za-z0-9]*)",\s*"="')
SIMPLE_OUTPUT_RE = re.compile(r'Cell\[BoxData\["([^"]+)"\],\s*"Output"', re.S)
@dataclass(frozen=True)
class NbRecord:
order: int
input_number: int
expression: str
assignments: tuple[str, ...]
value: float | None
@dataclass(frozen=True)
class Target:
pair: str
model: str
qw_side: int
@dataclass
class MethodValues:
Qb: float | None = None
Qw: float | None = None
hb: float | None = None
hw: float | None = None
L: float | None = None
Q: float | None = None
int1: float | None = None
int2: float | None = None
def as_dict(self) -> dict[str, float | None]:
return {
"Qb": self.Qb,
"Qw": self.Qw,
"hb": self.hb,
"hw": self.hw,
"L": self.L,
"Q": self.Q,
"int1": self.int1,
"int2": self.int2,
}
def mathematica_number_to_float(raw: str) -> float | None:
"""Преобразует 1.23`*^-4 и подобные числа Mathematica в float."""
value = raw.strip()
value = re.sub(r"`(?:\d+(?:\.\d*)?)?", "", value)
value = value.replace("*^", "e")
special = {
"Infinity": math.inf,
"-Infinity": -math.inf,
"Indeterminate": None,
}
if value in special:
return special[value]
try:
return float(value)
except ValueError:
return None
def parse_records(nb_path: Path) -> list[NbRecord]:
text = nb_path.read_text(encoding="utf-8", errors="replace")
records: list[NbRecord] = []
for order, group_match in enumerate(GROUP_RE.finditer(text)):
group = group_match.group(1)
input_label = INPUT_LABEL_RE.search(group)
if input_label is None:
continue
input_end = group.find(']], "Input"')
if input_end == -1:
input_end = group.find('], "Input"')
expression = group[:input_end] if input_end != -1 else group
outputs = SIMPLE_OUTPUT_RE.findall(group)
value = mathematica_number_to_float(outputs[-1]) if outputs else None
assignments = tuple(ASSIGNMENT_RE.findall(expression))
records.append(
NbRecord(
order=order,
input_number=int(input_label.group(1)),
expression=expression,
assignments=assignments,
value=value,
)
)
if not records:
raise ValueError(f"В файле {nb_path.name} не найдены группы Input/Output.")
return records
def is_q_result(record: NbRecord) -> bool:
expr = record.expression
return (
'FractionBox["q"' in expr
and "SqrtBox" in expr
and 'RowBox[{"S1", "*", "S2"}]' in expr
)
def is_integral(record: NbRecord) -> bool:
return 'SubsuperscriptBox["\\[Integral]"' in record.expression
def first_value_for_assignment(
records: Iterable[NbRecord], assignment_name: str
) -> float | None:
for record in records:
if assignment_name in record.assignments and record.value is not None:
return record.value
return None
def extract_methods(records: list[NbRecord]) -> dict[str, MethodValues]:
"""Выделяет heavy/new/light без привязки к номерам In[]/Out[]."""
q_indexes = [i for i, record in enumerate(records) if is_q_result(record)]
if len(q_indexes) < 3:
raise ValueError(
"Найдено меньше трёх выражений q/Sqrt[S1*S2]. "
"Проверьте, что notebook полностью вычислен и сохранён с Output-ячейками."
)
extracted: dict[str, MethodValues] = {}
previous_q_index = -1
for q_index in q_indexes:
search_window = records[previous_q_index + 1 : q_index + 1]
qb_candidates: list[tuple[int, str]] = []
for local_index, record in enumerate(search_window):
for qb_name in QB_TO_METHOD:
if qb_name in record.assignments and record.value is not None:
qb_candidates.append((local_index, qb_name))
if not qb_candidates:
previous_q_index = q_index
continue
qb_local_index, qb_name = qb_candidates[-1]
method = QB_TO_METHOD[qb_name]
section = search_window[qb_local_index:]
values = MethodValues()
values.Qb = first_value_for_assignment(section, qb_name)
values.hb = first_value_for_assignment(section, "hb")
values.Q = records[q_index].value
if values.Q is None:
q_value = first_value_for_assignment(section, "q")
s1_value = first_value_for_assignment(section, "S1")
s2_value = first_value_for_assignment(section, "S2")
if (
q_value is not None
and s1_value is not None
and s2_value is not None
and s1_value * s2_value > 0
):
values.Q = q_value / math.sqrt(s1_value * s2_value)
if method == "heavy":
values.Qw = first_value_for_assignment(section, "Qw3")
values.hw = first_value_for_assignment(section, "hw")
values.L = first_value_for_assignment(section, "L")
# Первый числовой интеграл после o1 — int1, после o2 — int2.
pending_integral: str | None = None
oscillator_values: dict[str, float | None] = {"int1": None, "int2": None}
for record in section:
if "o1" in record.assignments:
pending_integral = "int1"
oscillator_values["int1"] = record.value
continue
if "o2" in record.assignments:
pending_integral = "int2"
oscillator_values["int2"] = record.value
continue
if pending_integral and is_integral(record) and record.value is not None:
setattr(values, pending_integral, record.value)
pending_integral = None
# Аналитический запасной вариант:
# integral(0..inf) 2*x^2*o*exp(-o*x^2) dx = sqrt(pi)/(2*sqrt(o)).
for integral_name, oscillator_value in oscillator_values.items():
if (
getattr(values, integral_name) is None
and oscillator_value is not None
and oscillator_value > 0
):
setattr(
values,
integral_name,
math.sqrt(math.pi) / (2.0 * math.sqrt(oscillator_value)),
)
extracted[method] = values
previous_q_index = q_index
required_methods = {"heavy", "new", "light"}
missing_methods = required_methods - extracted.keys()
if missing_methods:
raise ValueError(
"Не удалось выделить методы: " + ", ".join(sorted(missing_methods))
)
required_by_method = {
"heavy": ("Qb", "Qw", "hb", "hw", "L", "Q", "int1", "int2"),
"new": ("Qb", "hb", "Q", "int1", "int2"),
"light": ("Qb", "hb", "Q", "int1", "int2"),
}
for method, required_names in required_by_method.items():
missing_values = [
name
for name in required_names
if getattr(extracted[method], name) is None
]
if missing_values:
raise ValueError(
f"Для метода {method} не найдены значения: "
+ ", ".join(missing_values)
)
return extracted
def extract_metadata(records: list[NbRecord]) -> dict[str, float | None]:
return {
name: first_value_for_assignment(records, name)
for name in ("Z1", "Z2", "d", "b1", "b2")
}
def detect_target(nb_path: Path) -> Target:
name = nb_path.stem
lower = name.lower()
pair_candidates = re.findall(r"(?:^|_)(\d+-\d+)(?=_|$)", name)
pair = next((item for item in pair_candidates if item in PAIR_START_ROW), None)
if pair is None:
raise ValueError(
f"В имени {nb_path.name} не найден поддерживаемый блок 2-3 или 4-5."
)
model: str | None = None
for candidate in MODEL_COLUMN_OFFSET:
if re.search(rf"(?:^|[_-]){candidate}(?:[_-]|$)", lower):
model = candidate
break
if model is None:
raise ValueError(
f"В имени {nb_path.name} не найден тип расчёта osc, pr или gid."
)
qw_match = re.search(r"qw[_-]?([12])(?:[_-]|$)", lower)
if qw_match is None:
raise ValueError(f"В имени {nb_path.name} не найдено Qw1 или Qw2.")
return Target(pair=pair, model=model, qw_side=int(qw_match.group(1)))
def format_metadata(pair: str, metadata: dict[str, float | None]) -> str:
def integer_or_blank(value: float | None) -> str:
if value is None:
return ""
if float(value).is_integer():
return str(int(value))
return f"{value:.3g}"
def compact_decimal(value: float | None) -> str:
if value is None:
return ""
text = f"{value:.2f}"
if text.startswith("0."):
text = text[1:]
elif text.startswith("-0."):
text = "-" + text[2:]
return text
return (
f"{pair}\n"
f"Z1={integer_or_blank(metadata['Z1'])}\n"
f"Z2={integer_or_blank(metadata['Z2'])}\n"
f"d={integer_or_blank(metadata['d'])}\n"
f"b1={compact_decimal(metadata['b1'])}\n"
f"b2={compact_decimal(metadata['b2'])}"
)
def write_notebook_to_sheet(
worksheet,
target: Target,
methods: dict[str, MethodValues],
metadata: dict[str, float | None],
written_cells: dict[str, str],
source_name: str,
) -> list[str]:
start_row = PAIR_START_ROW[target.pair]
# Qw1: подпись в A, данные с C. Qw2: подпись в M, данные с N.
if target.qw_side == 1:
metadata_column = 1
first_data_column = 3
else:
metadata_column = 13
first_data_column = 14
worksheet.cell(row=start_row, column=metadata_column).value = format_metadata(
target.pair, metadata
)
report: list[str] = []
model_start_column = first_data_column + MODEL_COLUMN_OFFSET[target.model]
for method_name in ("heavy", "new", "light"):
column = model_start_column + METHOD_COLUMN_OFFSET[method_name]
values = methods[method_name].as_dict()
for value_name, row_offset in VALUE_ROW_OFFSET.items():
value = values[value_name]
if value is None:
continue
row = start_row + row_offset
coordinate = worksheet.cell(row=row, column=column).coordinate
previous_source = written_cells.get(coordinate)
if previous_source is not None and previous_source != source_name:
print(
f"Предупреждение: {coordinate} уже заполнялась из "
f"{previous_source}; новое значение берётся из {source_name}.",
file=sys.stderr,
)
cell = worksheet.cell(row=row, column=column)
cell.value = value
cell.number_format = "0.000"
written_cells[coordinate] = source_name
report.append(f"{coordinate}={value:.12g}")
return report
def expand_notebook_arguments(patterns: list[str]) -> list[Path]:
paths: list[Path] = []
for pattern in patterns:
matches = [Path(p) for p in glob.glob(pattern)]
if matches:
paths.extend(matches)
else:
path = Path(pattern)
if path.exists():
paths.append(path)
else:
raise FileNotFoundError(f"Не найден notebook: {pattern}")
unique_paths: list[Path] = []
seen: set[Path] = set()
for path in paths:
resolved = path.resolve()
if resolved not in seen:
seen.add(resolved)
unique_paths.append(path)
return unique_paths
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--template", required=True, type=Path, help="Excel-шаблон")
parser.add_argument(
"--notebooks",
required=True,
nargs="+",
help="Один или несколько .nb файлов; разрешены маски *.nb",
)
parser.add_argument("--output", required=True, type=Path, help="Итоговый .xlsx")
parser.add_argument("--sheet", default="Лист1", help="Имя листа шаблона")
parser.add_argument(
"--dry-run",
action="store_true",
help="Только разобрать notebook и показать найденные значения",
)
args = parser.parse_args()
notebook_paths = expand_notebook_arguments(args.notebooks)
parsed_items = []
for nb_path in notebook_paths:
records = parse_records(nb_path)
target = detect_target(nb_path)
methods = extract_methods(records)
metadata = extract_metadata(records)
parsed_items.append((nb_path, target, methods, metadata))
print(
f"{nb_path.name}: pair={target.pair}, model={target.model}, "
f"Qw{target.qw_side}"
)
for method_name in ("heavy", "new", "light"):
values = methods[method_name].as_dict()
visible = ", ".join(
f"{key}={value:.12g}"
for key, value in values.items()
if value is not None
)
print(f" {method_name}: {visible}")
if args.dry_run:
return 0
try:
from openpyxl import load_workbook
except ImportError as exc:
raise SystemExit(
"Не установлен openpyxl. Выполните: python -m pip install openpyxl"
) from exc
if not args.template.exists():
raise FileNotFoundError(f"Не найден шаблон: {args.template}")
# Открытие через бинарный поток позволяет использовать шаблон без расширения .xlsx.
with args.template.open("rb") as template_file:
workbook = load_workbook(template_file)
if args.sheet not in workbook.sheetnames:
raise KeyError(
f"В шаблоне нет листа {args.sheet!r}. Доступны: "
+ ", ".join(workbook.sheetnames)
)
worksheet = workbook[args.sheet]
written_cells: dict[str, str] = {}
for nb_path, target, methods, metadata in parsed_items:
report = write_notebook_to_sheet(
worksheet=worksheet,
target=target,
methods=methods,
metadata=metadata,
written_cells=written_cells,
source_name=nb_path.name,
)
print(f" Записано {len(report)} ячеек: " + ", ".join(report))
args.output.parent.mkdir(parents=True, exist_ok=True)
workbook.save(args.output)
print(f"Готово: {args.output}")
return 0
if __name__ == "__main__":
raise SystemExit(main())