add ml output capabilities

CMakeLists.txt

@@ -31,6 +31,7 @@ option(ENABLE_TCHEM "Enable TChem chemistry support" OFF)
 option(ENABLE_MPI "Enable MPI parallel support" OFF)
 option(ENABLE_SUNDIALS "Enable SUNDIALS solver for condensation support" OFF)
 option(ENABLE_C_SORT "Enable C sorting routines" OFF)
+option(ENABLE_ML_PROCESS_OUTPUT "Enable output of data after each physical process" OFF)
 
 ######################################################################
 # CPack
@@ -179,6 +180,13 @@ if(ENABLE_SUNDIALS)
   add_definitions(-DPMC_USE_SUNDIALS)
 endif()
 
+######################################################################
+# Output intermediate data during timestep for ML
+
+if(ENABLE_ML_PROCESS_OUTPUT)
+  add_definitions(-DPMC_ML_OUTPUT)
+endif()
+
 ######################################################################
 # tests
 

src/mosaic.F90

@@ -310,6 +310,9 @@ subroutine mosaic_to_partmc(env_state, aero_data, aero_state, gas_data, &
     integer :: i_part, i_spec, i_spec_mosaic
     real(kind=dp), allocatable :: reweight_num_conc(:)
 
+    character(len=PMC_MAX_FILENAME_LEN) :: group_name
+    integer :: ncid_group_ml
+
 #ifdef PMC_USE_WRF
     if (any(particle_error) .eqv. .true.) then
        call output_state('before_error', OUTPUT_TYPE_DIST, aero_data, &
@@ -392,6 +395,17 @@ subroutine mosaic_to_partmc(env_state, aero_data, aero_state, gas_data, &
     end if
 #endif
 
+#ifdef PMC_ML_OUTPUT
+    if (do_ml_output) then
+       call pmc_nc_check(nf90_redef(ncid_ml))
+       write(group_name,'(a)') '4_after_aero_chemistry'
+       call pmc_nc_check(nf90_def_grp(ncid_ml, group_name, ncid_group_ml))
+       call pmc_nc_check(nf90_enddef(ncid_ml))
+       call output_ml_gas_state(gas_state, gas_data, ncid_group_ml)
+       call output_ml_aero_state(aero_state, aero_data, ncid_group_ml)
+    end if
+#endif
+
     ! adjust particles to account for weight changes
     call aero_state_reweight(aero_state, aero_data, reweight_num_conc)
 
@@ -409,10 +423,11 @@ end subroutine mosaic_to_partmc
   !! really matters, however. Because of this mosaic_aero_optical() is
   !! currently disabled.
   subroutine mosaic_timestep(env_state, aero_data, aero_state, gas_data, &
-       gas_state, do_optical, uuid)
+       gas_state, do_optical, uuid, do_process_output)
 
 #ifdef PMC_USE_MOSAIC
-    use module_data_mosaic_main, only: msolar
+    use module_data_mosaic_main, only: msolar, dt_sec, &
+        cair_mlc, cnn, pr_atm, te, avogad
 #endif
 
     !> Environment state.
@@ -429,6 +444,10 @@ subroutine mosaic_timestep(env_state, aero_data, aero_state, gas_data, &
     logical, intent(in) :: do_optical
     !> UUID for this simulation.
     character(len=PMC_UUID_LEN), intent(in) :: uuid
+    !> Whether to output intermidiate data for learning.
+    logical, intent(in), optional :: do_process_output
+
+    character(len=PMC_MAX_FILENAME_LEN) :: group_name
 
 #ifdef PMC_USE_MOSAIC
     ! MOSAIC function interfaces
@@ -437,6 +456,14 @@ subroutine SolarZenithAngle()
        end subroutine SolarZenithAngle
        subroutine IntegrateChemistry()
        end subroutine IntegrateChemistry
+       subroutine GasChemistry(t_in, t_out)
+         use module_data_mosaic_kind, only:  r8
+         real(kind=r8) :: t_in, t_out
+       end subroutine
+       subroutine AerChemistry(t_in, t_out)
+         use module_data_mosaic_kind, only:  r8
+         real(kind=r8) :: t_in, t_out
+       end subroutine
 #ifdef PMC_USE_MOSAIC_MULTI_OPT
        subroutine aerosol_optical(i_wavelength)
          integer, optional :: i_wavelength
@@ -447,6 +474,10 @@ end subroutine aerosol_optical
 #endif
     end interface
 
+#ifdef PMC_ML_OUTPUT
+    integer :: i_spec, i_spec_mosaic, ncid_group_ml
+#endif
+
     ! map PartMC -> MOSAIC
     call mosaic_from_partmc(env_state, aero_data, aero_state, gas_data, &
          gas_state)
@@ -455,7 +486,31 @@ end subroutine aerosol_optical
       call SolarZenithAngle
     end if
 
-    call IntegrateChemistry
+!    call IntegrateChemistry
+
+    call GasChemistry(0d0, dt_sec)
+
+#ifdef PMC_ML_OUTPUT
+    if (do_ml_output) then
+       ! Map back to gas_state and utput changed gas_state
+       cair_mlc = avogad*pr_atm/(82.056d0*te)   ! air conc [molec/cc]
+       ! gas chemistry: map MOSAIC -> PartMC
+       do i_spec = 1,gas_data_n_spec(gas_data)
+          i_spec_mosaic = gas_data%mosaic_index(i_spec)
+          if (i_spec_mosaic > 0) then
+             ! convert molec/cc to ppbv
+             gas_state%mix_rat(i_spec) = cnn(i_spec_mosaic) / cair_mlc * 1d9
+          end if
+       end do
+       call pmc_nc_check(nf90_redef(ncid_ml))
+       write(group_name,'(a)') '3_after_gas_chemistry'
+       call pmc_nc_check(nf90_def_grp(ncid_ml, group_name, ncid_group_ml))
+       call pmc_nc_check(nf90_enddef(ncid_ml))
+       call output_ml_gas_state(gas_state, gas_data, ncid_group_ml)
+    end if
+#endif
+
+    call AerChemistry(0d0, dt_sec)
 
     ! map MOSAIC -> PartMC
     if (do_optical) then

src/output.F90

@@ -101,6 +101,11 @@ module pmc_output
   !> Internal-use variable only.
   integer, parameter :: TAG_OUTPUT_STATE_SINGLE  = 4342
 
+#ifdef PMC_ML_OUTPUT
+    integer, save :: ncid_ml
+    logical, save :: do_ml_output
+#endif
+
 contains
 
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
@@ -962,6 +967,79 @@ subroutine input_column_to_file(filename, index, time, del_t, i_repeat, nz, &
 
   end subroutine input_column_to_file
 
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+  !> Output desired env_state data to a NetCDF group.
+  subroutine output_ml_env_state(env_state, ncid_group)
+
+    !> Environmental state.
+    type(env_state_t), intent(in) :: env_state
+    !> NetCDF group ID.
+    integer, intent(in) :: ncid_group
+
+    call env_state_output_netcdf(env_state, ncid_group)
+
+  end subroutine output_ml_env_state
+
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+  !> Output desired gas_state outputs to a NetCDF group.
+  subroutine output_ml_gas_state(gas_state, gas_data, ncid_group)
+
+    !> Gas state.
+    type(gas_state_t), intent(in) :: gas_state
+    !> Gas data.
+    type(gas_data_t), intent(in) :: gas_data
+    !> NetCDF group ID.
+    integer, intent(in) :: ncid_group
+
+    call gas_state_output_netcdf(gas_state, ncid_group, gas_data)
+
+  end subroutine output_ml_gas_state
+
+!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+
+  !> Output desired aero_state outputs to a NetCDF group.
+  subroutine output_ml_aero_state(aero_state, aero_data, ncid_group)
+
+    !> Aerosol state.
+    type(aero_state_t), intent(in) :: aero_state
+    !> Aerosol data.
+    type(aero_data_t), intent(in) :: aero_data
+    !> NetCDF group ID.
+    integer, intent(in) :: ncid_group
+
+    real(kind=dp), allocatable :: aero_num_conc(:)
+    real(kind=dp), allocatable :: aero_particle_mass(:,:)
+    integer :: dimid_aero_species, dimid_aero_particle
+    integer :: i_part
+
+    ! only need masses and number concentration
+    call aero_data_netcdf_dim_aero_species(aero_data, ncid_group, &
+         dimid_aero_species)
+    call aero_state_netcdf_dim_aero_particle(aero_state, ncid_group, &
+         dimid_aero_particle)
+
+    aero_num_conc = aero_state_num_concs(aero_state, aero_data)
+    allocate(aero_particle_mass(aero_state_n_part(aero_state), &
+         aero_data_n_spec(aero_data)))
+    aero_particle_mass = 0.0d0
+    do i_part = 1,aero_state_n_part(aero_state)
+       aero_particle_mass(i_part, :) &
+            = aero_state%apa%particle(i_part)%vol * aero_data%density
+    end do
+
+    call pmc_nc_write_real_2d(ncid_group, aero_particle_mass, &
+         "aero_particle_mass", (/ dimid_aero_particle, &
+         dimid_aero_species /), unit="kg", &
+         long_name="constituent masses of each aerosol particle")
+
+    call pmc_nc_write_real_1d(ncid_group, aero_num_conc, &
+         "aero_num_conc", (/ dimid_aero_particle /), unit="m^{-3}", &
+         long_name="number concentration for each particle")
+
+  end subroutine output_ml_aero_state
+
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 
 end module pmc_output

src/run_part.F90

@@ -816,7 +816,8 @@ subroutine run_part_timestep(scenario, env_state, aero_data, aero_state, &
     type(camp_state_t), pointer :: camp_pre_aero_state, camp_post_aero_state
     type(camp_env_state_t) :: camp_env_state
 #endif
-
+    character(len=PMC_MAX_FILENAME_LEN) :: filename, group_name
+
     time = i_time * run_part_opt%del_t
 
 #ifdef PMC_USE_CAMP
@@ -844,6 +845,22 @@ subroutine run_part_timestep(scenario, env_state, aero_data, aero_state, &
        progress_n_nuc = progress_n_nuc + n_nuc
     end if
 
+#ifdef PMC_ML_OUTPUT
+    call check_event(time, run_part_opt%del_t, run_part_opt%t_output, &
+            last_output_time, do_ml_output, .false.)
+    if (do_ml_output) then
+       call make_filename(filename, 'ml', ".nc", i_output, 1, 0, 0)
+       call pmc_nc_open_write(filename, ncid_ml)
+       call pmc_nc_check(nf90_redef(ncid_ml))
+       write(group_name,'(a)') '1_before_coagulation'
+       call pmc_nc_check(nf90_def_grp(ncid_ml, group_name, ncid_group_ml))
+       call pmc_nc_check(nf90_enddef(ncid_ml))
+       call output_ml_env_state(env_state, ncid_group_ml)
+       call output_ml_gas_state(gas_state, gas_data, ncid_group_ml)
+       call output_ml_aero_state(aero_state, aero_data, ncid_group_ml)
+    end if
+#endif
+
     if (run_part_opt%do_coagulation) then
        if (run_part_opt%parallel_coag_type &
             == PARALLEL_COAG_TYPE_LOCAL) then
@@ -861,6 +878,17 @@ subroutine run_part_timestep(scenario, env_state, aero_data, aero_state, &
        progress_n_coag = progress_n_coag + n_coag
     end if
 
+
+#ifdef PMC_ML_OUTPUT
+    if (do_ml_output) then
+       call pmc_nc_check(nf90_redef(ncid_ml))
+       write(group_name,'(a)') '2_after_coagulation'
+       call pmc_nc_check(nf90_def_grp(ncid_ml, group_name, ncid_group_ml))
+       call pmc_nc_check(nf90_enddef(ncid_ml))
+       call output_ml_aero_state(aero_state, aero_data, ncid_group_ml)
+    end if
+#endif
+
 #ifdef PMC_USE_SUNDIALS
     if (run_part_opt%do_condensation) then
        call condense_particles(aero_state, aero_data, old_env_state, &
@@ -973,6 +1001,12 @@ subroutine run_part_timestep(scenario, env_state, aero_data, aero_state, &
        end if
     end if
 
+#ifdef PMC_ML_OUTPUT
+    if (do_ml_output) then
+  call pmc_nc_check(nf90_close(ncid_ml))
+  end if
+#endif
+
   end subroutine run_part_timestep
 
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

src/util.F90

@@ -404,7 +404,7 @@ end subroutine check_time_multiple
   !! the next call is guaranteed to do the event. Otherwise the
   !! timestep is used to guess whether to do the event.
   subroutine check_event(time, timestep, interval, last_time, &
-       do_event)
+       do_event, do_increment)
 
     !> Current time.
     real(kind=dp), intent(in) :: time
@@ -416,12 +416,15 @@ subroutine check_event(time, timestep, interval, last_time, &
     real(kind=dp), intent(inout) :: last_time
     !> Whether the event should be done.
     logical, intent(out) :: do_event
-
+    !> Whether to increment the event or just check it.
+    logical, intent(in), optional :: do_increment
     !> Fuzz for event occurance.
     real(kind=dp), parameter :: tolerance = 1d-6
 
     real(kind=dp) closest_interval_time
 
+    if (.not. present(do_increment)) do_increment = .false.
+
     ! if we are at time 0 then do the event unconditionally
     if (time .eq. 0d0) then
        do_event = .true.
@@ -448,8 +451,10 @@ subroutine check_event(time, timestep, interval, last_time, &
        end if
     end if
 
-    if (do_event) then
-       last_time = time
+    if (do_increment) then
+       if (do_event) then
+          last_time = time
+       end if
     end if
 
   end subroutine check_event