neuralqx.hilbert.constraints.utils package

avmapr(f, x, key)

Apply a PRNG-keyed function to a state or a batch of states, recursively over leading axes.

This helper is useful for “randomized per-state updates” where a function needs a JAX PRNGKey and should be applied independently to each element of a batch (and also to nested batches).

Behaviour: - If x.ndim == 1: call f(x, key) and return the result. - Otherwise: split key into x.shape[0] subkeys and vmap recursively over axis 0.

This recursion makes it work for shapes like (B, N) as well as nested shapes such as (B, G, E) where you want independent randomness per element of the outermost batch.

Parameters:

  • f – Callable with signature f(x_leaf, key_leaf) returning an updated array.

  • x – Input array. If 1D, treated as a single state; otherwise treated as batched on axis 0.

  • key – JAX PRNGKey used (and split) to provide independent randomness per batch element.

Returns:

Output of applying f to x with appropriate key splitting and vectorization.

Raises:

ValueError – If key cannot be split into the required number of subkeys.

ensure_batchdim(x)

Ensure that an array has an explicit leading batch dimension.

This utility normalizes state arrays to the common shape (B, N):

  • If x.ndim == 2: return (x, False) (already batched).

  • Otherwise: return (x[None, :], True) (a single state promoted to batch size 1).

The boolean flag is intended to support round-tripping: callers can later remove the added batch dimension if was_single is True.

Parameters:

x – State array of shape (N,) or batched states of shape (B, N).

Returns:

Tuple (xb, was_single) where xb has shape (B, N) and was_single indicates whether a batch dimension was added.

Raises:

ValueError – If x is not 1D or 2D and cannot be interpreted as (single|batched) states.

preprocess_constraint(ce)

Preprocess a string-encoded gauge-fixing constraint list into integer/sign data.

This function converts constraints of the form:

[
    [["1"], ["2", "-0"]],
    [["3"], ["N", "-2", "5"]],
    ...
]

into a list of (inc_data, out_data) pairs where: - inc_data is [(idx, sign), ...] for incoming edges, - out_data is [(idx, sign), ...] for outgoing edges, with idx as an integer edge index and sign in {+1, -1}.

Sign parsing: - A token like "-7" contributes (7, -1). - A token like "2" contributes (2, +1). - If the first outgoing token is exactly "N", it acts as a global minus for that constraint:

all outgoing signs are multiplied by -1 and the "N" token is consumed/removed.

This representation is intended to be friendly to compiled kernels (e.g., Numba) and avoids any string parsing during hot loops.

Parameters:

ce – List of constraint pairs [[inc_tokens], [out_tokens]] where tokens are strings like "2", "-0", or "N".

Returns:

List of processed constraints [(inc_data, out_data), ...] with integer indices and explicit ±1 signs.

Raises:

ValueError – If any non-"N" token cannot be converted to an integer index.

preprocess_constraints_for_jax(ce)

Preprocess a string-encoded gauge-fixing constraint list into a JAX-friendly Python structure.

This function performs the same sign/index conversion as preprocess_constraint(), but returns a structure designed to be consumed by JAX code that iterates in Python over a fixed-size list:

  • The returned value is a plain Python list of pairs (inc_data, out_data).

  • Each of inc_data and out_data is a Python list of (idx, sign) tuples.

The outgoing global-minus token "N" is supported: - If the first outgoing token is "N", all outgoing signs are multiplied by -1 and "N"

is removed from the outgoing list.

This representation is typically fed into a JAX-jitted constraint checker that loops over the list and builds small JAX arrays inside the loop.

Parameters:

ce – List of constraint pairs [[inc_tokens], [out_tokens]] where tokens are strings like "2", "-0", or "N".

Returns:

List of processed constraints suitable for JAX-friendly loops: [(inc_data, out_data), ...] with inc_data/out_data as lists of (idx, ±1).

Raises:

ValueError – If any non-"N" token cannot be converted to an integer index.

generate_constraint_array(edges)

Generate a gauge-fixing constraint array from an oriented multigraph edge list.

