Commits

Britton Smith committed 8c7d595

Finished implementing freefall.py script and wrappers for library functions.

Comments (0)

Files changed (5)

src/clib/freefall.C

 
   chemistry_data my_chemistry = set_default_chemistry_parameters();
   my_chemistry.use_chemistry = 1;
-  my_chemistry.primordial_chemistry = 2;
+  my_chemistry.primordial_chemistry = 3;
   my_chemistry.metal_cooling = 1;
   my_chemistry.cloudy_table_file = (char*) "solar_2008_3D_metals.h5";
 
   gr_float gravitational_constant = 4.0 * 3.1415926 * 6.6726e-8 * 
     my_units.density_units * POW(my_units.time_units, 2);
 
-  // gr_float a_value = 1.0;
-  gr_float a_value = 1.0/(1.0+10.0);
+  gr_float a_value = 1.0;
+  //  gr_float a_value = 1.0/(1.0+10.0);
 
   if (initialize_chemistry_data(my_chemistry, my_units, a_value) == FAIL) {
     fprintf(stderr, "Error in initialize_chemistry_data.\n");
     return FAIL;
   }
 
-  my_chemistry.UVbackground = 1;
-  my_chemistry.UVbackground_file = (char*) "UVB_rates_HM2012.hdf5";
-  if (initialize_UVbackground_data(my_chemistry, my_units, a_value) == FAIL) {
-    fprintf(stderr, "Error in initialize_UVbackground_data.\n");
-    return FAIL;
-  }
+  // my_chemistry.UVbackground = 1;
+  // my_chemistry.UVbackground_file = (char*) "UVB_rates_HM2012.hdf5";
+  // if (initialize_UVbackground_data(my_chemistry, my_units, a_value) == FAIL) {
+  //   fprintf(stderr, "Error in initialize_UVbackground_data.\n");
+  //   return FAIL;
+  // }
 
-  update_UVbackground_rates(my_chemistry, my_units, a_value);
+  // update_UVbackground_rates(my_chemistry, my_units, a_value);
 
   gr_float *density, *energy, *x_velocity, *y_velocity, *z_velocity;
   gr_float *HI_density, *HII_density, *HM_density,

src/python/examples/freefall.py

 
 my_chemistry = chemistry_data()
 my_chemistry.use_chemistry = 1
-my_chemistry.primordial_chemistry = 2
-#my_chemistry.metal_cooling = 1
-#my_chemistry.cloudy_table_file = "solar_2008_3D_metals.h5";
+my_chemistry.primordial_chemistry = 3
+my_chemistry.metal_cooling = 1
+my_chemistry.cloudy_table_file = "solar_2008_3D_metals.h5";
 
 my_chemistry.comoving_coordinates = 0
 my_chemistry.density_units = 1.67e-24
 my_chemistry.time_units = 1.0e12
 my_chemistry.a_units = 1.0
 
-energy_units = (my_chemistry.length_units / my_chemistry.time_units)**2.0
+energy_units = (my_chemistry.length_units /
+                my_chemistry.time_units)**2.0
 
 gravitational_constant = (4.0 * 3.1415926 * 6.6726e-8 * 
   my_chemistry.density_units * my_chemistry.time_units**2)
 
-a_value = 1.0/(1.0+10.0)
+a_value = 1.0
 
 my_chemistry.initialize(a_value)
 
 fc["de"][:] = tiny_number * fc["density"]
 fc["metal"][:] = 1.e-5 * fc["density"]
 
-temperature_units = mass_h*((my_chemistry.length_units/my_chemistry.time_units)**2)/kboltz
+freefall_constant = np.power(fc["density"][0], -0.5)
+freefall_time_constant = np.power(((32. * gravitational_constant) / 
+                                   (3. * np.pi)), 0.5)
+
+
+temperature_units = mass_h * ((my_chemistry.length_units/
+                              my_chemistry.time_units)**2) / kboltz
 
 fc["energy"][:] = 1000. / temperature_units
 fc["x-velocity"][:] = 0.0
 fc["y-velocity"][:] = 0.0
 fc["z-velocity"][:] = 0.0
-  
-calculate_temperature(fc)
-print fc["temperature"]
+
+timestep_fraction = 0.1
+current_time = 0.0
+while fc["density"][0] < 1.e10:
+    dt = timestep_fraction * \
+      np.power(((3. * np.pi) / 
+                (32. * gravitational_constant * fc["density"][0])), 0.5)
+
+    calculate_temperature(fc)
+    print fc["density"][0], fc["temperature"][0]
+    
+    density_ratio = np.power((freefall_constant - 
+                              (0.5 * freefall_time_constant * current_time)), -2.) / \
+                              fc["density"][0]
+                              
+    for field in fc.density_fields:
+        fc[field] *= density_ratio
+
+    fc["energy"][0] += (my_chemistry.Gamma - 1.) * fc["energy"][0] * \
+      freefall_time_constant * np.power(fc["density"][0], 0.5) * dt
+
+    solve_chemistry(fc, a_value, dt)
+
+    current_time += dt

