Added script trying to decompose the stellar population

three_components.py

@@ -0,0 +1,140 @@
+#pylint: disable=W0401,W0614,W0622
+#All pylint codes: http://pylint.pycqa.org/en/latest/technical_reference/features.html
+
+from pylab import *
+import h5py
+from density_center import def_dc
+import ellipsoids
+import scipy.optimize
+
+# Simumation parameters
+h0 = 0.6774
+
+# Halo parameters
+file_name = 'data/subhalo_411321.hdf5'
+# centers_file_name = 'data/centers_411321.hdf5'
+
+### Snapshot parameters ###
+snapshot = 99
+a        = 1.0 # Should read this value from somewhere!
+
+# Read the centre from separate file
+# DISABLED: we calculate on our own.
+# f = h5py.File(centers_file_name, 'r')
+# X_center = f[str(snapshot)]['Coordinates'][...]
+# V_center = f[str(snapshot)]['Velocities'][...]
+# f.close()
+
+# Dictionary of particle types
+particle_types = {}
+particle_types['gas']   = '0'
+particle_types['dm']    = '1'
+particle_types['stars'] = '4'
+particle_types['bhs']   = '5'
+
+# Read stars
+particle_type = 'stars'
+with h5py.File(file_name, 'r') as f:
+    m = f[str(snapshot)][particle_types[particle_type]]['Masses'][...]
+    X = f[str(snapshot)][particle_types[particle_type]]['Coordinates'][...] * a / h0
+    V = f[str(snapshot)][particle_types[particle_type]]['Velocities'][...] * sqrt(a)
+M_tot = sum(m)
+
+# Calculate density centre and shift appropriately
+X_center_new, V_center_new = def_dc(m, X, V)
+X -= X_center_new
+V -= V_center_new
+
+# Rotate such that the short axis is z and medium axis is y
+r = linalg.norm(X, axis=1)
+rh = median(r)
+mask = r < 2*rh
+Q = ellipsoids.quadrupole_tensor(*X[mask].T, m[mask])
+eigenvalues, eigenvectors = np.linalg.eig(Q)
+R = ellipsoids.rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues)
+X_new = (R @ X.T).T
+x, y, z = X_new.T
+
+# d is the axial distance
+d = sqrt(x**2 + y**2)
+
+# Cread a two-dimensional grid
+# The grid sizes don't _have to_ be equal in both directions, but there is some
+# logic in keeping them the same.
+n_grid = 16
+d_max = 2*median(d)
+z_max = 2*median(abs(z))
+d_grid = linspace(0, d_max, n_grid+1)
+z_grid = linspace(0, z_max, n_grid+1)
+
+# Calculates the mass in each (d,z)-cell
+# We fold negative z-values assuming symmetry with respect to the xy-plane
+values, _, _ = histogram2d(d, abs(z), bins=[d_grid, z_grid], weights=m)
+values = values.T # Needed because how histogram2d workds
+
+# Calculate the volume of each (d,z)-cell (cylinder subtraction)
+d_edges, z_edges = meshgrid(d_grid, z_grid, indexing='xy')
+volumes = pi*(d_edges[1:,1:]**2 - d_edges[1:,:-1]**2)*(z_edges[1:,1:] - z_edges[:-1,1:])
+
+# Finally we have the density as a function of d and z. The normalization is to
+# make the numbers easier to work with for the minimization routine.
+rho_measured_normalized = values/volumes/M_tot
+
+# Define Plummer and Miyamoto-Nagai density
+rho_plummer = lambda r, M, b: (3*M/(4*pi*b**3))*(1+(r/b)**2)**(-2.5)
+rho_mn = lambda d, z, M, a, b: \
+    (b**2 * M / (4*pi)) * \
+    (a*d**2 + (a + 3*sqrt(z**2+b**2))*(a+sqrt(z**2+b**2))**2) / \
+    ((d**2 + (a+sqrt(z**2+b**2))**2)**2.5 * (z**2+b**2)**1.5)
+
+# The minimization procedure
+means = lambda arr: .5*(arr[:-1]+arr[1:]) # small helper function
+# Define a grid of the centre of each cell from the histogram we created earlier
+dd, zz = meshgrid(means(d_grid), means(z_grid))
+def cost(args):
+    f_plummer, b_plummer, a_mn, b_mn = args
+    f_mn = 1 - f_plummer
+    rho = rho_plummer(sqrt(dd**2+zz**2), f_plummer, b_plummer) + rho_mn(dd, zz, f_mn, a_mn, b_mn)
+    square_diff = (rho - rho_measured_normalized)**2
+    return sum(square_diff)
+minimization_result = \
+    scipy.optimize.minimize(cost, [0.5, rh, median(abs(z)), median(d)],
+    method='Nelder-Mead', tol=1e-6, options={'maxiter':5000})
+f_plummer, b_plummer, a_mn, b_mn = minimization_result.x
+print(f'f_plummer = {f_plummer:.4f}   b_plummer = {b_plummer:.4f} kpc   a_mn = {a_mn:.4f} kpc   b_mn = {b_mn:.4f} kpc')
+print(f'M_plummer = {f_plummer*M_tot:.2e} MSun   M_mn = {(1-f_plummer)*M_tot:.2e} MSun')
+
+# Compare with Matteo's results
+with h5py.File(file_name, 'r') as f:
+    particle_id = f[str(snapshot)][particle_types[particle_type]]['ParticleIDs'][...]
+    matteo_id = f[str(snapshot)][particle_types[particle_type]]['IDs_truncated'][...]
+    matteo_component_tag = f[str(snapshot)][particle_types[particle_type]]['Component'][...]
+
+# Find the array indices in the original arrays that appear in Matteo's list
+i = particle_id.argsort()
+matteo_disk_ids = matteo_id[matteo_component_tag==1]
+i_matteo_disk = i[searchsorted(particle_id[i], matteo_disk_ids)]
+matteo_bulge_ids = matteo_id[matteo_component_tag==2]
+i_matteo_bulge = i[searchsorted(particle_id[i], matteo_bulge_ids)]
+
+# Print Matteo's results
+M_matteo_bulge = sum(m[i_matteo_bulge])
+M_matteo_disk = sum(m[i_matteo_disk])
+rh_matteo_bulge = median(r[i_matteo_bulge])
+print('=== Matteo\'s results ===')
+print(f'M_bulge = {M_matteo_bulge:.2e} MSun   M_disk = {M_matteo_disk:.2e} MSun')
+print(f'rh_bulge = {rh_matteo_bulge:.4f} kpc (equivalent Plummer radius: {0.76642*rh_matteo_bulge:.4f}) kpc')
+
+# Plot the density as a function of d for three values of z.
+f_mn = 1 - f_plummer
+rho = lambda d, z: rho_plummer(sqrt(d**2+z**2), f_plummer, b_plummer) + rho_mn(d, z, f_mn, a_mn, b_mn)
+for c, i in enumerate([0, 3, 6, 9]):
+    semilogy(dd[i,:], M_tot*rho_measured_normalized[i,:], c=f'C{c}', ls='-', label=f'h={1000*zz[i,0]:.0f} pc')
+    semilogy(dd[i,:], M_tot*rho(dd[i,:], zz[i,0]), c=f'C{c}', ls='--', alpha=0.5)
+
+legend()
+xlabel('x [kpc]')
+ylabel(r'$\rho\ [\rm M_\odot\ kpc^{-3}]$')
+
+savefig('subhalo_411321_stellar_fit.png')
+show()