GraphIR and Scheduler

Graph-based intermediate representation for model fitting.

The GraphIR is a DAG (directed acyclic graph) of typed nodes connected by explicit dependency edges. It captures the full semantics of a model without prescribing an evaluation strategy.

Three-layer design:

OOP tree – the user-facing Model/Component/Par objects. Handles parsing, validation, user interaction. Unchanged.
GraphIR – a directed acyclic graph of typed nodes with explicit data-dependency edges. Axis-agnostic: works for 1D and 2D models.
ScheduledPlan2D – a flat, packed-array execution schedule compiled from the graph by the 2D backend. No Python objects, no strings, no dicts in the hot path.

class trspecfit.graph_ir.ConvKernelKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Registry IDs for convolution kernel functions.

Only kernel functions listed here can be lowered; other kernels (or non-time packages) fall back to MCP.

BOXCONV = 6

EXPDECAYCONV = 4

EXPRISECONV = 5

EXPSYMCONV = 3

GAUSSCONV = 0

LORENTZCONV = 1

VOIGTCONV = 2

class trspecfit.graph_ir.DomainKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Model domain classification.

Determined by which axes the model operates on:

ENERGY_1D: model has energy axis only.
TIME_1D: model has time axis only.
ENERGY_TIME_2D: model has both axes.

ENERGY_1D = 0

ENERGY_TIME_2D = 2

TIME_1D = 1

class trspecfit.graph_ir.DynFuncKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Dynamics function op codes for the 2D backend.

Backend-specific lowered enum. schedule_2d maps GraphNode.function_name (for DYNAMICS_TRACE nodes) to DynFuncKind during compilation.

ERFFUN = 4

EXPFUN = 0

LINFUN = 2

SINDIVX = 3

SINFUN = 1

SQRTFUN = 5

class trspecfit.graph_ir.EdgeKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Edge types in the model graph.

ADDEND = 3

BASE_INPUT = 2

EXPR_REF = 5

PARAM_INPUT = 0

SPECTRUM_INPUT = 4

TRACE_INPUT = 1

class trspecfit.graph_ir.ExprNodeKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

RPN instruction types for compiled expressions.

ADD = 2

CONST = 0

DIV = 5

MUL = 4

NEG = 6

PARAM_REF = 1

POW = 7

SUB = 3

class trspecfit.graph_ir.ExprProgram(instructions: ndarray)[source]

Bases: object

Compiled expression: flat int array encoding an RPN program.

Encoding: pairs of (node_kind, operand).

CONST: operand is float bits (np.float64.view(np.int64))
PARAM_REF: operand is row index into trace matrix
Operators: operand is 0 (unused)

instructions: ndarray

class trspecfit.graph_ir.GraphEdge(source: int, target: int, kind: EdgeKind, position: int | None = None)[source]

Bases: object

One edge in the model graph.

kind: EdgeKind

position: int | None = None

source: int

target: int

class trspecfit.graph_ir.GraphIR(nodes: list[~trspecfit.graph_ir.GraphNode], edges: list[~trspecfit.graph_ir.GraphEdge], domain: ~trspecfit.graph_ir.DomainKind, energy: ~numpy.ndarray | None = None, time: ~numpy.ndarray | None = None, node_by_name: dict[str, int] = <factory>)[source]

Bases: object

Directed acyclic graph representing a model.

Axis-agnostic: works for 1D and 2D models. A 1D energy model has time=None; adding dynamics populates time and promotes domain to ENERGY_TIME_2D.

domain: DomainKind

edges: list[GraphEdge]

energy: ndarray | None = None

node_by_name: dict[str, int]

nodes: list[GraphNode]

time: ndarray | None = None

to_dot(*, collapse_profiles: bool = True) → str[source]

Return a Graphviz DOT string for this graph.

Node shapes and colours encode NodeKind; edge labels encode EdgeKind. The output can be rendered with dot -Tpng or any Graphviz viewer.

Parameters:: collapse_profiles (bool, default=True) – When True, per-sample profile nodes (PROFILE_SAMPLE, per-sample COMPONENT_EVAL, per-sample EXPRESSION) are collapsed into single representative nodes showing the sample count. This keeps profile models readable.

class trspecfit.graph_ir.GraphNode(id: int, kind: ~trspecfit.graph_ir.NodeKind, name: str, source_order: int, value: float | None = None, function_name: str | None = None, package: str | None = None, expr_string: str | None = None, vary: bool = False, bounds: tuple[float, float] | None = None, arrays: dict[str, ~numpy.ndarray] = <factory>)[source]

Bases: object

One node in the model graph.

arrays: dict[str, ndarray]

bounds: tuple[float, float] | None = None

expr_string: str | None = None

function_name: str | None = None

id: int

kind: NodeKind

name: str

package: str | None = None

source_order: int

value: float | None = None

vary: bool = False

class trspecfit.graph_ir.NodeKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Node types in the model graph.

COMPONENT_EVAL = 5

CONVOLUTION = 100

DYNAMICS_TRACE = 2

EXPRESSION = 4

OPT_PARAM = 1

