More set_twiss and dispersion devel. (#990)

* More set_twiss devel.
* More dispersion devel.

bmad/code/reverse_lat.f90

@@ -71,11 +71,7 @@ subroutine reverse_lat (lat_in, lat_rev, track_antiparticle)
   ele_rev%s = 0
   ele_rev%floor = branch_in%ele(nt)%floor
 
-  call transfer_twiss (branch_in%ele(nt), ele_rev)
-  ele_rev%a%alpha = -ele_rev%a%alpha
-  ele_rev%b%alpha = -ele_rev%b%alpha
-  ele_rev%c_mat(1,2) = -ele_rev%c_mat(1,2)
-  ele_rev%c_mat(2,1) = -ele_rev%c_mat(2,1)
+  call transfer_twiss (branch_in%ele(nt), ele_rev, .true.)
 
   ! All other elements
 

bmad/code/set_twiss.f90

@@ -1,20 +1,27 @@
 !+
-! Subroutine set_twiss(branch, twiss_ele, ix_ele, err_flag, print_err)
+! Subroutine set_twiss(branch, twiss_ele, ix_ele, match_deta_ds, err_flag, print_err)
 !
-! Routine to set the beginning Twiss of a lattice branch so that the Twiss parameters at branch%ele(ix_ele) are
-! the same as the Twiss parameters in twiss_ele.
+! For an open lattice branch, set the beginning Twiss, coupling, and dispersion of the branch so that the 
+! Twiss parameters at the "target" element branch%ele(ix_ele) are the same as the Twiss parameters in twiss_ele.
 !
 ! It is OK for twiss_ele to be an element in the branch (this routine makes a copy to be safe).
 !
+! Since the the reference orbit at twiss_ele may be different from the reference orbit at the target ele,
+! it is not possible, in general, to match both etap and deta_ds. Which is matched is set by match_deta_ds.
+!
+! Additionally, It is not possible to simultaneously match both (x,y) deta_ds and (a,b) normal mode deta_ds.
+! This routine matches (x,y) deta_ds only (when match_deta_ds = True).
+!
 ! Input:
-!   branch      -- branch_struct: Branch to modify.
-!   twiss_ele   -- ele_struct: Element with desired Twiss parameters.
-!   ix_ele      -- integer: Match branch%ele(ix_ele) Twiss to twiss_ele.
-!   err_flag    -- logical: Set True if there is an error. False otherwise.
-!   print_err   -- logical, optional: Print an error message if there is an error? Default is True.
+!   branch        -- branch_struct: Branch to modify.
+!   twiss_ele     -- ele_struct: Element with desired Twiss parameters.
+!   ix_ele        -- integer: Match branch%ele(ix_ele) Twiss to twiss_ele.
+!   match_deta_ds -- logical: If True, match deta_ds. If False, match etap.
+!   err_flag      -- logical: Set True if there is an error. False otherwise.
+!   print_err     -- logical, optional: Print an error message if there is an error? Default is True.
 !-
 
-subroutine set_twiss(branch, twiss_ele, ix_ele, err_flag, print_err)
+subroutine set_twiss(branch, twiss_ele, ix_ele, match_deta_ds, err_flag, print_err)
 
 use bmad, dummy => set_twiss
 
@@ -24,12 +31,12 @@ subroutine set_twiss(branch, twiss_ele, ix_ele, err_flag, print_err)
 type (ele_struct) twiss_ele, ele0, ele1
 type (ele_struct), pointer :: ele
 
-real(rp) xmat(6,6), xvec(6)
+real(rp) xmat(6,6), xvec(6), rel_p
 
 integer ix_ele
 integer ie
 
-logical err_flag
+logical match_deta_ds, err_flag
 logical, optional :: print_err
 
 character(*), parameter :: r_name = 'set_twiss'
@@ -52,9 +59,18 @@ subroutine set_twiss(branch, twiss_ele, ix_ele, err_flag, print_err)
   return
 endif
 
-!
+! Note: twiss_propagate1 treats deta_ds as a dependent variable to etap
 
 call transfer_twiss(twiss_ele, ele0)
+ele0%map_ref_orb_out = twiss_ele%map_ref_orb_out
+
+if (match_deta_ds) then
+  ele => branch%ele(ix_ele)
+  rel_p = 1.0_rp + ele%map_ref_orb_out%vec(6)
+  ele0%x%etap = ele0%x%deta_ds * rel_p + ele%map_ref_orb_out%vec(2) / rel_p
+  ele0%y%etap = ele0%y%deta_ds * rel_p + ele%map_ref_orb_out%vec(4) / rel_p
+endif
+
 do ie = ix_ele, 1, -1
   ele => branch%ele(ie)
   call mat_inverse(ele%mat6, ele1%mat6)

bmad/code/twiss_at_start.f90

@@ -37,7 +37,7 @@ subroutine twiss_at_start (lat, status, ix_branch, type_out)
 type (ele_struct), target :: drift_ele
 type (branch_struct), pointer :: branch
 
-real(rp) eta_vec(4), t0_4(4,4), mat6(6,6), map0(4), m56
+real(rp) eta_vec(5), t0_5(5,5), mat6(6,6), map0(5), m56
 
 integer, optional, intent(in) :: ix_branch
 integer, optional, intent(out) ::status
@@ -51,7 +51,7 @@ subroutine twiss_at_start (lat, status, ix_branch, type_out)
 
 ! Init one turn. T0 is the transverse part of the matrix
 
-call mat_make_unit (t0_4)       ! form unit matrix
+call mat_make_unit (t0_5)       ! form unit matrix
 eta_vec = 0
 map0 = 0
 m56 = 0
@@ -61,7 +61,7 @@ subroutine twiss_at_start (lat, status, ix_branch, type_out)
 
 ! Propagate the transfer map around branch. 
 ! Since the RF is taken to be off we use a trick so we only have to multiply
-! 4x4 matrices.
+! 5x5 matrices.
 
 if (debug) then
   iu = lunget()
@@ -80,23 +80,22 @@ subroutine twiss_at_start (lat, status, ix_branch, type_out)
     ele => drift_ele
   endif
 
-  m56 = m56 + ele%mat6(5,6) + dot_product(ele%mat6(5,1:4), eta_vec)
-  eta_vec = matmul (ele%mat6(1:4,1:4), eta_vec) + ele%mat6(1:4,6)
-  map0 = matmul (ele%mat6(1:4,1:4), map0) + ele%vec0(1:4)
-  t0_4 = matmul (ele%mat6(1:4,1:4), t0_4)
+  eta_vec = matmul (ele%mat6(1:5,1:5), eta_vec) + ele%mat6(1:5,6)
+  map0 = matmul (ele%mat6(1:5,1:5), map0) + ele%vec0(1:5)
+  t0_5 = matmul (ele%mat6(1:5,1:5), t0_5)
   if (debug) then
     write (iu, *) '!------------------------------------', n
     call type_ele (ele, .false., 0, .false., 0, lines = lines, n_lines = n_lines)
     do i = 1, n_lines
       write (iu, '(a)') trim(lines(i))
     enddo
-    write (iu, *) 'Symplectic Check:', mat_symp_error(t0_4)
-    do i = 1, 4
-      write (iu, '(5x, 4f14.8, 6x, 4f14.8)') t0_4(i,:), ele%mat6(i,1:4)
+    write (iu, *) 'Symplectic Check:', mat_symp_error(t0_5(1:4,1:4))
+    do i = 1, 5
+      write (iu, '(5x, 5f14.8, 6x, 5f14.8)') t0_5(i,:), ele%mat6(i,1:5)
     enddo
     write (iu, '(a, 6es18.10)') 'Map_ref_orb:', ele%map_ref_orb_in%vec
-    write (iu, '(a, 4es20.12)') 'Eta: ', eta_vec
-    write (iu, '(a, 4es20.12)') 'Map0:', map0
+    write (iu, '(a, 5es20.12)') 'Eta: ', eta_vec
+    write (iu, '(a, 5es20.12)') 'Map0:', map0
   endif
 enddo
 
@@ -105,10 +104,8 @@ subroutine twiss_at_start (lat, status, ix_branch, type_out)
 ! Put 1-turn matrix into branch%param%t1_no_RF
 
 call mat_make_unit (mat6)
-mat6(1:4,1:4) = t0_4
-
-call mat6_dispersion (eta_vec, mat6) ! dispersion to %mat6
-mat6(5,6) = m56
+mat6(1:5,1:5) = t0_5
+mat6(1:5, 6) = eta_vec
 branch%param%t1_no_RF = mat6
 
 ! Compute twiss parameters
@@ -125,45 +122,4 @@ subroutine twiss_at_start (lat, status, ix_branch, type_out)
 branch%ele(0)%a%phi = 0
 branch%ele(0)%b%phi = 0
 
-contains
-
-!--------------------------------------------------------------------------
-!--------------------------------------------------------------------------
-!--------------------------------------------------------------------------
-!+
-! Subroutine mat6_dispersion (m_i6, mat6)
-!
-! Subroutine to set the mat6(5, 1:4) terms given the vector mat6(1:4, 6)
-! which is a measure of the dispersion.
-!
-! Input:
-!   m_i6(4)   -- Real(rp): mat6(1:4, 6) components.
-!   mat6(6,6) -- Real(rp): Matrix with 4x4 x-y submatrix already made.
-!
-! Output:
-!   mat6(6,6) -- Real(rp): mat6(5, 1:4) components set. 
-!-
-
-subroutine mat6_dispersion (m_i6, mat6)
-
-implicit none
-
-real(rp), intent(inout) :: mat6(:,:)
-real(rp), intent(in) :: m_i6(:)
-
-real(rp) vec4(4)
-
-!
-
-mat6(1:4, 6) = m_i6(1:4)
-
-vec4(1) = -m_i6(2)
-vec4(2) =  m_i6(1)
-vec4(3) = -m_i6(4)
-vec4(4) =  m_i6(3)
-
-mat6(5,1:4) = matmul (vec4, mat6(1:4,1:4))
-
-end subroutine mat6_dispersion
-
 end subroutine

bmad/code/twiss_from_mat6.f90

@@ -47,12 +47,7 @@ subroutine twiss_from_mat6 (mat6, orb0, ele, stable, growth_rate, status, type_o
 
 stable = .false.
 m6 = mat6
-
-if (bmad_private%normalize_twiss) then
-  rel_p = 1 + orb0(6)
-  m6(1:5:2, 2:6:2) = m6(1:5:2, 2:6:2) * rel_p
-  m6(2:6:2, 1:5:2) = m6(2:6:2, 1:5:2) / rel_p
-endif
+rel_p = 1 + orb0(6)
 
 mat4 = m6(1:4, 1:4)
 
@@ -110,43 +105,31 @@ subroutine twiss_from_mat6 (mat6, orb0, ele, stable, growth_rate, status, type_o
 forall (i = 1:4) mat4(i,i) = mat4(i,i) + 1
 call mat_inverse (mat4, mat4)
 
-vec(1) = m6(1,6) + (m6(1,2) * orb0(2) + m6(1,4) * orb0(4))
-vec(2) = m6(2,6) + (m6(2,2) * orb0(2) + m6(2,4) * orb0(4) - orb0(2))
-vec(3) = m6(3,6) + (m6(3,2) * orb0(2) + m6(3,4) * orb0(4))
-vec(4) = m6(4,6) + (m6(4,2) * orb0(2) + m6(4,4) * orb0(4) - orb0(4))
-
-eta_vec = matmul(mat4, vec)
+eta_vec = matmul(mat4, mat6(1:4,6))
 
 ele%x%eta     = eta_vec(1)
-ele%x%deta_ds = eta_vec(2)
+ele%x%etap    = eta_vec(2)
+ele%x%deta_ds = ele%x%etap / rel_p - orb0(2) / rel_p**2 
+
 ele%y%eta     = eta_vec(3)
-ele%y%deta_ds = eta_vec(4)
+ele%y%etap    = eta_vec(4)
+ele%y%deta_ds = eta_vec(4) / rel_p - orb0(4) / rel_p*2
 
 ele%z%eta     = 0
+ele%z%etap    = 1
 ele%z%deta_ds = 1
 
-eta_vec = matmul(mat_symp_conj(v), eta_vec)
-ele%a%eta     = eta_vec(1)
-ele%a%deta_ds = eta_vec(2)
-ele%b%eta     = eta_vec(3)
-ele%b%deta_ds = eta_vec(4)
-
 !
 
-eta_vec = matmul(mat4, mat6(1:4,6))
-
-ele%x%eta   = eta_vec(1)
-ele%x%etap  = eta_vec(2)
-ele%y%eta   = eta_vec(3)
-ele%y%etap  = eta_vec(4)
+eta_vec = matmul(mat_symp_conj(v), eta_vec)
+vec = matmul(mat_symp_conj(v), orb0(1:4))
 
-ele%z%eta     = 0
-ele%z%etap    = 1
+ele%a%eta     = eta_vec(1)
+ele%a%etap    = eta_vec(2)
+ele%a%deta_ds = ele%a%etap / rel_p - vec(2) / rel_p**2 
 
-eta_vec = matmul(mat_symp_conj(v), eta_vec)
-ele%a%eta  = eta_vec(1)
-ele%a%etap = eta_vec(2)
-ele%b%eta  = eta_vec(3)
-ele%b%etap = eta_vec(4)
+ele%b%eta     = eta_vec(3)
+ele%b%etap    = eta_vec(4)
+ele%b%deta_ds = eta_vec(4) / rel_p - vec(4) / rel_p*2
 
 end subroutine

bmad/code/twiss_propagate1.f90

@@ -25,14 +25,15 @@ subroutine twiss_propagate1 (ele1, ele2, err_flag)
 type (ele_struct), target :: ele1, ele2
 type (twiss_struct) twiss_a
 type (lat_param_struct) param
-type (coord_struct), pointer :: orb, orb_out
-integer key2
+type (coord_struct), pointer :: orb_in, orb_out
+
+integer key2, geometry
 
 real(rp) :: mat6(6,6)
 real(rp) v_mat(4,4), v_inv_mat(4,4), det, mat2_a(2,2), mat2_b(2,2)
 real(rp) big_M(2,2), small_m(2,2), big_N(2,2), small_n(2,2)
 real(rp) c_conj_mat(2,2), E_inv_mat(2,2), F_inv_mat(2,2)
-real(rp) mat2(2,2), eta1_vec(6), eta_vec(6), vec(6), dpz2_dpz1, rel_p1, rel_p2
+real(rp) mat2(2,2), eta1_vec(6), eta_vec(6), vec(6), dpz2_dpz1, rel_p1, rel_p21, rel_p2
 real(rp) det_factor, deriv_rel, gamma2_c, df
 
 logical error
@@ -65,6 +66,9 @@ subroutine twiss_propagate1 (ele1, ele2, err_flag)
   if (error) return
 endif
 
+geometry = open$
+if (associated(ele1%branch)) geometry = ele1%branch%param%geometry
+
 !---------------------------------------------------------------------
 ! init
 
@@ -98,20 +102,12 @@ subroutine twiss_propagate1 (ele1, ele2, err_flag)
 
 !
 
-orb  => ele2%map_ref_orb_in
+orb_in  => ele2%map_ref_orb_in
 orb_out => ele2%map_ref_orb_out
-rel_p1 = 1 + orb%vec(6)               ! reference energy 
+rel_p1 = 1 + orb_in%vec(6)    ! Map reference momentum
 rel_p2 = 1 + orb_out%vec(6)
 
 mat6 = ele2%mat6
-if (ele2%key /= e_gun$) then   ! Energy change normalization is not applied to an e-gun
-  if (bmad_private%normalize_twiss) then
-    mat6(:, 2:6:2) = mat6(:, 2:6:2) * rel_p1
-    mat6(2:6:2, :) = mat6(2:6:2, :) / rel_p2
-  endif
-  rel_p2 = rel_p2 / rel_p1
-  rel_p1 = 1
-endif
 
 !---------------------------------------------------------------------
 ! det_factor is a renormalization factor to handle non-symplectic errors.
@@ -208,34 +204,27 @@ subroutine twiss_propagate1 (ele1, ele2, err_flag)
 ! p_z2 is p_z at end of ele2 assuming p_z = 1 at end of ele1.
 ! This is just 1.0 (except for RF cavities).
 
-eta1_vec = [ele1%x%eta, ele1%x%etap, ele1%y%eta, ele1%y%etap, ele1%z%eta, 1.0_rp]
-
 ! For a circular ring, defining the dependence of z on pz is problematical.
 ! With the RF off, z is not periodic so dz/dpz is dependent upon turn number.
 ! With the RF on, dz/dpz, along with the other dispersion components, is not well defined.
-! This being the case, eta1_vec(5) is just treated as zero for an rfcavity.
+! This being the case, eta1_vec(5) is just treated as zero for a closed branch
 
-if (key2 == rfcavity$) eta1_vec(5) = 0
+eta1_vec = [ele1%x%eta, ele1%x%etap, ele1%y%eta, ele1%y%etap, ele1%z%eta, 1.0_rp]
+if (geometry == closed$) eta1_vec(5) = 0
 
 ! Must avoid 0/0 divide at zero reference momentum. 
-! If rel_p1 = 0 then total momentum is zero and orb%vec(2) and orb%vec(4) must be zero.
+! If rel_p1 = 0 then total momentum is zero and orb_in%vec(2) and orb_in%vec(4) must be zero.
 
 ! Also for elements with a static electric field, pz should include the potential energy and so the 
 ! mat6(6,:) terms should be all zero. However Bmad is not using proper canonical coords so the
 ! mat6(6,:) terms are forced to zero.
 
-if (rel_p1 == 0) then
-  dpz2_dpz1 = dot_product(mat6(6,:), eta1_vec) 
+if (rel_p1 == 0 .or. key2 == rfcavity$ .or. key2 == lcavity$) then
+  dpz2_dpz1 = dot_product(mat6(6,:), eta1_vec)
 else
-  if (key2 == rfcavity$ .or. key2 == lcavity$) then
-    dpz2_dpz1 = dot_product(mat6(6,:), eta1_vec) 
-  else
-    dpz2_dpz1 = 1 / rel_p2
-  endif
+  dpz2_dpz1 = rel_p1 / rel_p2
 endif
 
-!
-
 eta_vec(1:5) = matmul (mat6(1:5,:), eta1_vec) / dpz2_dpz1
 
 ele2%x%eta     = eta_vec(1)
@@ -246,9 +235,14 @@ subroutine twiss_propagate1 (ele1, ele2, err_flag)
 ele2%y%etap    = eta_vec(4)
 ele2%y%deta_ds = eta_vec(4) / rel_p2 - orb_out%vec(4) / rel_p2**2
 
-ele2%z%eta     = eta_vec(5)
-ele2%z%etap    = ele1%z%etap * dpz2_dpz1
-ele2%z%deta_ds = ele1%z%deta_ds * dpz2_dpz1
+if (geometry == closed$) then
+  ele2%z%eta     = 0
+else
+  ele2%z%eta     = eta_vec(5)
+endif
+
+ele2%z%etap    = 1
+ele2%z%deta_ds = 1
 
 call make_v_mats (ele2, v_mat, v_inv_mat)
 eta_vec(1:4) = matmul (v_inv_mat, eta_vec(1:4))