Merge pull request #338 from orbitfold/openmp

OpenMP WIP

setup.py

@@ -15,8 +15,9 @@
     import __builtin__ as builtins
 builtins._ASTROPY_SETUP_ = True
 
-from astropy_helpers.setup_helpers import (
-    register_commands, adjust_compiler, get_debug_option, get_package_info)
+from astropy_helpers.setup_helpers import (register_commands, adjust_compiler, 
+                                           get_debug_option, get_package_info,
+                                           add_command_option)
 from astropy_helpers.git_helpers import get_git_devstr
 from astropy_helpers.version_helpers import generate_version_py
 
@@ -55,6 +56,10 @@
 # invoking any other functionality from distutils since it can potentially
 # modify distutils' behavior.
 cmdclassd = register_commands(PACKAGENAME, VERSION, RELEASE)
+add_command_option('install', 'with-openmp', 'compile TARDIS without OpenMP',
+                   is_bool=True)
+add_command_option('build', 'with-openmp', 'compile TARDIS without OpenMP',
+                   is_bool=True)
 
 # Adjust the compiler in case the default on this platform is to use a
 # broken one.

tardis/data/tardis_config_definition.yml

@@ -262,6 +262,12 @@ model:
 
 
 montecarlo:
+    nthreads:
+        property_type: int
+        default: 1
+        mandatory: False
+        help: The number of OpenMP threads.
+
     seed:
         property_type: int
         default: 23111963

tardis/model.py

@@ -352,12 +352,12 @@ def simulate(self, update_radiation_field=True, enable_virtual=False, initialize
 
         self.j_blue_estimators = np.zeros((len(self.t_rads), len(self.atom_data.lines)))
         self.montecarlo_virtual_luminosity = np.zeros_like(self.spectrum.frequency.value)
-
+        
         montecarlo_nu, montecarlo_energies, self.j_estimators, self.nubar_estimators, \
         last_line_interaction_in_id, last_line_interaction_out_id, \
         self.last_interaction_type, self.last_line_interaction_shell_id = \
             montecarlo.montecarlo_radial1d(self,
-                                                     virtual_packet_flag=no_of_virtual_packets)
+                                                     virtual_packet_flag=no_of_virtual_packets, nthreads=self.tardis_config.montecarlo.nthreads)
 
         if np.sum(montecarlo_energies < 0) == len(montecarlo_energies):
             logger.critical("No r-packet escaped through the outer boundary.")

tardis/montecarlo/montecarlo.pyx

@@ -17,30 +17,6 @@ np.import_array()
 ctypedef np.int64_t int_type_t
 
 cdef extern from "src/cmontecarlo.h":
-    ctypedef enum rpacket_status_t:
-        TARDIS_PACKET_STATUS_IN_PROCESS = 0
-        TARDIS_PACKET_STATUS_EMITTED = 1
-        TARDIS_PACKET_STATUS_REABSORBED = 2
-
-    ctypedef struct rpacket_t:
-        double nu
-        double mu
-        double energy
-        double r
-        double tau_event
-        double nu_line
-        int_type_t current_shell_id
-        int_type_t next_line_id
-        int_type_t last_line
-        int_type_t close_line
-        int_type_t recently_crossed_boundary
-        int_type_t virtual_packet_flag
-        int_type_t virtual_packet
-        double d_line
-        double d_electron
-        double d_boundary
-        rpacket_status_t next_shell_id
-
     ctypedef struct storage_model_t:
         double *packet_nus
         double *packet_mus
@@ -86,17 +62,10 @@ cdef extern from "src/cmontecarlo.h":
         double inverse_sigma_thomson
         double inner_boundary_albedo
         int_type_t reflective_inner_boundary
-        int_type_t current_packet_id
-
-    int_type_t montecarlo_one_packet(storage_model_t *storage, rpacket_t *packet, int_type_t virtual_mode)
-    int rpacket_init(rpacket_t *packet, storage_model_t *storage, int packet_index, int virtual_packet_flag)
-    double rpacket_get_nu(rpacket_t *packet)
-    double rpacket_get_energy(rpacket_t *packet)
-    void initialize_random_kit(unsigned long seed)
-
 
+    void montecarlo_main_loop(storage_model_t * storage, int_type_t virtual_packet_flag, int nthreads, unsigned long seed)
 
-def montecarlo_radial1d(model, int_type_t virtual_packet_flag=0):
+def montecarlo_radial1d(model, int_type_t virtual_packet_flag=0, int nthreads=4):
     """
     Parameters
     ----------
@@ -125,8 +94,6 @@ def montecarlo_radial1d(model, int_type_t virtual_packet_flag=0):
                     int_type_t do_scatter
     """
     cdef storage_model_t storage
