flip angular momentum

flip_AM.py

import astropy.units as u
import numpy as np
from matplotlib.pylab import *

import agama
agama.setUnits(length=1, velocity=1, mass=1)


def sigmoid(x, fmax=1, scale=0.5):
    """
    fmax: [0, 1]
    return y in [1-fmax, fmax]
    sharper transition with smaller scale 
    """
    if not (0 <= fmax <= 1):
        raise ValueError('y should be in [0, 1]')
    A = (fmax - 0.5) / np.tanh(1 / scale)
    y = A * np.tanh(x / scale) + 0.5
    return y


def dice(x, fmax=1, scale=0.5):
    """
    throw dices to decide if flip
    -1 for flip, 1 for no flip
    """
    r = np.random.rand(*np.shape(x)) * 0.5
    y = sigmoid(x, fmax, scale)
    flip = np.sign(r - (0.5 - y))
    return flip


def sample_sph(param, n, spin_vec=None, flip_fmax=1, flip_scale=0.5):
    """
    param :
        agama.Density parameters
    n :
        Number of particles
    spin_vec : None or array shape (3,)
        Vector of spin direction. None for no spin.
    flip_fmax : float in [0, 1]
    flip_scale : positive float
        Parameters to set the level of spin.
        Larger fmax and smaller scale for stronger spins.
    """
    den = agama.Density(param)
    pot = agama.Potential(param)
    af = agama.ActionFinder(pot)
    df = agama.DistributionFunction(
        type='quasispherical', density=den, potential=pot, beta0=0, r_a=np.inf)
    mod = agama.GalaxyModel(potential=pot, df=df, af=af)
    xv, m = mod.sample(n)

    if spin_vec is not None:
        spin_vec = np.asfarray(spin_vec)
        spin_vec = spin_vec / (spin_vec**2).sum()**0.5

        pos, vel = xv[:, :3], xv[:, 3:]
        L_vec = np.cross(pos, vel)
        L = (L_vec**2).sum(1)**0.5
        L_prj = (L_vec * spin_vec).sum(1)  # L_prj/L ~ U[-1, 1]
        flip = dice(L_prj / L, fmax=flip_fmax, scale=flip_scale)  # flip according to L_proj

        vel_flip = vel * flip.reshape(-1, 1)
        xv_flip = np.hstack([pos, vel_flip])

        return xv_flip, m
    else:
        return xv, m

    
param = dict(type='King', mass=1e4, scaleRadius=0.02, W0=8)
n = 10000
# mass: total mass;
# scaleRadius: core radius, Rtide/c, with c decided by W0
# W0: dimensionless potential depth, Phi0/sig2, typically 8 for GCs

seed = 42
np.random.seed(seed)
xv, m = sample_sph(param, n=n)  # no spin

# with spin
np.random.seed(seed)
spin_vec = np.array([0, 0, 1])
xv_spin, m = sample_sph(param, n=n, spin_vec=[0, 0, 1], flip_fmax=0.8, flip_scale=0.5)  


for xv_ in [xv, xv_spin]:
    pos, vel = xv_[:, :3], xv_[:, 3:]

    L_vec = np.cross(pos, vel)
    L = (L_vec**2).sum(1)**0.5
    L_proj = (L_vec * spin_vec).sum(1) / (spin_vec**2).sum()**0.5

    hist(L_proj/L, linspace(-1, 1, 31), density=True, alpha=0.5)

x_ = np.linspace(-1, 1, 501)
plt.plot(x_, sigmoid(x_, fmax=0.5, scale=0.5))
plt.plot(x_, sigmoid(x_, fmax=0.8, scale=0.5))

xlabel(r'$L_z/L$')
ylabel('PDF')