Plotting survival curves from stpm2

[Craig Knott]

I'm looking to plot differences in survival between treatment groups. The KM curves are far from proportional, so I've started down the route of using stpm2, which I understand is a useful means of calculating hazards and survival in the presence of non-proportionality.

Following some of Paul Lambert's guidance (https://pclambert.net/software/stpm2/stpm2_timevar/), and beginning with things at their most basic, I've been able to plot survival for each treatment group without adjustment for any covariates. The resulting figure matches what I get by running sts graph, which is reassuring.

 // Run the model      
xi: stpm2 i.treatment_group, scale(hazard) tvc(i.treatment_group) df(3) dftvc(1) eform  

// Predict the survival curve for each treatment group
predict s1 if treatment_group==1, survival timevar(_t) ci
predict s2 if treatment_group==2, survival timevar(_t) ci
predict s3 if treatment_group==3, survival timevar(_t) ci
predict s4 if treatment_group==4, survival timevar(_t) ci
predict s5 if treatment_group==5, survival timevar(_t) ci  

// Plot the survival curves
twoway (rarea s1_lci s1_uci _t, sort) (line s1 _t, sort) ///        
(rarea s2_lci s2_uci _t, sort) (line s2 _t, sort) ///        
(rarea s3_lci s3_uci _t, sort) (line s3 _t, sort) ///        
(rarea s4_lci s4_uci _t, sort) (line s4 _t, sort) ///        
(rarea s5_lci s5_uci _t, sort) (line s5 _t, sort)

Things prove problematic when I add covariates and try to plot the adjusted survival curves for each treatment group. I've tried various combinations of the at() and zero options, but having trouble producing graphs that make sense. Any help would be much appreciated. From the example code below, I end up with five identical curves, which is far from consistent with the output from the model itself.

 
// Run the model      
xi: stpm2 i.treatment_group i.morph_logistic i.stage_best i.age_group i.imd charlson_score i.can_alliance, scale(hazard) tvc(i.treatment_group) df(3) dftvc(1) eform  

// Predict the survival curve for each treatment group
predict s1 if treatment_group==1, survival timevar(_t)  ci at(morph_logistic 6 stage_best 4 age_group 1 imd 1 charlson_score 3 can_alliance 3) zeros
predict s2 if treatment_group==2, survival timevar(_t) ci at(morph_logistic 6 stage_best 4 age_group 1 imd 1 charlson_score 3 can_alliance 3) zeros
predict s3 if treatment_group==3, survival timevar(_t) ci at(morph_logistic 6 stage_best 4 age_group 1 imd 1 charlson_score 3 can_alliance 3) zeros
predict s4 if treatment_group==4, survival timevar(_t) ci at(morph_logistic 6 stage_best 4 age_group 1 imd 1 charlson_score 3 can_alliance 3) zeros
predict s5 if treatment_group==5, survival timevar(_t) ci at(morph_logistic 6 stage_best 4 age_group 1 imd 1 charlson_score 3 can_alliance 3) zeros  

// Plot the adjusted survival curves
twoway (rarea s1_lci s1_uci _t, sort) (line s1 _t, sort) ///        
(rarea s2_lci s2_uci _t, sort) (line s2 _t, sort) ///        
(rarea s3_lci s3_uci _t, sort) (line s3 _t, sort) ///        
(rarea s4_lci s4_uci _t, sort) (line s4 _t, sort) ///        
(rarea s5_lci s5_uci _t, sort) (line s5 _t, sort)

I'm not sure where I'm going wrong.