/*MC
TSRK - ODE and DAE solver using Runge-Kutta schemes
The user should provide the right hand side of the equation
using TSSetRHSFunction().
Notes:
The default is TSRK3, it can be changed with TSRKSetType() or -ts_rk_type
Level: beginner
.seealso: TSCreate(), TS, TSSetType(), TSRKSetType(), TSRKGetType(), TSRKSetFullyImplicit(), TSRK2D, TTSRK2E, TSRK3,
TSRK4, TSRK5, TSRKPRSSP2, TSRKBPR3, TSRKType, TSRKRegister()
M*/
PETSC_EXTERN PetscErrorCode TSCreate_RK(TS ts)
{
TS_RK *rk;
PetscErrorCode ierr;
PetscFunctionBegin;
ierr = TSRKInitializePackage();CHKERRQ(ierr);
ts->ops->reset = TSReset_RK;
ts->ops->destroy = TSDestroy_RK;
ts->ops->view = TSView_RK;
ts->ops->load = TSLoad_RK;
ts->ops->setup = TSSetUp_RK;
ts->ops->adjointsetup = TSAdjointSetUp_RK;
ts->ops->step = TSStep_RK;
ts->ops->interpolate = TSInterpolate_RK;
ts->ops->evaluatestep = TSEvaluateStep_RK;
ts->ops->rollback = TSRollBack_RK;
ts->ops->setfromoptions = TSSetFromOptions_RK;
ts->ops->getstages = TSGetStages_RK;
ts->ops->adjointstep = TSAdjointStep_RK;
ts->ops->adjointintegral = TSAdjointCostIntegral_RK;
ts->ops->forwardintegral = TSForwardCostIntegral_RK;
ierr = PetscNewLog(ts,&rk);CHKERRQ(ierr);
ts->data = (void*)rk;
ierr = PetscObjectComposeFunction((PetscObject)ts,"TSRKGetType_C",TSRKGetType_RK);CHKERRQ(ierr);
ierr = PetscObjectComposeFunction((PetscObject)ts,"TSRKSetType_C",TSRKSetType_RK);CHKERRQ(ierr);
ierr = TSRKSetType(ts,TSRKDefault);CHKERRQ(ierr);
PetscFunctionReturn(0);
}
PetscErrorCode TSRKGetType_RK(TS ts,TSRKType *rktype)
{
TS_RK *rk = (TS_RK*)ts->data;
PetscErrorCode ierr;
PetscFunctionBegin;
if (!rk->tableau) {
ierr = TSRKSetType(ts,TSRKDefault);CHKERRQ(ierr);
}
*rktype = rk->tableau->name;
PetscFunctionReturn(0);
}
static PetscErrorCode TSSetFromOptions_RK(PetscOptionItems *PetscOptionsObject,TS ts)
{
TS_RK *rk = (TS_RK*)ts->data;
PetscErrorCode ierr;
PetscFunctionBegin;
ierr = PetscOptionsHead(PetscOptionsObject,"RK ODE solver options");CHKERRQ(ierr);
{
RKTableauLink link;
PetscInt count,choice;
PetscBool flg;
const char **namelist;
for (link=RKTableauList,count=0; link; link=link->next,count++) ;
ierr = PetscMalloc1(count,&namelist);CHKERRQ(ierr);
for (link=RKTableauList,count=0; link; link=link->next,count++) namelist[count] = link->tab.name;
ierr = PetscOptionsEList("-ts_rk_type","Family of RK method","TSRKSetType",(const char*const*)namelist,count,rk->tableau->name,&choice,&flg);CHKERRQ(ierr);
if (flg) {ierr = TSRKSetType(ts,namelist[choice]);CHKERRQ(ierr);}
ierr = PetscFree(namelist);CHKERRQ(ierr);
}
ierr = PetscOptionsTail();CHKERRQ(ierr);
PetscFunctionReturn(0);
}
static PetscErrorCode TSSetFromOptions_RK(PetscOptions *PetscOptionsObject,TS ts)
{
PetscErrorCode ierr;
char rktype[256];
PetscFunctionBegin;
ierr = PetscOptionsHead(PetscOptionsObject,"RK ODE solver options");CHKERRQ(ierr);
{
RKTableauLink link;
PetscInt count,choice;
PetscBool flg;
const char **namelist;
ierr = PetscStrncpy(rktype,TSRKDefault,sizeof(rktype));CHKERRQ(ierr);
for (link=RKTableauList,count=0; link; link=link->next,count++) ;
ierr = PetscMalloc1(count,&namelist);CHKERRQ(ierr);
for (link=RKTableauList,count=0; link; link=link->next,count++) namelist[count] = link->tab.name;
ierr = PetscOptionsEList("-ts_rk_type","Family of RK method","TSRKSetType",(const char*const*)namelist,count,rktype,&choice,&flg);CHKERRQ(ierr);
ierr = TSRKSetType(ts,flg ? namelist[choice] : rktype);CHKERRQ(ierr);
ierr = PetscFree(namelist);CHKERRQ(ierr);
}
ierr = PetscOptionsTail();CHKERRQ(ierr);
PetscFunctionReturn(0);
}
static PetscErrorCode TSSetUp_RK(TS ts)
{
TS_RK *rk = (TS_RK*)ts->data;
RKTableau tab;
PetscInt s;
PetscErrorCode ierr;
DM dm;
PetscFunctionBegin;
if (!rk->tableau) {
ierr = TSRKSetType(ts,TSRKDefault);CHKERRQ(ierr);
}
tab = rk->tableau;
s = tab->s;
ierr = VecDuplicateVecs(ts->vec_sol,s,&rk->Y);CHKERRQ(ierr);
ierr = VecDuplicateVecs(ts->vec_sol,s,&rk->YdotRHS);CHKERRQ(ierr);
ierr = PetscMalloc1(s,&rk->work);CHKERRQ(ierr);
ierr = TSGetDM(ts,&dm);CHKERRQ(ierr);
if (dm) {
ierr = DMCoarsenHookAdd(dm,DMCoarsenHook_TSRK,DMRestrictHook_TSRK,ts);CHKERRQ(ierr);
ierr = DMSubDomainHookAdd(dm,DMSubDomainHook_TSRK,DMSubDomainRestrictHook_TSRK,ts);CHKERRQ(ierr);
}
PetscFunctionReturn(0);
}
/*
* time_step solves for the time_dependence of the system
* that was previously setup using the add_to_ham and add_lin
* routines. Solver selection and parameters can be controlled via PETSc
* command line options. Default solver is TSRK3BS
*
* Inputs:
* Vec x: The density matrix, with appropriate inital conditions
* double dt: initial timestep. For certain explicit methods, this timestep
* can be changed, as those methods have adaptive time steps
* double time_max: the maximum time to integrate to
* int steps_max: max number of steps to take
*/
void time_step(Vec x, PetscReal init_time, PetscReal time_max,PetscReal dt,PetscInt steps_max){
PetscViewer mat_view;
TS ts; /* timestepping context */
PetscInt i,j,Istart,Iend,steps,row,col;
PetscScalar mat_tmp;
PetscReal tmp_real;
Mat AA;
PetscInt nevents,direction;
PetscBool terminate;
operator op;
int num_pop;
double *populations;
Mat solve_A,solve_stiff_A;
PetscLogStagePop();
PetscLogStagePush(solve_stage);
if (_lindblad_terms) {
if (nid==0) {
printf("Lindblad terms found, using Lindblad solver.\n");
}
solve_A = full_A;
if (_stiff_solver) {
if(nid==0) printf("ERROR! Lindblad-stiff solver untested.");
exit(0);
}
} else {
if (nid==0) {
printf("No Lindblad terms found, using (more efficient) Schrodinger solver.\n");
}
solve_A = ham_A;
solve_stiff_A = ham_stiff_A;
if (_num_time_dep&&_stiff_solver) {
if(nid==0) printf("ERROR! Schrodinger-stiff + timedep solver untested.");
exit(0);
}
}
/* Possibly print dense ham. No stabilization is needed? */
if (nid==0) {
/* Print dense ham, if it was asked for */
if (_print_dense_ham){
FILE *fp_ham;
fp_ham = fopen("ham","w");
if (nid==0){
for (i=0;i<total_levels;i++){
for (j=0;j<total_levels;j++){
fprintf(fp_ham,"%e %e ",PetscRealPart(_hamiltonian[i][j]),PetscImaginaryPart(_hamiltonian[i][j]));
}
fprintf(fp_ham,"\n");
}
}
fclose(fp_ham);
for (i=0;i<total_levels;i++){
free(_hamiltonian[i]);
}
free(_hamiltonian);
_print_dense_ham = 0;
}
}
/* Remove stabilization if it was previously added */
if (stab_added){
if (nid==0) printf("Removing stabilization...\n");
/*
* We add 1.0 in the 0th spot and every n+1 after
*/
if (nid==0) {
row = 0;
for (i=0;i<total_levels;i++){
col = i*(total_levels+1);
mat_tmp = -1.0 + 0.*PETSC_i;
MatSetValue(full_A,row,col,mat_tmp,ADD_VALUES);
}
}
}
MatGetOwnershipRange(solve_A,&Istart,&Iend);
/*
* Explicitly add 0.0 to all diagonal elements;
* this fixes a 'matrix in wrong state' message that PETSc
* gives if the diagonal was never initialized.
*/
//if (nid==0) printf("Adding 0 to diagonal elements...\n");
for (i=Istart;i<Iend;i++){
mat_tmp = 0 + 0.*PETSC_i;
MatSetValue(solve_A,i,i,mat_tmp,ADD_VALUES);
}
if(_stiff_solver){
MatGetOwnershipRange(solve_stiff_A,&Istart,&Iend);
for (i=Istart;i<Iend;i++){
mat_tmp = 0 + 0.*PETSC_i;
MatSetValue(solve_stiff_A,i,i,mat_tmp,ADD_VALUES);
}
}
/* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -*
* Create the timestepping solver and set various options *
*- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
/*
* Create timestepping solver context
*/
TSCreate(PETSC_COMM_WORLD,&ts);
TSSetProblemType(ts,TS_LINEAR);
/*
* Set function to get information at every timestep
*/
if (_ts_monitor!=NULL){
TSMonitorSet(ts,_ts_monitor,_tsctx,NULL);
}
/*
* Set up ODE system
*/
TSSetRHSFunction(ts,NULL,TSComputeRHSFunctionLinear,NULL);
if(_stiff_solver) {
/* TSSetIFunction(ts,NULL,TSComputeRHSFunctionLinear,NULL); */
if (nid==0) {
printf("Stiff solver not implemented!\n");
exit(0);
}
if(nid==0) printf("Using stiff solver - TSROSW\n");
}
if(_num_time_dep+_num_time_dep_lin) {
for(i=0;i<_num_time_dep;i++){
tmp_real = 0.0;
_add_ops_to_mat_ham(tmp_real,solve_A,_time_dep_list[i].num_ops,_time_dep_list[i].ops);
}
for(i=0;i<_num_time_dep_lin;i++){
tmp_real = 0.0;
_add_ops_to_mat_lin(tmp_real,solve_A,_time_dep_list_lin[i].num_ops,_time_dep_list_lin[i].ops);
}
/* Tell PETSc to assemble the matrix */
MatAssemblyBegin(solve_A,MAT_FINAL_ASSEMBLY);
MatAssemblyEnd(solve_A,MAT_FINAL_ASSEMBLY);
if (nid==0) printf("Matrix Assembled.\n");
MatDuplicate(solve_A,MAT_COPY_VALUES,&AA);
MatAssemblyBegin(AA,MAT_FINAL_ASSEMBLY);
MatAssemblyEnd(AA,MAT_FINAL_ASSEMBLY);
TSSetRHSJacobian(ts,AA,AA,_RHS_time_dep_ham_p,NULL);
} else {
/* Tell PETSc to assemble the matrix */
MatAssemblyBegin(solve_A,MAT_FINAL_ASSEMBLY);
MatAssemblyEnd(solve_A,MAT_FINAL_ASSEMBLY);
if (_stiff_solver){
MatAssemblyBegin(solve_stiff_A,MAT_FINAL_ASSEMBLY);
MatAssemblyEnd(solve_stiff_A,MAT_FINAL_ASSEMBLY);
/* TSSetIJacobian(ts,solve_stiff_A,solve_stiff_A,TSComputeRHSJacobianConstant,NULL); */
if (nid==0) {
printf("Stiff solver not implemented!\n");
exit(0);
}
}
if (nid==0) printf("Matrix Assembled.\n");
TSSetRHSJacobian(ts,solve_A,solve_A,TSComputeRHSJacobianConstant,NULL);
}
/* Print information about the matrix. */
PetscViewerASCIIOpen(PETSC_COMM_WORLD,NULL,&mat_view);
PetscViewerPushFormat(mat_view,PETSC_VIEWER_ASCII_INFO);
/* PetscViewerPushFormat(mat_view,PETSC_VIEWER_ASCII_MATLAB); */
/* MatView(solve_A,mat_view); */
/* PetscInt ncols; */
/* const PetscInt *cols; */
/* const PetscScalar *vals; */
/* for(i=0;i<total_levels*total_levels;i++){ */
/* MatGetRow(solve_A,i,&ncols,&cols,&vals); */
/* for (j=0;j<ncols;j++){ */
/* if(PetscAbsComplex(vals[j])>1e-5){ */
/* printf("%d %d %lf %lf\n",i,cols[j],vals[j]); */
/* } */
/* } */
/* MatRestoreRow(solve_A,i,&ncols,&cols,&vals); */
/* } */
if(_stiff_solver){
MatView(solve_stiff_A,mat_view);
}
PetscViewerPopFormat(mat_view);
PetscViewerDestroy(&mat_view);
TSSetTimeStep(ts,dt);
/*
* Set default options, can be changed at runtime
*/
TSSetMaxSteps(ts,steps_max);
TSSetMaxTime(ts,time_max);
TSSetTime(ts,init_time);
TSSetExactFinalTime(ts,TS_EXACTFINALTIME_STEPOVER);
if (_stiff_solver) {
TSSetType(ts,TSROSW);
} else {
TSSetType(ts,TSRK);
TSRKSetType(ts,TSRK3BS);
}
/* If we have gates to apply, set up the event handler. */
if (_num_quantum_gates > 0) {
nevents = 1; //Only one event for now (did we cross a gate?)
direction = -1; //We only want to count an event if we go from positive to negative
terminate = PETSC_FALSE; //Keep time stepping after we passed our event
/* Arguments are: ts context, nevents, direction of zero crossing, whether to terminate,
* a function to check event status, a function to apply events, private data context.
*/
TSSetEventHandler(ts,nevents,&direction,&terminate,_QG_EventFunction,_QG_PostEventFunction,NULL);
}
if (_num_circuits > 0) {
nevents = 1; //Only one event for now (did we cross a gate?)
direction = -1; //We only want to count an event if we go from positive to negative
terminate = PETSC_FALSE; //Keep time stepping after we passed our event
/* Arguments are: ts context, nevents, direction of zero crossing, whether to terminate,
* a function to check event status, a function to apply events, private data context.
*/
TSSetEventHandler(ts,nevents,&direction,&terminate,_QC_EventFunction,_QC_PostEventFunction,NULL);
}
if (_discrete_ec > 0) {
nevents = 1; //Only one event for now (did we cross an ec step?)
direction = -1; //We only want to count an event if we go from positive to negative
terminate = PETSC_FALSE; //Keep time stepping after we passed our event
/* Arguments are: ts context, nevents, direction of zero crossing, whether to terminate,
* a function to check event status, a function to apply events, private data context.
*/
TSSetEventHandler(ts,nevents,&direction,&terminate,_DQEC_EventFunction,_DQEC_PostEventFunction,NULL);
}
/* if (_lindblad_terms) { */
/* nevents = 1; //Only one event for now (did we cross a gate?) */
/* direction = 0; //We only want to count an event if we go from positive to negative */
/* terminate = PETSC_FALSE; //Keep time stepping after we passed our event */
/* TSSetEventHandler(ts,nevents,&direction,&terminate,_Normalize_EventFunction,_Normalize_PostEventFunction,NULL); */
/* } */
TSSetFromOptions(ts);
TSSolve(ts,x);
TSGetStepNumber(ts,&steps);
num_pop = get_num_populations();
populations = malloc(num_pop*sizeof(double));
get_populations(x,&populations);
/* if(nid==0){ */
/* printf("Final populations: "); */
/* for(i=0;i<num_pop;i++){ */
/* printf(" %e ",populations[i]); */
/* } */
/* printf("\n"); */
/* } */
/* PetscPrintf(PETSC_COMM_WORLD,"Steps %D\n",steps); */
/* Free work space */
TSDestroy(&ts);
if(_num_time_dep+_num_time_dep_lin){
MatDestroy(&AA);
}
free(populations);
PetscLogStagePop();
PetscLogStagePush(post_solve_stage);
return;
}