[WIP] compute and save the ISRF in every cell

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_nu = 'last'
         self.output_n_photons = 'none'
         self._freeze()
 
@@ -27,13 +28,15 @@ 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')
+        self.output_specific_energy_nu = group.attrs['output_specific_energy_nu'].decode('utf-8')
         self.output_n_photons = group.attrs['output_n_photons'].decode('utf-8')
         return self
 
     def write(self, group):
         group.attrs['output_density'] = np.string_(self.output_density.encode('utf-8'))
         group.attrs['output_density_diff'] = np.string_(self.output_density_diff.encode('utf-8'))
         group.attrs['output_specific_energy'] = np.string_(self.output_specific_energy.encode('utf-8'))
+        group.attrs['output_specific_energy_nu'] = np.string_(self.output_specific_energy_nu.encode('utf-8'))
         group.attrs['output_n_photons'] = np.string_(self.output_n_photons.encode('utf-8'))
 
 
@@ -52,6 +55,8 @@ def __init__(self):
         self.set_max_reabsorptions(1000000)
         self.set_pda(False)
         self.set_mrw(False)
+        self.compute_isrf(False)
+
         self.set_convergence(False)
         self.set_kill_on_absorb(False)
         self.set_kill_on_scatter(False)
@@ -391,6 +396,25 @@ def _read_pda(self, group):
     def _write_pda(self, group):
         group.attrs['pda'] = bool2str(self.pda)
 
+    def _read_isrf(self,group):
+        self.isrf = str2bool(group.attrs['isrf'])
+
+    def _write_isrf(self,group):
+        group.attrs['isrf'] = bool2str(self.isrf)
+
+    def compute_isrf(self,isrf):
+
+        '''
+
+        Set whether or not to compute and save the ISRF in each cell
+
+        If enabled, the ISRF is computed in every cell at the
+        frequencies of the dust opacity tables.
+
+        '''
+        self.isrf = isrf
+
+
     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 +743,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._reada_isrf(group)
         self._read_convergence(group)
         self._read_kill_on_absorb(group)
         self._read_kill_on_scatter(group)
@@ -747,6 +772,7 @@ def write_run_conf(self, group):
         self._write_max_reabsorptions(group)
         self._write_pda(group)
         self._write_mrw(group)
+        self._write_isrf(group)
         self._write_convergence(group)
         self._write_kill_on_absorb(group)
         self._write_kill_on_scatter(group)

hyperion/conf/tests/test_conf_io.py

@@ -249,6 +249,16 @@ def test_io_run_conf_mrw(value):
     r2.read_run_conf(v)
     assert r2.mrw == r1.mrw
 
