geope.engine

engine.py is a collection of pure function factories. Each returns an un-jitted callable; JIT compilation happens once, lazily, when the optimiser's top-level update_step is first traced. The optimisers read these (lazily built and cached) off the Parameters object — there is no engine class.

Fidelity / infidelity

geope.engine.fidelity(unitary, target_unitary)

Compute the fidelity between a unitary and a target unitary.

The fidelity is defined as the normalised absolute value of the Hilbert-Schmidt inner product between the two matrices.

Parameters:

Name	Type	Description	Default
unitary	Array	The unitary Array to evaluate.	required
target_unitary	Array	The target unitary Array.	required

Returns:

Type	Description
Array	A scalar fidelity Array in the range \([0, 1]\).

Unitary computation

geope.engine.get_compute_matrices_params_list_fn(basis)

Create a partial unitary-computation function with a fixed basis.

Parameters:

Name	Type	Description	Default
basis	ndarray	Array of shape (K, d, d) of Hermitian basis matrices.	required

Returns:

Type	Description
Callable[[Array], Array]	A Callable[[Array], Array] that accepts a parameter list
Callable[[Array], Array]	and returns the product unitary.

Geodesic, Jacobian, Hessian, gammas & omegas

geope.engine.get_geodesic_hamiltonian_fn(target_unitary, projective=True)

Create a partial geodesic Hamiltonian function with a fixed target.

Parameters:

Name	Type	Description	Default
target_unitary	Array	The target unitary Array to bind.	required
projective	bool	If True, return the projective (SU) geodesic. Defaults to True.	True

Returns:

Type	Description
Callable[[Array, Array], Array]	A Callable[[Array, Array], Array] that accepts a unitary and a
Callable[[Array, Array], Array]	JAX random key and returns the geodesic Hamiltonian.

geope.engine.get_jacobian_fn(compute_U_fn)

Build the autodiff Jacobian of the unitary w.r.t. parameters.

