Merge pull request #544 from bmad-sim/devel/step21

Better error handling in tao_spin_polarization_calc.

bmad/modules/bmad_struct.f90

@@ -280,6 +280,7 @@ module bmad_struct
  real(rp) :: M_1turn(8,8) = 0 ! 1-turn matrix
  real(rp) :: M_ele(8,8) = 0 ! Transfer matrix through element.
  real(rp) :: sq_ele(0:3) = 0, sq_1turn(0:3) = 0
+ logical :: valid = .false.
 end type
 
 ! Note: Polarization is not 1 when the spin_polar struct represents an ensamble of spins.

tao/code/tao_data_and_eval_mod.f90

@@ -498,7 +498,7 @@ recursive subroutine tao_evaluate_a_datum (datum, u, tao_lat, datum_value, valid
 character(100) data_type, str
 character(:), allocatable :: e_str
 
-logical found, valid_value, err, taylor_is_complex, use_real_part, term_found
+logical found, valid_value, err, taylor_is_complex, use_real_part, term_found, ok
 logical particle_lost, exterminate, printit
 logical, allocatable :: good(:)
 
@@ -3027,8 +3027,8 @@ recursive subroutine tao_evaluate_a_datum (datum, u, tao_lat, datum_value, valid
  call tao_load_this_datum (value_vec, ele_ref, ele_start, ele, datum_value, valid_value, datum, branch, why_invalid)
 
  case ('spin.depolarization_rate', 'spin.polarization_rate', 'spin.polarization_limit')
- if (.not. tao_branch%spin_valid) call tao_spin_polarization_calc(branch, tao_branch)
- valid_value = tao_branch%spin_valid
+ if (.not. tao_branch%spin%valid) call tao_spin_polarization_calc(branch, tao_branch, err_flag = err)
+ valid_value = tao_branch%spin%valid
 
  if (.not. valid_value) then
  call tao_set_invalid (datum, 'ERROR IN SPIN POLARIZAITON CALC.', why_invalid, .false.)
@@ -3060,8 +3060,11 @@ recursive subroutine tao_evaluate_a_datum (datum, u, tao_lat, datum_value, valid
  return
  endif
 
- if (.not. tao_branch%spin_valid) call tao_spin_polarization_calc(branch, tao_branch)
- valid_value = tao_branch%spin_valid
+ call tao_spin_polarization_calc(branch, tao_branch)
+
+ valid_value = (tao_branch%spin_ele(ix_start)%valid .and. tao_branch%spin_ele(ix_ele)%valid)
+ if (ix_ref > -1) valid_value = (valid_value .and. tao_branch%spin_ele(ix_ref)%valid)
+
  if (.not. valid_value) then
  call tao_set_invalid (datum, 'ERROR IN SPIN POLARIZAITON CALC.', why_invalid, .false.)
  return

tao/code/tao_evaluate_element_parameters.f90

@@ -237,27 +237,33 @@ subroutine evaluate_this_parameter(ele, parameter, component, where, u, n_tot, v
  endif
 
 else
- select case (parameter)
- case ('spin_dn_dpz.x')
- if (.not. tao_branch%spin_valid) call tao_spin_polarization_calc(branch, tao_branch); err = .not. tao_branch%spin_valid
- values(n_tot) = tao_branch%spin_ele(ixe)%dn_dpz%vec(1)
- case ('spin_dn_dpz.y')
- if (.not. tao_branch%spin_valid) call tao_spin_polarization_calc(branch, tao_branch); err = .not. tao_branch%spin_valid
- values(n_tot) = tao_branch%spin_ele(ixe)%dn_dpz%vec(2)
- case ('spin_dn_dpz.z')
- if (.not. tao_branch%spin_valid) call tao_spin_polarization_calc(branch, tao_branch) ;err = .not. tao_branch%spin_valid
- values(n_tot) = tao_branch%spin_ele(ixe)%dn_dpz%vec(3)
- case ('spin_dn_dpz.amp')
- if (.not. tao_branch%spin_valid) call tao_spin_polarization_calc(branch, tao_branch); err = .not. tao_branch%spin_valid
- values(n_tot) = norm2(tao_branch%spin_ele(ixe)%dn_dpz%vec)
- case default
+ if (parameter(1:12) == 'spin_dn_dpz.') then
+ if (.not. tao_branch%spin_ele(ixe)%valid) call tao_spin_polarization_calc(branch, tao_branch)
+ err = tao_branch%spin_ele(ixe)%valid
+ if (err) return
+
+ select case (parameter)
+ case ('spin_dn_dpz.x')
+ values(n_tot) = tao_branch%spin_ele(ixe)%dn_dpz%vec(1)
+ case ('spin_dn_dpz.y')
+ values(n_tot) = tao_branch%spin_ele(ixe)%dn_dpz%vec(2)
+ case ('spin_dn_dpz.z')
+ values(n_tot) = tao_branch%spin_ele(ixe)%dn_dpz%vec(3)
+ case ('spin_dn_dpz.amp')
+ values(n_tot) = norm2(tao_branch%spin_ele(ixe)%dn_dpz%vec)
+ case default
+ err = .true.
+ return 
+ end select
+
+ else
  values(n_tot) = tao_param_value_at_s (parameter, branch%ele(ixe), tao_branch%orbit(ixe), err, print_err = print_err)
  if (err) then
  call pointer_to_attribute (branch%ele(ixe), parameter, .true., a_ptr, err, print_err)
  if (err) return
  values(n_tot) = value_of_all_ptr(a_ptr)
  endif
- end select
+ endif
 endif
 
 err = .false.

tao/code/tao_interface.f90

@@ -922,12 +922,13 @@ subroutine tao_spin_matrix_calc (datum, u, ele_ref, ele, excite_zero)
  character(*), optional :: excite_zero(3)
 end subroutine
 
-subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_kinetic)
+subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_kinetic, err_flag)
  import
  implicit none
  type (branch_struct), target :: branch
  type (tao_lattice_branch_struct), target :: tao_branch
  character(*), optional :: excite_zero(3), ignore_kinetic
+ logical, optional :: err_flag
 end subroutine
 
 function tao_spin_matrices_calc_needed (data_type, data_source) result (do_calc)

tao/code/tao_lattice_calc.f90

@@ -104,7 +104,9 @@ subroutine tao_lattice_calc (calc_ok, print_err)
 
  branch => tao_lat%lat%branch(ib)
  tao_branch => tao_lat%tao_branch(ib)
- tao_branch%spin_valid = .false.
+ tao_branch%spin_map_valid = .false.
+ tao_branch%spin%valid = .false.
+ tao_branch%spin_ele(:)%valid = .false.
 
  u%model%tao_branch(:)%plot_cache_valid = .false.
  u%design%tao_branch(:)%plot_cache_valid = .false.

tao/code/tao_python_cmd.f90

@@ -7061,7 +7061,7 @@ subroutine tao_python_cmd (input_str)
  tao_branch => tao_lat%tao_branch(ix_branch)
  branch => tao_lat%lat%branch(ix_branch)
 
- if (.not. tao_branch%spin_valid) call tao_spin_polarization_calc (branch, tao_branch)
+ if (.not. tao_branch%spin%valid) call tao_spin_polarization_calc (branch, tao_branch)
 
  z = anomalous_moment_of(branch%param%particle) * branch%ele(0)%value(e_tot$) / mass_of(branch%param%particle)
  nl=incr(nl); write (li(nl), rmt) 'anom_moment_times_gamma;REAL;F;', z

tao/code/tao_show_this.f90

@@ -4419,7 +4419,7 @@ subroutine tao_show_this (what, result_id, lines, nl)
  call tao_spin_polarization_calc (branch, tao_branch, excite_zero, veto)
  nl=nl+1; lines(nl) = ''
  nl=nl+1; write (lines(nl), '(a, es18.7)') 'spin_tune: ', tao_branch%spin%tune / twopi
- if (tao_branch%spin_valid) then
+ if (tao_branch%spin%valid) then
  if (tao_branch%spin%pol_rate_bks == 0) then
  nl=nl+1; lines(nl) = 'No bends or other radiation producing lattice elements detected!'
  else

tao/code/tao_spin_polarization_calc.f90

@@ -1,5 +1,5 @@
 !+
-! Subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_kinetic)
+! Subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_kinetic, err_flag)
 !
 ! Routine to calculate the spin equalibrium polarization in a ring along with the polarization rate and
 ! the depolarization rate due to emission of synchrotron radiation photons.
