Example #1
0
def _init_ptrace_stuff(sel):
    psi0 = Qobj(odeconfig.psi0,
                dims=odeconfig.psi0_dims,
                shape=odeconfig.psi0_shape)
    qtf90.qutraj_run.init_ptrace_stuff(odeconfig.psi0_dims[0],
                                       np.array(sel) + 1,
                                       psi0.ptrace(sel).shape[0])
Example #2
0
def _init_ptrace_stuff(sel):
    psi0 = Qobj(config.psi0,
                dims=config.psi0_dims,
                shape=config.psi0_shape)
    qtf90.qutraj_run.init_ptrace_stuff(config.psi0_dims[0],
                                       np.array(sel) + 1,
                                       psi0.ptrace(sel).shape[0])
def generic_ode_solve_checkpoint(r, rho0, tlist, e_ops, opt, progress_bar,
                                 save, subdir):
    """
    Internal function for solving ME. Solve an ODE which solver parameters
    already setup (r). Calculate the required expectation values or invoke
    callback function at each time step.
    """

    #
    # prepare output array
    #
    n_tsteps = len(tlist)
    e_sops_data = []

    output = Result()
    output.solver = "mesolve"
    output.times = tlist

    if opt.store_states:
        output.states = []

    e_ops_dict = e_ops
    e_ops = [e for e in e_ops_dict.values()]
    headings = [key for key in e_ops_dict.keys()]

    if isinstance(e_ops, types.FunctionType):
        n_expt_op = 0
        expt_callback = True

    elif isinstance(e_ops, list):

        n_expt_op = len(e_ops)
        expt_callback = False

        if n_expt_op == 0:
            # fall back on storing states
            output.states = []
            opt.store_states = True
        else:
            output.expect = []
            output.num_expect = n_expt_op
            for op in e_ops:
                e_sops_data.append(spre(op).data)
                if op.isherm and rho0.isherm:
                    output.expect.append(np.zeros(n_tsteps))
                else:
                    output.expect.append(np.zeros(n_tsteps, dtype=complex))
    else:
        raise TypeError("Expectation parameter must be a list or a function")

    results_row = np.zeros(n_expt_op)

    #
    # start evolution
    #
    progress_bar.start(n_tsteps)

    rho = Qobj(rho0)
    dims = rho.dims

    dt = np.diff(tlist)

    end_time = tlist[-1]

    for t_idx, t in tqdm(enumerate(tlist)):
        progress_bar.update(t_idx)

        if not r.successful():
            raise Exception("ODE integration error: Try to increase "
                            "the allowed number of substeps by increasing "
                            "the nsteps parameter in the Options class.")

        if opt.store_states or expt_callback:
            rho.data = vec2mat(r.y)

            if opt.store_states:
                output.states.append(Qobj(rho))

            if expt_callback:
                # use callback method
                e_ops(t, rho)

        for m in range(n_expt_op):
            if output.expect[m].dtype == complex:
                output.expect[m][t_idx] = expect_rho_vec(
                    e_sops_data[m], r.y, 0)
                results_row[m] = output.expect[m][t_idx]
            else:
                output.expect[m][t_idx] = expect_rho_vec(
                    e_sops_data[m], r.y, 1)
                results_row[m] = output.expect[m][t_idx]

        results = pd.DataFrame(results_row).T
        results.columns = headings
        results.index = [t]
        results.index.name = 'times'
        if t == 0:
            first_row = True
        else:
            first_row = False
        if save:

            rho_checkpoint = Qobj(vec2mat(r.y))
            rho_checkpoint.dims = dims

            if t_idx % 200 == 0:
                rho_c = rho_checkpoint.ptrace(0)
                with open('./cavity_states.pkl', 'ab') as f:
                    pickle.dump(rho_c, f)

            with open('./results.csv', 'a') as file:
                results.to_csv(file, header=first_row, float_format='%.15f')

            qsave(rho_checkpoint, './state_checkpoint')

        save = True

        if t_idx < n_tsteps - 1:
            r.integrate(r.t + dt[t_idx])

    progress_bar.finished()

    if not opt.rhs_reuse and config.tdname is not None:
        _cython_build_cleanup(config.tdname)

    return output