fgrape (feedback GRAPE)

GRadient Ascent Pulse Engineering (GRAPE) with feedback.

Note

The below code is the default predefined RNN used in feedback GRAPE. For another example of how to use a custom RNN, see: E: State stabilization with SNAP gates and displacement gates

class RNN(nn.Module):
    hidden_size: int  # number of features in the hidden state
    output_size: int  # number of features in the output (inferred from the number of parameters) just provide those attributes to the class

    @nn.compact
    def __call__(self, measurement, hidden_state):

        gru_cell = nn.GRUCell(features=self.hidden_size)

        if measurement.ndim == 1:
            measurement = measurement.reshape(1, -1)

        ###############
        ### Free to change whatever you want below as long as hidden layers have size self.hidden_size
        ### and output layer has size self.output_size
        ###############

        new_hidden_state, _ = gru_cell(hidden_state, measurement)
        # this returns the povm_params after linear regression through the hidden state which contains
        # the information of the previous time steps and this is optimized to output best povm_params
        # new_hidden_state = nn.Dense(features=self.hidden_size)(new_hidden_state)
        output = nn.Dense(
            features=self.output_size,
            kernel_init=nn.initializers.glorot_uniform(),
            bias_init=nn.initializers.constant(jnp.pi),
        )(new_hidden_state)
        output = nn.relu(output)

        ###############
        ### Do not change the return statement
        ###############

        return output[0], new_hidden_state

class feedback_grape.fgrape.Decay(c_ops: list[Array])

Bases: NamedTuple

Decay class to store the parameters for each decay.

c_ops: list[Array]

Collapse operators for the decay.

class feedback_grape.fgrape.FgResult(optimized_trainable_parameters: Array, iterations: int, final_state: Array, returned_params: List[Array], final_purity: Array | None, final_fidelity: Array | None, fidelity_each_timestep: list[Array], purity_each_timestep: list[Array])

Bases: NamedTuple

result class to store the results of the optimization process.

fidelity_each_timestep: list[Array]

Fidelity of the optimized control along each timestep and batch.

final_fidelity: Array | None

Final fidelity of the optimized control.

final_purity: Array | None

Final purity of the optimized control.

final_state: Array

Final operator after applying the optimized control amplitudes.

iterations: int

Number of iterations taken for optimization.

optimized_trainable_parameters: Array

Optimized control amplitudes.

purity_each_timestep: list[Array]

Purity of the optimized control along each timestep and batch.

returned_params: List[Array]

Array of final parameters for each time step.

class feedback_grape.fgrape.Gate(gate: callable, initial_params: list[float], measurement_flag: bool, param_constraints: list[float] | None = None)

Bases: NamedTuple

Gate class to store the parameters for each gate.

gate: callable

Function that applies the gate to the state.

initial_params: list[float]

Initial parameters for the gate.

measurement_flag: bool

Flag indicating if the gate is used for measurement.

param_constraints: list[float] | None

Serves two distinct purposes.

First, it bounds the parameters applied to the gates during the forward pass: at every gate application, parameters lying outside their specified (min, max) range are softly mapped back into it via a sigmoid function, while parameters already within bounds are left unchanged.

Second, it controls the initialization of the trainable variables, though the precise effect depends on the chosen feedback mode.

• In ‘no-measurement’ mode, the user-supplied initial_params seed the first time step, while the variables of the remaining time steps are initialized by sampling uniformly within the specified ranges; if param_constraints is omitted, the same initial_params are instead copied to every time step.

• In ‘lookup’ mode, the analogous behavior applies to the rows of the lookup table: each row is either sampled within the specified ranges or, in the absence of constraints, initialized to the same initial_params.

• In ‘nn’ mode, by contrast, param_constraints plays no role in initialization. The per-step parameters are generated on the fly by the neural network, whose weights are initialized independently, so the parameter serves only its bounding role.

In all modes, initial_params remains required, since it defines the structure (shapes and number) of the gate parameters and seeds the gate(s) applied before any feedback signal is available.

feedback_grape.fgrape.calculate_trajectory(*, rho_cav, parameterized_gates, measurement_indices, initial_params, param_shapes, param_constraints, c_ops, decay_indices, time_steps, rnn_model=None, rnn_params=None, rnn_state=None, lut=None, evo_type, batch_size, rng_key)

Calculate a complete quantum trajectory with feedback.

Parameters:

• rho_cav – Initial density matrix of the cavity.

• parameterized_gates – List of parameterized gates.

• measurement_indices – Indices of gates used for measurements.

• initial_params – Initial parameters for all gates.

• param_shapes – List of shapes for each gate’s parameters.

• time_steps – Number of time steps within a trajectory.

• rnn_model – rnn model for feedback.

• rnn_params – Parameters of the rnn model.

• rnn_state – Initial state of the rnn model.

• evo_type – Type of quantum system representation (e.g., “density”).

Returns:

Final state, log probability, array of POVM parameters

feedback_grape.fgrape.optimize_pulse(init_state: ~jax.jaxlib._jax.Array, target_state: ~jax.jaxlib._jax.Array, system_params: list[~feedback_grape.fgrape.Gate], num_time_steps: int, max_iter: int, convergence_threshold: float, learning_rate: float, evo_type: str, reward_weights: list[float | int] | tuple[float | int, ...] | None = None, goal: str = 'fidelity', batch_size: int = 1, eval_batch_size: int = 10, eval_time_steps: int | None = None, mode: str = 'lookup', rnn: callable = <class 'feedback_grape.utils.fgrape_helpers.RNN'>, rnn_hidden_size: int = 30, progress: bool = False) → FgResult

Optimizes pulse parameters for quantum systems based on the specified configuration using ADAM.

Parameters:

• init_state – Initial state or density matrix.

• target_state – Target state or density matrix.

• system_params – List of Gate objects containing gate functions, initial parameters, measurement flags, and parameter constraints.

• num_time_steps (int) – The number of time steps for the optimization process.

• max_iter (int) – The maximum number of iterations for the optimization process.

• convergence_threshold (float) – The threshold for convergence to determine when to stop optimization provide None to enforce max iterations.

• learning_rate (float) – The learning rate for the optimization algorithm.

• evo_type (str) – The evo_type of quantum system representation, such as ‘state’, ‘density’.

• reward_weights (list[float] | None) –

  Weights for the reward at each time step. The 0th index refers to the point after the gates from the first time step are applied. If None, only the final time step is weighted.

  • (default: None)

• goal (str) –

  The optimization goal, which can be purity, fidelity, or both

  • (default: fidelity)

• batch_size (int) –

  The number of trajectories to process in parallel

  • (default: 1)

• eval_batch_size (int) –

  The number of trajectories to process in parallel during evaluation

  • (default: 10)

• eval_time_steps (int) –

  The number of time steps to use during evaluation.

  • (default: None) If None, it will be set to num_time_steps.

• mode (str) –

  The mode of operation, either ‘nn’ (neural network) or ‘lookup’ (lookup table)

  • (default: lookup)

• rnn (callable) –

  The rnn model to use for the optimization process. Defaults to a predefined rnn class. Only used if mode is ‘nn’.

  • (default: RNN)

• rnn_hidden_size (int) –

  The hidden size of the rnn model. Only used if mode is ‘nn’. (output size is inferred from the number of parameters)

  • (default: 30)

• progress –

  Whether to show progress (cost every 10 iterations) during optimization. (for debugging purposes). This may significantly slow down the optimization process

  • (default: False).

Returns:

Dictionary containing optimized pulse and convergence data.

Return type:

result