@@ -16,9 +16,10 @@
 ! tao_branch -- tao_lattice_branch_struct: Calculated is:
 ! %spin%dn_dpz(:) 
 ! %spin
+! err_flag -- logical, optional: Set True if there is an error (EG: invalid orbit). False otherwise.
 !-
 
-subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_kinetic)
+subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_kinetic, err_flag)
 
 use tao_data_and_eval_mod, dummy => tao_spin_polarization_calc
 use radiation_mod
@@ -46,11 +47,14 @@ subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_k
 integer ix1, ix2, n_loc
 integer i, j, k, kk, n, p, ie
 
+logical, optional :: err_flag
 logical valid_value, err
 character(*), optional :: excite_zero(3), ignore_kinetic
 
 !
 
+if (present(err_flag)) err_flag = .true.
+
 g_tol = 1e-4_rp * branch%param%g1_integral / branch%param%total_length
 g2_tol = 1e-4_rp * branch%param%g2_integral / branch%param%total_length
 g3_tol = 1e-4_rp * branch%param%g3_integral / branch%param%total_length
@@ -66,6 +70,14 @@ subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_k
 
 !
 
+orbit => tao_branch%orbit
+if (.not. allocated(tao_branch%spin_ele)) allocate (tao_branch%spin_ele(0:branch%n_ele_track))
+tao_branch%spin_ele(:)%valid = .false.
+tao_branch%spin%valid = .false.
+if (branch%param%geometry == closed$ .and. orbit(branch%n_ele_track)%state /= alive$) return
+
+!
+
 call tao_spin_tracking_turn_on()
 
 branch%ele%logic = .true. ! False => Ignore kinetic term
