data_gen.py

import pprint
import argparse
import sys
import atexit
import json
import multiprocessing as mp
from multiprocessing import shared_memory
from pathlib import Path
import networkx
import numpy as np
import qutip as qt
from tqdm import tqdm
from threadpoolctl import threadpool_limits

from utils.cs_utils import *
from utils.pauli_utils import *
from utils.shadow_utils import *
from utils.misc_utils import *
from utils import misc_utils

# ====== color preamble ======
RESET   = "\033[0m"
RED     = "\033[91m"
GREEN   = "\033[92m"
YELLOW  = "\033[93m"
CYAN    = "\033[96m"
# ============================

_FLUSH_EVERY_SHADOWS = 16
_FLUSH_EVERY_EST = 64

######### MP FUNCTIONS ##########
def _init_states(shm_name, shape, dtype_str):
    global _shm_states, _states
    _shm_states = shared_memory.SharedMemory(name=shm_name, create=False)  # keep reference!
    _states = np.ndarray(shape, dtype=np.dtype(dtype_str), buffer=_shm_states.buf)
    _states.setflags(write=False)
    atexit.register(_shm_states.close)  # parent should unlink() when all done

def _init_ops(ostrings):
    global _ops, _NQ, _ostrings
    _NQ = len(ostrings[0])
    _ops = [p2op(x, norm=False) for x in ostrings]
    _ostrings = ostrings

def _init_worker_exact(state_data, ostrings, tplimit):
    threadpool_limits(limits=tplimit)
    _init_states(*state_data)
    _init_ops(ostrings)

def _worker_exact(itt):
    rho = qt.Qobj(_states[itt], dims=[[2]*_NQ]*2, copy=False)
    rho = rho if rho.isherm else (rho + rho.dag())/2  # enforce hermiticity
    expvals = qt.expect(_ops, rho)
    expvals = np.real_if_close(expvals, tol=1e8) # drops imag if |Im| <= 10^-12-ish
    # print(np.max(np.abs(np.imag(expvals))))
    return itt, expvals

def _init_worker_shadow(state_data, all_mm_data, nq, nsmax, tplimit):
    threadpool_limits(limits=tplimit)
    global _NQ, _NSMAX
    _NQ, _NSMAX = nq, nsmax
    _init_states(*state_data)
    for mm_data in all_mm_data:
        init_mm(*mm_data)

def _worker_shadow(args):
    itt, seed = args
    rho = qt.Qobj(_states[itt], dims=[[2]*_NQ]*2, copy=False)
    rho = rho if rho.isherm else (rho + rho.dag())/2  # enforce hermiticity

    misc_utils._mm['shadow'][itt] = get_shadows(rho, _NSMAX, seed=seed)  # (NSMAX, Nq, 2)
    misc_utils._flush['shadow']['cnt'] += 1
    if misc_utils._flush['shadow']['cnt'] % misc_utils._flush['shadow']['every'] == 0:
        misc_utils._mm['shadow'].flush()

def _init_worker_est(all_mm_data, ostrings, shadow_subs, tplimit):
    global _shadow_subs
    _shadow_subs = shadow_subs
    threadpool_limits(limits=tplimit)
    for mm_data in all_mm_data:
        init_mm(*mm_data)
    _init_ops(ostrings)

def _worker_est(itt):
    misc_utils._mm['est'][itt] = estimate_batch(misc_utils._mm['shadow'][itt], _ostrings, _shadow_subs)
    misc_utils._flush['est']['cnt'] += 1
    if misc_utils._flush['est']['cnt'] % misc_utils._flush['est']['every'] == 0:
        misc_utils._mm['est'].flush()

