Cooling rate example ready

examples/cooling_rate/cooling.yml

@@ -45,6 +45,8 @@ GrackleCooling:
   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.

examples/cooling_rate/generate_grackle_data.py

@@ -32,6 +32,8 @@ from pygrackle.utilities.physical_constants import \
 
 filename = "grackle.hdf5"
 
+debug = 1
+
 
 def generate_data(primordial_chemistry):
     current_redshift = 0.
@@ -67,16 +69,44 @@ def generate_data(primordial_chemistry):
                                metal_mass_fraction=0.01295,
                                temperature=temperature,
                                density=density,
-                               converge=True)
+                               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"
+        print("Updating Dataset")
         data = f[data_name]
         f.pop(data_name)
     else:
-        print "Creating Dataset"
+        print("Creating Dataset")
 
     data = f.create_group(data_name)
 
@@ -107,30 +137,54 @@ def generate_data(primordial_chemistry):
     dset = tmp.create_dataset("Density", fc["density"].shape,
                               dtype=fc["density"].dtype)
     dset[:] = fc["density"]
-    # Temperautre
+    # Temperature
     dset = tmp.create_dataset("Temperature", temperature.shape,
                               dtype=temperature.dtype)
     dset[:] = temperature
 
-    old = deepcopy(fc)
-    fc.solve_chemistry(dt)
-    rate = (fc["energy"] - old["energy"]) / dt
+    write_fractions(tmp, fc, primordial_chemistry)
 
-    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"]
+    return f, data
 
-    f.close()
+
+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
+        print("Computing Primordial Chemistry %i" % i)
         generate_data(i)

examples/cooling_rate/plot_cooling.py

@@ -14,10 +14,11 @@ plt.rc('text', usetex=True)
 # PARAMETERS
 
 # grackle primordial chemistry
-primordial_chemistry = 0
+primordial_chemistry = 1
 
 # reference data
 grackle_filename = "grackle.hdf5"
+compute_equilibrium = True
 
 # swift params filename
 filename = "cooling.yml"
@@ -35,7 +36,7 @@ else:
     default_density = np.logspace(-3, 1, N_rho)
 
 # temperature in K
-N_T = 1000
+N_T = 10
 default_temperature = np.logspace(1, 9, N_T)
 
 # time step in s
@@ -48,6 +49,35 @@ 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 = {}
@@ -98,6 +128,10 @@ def read_grackle_data(filename, us, primordial_chemistry):
     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"]
 
@@ -112,24 +146,26 @@ def read_grackle_data(filename, us, primordial_chemistry):
 
 
 def initialize_swift(filename):
+    print("Initialization of SWIFT")
     d = {}
+
     # parse swift params
-    print("Reading parameters")
     params = wrapper.parserReadFile(filename)
     d["params"] = params
+
     # init units
-    print("Initialization of the unit system")
     us, pconst = wrapper.unitSystemInit(params)
     d["us"] = us
     d["pconst"] = pconst
+
     # init cooling
-    print("Initialization of the 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"]
@@ -161,10 +197,12 @@ def plot_solution(rate, data, us):
     if N_rho == 1:
         # plot Lambda vs T
         plt.figure()
-        plt.loglog(T, np.abs(lam), label="SWIFT")
+        plt.loglog(T, np.abs(lam),
+                   label="SWIFT, %s" % wrapper.configGetCooling())
         if ref:
             # plot reference
-            plt.loglog(T, np.abs(grackle_lam), label="Ref.")
+            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}$]")
@@ -222,15 +260,30 @@ if __name__ == "__main__":
     cooling = d["cooling"]
 
     if isfile(grackle_filename):
-        d = read_grackle_data(filename, us, primordial_chemistry)
+        d = read_grackle_data(grackle_filename, us, primordial_chemistry)
     else:
         d = generate_default_initial_condition(us, pconst)
 
     # du / dt
     print("Computing cooling...")
-    rate = wrapper.coolingRate(pconst, us, cooling,
-                               d["density"].astype(np.float32),
-                               d["energy"].astype(np.float32),
-                               d["time_step"])
+    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)

examples/cooling_rate/run.sh

@@ -14,4 +14,4 @@ then
     ./generate_grackle_data.py
 fi
 
-./test_cooling.py
+./plot_cooling.py