Given an oriented, keyed edge list describing a connected undirected graph, this function builds a spanning tree and expresses each tree edge as an integer-signed sum of chord edges (non-tree edges). The resulting constraints can be used as a gauge fixing that enforces vertex-wise conservation (Kirchhoff / Gauß law) when interpreted as modular equalities in the U(1) Hilbert space.

Output format: - Returns a list of constraints, each of the form [[lhs_edge_str], rhs_edge_strs] where:

  • lhs_edge_str is a single string representing the fixed (tree) edge,

  • rhs_edge_strs is a list of signed strings representing chord edges contributing to the RHS.

Each edge is rendered as "(u, v, key)" using the provided edge tuples.

Algorithm overview: 1) Collect and index vertices, build the incidence matrix B (rows: vertices, cols: edges). 2) Drop one row to obtain full rank (connected graphs have rank n-1). 3) Build a spanning tree via BFS to select n-1 tree edges. 4) Partition edges into tree edges T and chord edges C. 5) Solve the linear system to express f_T in terms of f_C and round coefficients to integers. 6) Emit one constraint per tree edge.

Parameters:

edges (List[Tuple]) – List of oriented keyed edges (u, v, key). The graph is treated as undirected for connectivity/spanning-tree construction, but orientation is used in the incidence matrix.

Return type:

List[List[List[str]]]

Returns:

Gauge-fixing constraint array as list[list[list[str]]] suitable for downstream parsing and pretty-printing.

Raises:

  • TypeError – If vertex labels are not sortable/comparable (required by the current implementation’s use of sorted(...)).

  • ValueError – If the underlying graph is disconnected or the spanning tree cannot be built (may lead to an invalid solve).

  • numpy.linalg.LinAlgError – If the tree incidence submatrix is singular (typically indicates the graph is not connected or the constructed tree is invalid).

map_edge_string(edge_str, coords_to_int)

Rewrite vertex-coordinate tuples inside an edge string using a provided coordinate mapping.

This helper supports pretty-printing constraints for non-planar embeddings where vertices may be represented by coordinate triples. It scans an edge_str for parenthesized tuples, attempts to parse them as comma-separated floats, and replaces any tuple present in coords_to_int by its mapped vertex identifier.

The function is intentionally conservative: - If a matched tuple cannot be parsed as floats, it is left unchanged. - Only tuples that appear as keys in coords_to_int are replaced.

Parameters:

  • edge_str – Edge string potentially containing coordinate tuples, e.g. "((1.1, 2.2, 3.3), (4.4, 5.5, 6.6), 0)".

  • coords_to_int – Mapping from coordinate tuples to vertex identifiers (e.g. ints or names).

Returns:

A rewritten string where recognized coordinate tuples have been replaced by their mapped identifiers.

Raises:

None – This function does not intentionally raise exceptions; unparsable tuples are skipped.

pretty_format_constraints(constraints, coordinates_map=None)

Format gauge-fixing constraints into a human-readable multi-line string.

Each constraint is expected to have the form:

[
    [fixed_edge_as_string],
    [signed_edge_strings...]
]

This function formats each constraint as a short description:

  • The LHS edge is shown as the “fixed” edge.

  • The RHS is printed as a signed expression, one line per constraint.

  • If coordinates_map is provided, coordinate tuples in edge strings are rewritten using map_edge_string() to recover original vertex identifiers.

Robustness: - If the constraint list is empty, a short message is returned. - Malformed or incomplete constraint entries are reported inline in the output.

Parameters:

  • constraints (list[list[list[str]]]) – List of constraints in the string-array format produced by generate_constraint_array() (or equivalent).

  • coordinates_map (Optional[Dict]) – Optional mapping from vertex identifiers to coordinate tuples. If provided, a reverse map is built and used to replace coordinate tuples in edge strings with vertex ids.

Return type:

str

Returns:

A formatted string suitable for logging or displaying autogenerated gauge fixings.

Raises:

TypeError – If coordinates_map contains unhashable coordinate values when building the reverse lookup map.