Merge branch 'init' into 'master'

Update cooling with Grackle3

See merge request !2

pyswiftsim/dev.py

+#!/usr/bin/env python3
+
+symbol_table = {
+    " int ": "i",
+    " double ": "d",
+    " char ": "c",
+}
+
+def _generateFormatString(filename):
+    """
+    Open a file containing a single struct and generate a format string
+    from it. (for dev, no guarantee of working).
+    """
+
+    struct_format = ""
+    struct_attribute = []
+    
+    with open(filename, "r") as f:
+        line = ""
+        # skip until reaching struct
+        while ("struct" not in line):
+            line = f.readline()
+
+        count = 0
+        in_struct = False
+
+        while (not in_struct or count != 0):
+            # count { and }
+            count += line.count("{")
+            if (count > 0):
+                in_struct = True
+            count -= line.count("}")
+            line = f.readline()
+
+            for k in symbol_table.keys():
+                if k not in line:
+                    continue
+
+                i = line.index(k) + len(k)
+
+                if (line[i] == "*"):
+                    struct_format += "p"
+                    i += 1
+                else:
+                    struct_format += symbol_table[k]
+
+                j = line.index(";")
+                struct_attribute.append(line[i:j])
+                
+
+    return struct_format, struct_attribute
+
+            
+if __name__ == "__main__":
+
+    filename = "chemistry_data.h"
+    struct_format, struct_attribute = _generateFormatString(filename)
+
+    attr = "[\n"
+    for i in struct_attribute:
+        attr += "\t'%s',\n" % i
+    attr += "]"
+        
+    txt = """
+    _format = "{form}"
+    _name = {attr}
+    """.format(attr=attr,
+               form=struct_format)
+    print(txt)

pyswiftsim/structure.py

@@ -148,17 +148,15 @@ class SwiftStruct(struct.Struct):
         return out_nber, out_form
 
 
-    def __str__(self):
-        tab = ""
-        parent = self.parent
-        while parent is not None:
-            tab += "\t"
-            parent = parent.parent
-             
+    def __str__(self, tab=""):
         txt = tab + "%s:\n" % type(self)
         for name in self.struct_name:
             d = getattr(self, name)
-            txt += tab + "\t%s: %s\n" % (name, d)
+            txt += tab + "\t%s: " % name
+            if isinstance(d, SwiftStruct):
+                txt += d.__str__(tab+"\t")
+            else:
+                txt += "%s\n" % d
         return txt
 
     def __getattr__(self, name):
@@ -181,7 +179,11 @@ class SwiftStruct(struct.Struct):
                     for j in range(d["size"]):
                         data_tmp = data[i][size*j:(j+1)*size]
                         tmp.append(cl(data_tmp, parent=self))
-                    return tmp
+
+                    if d["size"] == 1:
+                        return tmp[0]
+                    else:
+                        return tmp
                 # other case => array
                 else:
                     return data[i]
@@ -248,10 +250,10 @@ class ArrayStruct(SwiftStruct):
 
         return data
 
-    def __str__(self):
+    def __str__(self, tab=""):
         data = self.unpack(self.data)
         data = self._clean(data)
-        return str(data)
+        return tab + str(data) + "\n"
         
         
     def getArray(self):
@@ -414,6 +416,71 @@ class PhysConst(SwiftStruct):
         super().__init__(self.struct_format, data, parent)
 
 
+class GrackleUnits(SwiftStruct):
+    _format = "idddddd"
+    _name = [
+        "comoving_coordinates",
+        "density_units",
+        "length_units",
+        "time_units",
+        "velocity_units",
+        "a_units",
+        "a_value"
+    ]
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+    
+class GrackleChemistryData(SwiftStruct):
+    _format = "iiiiiPidiidiiiiiiidddiiddiddiiddddddiiiiii"
+    _name = [
+        'use_grackle',
+        'with_radiative_cooling',
+        'primordial_chemistry',
+        'metal_cooling',
+        'UVbackground',
+        'grackle_data_file',
+        'cmb_temperature_floor',
+        'Gamma',
+        'h2_on_dust',
+        'photoelectric_heating',
+        'photoelectric_heating_rate',
+        'use_volumetric_heating_rate',
+        'use_specific_heating_rate',
+        'three_body_rate',
+        'cie_cooling',
+        'h2_optical_depth_approximation',
+        'ih2co',
+        'ipiht',
+        'HydrogenFractionByMass',
+        'DeuteriumToHydrogenRatio',
+        'SolarMetalFractionByMass',
+        'NumberOfTemperatureBins',
+        'CaseBRecombination',
+        'TemperatureStart',
+        'TemperatureEnd',
+        'NumberOfDustTemperatureBins',
+        'DustTemperatureStart',
+        'DustTemperatureEnd',
+        'Compton_xray_heating',
+        'LWbackground_sawtooth_suppression',
+        'LWbackground_intensity',
+        'UVbackground_redshift_on',
+        'UVbackground_redshift_off',
+        'UVbackground_redshift_fullon',
+        'UVbackground_redshift_drop',
+        'cloudy_electron_fraction_factor',
+        'use_radiative_transfer',
+        'radiative_transfer_coupled_rate_solver',
+        'radiative_transfer_intermediate_step',
+        'radiative_transfer_hydrogen_only',
+        'self_shielding_method',
+        'H2_self_shielding',
+    ]
+
+    def __init__(self, data, parent=None):
+        super().__init__(self.struct_format, data, parent)
+
 class CoolingFunctionData(SwiftStruct):
     cooling_type = wrapper.configGetCooling()
     if cooling_type == "const_lambda":
@@ -426,16 +493,39 @@ class CoolingFunctionData(SwiftStruct):
             "cooling_tstep_mult"
         ]
     elif cooling_type == "grackle":
-        _format = "200cidd"
+        _format = "200cidd{units}s{chemistry}s".format(
+            units=struct.calcsize(GrackleUnits._format),
+            chemistry=struct.calcsize(GrackleChemistryData._format)
+        )
         _name = [
-            "GrackleCloudyTable",
-            "UVbackground",
-            "GrackleRedshift",
-            "GrackleHSShieldingDensityThreshold"
+            "cloudy_table",
+            "uv_background",
+            "redshift",
+            "density_self_shielding",
+            "units",
+            "chemistry"
         ]
+
+        @property
+        def struct_substruct(self):
+            chemistry = {
+                "class": GrackleChemistryData,
+                "size": 1
+            }
+
+            units = {
+                "class": GrackleUnits,
+                "size": 1
+            }
+            return {
+                "units": units,
+                "chemistry": chemistry
+            }
+
     else:
         raise ValueError(
             "Cooling Type %s not implemented" % cooling_type)
 
+
     def __init__(self, data, parent=None):
         super().__init__(self.struct_format, data, parent)

src/cooling_wrapper.c

@@ -99,6 +99,10 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) {
 
   struct part p;
 
+#ifdef COOLING_GRACKLE
+  grackle_data = &cooling->chemistry;
+#endif
+
   /* return object */
   PyArrayObject *rate = PyArray_NewLikeArray(energy, NPY_ANYORDER, NULL, 1);