-    cdef rpacket_t packet
-    initialize_random_kit(model.tardis_config.montecarlo.seed)
     cdef np.ndarray[double, ndim=1] packet_nus = model.packet_src.packet_nus
     storage.packet_nus = <double*> packet_nus.data
     cdef np.ndarray[double, ndim=1] packet_mus = model.packet_src.packet_mus
@@ -220,20 +187,9 @@ def montecarlo_radial1d(model, int_type_t virtual_packet_flag=0):
     storage.inverse_sigma_thomson = 1.0 / storage.sigma_thomson
     storage.reflective_inner_boundary = model.tardis_config.montecarlo.enable_reflective_inner_boundary
     storage.inner_boundary_albedo = model.tardis_config.montecarlo.inner_boundary_albedo
-    storage.current_packet_id = -1
     ######## Setting up the output ########
     #cdef np.ndarray[double, ndim=1] output_nus = np.zeros(storage.no_of_packets, dtype=np.float64)
     #cdef np.ndarray[double, ndim=1] output_energies = np.zeros(storage.no_of_packets, dtype=np.float64)
-    cdef int_type_t reabsorbed = 0
-    for packet_index in range(storage.no_of_packets):
-        storage.current_packet_id = packet_index
-        rpacket_init(&packet, &storage, packet_index, virtual_packet_flag)
-        if (virtual_packet_flag > 0):
-            #this is a run for which we want the virtual packet spectrum. So first thing we need to do is spawn virtual packets to track the input packet
-            reabsorbed = montecarlo_one_packet(&storage, &packet, -1)
-        #Now can do the propagation of the real packet
-        reabsorbed = montecarlo_one_packet(&storage, &packet, 0)
-        storage.output_nus[packet_index] = rpacket_get_nu(&packet)
-        storage.output_energies[packet_index] = -rpacket_get_energy(&packet) if reabsorbed == 1 else rpacket_get_energy(&packet)
+    montecarlo_main_loop(&storage, virtual_packet_flag, nthreads, model.tardis_config.montecarlo.seed)
     return output_nus, output_energies, js, nubars, last_line_interaction_in_id, last_line_interaction_out_id, last_interaction_type, last_line_interaction_shell_id
 

tardis/montecarlo/setup_package.py

@@ -2,9 +2,18 @@
 from setuptools import Extension
 import numpy as np
 import os
+from astropy_helpers.setup_helpers import get_distutils_option
 
 from glob import glob
 
+if get_distutils_option('with_openmp', ['build', 'install']) is not None:
+    compile_args = ['-fopenmp']
+    link_args = ['-fopenmp']
+    define_macros = [('WITHOPENMP', None)]
+else:
+    compile_args = []
+    link_args = []
+    define_macros = []
 
 def get_extensions():
     sources = ['tardis/montecarlo/montecarlo.pyx']
@@ -16,4 +25,7 @@ def get_extensions():
     return [Extension('tardis.montecarlo.montecarlo', sources,
                       include_dirs=['tardis/montecarlo/src',
                                     'tardis/montecarlo/src/randomkit',
-                                    np.get_include()])]
+                                    np.get_include()],
+                      extra_compile_args=compile_args,
+                      extra_link_args=link_args,
+                      define_macros=define_macros)]

tardis/montecarlo/src/cmontecarlo.c

@@ -1,3 +1,6 @@
+#ifdef WITHOPENMP
+#include <omp.h>
+#endif
 #include "cmontecarlo.h"
 
 rk_state mt_state;
@@ -292,8 +295,14 @@ move_packet (rpacket_t * packet, storage_model_t * storage, double distance)
 	{
 	  comov_energy = rpacket_get_energy (packet) * doppler_factor;
 	  comov_nu = rpacket_get_nu (packet) * doppler_factor;
+#ifdef WITHOPENMP
+#pragma omp atomic
+#endif
 	  storage->js[rpacket_get_current_shell_id (packet)] +=
 	    comov_energy * distance;
+#ifdef WITHOPENMP
+#pragma omp atomic
+#endif
 	  storage->nubars[rpacket_get_current_shell_id (packet)] +=
 	    comov_energy * distance * comov_nu;
 	}
@@ -314,6 +323,9 @@ increment_j_blue_estimator (rpacket_t * packet, storage_model_t * storage,
   doppler_factor = 1.0 - mu_interaction * r_interaction *
     storage->inverse_time_explosion * INVERSE_C;
   comov_energy = rpacket_get_energy (packet) * doppler_factor;
+#ifdef WITHOPENMP
+#pragma omp atomic
+#endif
   storage->line_lists_j_blues[j_blue_idx] +=
     comov_energy / rpacket_get_nu (packet);
 }
