diff --git a/README b/Readme.md similarity index 100% rename from README rename to Readme.md diff --git a/test/cooling.yml b/examples/cooling_rate/cooling.yml similarity index 80% rename from test/cooling.yml rename to examples/cooling_rate/cooling.yml index 96555e68c238897b81b856a624468385f0f9af90..54c4cf3c2ecdfe33e16cb28ebe16243206ecffac 100644 --- a/test/cooling.yml +++ b/examples/cooling_rate/cooling.yml @@ -37,10 +37,16 @@ InitialConditions: # Cooling with Grackle 2.0 GrackleCooling: - GrackleCloudyTable: CloudyData_UVB=HM2012.h5 # Name of the Cloudy Table - UVbackground: 1 # Enable or not the UV background - GrackleRedshift: 0 # Redshift to use (-1 means time based redshift) - GrackleHSShieldingDensityThreshold: 1.1708e-26 # self shielding threshold in internal system of units + CloudyTable: CloudyData_UVB=HM2012.h5 # Name of the Cloudy Table + WithUVbackground: 1 # Enable or not the UV background + Redshift: 0 # Redshift to use (-1 means time based redshift) + WithMetalCooling: 1 # Enable or not the metal cooling + ProvideVolumetricHeatingRates: 0 # User provide volumetric heating rates + ProvideSpecificHeatingRates: 0 # User provide specific heating rates + SelfShieldingMethod: 0 # Grackle (<= 3) or Gear self shielding method + OutputMode: 0 # Write in output corresponding primordial chemistry mode + ConvergenceLimit: 1e-3 + MaxSteps: 100000 Gravity: eta: 0.025 # Constant dimensionless multiplier for time integration. diff --git a/examples/cooling_rate/generate_grackle_data.py b/examples/cooling_rate/generate_grackle_data.py new file mode 100644 index 0000000000000000000000000000000000000000..3ebb364bf9d227113b781cbf068ea2c86290219f --- /dev/null +++ b/examples/cooling_rate/generate_grackle_data.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python +######################################################################## +# +# Cooling rate data generation. This code is a modified version +# of a code developed by the Grackle Team. +# +# +# Copyright (c) 2013-2016, Grackle Development Team. +# +# Distributed under the terms of the Enzo Public Licence. +# +# The full license is in the file LICENSE, distributed with this +# software. +######################################################################## + +""" +Generate (or update) a hdf5 file containing +the grackle data for each primordial chemistry +""" + +import numpy as np +from h5py import File +from copy import deepcopy + +from pygrackle import \ + chemistry_data, \ + setup_fluid_container + +from pygrackle.utilities.physical_constants import \ + mass_hydrogen_cgs, \ + sec_per_Myr + +filename = "grackle.hdf5" + +debug = 1 + + +def generate_data(primordial_chemistry): + current_redshift = 0. + + # Set solver parameters + my_chemistry = chemistry_data() + my_chemistry.use_grackle = 1 + my_chemistry.with_radiative_cooling = 1 + my_chemistry.primordial_chemistry = primordial_chemistry + my_chemistry.metal_cooling = 1 + my_chemistry.UVbackground = 1 + my_chemistry.self_shielding_method = 0 + my_chemistry.H2_self_shielding = 0 + my_chemistry.grackle_data_file = "CloudyData_UVB=HM2012.h5" + + # Set units + my_chemistry.comoving_coordinates = 0 # proper units + my_chemistry.a_units = 1.0 + my_chemistry.a_value = 1.0 / (1.0 + current_redshift) / \ + my_chemistry.a_units + my_chemistry.length_units = 3.085e21 + my_chemistry.density_units = 1.989e43 / my_chemistry.length_units**3 + my_chemistry.velocity_units = 20725573.785998672 + my_chemistry.time_units = my_chemistry.length_units / \ + my_chemistry.velocity_units + + density = mass_hydrogen_cgs + dt = 1e-8 * sec_per_Myr / my_chemistry.time_units + # Call convenience function for setting up a fluid container. + # This container holds the solver parameters, units, and fields. + temperature = np.logspace(1, 9, 1000) + fc = setup_fluid_container(my_chemistry, + metal_mass_fraction=0.01295, + temperature=temperature, + density=density, + converge=True, + tolerance=1e-5, + max_iterations=1e8) + + f, data = write_input( + fc, my_chemistry, temperature, dt, primordial_chemistry) + + old = deepcopy(fc) + fc.solve_chemistry(dt) + rate = (fc["energy"] - old["energy"]) / dt + + write_output(f, data, fc, rate) + + +def write_output(f, data, fc, rate): + data.create_group("Output") + tmp = data["Output"] + # Rate + dset = tmp.create_dataset("Rate", rate.shape, + dtype=rate.dtype) + dset[:] = rate + # Energy + dset = tmp.create_dataset("Energy", fc["energy"].shape, + dtype=fc["energy"].dtype) + dset[:] = fc["energy"] + + f.close() + + +def write_input(fc, my_chemistry, temperature, dt, primordial_chemistry): + f = File(filename, "a") + data_name = "PrimordialChemistry%i" % primordial_chemistry + if data_name in f: + print("Updating Dataset") + data = f[data_name] + f.pop(data_name) + else: + print("Creating Dataset") + + data = f.create_group(data_name) + + # Units + data.create_group("Units") + tmp = data["Units"] + tmp.attrs["Length"] = my_chemistry.length_units + tmp.attrs["Density"] = my_chemistry.density_units + tmp.attrs["Velocity"] = my_chemistry.velocity_units + tmp.attrs["Time"] = my_chemistry.time_units + + # Parameters + data.create_group("Params") + tmp = data["Params"] + tmp.attrs["MetalCooling"] = my_chemistry.metal_cooling + tmp.attrs["UVBackground"] = my_chemistry.UVbackground + tmp.attrs["SelfShieldingMethod"] = my_chemistry.self_shielding_method + tmp.attrs["TimeStep"] = dt + + # Inputs + data.create_group("Input") + tmp = data["Input"] + # energy + dset = tmp.create_dataset("Energy", fc["energy"].shape, + dtype=fc["energy"].dtype) + dset[:] = fc["energy"] + # density + dset = tmp.create_dataset("Density", fc["density"].shape, + dtype=fc["density"].dtype) + dset[:] = fc["density"] + # Temperature + dset = tmp.create_dataset("Temperature", temperature.shape, + dtype=temperature.dtype) + dset[:] = temperature + + write_fractions(tmp, fc, primordial_chemistry) + + return f, data + + +def write_fractions(tmp, fc, primordial_chemistry): + fields = get_fields(primordial_chemistry) + for i in fields: + dset = tmp.create_dataset(i, fc[i].shape, + dtype=fc[i].dtype) + dset[:] = fc[i] / fc["density"] + + +def get_fields(primordial_chemistry): + fields = [ + "metal" + ] + if primordial_chemistry > 0: + fields.extend([ + "HI", + "HII", + "HeI", + "HeII", + "HeIII", + "de"]) + + if primordial_chemistry > 1: + fields.extend([ + "HM", + "H2I", + "H2II" + ]) + + if primordial_chemistry > 2: + fields.extend([ + "DI", + "DII", + "HDI" + ]) + return fields + + +if __name__ == "__main__": + for i in range(4): + print("Computing Primordial Chemistry %i" % i) + generate_data(i) diff --git a/examples/cooling_rate/plot_cooling.py b/examples/cooling_rate/plot_cooling.py new file mode 100644 index 0000000000000000000000000000000000000000..1cf9aae0de8a26b35591506e3ab289d17933c53a --- /dev/null +++ b/examples/cooling_rate/plot_cooling.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 + +from pyswiftsim import wrapper + +from copy import deepcopy +import numpy as np +import matplotlib.pyplot as plt +from astropy import units +from os.path import isfile +from h5py import File + +plt.rc('text', usetex=True) + +# PARAMETERS + +# grackle primordial chemistry +primordial_chemistry = 1 + +# reference data +grackle_filename = "grackle.hdf5" +compute_equilibrium = True + +# swift params filename +filename = "cooling.yml" + +# if grackle_filename does not exist +# use following values + +# density in atom / cm3 +N_rho = 1 +# with N_rho > 1, the code is not implemented to deal +# with the reference +if N_rho == 1: + default_density = np.array([1.]) +else: + default_density = np.logspace(-3, 1, N_rho) + +# temperature in K +N_T = 10 +default_temperature = np.logspace(1, 9, N_T) + +# time step in s +default_dt = units.Myr * 1e-8 +default_dt = default_dt.to("s") / units.s + +# adiabatic index +gamma = 5. / 3. + +# SCRIPT + + +def get_fields(primordial_chemistry): + fields = [ + "metal" + ] + if primordial_chemistry > 0: + fields.extend([ + "HI", + "HII", + "HeI", + "HeII", + "HeIII", + "de"]) + + if primordial_chemistry > 1: + fields.extend([ + "HM", + "H2I", + "H2II" + ]) + + if primordial_chemistry > 2: + fields.extend([ + "DI", + "DII", + "HDI" + ]) + return fields + + +def generate_default_initial_condition(us, pconst): + print("Generating default initial conditions") + d = {} + # generate grid + rho, T = np.meshgrid(default_density, default_temperature) + rho = deepcopy(rho.reshape(-1)) + T = T.reshape(-1) + d["temperature"] = T + + # Deal with units + rho *= us.UnitLength_in_cgs**3 * pconst.const_proton_mass + d["density"] = rho + + energy = pconst.const_boltzmann_k * T / us.UnitTemperature_in_cgs + energy /= (gamma - 1.) * pconst.const_proton_mass + d["energy"] = energy + + dt = default_dt / us.UnitTime_in_cgs + d["time_step"] = dt + + return d + + +def read_grackle_data(filename, us, primordial_chemistry): + print("Reading initial conditions") + f = File(filename, "r") + data = f["PrimordialChemistry%i" % primordial_chemistry] + + # read units + tmp = data["Units"].attrs + + u_len = tmp["Length"] / us.UnitLength_in_cgs + u_den = tmp["Density"] * us.UnitLength_in_cgs**3 / us.UnitMass_in_cgs + u_time = tmp["Time"] / us.UnitTime_in_cgs + + # read input + tmp = data["Input"] + + energy = tmp["Energy"][:] * u_len**2 / u_time**2 + d["energy"] = energy + + T = tmp["Temperature"][:] / us.UnitTemperature_in_cgs + d["temperature"] = T + + density = tmp["Density"][:] * u_den + d["density"] = density + + dt = data["Params"].attrs["TimeStep"] * u_time + d["time_step"] = dt + + # read fractions + for i in get_fields(primordial_chemistry): + d[i] = tmp[i][:] + + # read output + tmp = data["Output"] + + energy = tmp["Energy"][:] * u_len**2 / u_time**2 + d["out_energy"] = energy + + rate = tmp["Rate"][:] * u_len**2 / (u_time**3) + d["rate"] = rate + + f.close() + return d + + +def initialize_swift(filename): + print("Initialization of SWIFT") + d = {} + + # parse swift params + params = wrapper.parserReadFile(filename) + d["params"] = params + + # init units + us, pconst = wrapper.unitSystemInit(params) + d["us"] = us + d["pconst"] = pconst + + # init cooling + cooling = wrapper.coolingInit(params, us, pconst) + d["cooling"] = cooling + return d + + +def plot_solution(rate, data, us): + print("Plotting solution") + energy = data["energy"] + rho = data["density"] + T = data["temperature"] + + ref = False + if "rate" in data: + ref = True + grackle_rate = data["rate"] + + # change units => cgs + rho *= us.UnitMass_in_cgs / us.UnitLength_in_cgs**3 + + T *= us.UnitTemperature_in_cgs + + energy *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**2 + + rate *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**3 + + if ref: + grackle_rate *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**3 + + # lambda cooling + lam = rate * rho + + if ref: + grackle_lam = grackle_rate * rho + + # do plot + if N_rho == 1: + # plot Lambda vs T + plt.figure() + plt.loglog(T, np.abs(lam), + label="SWIFT, %s" % wrapper.configGetCooling()) + if ref: + # plot reference + plt.loglog(T, np.abs(grackle_lam), + label="Grackle, Prim. Chem. %i" % primordial_chemistry) + plt.legend() + plt.xlabel("Temperature [K]") + plt.ylabel("$\\Lambda$ [erg s$^{-1}$ cm$^{3}$]") + + # plot error vs T + plt.figure() + plt.plot(T, (lam - grackle_lam) / grackle_lam) + plt.gca().set_xscale("log") + plt.xlabel("Temperature [K]") + plt.ylabel( + r"$(\Lambda - \Lambda_\textrm{ref}) / \Lambda_\textrm{ref}$") + + plt.show() + + else: + shape = [N_rho, N_T] + cooling_time = energy / rate + cooling_length = np.sqrt(gamma * (gamma-1.) * energy) * cooling_time + + cooling_length = np.log10(np.abs(cooling_length) / units.kpc.to('cm')) + + # reshape + rho = rho.reshape(shape) + T = T.reshape(shape) + energy = energy.reshape(shape) + cooling_length = cooling_length.reshape(shape) + + _min = -7 + _max = 7 + N_levels = 100 + levels = np.linspace(_min, _max, N_levels) + plt.figure() + plt.contourf(rho, T, cooling_length, levels) + plt.xlabel("Density [atom/cm3]") + plt.ylabel("Temperature [K]") + + ax = plt.gca() + ax.set_xscale("log") + ax.set_yscale("log") + + cbar = plt.colorbar() + tc = np.arange(_min, _max, 1.5) + cbar.set_ticks(tc) + cbar.set_ticklabels(tc) + + plt.show() + + +if __name__ == "__main__": + + d = initialize_swift(filename) + pconst = d["pconst"] + us = d["us"] + params = d["params"] + cooling = d["cooling"] + + if isfile(grackle_filename): + d = read_grackle_data(grackle_filename, us, primordial_chemistry) + else: + d = generate_default_initial_condition(us, pconst) + + # du / dt + print("Computing cooling...") + rate = np.zeros(d["density"].shape) + + if compute_equilibrium: + rate = wrapper.coolingRate(pconst, us, cooling, + d["density"].astype(np.float32), + d["energy"].astype(np.float32), + d["time_step"]) + else: + fields = get_fields(primordial_chemistry) + N = len(fields) + frac = np.zeros([len(d["density"]), N]) + for i in range(N): + frac[:, i] = d[fields[i]] + + rate = wrapper.coolingRate(pconst, us, cooling, + d["density"].astype(np.float32), + d["energy"].astype(np.float32), + d["time_step"], + frac.astype(np.float32)) + + plot_solution(rate, d, us) diff --git a/examples/cooling_rate/run.sh b/examples/cooling_rate/run.sh new file mode 100644 index 0000000000000000000000000000000000000000..3564de07f4f0751a51f025fc018bf5ed5c07e64c --- /dev/null +++ b/examples/cooling_rate/run.sh @@ -0,0 +1,17 @@ +#!/bin/bash + +# Get the Grackle cooling table +if [ ! -e CloudyData_UVB=HM2012.h5 ] +then + echo "Fetching the Cloudy tables required by Grackle..." + ./getCoolingTable.sh +fi + +# Generate Grackle data if not present +if [ ! -e grackle.hdf5 ] +then + echo "Generating Grackle Data..." + ./generate_grackle_data.py +fi + +./plot_cooling.py diff --git a/pyswiftsim/structure.py b/pyswiftsim/structure.py index 2e89c0c5e6a3427d0ec38b0ede0d8903160e4376..78a35b781e6062e2f133bbcdccb5d726ae67b440 100644 --- a/pyswiftsim/structure.py +++ b/pyswiftsim/structure.py @@ -1,13 +1,13 @@ from pyswiftsim import wrapper import struct -import numpy from ctypes import * PARSER_MAX_LINE_SIZE = 256 PARSER_MAX_NO_OF_PARAMS = 256 PARSER_MAX_NO_OF_SECTIONS = 64 + ###################################################################### # # # SwiftStruct # @@ -33,10 +33,10 @@ class SwiftStruct(struct.Struct): """ def __init__(self, struct_format, data, parent): super().__init__(struct_format) - self.parent = parent # parent for ArrayStruct + self.parent = parent # parent for ArrayStruct if isinstance(data, bytes): - self.data = data # bytes string data + self.data = data # bytes string data elif isinstance(data, dict): tmp = [] for i in self.struct_name: @@ -46,7 +46,7 @@ class SwiftStruct(struct.Struct): else: raise ValueError("Data should be either bytes or dict, " "received ", type(data)) - + @property def struct_name(self): """ @@ -71,7 +71,6 @@ class SwiftStruct(struct.Struct): """ return {} - def _getInfoFromName(self, name): """ Compute index, format and size from an attribute name. @@ -102,9 +101,8 @@ class SwiftStruct(struct.Struct): if n > 1: i = slice(i, i+n) - + return i, form, n - @property def struct_size_format(self): @@ -116,7 +114,7 @@ class SwiftStruct(struct.Struct): out_nber: list number of element for each attribute - + out_form: list format for each attribute """ @@ -139,7 +137,7 @@ class SwiftStruct(struct.Struct): if v == "s": count = 1 - + count = int(count) out_nber.append(count) out_form.append(v) @@ -147,7 +145,6 @@ class SwiftStruct(struct.Struct): return out_nber, out_form - def __str__(self, tab=""): txt = tab + "%s:\n" % type(self) for name in self.struct_name: @@ -163,7 +160,7 @@ class SwiftStruct(struct.Struct): # case where the attribute is not in the structure if name not in self.struct_name: return object.__getattribute__(self, name) - + # case where the attribute is in the structure else: i, form, n = self._getInfoFromName(name) @@ -187,7 +184,7 @@ class SwiftStruct(struct.Struct): # other case => array else: return data[i] - + else: # transform scalar -> vector nform = str(n) + form @@ -198,7 +195,6 @@ class SwiftStruct(struct.Struct): data = struct.pack(nform, *data) return ArrayStruct(nform, data, self, name) - def __setattr__(self, name, value): # case where the attribute is not in the structure if name not in self.struct_name: @@ -225,7 +221,7 @@ class ArrayStruct(SwiftStruct): _name = [ "array_data" ] - + def __init__(self, struct_format, data, parent, name): super().__init__(struct_format, data, parent) self._format = struct_format @@ -235,7 +231,7 @@ class ArrayStruct(SwiftStruct): data = self.unpack(self.data) data = self._clean(data) return data[ii] - + def __setitem__(self, ii, value): data = list(self.unpack(self.data)) data[ii] = value @@ -254,8 +250,7 @@ class ArrayStruct(SwiftStruct): data = self.unpack(self.data) data = self._clean(data) return tab + str(data) + "\n" - - + def getArray(self): return self.unpack(self.data) @@ -263,7 +258,7 @@ class ArrayStruct(SwiftStruct): def struct_format(self): return self._format - + ###################################################################### # # # UnitSystem # @@ -297,6 +292,7 @@ class ChemistryPartData(SwiftStruct): def __init__(self, data, parent=None): super().__init__(self.struct_format, data, parent) + ###################################################################### # # # Part # @@ -326,8 +322,8 @@ class Part(SwiftStruct): def __init__(self, data, parent=None): super().__init__(self.struct_format, data, parent) print("ERROR, need to fix density/time_bin") - - + + ###################################################################### # # # Parameter # @@ -343,10 +339,10 @@ class Parameter(SwiftStruct): "value" ] - def __init__(self, data, parent=None): super().__init__(self.struct_format, data, parent) - + + ###################################################################### # # # Section # @@ -363,6 +359,7 @@ class Section(SwiftStruct): def __init__(self, data, parent=None): super().__init__(self.struct_format, data, parent) + ###################################################################### # # # SwiftParams # @@ -402,6 +399,7 @@ class SwiftParams(SwiftStruct): "data_params": param } + ###################################################################### # # # PhysConst # @@ -428,13 +426,15 @@ class PhysConst(SwiftStruct): "const_earth_mass", ] - def __init__(self, data, parent=None): super().__init__(self.struct_format, data, parent) class GrackleCodeUnits(SwiftStruct): + cooling_type = wrapper.configGetCooling() _format = "idddddd" + if cooling_type == "grackle_float": + _format = _format.replace("d", "f") _name = [ "comoving_coordinates", "density_units", @@ -447,9 +447,13 @@ class GrackleCodeUnits(SwiftStruct): def __init__(self, data, parent=None): super().__init__(self.struct_format, data, parent) - + + class GrackleChemistryData(SwiftStruct): + cooling_type = wrapper.configGetCooling() _format = "iiiiiPidiidiiiiiiidddiiddiddiiddddddiiiiii" + if cooling_type == "grackle_float": + _format = _format.replace("d", "f") _name = [ 'use_grackle', 'with_radiative_cooling', @@ -498,6 +502,7 @@ class GrackleChemistryData(SwiftStruct): 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": @@ -509,7 +514,7 @@ class CoolingFunctionData(SwiftStruct): "min_energy", "cooling_tstep_mult" ] - elif cooling_type == "grackle": + elif "grackle" in cooling_type: _format = "200cidd{code_units}s{chemistry}s".format( code_units=struct.calcsize(GrackleCodeUnits._format), chemistry=struct.calcsize(GrackleChemistryData._format) @@ -543,6 +548,5 @@ class CoolingFunctionData(SwiftStruct): raise ValueError( "Cooling Type %s not implemented" % cooling_type) - def __init__(self, data, parent=None): super().__init__(self.struct_format, data, parent) diff --git a/setup.py b/setup.py index b87b13a808514f6cf666837b5570337f8a721e28..6f8712d727e90d1af93e599ed618d67ad68da860 100644 --- a/setup.py +++ b/setup.py @@ -1,28 +1,36 @@ #!/usr/bin/env python3 -descr = """ -Wrapper around the SPH cosmological simulation code SWIFT -""" - from setuptools import setup, find_packages, Extension import sys import os from glob import glob import numpy +descr = """ +Wrapper around the SPH cosmological simulation code SWIFT +""" + +with_omp = True os.environ["CC"] = "mpicc" -cflags = ["-Werror", - "-Wall", - "-Wextra", - # disables some warnings due to python - "-Wno-unused-parameter", - "-Wno-strict-prototypes", - "-Wno-unused-function", - "-Wno-incompatible-pointer-types", - "-Wno-missing-field-initializers", +cflags = [ + "-Werror", + "-Wall", + "-Wextra", + # disables some warnings due to python + "-Wno-unused-parameter", + "-Wno-strict-prototypes", + "-Wno-unused-function", + "-Wno-incompatible-pointer-types", + "-Wno-missing-field-initializers", + "-fopenmp" ] +lflags = [ + "-fopenmp" + ] + + # deal with arguments def parseCmdLine(arg, store=False): ret = False @@ -37,11 +45,12 @@ def parseCmdLine(arg, store=False): sys.argv.remove(arg) return ret - + + swift_path = parseCmdLine("--with-swift", store=True) # python lib dependency -install_requires=["numpy"] +install_requires = ["numpy"] def getValueFromMakefile(swift_root, value): @@ -51,7 +60,7 @@ def getValueFromMakefile(swift_root, value): with open(makefile, "r") as f: for line in f.readlines(): if value == line[:N]: - return line[N:-1] # remove \n + return line[N:-1] # remove \n raise ValueError("Value %s not found in Makefile" % value) @@ -80,9 +89,10 @@ if swift_path: include.append(grackle_inc) # C libraries -lib = ["m", - "swiftsim", - "hdf5", +lib = [ + "m", + "swiftsim", + "hdf5", ] lib_dir = [] @@ -90,7 +100,7 @@ lib_dir = [] if swift_path: lib_dir.append(swift_path + "/src/.libs") lib_dir.append(hdf5_root + "/lib") - + # src files c_src = [] @@ -99,7 +109,7 @@ data_files = [] ############## -## C ext +# C ext ############## c_src = glob("src/*.c") @@ -108,45 +118,41 @@ ext_modules = Extension("pyswiftsim.wrapper", include_dirs=include, libraries=lib, library_dirs=lib_dir, - extra_compile_args=cflags) + extra_compile_args=cflags, + extra_link_args=lflags) ext_modules = [ext_modules] - + ############## -## data +# data ############## data_files = [] ############## -## scripts +# scripts ############## list_scripts = [] ############## -## Setup +# Setup ############## setup( - name = "pyswiftsim", - version = "0.1", - author = "Hausammann Loic", - author_email = "loic.hausammann@epfl.ch", - description = descr, - license = "GPLv3", - keywords = "nbody sph simulation hpc", - url = "", - - packages = find_packages(), - - data_files = data_files, - - scripts = list_scripts, - - install_requires = install_requires, - - dependency_links = dependency_links, - - ext_modules = ext_modules, + name="pyswiftsim", + version="0.1", + author="Hausammann Loic", + author_email="loic.hausammann@epfl.ch", + description=descr, + license="GPLv3", + keywords="nbody sph simulation hpc", + url="", + + packages=find_packages(), + data_files=data_files, + scripts=list_scripts, + install_requires=install_requires, + dependency_links=dependency_links, + ext_modules=ext_modules, ) diff --git a/src/config_wrapper.h b/src/config_wrapper.h index 6b9b2a8e752bba7826b699d76050e946c57c4f1c..5a17aeb3c933dc9b36a7884bc11125c7a2b5c9ef 100644 --- a/src/config_wrapper.h +++ b/src/config_wrapper.h @@ -10,11 +10,25 @@ */ PyObject* config_get_cooling() { char *cooling_name; + /* lambda */ #ifdef COOLING_CONST_LAMBDA cooling_name = "const_lambda"; + + /* grackle */ #elif defined(COOLING_GRACKLE) +#if COOLING_GRACKLE_MODE == 0 cooling_name = "grackle"; -#endif +#elif COOLING_GRACKLE_MODE == 1 + cooling_name = "grackle1"; +#elif COOLING_GRACKLE_MODE == 2 + cooling_name = "grackle2"; +#elif COOLING_GRACKLE_MODE == 3 + cooling_name = "grackle3"; +#else + error("Grackle mode unknown"); +#endif // COOLING_GRACKLE_MODE +#endif // COOLING_GRACKLE + return PyUnicode_FromString(cooling_name); }; diff --git a/src/cooling_wrapper.c b/src/cooling_wrapper.c index cb23170b544e455b181f8007277b7d55ace4ffa8..11f5bdb97a6d694e8f3c396553034d31bc723015 100644 --- a/src/cooling_wrapper.c +++ b/src/cooling_wrapper.c @@ -1,7 +1,18 @@ #include "pyswiftsim_tools.h" #include "cooling_wrapper.h" - - +#include <omp.h> + + +/** + * @brief Initialize the cooling + * + * args is expecting pyswiftsim classes in the following order: + * SwiftParams, UnitSystem and PhysConst. + * + * @param self calling object + * @param args arguments + * @return CoolingFunctionData + */ PyObject* pycooling_init(PyObject* self, PyObject* args) { PyObject *pyparams; PyObject *pyus; @@ -35,6 +46,52 @@ PyObject* pycooling_init(PyObject* self, PyObject* args) { return pycooling; } +/** + * @brief Set the cooling element fractions + * + * @param xp The #xpart to set + * @param frac The numpy array containing the fractions (id, element) + * @param idx The id (in frac) of the particle to set + */ +void pycooling_set_fractions(struct xpart *xp, PyArrayObject* frac, const int idx) { + struct cooling_xpart_data *data = &xp->cooling_data; + data->metal_frac = *(float*)PyArray_GETPTR2(frac, idx, 0); + +#ifdef COOLING_GRACKLE +#if COOLING_GRACKLE_MODE > 0 + data->HI_frac = *(float*)PyArray_GETPTR2(frac, idx, 1); + data->HII_frac = *(float*)PyArray_GETPTR2(frac, idx, 2); + data->HeI_frac = *(float*)PyArray_GETPTR2(frac, idx, 3); + data->HeII_frac = *(float*)PyArray_GETPTR2(frac, idx, 4); + data->HeIII_frac = *(float*)PyArray_GETPTR2(frac, idx, 5); + data->e_frac = *(float*)PyArray_GETPTR2(frac, idx, 6); +#endif // COOLING_GRACKLE_MODE +#if COOLING_GRACKLE_MODE > 1 + data->HM_frac = *(float*)PyArray_GETPTR2(frac, idx, 7); + data->H2I_frac = *(float*)PyArray_GETPTR2(frac, idx, 8); + data->H2II_frac = *(float*)PyArray_GETPTR2(frac, idx, 9); +#endif // COOLING_GRACKLE_MODE +#if COOLING_GRACKLE_MODE > 2 + data->DI_frac = *(float*)PyArray_GETPTR2(frac, idx, 10); + data->DII_frac = *(float*)PyArray_GETPTR2(frac, idx, 11); + data->HDI_frac = *(float*)PyArray_GETPTR2(frac, idx, 12); +#endif // COOLING_GRACKLE_MODE +#endif // COOLING_GRACKLE + +} + +/** + * @brief Compute cooling rate + * + * args is expecting pyswiftsim classes in the following order: + * PhysConst, UnitSystem and CoolingFunctionData. + * Then two numpy arrays (density and specific energy) and an optional + * float for the time step + * + * @param self calling object + * @param args arguments + * @return cooling rate + */ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { import_array(); @@ -44,35 +101,40 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { PyArrayObject *rho; PyArrayObject *energy; + PyArrayObject *fractions = NULL; float dt = 1e-3; /* parse argument */ if (!PyArg_ParseTuple(args, - "OOOOO|f", + "OOOOO|fOO", &pypconst, &pyus, &pycooling, &rho, &energy, - &dt)) + &dt, + &fractions + )) return NULL; /* check numpy array */ if (pytools_check_array(energy, 1, NPY_FLOAT) != SUCCESS) - { - return NULL; - } + return NULL; if (pytools_check_array(rho, 1, NPY_FLOAT) != SUCCESS) - { - return NULL; - } + return NULL; + + if (fractions != NULL && + pytools_check_array(fractions, 2, NPY_FLOAT) != SUCCESS) + return NULL; if (PyArray_DIM(energy, 0) != PyArray_DIM(rho, 0)) - { - pyerror("Density and energy should have the same dimension"); - } + pyerror("Density and energy should have the same dimension"); + + if (fractions != NULL && + PyArray_DIM(fractions, 0) != PyArray_DIM(rho,0)) + pyerror("Fractions should have the same first dimension than the others"); size_t N = PyArray_DIM(energy, 0); @@ -99,9 +161,13 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { #endif /* return object */ - PyArrayObject *rate = PyArray_NewLikeArray(energy, NPY_ANYORDER, NULL, 1); + PyArrayObject *rate = PyArray_SimpleNew(PyArray_NDIM(energy), PyArray_DIMS(energy), NPY_FLOAT); + /* Release GIL */ + Py_BEGIN_ALLOW_THREADS; + /* loop over all particles */ +#pragma omp for for(size_t i = 0; i < N; i++) { /* set particle data */ @@ -109,8 +175,11 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { float u = *(float*) PyArray_GETPTR1(energy, i); p.entropy = gas_entropy_from_internal_energy(p.rho, u); - cooling_first_init_part(&p, &xp, cooling); - + if (fractions != NULL) + pycooling_set_fractions(&xp, fractions, i); + else + cooling_first_init_part(&p, &xp, cooling); + /* compute cooling rate */ float *tmp = PyArray_GETPTR1(rate, i); #ifdef COOLING_GRACKLE @@ -120,6 +189,9 @@ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args) { #endif } + /* Acquire GIL */ + Py_END_ALLOW_THREADS; + return rate; } diff --git a/src/cooling_wrapper.h b/src/cooling_wrapper.h index 4fbdd8abbc746d5d42aa4e32baa707c7770aa602..95717e8ac5d2867b25ddcebe33548fe86ecfa6cc 100644 --- a/src/cooling_wrapper.h +++ b/src/cooling_wrapper.h @@ -3,30 +3,8 @@ #include "pyswiftsim_tools.h" -/** - * @brief Initialize the cooling - * - * args is expecting pyswiftsim classes in the following order: - * SwiftParams, UnitSystem and PhysConst. - * - * @param self calling object - * @param args arguments - * @return CoolingFunctionData - */ PyObject* pycooling_init(PyObject* self, PyObject* args); -/** - * @brief Compute cooling rate - * - * args is expecting pyswiftsim classes in the following order: - * PhysConst, UnitSystem and CoolingFunctionData. - * Then two numpy arrays (density and specific energy) and an optional - * float for the time step - * - * @param self calling object - * @param args arguments - * @return cooling rate - */ PyArrayObject* pycooling_rate(PyObject* self, PyObject* args); #endif // __PYSWIFTSIM_COOLING_H__ diff --git a/src/wrapper.c b/src/wrapper.c index 12e7ef2b12653c9648e91bdf1c727cddb6993395..46076c5cb627eaed51654bb02107fcaeefcbbfb6 100644 --- a/src/wrapper.c +++ b/src/wrapper.c @@ -30,7 +30,21 @@ static PyMethodDef wrapper_methods[] = { "Initialize cooling."}, {"coolingRate", pycooling_rate, METH_VARARGS, - "Compute the cooling rate."}, + "Compute the cooling rate.\n\n" + "Parameters\n" + "----------\n\n" + "pyconst: swift physical constant\n" + "pyus: swift unit system\n" + "cooling: swift cooling structure\n" + "rho: np.array\n" + "\t Mass density in pyus units\n" + "energy: np.array\n" + "\t Internal energy in pyus units\n" + "dt: float, optional\n" + "\t Time step in pyus units\n" + "fractions: np.array, optional\n" + "\t Fraction of each cooling element (including metals)\n" + }, {"configGetCooling", config_get_cooling, METH_VARARGS, "Get the cooling type."}, diff --git a/test/test_cooling.py b/test/test_cooling.py deleted file mode 100644 index f8bb2aeca7e0c7a2e9e71cf209681711c4aff66a..0000000000000000000000000000000000000000 --- a/test/test_cooling.py +++ /dev/null @@ -1,141 +0,0 @@ -#!/usr/bin/env python3 - -from pyswiftsim import wrapper -from pyswiftsim import structure - -from copy import deepcopy -import numpy as np -import matplotlib.pyplot as plt -from astropy import units - -# -# parameters -# - -# adiabatic index -gamma = 5. / 3. - -# swift params filename -filename = "test/cooling.yml" - -# number of points -N_rho = 1 -N_T = 100 - -# density in atom / cm3 -if N_rho == 1: - rho = np.array([1.]) -else: - #rho = np.array([1.]) - rho = np.logspace(-6, 4, N_rho) - -# temperature in K -T = np.logspace(1, 9, N_T) - -# time step -dt = units.Myr * 1e-5 -dt = dt.to("s") / units.s - - - -# generate grid -rho, T = np.meshgrid(rho, T) -shape = rho.shape -rho = deepcopy(rho.reshape(-1)) -T = T.reshape(-1) - -# -# swift init -# - -# parse swift params -print("Reading parameters") -params = wrapper.parserReadFile(filename) -# init units -print("Initialization of the unit system") -us, pconst = wrapper.unitSystemInit(params) -# init cooling -print("Initialization of the cooling") -cooling = wrapper.coolingInit(params, us, pconst) - - -# -# Deal with units -# - -# change units of rho and T -# rho -rho *= us.UnitLength_in_cgs**3 * pconst.const_proton_mass -# specific energy -energy = pconst.const_boltzmann_k * T / us.UnitTemperature_in_cgs -energy /= (gamma - 1.) * pconst.const_proton_mass -# time step -dt /= us.UnitTime_in_cgs - -# -# compute rate -# - -# du / dt -print("Computing cooling...") -rate = wrapper.coolingRate(pconst, us, cooling, - rho.astype(np.float32), - energy.astype(np.float32), - dt) -print("Computing done") - -# -# plot -# - - -# change units => cgs -energy *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**2 - -rate *= us.UnitLength_in_cgs**2 / us.UnitTime_in_cgs**3 - -if N_rho == 1 or N_T == 1: - plt.figure() - n = rho / (pconst.const_proton_mass * us.UnitLength_in_cgs**3) - proton_mass_in_cgs = pconst.const_proton_mass * us.UnitMass_in_cgs - lambda_ = rate * proton_mass_in_cgs / (n) - - if N_rho == 1: - plt.loglog(T, np.abs(lambda_)) - plt.xlabel("Temperature [K]") - else: - plt.loglog(rho, np.abs(lambda_)) - plt.xlabel("Density [atom / cm3]") - plt.ylabel("Rate") - -else: - cooling_time = energy / rate - cooling_length = np.sqrt(gamma * (gamma-1.) * energy) * cooling_time - - cooling_length = np.log10(np.abs(cooling_length) / units.kpc.to('cm')) - - # reshape - rho = rho.reshape(shape) - T = T.reshape(shape) - energy = energy.reshape(shape) - cooling_length = cooling_length.reshape(shape) - - _min = -7 - _max = 7 - N_levels = 100 - levels = np.linspace(_min, _max, N_levels) - plt.figure() - plt.contourf(rho, T, cooling_length, levels) - plt.xlabel("Density [atom/cm3]") - plt.ylabel("Temperature [K]") - - ax = plt.gca() - ax.set_xscale("log") - ax.set_yscale("log") - - cbar = plt.colorbar() - tc = np.arange(_min, _max, 1.5) - cbar.set_ticks(tc) - cbar.set_ticklabels(tc) - -plt.show()