fgrape (feedback GRAPE)

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#

This constraints the initialization of the parameters to be within the specified range. This also constraints the parameters that gets applied to the gates by clipping to your specified range using a sigmoid function.

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(U_0: ~jax.jaxlib._jax.Array, C_target: ~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:

  • U_0 – Initial state or density matrix.

  • C_target – 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. 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

Contents