519 lines
17 KiB
Python
Executable file
519 lines
17 KiB
Python
Executable file
#!/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())
|