Coverage for jaxquantum/circuits/simulate.py: 95%
255 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Circuit simulation methods."""
3from math import prod
5import jax
6import jax.numpy as jnp
7from flax import struct
8from jax import config, lax
10from jaxquantum.circuits.channels import apply_kraus_map
11from jaxquantum.circuits.circuits import Circuit, Layer
12from jaxquantum.circuits.constants import SimulateMode
13from jaxquantum.core.measurements import overlap
14from jaxquantum.core.qarray import DenseImpl, Qarray, Qtypes, ket2dm
15from jaxquantum.core.solvers import SolverOptions, mesolve, sesolve, solve
17config.update("jax_enable_x64", True)
20@struct.dataclass
21class Results:
22 results: list[Qarray] = struct.field(pytree_node=False)
24 @classmethod
25 def create(cls, results: list[Qarray]):
26 return Results(results=results)
28 def __getitem__(self, j: int):
29 return self.results[j]
31 def __str__(self):
32 return self.__repr__()
34 def __repr__(self):
35 return str(self.results)
37 def append(self, result: Qarray):
38 self.results.append(result)
40 def __len__(self):
41 return len(self.results)
44def _apply_matrix_to_axes(data, matrix, target_axes, system_shape):
45 """Apply a matrix to selected tensor axes without forming a full operator."""
46 n_system_axes = len(system_shape)
47 batch_shape = data.shape[:-n_system_axes]
48 target_axes = tuple(target_axes)
50 if target_axes == tuple(range(target_axes[0], target_axes[-1] + 1)):
51 start, stop = target_axes[0], target_axes[-1] + 1
52 left = prod(system_shape[:start])
53 target = prod(system_shape[start:stop])
54 right = prod(system_shape[stop:])
55 data = data.reshape(batch_shape + (left, target, right))
56 data = jnp.einsum("...ij,...ljr->...lir", matrix, data)
57 return data.reshape(data.shape[:-3] + tuple(system_shape))
59 other_axes = tuple(i for i in range(n_system_axes) if i not in target_axes)
60 order = other_axes + target_axes
61 n_batch_axes = len(batch_shape)
63 data = jnp.transpose(
64 data,
65 tuple(range(n_batch_axes))
66 + tuple(n_batch_axes + axis for axis in order),
67 )
68 other_shape = tuple(system_shape[axis] for axis in other_axes)
69 target_shape = tuple(system_shape[axis] for axis in target_axes)
70 data = data.reshape(batch_shape + (prod(other_shape), prod(target_shape)))
71 data = jnp.einsum("...ij,...kj->...ki", matrix, data)
73 out_batch_shape = data.shape[:-2]
74 data = data.reshape(out_batch_shape + other_shape + target_shape)
75 return jnp.transpose(
76 data,
77 tuple(range(len(out_batch_shape)))
78 + tuple(
79 len(out_batch_shape) + order.index(axis)
80 for axis in range(n_system_axes)
81 ),
82 )
85def _apply_local_unitary(state: Qarray, operation) -> Qarray:
86 dims = tuple(operation.register.dims)
87 unitary = operation.gate.U.to_dense().data
88 n_modes = len(dims)
89 state_data = state.to_dense().data
91 if state.qtype == Qtypes.ket:
92 data = state_data.reshape(state_data.shape[:-1] + dims)
93 data = _apply_matrix_to_axes(data, unitary, operation.indices, dims)
94 data = data.reshape(data.shape[:-n_modes] + (prod(dims),))
95 else:
96 system_shape = dims + dims
97 data = state_data.reshape(state_data.shape[:-2] + system_shape)
98 data = _apply_matrix_to_axes(data, unitary, operation.indices, system_shape)
99 bra_axes = tuple(n_modes + index for index in operation.indices)
100 data = _apply_matrix_to_axes(data, jnp.conj(unitary), bra_axes, system_shape)
101 data = data.reshape(data.shape[:-2 * n_modes] + (prod(dims), prod(dims)))
103 return Qarray._from_impl(DenseImpl._make(data), state._qdims)
106def _apply_local_kraus(state: Qarray, operation) -> Qarray:
107 """Apply a local Kraus map without promoting it to the full register."""
108 state = ket2dm(state)
109 dims = tuple(operation.register.dims)
110 n_modes = len(dims)
111 system_shape = dims + dims
112 data = state.to_dense().data.reshape(state.data.shape[:-2] + system_shape)
113 direct_apply = operation.gate.channel_apply
114 kraus = None if direct_apply is not None else operation.gate.KM.to_dense().data
115 if kraus is not None and kraus.shape[0] == 0:
116 return state
118 target_axes = tuple(operation.indices) + tuple(
119 n_modes + index for index in operation.indices
120 )
121 other_axes = tuple(
122 index for index in range(2 * n_modes) if index not in target_axes
123 )
124 order = other_axes + target_axes
125 n_batch_axes = data.ndim - 2 * n_modes
126 data = jnp.transpose(
127 data,
128 tuple(range(n_batch_axes))
129 + tuple(n_batch_axes + index for index in order),
130 )
131 other_shape = tuple(system_shape[index] for index in other_axes)
132 target_shape = tuple(dims[index] for index in operation.indices)
133 target_size = prod(target_shape)
134 data = data.reshape(
135 data.shape[:n_batch_axes]
136 + (prod(other_shape), target_size, target_size)
137 )
138 if direct_apply is not None:
139 data = direct_apply(data, operation.gate.params)
140 else:
141 data = apply_kraus_map(kraus[..., None, :, :], data)
142 out_batch_shape = data.shape[:-3]
143 data = data.reshape(out_batch_shape + other_shape + target_shape + target_shape)
144 data = jnp.transpose(
145 data,
146 tuple(range(len(out_batch_shape)))
147 + tuple(
148 len(out_batch_shape) + order.index(index)
149 for index in range(2 * n_modes)
150 ),
151 )
152 data = data.reshape(out_batch_shape + (prod(dims), prod(dims)))
153 return Qarray._from_impl(DenseImpl._make(data), state._qdims)
156def _local_batch_shape(layer, state_data, time, density_matrix):
157 state_rank = 2 if density_matrix else 1
158 batch_shape = state_data.shape[:-state_rank]
159 for operation in layer.operations:
160 if operation.gate.Ht is not None:
161 shape = operation.gate.Ht(time).data.shape[:-2]
162 batch_shape = jnp.broadcast_shapes(batch_shape, shape)
163 if len(operation.gate.c_ops):
164 shape = operation.gate.c_ops.data.shape[1:-2]
165 batch_shape = jnp.broadcast_shapes(batch_shape, shape)
166 return batch_shape
169def _solve_local_hamiltonian(
170 layer,
171 state,
172 times,
173 saveat_times,
174 solver_options,
175):
176 dims = tuple(layer.operations[0].register.dims)
177 n_modes = len(dims)
178 has_c_ops = any(len(operation.gate.c_ops) for operation in layer.operations)
179 density_matrix = state.is_dm() or has_c_ops
180 state = state.to_dm().to_dense() if density_matrix else state.to_ket().to_dense()
181 qdims = state.qdims
182 data = state.data
183 batch_shape = _local_batch_shape(layer, data, times[0], density_matrix)
184 state_rank = 2 if density_matrix else 1
185 data = jnp.broadcast_to(data, batch_shape + data.shape[-state_rank:])
187 def rhs(time, value, _):
188 system_shape = dims + dims if density_matrix else dims
189 tensor = value.reshape(value.shape[:-state_rank] + system_shape)
190 derivative = jnp.zeros_like(tensor)
192 for operation in layer.operations:
193 indices = tuple(operation.indices)
194 hamiltonian = operation.gate.Ht
195 if hamiltonian is not None:
196 matrix = hamiltonian(time).to_dense().data
197 left = _apply_matrix_to_axes(tensor, matrix, indices, system_shape)
198 if density_matrix:
199 bra_indices = tuple(n_modes + index for index in indices)
200 right = _apply_matrix_to_axes(
201 tensor,
202 jnp.swapaxes(matrix, -1, -2),
203 bra_indices,
204 system_shape,
205 )
206 derivative += -1.0j * (left - right)
207 else:
208 derivative += -1.0j * left
210 if not density_matrix or len(operation.gate.c_ops) == 0:
211 continue
212 collapse_ops = operation.gate.c_ops.to_dense().data
213 bra_indices = tuple(n_modes + index for index in indices)
215 def dissipator(
216 matrix,
217 indices=indices,
218 bra_indices=bra_indices,
219 ):
220 left = _apply_matrix_to_axes(
221 tensor,
222 matrix,
223 indices,
224 system_shape,
225 )
226 sandwich = _apply_matrix_to_axes(
227 left,
228 jnp.conj(matrix),
229 bra_indices,
230 system_shape,
231 )
232 product = jnp.swapaxes(jnp.conj(matrix), -1, -2) @ matrix
233 anti_left = _apply_matrix_to_axes(
234 tensor,
235 product,
236 indices,
237 system_shape,
238 )
239 anti_right = _apply_matrix_to_axes(
240 tensor,
241 jnp.swapaxes(product, -1, -2),
242 bra_indices,
243 system_shape,
244 )
245 return sandwich - 0.5 * (anti_left + anti_right)
247 def add_dissipator(index, total, collapse_ops=collapse_ops):
248 return total + dissipator(collapse_ops[index])
250 derivative += lax.fori_loop(
251 1,
252 collapse_ops.shape[0],
253 add_dissipator,
254 dissipator(collapse_ops[0]),
255 )
257 return derivative.reshape(value.shape)
259 ys = solve(
260 rhs,
261 data,
262 times,
263 saveat_times,
264 None,
265 solver_options=solver_options,
266 ).ys
267 return Qarray._from_impl(DenseImpl._make(ys), qdims)
270def _single_state_batch(state: Qarray) -> Qarray:
271 impl = type(state._impl).from_data(state.data.reshape(1, *state.data.shape))
272 return Qarray._from_impl(impl, state._qdims)
275def _evolve_circuit(circuit, state, mode, start_time=0.0, **kwargs):
276 for layer in circuit.layers:
277 output = _simulate_layer(
278 layer,
279 state,
280 mode=mode,
281 start_time=start_time,
282 **kwargs,
283 )
284 state = output["result"][-1]
285 start_time = output["start_time"]
286 return state, start_time
289def simulate(
290 circuit: Circuit,
291 initial_state: Qarray,
292 mode: SimulateMode = SimulateMode.DEFAULT,
293 save_states: bool = True,
294 **kwargs,
295) -> Results:
296 """Simulate a circuit and optionally retain each layer's states.
298 Args:
299 circuit: Circuit to simulate.
300 initial_state: Initial ket or density matrix.
301 mode: Simulation mode, or each layer's default mode.
302 save_states: Whether to retain intermediate layer states.
304 Returns:
305 Saved states, or only the final state when ``save_states=False``.
306 """
308 results = Results.create(
309 [_single_state_batch(initial_state)] if save_states else []
310 )
311 state = initial_state
312 start_time = 0
313 if not save_states:
314 kwargs.setdefault("saveat_tlist", jnp.array([]))
316 for layer in circuit.layers:
317 result_dict = _simulate_layer(
318 layer,
319 state,
320 mode=mode,
321 start_time=start_time,
322 **kwargs,
323 )
324 result = result_dict["result"]
325 start_time = result_dict["start_time"]
326 state = result[-1]
327 if save_states:
328 results.append(result)
330 if not save_states:
331 results.append(_single_state_batch(state))
332 return results
335def simulate_final(
336 circuit: Circuit,
337 initial_state: Qarray,
338 mode: SimulateMode = SimulateMode.DEFAULT,
339 **kwargs,
340) -> Qarray:
341 """Return only the final circuit state."""
342 kwargs.setdefault("saveat_tlist", jnp.array([]))
343 return _evolve_circuit(circuit, initial_state, mode, **kwargs)[0]
346def simulate_repeated(
347 circuit: Circuit,
348 initial_state: Qarray,
349 repetitions: int,
350 mode: SimulateMode = SimulateMode.DEFAULT,
351 **kwargs,
352) -> Qarray:
353 """Apply one circuit repeatedly with a compiled loop."""
354 if repetitions < 0:
355 raise ValueError("repetitions must be non-negative")
356 if repetitions == 0:
357 return initial_state
359 kwargs.setdefault("saveat_tlist", jnp.array([]))
360 state, start_time = _evolve_circuit(
361 circuit,
362 initial_state,
363 mode,
364 **kwargs,
365 )
367 def repeat(_, carry):
368 return _evolve_circuit(circuit, carry[0], mode, carry[1], **kwargs)
370 return lax.fori_loop(
371 1,
372 repetitions,
373 repeat,
374 (state, start_time),
375 )[0]
378def _expectations(state, observables):
379 return jnp.stack([overlap(state, observable) for observable in observables], -1)
382def simulate_expectations(
383 circuit: Circuit,
384 initial_state: Qarray,
385 observables: list[Qarray],
386 mode: SimulateMode = SimulateMode.DEFAULT,
387 include_initial: bool = True,
388 **kwargs,
389):
390 """Return the final state and per-layer expectation values."""
391 if not observables:
392 raise ValueError("observables must not be empty")
394 kwargs.setdefault("saveat_tlist", jnp.array([]))
395 state = initial_state
396 start_time = 0.0
397 values = [_expectations(state, observables)] if include_initial else []
398 for layer in circuit.layers:
399 output = _simulate_layer(layer, state, mode, start_time, **kwargs)
400 state = output["result"][-1]
401 start_time = output["start_time"]
402 values.append(_expectations(state, observables))
403 if not values:
404 return state, _expectations(state, observables)[None][:0]
405 return state, jnp.stack(values)
408def simulate_repeated_expectations(
409 circuit: Circuit,
410 initial_state: Qarray,
411 repetitions: int,
412 observables: list[Qarray],
413 mode: SimulateMode = SimulateMode.DEFAULT,
414 include_initial: bool = True,
415 **kwargs,
416):
417 """Return the final state and per-repetition expectation values."""
418 if repetitions < 0:
419 raise ValueError("repetitions must be non-negative")
420 if not observables:
421 raise ValueError("observables must not be empty")
422 if repetitions == 0:
423 values = _expectations(initial_state, observables)[None]
424 return initial_state, values if include_initial else values[:0]
426 kwargs.setdefault("saveat_tlist", jnp.array([]))
427 state, start_time = _evolve_circuit(
428 circuit,
429 initial_state,
430 mode,
431 **kwargs,
432 )
433 first = _expectations(state, observables)
435 def repeat(carry, _):
436 state, start_time = _evolve_circuit(
437 circuit,
438 carry[0],
439 mode,
440 carry[1],
441 **kwargs,
442 )
443 return (state, start_time), _expectations(state, observables)
445 (state, _), rest = lax.scan(
446 repeat,
447 (state, start_time),
448 None,
449 length=repetitions - 1,
450 )
451 values = jnp.concatenate((first[None], rest), axis=0)
452 if include_initial:
453 values = jnp.concatenate(
454 (_expectations(initial_state, observables)[None], values),
455 axis=0,
456 )
457 return state, values
460def _simulate_layer(
461 layer: Layer,
462 initial_state: Qarray,
463 mode: SimulateMode = SimulateMode.UNITARY,
464 start_time: float = 0,
465 **kwargs,
466) -> dict:
467 """Simulate one circuit layer and return its states and ending time."""
469 state = initial_state
471 if mode == SimulateMode.DEFAULT:
472 mode = layer._default_simulate_mode
474 if mode == SimulateMode.UNITARY:
475 for operation in layer.operations:
476 state = _apply_local_unitary(state, operation)
477 result = _single_state_batch(state)
479 elif mode == SimulateMode.HAMILTONIAN:
481 solver_options = kwargs.pop(
482 "solver_options",
483 SolverOptions(progress_meter=None),
484 )
485 ts = layer.gen_ts()
486 ts = ts + start_time
487 saveat_times = kwargs.pop("saveat_tlist", ts)
488 local_operators = kwargs.pop("local_operators", None)
489 has_c_ops = any(len(operation.gate.c_ops) for operation in layer.operations)
490 if local_operators is None:
491 local_operators = (
492 state.is_dm() or has_c_ops or jax.default_backend() == "cpu"
493 )
495 if local_operators:
496 intermediate_states = _solve_local_hamiltonian(
497 layer,
498 state,
499 ts,
500 saveat_times,
501 solver_options,
502 )
503 else:
504 hamiltonian = layer.gen_Ht()
505 collapse_ops = layer.gen_c_ops()
506 if state.is_dm() or len(collapse_ops):
507 intermediate_states = mesolve(
508 hamiltonian,
509 state,
510 ts,
511 saveat_tlist=saveat_times,
512 c_ops=collapse_ops,
513 solver_options=solver_options,
514 )
515 else:
516 intermediate_states = sesolve(
517 hamiltonian,
518 state,
519 ts,
520 saveat_tlist=saveat_times,
521 solver_options=solver_options,
522 )
524 result = intermediate_states
525 state = intermediate_states[-1]
526 start_time = ts[-1]
528 elif mode == SimulateMode.KRAUS:
529 for operation in layer.operations:
530 state = _apply_local_kraus(state, operation)
531 result = _single_state_batch(state)
533 else:
534 raise ValueError(f"Unsupported simulation mode: {mode}")
536 return {"result": result, "start_time": start_time}