+@pytest.mark.parametrize(('value'), [True, False])
+def test_io_run_conf_isrf(value):
+    r1 = RunConf()
+    r1.compute_isrf(value)
+    r1.set_n_photons(1, 2)
+    v = virtual_file()
+    r1.write_run_conf(v)
+    r2 = RunConf()
+    r2.read_run_conf(v)
+    assert r2.isrf == r1.isrf
 
 @pytest.mark.parametrize(('value'), [False, True])
 def test_io_run_conf_convergence(value):

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,22 @@ 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
+
+        #DN CRAZY ADDITIONS
+        elif data.ndim == ndim+2:
+            n_pop = data.shape[1]
+            shape = data.shape[-1]
+            shape = tuple(map(int,str(shape).split(' ')))
+
         else:
             raise Exception("Unexpected number of dimensions: %i" % data.ndim)
 

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_nu'],
                        use_minimum_specific_energy=True, use_dust=True, copy=True,
                        only_initial=False):
         '''
@@ -297,6 +297,14 @@ def use_quantities(self, filename, quantities=['density', 'specific_energy'],
                 else:
                     quantities_path['specific_energy'] = last_iteration
 
+            if self.compute_isrf == True:
+                if 'specific_energy_nu' in quantities:
+                    if only_initial or last_iteration is None:
+                        if 'specific_energy_nu' in f['/Input/Grid/Quantities']:
+                            quantities_path['specific_energy_nu'] = '/Input/Grid/Quantities'
+                    else:
+                        quantities_path['specific_energy_nu'] = last_iteration
+
             # Minimum specific energy
             if use_minimum_specific_energy:
                 minimum_specific_energy_path = '/Input/Grid/Quantities'
@@ -313,6 +321,11 @@ def use_quantities(self, filename, quantities=['density', 'specific_energy'],
                 if 'specific_energy' in f['/Grid/Quantities']:
                     quantities_path['specific_energy'] = '/Grid/Quantities'
 
+            if self.compute_isrf == True:
+                if 'specific_energy_nu' in quantities:
+                    if 'specific_energy_nu' in f['/Grid/Quantities']:
+                        quantities_path['specific_energy_nu'] = '/Grid/Quantities'
+
             # Minimum specific energy
             if use_minimum_specific_energy:
                 minimum_specific_energy_path = '/Grid/Quantities'
@@ -808,6 +821,9 @@ class (e.g. ``SphericalDust``, ``IsotropicDust``, ...).
             self.grid['density'] = []
             if specific_energy is not None:
                 self.grid['specific_energy'] = []
+
+                if self.compute_isrf == True:
+                    self.grid['specific_energy_nu'] = []
 
         # Check whether the density can be added to an existing one
         if merge_if_possible:
@@ -849,6 +865,9 @@ class (e.g. ``SphericalDust``, ``IsotropicDust``, ...).
         # Set specific energy if specified
         if specific_energy is not None:
             self.grid['specific_energy'].append(specific_energy)
+
+            if self.compute_isrf == True:
+                self.grid['specific_energy_nu'].append(specific_energy)
 
     def set_cartesian_grid(self, x_wall, y_wall, z_wall):
         self.set_grid(CartesianGrid(x_wall, y_wall, z_wall))

src/grid/grid_generic.f90

@@ -3,11 +3,14 @@ module grid_generic
   use core_lib, only : sp, dp, hid_t, warn, error
   use mpi_hdf5_io, only : mp_create_group
   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_nu, specific_energy, specific_energy_nu, density, density_original
+  use settings, only : output_n_photons, output_specific_energy, output_density, output_density_diff, physics_io_type, compute_isrf
+
+  use dust_main, only: d
+
 
   implicit none
   save
@@ -23,6 +26,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_nu)) specific_energy_sum_nu = 0._dp
   end subroutine grid_reset_energy
 
   subroutine output_grid(group, iter, n_iter)
@@ -31,6 +35,10 @@ subroutine output_grid(group, iter, n_iter)
 
     integer(hid_t),intent(in) :: group
     integer,intent(in) :: iter, n_iter
+    integer :: n_cells, n_dust, n_isrf_lam
+    integer :: i,j,k
+    real, dimension(d(1)%n_nu) :: energy_frequency_bins
+
 
     if(main_process()) write(*,'(" [output_grid] outputting grid arrays for iteration")')
 
@@ -61,6 +69,58 @@ subroutine output_grid(group, iter, n_iter)
        end if
     end if
 
+
+
+    ! WRITE THE ISRF
+    if (compute_isrf) then 
+
+       n_cells = size(specific_energy_nu, 1)
+       n_dust = size(specific_energy_nu, 2)
+       n_isrf_lam = size(specific_energy_nu,3)
+
+       do i=1,d(1)%n_nu
+          energy_frequency_bins(i) = d(1)%nu(i)
+       end do
+
+       if(trim(output_specific_energy)=='all' .or. (trim(output_specific_energy)=='last'.and.iter==n_iter)) then
+          !if(allocated(energy_frequency_bins)) then
+          call write_grid_3d(group, 'ISRF_frequency_bins',energy_frequency_bins, geo)
+          !else
+          !call warn("output_grid","energy_frequency bins [ISRF wavelengths] array is not allocated")
+          !end if
+       end if
+
+       if(trim(output_specific_energy)=='all' .or. (trim(output_specific_energy)=='last'.and.iter==n_iter)) then
+          if(allocated(specific_energy_nu)) then
+
+
+             if (compute_isrf) then
+                print *,'[GRID_GENERIC.F90] ENTERING BLOCK FOR WRITING SPECIFIC_ENERGY_NU'
+
+                select case(physics_io_type)
+
+                case(sp)
+                   print *,'[GRID_GENERIC.F90] diving into write_grid_5d sp'
+                   call write_grid_5d(group, 'specific_energy_nu', real(specific_energy_nu, sp), geo)
+
+                case(dp)
+                   print *,'[GRID_GENERIC.F90] diving into write_grid_5d dp'
+                   call write_grid_5d(group, 'specific_energy_nu', real(specific_energy_nu, dp), geo)
+                   print *,'[GRID_GENERIC.F90] finished with write_grid_5d dp'
+
+                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_nu array is not allocated")
+
+             end if
+          end if
+       end if
+
+    endif
+
+
     ! DENSITY
 
     if(trim(output_density)=='all' .or. (trim(output_density)=='last'.and.iter==n_iter)) then