src/python/pygrackle/fluid_container.py

 import numpy as np
 
+_base_fluids = ["density", "metal"]
+_nd_fields   = ["energy",
+                "x-velocity", "y-velocity", "z-velocity",
+                "temperature", "pressure",
+                "gamma", "cooling_time"]
 _fluid_names = {}
-_fluid_names[0] = ["HI", "HII", "HeI", "HeII", "HeIII", "de",
-    "x-velocity", "y-velocity", "z-velocity",
-    "temperature", "gamma", "energy", "density", "metal"]
-_fluid_names[1] = _fluid_names[0] + \
-    ["H2I", "H2II", "HM"]
+_fluid_names[1] = _base_fluids + \
+  ["HI", "HII", "HeI", "HeII", "HeIII", "de"]
 _fluid_names[2] = _fluid_names[1] + \
-    ["DI", "DII", "HDI"]
+  ["H2I", "H2II", "HM"]
+_fluid_names[3] = _fluid_names[2] + \
+  ["DI", "DII", "HDI"]
 
 class FluidContainer(dict):
-    def __init__(self, chemistry_data, n_vals, dtype="float64"):
+    def __init__(self, chemistry_data, n_vals, dtype="float64",
+                 itype="int64"):
         super(FluidContainer, self).__init__()
         self.dtype = dtype
         self.chemistry_data = chemistry_data
-        self.n_vals = n_vals
-        for fluid in _fluid_names[self.chemistry_data.primordial_chemistry]:
+        self.n_vals = n_vals        
+        for fluid in _fluid_names[self.chemistry_data.primordial_chemistry] + _nd_fields:
             self._setup_fluid(fluid)
 
     def _setup_fluid(self, fluid_name):
         self[fluid_name] = np.zeros(self.n_vals, self.dtype)
+
+    @property
+    def density_fields(self):
+        return _fluid_names[self.chemistry_data.primordial_chemistry]

src/python/pygrackle/grackle_defs.pxd

                 gr_float *internal_energy,
                 gr_float *x_velocity,
                 gr_float *y_velocity,
-                gr_float  *z_velocity,
+                gr_float *z_velocity,
                 gr_float *HI_density,
                 gr_float *HII_density,
                 gr_float *HM_density,
                 gr_float *internal_energy,
                 gr_float *x_velocity,
                 gr_float *y_velocity,
-                gr_float  *z_velocity,
+                gr_float *z_velocity,
                 gr_float *HI_density,
                 gr_float *HII_density,
                 gr_float *HM_density,

src/python/pygrackle/grackle_wrapper.pyx

 from grackle_defs cimport *
+import numpy as np
 cimport numpy as np
 
 cdef class chemistry_data:
         def __get__(self):
             return self.data.cloudy_table_file
         def __set__(self, val):
-            raise NotImplementedError
+            self.data.cloudy_table_file = val
 
     property comoving_coordinates:
         def __get__(self):
     else:
         return <gr_float *> rv.data
 
-
-def calculate_temperature(fc):
+def solve_chemistry(fc, my_a, my_dt):
     cdef gr_int grid_rank = 1
     cdef gr_int grid_dimension
+    grid_dimension = fc["density"].shape[0]
+
+    cdef np.ndarray ref_gs, ref_ge
+    ref_gs = np.zeros(3, dtype="int64")
+    ref_ge = np.zeros(3, dtype="int64")
+    cdef gr_int *grid_start, *grid_end
+    grid_start = <gr_int *> ref_gs.data
+    grid_end = <gr_int *> ref_ge.data
+
+    cdef gr_float a_value = <gr_float> my_a
+    cdef gr_float dt_value = <gr_float> my_dt
+
     cdef chemistry_data chem_data = fc.chemistry_data
     cdef c_chemistry_data my_chemistry = chem_data.data
     cdef c_code_units my_units = chem_data.units
