"""1D evaluator for the compiled backend.
All component functions live in ``trspecfit.functions.energy`` as the
single source of truth. For 1D evaluation, parameters are plain
scalars (no ``(n_time, 1)`` broadcasting needed).
"""
from __future__ import annotations
import numpy as np
from trspecfit.graph_ir import (
ScheduledPlan1D,
_eval_expr_scalar,
_evaluate_profile_expr_values,
_evaluate_profile_sample_values,
_evaluate_scheduled_op_1d,
)
# ---------------------------------------------------------------------------
# Core 1D evaluator
# ---------------------------------------------------------------------------
#
[docs]
def evaluate_1d(plan: ScheduledPlan1D, theta: np.ndarray) -> np.ndarray:
"""Evaluate the compiled 1D model at optimizer parameters *theta*.
Parameters
----------
plan
Immutable compiled execution schedule from ``schedule_1d``.
theta
``(n_opt,)`` optimizer parameter vector. Order must match
``plan.opt_param_names``.
Returns
-------
ndarray
``(n_energy,)`` model spectrum.
Raises
------
ValueError
If ``len(theta) != len(plan.opt_indices)``.
"""
if len(theta) != len(plan.opt_indices):
raise ValueError(
f"theta length {len(theta)} does not match "
f"plan.opt_indices length {len(plan.opt_indices)}"
)
# 1a. Copy parameter values -> scratch
values = plan.param_values_init.copy()
# 1b. Write optimizer params
values[plan.opt_indices] = theta
# 1c. Resolve expressions in topological order
for i in range(plan.n_expressions):
target = int(plan.expr_target_indices[i])
values[target] = _eval_expr_scalar(
plan.expr_instructions[plan.expr_indptr[i] : plan.expr_indptr[i + 1]],
values,
)
profile_sample_values = _evaluate_profile_sample_values(
plan.aux_axis,
values,
plan.profile_sample_base_indices,
plan.profile_sample_component_indptr,
plan.profile_component_func_ids,
plan.profile_component_param_indptr,
plan.profile_component_param_indices,
)
profile_expr_values = _evaluate_profile_expr_values(
values,
profile_sample_values,
plan.n_params,
plan.profile_expr_instructions,
plan.profile_expr_indptr,
)
# 2. Component evaluation
energy = plan.energy
result = plan.cached_result.copy()
peak_sum = plan.cached_peak_sum.copy()
for op_idx in range(plan.n_ops):
if plan.op_is_constant[op_idx]:
continue
kind = int(plan.op_kinds[op_idx])
start = int(plan.op_param_indptr[op_idx])
end = int(plan.op_param_indptr[op_idx + 1])
needs_spectrum = bool(plan.op_needs_spectrum[op_idx])
is_pre = bool(plan.op_is_pre_spectrum[op_idx])
component = _evaluate_scheduled_op_1d(
energy,
kind,
plan.op_param_source_kinds[start:end],
plan.op_param_indices[start:end],
values,
profile_sample_values,
profile_expr_values,
peak_sum,
needs_spectrum=needs_spectrum,
is_profiled=bool(plan.op_is_profiled[op_idx]),
n_aux=plan.n_aux,
)
result += component
if is_pre:
peak_sum += component
return result