@@ -77,14 +89,13 @@ subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_k
  enddo
 endif
 
-orbit => tao_branch%orbit
-if (.not. allocated(tao_branch%spin_ele)) allocate (tao_branch%spin_ele(0:branch%n_ele_track))
-tao_branch%spin_valid = .true.
-
-call spin_concat_linear_maps (err, q_1turn, branch, 0, branch%n_ele_track, q_ele, orbit, excite_zero)
-if (err) return
+if (.not. tao_branch%spin_map_valid) then
+ call spin_concat_linear_maps (err, q_1turn, branch, 0, branch%n_ele_track, q_ele, orbit, excite_zero)
+ if (err) return
+ tao_branch%spin_map_valid = .true.
+endif
 
-if (branch%param%geometry == closed$) then
+if (branch%param%geometry == closed$) then 
  tao_branch%spin%tune = 2.0_rp * atan2(norm2(q_1turn%spin_q(1:3,0)), q_1turn%spin_q(0,0))
  n0 = 0
  s_vec = [0.0_rp, 0.0_rp, 1.0_rp]
@@ -97,6 +108,7 @@ subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_k
  dn_dpz = branch%ele(0)%value(spin_dn_dpz_x$:spin_dn_dpz_z$)
  partial = 0 ! Not sure this is computable.
  partial2 = 0
+ tao_branch%spin_ele(0)%valid = .true.
  tao_branch%spin_ele(0)%dn_dpz%vec = dn_dpz
  tao_branch%spin_ele(0)%dn_dpz%partial = partial
  tao_branch%spin_ele(0)%dn_dpz%partial2 = partial2
