module sc_track_mod use bmad use bsim_cesr_interface TYPE simple_coord_struct SEQUENCE REAL(rp) vec(1:6) END TYPE contains subroutine sc_track(ring, n_part_track, n_turn, n_turns_recalc, space_charge_on, current, coupling, ele_noise, obs_point_idx, in_file) use bmadz_interface use bsim_cesr_interface use gfit3d_mod use sc_track_mod_mpi use ele_noise_mod use beam_mod implicit none type (lat_struct), target :: ring type (lat_struct) :: ring_ideal type (coord_struct), allocatable :: co(:) !holds closed orbit type (coord_struct), allocatable :: orb(:) !holds trajectory of one turn on ring type (simple_coord_struct), allocatable :: start_coord(:) !allocated to one element for each test particle type (simple_coord_struct), allocatable :: end_coord(:) !allocated to one element for each test particle type (simple_coord_struct), allocatable :: sbuf(:) !send buffer type (normal_modes_struct) mode type (beam_init_struct) beam_init type(ele_noise_struct) :: ele_noise(:) type (beam_struct) :: beam ! foo - debug: type (ele_struct), pointer :: ele integer :: obs_point_idx(:) integer, allocatable :: param_lun_list(:), norm_lun_list(:) integer :: ix_start, ix_end, n_obs_points integer :: n_part_start, n_part_end real(rp) current real(rp) :: orb0(6) = 0 !real(rp) :: x_offset, y_offset real(rp) :: coupling ! artificially inflate the bunch initially integer ix, i, jx, kx, lx integer n_turns_recalc, n_turn, turns integer n_part_track, n_recalc real(rp) :: n_part character(200), optional:: in_file character(200) file_name, params_file, norm_file character(4) numnodes_str character(20) fmt1 character(200) :: temp_suffix integer track_state logical space_charge_on !mpi related variables integer mpierr integer, allocatable :: mpi_counts(:) integer, allocatable :: mpi_offsets(:) integer myrank integer numnodes integer remainder integer*8 time_start, time_finish integer my_time integer, allocatable :: perf_info(:) real, allocatable :: time_per_part(:) integer loc, buffstart, buffend logical spokesnode ! for describing particle distributions: real(rp) :: sigma_mat(1:6,1:6) = 0, sigma_beta_mat(1:6,1:6) = 0, normal(1:3) = 0 !###################### !get my rank and find out if I should talk call mympi_get_rank(myrank,numnodes,spokesnode) write(numnodes_str,'(I4)') numnodes if (present(in_file)) then file_name = in_file else file_name = 'space_charge_mc.' endif allocate(start_coord(1:n_part_track)) allocate(end_coord(1:n_part_track)) allocate(sbuf(1:n_part_track)) !allocate(co(0:ring%n_ele_track)) !allocate(orb(0:ring%n_ele_track)) ring%param%particle = positron$ call set_on_off (rfcavity$, ring, on$) call closed_orbit_calc(ring, co, 6) call closed_orbit_calc(ring, orb, 6) call set_on_off (rfcavity$, ring, off$) call lat_make_mat6(ring,-1,co) call twiss_at_start(ring) call twiss_propagate_all(ring) call set_on_off (rfcavity$, ring, on$) ! save the ideal ring, before adding any kicks due to magnet jitter, etc. ring_ideal = ring ! note: n_part != n_part_track; we can't feasibly track 10^10 particles n_part = current * ring%ele(ring%n_ele_track)%s / (e_charge * c_light) call radiation_integrals (ring, co, mode) call longitudinal_beta (ring, mode) !------------ ! generate initial gaussian dist & calculate its size ! note: No longer need to move bunch onto closed-orbit; ! init_beam_distribution does this automatically. if(spokesnode) then write(*,'(a,3es12.4)') "Beam emittances used for init_beam_distribution: ", mode%a%emittance, & sqrt(mode%b%emittance**2 + (mode%a%emittance * coupling)**2), mode%z%emittance endif ! initialize some things first: beam_init%distribution_type(:) = 'RAN_GAUSS' beam_init%n_particle = n_part_track beam_init%a_emit = mode%a%emittance beam_init%b_emit = sqrt(mode%b%emittance**2 + (mode%a%emittance * coupling)**2) beam_init%sig_z = mode%sig_z beam_init%sig_pz = mode%sigE_E beam_init%center(:) = co(0)%vec(:) call init_beam_distribution (ring%ele(0), ring%param, beam_init, beam) !------------- OPEN(45,FILE='load_balance.info',ACTION='WRITE',STATUS='REPLACE') !from here on out we will be working with start_coord, which is a simple_start_coord datatype start_coord(1:n_part_track)%vec(1) = beam%bunch(1)%particle(1:n_part_track)%vec(1) start_coord(1:n_part_track)%vec(2) = beam%bunch(1)%particle(1:n_part_track)%vec(2) start_coord(1:n_part_track)%vec(3) = beam%bunch(1)%particle(1:n_part_track)%vec(3) start_coord(1:n_part_track)%vec(4) = beam%bunch(1)%particle(1:n_part_track)%vec(4) start_coord(1:n_part_track)%vec(5) = beam%bunch(1)%particle(1:n_part_track)%vec(5) start_coord(1:n_part_track)%vec(6) = beam%bunch(1)%particle(1:n_part_track)%vec(6) orb0(1) = sum(start_coord(:)%vec(1))/size(start_coord(:)%vec(1)) orb0(3) = sum(start_coord(:)%vec(3))/size(start_coord(:)%vec(3)) !-------------------------------------------------------- ! prep output files: n_obs_points = 0 do ix = 1, size(obs_point_idx) if (obs_point_idx(ix) .gt. -1) n_obs_points = n_obs_points + 1 enddo if(spokesnode) then allocate(param_lun_list(1:n_obs_points), norm_lun_list(1:n_obs_points)) param_lun_list(:) = -1 ! initialize norm_lun_list(:) = -1 ! initialize do ix = 1, n_obs_points write(temp_suffix, '(a8,a)') '.params.', trim(ring%ele(obs_point_idx(ix))%name) call file_suffixer (file_name, params_file, trim(temp_suffix), .true.) param_lun_list(ix) = lunget() open(unit=param_lun_list(ix), file=params_file, status="replace") write(param_lun_list(ix), '(a,i0,2a)') '! Lab-frame output file for ring%ele(', & obs_point_idx(ix),') = ', trim(ring%ele(obs_point_idx(ix))%name) write(param_lun_list(ix), '(9a12)') '!turn', "", "sigma.x", "emit.x", "", "sigma.y", "emit.y", "sigma.z", "emit.z" write(temp_suffix, '(a6,a)') '.norm.', trim(ring%ele(obs_point_idx(ix))%name) norm_lun_list(ix) = lunget() call file_suffixer (file_name, norm_file, trim(temp_suffix), .true.) open(unit=norm_lun_list(ix), file=norm_file, status="replace") write(norm_lun_list(ix), '(a,i0,2a)') '! Normal-mode output file for ring%ele(', & obs_point_idx(ix),') = ', trim(ring%ele(obs_point_idx(ix))%name) write(norm_lun_list(ix), '(9a12)') '!turn', "", "sigma.x", "emit.a", "", "sigma.y", "emit.b", "sigma.z", "emit.c" enddo ! iterate over obs_point_idx endif !----------------------------------------------------- end_coord(:) = start_coord(:) ! initialize !determine which particles go to which nodes !for now this is a flat distribution. it will be replaced with something more intelligent !that is based on relative performance of the nodes. ALLOCATE(mpi_counts(1:numnodes)) ALLOCATE(mpi_offsets(1:numnodes)) ALLOCATE(perf_info(1:numnodes)) ALLOCATE(time_per_part(1:numnodes)) mpi_counts(:) = int(n_part_track/numnodes) remainder = mod(n_part_track,numnodes) if (remainder .gt. 0) then mpi_counts(1:remainder) = mpi_counts(1:remainder) + 1 endif mpi_offsets(1) = 0 do i=2,numnodes mpi_offsets(i) = mpi_offsets(i-1) + mpi_counts(i-1) enddo ! how many times do we need to recompute the space charge kick? if (n_turns_recalc .eq. 1) then write(*,*) 'n_turns_recalc must be greater than 1 (software limitation). Stopping here...' stop endif n_recalc = n_turn/n_turns_recalc turns = 0 ! loop over # of times we recompute space charge kick n_recalc_loop: do ix = 1, n_recalc !------------------------------------------ ! For one turn, have all nodes track all particles to observation points. ! Since this is only for one turn, we shouldn't lose too much time having ! all nodes do this in unison. obs_calc_loop: do lx = 1, n_obs_points ! determine start- and stop-points for tracking: if (lx .eq. 1) then ix_start = 0 else ix_start = obs_point_idx(lx-1) endif ix_end = obs_point_idx(lx) ! move all particles to the observation point: n_part_calc_loop: do jx = 1, n_part_track if (start_coord(jx)%vec(1) .eq. 999.) cycle n_part_calc_loop orb(ix_start)%vec = start_coord(jx)%vec !set initial particle coordinates call track_many(ring, orb, ix_start, ix_end, +1, 0, track_state) if (track_state /= moving_forward$) then orb(ix_end)%vec(1:6) = 999. ! don't need to 'exit' here; only doing one partial turn else end_coord(jx)%vec = orb(ix_end)%vec endif enddo n_part_calc_loop start_coord(:) = end_coord(:) ! get beam parameters at this observation point: call size_beam(start_coord, sigma_mat, sigma_beta_mat, normal) !compute the emittances based on normal-mode analysis on sigma_beta_mat: mode%a%emittance = normal(1) mode%b%emittance = normal(2) mode%z%emittance = normal(3) orb0(1) = sum(start_coord(:)%vec(1))/size(start_coord(:)%vec(1)) orb0(3) = sum(start_coord(:)%vec(3))/size(start_coord(:)%vec(3)) if(spokesnode) then write(*,'(3a,3es12.4)') "New beam emittance at ", trim(ring%ele(obs_point_idx(lx))%name), ": ", & mode%a%emittance, mode%b%emittance, mode%z%emittance ! write to .params file write(norm_lun_list(lx),'(i12,8es12.4)') ((ix-1)*n_turns_recalc+1), orb0(1), sqrt(sigma_mat(1,1)), normal(1), & orb0(3), sqrt(sigma_mat(3,3)), normal(2), sqrt(sigma_mat(5,5)), normal(3) write(param_lun_list(lx),'(i12,8es12.4)') ((ix-1)*n_turns_recalc+1), orb0(1), sqrt(sigma_mat(1,1)), & sqrt(sigma_beta_mat(1,1)*sigma_beta_mat(2,2) - sigma_beta_mat(1,2)**2), & orb0(3), sqrt(sigma_mat(3,3)), sqrt(sigma_beta_mat(3,3)*sigma_beta_mat(4,4) - sigma_beta_mat(3,4)**2), & sqrt(sigma_mat(5,5)), sqrt(sigma_mat(5,5)*sigma_mat(6,6) - sigma_mat(5,6)**2) endif enddo obs_calc_loop !------------------------------------------ ! move all particles from the final observation point to the end of the lattice: if (n_obs_points .gt. 0) then ix_start = obs_point_idx(n_obs_points) else ix_start = 0 endif ix_end = ring%n_ele_track n_part_end_loop: do jx = 1, n_part_track if (start_coord(jx)%vec(1) .eq. 999.) cycle n_part_end_loop orb(ix_start)%vec = start_coord(jx)%vec !set initial particle coordinates call track_many(ring, orb, ix_start, ix_end, +1, 0, track_state) if (track_state /= moving_forward$) then orb(ix_end)%vec(1:6) = 999. ! again, don't need to 'exit' since we're only tracking a partial turn else end_coord(jx)%vec = orb(ix_end)%vec endif enddo n_part_end_loop start_coord(:) = end_coord(:) ! End of computing emittances at observation points !------------------- if (space_charge_on) then call setup_ultra_rel_space_charge_calc(space_charge_on, ring, n_part, mode, co) endif call get_seconds(time_start) !------------------------------------------------------- !------------------------------------------------------- !(non-mpi code) n_part_loop: do jx = 1, n_part_track ! (mpi code): n_part_start, n_part_end computed for each node n_part_start = mpi_offsets(myrank+1) + 1 n_part_end = n_part_start + mpi_counts(myrank+1) - 1 n_part_loop: do jx = n_part_start, n_part_end ! if particle was previously lost, don't track it here if (start_coord(jx)%vec(1) .eq. 999.) cycle n_part_loop orb(0)%vec = start_coord(jx)%vec !set initial particle coordinates ! We finished tracking one turn for observing the beam; need to check if we ! are due for updating ele_noise: if (jx .eq. n_part_start) then turns = (ix-1)*n_turns_recalc + 1 call apply_ele_noise(ring, ring_ideal, ele_noise, turns) endif ! now track one particle for n_turns_recalc ! Note that we already tracked one turn during beam sizing; only need ! to track n_turns_recalc - 1 now. n_turns_loop: do kx = 2, n_turns_recalc call track_all(ring, orb, 0, track_state) if (track_state /= moving_forward$) then orb(ring%n_ele_track)%vec(1:6) = 999. exit else orb(0) = orb(ring%n_ele_track) endif ! check if any noise spectra are ready to be updated, and update as necessary. ! note: only update for first particle in subset; other particles will be on ! the same setting anyway. if (jx .eq. n_part_start) then turns = (ix-1)*n_turns_recalc + kx call apply_ele_noise(ring, ring_ideal, ele_noise, turns) endif enddo n_turns_loop end_coord(jx)%vec = orb(ring%n_ele_track)%vec enddo n_part_loop !------------------------------------------------------- !------------------------------------------------------- CALL get_seconds(time_finish) buffstart = mpi_offsets(myrank+1) + 1 buffend = buffstart + mpi_counts(myrank+1) - 1 sbuf(1:mpi_counts(myrank+1)) = end_coord(buffstart:buffend) call MPI_allgatherv(sbuf,mpi_counts(myrank+1),MPI_DATATYPE_COORD, & end_coord(1),mpi_counts,mpi_offsets,MPI_DATATYPE_COORD, & MPI_COMM_WORLD,mpierr) !- Begin Redistribute Load if(time_finish .lt. time_start) then time_finish = time_finish + 24*60*60 endif my_time = time_finish - time_start if (my_time .eq. 0) my_time = 1 ! duration is below integer resolution; set to minimum perf_info(myrank+1) = my_time call MPI_allgather(perf_info(myrank+1),1,MPI_INTEGER, & perf_info(1),1,MPI_INTEGER, & MPI_COMM_WORLD,mpierr) time_per_part = perf_info / (1.0*mpi_counts) fmt1 = '(A,I5,A,'//trim(numnodes_str)//'I5)' if(spokesnode) WRITE(45,fmt1) "Round ", ix, " allocation: ", mpi_counts if(spokesnode) WRITE(45,fmt1) "Round ", ix, " time per node: ", perf_info mpi_counts(1:numnodes) = 0.0 do jx = 1,n_part_track loc = minloc( (mpi_counts(1:numnodes)+1.0)*time_per_part(1:numnodes), 1 ) mpi_counts(loc) = mpi_counts(loc) + 1 enddo mpi_offsets(1) = 0 do i=2,numnodes mpi_offsets(i) = mpi_offsets(i-1) + mpi_counts(i-1) enddo !- End Redistribute Load start_coord(:) = end_coord(:) if(spokesnode) then write(*,*) "" write(*,'(a,i0,a,i0)') "Tracking turns: ", (ix*n_turns_recalc), "/", n_turn endif enddo n_recalc_loop do ix = 1, size(param_lun_list) if (param_lun_list(ix) .gt. 0) close(param_lun_list(ix)) if (norm_lun_list(ix) .gt. 0) close(norm_lun_list(ix)) enddo close(45) end subroutine sc_track !################################################# subroutine size_beam(particle_coords, sigma_mat, sigma_beta_mat, normal) use bmad_interface use bmadz_mod use bmadz_interface use scan_parameters use gfit3d_mod implicit none type (simple_coord_struct) :: particle_coords(:) type(coord_struct) :: orb type(coord_struct) :: adjusted_coords(1:size(particle_coords)) integer i,n_ok real(rp) :: sigma_mat(:,:), sigma_beta_mat(:,:), normal(:) n_ok = count(abs(particle_coords(:)%vec(1)) <= 900.) do i = 1, 6 orb%vec(i) = sum(particle_coords(:)%vec(i), mask=(abs(particle_coords(:)%vec(i)) <= 900.))/n_ok enddo ! move distribution of particles back to the origin for fitting n_ok = 0 do i = 1, size(particle_coords) if(abs(particle_coords(i)%vec(1)) <= 900.) then n_ok = n_ok + 1 adjusted_coords(n_ok)%vec(1) = particle_coords(i)%vec(1) - orb%vec(1) ! - x_offset adjusted_coords(n_ok)%vec(2) = particle_coords(i)%vec(2) - orb%vec(2) adjusted_coords(n_ok)%vec(3) = particle_coords(i)%vec(3) - orb%vec(3) ! - y_offset adjusted_coords(n_ok)%vec(4) = particle_coords(i)%vec(4) - orb%vec(4) adjusted_coords(n_ok)%vec(5) = particle_coords(i)%vec(5) - orb%vec(5) adjusted_coords(n_ok)%vec(6) = particle_coords(i)%vec(6) - orb%vec(6) endif enddo !if there are no more good particles or we have reached the designated number of turns, end the simulation if(n_ok.eq.0) then print *,"All particles lost, halting simulation" stop endif !call gfit3D(adjusted_coords(1:n_ok), parameters) call make_sigma_mat(adjusted_coords, sigma_mat, sigma_beta_mat) call normal_sigma_mat(sigma_mat, normal) end subroutine size_beam !-------------------------- subroutine get_seconds(seconds) IMPLICIT NONE INTEGER*8 seconds INTEGER*8 values(8) INTEGER*8 s,mn,h CALL date_and_time(values=values) !y = values(1) !mo = values(2) !d = values(3) !values(4) is hours from UTC and is not needed here h = values(5) mn = values(6) s = values(7) !ms = values(8) !get number of seconds since start of day seconds = h*3600 + mn*60 + s end subroutine get_seconds end module sc_track_mod