Dataset

This module contain the implementation of the Dataset class

Overview

This module defines dataset utilities. It includes functions for sampling quantum spin configurations and classical shadows, plus the Dataset class for organizing quantum dataset within the qdisc pipeline.

Helper functions for data creation

These functions help build quantum training data from exact wavefunctions. The helpers include:

sample_spin_configurations(...): sample bitstring configurations from a wavefunction probability distribution.
get_classical_shadow(...): generate classical shadow measurements in X/Y/Z bases.

source

get_classical_shadow


def get_classical_shadow(
    psi, num_shots, N, rng_key, bases:str='XYZ'
):

generate classical shadows from an exact wavefunction

Args:

psi: (2**N,) complex64 from, f. ex., exact diagonalization num_shots: int N: int rng_key: a JAX PRNGKey bases: which measurement basis, default ‘XYZ’ -> {X,Y,Z} randomly, other available option: ‘X’, ‘Y’, ‘Z’, ‘XY’, ‘YZ’

Returns:

shadows: (num_shots, 2*N) int32 interleaved [basis,outcome] pairs basis in {2,3,4} = {X,Y,Z}, outcome in {0,1} = {+1,-1}

source

sample_spin_configurations


def sample_spin_configurations(
    wave_function, num_samples, N, key
):

samples spin config. from an exact wavefunction

Args:

wave_function: (2**N,) complex64 num_samples: int, nbr of samples N: int, nbr of sites key: a JAX PRNGKey

Returns:

spin_configurations: (num_samples, N) int32, samples of spin configurations

Dataset class

The Dataset class centralizes quantum data, parameter grids, and help create training sets. It supports:

validating input tensor shapes and parameter grids
creating shuffled training batches via create_batches(...)

source

Dataset


def Dataset(
    data:Array, thetas:Sequence, data_type:str, local_dimension:int=0, local_states:Array=None
):

source

Dataset.create_batches


def create_batches(
    batch_size:int, key:Array
)->list:

create batches of data, shuffled