makeIC.py

Go to the documentation of this file.

import h5py
import argparse
import numpy as np
 
"""!
Generates 3D uniform particle distribution with randomized small displacements.
"""
 
fileName = "test_sml_3D.hdf5"
boxSize = 1.0
 
parser = argparse.ArgumentParser(description="SPH benchmarking script")
parser.add_argument(
    "-np",
    "--particle-number",
    dest="particle_number",
    required=True,
    type=int,
    help="Particle number per dimension",
)
parser.add_argument(
    "-s",
    "--seed",
    dest="random_seed",
    required=True,
    type=int,
    help="Seed for RNG",
)
parser.add_argument(
    "-p",
    "--plot",
    dest="plot",
    action="store_true",
    default=False,
    help="Plot historgram of particle positions",
)
parser.add_argument(
    "-dx",
    "--displacement",
    dest="max_displacement",
    required=False,
    default=0.2499,
    type=float,
    help="Maximal displacement in units of interparticle distance",
)
 
 
args = parser.parse_args()
npart = args.particle_number
seed = args.random_seed
max_displacement = args.max_displacement
 
 
# ---------------------------------------------------
 
# Build the arrays
coords = np.zeros((npart**3, 3))
v = np.zeros((npart**3, 3))
ids = np.linspace(1, npart**3, npart**3)
m = np.ones(npart**3)
dx = boxSize / npart
h = np.ones(npart**3) * 2.0 * dx
u = np.ones(npart**3)
 
# Get uniform coordinates.
ind = 0
for i in range(npart):
    x = (i + 0.5) * dx
    for j in range(npart):
        y = (j + 0.5) * dx
        for k in range(npart):
            z = (k + 0.5) * dx
 
            coords[ind, 0] = x
            coords[ind, 1] = y
            coords[ind, 2] = z
            ind += 1
 
# Now randomize them a bit.
# Taking a completely randomly sampled particle distribution may lead
# to trouble the codes aren't set up to handle, such as particles which
# are either too close or too far from each other. So use a somewhat
# regular distribution instead, by randomly displacing an initially
# uniform particle distribution.
 
rng = np.random.RandomState(seed=seed)
# rng.random returns random numbers in [0, 1). Scale that.
displacement = (2.0 * rng.random(coords.shape) - 1.0) * max_displacement * dx
#  displacement = np.zeros(coords.shape)
coords[:, 0] += displacement[:, 0]
coords[:, 1] += displacement[:, 1]
coords[:, 2] += displacement[:, 2]
 
 
# too many particles mean that the comparison python script takes ages.
# too few particles mean that the search radii are going to be too big
# for the cell size.
# So let's scale the particles down into a smaller box than the actual
# box size.
# Also, at full scale, the particles along the boundary/in corners start
# disagreeing stupidly. So let's not do that.
scale = 0.6
coords *= scale
h *= scale
# and shift them into the center
coords += 0.2
 
if args.plot:
    from matplotlib import pyplot as plt
 
    plt.figure()
    plt.scatter(coords[:, 0], coords[:, 1], s=1, marker=",")
    plt.show()
 
 
# File
file = h5py.File(fileName, "w")
 
# Header
grp = file.create_group("/Header")
grp.attrs["BoxSize"] = boxSize
grp.attrs["NumPart_Total"] = [npart**3, 0, 0, 0, 0, 0]
grp.attrs["NumPart_Total_HighWord"] = [0, 0, 0, 0, 0, 0]
grp.attrs["NumPart_ThisFile"] = [npart**3, 0, 0, 0, 0, 0]
# This is an addition for this test to be able to estimate the search radius.
grp.attrs["NumPart_base"] = [npart]
grp.attrs["Time"] = 0.0
grp.attrs["NumFilesPerSnapshot"] = 1
grp.attrs["MassTable"] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
grp.attrs["Flag_Entropy_ICs"] = 0
grp.attrs["Dimension"] = 3
# Non-swift additions
grp.attrs["RandomSeed"] = seed
grp.attrs["MaxDisplacement"] = max_displacement
grp.attrs["NumpyVersion"] = np.__version__
 
# Units
grp = file.create_group("/Units")
grp.attrs["Unit length in cgs (U_L)"] = 1.0
grp.attrs["Unit mass in cgs (U_M)"] = 1.0
grp.attrs["Unit time in cgs (U_t)"] = 1.0
grp.attrs["Unit current in cgs (U_I)"] = 1.0
grp.attrs["Unit temperature in cgs (U_T)"] = 1.0
 
# Particle group
grp = file.create_group("/PartType0")
grp.create_dataset("Coordinates", data=coords, dtype="d")
grp.create_dataset("Velocities", data=v, dtype="f")
grp.create_dataset("Masses", data=m, dtype="f")
grp.create_dataset("SmoothingLength", data=h, dtype="f")
grp.create_dataset("InternalEnergy", data=u, dtype="f")
grp.create_dataset("ParticleIDs", data=ids, dtype="L")
 
 
file.close()