Added script trying to decompose the stellar population

.gitignore

@@ -1,3 +1,6 @@
 .*
 main
 libmain.so
+data
+__pycache__
+*.png

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

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.')

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()