diff --git a/bmad/modules/bmad_struct.f90 b/bmad/modules/bmad_struct.f90
index 92052bf415..c5171a4114 100644
--- a/bmad/modules/bmad_struct.f90
+++ b/bmad/modules/bmad_struct.f90
@@ -19,7 +19,7 @@ module bmad_struct
 ! IF YOU CHANGE THE LAT_STRUCT OR ANY ASSOCIATED STRUCTURES YOU MUST INCREASE THE VERSION NUMBER !!!
 ! THIS IS USED BY BMAD_PARSER TO MAKE SURE DIGESTED FILES ARE OK.
 
-integer, parameter :: bmad_inc_version$ = 317
+integer, parameter :: bmad_inc_version$ = 318
 
 !+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 !+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
diff --git a/tao/code/tao_graph_setup_mod.f90 b/tao/code/tao_graph_setup_mod.f90
index 2cade5839e..29f2e6e226 100644
--- a/tao/code/tao_graph_setup_mod.f90
+++ b/tao/code/tao_graph_setup_mod.f90
@@ -449,7 +449,7 @@ subroutine tao_graph_phase_space_setup (plot, graph)
 logical err, same_uni
 
 character(40) name
-character(40) :: r_name = 'tao_graph_phase_space_setup'
+character(*), parameter :: r_name = 'tao_graph_phase_space_setup'
 
 ! Set up the graph suffix
 
@@ -504,6 +504,24 @@ subroutine tao_graph_phase_space_setup (plot, graph)
   curve => graph%curve(k)
   u => tao_pointer_to_universe (tao_curve_ix_uni(curve))
 
+  if (curve%ix_ele_ref_track < 0) then
+    call out_io (s_error$, r_name, 'CURVE REFERENCE ELEMENT IS NOT AN ELEMENT THAT IS TRACKED THROUGH: ' // curve%ele_ref_name, &
+                                   'FOR CURVE: ' // tao_curve_name(curve))
+  endif
+
+  ib = tao_branch_index(curve%ix_branch)
+  select case (curve%component)
+  case('model', '')
+    ele => u%model%lat%branch(ib)%ele(curve%ix_ele_ref_track)
+  case('design')
+    ele => u%design%lat%branch(ib)%ele(curve%ix_ele_ref_track)
+  case('base')
+    ele => u%base%lat%branch(ib)%ele(curve%ix_ele_ref_track)
+  case default
+    call out_io (s_error$, r_name, 'IF THE CURVE COMPONENT IS SET, IT MUST BE SET TO ONE OF: "model", "design", OR "base".', &
+                                   'FOR CURVE: ' // tao_curve_name(curve))
+  end select
+
   ! fill the curve data arrays
 
   if (allocated (curve%ix_symb) .and. curve%data_source /= 'multi_turn_orbit' .and. &
@@ -522,7 +540,6 @@ subroutine tao_graph_phase_space_setup (plot, graph)
 
     ib = tao_branch_index(curve%ix_branch)
     beam => u%model_branch(ib)%ele(curve%ix_ele_ref_track)%beam
-    ele => u%model%lat%branch(ib)%ele(curve%ix_ele_ref_track)
     if (.not. allocated(beam%bunch)) then
       call out_io (s_warn$, r_name, 'NO BEAM AT ELEMENT: ' // trim(ele%name), &
                     'CANNOT DO PHASE_SPACE PLOTTING FOR CURVE: ' // tao_curve_name(curve))