diff --git a/hyperion/model/tests/test_specific_energy_spectrum.py b/hyperion/model/tests/test_specific_energy_spectrum.py index 61ccc638..145fc545 100644 --- a/hyperion/model/tests/test_specific_energy_spectrum.py +++ b/hyperion/model/tests/test_specific_energy_spectrum.py @@ -10,7 +10,7 @@ from .. import Model from ...grid import AMRGrid from ...util.functions import random_id -from .test_helpers import get_test_dust +from .test_helpers import get_test_dust, get_realistic_test_dust def _read_dataset(filename, suffix): @@ -25,17 +25,22 @@ def _read_dataset(filename, suffix): return np.asarray(f[matches[-1]][()], dtype=float) -def _assert_spectrum_sums_to_specific_energy(filename): +def _assert_spectrum_sums_to_specific_energy(filename, exclude_empty_spectrum=False): # Summed over frequency, specific_energy_spectrum must reproduce specific_energy. # Cells with no absorption are clamped up to the minimum specific energy # (which only affects specific_energy, not the unclamped specific_energy_spectrum), # so we compare only cells that were genuinely heated above that floor. + # With exclude_empty_spectrum, cells whose spectrum is identically zero are + # also skipped: the PDA can heat cells that received no photons at all, and + # for those there is no Monte-Carlo spectral shape to rescale. se = _read_dataset(filename, '/specific_energy') se_nu = _read_dataset(filename, '/specific_energy_spectrum') assert se_nu is not None nu_sum = se_nu.sum(axis=0) floor = se.min() heated = se > floor * (1. + 1.e-6) + if exclude_empty_spectrum: + heated &= nu_sum > 0. assert np.count_nonzero(heated) >= 1 np.testing.assert_allclose(nu_sum[heated], se[heated], rtol=1.e-6) @@ -85,6 +90,19 @@ def test_specific_energy_spectrum_sums_to_specific_energy(tmpdir): _assert_spectrum_sums_to_specific_energy(out.filename) +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_sums_to_specific_energy_with_pda(tmpdir): + # The partial diffusion approximation overwrites specific_energy in cells + # with few photon deposits, so the frequency-resolved values in those cells + # have to be rescaled to match, rather than keeping the noisy Monte-Carlo + # values. Few photons over many cells guarantees PDA cells exist. + m = _cartesian_model('last', n_cells=8, n_photons=2000) + m.set_pda(True) + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + _assert_spectrum_sums_to_specific_energy(out.filename, exclude_empty_spectrum=True) + + @pytest.mark.requires_hyperion_binaries def test_specific_energy_spectrum_frequencies_written(tmpdir): # specific_energy_spectrum_frequencies is a per-frequency (1-D) quantity, not @@ -162,6 +180,33 @@ def test_specific_energy_spectrum_custom_frequency_grid(tmpdir): _assert_spectrum_sums_to_specific_energy(out.filename) +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_with_mrw(tmpdir): + # Energy deposited during modified-random-walk steps must appear in the + # frequency-resolved spectrum as well as in the scalar specific energy. + # This guards against the bug where grid_do_mrw updated only + # specific_energy_sum, so the spectrum was biased low in the optically + # thick cells where the MRW handles most of the absorption. + m = Model() + m.set_cartesian_grid([-1., 0., 1.], [-1., 0., 1.], [-1., 0., 1.]) + # The density is chosen so that the cells are optically thick enough to + # the local Planck-mean opacity for MRW steps to actually occur. + m.add_density_grid(np.ones((2, 2, 2)) * 1.e5, get_realistic_test_dust()) + s = m.add_point_source() + s.luminosity = 1. + s.temperature = 6000. + m.set_n_initial_iterations(3) + m.set_n_photons(initial=1000, imaging=0) + m.set_seed(-12345) + m.set_mrw(True, gamma=2.) + m.set_max_interactions(1000000000) + m.conf.output.output_specific_energy = 'last' + m.conf.output.output_specific_energy_spectrum = 'last' + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + _assert_spectrum_sums_to_specific_energy(out.filename) + + @pytest.mark.requires_hyperion_binaries def test_specific_energy_spectrum_with_sublimation_cap(tmpdir): # With a dust type using the 'cap' sublimation mode, cells below the @@ -292,3 +337,35 @@ def test_specific_energy_spectrum_mpi_matches_serial(tmpdir): total_serial = np.nansum(_read_dataset(out_serial.filename, '/specific_energy_spectrum')) total_mpi = np.nansum(_read_dataset(out_mpi.filename, '/specific_energy_spectrum')) np.testing.assert_allclose(total_mpi, total_serial, rtol=2.e-2) + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_sums_to_specific_energy_with_mrw(tmpdir): + # In cells heated via the Modified Random Walk, the deposited energy is + # distributed over the frequency bins according to the local emissivity + # (the radiation field is Planckian in the diffusion regime), so the sum + # over frequency must still reproduce the scalar specific energy. The + # high density here makes the model optically thick enough for the MRW + # to handle most of the energy deposition. + m = _cartesian_model('last', n_photons=2000, density=8.) + m.set_mrw(True, gamma=2.) + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + _assert_spectrum_sums_to_specific_energy(out.filename) + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_is_passive_with_mrw(tmpdir): + # The MRW spectrum deposit distributes energy deterministically instead + # of sampling a frequency, so enabling the spectrum must not consume + # random numbers and therefore must not change the specific energy at + # all, even in MRW-dominated models. + se = {} + for output in ('none', 'last'): + m = _cartesian_model(output_specific_energy_spectrum=output, + n_photons=2000, density=8.) + m.set_mrw(True, gamma=2.) + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + se[output] = _read_dataset(out.filename, '/specific_energy') + np.testing.assert_array_equal(se['last'], se['none']) diff --git a/src/dust/dust_type_4elem.f90 b/src/dust/dust_type_4elem.f90 index 5558f513..ff90ce89 100644 --- a/src/dust/dust_type_4elem.f90 +++ b/src/dust/dust_type_4elem.f90 @@ -20,6 +20,7 @@ module type_dust public :: get_j_nu_interp public :: get_j_nu_binned + public :: get_j_nu_bin_fractions public :: get_chi_nu_interp public :: get_chi_nu_binned @@ -748,6 +749,32 @@ function get_j_nu_binned(d, n_nu, nu_min, nu_max, jnu_var_id) result(j_nu) end function get_j_nu_binned + function get_j_nu_bin_fractions(d, jnu_var_id, nu_edges) result(frac) + + ! Fraction of the emissivity falling into each of the bins whose edges + ! are given by nu_edges (in ascending order). The fractions are + ! normalized so that they add up to exactly one. + + implicit none + + type(dust),intent(in) :: d + integer,intent(in) :: jnu_var_id + real(dp),intent(in) :: nu_edges(:) + real(dp) :: frac(size(nu_edges)-1) + + integer :: inu + real(dp) :: norm + + do inu=1, size(nu_edges) - 1 + frac(inu) = integral_loglog(d%j_nu(jnu_var_id)%x, d%j_nu(jnu_var_id)%pdf, & + & nu_edges(inu), nu_edges(inu+1)) + end do + + norm = sum(frac) + if(norm > 0._dp) frac = frac / norm + + end function get_j_nu_bin_fractions + function get_chi_nu_interp(d, nu) result(chi_nu) implicit none diff --git a/src/grid/grid_mrw_3d.f90 b/src/grid/grid_mrw_3d.f90 index 6a6e4d26..1ac3c1d8 100644 --- a/src/grid/grid_mrw_3d.f90 +++ b/src/grid/grid_mrw_3d.f90 @@ -79,6 +79,10 @@ subroutine grid_do_mrw(p) ! Insert ct into (9), get energy deposited for Lucy method e = p%energy * ct * kappa_planck(id, specific_energy(p%icell%ic, id)) specific_energy_sum(p%icell%ic, id) = specific_energy_sum(p%icell%ic, id) + e + ! Deposit the same energy in the frequency-resolved spectrum, + ! distributed over the frequency bins according to the local + ! emissivity (see deposit_specific_energy_spectrum). + call deposit_specific_energy_spectrum(p%icell%ic, id, e) end if end do diff --git a/src/grid/grid_pda_3d.f90 b/src/grid/grid_pda_3d.f90 index dbe237b9..925b781b 100644 --- a/src/grid/grid_pda_3d.f90 +++ b/src/grid/grid_pda_3d.f90 @@ -36,10 +36,11 @@ subroutine update_specific_energy(ic) implicit none integer,intent(in) :: ic integer :: id - real(dp) :: s_prev, s, smin, smax + real(dp) :: s_prev, s, s_old, smin, smax do id=1,n_dust s = specific_energy(ic, id) + s_old = s smin = d(id)%specific_energy(1) smax = d(id)%specific_energy(d(id)%n_e) @@ -58,6 +59,12 @@ subroutine update_specific_energy(ic) end do end if specific_energy(ic, id) = s + + ! Keep the frequency-resolved specific energy consistent with the new + ! scalar value by rescaling the spectrum bins, preserving their shape. + ! If the cell had no energy before, the bins are all zero and there is + ! no shape to rescale, so they are left unchanged. + if(s_old > 0._dp) call scale_specific_energy_spectrum(ic, id, s / s_old) end do end subroutine update_specific_energy diff --git a/src/grid/grid_physics_3d.f90 b/src/grid/grid_physics_3d.f90 index ffc014a4..3a486186 100644 --- a/src/grid/grid_physics_3d.f90 +++ b/src/grid/grid_physics_3d.f90 @@ -16,6 +16,8 @@ module grid_physics private public :: setup_grid_physics + public :: scale_specific_energy_spectrum + public :: deposit_specific_energy_spectrum public :: sublimate_dust public :: update_alpha_inv_planck public :: check_energy_abs @@ -42,6 +44,11 @@ module grid_physics real(dp),allocatable, public :: nu_bins(:) real(dp),allocatable, public :: log_nu_bins(:) + ! Fraction of the emissivity falling into each frequency bin, for each + ! emissivity state and dust type - used to distribute energy deposited by + ! the MRW over the specific energy spectrum. Indexed (bin, state, dust). + real(dp),allocatable, public :: j_nu_bin_frac(:,:,:) + real(dp),allocatable, public :: specific_energy_additional(:,:) real(dp),allocatable, public :: specific_energy_additional_spectrum(:,:,:) real(dp),allocatable, public :: energy_abs_tot(:) @@ -279,6 +286,8 @@ subroutine setup_grid_physics(group, use_mrw, use_pda, compute_specific_energy_s end if allocate(log_nu_bins(n_nu_bins)) log_nu_bins = log10(nu_bins) + + call setup_j_nu_bin_fractions() end if ! Total energy absorbed @@ -319,6 +328,93 @@ subroutine setup_grid_physics(group, use_mrw, use_pda, compute_specific_energy_s end subroutine setup_grid_physics + subroutine setup_j_nu_bin_fractions() + + ! Pre-compute, for each dust type and emissivity state, the fraction of + ! the emissivity falling into each frequency bin, stored in + ! j_nu_bin_frac. This is used to distribute the energy deposited by the + ! MRW over the specific energy spectrum without sampling. The bins are + ! defined by the same nearest-neighbour convention in log-frequency as + ! the photon binning in grid_propagate, so the edges are placed half-way + ! (in log space) between the bin frequencies, with the outer bins + ! extending to all lower/higher frequencies. + + implicit none + + integer :: inu, id, iv, n_nu_bins, n_jnu_max + real(dp),allocatable :: nu_bin_edges(:) + + n_nu_bins = size(log_nu_bins) + + allocate(nu_bin_edges(n_nu_bins+1)) + nu_bin_edges(1) = 10._dp**(log_nu_bins(1) - 99._dp) + do inu=2,n_nu_bins + nu_bin_edges(inu) = 10._dp**(0.5_dp * (log_nu_bins(inu-1) + log_nu_bins(inu))) + end do + nu_bin_edges(n_nu_bins+1) = 10._dp**(log_nu_bins(n_nu_bins) + 99._dp) + + n_jnu_max = 0 + do id=1,n_dust + n_jnu_max = max(n_jnu_max, d(id)%n_jnu) + end do + allocate(j_nu_bin_frac(n_nu_bins, n_jnu_max, n_dust)) + j_nu_bin_frac = 0._dp + do id=1,n_dust + do iv=1,d(id)%n_jnu + j_nu_bin_frac(:, iv, id) = get_j_nu_bin_fractions(d(id), iv, nu_bin_edges) + end do + end do + deallocate(nu_bin_edges) + + end subroutine setup_j_nu_bin_fractions + + subroutine scale_specific_energy_spectrum(ic, id, factor) + + ! Multiply the frequency-resolved specific energy of a cell and dust + ! type by the given factor, preserving the spectral shape. Does nothing + ! if the frequency-resolved specific energy is not being computed. + + implicit none + + integer,intent(in) :: ic, id + real(dp),intent(in) :: factor + + if(.not.compute_specific_energy_spectrum) return + + specific_energy_spectrum(ic, id, :) = specific_energy_spectrum(ic, id, :) * factor + + end subroutine scale_specific_energy_spectrum + + subroutine deposit_specific_energy_spectrum(ic, id, energy) + + ! Add energy to the frequency-resolved specific energy of a cell and + ! dust type, distributed over the frequency bins according to the local + ! emissivity (pre-computed in j_nu_bin_frac), interpolating between the + ! two adjacent emissivity states. This is used by the MRW, for which + ! the radiation field is Planckian, so the absorbed energy is + ! distributed in frequency as kappa_nu * B_nu, which is the local + ! emissivity. No frequency sampling is involved, so computing the + ! spectrum does not affect the random number stream. Does nothing if + ! the frequency-resolved specific energy is not being computed. + + implicit none + + integer,intent(in) :: ic, id + real(dp),intent(in) :: energy + + integer :: iv + real(dp) :: fr + + if(.not.compute_specific_energy_spectrum) return + + iv = jnu_var_id(ic, id) + fr = jnu_var_frac(ic, id) + + specific_energy_sum_spectrum(ic, id, :) = specific_energy_sum_spectrum(ic, id, :) & + & + energy * ((1._dp - fr) * j_nu_bin_frac(:, iv, id) + fr * j_nu_bin_frac(:, iv + 1, id)) + + end subroutine deposit_specific_energy_spectrum + subroutine update_alpha_inv_planck() ! Optimization: could pre-compute alpha_inv_planck just for masked (valid) cells @@ -384,16 +480,13 @@ subroutine sublimate_dust() & * d(id)%sublimation_specific_energy / specific_energy(ic, id) & & * (chi_rosseland(id, specific_energy(ic,id)) & & / chi_rosseland(id, d(id)%sublimation_specific_energy))**2 + + call scale_specific_energy_spectrum(ic, id, & + & d(id)%sublimation_specific_energy / specific_energy(ic, id)) + specific_energy(ic,id) = d(id)%sublimation_specific_energy reset = reset + 1 - - if (compute_specific_energy_spectrum) then - do idx=1,n_nu_bins - specific_energy_spectrum(ic,id,idx) = minimum_specific_energy(id) - end do - end if - end if end do @@ -403,15 +496,13 @@ subroutine sublimate_dust() do ic=1,geo%n_cells if(specific_energy(ic, id) > d(id)%sublimation_specific_energy) then + + call scale_specific_energy_spectrum(ic, id, & + & d(id)%sublimation_specific_energy / specific_energy(ic, id)) + specific_energy(ic, id) = d(id)%sublimation_specific_energy reset = reset + 1 - if (compute_specific_energy_spectrum) then - do idx=1,n_nu_bins - specific_energy_spectrum(ic,id,idx) = minimum_specific_energy(id) - end do - end if - end if end do