+    cdef gr_float *density = get_field(fc, "density")
+    cdef gr_float *internal_energy = get_field(fc, "energy")
+    cdef gr_float *x_velocity = get_field(fc, "x-velocity")
+    cdef gr_float *y_velocity = get_field(fc, "y-velocity")
+    cdef gr_float *z_velocity = get_field(fc, "z-velocity")
+    cdef gr_float *HI_density = get_field(fc, "HI")
+    cdef gr_float *HII_density = get_field(fc, "HII")
+    cdef gr_float *HM_density = get_field(fc, "HM")
+    cdef gr_float *HeI_density = get_field(fc, "HeI")
+    cdef gr_float *HeII_density = get_field(fc, "HeII")
+    cdef gr_float *HeIII_density = get_field(fc, "HeIII")
+    cdef gr_float *H2I_density = get_field(fc, "H2I")
+    cdef gr_float *H2II_density = get_field(fc, "H2II")
+    cdef gr_float *DI_density = get_field(fc, "DI")
+    cdef gr_float *DII_density = get_field(fc, "DII")
+    cdef gr_float *HDI_density = get_field(fc, "HDI")
+    cdef gr_float *e_density = get_field(fc, "de")
+    cdef gr_float *metal_density = get_field(fc, "metal")
+
+    c_solve_chemistry (
+                my_chemistry,
+                my_units,
+                a_value,
+                dt_value,
+                grid_rank,
+                &grid_dimension,
+                grid_start,
+                grid_end,
+                density,
+                internal_energy,
+                x_velocity,
+                y_velocity,
+                z_velocity,
+                HI_density,
+                HII_density,
+                HM_density,
+                HeI_density,
+                HeII_density,
+                HeIII_density,
+                H2I_density,
+                H2II_density,
+                DI_density,
+                DII_density,
+                HDI_density,
+                e_density,
+                metal_density)
+    
+def calculate_cooling_time(fc, my_a, my_dt):
+    cdef gr_int grid_rank = 1
+    cdef gr_int grid_dimension
     grid_dimension = fc["density"].shape[0]
+
+    cdef np.ndarray ref_gs, ref_ge
+    ref_gs = np.zeros(3, dtype="int64")
+    ref_ge = np.zeros(3, dtype="int64")
+    cdef gr_int *grid_start, *grid_end
+    grid_start = <gr_int *> ref_gs.data
+    grid_end = <gr_int *> ref_ge.data
+
+    cdef gr_float a_value = <gr_float> my_a
+    cdef gr_float dt_value = <gr_float> my_dt
+
+    cdef chemistry_data chem_data = fc.chemistry_data
+    cdef c_chemistry_data my_chemistry = chem_data.data
+    cdef c_code_units my_units = chem_data.units
+    cdef gr_float *density = get_field(fc, "density")
+    cdef gr_float *internal_energy = get_field(fc, "energy")
+    cdef gr_float *x_velocity = get_field(fc, "x-velocity")
+    cdef gr_float *y_velocity = get_field(fc, "y-velocity")
+    cdef gr_float *z_velocity = get_field(fc, "z-velocity")
+    cdef gr_float *HI_density = get_field(fc, "HI")
+    cdef gr_float *HII_density = get_field(fc, "HII")
+    cdef gr_float *HM_density = get_field(fc, "HM")
+    cdef gr_float *HeI_density = get_field(fc, "HeI")
+    cdef gr_float *HeII_density = get_field(fc, "HeII")
+    cdef gr_float *HeIII_density = get_field(fc, "HeIII")
+    cdef gr_float *H2I_density = get_field(fc, "H2I")
+    cdef gr_float *H2II_density = get_field(fc, "H2II")
+    cdef gr_float *DI_density = get_field(fc, "DI")
+    cdef gr_float *DII_density = get_field(fc, "DII")
+    cdef gr_float *HDI_density = get_field(fc, "HDI")
+    cdef gr_float *e_density = get_field(fc, "de")
+    cdef gr_float *metal_density = get_field(fc, "metal")
+    cdef gr_float *cooling_time = get_field(fc, "cooling_time")
+    
+    c_calculate_cooling_time (
+                my_chemistry,
+                my_units,
+                a_value,
+                dt_value,
+                grid_rank,
+                &grid_dimension,
+                grid_start,
+                grid_end,
+                density,
+                internal_energy,
+                x_velocity,
+                y_velocity,
+                z_velocity,
+                HI_density,
+                HII_density,
+                HM_density,
+                HeI_density,
+                HeII_density,
+                HeIII_density,
+                H2I_density,
+                H2II_density,
+                DI_density,
+                DII_density,
+                HDI_density,
+                e_density,
+                metal_density,
+                cooling_time)
+    
+def calculate_gamma(fc):
+    cdef gr_int grid_rank = 1
+    cdef gr_int grid_dimension
+    grid_dimension = fc["density"].shape[0]
+    cdef chemistry_data chem_data = fc.chemistry_data
+    cdef c_chemistry_data my_chemistry = chem_data.data
+    cdef c_code_units my_units = chem_data.units
     cdef gr_float *density = get_field(fc, "density")
     cdef gr_float *internal_energy = get_field(fc, "energy")
     cdef gr_float *HI_density = get_field(fc, "HI")
     cdef gr_float *DII_density = get_field(fc, "DII")
     cdef gr_float *HDI_density = get_field(fc, "HDI")
     cdef gr_float *e_density = get_field(fc, "de")
