diff --git a/pyswiftsim/structure.py b/pyswiftsim/structure.py index 6197a2241d1dfa038cbc64d5836fe261b1618dc5..666bd4d0b3028a90b539180b50561b0aff1b67c5 100644 --- a/pyswiftsim/structure.py +++ b/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)