Merge pull request #241 from harmoniqs/refactor/template-redesign

Refactor/template redesign

Author: aarontrowbridge

CONTEXT.md

@@ -0,0 +1,316 @@
+# QuantumCollocation.jl Context
+
+> AI-friendly context for maintaining consistency. Update this when making significant changes.
+
+## Package Purpose
+
+QuantumCollocation.jl provides **problem templates** for quantum optimal control using direct collocation. It builds on:
+- **DirectTrajOpt.jl** - Core optimization infrastructure (objectives, constraints, integrators, NLP solver)
+- **PiccoloQuantumObjects.jl** - Quantum systems, trajectories, pulses, and utilities
+- **NamedTrajectories.jl** - Named variable trajectory data structure
+
+The main abstractions are:
+1. **QuantumControlProblem** - Wrapper combining quantum trajectory info with optimization problem
+2. **Problem Templates** - Constructors that build complete optimization problems
+
+## Canonical Workflow
+
+### Basic Gate Synthesis
+
+```julia
+using QuantumCollocation
+using PiccoloQuantumObjects
+
+# 1. Define quantum system (no T_max - duration comes from trajectory)
+H_drift = GATES[:Z]
+H_drives = [GATES[:X], GATES[:Y]]
+drive_bounds = [1.0, 1.0]  # symmetric bounds
+sys = QuantumSystem(H_drift, H_drives, drive_bounds)
+
+# 2. Create quantum trajectory (defines the control problem)
+U_goal = GATES[:H]  # target: Hadamard gate
+T = 10.0  # duration
+qtraj = UnitaryTrajectory(sys, U_goal, T)  # creates zero pulse internally
+
+# 3. Build optimization problem
+N = 51  # number of timesteps
+qcp = SmoothPulseProblem(qtraj, N; Q=100.0, R=1e-2)
+
+# 4. Solve (use 'options' keyword for IpoptOptions)
+solve!(qcp; options=IpoptOptions(max_iter=200))
+# OR pass options during problem construction:
+# qcp = SmoothPulseProblem(qtraj, N; Q=100.0, R=1e-2, ipopt_options=opts)
+# solve!(qcp)  # uses options from construction
+
+# 5. Extract results
+traj = get_trajectory(qcp)
+U_final = iso_vec_to_operator(traj[end][:Ũ⃗])
+fid = unitary_fidelity(U_final, U_goal)
+```
+
+### State Transfer (Ket)
+
+```julia
+sys = QuantumSystem(H_drift, H_drives, drive_bounds)
+
+ψ_init = ComplexF64[1.0, 0.0]  # |0⟩
+ψ_goal = ComplexF64[0.0, 1.0]  # |1⟩
+T = 10.0
+qtraj = KetTrajectory(sys, ψ_init, ψ_goal, T)
+
+qcp = SmoothPulseProblem(qtraj, N; Q=50.0, R=1e-3)
+solve!(qcp; max_iter=100)
+```
+
+### Ensemble Control (Multiple Initial States)
+
+```julia
+sys = QuantumSystem(H_drift, H_drives, drive_bounds)
+
+# X gate: |0⟩→|1⟩ and |1⟩→|0⟩
+initials = [ComplexF64[1,0], ComplexF64[0,1]]
+goals = [ComplexF64[0,1], ComplexF64[1,0]]
+T = 10.0
+qtraj = MultiKetTrajectory(sys, initials, goals, T)
+
+qcp = SmoothPulseProblem(qtraj, N; Q=100.0, R=1e-2)
+solve!(qcp; max_iter=200)
+```
+
+### With Explicit Pulse (Spline-Based)
+
+```julia
+sys = QuantumSystem(H_drift, H_drives, drive_bounds)
+
+# Create explicit pulse with initial guess
+T = 10.0
+times = collect(range(0.0, T, length=N))
+controls = 0.1 * randn(sys.n_drives, N)
+pulse = LinearSplinePulse(controls, times)
+
+qtraj = UnitaryTrajectory(sys, pulse, U_goal)
+
+# Use SplinePulseProblem for spline pulses
+qcp = SplinePulseProblem(qtraj, N; Q=100.0, R=1e-2, du_bound=10.0)
+solve!(qcp; max_iter=200)
+```
+
+### Minimum Time Optimization
+
+```julia
+# First solve with fixed time
+qcp = SmoothPulseProblem(qtraj, N; Q=100.0, Δt_bounds=(0.01, 0.5))
+solve!(qcp; max_iter=100)
+
+# Convert to minimum-time problem
+qcp_mintime = MinimumTimeProblem(qcp; final_fidelity=0.99, D=50.0)
+solve!(qcp_mintime; max_iter=100)
+```
+
+### Bootstrapping Across System Sizes
+
+**Best Practice:** When optimizing similar systems of different sizes, **pass the pulse object** between solves to preserve both controls and derivatives:
+
+```julia
+# Optimize for N=3 atoms
+sys3 = build_system(N=3)
+qtraj3 = UnitaryTrajectory(sys3, pulse_init, U_goal)
+qcp3 = SplinePulseProblem(qtraj3, N_timesteps; Q=100.0, R_du=1e-4)
+solve!(qcp3; max_iter=200)
+
+# Bootstrap to N=4 using optimized pulse directly
+optimized_pulse = qcp3.qtraj.pulse  # Contains both u and du
+sys4 = build_system(N=4)
+qtraj4 = UnitaryTrajectory(sys4, optimized_pulse, U_goal)  # Reuse pulse
+qcp4 = SplinePulseProblem(qtraj4, N_timesteps; Q=100.0, R_du=1e-4)
+solve!(qcp4; max_iter=200)
+
+# Save pulse for next bootstrap
+using JLD2
+save("optimized_N4.jld2", "pulse", qcp4.qtraj.pulse)
+```
+
+**Why this works:**
+- `Pulse` objects are system-independent (just time → control mappings)
+- Preserves full optimization state: `u` (controls) + `du` (derivatives/tangents)
+- The new system's dynamics are enforced during `solve!` via rollout constraints
+
+**Don't:**
+```julia
+# ❌ Bad: Extracting controls from NamedTrajectory loses du information
+u_old = get_trajectory(qcp3)[:u]
+du_zeros = zeros(size(u_old))  # Throws away optimized tangents!
+pulse_new = CubicSplinePulse(u_old, du_zeros, times)
+
+# ❌ Bad: Trying to transfer NamedTrajectory (has system-specific states)
+qtraj4 = bootstrap_from_trajectory(traj3, sys4)  # Incompatible states
+```
+
+**Do:**
+```julia
+# ✓ Good: Pass pulse object directly
+pulse_optimized = qcp3.qtraj.pulse  # Preserves u and du
+qtraj4 = UnitaryTrajectory(sys4, pulse_optimized, U_goal)
+```
+
+### Robust Control (Parameter Sampling)
+
+```julia
+# Create base problem
+qtraj = UnitaryTrajectory(sys_nominal, pulse, U_goal)
+qcp = SmoothPulseProblem(qtraj, N; Q=100.0)
+
+# Add robustness over parameter variations
+sampling_prob = SamplingProblem(qcp, [sys_nominal, sys_perturbed]; Q=100.0)
+solve!(sampling_prob; max_iter=200)
+```
+
+## Key Abstractions
+
+### Quantum Trajectories (defined in PiccoloQuantumObjects.jl)
+
+Parametric type hierarchy (see PiccoloQuantumObjects.jl/CONTEXT.md for details):
+```
+AbstractQuantumTrajectory{P<:AbstractPulse}
+├── UnitaryTrajectory{P}      # Full unitary evolution
+├── KetTrajectory{P}          # Single state evolution  
+├── MultiKetTrajectory{P}  # Multiple states, shared controls
+├── DensityTrajectory{P}      # Density matrix evolution (WIP)
+└── SamplingTrajectory{P,Q}   # Robustness over system parameters
+```
+
+QuantumCollocation.jl **uses** these types and provides problem templates that build on them.
+
+### Problem Templates
+
+| Template | Use Case | Pulse Type |
+|----------|----------|------------|
+| `SmoothPulseProblem` | Standard optimization with smooth pulses | `ZeroOrderPulse` (piecewise constant) |
+| `SplinePulseProblem` | Spline-based pulse optimization | `LinearSplinePulse`, `CubicSplinePulse` |
+| `MinimumTimeProblem` | Time-optimal control | Any (converts existing problem) |
+| `SamplingProblem` | Robust control over parameters | Any (wraps existing problem) |
+
+**SmoothPulseProblem** (for ZeroOrderPulse):
+- Adds derivative variables `:du`, `:ddu` for smoothness
+- Creates `DerivativeIntegrator` constraints enforcing `u[k+1] - u[k] = Δt * du[k]`
+- Applies quadratic regularization on `u`, `du`, `ddu`
+
+**SplinePulseProblem** (for spline pulses):
+- For `LinearSplinePulse`: `:du` represents slopes (added automatically)
+- For `CubicSplinePulse`: `:du` represents Hermite tangents (built into pulse)
+- Uses `DerivativeIntegrator` with spline semantics
+
+**Adding constraints to SplinePulseProblem:** When working with `SplinePulseProblem`, especially with dynamical timesteps (`Δt_bounds`), add constraints to the existing problem rather than recreating it:
+
+```julia
+# Create problem with dynamical timesteps
+qcp = SplinePulseProblem(qtraj, N; 
+    Q=100.0, R_u=1e-2, R_du=1e-4,
+    Δt_bounds=(0.01, 0.5))
+
+# Get trajectory and build additional constraints
+traj = get_trajectory(qcp)
+constraint = MyCustomConstraint(traj, ...)
+
+# Add to existing problem (preserves all problem structure)
+push!(qcp.prob.constraints, constraint)
+
+# Solve the quantum collocation problem
+solve!(qcp; options=ipopt_options)
+
+# Access results
+final_fidelity = fidelity(qcp.qtraj)       # Use qcp.qtraj for quantum operations
+u_optimized = qcp.prob.trajectory[:u]       # Use qcp.prob.trajectory for variables
+optimized_pulse = qcp.qtraj.pulse           # For saving or bootstrapping
+```
+
+**Important:** Do NOT recreate `DirectTrajOptProblem` when modifying a `SplinePulseProblem`, as this breaks the dynamical timestep mechanism and integration with the quantum trajectory.
+
+### Quantum Integrators (`quantum_integrators.jl`)
+
+`BilinearIntegrator` dispatches on trajectory type:
+```julia
+BilinearIntegrator(qtraj::UnitaryTrajectory, N)      # → single integrator
+BilinearIntegrator(qtraj::KetTrajectory, N)          # → single integrator
+BilinearIntegrator(qtraj::MultiKetTrajectory, N)  # → Vector of integrators
+BilinearIntegrator(qtraj::SamplingTrajectory, N)     # → Vector of integrators
+```
+
+**Signature pattern:** `(qtraj, N::Int)` - creates NamedTrajectory internally
+
+### Quantum Constraints (`quantum_constraints.jl`)
+
+Fidelity constraints for minimum-time problems:
+- `FinalUnitaryFidelityConstraint(qtraj, traj, fidelity)`
+- `FinalKetFidelityConstraint(qtraj, traj, fidelity)`
+- `FinalCoherentKetFidelityConstraint(qtraj, traj, fidelity)` - For ensembles, preserves phase coherence
+
+## Component Naming Conventions
+
+| Component | Symbol | Description |
+|-----------|--------|-------------|
+| Unitary state | `:Ũ⃗` | Isomorphism-vectorized unitary |
+| Ket state | `:ψ̃` | Isomorphism-vectorized ket |
+| Ensemble kets | `:ψ̃1`, `:ψ̃2`, ... | Multiple states |
+| Sampling states | `:Ũ⃗1`, `:Ũ⃗2`, ... | States for each system sample |
+| Controls | `:u` | Pulse amplitudes (canonical name) |
+| Control derivative | `:du` | First derivative / slope |
+| Control 2nd derivative | `:ddu` | Second derivative / acceleration |
+| Timestep | `:Δt` | Duration of each timestep |
+| Time | `:t` | Accumulated time (always present) |
+
+## Testing Conventions
+
+Tests use `@testitem` blocks in source files:
+```julia
+@testitem "descriptive name" begin
+    using QuantumCollocation
+    using PiccoloQuantumObjects
+    using DirectTrajOpt
+    using LinearAlgebra
+    # ... test code
+end
+```
+
+Run tests with `TestItemRunner.@run_package_tests`.
+
+## Recent Changes (Update This!)
+
+### January 2026
+- Removed `time_dependent=true` from test QuantumSystem constructions (`:t` is now always in trajectories)
+- Removed `adapt_trajectory`/`unadapt_trajectory` usage (control scaling removed)
+- Updated `BilinearIntegrator` signatures from `(qtraj, traj)` to `(qtraj, N)`
+- Added `FinalCoherentKetFidelityConstraint` for ensemble minimum-time problems
+
+### Removed Patterns (Don't Reintroduce!)
+- ❌ `adapt_trajectory` / `unadapt_trajectory` - Control scaling was removed
+- ❌ `@test sys.time_dependent` in tests - Not needed since `:t` is always present
+- ❌ `QuantumSystem(...; time_dependent=true)` in tests - Redundant
+- ❌ Control name `:a` - Canonical control name is `:u`
+
+## File Structure
+
+```
+src/
+├── QuantumCollocation.jl      # Main module, reexports
+├── piccolo_options.jl         # PiccoloOptions configuration
+├── quantum_control_problem.jl # QuantumControlProblem wrapper
+├── quantum_integrators.jl     # BilinearIntegrator dispatch
+├── quantum_constraints.jl     # Fidelity constraints
+├── quantum_objectives.jl      # Fidelity objectives
+└── problem_templates/
+    ├── _problem_templates.jl  # Submodule definition
+    ├── smooth_pulse_problem.jl  # ZeroOrderPulse optimization
+    ├── spline_pulse_problem.jl  # Spline pulse optimization
+    ├── minimum_time_problem.jl  # Time-optimal control
+    └── sampling_problem.jl      # Robust control
+```
+
+## Common Gotchas
+
+1. **Pulse type determines problem template**: Use `SmoothPulseProblem` for `ZeroOrderPulse`, `SplinePulseProblem` for spline pulses
+2. **N is timesteps, not knot points**: For `N=51`, you get 50 intervals
+3. **Trajectories always have `:t`**: Time is accumulated automatically, no need for `time_dependent` flag
+4. **MultiKetTrajectory vs SamplingTrajectory**: Ensemble = same system, different initial states; Sampling = different systems, same initial state
+5. **Fidelity vs Infidelity**: Objectives minimize infidelity, constraints bound infidelity (e.g., `1 - fidelity ≤ 1 - 0.99`)
+6. **Bootstrapping between system sizes**: Pass the **pulse object** (`qcp.qtraj.pulse`), not the `NamedTrajectory` or extracted controls. This preserves both `u` and `du` optimization state.