@@ -119,6 +131,7 @@ subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_k
  q_1turn = q_ele(ie) * q_1turn * map1_inverse(q_ele(ie))
  dn_dpz = spin_dn_dpz_from_qmap(q_1turn%orb_mat, q_1turn%spin_q, partial, partial2, err)
  if (err) exit
+ tao_branch%spin_ele(ie)%valid = .true.
  tao_branch%spin_ele(ie)%dn_dpz%vec = dn_dpz
  tao_branch%spin_ele(ie)%dn_dpz%partial = partial
  tao_branch%spin_ele(ie)%dn_dpz%partial2 = partial2
@@ -132,11 +145,17 @@ subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_k
  n0 = quat_rotate(q_ele(ie)%spin_q(:,0), n0)
  partial = 0
  partial2 = 0
+ tao_branch%spin_ele(ie)%valid = .true.
  tao_branch%spin_ele(ie)%dn_dpz%vec = dn_dpz
  tao_branch%spin_ele(ie)%dn_dpz%partial = partial
  tao_branch%spin_ele(ie)%dn_dpz%partial2 = partial2
  endif
 
+ if (orbit(ie)%state /= alive$) then
+ err = .true.
+ exit
+ endif
+
  del_p = 1 + orbit(ie)%vec(6)
  s_vec(1:2) = [orbit(ie)%vec(2)/del_p, orbit(ie)%vec(4)/del_p]
  s_vec(3) = sqrt(1.0_rp - s_vec(1)**2 - s_vec(2)**2)
@@ -199,6 +218,8 @@ subroutine tao_spin_polarization_calc (branch, tao_branch, excite_zero, ignore_k
  tao_branch%spin%integral_bdn = integral_bdn
  tao_branch%spin%integral_1ns = integral_1ns
  tao_branch%spin%integral_dn2 = integral_dn2
+ tao_branch%spin%valid = .true.
+ if (present(err_flag)) err_flag = .false.
 endif
 
 !--------------------------------------

tao/code/tao_struct.f90

@@ -871,9 +871,10 @@ module tao_struct
 
 type tao_spin_ele_struct
  type (tao_spin_dn_dpz_struct) dn_dpz
- real(rp) orb_eigen_val(6)
- real(rp) orb_eigen_vec(6,6) ! (j,:) is j^th vector
- real(rp) spin_eigen_vec(6,3) ! (j,:) is j^th vector
+ real(rp) :: orb_eigen_val(6) = 0
+ real(rp) :: orb_eigen_vec(6,6) = 0 ! (j,:) is j^th vector
+ real(rp) :: spin_eigen_vec(6,3) = 0 ! (j,:) is j^th vector
+ logical :: valid = .false.
 end type
 
 type tao_spin_polarization_struct
@@ -890,6 +891,7 @@ module tao_struct
  real(rp) :: integral_bdn = real_garbage$ ! Integral of g^3 * b_hat * dn/ddelta
  real(rp) :: integral_1ns = real_garbage$ ! Integral of g^3 (1 - 2(n * s_hat)/9)
  real(rp) :: integral_dn2 = real_garbage$ ! Integral of g^3 * 11 (dn/ddelta)^2 / 9
+ logical :: valid = .false.
 end type
 
 ! For caching lattice calculations associated with plotting.
@@ -929,7 +931,7 @@ module tao_struct
  integer :: n_hterms = 0 ! Number of distinct res driving terms to evaluate.
  logical has_open_match_element
  logical :: plot_cache_valid = .false. ! Valid plotting data cache?
- logical :: spin_valid = .false.
+ logical :: spin_map_valid = .false.
 end type
 
 ! Structure to hold a single lat_struct (model, base, or design) in

tao/version/tao_version_mod.f90

@@ -6,5 +6,5 @@
 !-
 
 module tao_version_mod
-character(*), parameter :: tao_version_date = "2023/10/02 02:28:01"
+character(*), parameter :: tao_version_date = "2023/10/05 02:03:56"
 end module