################################################################################
# Copyright (c) 2022 Patrick Hirling (patrick.hirling@epfl.ch)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see .
#
################################################################################
import numpy as np
import scipy.special as sci
from scipy.optimize import minimize
import time
import argparse
from swiftsimio import Writer
import unyt
parser = argparse.ArgumentParser(
description="Generate discrete realization of Anisotropic Plummer Model"
)
# Parse Main Arguments
parser.add_argument("-a", type=float, default=0.05, help="Softening Length (in kpc)")
parser.add_argument(
"-M", type=float, default=1.0e-5, help="Total Mass of Model (in 10^10 solar masses)"
)
parser.add_argument(
"-q", type=float, default=1.0, help="Anisotropy parameter (between -inf and +2)"
)
parser.add_argument(
"-N", type=int, default=1e4, help="Number of particles in the simulation"
)
parser.add_argument(
"-bound", type=float, default=2.0, help="Upper bound of radius sampling"
)
parser.add_argument("-boxsize", type=float, default=4.0, help="Box Size of simulation")
parser.add_argument(
"--noparallel",
action="store_false",
help="Do not use multiprocessing library (default: false)",
)
parser.add_argument(
"-nbthreads",
type=int,
default=6,
help="Number of threads to use for multiprocessing (default: 6)",
)
args = parser.parse_args()
if args.q > 2.0:
raise ValueError(
"Anisotropy parameter must be smaller than 2.0 (q = {:.2f} given)".format(
args.q
)
)
#### Parameters
# Plummer Model
G = 4.299581e04 # Gravitational constant [kpc / 10^10 M_s * (kms)^2]
a = args.a # Plummer softening length [kpc]
M = args.M # Total Mass [10^10 M_s]
q = args.q # Anisotropy Parameter (-inf,2]
# IC File
N = int(args.N) # Number of Particles
fname = "plummer.hdf5" # Name of the ic file (dont forget .hdf5)
pickle_ics = 0 # Optional: Pickle ics (pos,vel,mass) for future use
periodic_bdry = 0
# Parallelism for velocity sampling
use_parallel = args.noparallel
nb_threads = int(args.nbthreads) # Set to None to use os.cpucount()
### Print parameters
print(
"\n############################################################################################# \n"
)
print("Generating Plummer ICs with the following parameters::\n")
print(
"Softening length: a = " + "{:.3e} kpc".format(a)
)
print(
"Total Mass: M = "
+ "{:.3e} Solar Masses".format(M * 1e10)
)
print("Anisotropy parameter: q = " + str(q))
print("Output file: " + fname)
print("Number of particles: N = " + "{:.1e}".format(N))
print(
"\n############################################################################################# \n"
)
# Estimate good softening length (density at origin, sphere s.t. contains on avg 1 particle)
epsilon0 = a * N ** (-1.0 / 3.0)
print(
"Recommended Softening length (times conventional factor): "
+ "{:.4e}".format(epsilon0)
+ " kpc"
)
### Generate Positions
# Inverse transform sampling from analytical inverse of M(=0)
vmax = np.sqrt(2.0 * relative_potential(r))
# Compute max of DF at this radius
fm = 1.1 * fmax(r, vmax) # 1.1 factor to be sure to include max
while True:
# Select candidates for vr,vt based on max full velocity
while True:
vr = np.random.uniform(0.0, vmax)
vt = np.random.uniform(0.0, vmax)
if vr ** 2 + vt ** 2 <= vmax ** 2:
break
# Rejection Sampling on Fq
f = np.random.uniform(0.0, fm)
if Fq(relative_energy(r, vr, vt), r * vt) * vt >= f:
return vr, vt
print("Sampling velocities...")
ti = time.time()
if use_parallel:
from multiprocessing import Pool
with Pool(nb_threads) as p:
vels = np.array(p.map(sample_vel, r_rand)).transpose()
else:
from tqdm import tqdm
vels = np.empty((2, N))
for j, r in enumerate(tqdm(r_rand)):
vels[:, j] = sample_vel(r)
tf = time.time()
print("Sampling took " + "{:.2f}".format(tf - ti) + " seconds.")
# Convert to Cartesian
# First: project vt on e_theta, e_phi with random orientation
alph = np.random.uniform(0, 2 * np.pi, N)
sgn = np.random.choice([-1, 1], size=N)
vphi = vels[1] * np.cos(alph)
vtheta = vels[1] * np.sin(alph)
# project vr on e_r (random sign)
vr = sgn * vels[0]
# Convert Spherical to cartesian coordinates
v_x = (
np.sin(theta_rand) * np.cos(phi_rand) * vr
+ np.cos(theta_rand) * np.cos(phi_rand) * vtheta
- np.sin(phi_rand) * vphi
)
v_y = (
np.sin(theta_rand) * np.sin(phi_rand) * vr
+ np.cos(theta_rand) * np.sin(phi_rand) * vtheta
+ np.cos(phi_rand) * vphi
)
v_z = np.cos(theta_rand) * vr - np.sin(theta_rand) * vtheta
# Create velocity array
V = np.array([v_x, v_y, v_z]).transpose()
### Generate masses
m = M / N * np.ones(N)
### Exclude extreme outliers
idx = r_rand < args.bound
X = X[idx]
V = V[idx]
m = m[idx]
new_N = len(m)
### Write to hdf5
galactic_units = unyt.UnitSystem(
"galactic",
unyt.kpc,
unyt.unyt_quantity(1e10, units=unyt.Solar_Mass),
unyt.unyt_quantity(1.0, units=unyt.s * unyt.Mpc / unyt.km).to(unyt.Gyr),
)
wrt = Writer(galactic_units, args.boxsize * unyt.kpc)
wrt.dark_matter.coordinates = X * unyt.kpc
wrt.dark_matter.velocities = V * (unyt.km / unyt.s)
wrt.dark_matter.masses = m * 1e10 * unyt.msun
wrt.dark_matter.particle_ids = np.arange(new_N)
wrt.write(fname)
print("Writing IC file...")
### Optional: Pickle ic arrays for future use
if pickle_ics:
import pickle as pkl
with open("X.pkl", "wb") as f:
pkl.dump(X[:], f)
with open("V.pkl", "wb") as f:
pkl.dump(V[:], f)
with open("M.pkl", "wb") as f:
pkl.dump(m[:], f)