Add Grackle struct (still need some debugging) and improved str

pyswiftsim/structure.py

@@ -148,17 +148,14 @@ 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)
+            if isinstance(d, SwiftStruct):
+                txt += d.__str__(tab+"\t")
+            else:
+                txt += tab + "\t%s: %s\n" % (name, d)
         return txt
 
     def __getattr__(self, name):
@@ -181,7 +178,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 +249,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)
         
         
     def getArray(self):
@@ -414,6 +415,75 @@ 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):
+    """
+    WARNING, class quickly done, may contain error when accessing it directly
+    """
+    _format = "iiiiipidiidiiiiiiidddiiddiddiiddddddiiiiiid"
+    _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',
+        "omp"
+    ]
+
+    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 +496,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)