Project.toml

@@ -1,6 +1,6 @@
 name = "QuantumCollocation"
 uuid = "0dc23a59-5ffb-49af-b6bd-932a8ae77adf"
-version = "0.9.0"
+version = "0.10.0"
 authors = ["Aaron Trowbridge <aaron.j.trowbridge@gmail.com> and contributors"]
 
 [deps]
@@ -19,12 +19,12 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 TrajectoryIndexingUtils = "6dad8b7f-dd9a-4c28-9b70-85b9a079bfc8"
 
 [compat]
-DirectTrajOpt = "0.5"
+DirectTrajOpt = "0.8"
 Distributions = "0.25"
 ExponentialAction = "0.2"
 LinearAlgebra = "1.10, 1.11, 1.12"
-NamedTrajectories = "0.6"
-PiccoloQuantumObjects = "0.7"
+NamedTrajectories = "0.8"
+PiccoloQuantumObjects = "0.10"
 Random = "1.10, 1.11, 1.12"
 Reexport = "1.2"
 SparseArrays = "1.10, 1.11, 1.12"

docs/dev/ipopt_callbacks.jl

@@ -42,7 +42,7 @@ T_max = 1.0
 u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
 T = 50
 Δt = 0.2
-sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
+sys = QuantumSystem(0.1 * GATES[:Z], [GATES[:X], GATES[:Y]], u_bounds)
 ψ_init =  Vector{ComplexF64}([1.0, 0.0])
 ψ_target =  Vector{ComplexF64}([0.0, 1.0])
 
@@ -70,7 +70,7 @@ end
 # Using a callback to get the best trajectory from all the optimization iterations
 T_max = 1.0
 u_bounds = [(-1.0, 1.0), (-1.0, 1.0)]
-sys2 = QuantumSystem(0.15 * GATES[:Z], [GATES[:X], GATES[:Y]], T_max, u_bounds)
+sys2 = QuantumSystem(0.15 * GATES[:Z], [GATES[:X], GATES[:Y]], u_bounds)
 ψ_init2 =  Vector{ComplexF64}([0.0, 1.0])
 ψ_target2 =  Vector{ComplexF64}([1.0, 0.0])