-    cdef gr_float *metal_density = get_field(fc, "metal_density")
+    cdef gr_float *metal_density = get_field(fc, "metal")
+    cdef gr_float *gamma = get_field(fc, "gamma")
+    
+    c_calculate_gamma (
+                my_chemistry,
+                my_units,
+                grid_rank,
+                &grid_dimension,
+                density,
+                internal_energy,
+                HI_density,
+                HII_density,
+                HM_density,
+                HeI_density,
+                HeII_density,
+                HeIII_density,
+                H2I_density,
+                H2II_density,
+                DI_density,
+                DII_density,
+                HDI_density,
+                e_density,
+                metal_density,
+                gamma)
+    
+def calculate_pressure(fc):
+    cdef gr_int grid_rank = 1
+    cdef gr_int grid_dimension
+    grid_dimension = fc["density"].shape[0]
+    cdef chemistry_data chem_data = fc.chemistry_data
+    cdef c_chemistry_data my_chemistry = chem_data.data
+    cdef c_code_units my_units = chem_data.units
+    cdef gr_float *density = get_field(fc, "density")
+    cdef gr_float *internal_energy = get_field(fc, "energy")
+    cdef gr_float *HI_density = get_field(fc, "HI")
+    cdef gr_float *HII_density = get_field(fc, "HII")
+    cdef gr_float *HM_density = get_field(fc, "HM")
+    cdef gr_float *HeI_density = get_field(fc, "HeI")
+    cdef gr_float *HeII_density = get_field(fc, "HeII")
+    cdef gr_float *HeIII_density = get_field(fc, "HeIII")
+    cdef gr_float *H2I_density = get_field(fc, "H2I")
+    cdef gr_float *H2II_density = get_field(fc, "H2II")
+    cdef gr_float *DI_density = get_field(fc, "DI")
+    cdef gr_float *DII_density = get_field(fc, "DII")
+    cdef gr_float *HDI_density = get_field(fc, "HDI")
+    cdef gr_float *e_density = get_field(fc, "de")
+    cdef gr_float *metal_density = get_field(fc, "metal")
+    cdef gr_float *pressure = get_field(fc, "pressure")
+    
+    c_calculate_pressure (
+                my_chemistry,
+                my_units,
+                grid_rank,
+                &grid_dimension,
+                density,
+                internal_energy,
+                HI_density,
+                HII_density,
+                HM_density,
+                HeI_density,
+                HeII_density,
+                HeIII_density,
+                H2I_density,
+                H2II_density,
+                DI_density,
+                DII_density,
+                HDI_density,
+                e_density,
+                metal_density,
+                pressure)
+
+def calculate_temperature(fc):
+    cdef gr_int grid_rank = 1
+    cdef gr_int grid_dimension
+    grid_dimension = fc["density"].shape[0]
+    cdef chemistry_data chem_data = fc.chemistry_data
+    cdef c_chemistry_data my_chemistry = chem_data.data
+    cdef c_code_units my_units = chem_data.units
+    cdef gr_float *density = get_field(fc, "density")
+    cdef gr_float *internal_energy = get_field(fc, "energy")
+    cdef gr_float *HI_density = get_field(fc, "HI")
+    cdef gr_float *HII_density = get_field(fc, "HII")
+    cdef gr_float *HM_density = get_field(fc, "HM")
+    cdef gr_float *HeI_density = get_field(fc, "HeI")
+    cdef gr_float *HeII_density = get_field(fc, "HeII")
+    cdef gr_float *HeIII_density = get_field(fc, "HeIII")
+    cdef gr_float *H2I_density = get_field(fc, "H2I")
+    cdef gr_float *H2II_density = get_field(fc, "H2II")
+    cdef gr_float *DI_density = get_field(fc, "DI")
+    cdef gr_float *DII_density = get_field(fc, "DII")
+    cdef gr_float *HDI_density = get_field(fc, "HDI")
+    cdef gr_float *e_density = get_field(fc, "de")
+    cdef gr_float *metal_density = get_field(fc, "metal")
     cdef gr_float *temperature = get_field(fc, "temperature")
 
     c_calculate_temperature(