diff --git a/docs/setup/setup_conf.rst b/docs/setup/setup_conf.rst index b4b7908f..21c7747c 100644 --- a/docs/setup/setup_conf.rst +++ b/docs/setup/setup_conf.rst @@ -195,6 +195,58 @@ only after last iteration), or ``none`` (do not output). The default is to output only the last iteration of ``specific_energy``. To find out how to view these values, see :doc:`../postprocessing/postprocessing` +Frequency-resolved specific energy +---------------------------------- + +In addition to the total specific energy absorbed in each cell +(``output_specific_energy``), Hyperion can save the specific energy absorbed as +a function of frequency. This is controlled in the same way as the other output +quantities:: + + m.conf.output.output_specific_energy_spectrum = 'all' # every iteration + # 'last' -> only final iteration + # 'none' -> not computed or saved (default) + +When it is not ``'none'``, two extra arrays are written out: + +* ``specific_energy_spectrum`` -- the specific energy absorbed in each cell as a + function of frequency, in erg/s/g. +* ``specific_energy_spectrum_frequencies`` -- the frequencies (in Hz) corresponding to + the leading axis of ``specific_energy_spectrum``. These are bin centers, not edges, + so this array has exactly the same length as that axis (one entry per bin). + +This is the *absorbed* (deposited) energy spectrum, not the mean intensity: each +contribution is weighted by the dust absorption opacity, so it is proportional +to :math:`\kappa_\nu J_\nu`. To recover the mean intensity :math:`J_\nu` (the +radiation field), divide by the dust absorption opacity at each frequency. +Summed over frequency, ``specific_energy_spectrum`` recovers the total +``specific_energy``. + +By default the binning uses the frequency grid of the first dust type. You can +instead provide your own frequency grid (in Hz), independent of the dust +properties:: + + import numpy as np + m.set_specific_energy_spectrum_frequencies(np.logspace(11., 16., 100)) + +Photons are binned to the nearest of these frequencies in log space (so the +supplied values act as bin centers). This works for all +grid types, including AMR and Voronoi. + +.. warning:: The binning has no outer edges: the first and last bins collect + *all* photons below and above the outermost bin centers + respectively, however far away in frequency. This keeps the spectrum + energy-conserving (no photons are dropped), but it means a grid that + does not span the full frequency range will pile up out-of-range + flux in its end bins. When supplying a custom grid, make sure it + brackets the full range of frequencies present in the simulation + (the default dust-based grid already does this). + +``specific_energy_spectrum`` can be retrieved like other grid quantities, as an array +with an extra leading frequency axis:: + + senu = m.get_quantities()['specific_energy_spectrum'] + Advanced parameters ------------------- diff --git a/hyperion/conf/conf_files.py b/hyperion/conf/conf_files.py index 59c62a64..20cd1b14 100644 --- a/hyperion/conf/conf_files.py +++ b/hyperion/conf/conf_files.py @@ -18,6 +18,7 @@ def __init__(self): self.output_density = 'none' self.output_density_diff = 'none' self.output_specific_energy = 'last' + self.output_specific_energy_spectrum = 'none' self.output_n_photons = 'none' self._freeze() @@ -27,6 +28,10 @@ def read(cls, group): self.output_density = group.attrs['output_density'].decode('utf-8') self.output_density_diff = group.attrs['output_density_diff'].decode('utf-8') self.output_specific_energy = group.attrs['output_specific_energy'].decode('utf-8') + if 'output_specific_energy_spectrum' in group.attrs: + self.output_specific_energy_spectrum = group.attrs['output_specific_energy_spectrum'].decode('utf-8') + else: + self.output_specific_energy_spectrum = 'none' self.output_n_photons = group.attrs['output_n_photons'].decode('utf-8') return self @@ -34,6 +39,7 @@ def write(self, group): group.attrs['output_density'] = np.bytes_(self.output_density.encode('utf-8')) group.attrs['output_density_diff'] = np.bytes_(self.output_density_diff.encode('utf-8')) group.attrs['output_specific_energy'] = np.bytes_(self.output_specific_energy.encode('utf-8')) + group.attrs['output_specific_energy_spectrum'] = np.bytes_(self.output_specific_energy_spectrum.encode('utf-8')) group.attrs['output_n_photons'] = np.bytes_(self.output_n_photons.encode('utf-8')) @@ -52,6 +58,8 @@ def __init__(self): self.set_max_reabsorptions(1000000) self.set_pda(False) self.set_mrw(False) + self.specific_energy_spectrum_frequencies = None + self.set_convergence(False) self.set_kill_on_absorb(False) self.set_kill_on_scatter(False) @@ -391,6 +399,43 @@ def _read_pda(self, group): def _write_pda(self, group): group.attrs['pda'] = bool2str(self.pda) + def _read_specific_energy_spectrum_frequencies(self, group): + if 'specific_energy_spectrum_frequencies' in group: + self.specific_energy_spectrum_frequencies = np.array(group['specific_energy_spectrum_frequencies']['nu']) + else: + self.specific_energy_spectrum_frequencies = None + + def _write_specific_energy_spectrum_frequencies(self, group): + if self.specific_energy_spectrum_frequencies is not None: + group.create_dataset('specific_energy_spectrum_frequencies', + data=np.array(list(zip(self.specific_energy_spectrum_frequencies)), + dtype=[('nu', float)])) + + def set_specific_energy_spectrum_frequencies(self, frequencies): + + ''' + Set the frequency grid onto which the frequency-resolved specific energy + (``specific_energy_spectrum``) is binned. + + This is only relevant if ``conf.output.output_specific_energy_spectrum`` is set + to ``'all'`` or ``'last'``. If this method is not called, the frequency + grid of the first dust type is used. Photons are binned to the nearest + frequency in log space, so the supplied values are bin centers (not + edges) and ``specific_energy_spectrum`` has one entry per supplied frequency. + + Parameters + ---------- + frequencies : iterable of float + The frequencies (in Hz) onto which to bin ``specific_energy_spectrum``. + ''' + frequencies = np.asarray(frequencies, dtype=float) + if frequencies.ndim != 1 or frequencies.size < 1: + raise ValueError("frequencies should be a 1-d array of at least one value") + if np.any(frequencies <= 0.): + raise ValueError("frequencies should be positive (in Hz)") + self.specific_energy_spectrum_frequencies = np.sort(frequencies) + + def set_mrw(self, mrw, gamma=1.0, inter_max=1000, warn=True): ''' Set whether to use the Modified Random Walk (MRW) approximation @@ -719,6 +764,7 @@ def read_run_conf(self, group): # not a class method because inherited self._read_max_reabsorptions(group) self._read_pda(group) self._read_mrw(group) + self._read_specific_energy_spectrum_frequencies(group) self._read_convergence(group) self._read_kill_on_absorb(group) self._read_kill_on_scatter(group) @@ -747,6 +793,7 @@ def write_run_conf(self, group): self._write_max_reabsorptions(group) self._write_pda(group) self._write_mrw(group) + self._write_specific_energy_spectrum_frequencies(group) self._write_convergence(group) self._write_kill_on_absorb(group) self._write_kill_on_scatter(group) diff --git a/hyperion/conf/tests/test_conf_io.py b/hyperion/conf/tests/test_conf_io.py index da688d32..9cfe7ed5 100644 --- a/hyperion/conf/tests/test_conf_io.py +++ b/hyperion/conf/tests/test_conf_io.py @@ -4,6 +4,7 @@ from itertools import product +import numpy as np from numpy.testing import assert_equal import pytest @@ -13,7 +14,8 @@ @pytest.mark.parametrize(('attribute', 'value'), list(product(['output_density', 'output_density_diff', - 'output_specific_energy', 'output_n_photons'], + 'output_specific_energy', 'output_specific_energy_spectrum', + 'output_n_photons'], ['none', 'last', 'all']))) def test_io_output_conf(attribute, value): o1 = OutputConf() @@ -249,6 +251,16 @@ def test_io_run_conf_mrw(value): r2.read_run_conf(v) assert r2.mrw == r1.mrw +def test_io_run_conf_specific_energy_spectrum_frequencies(): + r1 = RunConf() + r1.set_specific_energy_spectrum_frequencies(np.logspace(11., 16., 10)) + r1.set_n_photons(1, 2) + v = virtual_file() + r1.write_run_conf(v) + r2 = RunConf() + r2.read_run_conf(v) + np.testing.assert_allclose(r2.specific_energy_spectrum_frequencies, + r1.specific_energy_spectrum_frequencies) @pytest.mark.parametrize(('value'), [False, True]) def test_io_run_conf_convergence(value): diff --git a/hyperion/grid/cartesian_grid.py b/hyperion/grid/cartesian_grid.py index 7676162f..b1f00bde 100644 --- a/hyperion/grid/cartesian_grid.py +++ b/hyperion/grid/cartesian_grid.py @@ -279,6 +279,8 @@ def read_quantities(self, group, quantities='all'): # Read in physical quantities if quantities is not None: for quantity in group: + if quantity == 'specific_energy_spectrum_frequencies': + continue # per-frequency metadata, not a per-cell grid quantity if quantities == 'all' or quantity in quantities: array = np.array(group[quantity]) if array.ndim == 4: # if array is 4D, it is a list of 3D arrays diff --git a/hyperion/grid/cylindrical_polar_grid.py b/hyperion/grid/cylindrical_polar_grid.py index f78342ba..2dc5b149 100644 --- a/hyperion/grid/cylindrical_polar_grid.py +++ b/hyperion/grid/cylindrical_polar_grid.py @@ -307,6 +307,8 @@ def read_quantities(self, group, quantities='all'): # Read in physical quantities if quantities is not None: for quantity in group: + if quantity == 'specific_energy_spectrum_frequencies': + continue # per-frequency metadata, not a per-cell grid quantity if quantities == 'all' or quantity in quantities: array = np.array(group[quantity]) if array.ndim == 4: # if array is 4D, it is a list of 3D arrays diff --git a/hyperion/grid/grid_helpers.py b/hyperion/grid/grid_helpers.py index cadb3496..a60dd8d8 100644 --- a/hyperion/grid/grid_helpers.py +++ b/hyperion/grid/grid_helpers.py @@ -22,7 +22,7 @@ def single_grid_dims(data, ndim=3): ''' if type(data) in [list, tuple]: - + n_pop = len(data) shape = None for item in data: @@ -34,14 +34,21 @@ def single_grid_dims(data, ndim=3): if shape is not None and len(shape) != ndim: raise ValueError("Grids should be %i-dimensional" % ndim) - elif isinstance(data, np.ndarray): + + elif isinstance(data, np.ndarray): if data.ndim == ndim: n_pop = None shape = data.shape elif data.ndim == ndim + 1: n_pop = data.shape[0] shape = data[0].shape + + elif data.ndim == ndim + 2: + # Frequency-resolved quantity, stored as (n_freq, n_pop, *grid) + n_pop = data.shape[1] + shape = data.shape[-ndim:] + else: raise Exception("Unexpected number of dimensions: %i" % data.ndim) @@ -56,6 +63,9 @@ def single_grid_dims(data, ndim=3): elif len(shape) == ndim + 1: n_pop = shape[0] shape = shape[1:] + elif len(shape) == ndim + 2: + n_pop = shape[1] + shape = shape[-ndim:] else: raise Exception("Unexpected number of dimensions: %i" % len(shape)) else: diff --git a/hyperion/grid/octree_grid.py b/hyperion/grid/octree_grid.py index 8396db9d..38485b16 100644 --- a/hyperion/grid/octree_grid.py +++ b/hyperion/grid/octree_grid.py @@ -372,6 +372,8 @@ def read_quantities(self, group, quantities='all'): # Read in physical quantities if quantities is not None: for quantity in group: + if quantity == 'specific_energy_spectrum_frequencies': + continue # per-frequency metadata, not a per-cell grid quantity if quantities == 'all' or quantity in quantities: array = np.array(group[quantity]) if array.ndim == 2: # if array is 2D, it is a list of 1D arrays diff --git a/hyperion/grid/spherical_polar_grid.py b/hyperion/grid/spherical_polar_grid.py index 5d4e7eba..e6d925d4 100644 --- a/hyperion/grid/spherical_polar_grid.py +++ b/hyperion/grid/spherical_polar_grid.py @@ -317,6 +317,8 @@ def read_quantities(self, group, quantities='all'): # Read in physical quantities if quantities is not None: for quantity in group: + if quantity == 'specific_energy_spectrum_frequencies': + continue # per-frequency metadata, not a per-cell grid quantity if quantities == 'all' or quantity in quantities: array = np.array(group[quantity]) if array.ndim == 4: # if array is 4D, it is a list of 3D arrays diff --git a/hyperion/grid/voronoi_grid.py b/hyperion/grid/voronoi_grid.py index 2b03a62f..d15a70a3 100644 --- a/hyperion/grid/voronoi_grid.py +++ b/hyperion/grid/voronoi_grid.py @@ -396,6 +396,8 @@ def read_quantities(self, group, quantities='all'): # Read in physical quantities if quantities is not None: for quantity in group: + if quantity == 'specific_energy_spectrum_frequencies': + continue # per-frequency metadata, not a per-cell grid quantity if quantities == 'all' or quantity in quantities: array = np.array(group[quantity]) if array.ndim == 2: # if array is 2D, it is a list of 1D arrays diff --git a/hyperion/grid/yt3_wrappers.py b/hyperion/grid/yt3_wrappers.py index 5c9aa426..ef2eb1ce 100644 --- a/hyperion/grid/yt3_wrappers.py +++ b/hyperion/grid/yt3_wrappers.py @@ -79,6 +79,9 @@ def amr_grid_to_yt_stream(levels, dust_id=0): grid_dict['level'] = ilevel for field in grid.quantities: + if not isinstance(grid.quantities[field], list): + logger.warning("Skipping frequency-resolved quantity '{0}' in yt export (select a single frequency first)".format(field)) + continue grid_dict[('gas', field)] = grid.quantities[field][dust_id].transpose() grid_data.append(grid_dict) @@ -158,6 +161,9 @@ def octree_grid_to_yt_stream(grid, dust_id=0): quantities = {} for field in grid.quantities: + if not isinstance(grid.quantities[field], list): + logger.warning("Skipping frequency-resolved quantity '{0}' in yt export (select a single frequency first)".format(field)) + continue quantities[('gas', field)] = np.atleast_2d(grid.quantities[field][dust_id][order][~refined]).transpose() bbox = np.array([[xmin, xmax], [ymin, ymax], [zmin, zmax]]) @@ -180,6 +186,9 @@ def cartesian_grid_to_yt_stream(grid, xmin, xmax, ymin, ymax, zmin, zmax, dust_i # Make data dict which should contain (array, unit) tuples data = {} for field in grid.quantities: + if not isinstance(grid.quantities[field], list): + logger.warning("Skipping frequency-resolved quantity '{0}' in yt export (select a single frequency first)".format(field)) + continue data[field] = grid.quantities[field][dust_id].transpose(), '' # Load cartesian grid into yt diff --git a/hyperion/model/model.py b/hyperion/model/model.py index c781ee7b..7701d0d5 100644 --- a/hyperion/model/model.py +++ b/hyperion/model/model.py @@ -223,7 +223,7 @@ def use_geometry(self, filename): # Close the file f.close() - def use_quantities(self, filename, quantities=['density', 'specific_energy'], + def use_quantities(self, filename, quantities=['density', 'specific_energy', 'specific_energy_spectrum'], use_minimum_specific_energy=True, use_dust=True, copy=True, only_initial=False): ''' @@ -297,6 +297,13 @@ def use_quantities(self, filename, quantities=['density', 'specific_energy'], else: quantities_path['specific_energy'] = last_iteration + if 'specific_energy_spectrum' in quantities: + if only_initial or last_iteration is None: + if 'specific_energy_spectrum' in f['/Input/Grid/Quantities']: + quantities_path['specific_energy_spectrum'] = '/Input/Grid/Quantities' + elif 'specific_energy_spectrum' in f[last_iteration]: + quantities_path['specific_energy_spectrum'] = last_iteration + # Minimum specific energy if use_minimum_specific_energy: minimum_specific_energy_path = '/Input/Grid/Quantities' @@ -313,6 +320,10 @@ def use_quantities(self, filename, quantities=['density', 'specific_energy'], if 'specific_energy' in f['/Grid/Quantities']: quantities_path['specific_energy'] = '/Grid/Quantities' + if 'specific_energy_spectrum' in quantities: + if 'specific_energy_spectrum' in f['/Grid/Quantities']: + quantities_path['specific_energy_spectrum'] = '/Grid/Quantities' + # Minimum specific energy if use_minimum_specific_energy: minimum_specific_energy_path = '/Grid/Quantities' diff --git a/hyperion/model/tests/test_specific_energy_spectrum.py b/hyperion/model/tests/test_specific_energy_spectrum.py new file mode 100644 index 00000000..78a639b4 --- /dev/null +++ b/hyperion/model/tests/test_specific_energy_spectrum.py @@ -0,0 +1,270 @@ +from __future__ import print_function, division + +import shutil + +import numpy as np +import pytest + +import h5py + +from .. import Model +from ...grid import AMRGrid +from ...util.functions import random_id +from .test_helpers import get_test_dust + + +def _read_dataset(filename, suffix): + """Return the last grid dataset whose HDF5 path ends with ``suffix``.""" + with h5py.File(filename, 'r') as f: + matches = [] + f.visititems(lambda name, obj: matches.append(name) + if name.endswith(suffix) and isinstance(obj, h5py.Dataset) + else None) + if not matches: + return None + return np.asarray(f[matches[-1]][()], dtype=float) + + +def _assert_spectrum_sums_to_specific_energy(filename): + # 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. + 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) + assert np.count_nonzero(heated) >= 1 + np.testing.assert_allclose(nu_sum[heated], se[heated], rtol=1.e-6) + + +def _cartesian_model(output_specific_energy_spectrum='last', n_cells=2, n_photons=100000, + density=1.e-18, frequencies=None): + m = Model() + edges = np.linspace(-1., 1., n_cells + 1) + m.set_cartesian_grid(edges, edges, edges) + dust = get_test_dust() + m.add_density_grid(np.ones((n_cells, n_cells, n_cells)) * density, dust) + s = m.add_point_source() + s.luminosity = 1. + s.temperature = 6000. + m.set_n_initial_iterations(3) + m.set_n_photons(initial=n_photons, imaging=0) + m.set_seed(-12345) + m.conf.output.output_specific_energy = 'last' + m.conf.output.output_specific_energy_spectrum = output_specific_energy_spectrum + if frequencies is not None: + m.set_specific_energy_spectrum_frequencies(frequencies) + return m + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_does_not_change_specific_energy(tmpdir): + # Computing the frequency-resolved specific energy must not perturb the + # ordinary specific-energy (temperature) calculation. This guards against the + # regression where specific_energy_sum accumulation was gated on the option. + se = {} + for output in ('none', 'last'): + m = _cartesian_model(output_specific_energy_spectrum=output) + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + se[output] = _read_dataset(out.filename, '/specific_energy') + # Same seed, and the extra code path does not touch the random stream, so + # the specific energy should be identical with and without the option. + np.testing.assert_allclose(se['last'], se['none'], rtol=1.e-10) + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_sums_to_specific_energy(tmpdir): + # specific_energy_spectrum, summed over frequency, must reproduce specific_energy. + m = _cartesian_model('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_frequencies_written(tmpdir): + # specific_energy_spectrum_frequencies is a per-frequency (1-D) quantity, not + # per-cell. This guards against the regression where it was written as a grid + # array, which segfaulted whenever n_nu was smaller than n_cells. + m = _cartesian_model('last', n_cells=8) # 512 cells, dust has 2 frequencies + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + bins = _read_dataset(out.filename, '/specific_energy_spectrum_frequencies') + assert bins is not None + assert bins.shape == (2,) # == number of dust frequencies, not n_cells + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_dustless_model_runs(tmpdir): + # A model with no dust must still run (and image) without crashing. This + # guards against the regression where the instrumentation dereferenced the + # first dust type unconditionally, segfaulting dustless models. + m = Model() + m.set_cartesian_grid([-1., 1.], [-1., 1.], [-1., 1.]) + s = m.add_point_source() + s.luminosity = 1. + s.temperature = 6000. + i = m.add_peeled_images(sed=True, image=False) + i.set_viewing_angles([45.], [45.]) + i.set_wavelength_range(5, 0.1, 100.) + m.set_n_initial_iterations(0) + m.set_n_photons(imaging=10000) + m.write(tmpdir.join(random_id()).strpath) + # Should complete without raising (previously segfaulted in the final iteration). + m.run(tmpdir.join(random_id()).strpath) + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_amr(tmpdir): + # specific_energy_spectrum must work on AMR grids: it should run and, summed over + # frequency, reproduce specific_energy in every cell. + m = Model() + amr = AMRGrid() + level = amr.add_level() + grid = level.add_grid() + grid.xmin, grid.xmax = -1., 1. + grid.ymin, grid.ymax = -1., 1. + grid.zmin, grid.zmax = -1., 1. + grid.nx, grid.ny, grid.nz = 4, 4, 4 + grid.quantities['density'] = [np.ones((4, 4, 4)) * 1.e-16] + m.set_amr_grid(amr) + m.add_density_grid(amr['density'][0], get_test_dust()) + s = m.add_point_source() + s.luminosity = 1. + s.temperature = 6000. + m.set_n_initial_iterations(3) + m.set_n_photons(initial=100000, imaging=0) + m.set_seed(-9) + 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_custom_frequency_grid(tmpdir): + # A user-specified frequency grid (independent of the dust frequency grid) + # should be used for the output bins, and energy must still be conserved + # when summing over frequency. + frequencies = np.logspace(10., 16., 12) + m = _cartesian_model('last', density=1.e-16, frequencies=frequencies) + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + bins = _read_dataset(out.filename, '/specific_energy_spectrum_frequencies') + np.testing.assert_allclose(bins, frequencies, rtol=1.e-6) + se_nu = _read_dataset(out.filename, '/specific_energy_spectrum') + assert se_nu.shape[0] == len(frequencies) + _assert_spectrum_sums_to_specific_energy(out.filename) + + +def test_specific_energy_spectrum_frequencies_roundtrip(tmpdir): + # The custom frequency grid should survive a write/read round-trip. + from ...conf.conf_files import RunConf + frequencies = np.logspace(10., 16., 8) + conf = RunConf() + conf.set_specific_energy_spectrum_frequencies(frequencies) + with h5py.File(tmpdir.join('conf.h5').strpath, 'w') as f: + conf._write_specific_energy_spectrum_frequencies(f) + conf2 = RunConf() + conf2._read_specific_energy_spectrum_frequencies(f) + np.testing.assert_allclose(conf2.specific_energy_spectrum_frequencies, frequencies) + + +def test_specific_energy_spectrum_frequencies_validation(): + from ...conf.conf_files import RunConf + conf = RunConf() + with pytest.raises(ValueError): + conf.set_specific_energy_spectrum_frequencies([[1.e10, 1.e12], [1.e14, 1.e16]]) + with pytest.raises(ValueError): + conf.set_specific_energy_spectrum_frequencies([1.e10, -1.e12]) + # Default: no custom grid set + assert conf.specific_energy_spectrum_frequencies is None + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_voronoi(tmpdir): + # specific_energy_spectrum must also work on (unstructured) Voronoi grids, where + # cells are a flat list rather than a structured grid. + rng = np.random.RandomState(12345) + n = 30 + x = rng.uniform(-1., 1., n) + y = rng.uniform(-1., 1., n) + z = rng.uniform(-1., 1., n) + m = Model() + m.set_voronoi_grid(x, y, z) + m.add_density_grid(np.ones(n) * 1.e-16, get_test_dust()) + s = m.add_point_source() + s.luminosity = 1. + s.temperature = 6000. + m.set_n_initial_iterations(3) + m.set_n_photons(initial=100000, imaging=0) + m.set_seed(-12345) + 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_get_quantities(tmpdir): + # specific_energy_spectrum should be retrievable through get_quantities, the + # per-frequency frequencies metadata should not pollute the grid quantities, + # and the retrieved array should still conserve energy. + m = _cartesian_model('last', density=1.e-16) + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + g = out.get_quantities() + assert 'specific_energy_spectrum' in g.quantities + assert 'specific_energy_spectrum_frequencies' not in g.quantities + se = np.array(g.quantities['specific_energy']) # (dust, nz, ny, nx) + se_nu = np.array(g.quantities['specific_energy_spectrum']) # (nu, dust, nz, ny, nx) + assert se_nu.shape[0] == 2 # number of dust frequencies + nu_sum = se_nu.sum(axis=0) + heated = se > se.min() * (1. + 1.e-6) + np.testing.assert_allclose(nu_sum[heated], se[heated], rtol=1.e-6) + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_yt_export_skips_frequency_resolved(tmpdir): + # The frequency-resolved specific_energy_spectrum is not a per-cell scalar, so it + # must be skipped (not exported as a malformed field) when converting to yt, + # while the ordinary scalar quantities are still exported. + pytest.importorskip("yt") + m = _cartesian_model('last', density=1.e-16) + m.write(tmpdir.join(random_id()).strpath) + out = m.run(tmpdir.join(random_id()).strpath) + g = out.get_quantities() + assert 'specific_energy_spectrum' in g.quantities + ds = g.to_yt() + fields = [str(f) for f in ds.field_list] + assert not any('specific_energy_spectrum' in f for f in fields) + assert any(f.endswith("'density')") for f in fields) + + +@pytest.mark.requires_hyperion_binaries +def test_specific_energy_spectrum_mpi_matches_serial(tmpdir): + # The saved spectrum must not depend on the number of MPI processes. We + # compare the total (which is conserved, hence robust to Monte Carlo noise) + # between a serial run and a 2-process MPI run. This is skipped unless both an + # MPI launcher and the MPI build of the binary are available (e.g. CI only + # builds the serial binaries). + if shutil.which('mpirun') is None and shutil.which('mpiexec') is None: + pytest.skip("no MPI launcher available") + if shutil.which('hyperion_car_mpi') is None: + pytest.skip("MPI build of the Hyperion binaries not available") + + m = _cartesian_model('last', n_photons=200000) + m.write(tmpdir.join(random_id()).strpath) + + out_serial = m.run(tmpdir.join(random_id()).strpath) + out_mpi = m.run(tmpdir.join(random_id()).strpath, mpi=True, n_processes=2) + + 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) diff --git a/src/grid/grid_generic.f90 b/src/grid/grid_generic.f90 index d050bca0..c7ec2641 100644 --- a/src/grid/grid_generic.f90 +++ b/src/grid/grid_generic.f90 @@ -1,13 +1,13 @@ module grid_generic use core_lib, only : sp, dp, hid_t, warn, error - use mpi_hdf5_io, only : mp_create_group + use mpi_hdf5_io, only : mp_create_group, mp_write_array use mpi_core, only : main_process - - use grid_io, only : write_grid_3d, write_grid_4d + + use grid_io, only : write_grid_3d, write_grid_4d, write_grid_5d use grid_geometry, only : geo - use grid_physics, only : n_photons, last_photon_id, specific_energy_sum, specific_energy, density, density_original - use settings, only : output_n_photons, output_specific_energy, output_density, output_density_diff, physics_io_type + use grid_physics, only : n_photons, last_photon_id, specific_energy_sum, specific_energy_sum_spectrum, specific_energy, specific_energy_spectrum, nu_bins, density, density_original + use settings, only : output_n_photons, output_specific_energy, output_specific_energy_spectrum, output_density, output_density_diff, physics_io_type, compute_specific_energy_spectrum implicit none save @@ -23,6 +23,7 @@ subroutine grid_reset_energy() if(allocated(n_photons)) n_photons = 0 if(allocated(last_photon_id)) last_photon_id = 0 if(allocated(specific_energy_sum)) specific_energy_sum = 0._dp + if(allocated(specific_energy_sum_spectrum)) specific_energy_sum_spectrum = 0._dp end subroutine grid_reset_energy subroutine output_grid(group, iter, n_iter) @@ -61,6 +62,35 @@ subroutine output_grid(group, iter, n_iter) end if end if + + + ! WRITE THE FREQUENCY-RESOLVED SPECIFIC ENERGY + if (compute_specific_energy_spectrum) then + + if(trim(output_specific_energy_spectrum)=='all' .or. (trim(output_specific_energy_spectrum)=='last'.and.iter==n_iter)) then + + ! The frequencies are a per-frequency quantity, not per-cell, so they + ! are written as a plain 1-D dataset rather than a grid array. + call mp_write_array(group, 'specific_energy_spectrum_frequencies', nu_bins) + + if(allocated(specific_energy_spectrum)) then + select case(physics_io_type) + case(sp) + call write_grid_5d(group, 'specific_energy_spectrum', real(specific_energy_spectrum, sp), geo) + case(dp) + call write_grid_5d(group, 'specific_energy_spectrum', real(specific_energy_spectrum, dp), geo) + case default + call error("output_grid","unexpected value of physics_io_type (should be sp or dp)") + end select + else + call warn("output_grid","specific_energy_spectrum array is not allocated") + end if + + end if + + endif + + ! DENSITY if(trim(output_density)=='all' .or. (trim(output_density)=='last'.and.iter==n_iter)) then diff --git a/src/grid/grid_io.f90 b/src/grid/grid_io.f90 index 8cfc82f5..18c4e343 100644 --- a/src/grid/grid_io.f90 +++ b/src/grid/grid_io.f90 @@ -14,6 +14,7 @@ module grid_io public :: read_grid_4d public :: write_grid_3d public :: write_grid_4d + public :: write_grid_5d interface read_grid_3d module procedure read_grid_3d_sp @@ -43,6 +44,14 @@ module grid_io module procedure write_grid_4d_int8 end interface write_grid_4d + interface write_grid_5d + module procedure write_grid_5d_sp + module procedure write_grid_5d_dp + module procedure write_grid_5d_int + module procedure write_grid_5d_int8 + end interface write_grid_5d + + contains logical function grid_exists(group, name) @@ -108,6 +117,28 @@ subroutine read_grid_3d_int8(group, path, array, geo) end subroutine read_grid_3d_int8 + + + + subroutine write_grid_5d_int8(group, path, array, geo) + + implicit none + + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + integer(idp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + + call mp_write_array(group, path, reshape(array, (/geo%n1, geo%n2, geo%n3, size(array,2), size(array,3)/))) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_int8 + + + + subroutine write_grid_4d_int8(group, path, array, geo) implicit none @@ -192,6 +223,24 @@ subroutine read_grid_3d_int(group, path, array, geo) end subroutine read_grid_3d_int + + + + subroutine write_grid_5d_int(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + integer, intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, reshape(array, (/geo%n1, geo%n2, geo%n3, size(array,2), size(array,3)/))) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_int + + subroutine write_grid_4d_int(group, path, array, geo) implicit none @@ -276,6 +325,24 @@ subroutine read_grid_3d_dp(group, path, array, geo) end subroutine read_grid_3d_dp + + + subroutine write_grid_5d_dp(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + real(dp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + + call mp_write_array(group, path, reshape(array, (/geo%n1, geo%n2, geo%n3, size(array,2), size(array,3)/))) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_dp + + subroutine write_grid_4d_dp(group, path, array, geo) implicit none @@ -360,6 +427,22 @@ subroutine read_grid_3d_sp(group, path, array, geo) end subroutine read_grid_3d_sp + + subroutine write_grid_5d_sp(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + real(sp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, reshape(array, (/geo%n1, geo%n2, geo%n3, size(array,2), size(array,3)/))) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_sp + + subroutine write_grid_4d_sp(group, path, array, geo) implicit none diff --git a/src/grid/grid_io_1d.f90 b/src/grid/grid_io_1d.f90 index 5d87aaa2..7388046d 100644 --- a/src/grid/grid_io_1d.f90 +++ b/src/grid/grid_io_1d.f90 @@ -14,6 +14,8 @@ module grid_io public :: read_grid_4d public :: write_grid_3d public :: write_grid_4d + public :: write_grid_5d + interface read_grid_3d module procedure read_grid_3d_sp @@ -43,6 +45,14 @@ module grid_io module procedure write_grid_4d_int8 end interface write_grid_4d + interface write_grid_5d + module procedure write_grid_5d_sp + module procedure write_grid_5d_dp + module procedure write_grid_5d_int + module procedure write_grid_5d_int8 + end interface write_grid_5d + + contains logical function grid_exists(group, name) @@ -95,6 +105,23 @@ subroutine read_grid_3d_int8(group, path, array, geo) end subroutine read_grid_3d_int8 + + + subroutine write_grid_5d_int8(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + integer(idp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, array) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_int8 + + subroutine write_grid_4d_int8(group, path, array, geo) implicit none @@ -166,6 +193,25 @@ subroutine read_grid_3d_int(group, path, array, geo) end subroutine read_grid_3d_int + + + + + subroutine write_grid_5d_int(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + integer, intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, array) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_int + + subroutine write_grid_4d_int(group, path, array, geo) implicit none @@ -194,7 +240,7 @@ subroutine write_grid_3d_int(group, path, array, geo) end subroutine write_grid_3d_int - + subroutine read_grid_4d_dp(group, path, array, geo) implicit none @@ -237,6 +283,23 @@ subroutine read_grid_3d_dp(group, path, array, geo) end subroutine read_grid_3d_dp + + + subroutine write_grid_5d_dp(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + real(dp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, array) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_dp + + subroutine write_grid_4d_dp(group, path, array, geo) implicit none @@ -265,7 +328,7 @@ subroutine write_grid_3d_dp(group, path, array, geo) end subroutine write_grid_3d_dp - + subroutine read_grid_4d_sp(group, path, array, geo) implicit none @@ -308,6 +371,23 @@ subroutine read_grid_3d_sp(group, path, array, geo) end subroutine read_grid_3d_sp + + + subroutine write_grid_5d_sp(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + real(sp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, array) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_sp + + subroutine write_grid_4d_sp(group, path, array, geo) implicit none diff --git a/src/grid/grid_io_1d_template.f90 b/src/grid/grid_io_1d_template.f90 index 1abbe776..f36aa0b7 100644 --- a/src/grid/grid_io_1d_template.f90 +++ b/src/grid/grid_io_1d_template.f90 @@ -41,6 +41,15 @@ module grid_io module procedure write_grid_4d_int module procedure write_grid_4d_int8 end interface write_grid_4d + + interface write_grid_5d + module procedure write_grid_5d_sp + module procedure write_grid_5d_dp + module procedure write_grid_5d_int + module procedure write_grid_5d_int8 + end interface write_grid_5d + + contains @@ -53,6 +62,7 @@ end function grid_exists !!@FOR real(sp):sp real(dp):dp integer:int integer(idp):int8 + subroutine read_grid_4d_(group, path, array, geo) implicit none @@ -95,6 +105,22 @@ subroutine read_grid_3d_(group, path, array, geo) end subroutine read_grid_3d_ + + subroutine write_grid_5d_(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + @T, intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, array) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_ + + subroutine write_grid_4d_(group, path, array, geo) implicit none diff --git a/src/grid/grid_io_amr.f90 b/src/grid/grid_io_amr.f90 index 832e1725..b49281c4 100644 --- a/src/grid/grid_io_amr.f90 +++ b/src/grid/grid_io_amr.f90 @@ -14,6 +14,8 @@ module grid_io public :: read_grid_4d public :: write_grid_3d public :: write_grid_4d + public :: write_grid_5d + interface read_grid_3d module procedure read_grid_3d_sp @@ -43,6 +45,14 @@ module grid_io module procedure write_grid_4d_int8 end interface write_grid_4d + interface write_grid_5d + module procedure write_grid_5d_sp + module procedure write_grid_5d_dp + module procedure write_grid_5d_int + module procedure write_grid_5d_int8 + end interface write_grid_5d + + contains logical function grid_exists(group, name) @@ -146,6 +156,52 @@ end subroutine test end subroutine read_grid_3d_int8 + + + subroutine write_grid_5d_int8(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + integer(idp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in),target :: geo + integer(hid_t) :: g_level, g_grid + character(len=100) :: name + integer :: ilevel, igrid + type(level_desc), pointer :: level + type(grid_desc), pointer :: grid + + do ilevel=1,size(geo%levels) + level => geo%levels(ilevel) + write(name, '("level_", I5.5)') ilevel + if(mp_path_exists(group, name)) then + g_level = mp_open_group(group, name) + else + g_level = mp_create_group(group, name) + end if + do igrid=1,size(level%grids) + grid => level%grids(igrid) + write(name, '("grid_", I5.5)') igrid + if(mp_path_exists(g_level, name)) then + g_grid = mp_open_group(g_level, name) + else + g_grid = mp_create_group(g_level, name) + end if + + call mp_write_array(g_grid, path, reshape(array(grid%start_id:grid%start_id + grid%n_cells - 1, :,:), & + & (/grid%n1, grid%n2, grid%n3, size(array,2), size(array,3)/))) + call mp_close_group(g_grid) + end do + call mp_close_group(g_level) + + end do + + end subroutine write_grid_5d_int8 + + + + subroutine write_grid_4d_int8(group, path, array, geo) implicit none @@ -226,6 +282,7 @@ subroutine write_grid_3d_int8(group, path, array, geo) end subroutine write_grid_3d_int8 + subroutine read_grid_4d_int(group, path, array, geo) implicit none @@ -304,6 +361,47 @@ end subroutine test end subroutine read_grid_3d_int + + subroutine write_grid_5d_int(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + integer, intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in),target :: geo + integer(hid_t) :: g_level, g_grid + character(len=100) :: name + integer :: ilevel, igrid + type(level_desc), pointer :: level + type(grid_desc), pointer :: grid + + do ilevel=1,size(geo%levels) + level => geo%levels(ilevel) + write(name, '("level_", I5.5)') ilevel + if(mp_path_exists(group, name)) then + g_level = mp_open_group(group, name) + else + g_level = mp_create_group(group, name) + end if + do igrid=1,size(level%grids) + grid => level%grids(igrid) + write(name, '("grid_", I5.5)') igrid + if(mp_path_exists(g_level, name)) then + g_grid = mp_open_group(g_level, name) + else + g_grid = mp_create_group(g_level, name) + end if + call mp_write_array(g_grid, path, reshape(array(grid%start_id:grid%start_id + grid%n_cells - 1, :,:), & + & (/grid%n1, grid%n2, grid%n3, size(array,2), size(array,3)/))) + call mp_close_group(g_grid) + end do + call mp_close_group(g_level) + end do + + end subroutine write_grid_5d_int + + subroutine write_grid_4d_int(group, path, array, geo) implicit none @@ -462,6 +560,48 @@ end subroutine test end subroutine read_grid_3d_dp + + + subroutine write_grid_5d_dp(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + real(dp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in),target :: geo + integer(hid_t) :: g_level, g_grid + character(len=100) :: name + integer :: ilevel, igrid + type(level_desc), pointer :: level + type(grid_desc), pointer :: grid + + do ilevel=1,size(geo%levels) + level => geo%levels(ilevel) + write(name, '("level_", I5.5)') ilevel + if(mp_path_exists(group, name)) then + g_level = mp_open_group(group, name) + else + g_level = mp_create_group(group, name) + end if + do igrid=1,size(level%grids) + grid => level%grids(igrid) + write(name, '("grid_", I5.5)') igrid + if(mp_path_exists(g_level, name)) then + g_grid = mp_open_group(g_level, name) + else + g_grid = mp_create_group(g_level, name) + end if + call mp_write_array(g_grid, path, reshape(array(grid%start_id:grid%start_id + grid%n_cells - 1, :, :), & + & (/grid%n1, grid%n2, grid%n3, size(array,2), size(array,3)/))) + call mp_close_group(g_grid) + end do + call mp_close_group(g_level) + end do + + end subroutine write_grid_5d_dp + + subroutine write_grid_4d_dp(group, path, array, geo) implicit none @@ -620,6 +760,49 @@ end subroutine test end subroutine read_grid_3d_sp + + + subroutine write_grid_5d_sp(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + real(sp), intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in),target :: geo + integer(hid_t) :: g_level, g_grid + character(len=100) :: name + integer :: ilevel, igrid + type(level_desc), pointer :: level + type(grid_desc), pointer :: grid + + do ilevel=1,size(geo%levels) + level => geo%levels(ilevel) + write(name, '("level_", I5.5)') ilevel + if(mp_path_exists(group, name)) then + g_level = mp_open_group(group, name) + else + g_level = mp_create_group(group, name) + end if + do igrid=1,size(level%grids) + grid => level%grids(igrid) + write(name, '("grid_", I5.5)') igrid + if(mp_path_exists(g_level, name)) then + g_grid = mp_open_group(g_level, name) + else + g_grid = mp_create_group(g_level, name) + end if + call mp_write_array(g_grid, path, reshape(array(grid%start_id:grid%start_id + grid%n_cells - 1, :, :), & + & (/grid%n1, grid%n2, grid%n3, size(array,2),size(array,3)/))) + call mp_close_group(g_grid) + end do + call mp_close_group(g_level) + end do + + end subroutine write_grid_5d_sp + + + subroutine write_grid_4d_sp(group, path, array, geo) implicit none diff --git a/src/grid/grid_io_amr_template.f90 b/src/grid/grid_io_amr_template.f90 index f80eb96c..a37d6354 100644 --- a/src/grid/grid_io_amr_template.f90 +++ b/src/grid/grid_io_amr_template.f90 @@ -13,6 +13,8 @@ module grid_io public :: read_grid_4d public :: write_grid_3d public :: write_grid_4d + public :: write_grid_5d + interface read_grid_3d module procedure read_grid_3d_sp @@ -42,6 +44,15 @@ module grid_io module procedure write_grid_4d_int8 end interface write_grid_4d + interface write_grid_5d + module procedure write_grid_5d_sp + module procedure write_grid_5d_dp + module procedure write_grid_5d_int + module procedure write_grid_5d_int8 + end interface write_grid_5d + + + contains logical function grid_exists(group, name) @@ -146,6 +157,48 @@ end subroutine test end subroutine read_grid_3d_ + + + subroutine write_grid_5d_(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + @T, intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in),target :: geo + integer(hid_t) :: g_level, g_grid + character(len=100) :: name + integer :: ilevel, igrid + type(level_desc), pointer :: level + type(grid_desc), pointer :: grid + + do ilevel=1,size(geo%levels) + level => geo%levels(ilevel) + write(name, '("level_", I5.5)') ilevel + if(mp_path_exists(group, name)) then + g_level = mp_open_group(group, name) + else + g_level = mp_create_group(group, name) + end if + do igrid=1,size(level%grids) + grid => level%grids(igrid) + write(name, '("grid_", I5.5)') igrid + if(mp_path_exists(g_level, name)) then + g_grid = mp_open_group(g_level, name) + else + g_grid = mp_create_group(g_level, name) + end if + call mp_write_array(g_grid, path, reshape(array(grid%start_id:grid%start_id + grid%n_cells - 1, :, :), & + & (/grid%n1, grid%n2, grid%n3, size(array,2), size(array,3)/))) + call mp_close_group(g_grid) + end do + call mp_close_group(g_level) + + end do + + end subroutine write_grid_5d_ + subroutine write_grid_4d_(group, path, array, geo) implicit none diff --git a/src/grid/grid_io_template.f90 b/src/grid/grid_io_template.f90 index f42e6eab..696575d0 100644 --- a/src/grid/grid_io_template.f90 +++ b/src/grid/grid_io_template.f90 @@ -13,6 +13,8 @@ module grid_io public :: read_grid_4d public :: write_grid_3d public :: write_grid_4d + public :: write_grid_5d + interface read_grid_3d module procedure read_grid_3d_sp @@ -42,6 +44,14 @@ module grid_io module procedure write_grid_4d_int8 end interface write_grid_4d + interface write_grid_5d + module procedure write_grid_5d_sp + module procedure write_grid_5d_dp + module procedure write_grid_5d_int + module procedure write_grid_5d_int8 + end interface write_grid_5d + + contains logical function grid_exists(group, name) @@ -53,6 +63,7 @@ end function grid_exists !!@FOR real(sp):sp real(dp):dp integer:int integer(idp):int8 + subroutine read_grid_4d_(group, path, array, geo) implicit none @@ -107,6 +118,22 @@ subroutine read_grid_3d_(group, path, array, geo) array = reshape(array3d, (/n_cells/)) end subroutine read_grid_3d_ + + subroutine write_grid_5d_(group, path, array, geo) + + implicit none + + integer(hid_t), intent(in) :: group + character(len=*), intent(in) :: path + @T, intent(in) :: array(:,:,:) + type(grid_geometry_desc),intent(in) :: geo + + call mp_write_array(group, path, reshape(array, (/geo%n1, geo%n2, geo%n3, size(array,2), size(array,3)/))) + call mp_write_keyword(group, path, 'geometry', geo%id) + + end subroutine write_grid_5d_ + + subroutine write_grid_4d_(group, path, array, geo) diff --git a/src/grid/grid_physics_3d.f90 b/src/grid/grid_physics_3d.f90 index 90674116..26578fc6 100644 --- a/src/grid/grid_physics_3d.f90 +++ b/src/grid/grid_physics_3d.f90 @@ -36,8 +36,14 @@ module grid_physics integer(idp),allocatable, public :: n_photons(:) integer(idp),allocatable, public :: last_photon_id(:) real(dp),allocatable, public :: specific_energy(:,:) + real(dp),allocatable, public :: specific_energy_spectrum(:,:,:) real(dp),allocatable, public :: specific_energy_sum(:,:) + real(dp),allocatable, public :: specific_energy_sum_spectrum(:,:,:) + real(dp),allocatable, public :: nu_bins(:) + real(dp),allocatable, public :: log_nu_bins(:) + real(dp),allocatable, public :: specific_energy_additional(:,:) + real(dp),allocatable, public :: specific_energy_additional_spectrum(:,:,:) real(dp),allocatable, public :: energy_abs_tot(:) real(dp),allocatable, public :: minimum_specific_energy(:) @@ -51,12 +57,16 @@ module grid_physics ! Temporary variables (used for convenience in external modules) real(dp), allocatable, public :: tmp_column_density(:) + !Counting indices integer :: id + integer :: idx logical :: debug = .false. type(pdf_discrete_dp) :: absorption + + contains real(dp) function tau_inv_planck_to_closest_wall(p) result(tau) @@ -89,16 +99,25 @@ integer function select_dust_specific_energy_rho(icell) result(id_select) id_select = sample_pdf(absorption) end function select_dust_specific_energy_rho - subroutine setup_grid_physics(group, use_mrw, use_pda) + subroutine setup_grid_physics(group, use_mrw, use_pda, compute_specific_energy_spectrum) implicit none integer(hid_t),intent(in) :: group - logical,intent(in) :: use_mrw, use_pda + logical,intent(in) :: use_mrw, use_pda, compute_specific_energy_spectrum + integer :: n_nu_bins + ! specific_energy_spectrum is binned onto a user-specified frequency grid if one was given, + ! otherwise onto the frequency grid of the first dust type. + if (allocated(specific_energy_spectrum_frequencies)) then + n_nu_bins = size(specific_energy_spectrum_frequencies) + else + n_nu_bins = d(1)%n_nu + end if ! Density allocate(density(geo%n_cells, n_dust)) allocate(specific_energy(geo%n_cells, n_dust)) + allocate(specific_energy_spectrum(geo%n_cells,n_dust,n_nu_bins)) if(n_dust > 0) then @@ -146,6 +165,8 @@ subroutine setup_grid_physics(group, use_mrw, use_pda) ! Check number of dust types for specific_energy if(size(specific_energy, 2).ne.n_dust) call error("setup_grid","specific_energy array has wrong number of dust types") + + ! Reset specific energy to zero in masked cells if(geo%masked) then if(main_process()) write(*, '(" [grid_physics] applying mask to specific_energy grid")') @@ -153,19 +174,39 @@ subroutine setup_grid_physics(group, use_mrw, use_pda) where(.not.geo%mask) specific_energy(:, id) = 0. end where + + if (compute_specific_energy_spectrum) then + do idx=1,n_nu_bins + where(.not.geo%mask) + specific_energy_spectrum(:, id, idx) = 0. + endwhere + end do + end if + end do end if if(trim(specific_energy_type) == 'additional') then allocate(specific_energy_additional(geo%n_cells, n_dust)) + if (compute_specific_energy_spectrum) then + allocate(specific_energy_additional_spectrum(geo%n_cells,n_dust,n_nu_bins)) + end if ! We store a copy of the initial specific energy in a separate ! array, and we set the specific energy to the minimum specific ! energy. After the first iteration, specific_energy will get ! re-calculated and we will then add specific_energy_additional specific_energy_additional = specific_energy + specific_energy_additional_spectrum = specific_energy_spectrum do id=1,n_dust specific_energy(:,id) = minimum_specific_energy(id) - end do + + if (compute_specific_energy_spectrum) then + do idx=1,n_nu_bins + specific_energy_spectrum(:,id,idx) = minimum_specific_energy(id) + end do + end if + + end do end if @@ -178,6 +219,13 @@ subroutine setup_grid_physics(group, use_mrw, use_pda) ! Set all specific_energy to minimum requested do id=1,n_dust specific_energy(:,id) = minimum_specific_energy(id) + + if (compute_specific_energy_spectrum) then + do idx = 1,n_nu_bins + specific_energy_spectrum(:,id,idx) = minimum_specific_energy(id) + end do + end if + end do end if @@ -191,6 +239,26 @@ subroutine setup_grid_physics(group, use_mrw, use_pda) allocate(specific_energy_sum(geo%n_cells, n_dust)) specific_energy_sum = 0._dp + ! Set up basics for the frequency-resolved specific energy + allocate(specific_energy_sum_spectrum(geo%n_cells, n_dust, n_nu_bins)) + specific_energy_sum_spectrum = 0._dp + + ! Cache the frequency grid (and its log) once so they do not have to be + ! rebuilt for every photon. Photons are binned to the nearest grid point in + ! log-frequency space (see grid_propagate). + if (compute_specific_energy_spectrum) then + allocate(nu_bins(n_nu_bins)) + if (allocated(specific_energy_spectrum_frequencies)) then + nu_bins = specific_energy_spectrum_frequencies + else + do idx=1,n_nu_bins + nu_bins(idx) = d(1)%nu(idx) + end do + end if + allocate(log_nu_bins(n_nu_bins)) + log_nu_bins = log10(nu_bins) + end if + ! Total energy absorbed allocate(energy_abs_tot(n_dust)) energy_abs_tot = 0._dp @@ -257,6 +325,9 @@ subroutine sublimate_dust() implicit none integer :: ic, id integer :: reset + integer :: n_nu_bins + + n_nu_bins = size(specific_energy_spectrum, 3) reset = 0 @@ -270,7 +341,15 @@ subroutine sublimate_dust() density(ic, id) = 0. specific_energy(ic, id) = minimum_specific_energy(id) 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 if(reset > 0) write(*,'(" [sublimate_dust] dust removed in ",I0," cells")') reset @@ -284,6 +363,14 @@ subroutine sublimate_dust() & / chi_rosseland(id, d(id)%sublimation_specific_energy))**2 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 @@ -296,6 +383,14 @@ subroutine sublimate_dust() specific_energy(ic, id) = d(id)%sublimation_specific_energy reset = reset + 1 end if + + + 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 do if(reset > 0) write(*,'(" [sublimate_dust] capping dust specific_energy in ",I0," cells")') reset @@ -316,12 +411,23 @@ subroutine update_energy_abs(scale) real(dp), intent(in) :: scale - integer :: id + integer :: id,idx + integer :: n_nu_bins + + n_nu_bins = size(specific_energy_spectrum, 3) if(main_process()) write(*,'(" [grid_physics] updating energy_abs")') do id=1,n_dust specific_energy(:,id) = specific_energy_sum(:,id) * scale / geo%volume + + if (compute_specific_energy_spectrum) then + do idx=1,n_nu_bins + specific_energy_spectrum(:,id,idx) = specific_energy_sum_spectrum(:,id,idx) * scale/geo%volume + end do + end if + + where(geo%volume == 0._dp) specific_energy(:,id) = 0._dp end where @@ -331,10 +437,17 @@ subroutine update_energy_abs(scale) write(*,'(" [update_energy_abs] ",I0," cells have no energy")') count(specific_energy==0.and.density>0.) end if + ! Add in additional source of heating if(trim(specific_energy_type) == 'additional') then if(main_process()) write(*,'(" [grid_physics] adding additional heating source")') specific_energy = specific_energy + specific_energy_additional + + + if (compute_specific_energy_spectrum) then + specific_energy_spectrum = specific_energy_spectrum + specific_energy_additional_spectrum + end if + end if call update_energy_abs_tot() @@ -357,7 +470,10 @@ subroutine check_energy_abs() if(main_process()) call warn("update_energy_abs","specific_energy below minimum requested in some cells - resetting") where(specific_energy(:,id) < minimum_specific_energy(id)) specific_energy(:,id) = minimum_specific_energy(id) + + end where + end if if(any(specific_energy(:,id) < d(id)%specific_energy(1))) then @@ -365,7 +481,8 @@ subroutine check_energy_abs() if(main_process()) call warn("update_energy_abs","specific_energy below minimum allowed in some cells - resetting") where(specific_energy(:,id) < d(id)%specific_energy(1)) specific_energy(:,id) = d(id)%specific_energy(1) - end where + + end where else if(main_process()) call warn("update_energy_abs","specific_energy below minimum allowed in some cells - will pick closest emissivities") end if @@ -376,6 +493,7 @@ subroutine check_energy_abs() if(main_process()) call warn("update_energy_abs","specific_energy above maximum allowed in some cells - resetting") where(specific_energy(:,id) > d(id)%specific_energy(d(id)%n_e)) specific_energy(:,id) = d(id)%specific_energy(d(id)%n_e) + end where else if(main_process()) call warn("update_energy_abs","specific_energy above maximum allowed in some cells - will pick closest emissivities") diff --git a/src/grid/grid_propagate_3d.f90 b/src/grid/grid_propagate_3d.f90 index 3c345161..3769704c 100644 --- a/src/grid/grid_propagate_3d.f90 +++ b/src/grid/grid_propagate_3d.f90 @@ -5,10 +5,10 @@ module grid_propagate use type_grid_cell use dust_main, only : n_dust use grid_geometry, only : escaped, find_wall, in_correct_cell, next_cell, opposite_wall - use grid_physics, only : specific_energy_sum, density, n_photons, last_photon_id + use grid_physics, only : specific_energy_sum, specific_energy_sum_spectrum, log_nu_bins, density, n_photons, last_photon_id use sources use counters - use settings, only : frac_check => propagation_check_frequency + use settings, only : frac_check => propagation_check_frequency, compute_specific_energy_spectrum implicit none save @@ -57,6 +57,8 @@ subroutine grid_integrate(p,tau_required,tau_achieved) integer :: source_id + integer :: idx + radial = (p%r .dot. p%v) > 0. if(debug) write(*,'(" [debug] start grid_integrate")') @@ -130,15 +132,22 @@ subroutine grid_integrate(p,tau_required,tau_achieved) p%r = p%r + tmin * p%v tau_achieved = tau_achieved + tau_cell - ! NEED INDIVIDUAL ALPHA HERE + + if (compute_specific_energy_spectrum) idx = minloc(abs(log_nu_bins - log10(p%nu)), DIM=1) do id=1,n_dust if(density(p%icell%ic, id) > 0._dp) then specific_energy_sum(p%icell%ic, id) = & & specific_energy_sum(p%icell%ic, id) + tmin * p%current_kappa(id) * p%energy + if (compute_specific_energy_spectrum) then + specific_energy_sum_spectrum(p%icell%ic, id, idx) = & + & specific_energy_sum_spectrum(p%icell%ic, id, idx) + tmin * p%current_kappa(id) * p%energy + end if end if end do + + p%on_wall = .true. p%icell = next_cell(p%icell, id_min, intersection=p%r) p%on_wall_id = opposite_wall(id_min) @@ -271,6 +280,8 @@ subroutine grid_integrate_noenergy(p,tau_required,tau_achieved) do id=1,n_dust chi_rho_total = chi_rho_total + p%current_chi(id) * density(p%icell%ic, id) end do + + tau_cell = chi_rho_total * tmin if(tau_cell < tau_needed) then diff --git a/src/main/settings.f90 b/src/main/settings.f90 index daf4ad23..5fd3bf65 100644 --- a/src/main/settings.f90 +++ b/src/main/settings.f90 @@ -20,7 +20,7 @@ module settings integer(idp),public :: n_last_photons_dust = 0 integer(idp),public :: n_raytracing_photons_sources = 0 integer(idp),public :: n_raytracing_photons_dust = 0 - logical,public :: use_raytracing, use_mrw, use_pda + logical,public :: use_raytracing, use_mrw, use_pda, compute_specific_energy_spectrum logical, public :: kill_on_absorb, kill_on_scatter real(dp),public :: mrw_gamma logical, public :: forced_first_interaction @@ -30,11 +30,16 @@ module settings real(dp) :: monochromatic_energy_threshold real(dp),allocatable :: frequencies(:) + ! Optional user-specified frequency grid for specific_energy_spectrum. If not + ! allocated, the frequency grid of the first dust type is used instead. + real(dp),allocatable :: specific_energy_spectrum_frequencies(:) + integer :: physics_io_type character(len=4) :: output_density character(len=4) :: output_density_diff character(len=4) :: output_specific_energy + character(len=4) :: output_specific_energy_spectrum character(len=4) :: output_n_photons logical :: check_convergence = .false. diff --git a/src/main/setup_rt.f90 b/src/main/setup_rt.f90 index 649799fa..9bee928c 100644 --- a/src/main/setup_rt.f90 +++ b/src/main/setup_rt.f90 @@ -74,6 +74,31 @@ subroutine setup_initial(input_handle) call mp_read_keyword(input_handle, '/', 'pda', use_pda) call mp_read_keyword(input_handle, '/', 'mrw', use_mrw) + ! The frequency-resolved specific energy is computed whenever it is to be + ! output, i.e. unless output_specific_energy_spectrum is 'none'. This is read here + ! (ahead of the main OUTPUT block below) because it is needed to set up the + ! grid physics arrays. + g_output = mp_open_group(input_handle, '/Output') + if(mp_exists_keyword(g_output, '.', 'output_specific_energy_spectrum')) then + call mp_read_keyword(g_output, '.', 'output_specific_energy_spectrum', output_specific_energy_spectrum) + else + output_specific_energy_spectrum = 'none' + end if + call mp_close_group(g_output) + + if(trim(output_specific_energy_spectrum).ne.'all' & + & .and.trim(output_specific_energy_spectrum).ne.'last' & + & .and.trim(output_specific_energy_spectrum).ne.'none') & + & call error("setup_initial","output_specific_energy_spectrum should be one of all/last/none") + + compute_specific_energy_spectrum = trim(output_specific_energy_spectrum).ne.'none' + + ! Optional user-specified frequency grid (otherwise the dust grid is used) + if(compute_specific_energy_spectrum) then + if(mp_path_exists(input_handle, 'specific_energy_spectrum_frequencies')) then + call mp_table_read_column_auto(input_handle, 'specific_energy_spectrum_frequencies', 'nu', specific_energy_spectrum_frequencies) + end if + end if if(use_mrw) then call mp_read_keyword(input_handle, '/', 'mrw_gamma', mrw_gamma) @@ -173,7 +198,7 @@ subroutine setup_initial(input_handle) call mp_close_group(g_geometry) g_physics = mp_open_group(input_handle, '/Grid/Quantities') - call setup_grid_physics(g_physics, use_mrw, use_pda) + call setup_grid_physics(g_physics, use_mrw, use_pda, compute_specific_energy_spectrum) call mp_close_group(g_physics) call mp_read_keyword(input_handle, '/', 'physics_io_bytes', physics_io_bytes) diff --git a/src/mpi/mpi_routines.f90 b/src/mpi/mpi_routines.f90 index fb0d6e7b..bb585432 100644 --- a/src/mpi/mpi_routines.f90 +++ b/src/mpi/mpi_routines.f90 @@ -266,6 +266,8 @@ subroutine mp_collect_physical_arrays() real(dp) :: tmp real(dp) :: dummy_dp integer(idp) :: dummy_idp + + real(dp),allocatable :: tmp_3d(:,:,:) real(dp),allocatable :: tmp_2d(:,:) integer(idp),allocatable :: tmp_int_1d(:) @@ -278,6 +280,15 @@ subroutine mp_collect_physical_arrays() call mpi_reduce(specific_energy_sum, dummy_dp, size(specific_energy_sum), mpi_real8, mpi_sum, rank_main, mpi_comm_world, ierr) end if + if(main_process()) then + allocate(tmp_3d(size(specific_energy_sum_spectrum,1),size(specific_energy_sum_spectrum,2),size(specific_energy_sum_spectrum,3))) + call mpi_reduce(specific_energy_sum_spectrum, tmp_3d, size(specific_energy_sum_spectrum), mpi_real8, mpi_sum, rank_main, mpi_comm_world, ierr) + specific_energy_sum_spectrum = tmp_3d + deallocate(tmp_3d) + else + call mpi_reduce(specific_energy_sum_spectrum, dummy_dp, size(specific_energy_sum_spectrum), mpi_real8, mpi_sum, rank_main, mpi_comm_world, ierr) + end if + if(allocated(n_photons)) then if(main_process()) then allocate(tmp_int_1d(size(n_photons,1)))