From e97437c3ccfbef0b82ef1d8c69e8144c9d9295bc Mon Sep 17 00:00:00 2001
From: David Sagan <david.sagan@gmail.com>
Date: Sun, 14 Jul 2024 14:20:37 -0400
Subject: [PATCH] Fix Tao python  and  output. (#1060)

* Fix Tao python `plot_lat_layout` and `floor_plan` output.
---
 tao/code/tao_init_data_mod.f90 |  3 ---
 tao/code/tao_python_cmd.f90    | 31 ++++++++++++++++++++++++++-----
 tao/code/tao_struct.f90        |  2 +-
 3 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/tao/code/tao_init_data_mod.f90 b/tao/code/tao_init_data_mod.f90
index 0320b7c50a..ae3ac0312a 100644
--- a/tao/code/tao_init_data_mod.f90
+++ b/tao/code/tao_init_data_mod.f90
@@ -611,9 +611,6 @@ subroutine d1_data_stuffit (i_d1, u, n_d2, datum, default_merit_type, &
 
   if (tao_rad_int_calc_needed(dat%data_type, dat%data_source)) then
     u%calc%rad_int_for_data = .true. 
-    if (dat%ix_branch /= 0) then
-      call out_io (s_error$, r_name, 'EVALUATING A DATUM OF TYPE: ' // dat%data_type, 'ON A BRANCH NOT YET IMPLEMENTED!')
-    endif
   endif
 
   if (tao_lat_sigma_calc_needed(dat%data_type, dat%data_source)) u%calc%lat_sigma_for_data = .true. 
diff --git a/tao/code/tao_python_cmd.f90 b/tao/code/tao_python_cmd.f90
index 552878f104..87b1280fb3 100644
--- a/tao/code/tao_python_cmd.f90
+++ b/tao/code/tao_python_cmd.f90
@@ -107,7 +107,7 @@ subroutine tao_python_cmd (input_str)
 type (bunch_track_struct), pointer :: bunch_params_comb(:)
 type (bunch_track_struct), pointer :: comb1
 type (ele_pointer_struct), allocatable :: eles(:), eles2(:)
-type (branch_struct), pointer :: branch
+type (branch_struct), pointer :: branch, branch2
 type (tao_model_branch_struct), pointer :: model_branch
 type (random_state_struct) ran_state
 type (ele_attribute_struct) attrib
@@ -4256,7 +4256,7 @@ subroutine tao_python_cmd (input_str)
 
 case ('floor_plan')
 
-  call tao_find_plots (err, line, 'BOTH', graph = graphs, only_visible = .false.)
+  call tao_find_plots (err, line, 'BOTH', graph = graphs, blank_means_all = .true., only_visible = .false.)
 
   if (err .or. size(graphs) /= 1) then
     call invalid ('Bad graph name')
@@ -5616,7 +5616,8 @@ subroutine tao_python_cmd (input_str)
 
   u => point_to_uni(line, .true., err); if (err) return
   ix_branch = parse_branch(line, u, .false., err); if (err) return
-  branch => u%model%lat%branch(ix_branch)
+  lat => u%model%lat
+  branch => lat%branch(ix_branch)
 
   do i = 1, branch%n_ele_track
     ele => branch%ele(i)
@@ -5625,12 +5626,32 @@ subroutine tao_python_cmd (input_str)
     ix_shape_min = 1
     do
       call tao_ele_shape_info (u%ix_uni, ele, s%plot_page%lat_layout%ele_shape, shape, label_name, y1, y2, ix_shape_min)
+      if (.not. associated(shape)) exit
+      if (.not. shape%draw) cycle
       y1 = y1 * s%plot_page%lat_layout_shape_scale
       y2 = y2 * s%plot_page%lat_layout_shape_scale
+      nl=incr(nl); write (li(nl), '(2(i0, a), 2(es22.14, a), (i0, a), 2a, 2(es10.2, a), 4a)') ele%ix_branch, ';', ele%ix_ele, &
+                ';', ele%s_start, ';', ele%s, ';', shape%line_width, ';', trim(shape%shape), ';', &
+                y1, ';', y2, ';', trim(shape%color), ';', trim(label_name)
+    enddo
+  enddo
+
+  do i = lat%n_ele_track+1, lat%n_ele_max
+    ele => lat%ele(i)
+    if (ele%lord_status == multipass_lord$) cycle
+    branch2 => pointer_to_branch(ele)
+    if (branch2%ix_branch /= branch%ix_branch) cycle
+
+    ix_shape_min = 1
+    do
+      call tao_ele_shape_info (u%ix_uni, ele, s%plot_page%lat_layout%ele_shape, shape, label_name, y1, y2, ix_shape_min)
       if (.not. associated(shape)) exit
       if (.not. shape%draw) cycle
-      nl=incr(nl); write (li(nl), '(i0, 2(a, es22.14), (a, i0), 2a, 2(a, es10.2), 4a)') i, ';', ele%s_start, ';', ele%s, ';', &
-                shape%line_width, ';', trim(shape%shape), ';', y1, ';', y2, ';', trim(shape%color), ';', trim(label_name)
+      y1 = y1 * s%plot_page%lat_layout_shape_scale
+      y2 = y2 * s%plot_page%lat_layout_shape_scale
+      nl=incr(nl); write (li(nl), '(2(i0, a), 2(es22.14, a), (i0, a), 2a, 2(es10.2, a), 4a)') ele%ix_branch, ';', ele%ix_ele, &
+                ';', ele%s_start, ';', ele%s, ';', shape%line_width, ';', trim(shape%shape), ';', &
+                y1, ';', y2, ';', trim(shape%color), ';', trim(label_name)
     enddo
   enddo
 
diff --git a/tao/code/tao_struct.f90 b/tao/code/tao_struct.f90
index f0b10c1cb8..44850fcca6 100644
--- a/tao/code/tao_struct.f90
+++ b/tao/code/tao_struct.f90
@@ -1008,8 +1008,8 @@ module tao_struct
 ! Keep data and plotting separate since when optimizing will only do a calc if the data needs it
 
 type tao_universe_calc_struct
-  logical :: rad_int_for_data = .false.           ! Do the radiation integrals need to be computed for
   integer :: srdt_for_data = 0                    ! 0 = false, 1 = 1st order, 2 = 1st & 2nd order
+  logical :: rad_int_for_data = .false.           ! Do the radiation integrals need to be computed for
   logical :: rad_int_for_plotting = .false.       !   data or plotting?
   logical :: chrom_for_data = .false.             ! Does the chromaticity need to be computed for
   logical :: chrom_for_plotting = .false.         !   data or plotting?