Returns the holomorphic jax.jacobian of compute_U_fn. This is the live Jacobian path for all system sizes: the manual Jacobian (geope.jax.jacobian.get_jacobian_propagator) exists and is independently tested, but is not currently wired into the optimisation pipeline (the autodiff path historically overwrote it for the >5-qubit branch — see issue #4). The returned function is left un-jitted so it fuses into the enclosing @jax.jit update step on first optimize().

Parameters:

Name	Type	Description	Default
compute_U_fn	Callable[[Array], Array]	Callable mapping a parameter list to the product unitary.	required

Returns:

Type	Description
Callable[[Array], Array]	A Callable[[Array], Array] returning the Jacobian of the unitary.

geope.engine.get_gammas_fn(compute_U_fn, geo_fn, project_omegas_fn)

Build the projected geodesic-Hamiltonian (gammas) function.

Computes the unitary, its geodesic Hamiltonian towards the target, and projects that onto the Pauli basis (normalised by the dimension). Returned un-jitted so it composes inside an enclosing @jax.jit.

Parameters:

Name	Type	Description	Default
compute_U_fn	Callable[[Array], Array]	Parameter-list -> unitary.	required
geo_fn	Callable[..., Array]	(unitary, key) -> geodesic Hamiltonian.	required
project_omegas_fn	Callable[[Array], Array]	Projection of matrices onto the Lie-algebra basis.	required

Returns:

Type	Description
Callable[[Array, Array], Array]	A Callable[[Array, Array], Array] gammas(free_params, key).

geope.engine.get_omegas_fn(jac_fn, project_omegas_fn, proj_indices, has_proj_drift)

Build the projected per-gate Jacobian (omegas) function.

Projects the Jacobian of each gate (w.r.t. each parameter) onto the Pauli basis, optionally restricting to the projected indices within the combined proj+drift basis. Returned un-jitted so it composes inside an enclosing @jax.jit.

Parameters:

Name	Type	Description	Default
jac_fn	Callable[[Array], Array]	Jacobian of the unitary w.r.t. the free parameters.	required
project_omegas_fn	Callable[[Array], Array]	Projection of matrices onto the Lie-algebra basis.	required
proj_indices	ndarray	Projected indices within the proj+drift basis.	required
has_proj_drift	bool	Whether the proj+drift basis is non-empty (gates the projected-index restriction; mirrors the legacy np.any(proj_drift_basis) check).	required

Returns:

Type	Description
Callable[[Array], Array]	A Callable[[Array], Array] omegas(free_params).

geope.engine.get_gammas_and_omegas_fn(compute_U_fn, jac_fn, geo_fn, project_omegas_fn, proj_indices, has_proj_drift)

Build the combined gammas-and-omegas function used by the GEOPE step.

Gammas are the projected geodesic Hamiltonian coefficients; omegas encode the Jacobian of each gate w.r.t. each parameter, projected onto the Pauli basis. This is the single combined body the GEOPE update step calls (one compute_U_fn and one jac_fn evaluation), matching the legacy numerics; :func:get_gammas_fn / :func:get_omegas_fn are the separately testable halves. Returned un-jitted so it fuses into the enclosing @jax.jit update step on first optimize().

Parameters:

Name	Type	Description	Default
compute_U_fn	Callable[[Array], Array]	Parameter-list -> unitary.	required
jac_fn	Callable[[Array], Array]	Jacobian of the unitary w.r.t. the free parameters.	required
geo_fn	Callable[..., Array]	(unitary, key) -> geodesic Hamiltonian.	required
project_omegas_fn	Callable[[Array], Array]	Projection of matrices onto the Lie-algebra basis.	required
proj_indices	ndarray	Projected indices within the proj+drift basis.	required
has_proj_drift	bool	Whether the proj+drift basis is non-empty.	required

Returns:

Type	Description
Callable[[Array, Array], tuple[Array, Array]]	A Callable[[Array, Array], tuple[Array, Array]]
Callable[[Array, Array], tuple[Array, Array]]	gammas_and_omegas(free_params, key) -> (gammaU_params, omegas).

geope.engine.get_hessian_fn(infid_fn)

Build the full Hessian function via forward-over-reverse HVPs.

Materialises the Hessian of infid_fn by mapping a Hessian-vector product over the identity matrix's columns. Returned un-jitted so it fuses into the enclosing @jax.jit update step.

Parameters:

Name	Type	Description	Default
infid_fn	Callable[[Array], Array]	Scalar-valued infidelity callable of the free parameters.	required

Returns:

Type	Description
Callable[[Array], Array]	A Callable[[Array], Array] hess(y) returning the Hessian.

param_transform helpers

geope.engine.wrap_compute_U_param_transform(params, raw_compute_U)

Wrap compute_U to honour params.param_transform.

The user-facing experimental parameters \(\phi^{\mathrm{exp}}\) are mapped to projected-basis coefficients via params.param_transform (possibly step-dependent), embedded into the proj+drift basis, and combined with the drift before the original raw_compute_U is called.

Returned un-jitted so it fuses into the enclosing @jax.jit update step on first optimize().

Parameters:

Name	Type	Description	Default
params	Parameters	The Parameters object carrying param_transform.	required
raw_compute_U	Callable[[Array], Array]	The projected-basis unitary-computation function.	required

Returns:

Type	Description
Callable[[Array], Array]	The wrapped experimental-space compute_U callable.

geope.engine.get_split_jacobian_fn(compute_U_fn)

Build a real/imag-split Jacobian of compute_U_fn.

Used on the param_transform path: differentiating through the real-valued user transform with a holomorphic Jacobian would discard the imaginary part of intermediates, so the unitary is split into real and imaginary parts, each differentiated, then recombined.

Returned un-jitted so it fuses into the enclosing @jax.jit update step.

Parameters:

Name	Type	Description	Default
compute_U_fn	Callable[[Array], Array]	The (wrapped) experimental-space unitary function.	required

Returns:

Type	Description
Callable[[Array], Array]	A Callable[[Array], Array] returning the complex Jacobian.