Optimising with frozen signal-chain parameters¶

In this tutorial, we demonstrate the freeze_keys feature of Optimiser.optimise(). When a signal chain component (such as a TransferFunc) is used directly as a Hamiltonian coefficient, it may contain parameters that represent fixed hardware properties — filter coefficients, timing boundaries, etc. — that should never be modified by the optimiser.

This tutorial covers:

1. System setup: PWCPulse → TransferFunc → Hamiltonian (direct, no wrapper)
2. Inspecting frozen parameters via opt_prob.inactive_params
3. Optimisation without freeze_keys — demonstrating the problem
4. Optimisation with freeze_keys — the correct approach
5. Summary and best practices

To ensure numerical precision throughout, we enable 64-bit floating point in JAX:

In [1]:

import jax

jax.config.update("jax_enable_x64", True)

import jax jax.config.update("jax_enable_x64", True)

1. System setup¶

We model a two-level qubit with a stationary detuning term ($\sigma_z$) and an external drive ($\sigma_x$) shaped by a piecewise-constant (PWC) pulse filtered through a Gaussian rise-time transfer function.

The signal chain is:

PWCPulse("drive")  →  TransferFunc("tf")  →  Hamiltonian coefficient

• PWCPulse provides the optimisable envelope (env, active) plus timing parameters (t0, dt) that are frozen — they encode the hardware sampling grid.
• TransferFunc applies a Gaussian filter whose coefficients (b, a) and time window (t0, t1) are frozen — they represent fixed hardware characteristics.
• The TransferFunc instance is passed directly to the Hamiltonian constructor via its Ht parameter, allowing Hamiltonian.inactive_params to discover all frozen keys automatically.

We begin by setting up the simulation grid and pulse parameters.

In [2]:

import jax.numpy as jnp
from qruise.toolset.utils import compute_rise_time_gaussian_coeffs

t0 = 0.0  # start time (s)
tfinal = 20e-9  # end time (s)
grid_size = 50  # simulation grid points
pulse_size = 10  # number of PWC steps
rise_time = 3e-9  # hardware rise time (s)

ts_pulse = jnp.linspace(t0, tfinal, pulse_size)
dt = float(ts_pulse[1] - ts_pulse[0])

import jax.numpy as jnp from qruise.toolset.utils import compute_rise_time_gaussian_coeffs t0 = 0.0 # start time (s) tfinal = 20e-9 # end time (s) grid_size = 50 # simulation grid points pulse_size = 10 # number of PWC steps rise_time = 3e-9 # hardware rise time (s) ts_pulse = jnp.linspace(t0, tfinal, pulse_size) dt = float(ts_pulse[1] - ts_pulse[0])

The optimisable input signal is a piecewise-constant (PWC) pulse. We mark its envelope samples (env) as active and its timing grid (t0, dt) as frozen — the timing encodes a fixed hardware sampling grid that the optimiser must not modify.

In [3]:

from qruise.toolset.control_stack import PWCPulse

pwc = PWCPulse(
    "drive",
    {
        "env": (jnp.ones(pulse_size) * 1e7, True),  # active — will be optimised
        "t0": (t0, False),  # frozen — hardware timing
        "dt": (dt, False),  # frozen — hardware timing
    },
)