PARAM_PLUS_TRACE = 3

PROFILE_AVERAGE = 102

PROFILE_SAMPLE = 101

SPECTRUM_FED_OP = 7

STATIC_PARAM = 0

SUBCYCLE_MASK = 103

SUBCYCLE_REMAP = 104

SUM = 6

trspecfit.graph_ir.OP_DISPATCH: dict[int, tuple[Callable[[...], Any], bool]] = {0: (<function Gauss>, False), 1: (<function GaussAsym>, False), 2: (<function Lorentz>, False), 3: (<function Voigt>, False), 4: (<function GLS>, False), 5: (<function GLP>, False), 6: (<function DS>, False), 10: (<function Offset>, False), 11: (<function LinBack>, False), 12: (<function Shirley>, True)}: Maps OpKind → (eval_function, needs_spectrum). Single source of truth for component dispatch – used by both the evaluator hot path and constant-component precomputation.

class trspecfit.graph_ir.OpKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

2D backend component function op codes.

Backend-specific lowered enum. schedule_2d maps GraphNode.function_name to OpKind during compilation.

DS = 6

GAUSS = 0

GAUSS_ASYM = 1

GLP = 5

GLS = 4

LINBACK = 11

LORENTZ = 2

OFFSET = 10

SHIRLEY = 12

VOIGT = 3

class trspecfit.graph_ir.ParamSourceKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Parameter source kinds for profiled component evaluation (1D and 2D).

PROFILE_EXPR = 2

PROFILE_SAMPLE = 1

SCALAR = 0

class trspecfit.graph_ir.ProfileFuncKind(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)[source]

Bases: IntEnum

Profile-function op codes for lowered 1D profile evaluation.

PEXPDECAY = 0

PGAUSS = 2

PLINEAR = 1

class trspecfit.graph_ir.ScheduledPlan1D(energy: ndarray, n_params: int, param_values_init: ndarray, opt_indices: ndarray, opt_param_names: list[str], n_expressions: int, expr_target_indices: ndarray, expr_programs: list[ExprProgram], n_aux: int, aux_axis: ndarray, n_profile_samples: int, profile_sample_base_indices: ndarray, profile_sample_component_indptr: ndarray, profile_component_func_ids: ndarray, profile_component_param_indptr: ndarray, profile_component_param_indices: ndarray, n_profile_exprs: int, profile_expr_programs: list[ExprProgram], n_ops: int, op_kinds: ndarray, op_param_indptr: ndarray, op_param_source_kinds: ndarray, op_param_indices: ndarray, op_needs_spectrum: ndarray, op_is_pre_spectrum: ndarray, op_is_profiled: ndarray, op_is_constant: ndarray, cached_result: ndarray, cached_peak_sum: ndarray)[source]

Bases: object

Compiled 1D execution schedule for ENERGY_1D models.

Simpler than ScheduledPlan2D: no time axis, no dynamics, no trace matrix. Parameters are stored as a flat (n_params,) scalar vector.

aux_axis: ndarray

cached_peak_sum: ndarray

cached_result: ndarray

energy: ndarray

expr_programs: list[ExprProgram]

expr_target_indices: ndarray

n_aux: int

n_expressions: int

n_ops: int

n_params: int

n_profile_exprs: int

n_profile_samples: int

op_is_constant: ndarray

op_is_pre_spectrum: ndarray

op_is_profiled: ndarray

op_kinds: ndarray

op_needs_spectrum: ndarray

op_param_indices: ndarray

op_param_indptr: ndarray

op_param_source_kinds: ndarray

opt_indices: ndarray

opt_param_names: list[str]

param_values_init: ndarray

profile_component_func_ids: ndarray

profile_component_param_indices: ndarray

profile_component_param_indptr: ndarray

profile_expr_programs: list[ExprProgram]

profile_sample_base_indices: ndarray

profile_sample_component_indptr: ndarray

class trspecfit.graph_ir.ScheduledPlan2D(energy: ndarray, time: ndarray, n_params: int, n_time: int, param_traces_init: ndarray, opt_indices: ndarray, opt_param_names: list[str], n_dyn_groups: int, dyn_group_target_row: ndarray, dyn_group_base_row: ndarray, dyn_group_indptr: ndarray, dyn_sub_func_id: ndarray, dyn_sub_param_rows: ndarray, dyn_sub_n_params: ndarray, dyn_sub_time_axes: ndarray, dyn_sub_masks: ndarray, n_expressions: int, expr_target_rows: ndarray, expr_programs: list[ExprProgram], resolution_kinds: ndarray, resolution_indices: ndarray, n_conv_steps: int, conv_target_rows: ndarray, conv_func_ids: ndarray, conv_param_indptr: ndarray, conv_param_rows: ndarray, conv_support_indptr: ndarray, conv_support_values: ndarray, n_aux: int, aux_axis: ndarray, n_profile_samples: int, profile_sample_base_rows: ndarray, profile_sample_component_indptr: ndarray, profile_component_func_ids: ndarray, profile_component_param_indptr: ndarray, profile_component_param_rows: ndarray, n_profile_exprs: int, profile_expr_programs: list[ExprProgram], n_ops: int, op_schedule: ndarray, op_kinds: ndarray, op_param_indptr: ndarray, op_param_source_kinds: ndarray, op_param_indices: ndarray, op_needs_spectrum: ndarray, op_is_pre_spectrum: ndarray, op_is_profiled: ndarray, op_is_constant: ndarray, cached_result: ndarray, cached_peak_sum: ndarray)[source]

