mirror of
https://github.com/Richard-Sti/csiborgtools_public.git
synced 2025-05-12 05:38:42 +00:00
Power spectrum stuff (#16)
* add autok1d * update TODO * add basic cross pk * del blank line * fix small bugs * fix loop bug * rm nbs * edit docstring * Update TODO * get rid of verbose flag * take all ICs * attempt better memory * simplify ASCII dump * rm verbose statement to get rid of bugs! * close fortran files * add MAS options * add boxsize setter * Move comment up * paths chagne * Update TODO * Update TODO * Forgotten units in Xpk * Attempt at fixing the distance bug * Update TODO * Update TODO * Remove comment * add summary reader * Update gitignore
This commit is contained in:
Richard Stiskalek
GitHub
parent
1bb4255874
commit
c3c686fa60
10 changed files with 439 additions and 5314 deletions
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
168
scripts/run_crosspk.py
Normal file
168
scripts/run_crosspk.py
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
# Copyright (C) 2022 Richard Stiskalek
|
||||
# This program is free software; you can redistribute it and/or modify it
|
||||
# under the terms of the GNU 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 General Public License along
|
||||
# with this program; if not, write to the Free Software Foundation, Inc.,
|
||||
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
|
||||
"""
|
||||
MPI script to calculate the matter cross power spectrum between CSiBORG
|
||||
IC realisations. Units are Mpc/h.
|
||||
"""
|
||||
from argparse import ArgumentParser
|
||||
import numpy
|
||||
import joblib
|
||||
from datetime import datetime
|
||||
from itertools import combinations
|
||||
from os.path import join
|
||||
from os import remove
|
||||
from gc import collect
|
||||
from mpi4py import MPI
|
||||
import Pk_library as PKL
|
||||
try:
|
||||
import csiborgtools
|
||||
except ModuleNotFoundError:
|
||||
import sys
|
||||
sys.path.append("../")
|
||||
import csiborgtools
|
||||
import utils
|
||||
|
||||
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--grid", type=int)
|
||||
parser.add_argument("--halfwidth", type=float, default=0.5)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get MPI things
|
||||
comm = MPI.COMM_WORLD
|
||||
rank = comm.Get_rank()
|
||||
nproc = comm.Get_size()
|
||||
MAS = "CIC" # mass asignment scheme
|
||||
|
||||
paths = csiborgtools.read.CSiBORGPaths()
|
||||
ics = paths.ic_ids
|
||||
n_sims = len(ics)
|
||||
|
||||
# File paths
|
||||
ftemp = join(utils.dumpdir, "temp_crosspk",
|
||||
"out_{}_{}" + "_{}".format(args.halfwidth))
|
||||
fout = join(utils.dumpdir, "crosspk",
|
||||
"out_{}_{}" + "_{}.p".format(args.halfwidth))
|
||||
|
||||
|
||||
jobs = csiborgtools.fits.split_jobs(n_sims, nproc)[rank]
|
||||
for n in jobs:
|
||||
print("Rank {}@{}: saving {}th delta.".format(rank, datetime.now(), n))
|
||||
# Set the paths
|
||||
n_sim = ics[n]
|
||||
paths.set_info(n_sim, paths.get_maximum_snapshot(n_sim))
|
||||
# Set reader and the box
|
||||
reader = csiborgtools.read.ParticleReader(paths)
|
||||
box = csiborgtools.units.BoxUnits(paths)
|
||||
# Read particles
|
||||
particles = reader.read_particle(["x", "y", "z", "M"], verbose=False)
|
||||
# Halfwidth -- particle selection
|
||||
if args.halfwidth < 0.5:
|
||||
hw = args.halfwidth
|
||||
mask = ((0.5 - hw < particles['x']) & (particles['x'] < 0.5 + hw)
|
||||
& (0.5 - hw < particles['y']) & (particles['y'] < 0.5 + hw)
|
||||
& (0.5 - hw < particles['z']) & (particles['z'] < 0.5 + hw))
|
||||
# Subselect the particles
|
||||
particles = particles[mask]
|
||||
# Rescale to range [0, 1]
|
||||
for p in ('x', 'y', 'z'):
|
||||
particles[p] = (particles[p] - 0.5 + hw) / (2 * hw)
|
||||
|
||||
length = box.box2mpc(2 * hw) * box.h
|
||||
else:
|
||||
mask = None
|
||||
length = box.box2mpc(1) * box.h
|
||||
# Calculate the overdensity field
|
||||
field = csiborgtools.field.DensityField(particles, length, box, MAS)
|
||||
delta = field.overdensity_field(args.grid, verbose=False)
|
||||
aexp = box._aexp
|
||||
|
||||
# Try to clean up memory
|
||||
del field, particles, box, reader, mask
|
||||
collect()
|
||||
|
||||
# Dump the results
|
||||
with open(ftemp.format(n_sim, "delta") + ".npy", "wb") as f:
|
||||
numpy.save(f, delta)
|
||||
joblib.dump([aexp, length], ftemp.format(n_sim, "lengths") + ".p")
|
||||
|
||||
# Try to clean up memory
|
||||
del delta
|
||||
collect()
|
||||
|
||||
|
||||
comm.Barrier()
|
||||
|
||||
# Get off-diagonal elements and append the diagoal
|
||||
combs = [c for c in combinations(range(n_sims), 2)]
|
||||
for i in range(n_sims):
|
||||
combs.append((i, i))
|
||||
prev_delta = [-1, None, None, None] # i, delta, aexp, length
|
||||
|
||||
jobs = csiborgtools.fits.split_jobs(len(combs), nproc)[rank]
|
||||
for n in jobs:
|
||||
i, j = combs[n]
|
||||
print("Rank {}@{}: combination {}.".format(rank, datetime.now(), (i, j)))
|
||||
|
||||
# If i same as last time then don't have to load it
|
||||
if prev_delta[0] == i:
|
||||
delta_i = prev_delta[1]
|
||||
aexp_i = prev_delta[2]
|
||||
length_i = prev_delta[3]
|
||||
else:
|
||||
with open(ftemp.format(ics[i], "delta") + ".npy", "rb") as f:
|
||||
delta_i = numpy.load(f)
|
||||
aexp_i, length_i = joblib.load(ftemp.format(ics[i], "lengths") + ".p")
|
||||
# Store in prev_delta
|
||||
prev_delta[0] = i
|
||||
prev_delta[1] = delta_i
|
||||
prev_delta[2] = aexp_i
|
||||
prev_delta[3] = length_i
|
||||
|
||||
# Get jth delta
|
||||
with open(ftemp.format(ics[j], "delta") + ".npy", "rb") as f:
|
||||
delta_j = numpy.load(f)
|
||||
aexp_j, length_j = joblib.load(ftemp.format(ics[j], "lengths") + ".p")
|
||||
|
||||
# Verify the difference between the scale factors! Say more than 1%
|
||||
daexp = abs((aexp_i - aexp_j) / aexp_i)
|
||||
if daexp > 0.01:
|
||||
raise ValueError(
|
||||
"Boxes {} and {} final snapshot scale factors disagree by "
|
||||
"`{}` percent!".format(ics[i], ics[j], daexp * 100))
|
||||
# Check how well the boxsizes agree
|
||||
dlength = abs((length_i - length_j) / length_i)
|
||||
if dlength > 0.001:
|
||||
raise ValueError("Boxes {} and {} box sizes disagree by `{}` percent!"
|
||||
.format(ics[i], ics[j], dlength * 100))
|
||||
|
||||
# Calculate the cross power spectrum
|
||||
Pk = PKL.XPk([delta_i, delta_j], length_i, axis=1, MAS=[MAS, MAS],
|
||||
threads=1)
|
||||
joblib.dump(Pk, fout.format(ics[i], ics[j]))
|
||||
|
||||
del delta_i, delta_j, Pk
|
||||
collect()
|
||||
|
||||
|
||||
# Clean up the temp files
|
||||
comm.Barrier()
|
||||
if rank == 0:
|
||||
print("Cleaning up the temporary files...")
|
||||
for ic in ics:
|
||||
remove(ftemp.format(ic, "delta") + ".npy")
|
||||
remove(ftemp.format(ic, "lengths") + ".p")
|
||||
|
||||
print("All finished!")
|
||||
Loading…
Add table
Add a link
Reference in a new issue