From efdb0ea42574ab36e02d0940de78a41849534f76 Mon Sep 17 00:00:00 2001
From: Yohai Meiron <sahmes@gmail.com>
Date: Fri, 24 Apr 2020 00:17:54 -0400
Subject: [PATCH] Wrote script to find disk and halo parameters in all
 snapshots

---
 parameterize-all.py | 103 ++++++++++++++++++++++++++++++++++++++++++++
 three_components.py |   4 +-
 2 files changed, 105 insertions(+), 2 deletions(-)
 create mode 100644 parameterize-all.py

diff --git a/parameterize-all.py b/parameterize-all.py
new file mode 100644
index 0000000..1c85587
--- /dev/null
+++ b/parameterize-all.py
@@ -0,0 +1,103 @@
+#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
+from astropy.cosmology import Planck15
+
+def miyamoto_nagai_params_from_medians(m_d, m_z):
+    m_d_pred = lambda x: exp(1.43475163 + (1.04827148*log(x)-1.09023112)*(arctan(log(x)/2.30939056)/pi+0.5))
+    m_z_pred = lambda x: 5.77340E-01*x
+    b_over_a = scipy.optimize.brentq(lambda x: m_d_pred(x)/m_z_pred(x)-m_d/m_z, 1/64, 64)
+    a1 = m_d/m_d_pred(b_over_a)
+    a2 = m_z/m_z_pred(b_over_a)
+    a = 0.5*(a1 + a2)
+    b = a*b_over_a
+    return a, b
+
+# Simumation parameters
+h0 = 0.6774
+snap, z = loadtxt('snapshots.dat', unpack=True)
+snap = snap.astype(int)
+a = 1/(1+z)
+t = Planck15.age(z).value
+
+# Halo parameters
+file_name = 'data/subhalo_411321.hdf5'
+
+# Dictionary of particle types
+particle_types = {}
+particle_types['gas']   = '0'
+particle_types['dm']    = '1'
+particle_types['stars'] = '4'
+particle_types['bhs']   = '5'
+
+def get_half_mass_radius(m, r):
+    i = argsort(r)
+    r_sorted = r[i]
+    m_sorted = m[i]
+    m_cum = cumsum(m_sorted)
+    j = searchsorted(m_cum, m_cum[-1]/2)
+    return r_sorted[j] # close enough
+
+def get_transformation(m, X):
+    X_center = def_dc(m, X)
+    X_shifted = X - X_center
+    r = linalg.norm(X_shifted, axis=1)
+    rh = get_half_mass_radius(m, r)
+    mask = r < 2*rh
+    Q = ellipsoids.quadrupole_tensor(*X_shifted[mask].T, m[mask])
+    eigenvalues, eigenvectors = linalg.eig(Q)
+    R = ellipsoids.rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues)
+    return X_center, R, mask
+
+f = h5py.File(file_name, 'r')
+
+# First get the global coordinate system
+i, snapshot = len(snap)-1, snap[-1]
+m = f[str(snapshot)][particle_types['stars']]['Masses'][...]
+X = f[str(snapshot)][particle_types['stars']]['Coordinates'][...] * a[i] / h0
+X_center_glob, R_glob, _ = get_transformation(m, X)
+
+for i in range(len(snap)):
+    snapshot = snap[0] + i
+
+    m_dm  = f[str(snapshot)][particle_types['dm']]['Masses'][...]
+    X_dm  = f[str(snapshot)][particle_types['dm']]['Coordinates'][...] * a[i] / h0
+    m_gas = f[str(snapshot)][particle_types['gas']]['Masses'][...]
+    X_gas = f[str(snapshot)][particle_types['gas']]['Coordinates'][...] * a[i] / h0
+    m = append(m_dm, m_gas)
+    X = vstack([X_dm, X_gas])
+
+    X -= X_center_glob # NOTE we don't rotate here because we assume spherical symmetry
+    X_center_halo = def_dc(m, X)
+    r = linalg.norm(X, axis=1)
+    rh = get_half_mass_radius(m, r)
+    b_halo = 0.76642093654*rh
+    M_halo = sum(m)
+
+    particle_type = 'stars'
+    m = f[str(snapshot)][particle_types[particle_type]]['Masses'][...]
+    X = f[str(snapshot)][particle_types[particle_type]]['Coordinates'][...] * a[i] / h0
+    X = (R_glob @ (X - X_center_glob).T).T
+
+    X_center_stars, R, mask = get_transformation(m, X)
+    direction = R[0,:]
+    if direction[2] < 0: direction = -direction
+    phi   = arctan2(direction[1], direction[0])
+    theta = arccos(direction[2]/linalg.norm(direction))
+
+    X_new = (R @ (X - X_center_stars).T).T
+    x, y, z = X_new.T
+    m_z = median(abs(z[mask]))
+    m_d = median(sqrt(x[mask]**2+y[mask]**2))
+    a_mn, b_mn = miyamoto_nagai_params_from_medians(m_d, m_z)
+    M_disk = sum(m)
+    print('%d %.8E   %.8E   %15.8E %15.8E %15.8E   %15.8E %15.8E   %.8E %.8E   %.8E   %15.8E %15.8E %15.8E   %15.8E' % (snapshot, t[i], M_disk, *X_center_stars, phi, theta, a_mn, b_mn, M_halo, *X_center_halo, b_halo))
+
+
+
+f.close()
\ No newline at end of file
diff --git a/three_components.py b/three_components.py
index c8f4843..3ce994f 100644
--- a/three_components.py
+++ b/three_components.py
@@ -139,5 +139,5 @@ for c, i in enumerate([0, 3, 6, 9]):
     xlabel('x [kpc]')
     ylabel(r'$\rho\ [\rm M_\odot\ kpc^{-3}]$')
 
-    savefig('subhalo_411321_stellar_fit.png')
-    show()
\ No newline at end of file
+savefig('subhalo_411321_stellar_fit.png')
+show()
\ No newline at end of file