@@ -337,63 +349,64 @@ montecarlo_one_packet (storage_model_t * storage, rpacket_t * packet,
   else
     {
       if ((rpacket_get_nu (packet) > storage->spectrum_virt_start_nu) && (rpacket_get_nu(packet) < storage->spectrum_virt_end_nu))
-    {
-      for (i = 0; i < rpacket_get_virtual_packet_flag (packet); i++)
 	{
-	  memcpy ((void *) &virt_packet, (void *) packet, sizeof (rpacket_t));
-	  if (virt_packet.r > storage->r_inner[0])
-	    {
-	      mu_min =
-		-1.0 * sqrt (1.0 -
-			     (storage->r_inner[0] / virt_packet.r) *
-			     (storage->r_inner[0] / virt_packet.r));
-	    }
-	  else
+	  for (i = 0; i < rpacket_get_virtual_packet_flag (packet); i++)
 	    {
-	      mu_min = 0.0;
-	    }
-	  mu_bin = (1.0 - mu_min) / rpacket_get_virtual_packet_flag (packet);
-	  virt_packet.mu = mu_min + (i + rk_double (&mt_state)) * mu_bin;
-	  switch (virtual_mode)
-	    {
-	    case -2:
-	      weight = 1.0 / rpacket_get_virtual_packet_flag (packet);
-	      break;
-	    case -1:
-	      weight =
-		2.0 * virt_packet.mu /
-		rpacket_get_virtual_packet_flag (packet);
-	      break;
-	    case 1:
-	      weight =
-		(1.0 -
-		 mu_min) / 2.0 / rpacket_get_virtual_packet_flag (packet);
-	      break;
-	    default:
-	      fprintf (stderr, "Something has gone horribly wrong!\n");
-	    }
-	  doppler_factor_ratio =
-	    rpacket_doppler_factor (packet, storage) /
-	    rpacket_doppler_factor (&virt_packet, storage);
-	  virt_packet.energy =
-	    rpacket_get_energy (packet) * doppler_factor_ratio;
-	  virt_packet.nu = rpacket_get_nu (packet) * doppler_factor_ratio;
-	  reabsorbed = montecarlo_one_packet_loop (storage, &virt_packet, 1);
-	  if ((virt_packet.nu < storage->spectrum_end_nu) &&
-	      (virt_packet.nu > storage->spectrum_start_nu))
-	    {
-	      virt_id_nu =
-		floor ((virt_packet.nu -
-			storage->spectrum_start_nu) /
-		       storage->spectrum_delta_nu);
-	      storage->spectrum_virt_nu[virt_id_nu] +=
-		virt_packet.energy * weight;
+	      memcpy ((void *) &virt_packet, (void *) packet, sizeof (rpacket_t));
+	      if (virt_packet.r > storage->r_inner[0])
+		{
+		  mu_min =
+		    -1.0 * sqrt (1.0 -
+				 (storage->r_inner[0] / virt_packet.r) *
+				 (storage->r_inner[0] / virt_packet.r));
+		}
+	      else
+		{
+		  mu_min = 0.0;
+		}
+	      mu_bin = (1.0 - mu_min) / rpacket_get_virtual_packet_flag (packet);
+	      virt_packet.mu = mu_min + (i + rk_double (&mt_state)) * mu_bin;
+	      switch (virtual_mode)
+		{
+		case -2:
+		  weight = 1.0 / rpacket_get_virtual_packet_flag (packet);
+		  break;
+		case -1:
+		  weight =
+		    2.0 * virt_packet.mu /
+		    rpacket_get_virtual_packet_flag (packet);
+		  break;
+		case 1:
+		  weight =
+		    (1.0 -
+		     mu_min) / 2.0 / rpacket_get_virtual_packet_flag (packet);
+		  break;
+		default:
+		  fprintf (stderr, "Something has gone horribly wrong!\n");
+		}
+	      doppler_factor_ratio =
+		rpacket_doppler_factor (packet, storage) /
+		rpacket_doppler_factor (&virt_packet, storage);
+	      virt_packet.energy =
+		rpacket_get_energy (packet) * doppler_factor_ratio;
+	      virt_packet.nu = rpacket_get_nu (packet) * doppler_factor_ratio;
+	      reabsorbed = montecarlo_one_packet_loop (storage, &virt_packet, 1);
+	      if ((virt_packet.nu < storage->spectrum_end_nu) &&
+		  (virt_packet.nu > storage->spectrum_start_nu))
+		{
+		  virt_id_nu =
+		    floor ((virt_packet.nu -
+			    storage->spectrum_start_nu) /
+			   storage->spectrum_delta_nu);
+		  storage->spectrum_virt_nu[virt_id_nu] +=
+		    virt_packet.energy * weight;
+		}
 	    }
 	}
-    }
-  else{
-    return 1;
-  }
+      else
+	{
+	  return 1;
+	}
     }
   return reabsorbed;
 }
