module m_snapshot_manager !! @brief Manages the creation of simulation snapshots for post-processing !! and visualisation. !! !! @details This module is responsible for periodically writing simulation !! data to files intended for analysis and visualisation !! Unlike checkpoints, which are always full-resolution for exact restarts, !! snapshots can be strided to reduce file size. use m_mpi, only: MPI_COMM_WORLD, MPI_Comm_rank, MPI_Allreduce, & MPI_IN_PLACE, MPI_MIN, MPI_SUCCESS, MPI_Bcast, & MPI_LOGICAL, MPI_Barrier use m_common, only: dp, i8, MPI_X3D2_DP, DIR_C, VERT, get_argument use m_field, only: field_t use m_solver, only: solver_t use m_io_session, only: writer_session_t use m_config, only: checkpoint_config_t, has_output_field use m_io_field_utils, only: field_buffer_map_t, field_ptr_t, & setup_field_arrays, cleanup_field_arrays, & stride_data_to_buffer, get_output_dimensions, & prepare_field_buffers, cleanup_field_buffers, & write_single_field_to_buffer, get_field_ptr use m_logger, only: log_info, abort_run, str implicit none private public :: snapshot_manager_t type :: snapshot_manager_t type(checkpoint_config_t) :: config integer, dimension(3) :: output_stride = [1, 1, 1] type(field_buffer_map_t), allocatable :: field_buffers(:) integer(i8), dimension(3) :: last_shape_dims = 0 integer, dimension(3) :: last_stride_factors = 0 integer(i8), dimension(3) :: last_output_shape = 0 character(len=:), allocatable :: vtk_xml logical :: is_snapshot_file_open = .false. type(writer_session_t) :: snapshot_writer logical :: convert_to_sp = .false. !! Flag for single precision snapshots !! Steps of an existing file kept on restart; negative starts a new file integer :: append_after_steps = -1 !! Iteration the run restarted from, and the layout of the existing file integer :: restart_iter = 0 character(len=:), allocatable :: restart_vtk_xml contains procedure :: init procedure :: handle_restart procedure :: handle_snapshot_step procedure :: finalise procedure, private :: write_snapshot procedure, private :: write_fields procedure, private :: cleanup_output_buffers procedure, private :: generate_vtk_xml procedure, private :: open_snapshot_file procedure, private :: retire_snapshot_file procedure, private :: close_snapshot_file end type snapshot_manager_t contains subroutine init(self, comm) !! Initialise snapshot manager class(snapshot_manager_t), intent(inout) :: self integer, intent(in) :: comm self%config = checkpoint_config_t() call self%config%read(nml_file=get_argument(1)) if (self%config%snapshot_freq > 0) then call configure_output(self, comm) end if end subroutine init subroutine configure_output(self, comm) !! Configure snapshot output settings use m_io_backend, only: get_default_backend, IO_BACKEND_DUMMY class(snapshot_manager_t), intent(inout) :: self integer, intent(in) :: comm self%output_stride = self%config%output_stride self%convert_to_sp = self%config%snapshot_sp if (get_default_backend() /= IO_BACKEND_DUMMY) then call log_info('Snapshot frequency: '//str(self%config%snapshot_freq)) call log_info('Snapshot prefix: '//trim(self%config%snapshot_prefix)) call log_info('Output stride: '//str(self%output_stride(1))//' ' & //str(self%output_stride(2))//' ' & //str(self%output_stride(3))) call log_info('Snapshot precision: ' & //merge('Single', 'Double', self%config%snapshot_sp)) if (any(self%config%output_fields /= '')) then call log_info('Additional output fields: ' & //trim(adjustl(join_output_fields( & self%config%output_fields)))) end if end if end subroutine configure_output subroutine handle_restart(self, restart_iter, dt, comm) !! Prepare to continue the snapshot file of the run being restarted. !! Steps up to the restart time are kept and later ones (written past !! the checkpoint) are dropped, so the restarted run appends without !! duplicating them, whatever snapshot frequency wrote the file. !! The file's layout is kept so the first snapshot can check it still !! matches the output settings. use m_io_backend, only: count_steps_up_to, read_string_attribute class(snapshot_manager_t), intent(inout) :: self integer, intent(in) :: restart_iter real(dp), intent(in) :: dt integer, intent(in) :: comm character(len=256) :: filename real(dp) :: limit logical :: file_exists integer :: myrank, ierr if (self%config%snapshot_freq <= 0) return self%restart_iter = restart_iter call MPI_Comm_rank(comm, myrank, ierr) write (filename, '(A,A)') trim(self%config%snapshot_prefix), '.bp' if (myrank == 0) inquire (file=trim(filename), exist=file_exists) call MPI_Bcast(file_exists, 1, MPI_LOGICAL, 0, comm, ierr) if (.not. file_exists) return ! Half a time step of slack absorbs rounding in the stored times limit = (real(restart_iter, dp) + 0.5_dp)*dt self%append_after_steps = count_steps_up_to(filename, "time", limit, comm) self%restart_vtk_xml = read_string_attribute(filename, "vtk.xml", comm) end subroutine handle_restart subroutine handle_snapshot_step(self, solver, timestep, comm) !! Handle snapshot writing at a given timestep class(snapshot_manager_t), intent(inout) :: self class(solver_t), intent(in) :: solver integer, intent(in) :: timestep integer, intent(in), optional :: comm integer :: comm_to_use comm_to_use = MPI_COMM_WORLD if (present(comm)) comm_to_use = comm call self%write_snapshot(solver, timestep, comm_to_use) end subroutine handle_snapshot_step subroutine write_snapshot(self, solver, timestep, comm) !! Write a snapshot file for visualisation !! Uses a persistent file that stays open across multiple snapshots !! Each snapshot is written as a separate timestep in the file class(snapshot_manager_t), intent(inout) :: self class(solver_t), intent(in) :: solver integer, intent(in) :: timestep integer, intent(in) :: comm character(len=32), allocatable :: field_names(:) integer :: ierr character(len=256) :: filename integer(i8), dimension(3) :: output_shape_dims integer, dimension(3) :: global_dims, output_dims type(field_ptr_t), allocatable :: field_ptrs(:), host_fields(:) real(dp), dimension(3) :: origin, original_spacing, output_spacing real(dp) :: simulation_time logical :: snapshot_uses_stride = .true. logical :: is_first_snapshot_of_run, use_device_write integer :: i if (self%config%snapshot_freq <= 0) return if (mod(timestep, self%config%snapshot_freq) /= 0) return if (has_output_field(self%config, 'species')) then if (solver%nspecies <= 0) then call abort_run('species snapshot output requested, but no transported & &species are configured.', collective=.true.) end if end if field_names = get_snapshot_fields(self%config, solver%nspecies) write (filename, '(A,A)') trim(self%config%snapshot_prefix), '.bp' ! Build identical VTK XML metadata on every rank when the dataset opens. ! It is built before the file opens, as a restart compares it with the ! layout of the existing file. The reduction finds the global origin ! without assuming how ranks map to mesh subdomains. is_first_snapshot_of_run = .not. self%is_snapshot_file_open if (is_first_snapshot_of_run) then global_dims = solver%mesh%get_global_dims(VERT) origin = solver%mesh%get_coordinates(1, 1, 1) call MPI_Allreduce( & MPI_IN_PLACE, origin, size(origin), MPI_X3D2_DP, MPI_MIN, comm, ierr & ) if (ierr /= MPI_SUCCESS) & call abort_run("Failed to determine global snapshot origin") original_spacing = solver%mesh%geo%d if (snapshot_uses_stride) then output_spacing = original_spacing*real(self%output_stride, dp) do i = 1, size(global_dims) output_dims(i) = (global_dims(i) + self%output_stride(i) - 1)/ & self%output_stride(i) end do else output_spacing = original_spacing output_dims = global_dims end if output_shape_dims = int(output_dims, i8) call self%generate_vtk_xml( & output_shape_dims, field_names, origin, output_spacing & ) end if ! Open snapshot file on first call (check for existence) if (is_first_snapshot_of_run) then call self%open_snapshot_file(filename, comm) call self%snapshot_writer%write_attribute("vtk.xml", self%vtk_xml) else ! For subsequent snapshots, begin a new step call self%snapshot_writer%begin_step() end if simulation_time = timestep*solver%dt if (self%snapshot_writer%is_session_functional()) then call log_info('Writing snapshot for time = '//str(simulation_time) & //' iteration = '//str(timestep)) end if call self%snapshot_writer%write_data("time", real(simulation_time, dp)) use_device_write = .false. if (all(self%output_stride == 1)) then use_device_write = all_fields_support_device_write( & self%snapshot_writer, solver, field_names) end if ! Only copy device->host when every field supports GPU-aware I/O. if (.not. use_device_write) then call setup_field_arrays(solver, field_names, field_ptrs, host_fields) end if call self%write_fields( & field_names, host_fields, & solver, self%snapshot_writer, solver%u%data_loc, use_device_write & ) call self%snapshot_writer%end_step() if (.not. use_device_write) then call cleanup_field_arrays(solver, field_ptrs, host_fields) end if deallocate (field_names) end subroutine write_snapshot function get_snapshot_fields(config, nspecies) result(names) !! Build the list of field names written to each snapshot. type(checkpoint_config_t), intent(in) :: config integer, intent(in) :: nspecies character(len=32), allocatable :: names(:) integer :: n, num_fields, is logical :: include_pressure, include_vorticity logical :: include_qcriterion, include_ibm, include_species include_pressure = has_output_field(config, 'pressure') include_vorticity = has_output_field(config, 'vorticity') include_qcriterion = has_output_field(config, 'qcriterion') include_ibm = has_output_field(config, 'ibm') include_species = has_output_field(config, 'species') num_fields = 3 if (include_pressure) num_fields = num_fields + 1 if (include_vorticity) num_fields = num_fields + 1 if (include_qcriterion) num_fields = num_fields + 1 if (include_ibm) num_fields = num_fields + 1 if (include_species) num_fields = num_fields + nspecies allocate (names(num_fields)) n = 3 names(1:3) = [character(len=32) :: "u", "v", "w"] if (include_pressure) then n = n + 1; names(n) = "p" end if if (include_vorticity) then n = n + 1; names(n) = "vort" end if if (include_qcriterion) then n = n + 1; names(n) = "qcrit" end if if (include_ibm) then n = n + 1; names(n) = "ibm" end if if (include_species) then do is = 1, nspecies n = n + 1 write (names(n), '(A,I0)') 'phi_', is end do end if end function get_snapshot_fields subroutine generate_vtk_xml(self, dims, fields, origin, spacing) !! Generate VTK XML string for ImageData format for ParaView's ADIOS2VTXReader class(snapshot_manager_t), intent(inout) :: self integer(i8), dimension(3), intent(in) :: dims character(len=*), dimension(:), intent(in) :: fields real(dp), dimension(3), intent(in) :: origin, spacing character(len=:), allocatable :: xml character(len=96) :: extent_str, origin_str, spacing_str integer :: i ! VTK uses (x,y,z) order, extent defines grid size from 0 to N-1 write (extent_str, '(A,I0,A,I0,A,I0,A,I0,A,I0,A,I0,A)') & '0 ', dims(3) - 1, ' 0 ', dims(2) - 1, ' 0 ', dims(1) - 1 write (origin_str, '(G0, 1X, G0, 1X, G0)') origin write (spacing_str, '(G0, 1X, G0, 1X, G0)') spacing xml = '<?xml version="1.0"?>'//new_line('a')// & '<VTKFile type="ImageData" version="0.1">'//new_line('a')// & ' <ImageData WholeExtent=" '//trim(adjustl(extent_str))//'" ' & //'Origin="'//trim(adjustl(origin_str))//'" '// & 'Spacing="'//trim(adjustl(spacing_str))//'">'//new_line('a')// & ' <Piece Extent="'//trim(adjustl(extent_str))//'">' & //new_line('a')//' <PointData>'//new_line('a') do i = 1, size(fields) xml = trim(xml)//' <DataArray Name="'//trim(fields(i))//'">'// & trim(fields(i))//'</DataArray>'//new_line('a') end do xml = trim(xml)//' <DataArray Name="TIME">time</DataArray>' & //new_line('a') xml = trim(xml)//' </PointData>'//new_line('a')// & ' </Piece>'//new_line('a')// & ' </ImageData>'//new_line('a')// & '</VTKFile>' self%vtk_xml = xml end subroutine generate_vtk_xml subroutine write_fields( & self, field_names, host_fields, solver, writer_session, data_loc, & use_device_write & ) !! Write field data with striding for snapshots class(snapshot_manager_t), intent(inout) :: self character(len=*), dimension(:), intent(in) :: field_names class(field_ptr_t), dimension(:), target, intent(in), optional :: & host_fields class(solver_t), intent(in) :: solver type(writer_session_t), intent(inout) :: writer_session integer, intent(in) :: data_loc logical, intent(in) :: use_device_write integer :: i_field integer(i8), dimension(3) :: output_start, output_count, shape_dims, & count_dims integer, dimension(3) :: output_dims_local class(field_t), pointer :: io_field ! Calculate dimensions for I/O shape_dims = int(solver%mesh%get_global_dims(data_loc), i8) output_start = int(solver%mesh%par%n_offset, i8) count_dims = int(solver%mesh%get_dims(data_loc), i8) ! No striding and GPU-aware I/O available - write directly from device. ! (Striding, and the explicit host override, fall through to the ! host-staged path below, which always slices to the true field ! extent.) if (use_device_write) then ! Let the solver's device work finish once before the batch call solver%backend%sync() do i_field = 1, size(field_names) io_field => get_field_ptr(solver, field_names(i_field)) call writer_session%write_field_from_solver( & trim(field_names(i_field)), io_field, solver%backend, & shape_dims, output_start, count_dims, & self%convert_to_sp & ) end do return end if ! Fallback: host-staged path with striding (requires host_fields) if (.not. present(host_fields)) then call abort_run("write_fields(snapshot): host_fields required for & &strided output") end if ! Prepare buffers with striding for snapshots call prepare_field_buffers( & solver, self%output_stride, field_names, data_loc, & self%field_buffers, self%last_shape_dims, self%last_stride_factors, & self%last_output_shape & ) ! Calculate output dimensions for writing call get_output_dimensions( & int(solver%mesh%get_global_dims(data_loc), i8), & int(solver%mesh%par%n_offset, i8), & int(solver%mesh%get_dims(data_loc), i8), & self%output_stride, & self%last_output_shape, output_start, output_count, & output_dims_local, & self%last_shape_dims, self%last_stride_factors, & self%last_output_shape & ) do i_field = 1, size(field_names) call write_single_field_to_buffer( & trim(field_names(i_field)), host_fields(i_field)%ptr, & solver, self%output_stride, data_loc, & self%field_buffers, self%last_shape_dims, self%last_stride_factors, & self%last_output_shape & ) call writer_session%write_data( & trim(field_names(i_field)), & self%field_buffers(i_field)%buffer, & self%last_output_shape, & output_start, output_count, & self%convert_to_sp) end do end subroutine write_fields logical function all_fields_support_device_write( & writer_session, solver, field_names & ) type(writer_session_t), intent(in) :: writer_session class(solver_t), intent(in) :: solver character(len=*), dimension(:), intent(in) :: field_names class(field_t), pointer :: io_field integer :: i_field all_fields_support_device_write = .false. do i_field = 1, size(field_names) io_field => get_field_ptr(solver, field_names(i_field)) if (.not. writer_session%supports_device_field_write( & io_field, solver%backend)) return end do all_fields_support_device_write = .true. end function all_fields_support_device_write subroutine cleanup_output_buffers(self) !! Clean up dynamic field buffers class(snapshot_manager_t), intent(inout) :: self call cleanup_field_buffers(self%field_buffers) end subroutine cleanup_output_buffers subroutine finalise(self) !! Clean up snapshot manager class(snapshot_manager_t), intent(inout) :: self call self%cleanup_output_buffers() call self%close_snapshot_file() end subroutine finalise subroutine open_snapshot_file(self, filename, comm) !! Open a persistent snapshot file. A restart appends to the file of the !! run it continues if the output layout is unchanged; otherwise a new !! file is created, replacing any old one. class(snapshot_manager_t), intent(inout) :: self character(len=*), intent(in) :: filename integer, intent(in) :: comm logical :: file_exists integer :: myrank, ierr call MPI_Comm_rank(comm, myrank, ierr) ! A changed stride or field list cannot share the file: ParaView reads ! every step with the one vtk.xml it holds if (self%append_after_steps > 0) then if (self%restart_vtk_xml /= self%vtk_xml) & call self%retire_snapshot_file(filename, comm) end if if (self%append_after_steps >= 0) then if (myrank == 0) then call log_info('Appending to existing snapshot file: ' & //trim(filename)//' after step ' & //str(self%append_after_steps)) end if call self%snapshot_writer%open( & filename, comm, append_after_steps=self%append_after_steps & ) else if (myrank == 0) then inquire (file=trim(filename), exist=file_exists) if (file_exists) then call log_info('Overwriting existing snapshot file: '//trim(filename)) else call log_info('Creating new snapshot file: '//trim(filename)) end if end if call self%snapshot_writer%open(filename, comm) end if self%is_snapshot_file_open = .true. end subroutine open_snapshot_file subroutine retire_snapshot_file(self, filename, comm) !! Keep the steps of the existing file up to the restart, under a name !! marking the iteration they end at, so a new file can take its name use m_io_backend, only: allocate_io_writer use m_io_base, only: io_writer_t, io_file_t, io_mode_append class(snapshot_manager_t), intent(inout) :: self character(len=*), intent(in) :: filename integer, intent(in) :: comm class(io_writer_t), allocatable :: writer class(io_file_t), allocatable :: file character(len=256) :: retired_filename logical :: file_exists integer :: myrank, ierr, stat write (retired_filename, '(A,A,I0.6,A)') & trim(self%config%snapshot_prefix), '_to_', self%restart_iter, '.bp' call MPI_Comm_rank(comm, myrank, ierr) if (myrank == 0) inquire (file=trim(retired_filename), exist=file_exists) call MPI_Bcast(file_exists, 1, MPI_LOGICAL, 0, comm, ierr) if (file_exists) then call abort_run('Snapshot output settings changed on restart, but '// & trim(retired_filename)//' already exists', & collective=.true.) end if ! Opening for append drops the steps written past the checkpoint call allocate_io_writer(writer) call writer%init(comm, "snapshot_retire") file = writer%open(filename, io_mode_append, comm, & self%append_after_steps) call file%close() call writer%finalise() if (myrank == 0) then call execute_command_line('mv '//trim(filename)//' '// & trim(retired_filename), exitstat=stat) if (stat /= 0) then call abort_run('Failed to rename '//trim(filename)//' to '// & trim(retired_filename)) end if call log_info('Snapshot output settings changed on restart: kept '// & str(self%append_after_steps)//' steps in '// & trim(retired_filename)) end if call MPI_Barrier(comm, ierr) self%append_after_steps = -1 end subroutine retire_snapshot_file subroutine close_snapshot_file(self) !! Close the persistent snapshot file class(snapshot_manager_t), intent(inout) :: self if (self%is_snapshot_file_open) then call self%snapshot_writer%close() self%is_snapshot_file_open = .false. end if end subroutine close_snapshot_file function join_output_fields(fields) result(str) !! Join non-empty output field names into a comma-separated string. use m_config, only: MAX_OUTPUT_FIELDS character(len=32), intent(in) :: fields(MAX_OUTPUT_FIELDS) character(len=256) :: str integer :: i str = '' do i = 1, MAX_OUTPUT_FIELDS if (fields(i) == '') cycle if (len_trim(str) > 0) str = trim(str)//', ' str = trim(str)//trim(fields(i)) end do end function join_output_fields end module m_snapshot_manager