Bases: object

Compiled 2D execution schedule.

No Python objects in the hot path (except expr_programs).

aux_axis: ndarray

cached_peak_sum: ndarray

cached_result: ndarray

conv_func_ids: ndarray

conv_param_indptr: ndarray

conv_param_rows: ndarray

conv_support_indptr: ndarray

conv_support_values: ndarray

conv_target_rows: ndarray

dyn_group_base_row: ndarray

dyn_group_indptr: ndarray

dyn_group_target_row: ndarray

dyn_sub_func_id: ndarray

dyn_sub_masks: ndarray

dyn_sub_n_params: ndarray

dyn_sub_param_rows: ndarray

dyn_sub_time_axes: ndarray

energy: ndarray

expr_programs: list[ExprProgram]

expr_target_rows: ndarray

n_aux: int

n_conv_steps: int

n_dyn_groups: int

n_expressions: int

n_ops: int

n_params: int

n_profile_exprs: int

n_profile_samples: int

n_time: int

op_is_constant: ndarray

op_is_pre_spectrum: ndarray

op_is_profiled: ndarray

op_kinds: ndarray

op_needs_spectrum: ndarray

op_param_indices: ndarray

op_param_indptr: ndarray

op_param_source_kinds: ndarray

op_schedule: ndarray

opt_indices: ndarray

opt_param_names: list[str]

param_traces_init: ndarray

profile_component_func_ids: ndarray

profile_component_param_indptr: ndarray

profile_component_param_rows: ndarray

profile_expr_programs: list[ExprProgram]

profile_sample_base_rows: ndarray

profile_sample_component_indptr: ndarray

resolution_indices: ndarray

resolution_kinds: ndarray

time: ndarray

class trspecfit.graph_ir.SymbolicRPN(instructions: list[tuple[ExprNodeKind, float | str | None]], referenced_names: list[str])[source]

Bases: object

Symbolic RPN program with parameter references by name.

This is the frontend output of the expression compiler. schedule_2d binds names to trace-matrix row indices and produces the final ExprProgram.

Each instruction is a (ExprNodeKind, operand) pair:

CONST: operand is the float value itself
PARAM_REF: operand is the parameter name (str)
Operators: operand is None

instructions: list[tuple[ExprNodeKind, float | str | None]]

referenced_names: list[str]

trspecfit.graph_ir.build_graph(model: Model) → GraphIR[source]

Walk the OOP Model tree and emit a GraphIR.

Parameters:: model (Model) – A fully-constructed model (components, pars, and any dynamics / profile models already attached).
Returns:: The semantic DAG representation.
Return type:: GraphIR

trspecfit.graph_ir.can_lower_1d(graph: GraphIR) → bool[source]

Check whether the 1D NumPy backend can compile this graph.

Parameters:: graph (GraphIR) – The model graph to check.
Returns:: True if schedule_1d can compile this graph.
Return type:: bool

trspecfit.graph_ir.can_lower_2d(graph: GraphIR) → bool[source]

Check whether the 2D NumPy backend can compile this graph.

Parameters:: graph (GraphIR) – The model graph to check.
Returns:: True if schedule_2d can compile this graph.
Return type:: bool

trspecfit.graph_ir.compile_expr_symbolic(expr_string: str) → SymbolicRPN[source]

Parse an arithmetic expression string into symbolic RPN.

Uses the Python ast module to walk the expression tree and emit a postfix instruction sequence. Parameter references are kept as name strings; the scheduler resolves them to row indices.

Parameters:: expr_string (str) – Arithmetic expression (e.g. "3/4*GLP_01_A").
Returns:: The symbolic RPN program.
Return type:: SymbolicRPN
Raises:: ValueError – If the expression contains unsupported AST nodes.

trspecfit.graph_ir.schedule_1d(graph: GraphIR) → ScheduledPlan1D[source]

Compile a GraphIR into a flat 1D execution schedule.

Parameters:: graph (GraphIR) – Must pass can_lower_1d(graph).
Returns:: Packed-array execution schedule for evaluate_1d.
Return type:: ScheduledPlan1D
Raises:: ValueError – If the graph cannot be lowered (domain, unsupported nodes, etc.).

trspecfit.graph_ir.schedule_2d(graph: GraphIR) → ScheduledPlan2D[source]

Compile a GraphIR into a flat 2D execution schedule.

Parameters:: graph (GraphIR) – Must pass can_lower_2d(graph).
Returns:: Packed-array execution schedule for evaluate_2d.
Return type:: ScheduledPlan2D
Raises:: ValueError – If the graph cannot be lowered (domain, unsupported nodes, etc.).