@@ -469,7 +482,7 @@ montecarlo_thomson_scatter (rpacket_t * packet, storage_model_t * storage,
   rpacket_set_energy (packet, comov_energy * inverse_doppler_factor);
   rpacket_reset_tau_event (packet);
   rpacket_set_recently_crossed_boundary (packet, 0);
-  storage->last_interaction_type[storage->current_packet_id] = 1;
+  storage->last_interaction_type[rpacket_get_id (packet)] = 1;
   if (rpacket_get_virtual_packet_flag (packet) > 0)
     {
       montecarlo_one_packet (storage, packet, 1);
@@ -522,11 +535,11 @@ montecarlo_line_scatter (rpacket_t * packet, storage_model_t * storage,
       inverse_doppler_factor = 1.0 / rpacket_doppler_factor (packet, storage);
       comov_energy = rpacket_get_energy (packet) * old_doppler_factor;
       rpacket_set_energy (packet, comov_energy * inverse_doppler_factor);
-      storage->last_line_interaction_in_id[storage->current_packet_id] =
+      storage->last_line_interaction_in_id[rpacket_get_id (packet)] =
 	rpacket_get_next_line_id (packet) - 1;
-      storage->last_line_interaction_shell_id[storage->current_packet_id] =
+      storage->last_line_interaction_shell_id[rpacket_get_id (packet)] =
 	rpacket_get_current_shell_id (packet);
-      storage->last_interaction_type[storage->current_packet_id] = 2;
+      storage->last_interaction_type[rpacket_get_id (packet)] = 2;
       if (storage->line_interaction_id == 0)
 	{
 	  emission_line_id = rpacket_get_next_line_id (packet) - 1;
@@ -535,7 +548,7 @@ montecarlo_line_scatter (rpacket_t * packet, storage_model_t * storage,
 	{
 	  emission_line_id = macro_atom (packet, storage);
 	}
-      storage->last_line_interaction_out_id[storage->current_packet_id] =
+      storage->last_line_interaction_out_id[rpacket_get_id (packet)] =
 	emission_line_id;
       rpacket_set_nu (packet,
 		      storage->line_list_nu[emission_line_id] *
@@ -670,3 +683,43 @@ montecarlo_one_packet_loop (storage_model_t * storage, rpacket_t * packet,
   return rpacket_get_status (packet) ==
     TARDIS_PACKET_STATUS_REABSORBED ? 1 : 0;
 }
+
+void
+montecarlo_main_loop(storage_model_t * storage, int64_t virtual_packet_flag, int nthreads, unsigned long seed)
+{
+  int64_t packet_index;
+#ifdef WITHOPENMP
+  omp_set_dynamic(0);
+  omp_set_num_threads(nthreads);
+#pragma omp parallel
+  {
+    initialize_random_kit(seed + omp_get_thread_num());
+#pragma omp for
+#else
+  initialize_random_kit(seed);
+#endif
+  for (packet_index = 0; packet_index < storage->no_of_packets; packet_index++)
+    {
+      int reabsorbed = 0;
+      rpacket_t packet;
+      rpacket_set_id(&packet, packet_index);
+      rpacket_init(&packet, storage, packet_index, virtual_packet_flag);
+      if (virtual_packet_flag > 0)
+	{
+	  reabsorbed = montecarlo_one_packet(storage, &packet, -1);
+	}
+      reabsorbed = montecarlo_one_packet(storage, &packet, 0);
+      storage->output_nus[packet_index] = rpacket_get_nu(&packet);
+      if (reabsorbed == 1)
+	{
+	  storage->output_energies[packet_index] = -rpacket_get_energy(&packet);
+	}
+      else
+	{
+	  storage->output_energies[packet_index] = rpacket_get_energy(&packet);
+	}
+    }
+#ifdef WITHOPENMP
+  }
+#endif
+}

tardis/montecarlo/src/cmontecarlo.h

@@ -104,4 +104,9 @@ int64_t montecarlo_one_packet_loop (storage_model_t * storage,
 				    rpacket_t * packet,
 				    int64_t virtual_packet);
 
+void montecarlo_main_loop(storage_model_t * storage, 
+			  int64_t virtual_packet_flag,
+			  int nthreads, 
+			  unsigned long seed);
+
 #endif // TARDIS_CMONTECARLO_H