Fix specific energy spectrum when PDA and MRW are used, as well as dust sublimation capping

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'])

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

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
 

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
 

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