From 9ec0633094bde62e9c023b07a09a880777c5a62c Mon Sep 17 00:00:00 2001 From: Yohai Meiron Date: Thu, 16 Apr 2020 20:51:01 -0400 Subject: [PATCH] Added script trying to decompose the stellar population --- .gitignore | 3 + density_center.py | 56 ++++++++++++++++++ ellipsoids.py | 65 ++++++++++++++++++++ three_components.py | 140 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 264 insertions(+) create mode 100644 density_center.py create mode 100644 ellipsoids.py create mode 100644 three_components.py diff --git a/.gitignore b/.gitignore index 6f8fe4b..60b2039 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ .* main libmain.so +data +__pycache__ +*.png \ No newline at end of file diff --git a/density_center.py b/density_center.py new file mode 100644 index 0000000..f01fa4b --- /dev/null +++ b/density_center.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +import numpy as np + +def def_dc(m, x, v=None, r_min=0.0045): + """ + Calculates the center of density. Translated to Python from phiGRAPE. + + Parameters + ---------- + m, x, v : array_like + Particle parameters. v is optional. + r_min : scalar + For star clusters, should be 0.1 pc in N-body units. + + Returns + ------- + xdc, vdc : ndarrays + """ + calc_vdc = not v is None + x_ = x.copy() + xdc = np.zeros(3) + if calc_vdc: + v_ = v.copy() + vdc = np.zeros(3) + else: + v_ = None + r_lim = np.sqrt(np.max(np.sum(x**2, axis=1))) + num_iter = 0 + + while (r_lim > r_min) and (num_iter < 100): + ncm, mcm, xcm, vcm = cenmas_lim(m, x_, v_, r_lim) + if((mcm > 0) and (ncm > 100)): + x_ -= xcm + xdc += xcm + if calc_vdc: + v_ -= vcm + vdc += vcm + else: + break + r_lim *= 0.8 + num_iter += 1 + if calc_vdc: + return xdc, vdc + else: + return xdc + +def cenmas_lim(m, x, v, r_lim): + r2 = np.sum(x**2, axis=1) + cond = r2 < r_lim**2 + ncm = np.sum(cond) + mcm = np.sum(m[cond]) + if mcm == 0: return ncm, 0., np.zeros(3), np.zeros(3) + xcm = np.sum(m[cond,None] * x[cond,:], axis=0) / mcm + if not v is None: vcm = np.sum(m[cond,None] * v[cond,:], axis=0) / mcm + else: vcm = None + return ncm, mcm, xcm, vcm diff --git a/ellipsoids.py b/ellipsoids.py new file mode 100644 index 0000000..7ca815a --- /dev/null +++ b/ellipsoids.py @@ -0,0 +1,65 @@ +import numpy as np + +def quadrupole_tensor(x, y, z, m=1): + Qxx = np.sum(m*x**2) + Qyy = np.sum(m*y**2) + Qzz = np.sum(m*z**2) + Qxy = Qyx = np.sum(m*x*y) + Qxz = Qzx = np.sum(m*x*z) + Qzy = Qyz = np.sum(m*z*y) + return np.array([[Qxx, Qxy, Qxz], + [Qyx, Qyy, Qyz], + [Qzx, Qzy, Qzz]]) + +def rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues): + """ + Calculates a rotation matrix such that the eigenvector with the largest + eigenvalue points in the x-direction, the eigenvector with the middle + eigenvalue points in the y-direction, and so on. + """ + return np.array([eigenvectors[:,i] for i in np.argsort(eigenvalues)[::-1]]) + +def generate_ellipsoid(n, power, a, b, c, euler_1, euler_2, euler_3, seed=None): + from scipy.spatial.transform import Rotation + if not (a >= b >= c): + raise ValueError('Requires a >= b >= c') + b_over_a = b/a + c_over_a = c/a + np.random.seed(seed) + + # Create a sphere of particles with radius a. + r = (np.random.rand(n)**power)*a + theta = np.arccos(np.random.rand(n)*2-1) + phi = np.random.rand(n)*2*np.pi + x = r*np.cos(phi)*np.sin(theta) + y = r*np.sin(phi)*np.sin(theta) + z = r* np.cos(theta) + + # Scale the intermediate and minor axes + y *= b_over_a + z *= c_over_a + + # Rotate + rotation = Rotation.from_euler('zxz', [euler_1, euler_2, euler_3]) + return rotation.apply(np.stack((x,y,z), axis=1)) + +if __name__ == '__main__': # Test if the above routines + n = 65536 + power = 1.5 + a, b, c = 2.3, 0.92, 0.345 # or sorted(np.random.rand(3))[::-1] + angles = 0.43, 1.81, 0.32 # or np.random.rand(3)*2*np.pi + + X = generate_ellipsoid(n, power, a, b, c, *angles, seed=None) + Q = quadrupole_tensor(*X.T) + eigenvalues, eigenvectors = np.linalg.eig(Q) + i = np.argsort(eigenvalues) + print('measured b/a = %.4f (expected %.4f)' % (np.sqrt(eigenvalues[i[1]]/eigenvalues[i[2]]), b/a)) + print('measured c/a = %.4f (expected %.4f)' % (np.sqrt(eigenvalues[i[0]]/eigenvalues[i[2]]), c/a)) + R = rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues) + print('Rotating vectors and recalculating quadrupole tensor...') + X_new = (R @ X.T).T + Q_new = quadrupole_tensor(*X_new.T) + print('%11.4e %11.4e %11.4e' % tuple(Q_new[0,:])) + print('%11.4e %11.4e %11.4e' % tuple(Q_new[1,:])) + print('%11.4e %11.4e %11.4e' % tuple(Q_new[2,:])) + print('... it should be almost diagonal.') diff --git a/three_components.py b/three_components.py new file mode 100644 index 0000000..dc1e6ec --- /dev/null +++ b/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() \ No newline at end of file