from qruise.toolset.control_stack import PWCPulse pwc = PWCPulse( "drive", { "env": (jnp.ones(pulse_size) * 1e7, True), # active — will be optimised "t0": (t0, False), # frozen — hardware timing "dt": (dt, False), # frozen — hardware timing }, )

The TransferFunc applies a Gaussian low-pass filter that models the hardware rise time. We compute its IIR filter coefficients using compute_rise_time_gaussian_coeffs:

In [4]:

abs_coeffs = compute_rise_time_gaussian_coeffs(rise_time, n_order=4)
b = jnp.array([1.0])
a = jnp.array(abs_coeffs)

abs_coeffs = compute_rise_time_gaussian_coeffs(rise_time, n_order=4) b = jnp.array([1.0]) a = jnp.array(abs_coeffs)

TransferFunc needs a time derivative of its source signal (dsource), which is computed with jax.grad. Since jax.grad requires a scalar-to-scalar function and PWCPulse.__call__ operates on a batched time array, we wrap it in a thin scalar helper first:

In [5]:

# scalar wrapper for jax.grad (scalar t → scalar output)
def pwc_signal(t, params):
    return pwc(t, params)


def scalar_pwc_signal(t, params):
    return pwc_signal(jnp.array([t]), params)[0]


# time derivative of the signal
dsource = jax.grad(scalar_pwc_signal, argnums=0)

# scalar wrapper for jax.grad (scalar t → scalar output) def pwc_signal(t, params): return pwc(t, params) def scalar_pwc_signal(t, params): return pwc_signal(jnp.array([t]), params)[0] # time derivative of the signal dsource = jax.grad(scalar_pwc_signal, argnums=0)

We can now construct the TransferFunc, marking all filter properties as frozen:

In [6]:

from qruise.toolset import TransferFunc

tf = TransferFunc(
    "tf",
    {
        "b": (b, False),  # frozen — filter numerator
        "a": (a, False),  # frozen — filter denominator
        "t0": (t0, False),  # frozen — filter start time
        "t1": (tfinal, False),  # frozen — filter end time
    },
    pwc_signal,
    dsource=dsource,
)

from qruise.toolset import TransferFunc tf = TransferFunc( "tf", { "b": (b, False), # frozen — filter numerator "a": (a, False), # frozen — filter denominator "t0": (t0, False), # frozen — filter start time "t1": (tfinal, False), # frozen — filter end time }, pwc_signal, dsource=dsource, )

---

Key point: tf (the TransferFunc instance) is passed directly to the Hamiltonian constructor via its Ht parameter — not wrapped in a plain Python function. This allows Hamiltonian.inactive_params to introspect its frozen keys automatically.

---

We now merge the PWCPulse and TransferFunc parameter dicts into a single flat collection using ParameterCollection, and define the Hamiltonian. Note that tf is passed directly as the Hamiltonian coefficient — not wrapped in a plain Python function. This is what enables Hamiltonian.inactive_params to introspect frozen keys automatically, as highlighted in the note above.

In [7]:

from qruise.toolset import Hamiltonian, ParameterCollection
from qutip import sigmaz, sigmax

# merge PWCPulse and TransferFunc parameters into a single flat dict
pc = ParameterCollection()
pc.add_dict(pwc.params | tf.params)
params = pc.get_collection()

delta = 1e8  # qubit detuning (Hz)

# tf passed directly as Hamiltonian coefficient
H = Hamiltonian(delta * sigmaz() / 2, [(sigmax() / 2, tf)])

from qruise.toolset import Hamiltonian, ParameterCollection from qutip import sigmaz, sigmax # merge PWCPulse and TransferFunc parameters into a single flat dict pc = ParameterCollection() pc.add_dict(pwc.params | tf.params) params = pc.get_collection() delta = 1e8 # qubit detuning (Hz) # tf passed directly as Hamiltonian coefficient H = Hamiltonian(delta * sigmaz() / 2, [(sigmax() / 2, tf)])

/opt/hostedtoolcache/Python/3.11.14/x64/lib/python3.11/site-packages/qruise/toolset/hamiltonian.py:46: UserWarning: 
                    You have provided a function for the envolpe that has
                    more than 2 arguments. Make sure the first argument is
                    reselved for time and the second one for the dictionary
                    of parameters. The rest should have default values.
                    
  elif len(C)>2:warnings.warn('\n                    You have provided a function for the envolpe that has\n                    more than 2 arguments. Make sure the first argument is\n                    reselved for time and the second one for the dictionary\n                    of parameters. The rest should have default values.\n                    ')

With the Hamiltonian in place, we define the QOCProblem as usual, specifying the initial state, target state, and state infidelity loss function.

In [8]:

from qruise.toolset import QOCProblem

y0 = jnp.array([1.0, 0.0], dtype=jnp.complex128)  # initial state |0>
y_t = jnp.array([0.0, 1.0], dtype=jnp.complex128)  # target state |1>


def loss(x, y):
    """State infidelity: 1 - |<x|y>|^2"""
    return jnp.real(1.0 - jnp.abs(jnp.vdot(x, y)) ** 2)


opt_prob = QOCProblem(H, y0, params, (t0, tfinal), y_t, loss)

from qruise.toolset import QOCProblem y0 = jnp.array([1.0, 0.0], dtype=jnp.complex128) # initial state |0> y_t = jnp.array([0.0, 1.0], dtype=jnp.complex128) # target state |1> def loss(x, y): """State infidelity: 1 - ||^2""" return jnp.real(1.0 - jnp.abs(jnp.vdot(x, y)) ** 2) opt_prob = QOCProblem(H, y0, params, (t0, tfinal), y_t, loss)

2. Inspecting frozen parameters¶

QOCProblem.inactive_params aggregates frozen keys from all signal-chain components that are used directly as Hamiltonian coefficients.

The expected frozen keys are those from TransferFunc: tf/a, tf/b, tf/t0, and tf/t1. Note that drive/t0 and drive/dt from PWCPulse are not surfaced here because PWCPulse is accessed via the pwc_signal closure, not directly as a Hamiltonian coefficient — see the Summary section for details.

In [9]:

print("Frozen parameter keys:")
for key in sorted(opt_prob.inactive_params):
    print(f"  {key}")

print("Frozen parameter keys:") for key in sorted(opt_prob.inactive_params): print(f" {key}")

Frozen parameter keys:
  tf/a
  tf/b
  tf/t0
  tf/t1

3. Optimisation without freeze_keys — why the search space matters¶

Passing all params to the minimiser unnecessarily inflates the search space and risks drifting hardware parameters over longer runs. The scipy parameter vector includes the array-valued filter coefficients tf/b (shape (1, 5)) and tf/a (shape (5,)), alongside the scalar timing parameters tf/t0 and tf/t1. The optimiser operates in a much larger search space than necessary — tf/b, tf/a, tf/t0, tf/t1 contribute extra dimensions that should be held fixed.

We show the number of parameters in the scipy vector with and without freeze_keys.

First, we record the initial values of the frozen filter parameters so we can check whether they are modified after optimisation.

In [10]:

from qruise.toolset import Optimiser, ScipyMinimiser, PWCSolver

# record initial filter param values for later comparison
b_initial = params["tf/b"].copy()
a_initial = params["tf/a"].copy()

from qruise.toolset import Optimiser, ScipyMinimiser, PWCSolver # record initial filter param values for later comparison b_initial = params["tf/b"].copy() a_initial = params["tf/a"].copy()

To see the difference concretely, we count how many scalar values are packed into the scipy minimiser's parameter vector in each case. Array-valued parameters such as tf/b (shape (1, 5)) and tf/a (shape (5,)) each contribute multiple elements, inflating the search space when freeze_keys is not used:

In [11]:

# total scalar elements across all params
n_all = sum(v.flatten().size if hasattr(v, "flatten") else 1 for v in params.values())

# active params only (frozen keys excluded)
active_keys = set(params.keys()) - opt_prob.inactive_params
active_params = {k: v for k, v in params.items() if k in active_keys}
n_active = sum(
    v.flatten().size if hasattr(v, "flatten") else 1 for v in active_params.values()
)

print(f"Parameters in scipy vector WITHOUT freeze_keys: {n_all}")
print(f"Parameters in scipy vector WITH    freeze_keys: {n_active}")
print(f"Frozen params excluded: {sorted(opt_prob.inactive_params)}")

# total scalar elements across all params n_all = sum(v.flatten().size if hasattr(v, "flatten") else 1 for v in params.values()) # active params only (frozen keys excluded) active_keys = set(params.keys()) - opt_prob.inactive_params active_params = {k: v for k, v in params.items() if k in active_keys} n_active = sum( v.flatten().size if hasattr(v, "flatten") else 1 for v in active_params.values() ) print(f"Parameters in scipy vector WITHOUT freeze_keys: {n_all}") print(f"Parameters in scipy vector WITH freeze_keys: {n_active}") print(f"Frozen params excluded: {sorted(opt_prob.inactive_params)}")

Parameters in scipy vector WITHOUT freeze_keys: 24
Parameters in scipy vector WITH    freeze_keys: 12
Frozen params excluded: ['tf/a', 'tf/b', 'tf/t0', 'tf/t1']

Now let's run the optimisation without freeze_keys to observe the behaviour:

In [12]:

minimiser_bad = ScipyMinimiser("L-BFGS-B", maxiter=5)
solver_bad = PWCSolver(n=grid_size, eq=opt_prob.sepwc())
opt_bad = Optimiser(minimiser_bad, solver_bad)
opt_bad.set_optimisation(opt_prob.loss)

# no freeze_keys — all params enter the scipy vector
result_bad, _ = opt_bad.optimise(*opt_prob.problem())

minimiser_bad = ScipyMinimiser("L-BFGS-B", maxiter=5) solver_bad = PWCSolver(n=grid_size, eq=opt_prob.sepwc()) opt_bad = Optimiser(minimiser_bad, solver_bad) opt_bad.set_optimisation(opt_prob.loss) # no freeze_keys — all params enter the scipy vector result_bad, _ = opt_bad.optimise(*opt_prob.problem())

4. Optimisation with freeze_keys — the correct approach¶

Passing freeze_keys=opt_prob.inactive_params tells the optimiser to exclude frozen parameters from the search space. They are held constant, and the minimiser only updates the active envelope samples (drive/env).

In [13]:

minimiser_good = ScipyMinimiser("L-BFGS-B", maxiter=20)
solver_good = PWCSolver(n=grid_size, eq=opt_prob.sepwc())
opt_good = Optimiser(minimiser_good, solver_good)
opt_good.set_optimisation(opt_prob.loss)

result_good, summary = opt_good.optimise(
    *opt_prob.problem(),
    freeze_keys=opt_prob.inactive_params,  # pass frozen keys here
)

print("tf/b initial:          ", b_initial)
print("tf/b after (unchanged):", result_good["tf/b"])
print("tf/a unchanged:", jnp.allclose(result_good["tf/a"], a_initial))
print(f"Optimiser converged: {summary.success}")

minimiser_good = ScipyMinimiser("L-BFGS-B", maxiter=20) solver_good = PWCSolver(n=grid_size, eq=opt_prob.sepwc()) opt_good = Optimiser(minimiser_good, solver_good) opt_good.set_optimisation(opt_prob.loss) result_good, summary = opt_good.optimise( *opt_prob.problem(), freeze_keys=opt_prob.inactive_params, # pass frozen keys here ) print("tf/b initial: ", b_initial) print("tf/b after (unchanged):", result_good["tf/b"]) print("tf/a unchanged:", jnp.allclose(result_good["tf/a"], a_initial)) print(f"Optimiser converged: {summary.success}")

tf/b initial:           [[0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
  4.85520219e+35]]
tf/b after (unchanged): [[0.00000000e+00 0.00000000e+00 0.00000000e+00 0.00000000e+00
  4.85520219e+35]]
tf/a unchanged: True
Optimiser converged: True

We can verify that the frozen parameters are bitwise identical to their initial values, while the active envelope has been updated by the optimiser:

In [14]:

assert jnp.allclose(result_good["tf/b"], params["tf/b"]), "tf/b drifted!"
assert jnp.allclose(result_good["tf/a"], params["tf/a"]), "tf/a drifted!"
assert result_good["tf/t0"] == params["tf/t0"], "tf/t0 drifted!"
assert result_good["tf/t1"] == params["tf/t1"], "tf/t1 drifted!"
assert "drive/env" in result_good, "drive/env missing from result!"

print("All assertions passed — frozen params are unchanged, active params optimised.")

assert jnp.allclose(result_good["tf/b"], params["tf/b"]), "tf/b drifted!" assert jnp.allclose(result_good["tf/a"], params["tf/a"]), "tf/a drifted!" assert result_good["tf/t0"] == params["tf/t0"], "tf/t0 drifted!" assert result_good["tf/t1"] == params["tf/t1"], "tf/t1 drifted!" assert "drive/env" in result_good, "drive/env missing from result!" print("All assertions passed — frozen params are unchanged, active params optimised.")

All assertions passed — frozen params are unchanged, active params optimised.

5. Summary and best practices¶

The freeze_keys workflow for signal-chain optimisation:

opt_prob = QOCProblem(H, y0, params, (t0, tfinal), y_t, loss)

result, summary = opt.optimise(
    *opt_prob.problem(),
    freeze_keys=opt_prob.inactive_params,  # automatically excludes frozen params
)

Key points:

• Declare frozen parameters with False and active ones with True when constructing signal-chain components (PWCPulse, TransferFunc, etc.).
• Pass signal-chain components directly to the Hamiltonian constructor via its Ht parameter (not wrapped in a plain function) so that Hamiltonian.inactive_params can discover frozen keys.
• Use opt_prob.inactive_params as the freeze_keys argument — it aggregates all frozen keys from the entire signal chain automatically.

Known limitation: If a signal-chain component is wrapped in a plain Python function before being passed to the Hamiltonian, inactive_params will return an empty set for that component. In that case, pass freeze_keys manually:

# When using a plain-function wrapper, specify frozen keys explicitly:
result, _ = opt.optimise(*opt_prob.problem(), freeze_keys={"drive/t0", "drive/dt"})