######################
def get_parser():
    """Build the CLI argument parser for CSST data generation."""
    parser = argparse.ArgumentParser(description="Command line arguments for CSST")
    parser.add_argument("--dir",    type=str,     default="data",       help="directory to save data files")
    parser.add_argument("--nx",     type=int,     default=2,            help="Number of sites in x direction")
    parser.add_argument("--ny",     type=int,     default=2,            help="Number of sites in y direction")
    parser.add_argument("--nb",     type=int,     default=4,            help="Max Pauli observable weight")
    parser.add_argument("--ham",    type=str,     default='heis',       help="Hamiltonian type",        choices=['tfim','heis'])
    parser.add_argument("--istate", type=str,     default='neel',       help="Initial state ('ghz','neel','w','r','rp','hr','hrp', or length NQ string of [0,1,+,-,>,<])")
    parser.add_argument("--n",      type=int,     default=500,          help="Number of time steps")
    parser.add_argument("--nsmin",  type=int,     default=10,           help="Min number of shadows")
    parser.add_argument("--nsmax",  type=int,     default=1000,         help="Max number of shadows")
    parser.add_argument("--nsnum",  type=int,     default=10,           help="Number of shadows to try (logspaced)")
    parser.add_argument("--nw",     type=int,     default=1,            help="Number of workers for multiprocessing")
    parser.add_argument("--tplimit",type=int,     default=1,            help="Threadpool limit for each worker process")
    parser.add_argument("--eps",    type=float,   default=0,            help="Max percentage of perturbation to Ham coeffs and Lindblad rates")
    parser.add_argument("--gamma",  type=float,   default=1e-2,         help="Common decay rate for all collapse operators")
    parser.add_argument("--seed",   type=int,     default=42,           help="Random seed for reproducibility")
    return parser

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    pprint.pprint(vars(args))
    print("\n")

    NX = args.nx               # lattice dimension x
    NY = args.ny               # lattice dimension y
    NQ = NX * NY               # number of qubits
    N = args.n                 # number of time steps
    NB = args.nb               # max Pauli observable weight
    HAM = args.ham             # Hamiltonian type: 'tfim', 'random', 'heis'
    NSMIN = args.nsmin         # min number of shadows
    NSMAX = args.nsmax         # max number of shadows
    NSNUM = args.nsnum         # number of shadow subsamples
    ISTATE = args.istate       # initial state
    EPS = args.eps             # max percentage of perturbation to Ham coeffs and Lindblad rates
    GAMMA = args.gamma         # common decay rate for all collapse operators
    NUM_WORKERS = args.nw      # number of workers
    TPLIMIT = args.tplimit     # threadpool limit for each worker process
    SEED = args.seed           # random seed for reproducibility

    is_valid_init_state(NQ, ISTATE)  # validate istate/nq before creating anything

    if EPS:
        print(f"{YELLOW}Perturbing Hamiltonian coefficients with max percentage {EPS}...{RESET}")
        LABEL = f"{NX}x{NY}_{HAM}_{ISTATE}_eps={EPS:.1e}"
    else:
        LABEL = f"{NX}x{NY}_{HAM}_{ISTATE}"

    DIR = Path(args.dir) / LABEL
    if not DIR.exists():
        print(f"Creating directory {DIR}")
    DIR.mkdir(parents=True, exist_ok=True)

    ########## MODEL ##########
    G = networkx.generators.lattice.grid_2d_graph(NX,NY)
    G = networkx.convert_node_labels_to_integers(G)
    H_ops = get_h_ops(NQ, model=HAM, graph=G, seed=SEED, eps=EPS)

    c_ops = []
    rng = np.random.default_rng(SEED + 1)
    # amplitude damping rates
    ad_gammas = ([np.sqrt(GAMMA)]*NQ) * (1 + EPS * rng.uniform(low=-1, high=1, size=NQ))
    ad_ops = amp_damp_ops(NQ)
    c_ops += [gamma * op for (gamma,op) in zip(ad_gammas, ad_ops)]

    # dephasing rates
    dz_gammas = ([np.sqrt(GAMMA)]*NQ) * (1 + EPS * rng.uniform(low=-1, high=1, size=NQ))
    dz_ops = [p2op(z) for z in pbw(NQ,nb=1,ptype='Z')]
    c_ops += [gamma * op for (gamma,op) in zip(dz_gammas, dz_ops)]

    H = ps2op(H_ops)
    # L = qt.liouvillian(H,c_ops=c_ops)

    re_bound, im_bound = get_eigval_bounds_loose(H_ops, ad_gammas + dz_gammas, ub=1)
    # print(f"\nReal part bound: {re_bound}, Imaginary part bound: {im_bound}")
    dt_nyq = np.pi/im_bound
    dt = dt_nyq/4

    ########## SIMULATION PARAMETERS ##########
    MAXT = dt * N
    times = np.linspace(0,MAXT,N)
    init_state = get_init_state(NQ, ISTATE, graph=G)
    ostrings = pbw(NQ, nb=NB, max=True)
    NUM_PAULIS = len(ostrings)

    # ODE solver options
    options = {'nsteps':1000,           #<-- make this bigger if you get an ODE error
               'progress_bar':True,
               'atol':1e-9,
               'rtol':1e-7
               }

    ########## MESOLVE SIMULATION ##########
    print("\nStarting mesolve...")
    result = qt.mesolve(H,
                        rho0=init_state,
                        tlist=times,
                        c_ops=c_ops,
                        e_ops=[],
                        options=options) #<-- if e_ops == [], then the density matrix
                                          #    is returned at each time
    print("\n")
    STATES = list(result.states)
    print(f"{'Size of states:':<20} {STATES.__sizeof__() / (1024**2):8.2f} MB")

    ########################################
    shadow_subs = np.round(np.logspace(np.log10(NSMIN),np.log10(NSMAX),NSNUM)).astype(int)

    #make rng
    rng = np.random.default_rng(SEED)

    #data arrays
    exact_path = DIR / "exacts.npy"
    exacts = create_memmap(exact_path, shape=(N, NUM_PAULIS), dtype=np.float32, fill=np.nan)
    print(f"{'Size of exacts:':<20} {exacts.nbytes / (1024**2):8.2f} MB")

    err_path = DIR / "errs.npy"
    errs = create_memmap(err_path, shape=(NUM_PAULIS, NSNUM, 8), dtype=np.float32, fill=np.nan)
    print(f"{'Size of errs:':<20} {errs.nbytes / (1024**2):8.2f} MB")

    shadow_path = DIR / "shadows.npy"
    # cant put nan in int8
    shadows = create_memmap(shadow_path, shape=(N, NSMAX, NQ, 2), dtype=np.int8, fill=0)
    print(f"{'Size of shadows:':<20} {shadows.nbytes / (1024**2):8.2f} MB")

    est_path = DIR / "ests.npy"
    ests = create_memmap(est_path, shape=(N, NUM_PAULIS, NSNUM), dtype=np.float32, fill=np.nan)
    print(f"{'Size of ests:':<20} {ests.nbytes / (1024**2):8.2f} MB")
    print("\n")

    ########################################
    info = {'nx':     NX,
            'ny':     NY,
            'nb':     NB,
            'ham':    HAM,
            'istate': ISTATE,
            'nsmin':  NSMIN,
            'nsmax':  NSMAX,
            'nsnum':  NSNUM,
            'tmin':   times[0],
            'tmax':   times[-1],
            'n':      N,
            'dt':     dt,
            'gamma':  GAMMA,
            'nw':     NUM_WORKERS,
            'tplimit':TPLIMIT,
            'eps':    EPS,
            'seed':   SEED
            }

    def to_builtin(x):
        """Convert numpy scalars to plain Python builtins so they're JSON-serializable."""
        return x.item() if hasattr(x, "item") else x
    
    info = {k: to_builtin(v) for k, v in info.items()}
    with (DIR / "info.json").open("w") as f:
        json.dump(info, f, indent=2)
    pprint.pprint(info)
    print("\n")
    ########################################

    mp.set_start_method("spawn", force=True)
    NUM_WORKERS = min(cpu_cap(), NUM_WORKERS)
    print(f"{CYAN}Using multiprocessing with {NUM_WORKERS} workers...{RESET}")
    print(f"{CYAN}Threadpool limit per worker: {TPLIMIT}{RESET}")
    print("\n")

    try:
        # convert to (N, d, d) complex128
        states_arr = np.stack([np.asarray(temp.full()) for temp in STATES])
        # create shared memory block
        shm = shared_memory.SharedMemory(create=True, size=states_arr.nbytes)
        state_data = [shm.name, states_arr.shape, str(states_arr.dtype)]

        shm_view = np.ndarray(states_arr.shape, dtype=states_arr.dtype, buffer=shm.buf)
        shm_view[:] = states_arr
        shm_view.setflags(write=False)

        #exacts
        with mp.Pool(processes=NUM_WORKERS,
                     initializer=_init_worker_exact,
                     initargs=(state_data, ostrings, TPLIMIT,)) as pool:

            for itt, row in tqdm(pool.imap_unordered(_worker_exact, range(N), chunksize=16),
                                 total=N,
                                 desc=f"{'exacts over time':<35}",
                                 mininterval=1,
                                 file=sys.stdout, leave=True, disable=False):
                exacts[itt] = row
        exacts.flush()

        #shadows
        with mp.Pool(processes=NUM_WORKERS,
                     initializer=_init_worker_shadow,
                     maxtasksperchild=50,
                     initargs=(state_data,
                               [['shadow', shadow_path, 'r+', _FLUSH_EVERY_SHADOWS]],
                               NQ,
                               NSMAX,
                               TPLIMIT)) as pool:

            task_iter = (
                (itt, int(rng.integers(0, 2**32 - 1, dtype=np.uint32)))
                for itt in range(N)
            )

            for _ in tqdm(pool.imap_unordered(_worker_shadow, task_iter, chunksize=1),
                          total=N,
                          desc=f"{'generating shadows':<35}",
                          miniters=max(1, N // 100),
                          mininterval=1.0,
                          file=sys.stdout, leave=True, disable=False):
                pass #workers write directly to memmap
        shadows.flush()

        #estimates
        task_iter = (
            itt
            for itt in range(N)
            )

        with mp.Pool(processes=NUM_WORKERS,
                     initializer=_init_worker_est,
                     maxtasksperchild=50, #trying to fix hung workers
                     initargs=([['shadow', shadow_path, 'r', None],
                                ['est', est_path, 'r+', _FLUSH_EVERY_EST]],
                                ostrings,
                                shadow_subs,
                                TPLIMIT)) as pool:

            for _ in tqdm(pool.imap_unordered(_worker_est, task_iter, chunksize=10),
                          total=N,
                          desc=f"{'estimating shadow batches':<35}",
                          mininterval=1.0,
                          file=sys.stdout, leave=True, disable=False):
                pass #workers write directly to memmap
        ests.flush()

        #errors
        #compute errors in main process since this is fast enough (by now, ests is populated)
        for ii in tqdm(range(NUM_PAULIS), desc=f"{'computing errors':<35}"):
            true = exacts[:,ii].copy()
            for jj in range(NSNUM):
                recon = ests[:,ii,jj].copy()
                errs[ii,jj,:] = get_errors(recon, true)

    finally:
        shm.close()
        shm.unlink()