#!/usr/bin/env python ############################################################################### # This file is part of SWIFT. # Copyright (c) 2018 Folkert Nobels (nobels@strw.leidenuniv.nl) # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Lesser General Public License as published # by the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU Lesser General Public License # along with this program. If not, see . # ################################################################################ import numpy as np import h5py import matplotlib.pyplot as plt from matplotlib.colors import BoundaryNorm from matplotlib.ticker import MaxNLocator import scipy.optimize as sco import os def linearfunc(x, a, b): return a * x + b def radialfunc(r, h, A): return A * np.exp(-r / h) * r def verticalfunc(z, A, z0, zoff): return 2 * A * np.exp(-(z - zoff) / z0) def verticalfunc2(z, A, z0): return 2 * A * np.exp(-(z) / z0) def verticalfunc3(z, A, z0, zoff, b): return 2 * A * np.exp(-(z - zoff) / z0) + b Nmax = 2001 steps = 10 storefits = False logfit = True normalfit = False # if the user wants to store the indivudal fits if storefits: if not os.path.exists("radial"): os.mkdir("radial") os.mkdir("vertical") os.mkdir("histsnap") # Initialize the arrays R_ideal = np.linspace(0, 40, 100) Z_ideal = np.linspace(0, 10, 100) iterarray = np.arange(0, Nmax + 1, steps) Z0t = np.zeros(len(iterarray)) Z0terr = np.zeros(len(iterarray)) h0t = np.zeros(len(iterarray)) h0terr = np.zeros(len(iterarray)) Ar = np.zeros(len(iterarray)) Arerr = np.zeros(len(iterarray)) Az = np.zeros(len(iterarray)) Azerr = np.zeros(len(iterarray)) time_array = np.zeros(len(iterarray)) ar = np.zeros(len(iterarray)) arerr = np.zeros(len(iterarray)) br = np.zeros(len(iterarray)) brerr = np.zeros(len(iterarray)) az = np.zeros(len(iterarray)) azerr = np.zeros(len(iterarray)) bz = np.zeros(len(iterarray)) bzerr = np.zeros(len(iterarray)) eps = 1e-6 for i in iterarray: # Getting the data from the snapshots f = h5py.File("output_%04d.hdf5" % i, "r") boxsize = f["Header"].attrs["BoxSize"] / 2.0 time_array[int(i / steps)] = f["Header"].attrs["Time"] particles = f["PartType4"] coordinates = particles["Coordinates"][:, :] masses = particles["Masses"][:] R = ( (coordinates[:, 0] - boxsize[0]) ** 2 + (coordinates[:, 1] - boxsize[1]) ** 2 ) ** 0.5 Z = np.abs(coordinates[:, 1] - boxsize[1]) # Bin the coordinates to make them suitable for fitting Rhist = np.histogram(R, bins=100, range=[0, 40], normed=True) Zhist = np.histogram(Z, bins=100, range=[0, 10.0], normed=True) # Create correct variables for fitting Ry = Rhist[0] Rx = (Rhist[1][1:] + Rhist[1][: len(Rhist[0])]) / 2.0 Zy = Zhist[0] Zx = (Zhist[1][1:] + Zhist[1][: len(Zhist[0])]) / 2.0 # Fit with two methods: non-linear LSQ and linear LSQ in log space bestsolR = sco.curve_fit(radialfunc, Rx[10:], Ry[10:], p0=[2.0, 0.2]) bestsolZ = sco.curve_fit(verticalfunc, Zx[40:], Zy[40:]) bestsolRlog = sco.curve_fit(linearfunc, Rx[10:], np.log10(Ry[10:] + eps)) bestsolZlog = sco.curve_fit(linearfunc, Zx[40:], np.log10(Zy[40:] + eps)) # Store variables h0t[int(i / steps)] = bestsolR[0][0] Z0t[int(i / steps)] = bestsolZ[0][1] Ar[int(i / steps)] = bestsolR[0][1] Az[int(i / steps)] = bestsolZ[0][0] Z0terr[int(i / steps)] = (bestsolZ[1][1, 1]) ** 0.5 h0terr[int(i / steps)] = (bestsolR[1][0, 0]) ** 0.5 Arerr[int(i / steps)] = (bestsolR[1][1, 1]) ** 0.5 Azerr[int(i / steps)] = (bestsolZ[1][0, 0]) ** 0.5 ar[int(i / steps)] = bestsolRlog[0][0] arerr[int(i / steps)] = (bestsolRlog[1][0, 0]) ** 0.5 br[int(i / steps)] = bestsolRlog[0][1] brerr[int(i / steps)] = (bestsolRlog[1][1, 1]) ** 0.5 az[int(i / steps)] = bestsolZlog[0][0] azerr[int(i / steps)] = (bestsolZlog[1][0, 0]) ** 0.5 bz[int(i / steps)] = bestsolZlog[0][1] bzerr[int(i / steps)] = (bestsolZlog[1][1, 1]) ** 0.5 if storefits: plt.step(Rx, Ry) plt.plot( R_ideal, radialfunc(R_ideal, bestsolR[0][0], bestsolR[0][1]), label="Non linear LSQ", ) plt.plot( R_ideal, 10 ** (linearfunc(R_ideal, bestsolRlog[0][0], bestsolRlog[0][1])), label="Linear LSQ", ) plt.xlim(0, 40) plt.ylim(0, 0.25) plt.xlabel("R (kpc)") plt.ylabel("Probability") plt.savefig("./radial/radialsnap%04d.png" % i) plt.close() plt.step(Zx, Zy) plt.plot( Z_ideal, verticalfunc(Z_ideal, bestsolZ[0][0], bestsolZ[0][1], bestsolZ[0][2]), label="Non linear LSQ", ) plt.plot( Z_ideal, 10 ** (linearfunc(Z_ideal, bestsolZlog[0][0], bestsolZlog[0][1])), label="Linear LSQ", ) plt.xlim(0, 10.0) plt.ylim(0, 0.6) plt.xlabel("z (kpc)") plt.ylabel("Probability") plt.savefig("./vertical/verticalsnap%04d.png" % i) plt.close() time_array[-1] = 2.0 ax = plt.subplot(111) ax.set_yscale("log") if logfit: plt.errorbar( time_array, np.absolute(az / (az[0]) - 1), yerr=azerr / (az[0]), label="z0 scale height (Log space)", ) plt.errorbar( time_array, np.absolute(ar / (ar[0]) - 1), yerr=arerr / (ar[0]), label="h scale lenght (Log space)", ) if normalfit: plt.errorbar( time_array, np.absolute(Z0t / (Z0t[0]) - 1), yerr=Z0terr / (Z0t[0]), label="z0 scale height (normal space)", ) plt.errorbar( time_array, np.absolute(h0t / (h0t[0]) - 1), yerr=h0terr / (h0t[0]), label="h scale height (normal space)", ) ax.set_xlabel("Time (Gyr)") ax.set_ylabel("Fractional difference") plt.legend() plt.savefig("Fitdifference-witherror.pdf") plt.close() ax = plt.subplot(111) ax.set_yscale("log") if logfit: plt.plot( time_array, np.absolute(az / (az[0]) - 1), label="z0 scale height (Log space)" ) plt.plot( time_array, np.absolute(ar / (ar[0]) - 1), label="h scale lenght (Log space)" ) if normalfit: plt.plot( time_array, np.absolute(Z0t / (Z0t[0]) - 1), label="z0 scale height (normal space)", ) plt.plot( time_array, np.absolute(h0t / (h0t[0]) - 1), label="h scale height (normal space)", ) ax.set_xlabel("Time (Gyr)") ax.set_ylabel("Fractional difference") plt.legend() plt.savefig("Fitdifference.pdf") plt.show()