csiborgtools/old/fit_profiles.py
Richard Stiskalek aaa14fc880
Add density field plot and start preparing CSiBORG2 (#94)
* Add RAMSES2HDF5 conversion

* Upload changes

* Clean up

* More clean up

* updates

* Little change

* pep9

* Add basic SPH calculation for a snapshot

* Add submit script

* Remove echo

* Little changes

* Send off changes

* Little formatting

* Little updates

* Add nthreads argument

* Upload chagnes

* Add nthreads arguemnts

* Some local changes..

* Update scripts

* Add submission script

* Update script

* Update params

* Rename CSiBORGBox to CSiBORG1box

* Rename CSiBORG1 reader

* Move files

* Rename folder again

* Add basic plotting here

* Add new skeletons

* Move def

* Update nbs

* Edit directories

* Rename files

* Add units to converted snapshots

* Fix empty dataset bug

* Delete file

* Edits to submission scripts

* Edit paths

* Update .gitignore

* Fix attrs

* Update weighting

* Fix RA/dec bug

* Add FORNAX cluster

* Little edit

* Remove boxes since will no longer need

* Move func back

* Edit to include sort by membership

* Edit paths

* Purge basic things

* Start removing

* Bring file back

* Scratch

* Update the rest

* Improve the entire file

* Remove old things

* Remove old

* Delete old things

* Fully updates

* Rename file

* Edit submit script

* Little things

* Add print statement

* Add here cols_to_structured

* Edit halo cat

* Remove old import

* Add comment

* Update paths manager

* Move file

* Remove file

* Add chains
2023-12-13 16:08:25 +00:00

110 lines
3.8 KiB
Python

# Copyright (C) 2023 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.
"""
A script to calculate the particle's separation from the CM and save it.
Currently MPI is not supported.
"""
from argparse import ArgumentParser
from datetime import datetime
from gc import collect
import numpy
from mpi4py import MPI
from tqdm import trange
try:
import csiborgtools
except ModuleNotFoundError:
import sys
sys.path.append("../")
import csiborgtools
parser = ArgumentParser()
parser.add_argument("--ics", type=int, nargs="+", default=None,
help="IC realisatiosn. If `-1` processes all simulations.")
args = parser.parse_args()
# Get MPI things
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
nproc = comm.Get_size()
if nproc > 1:
raise NotImplementedError("MPI is not implemented implemented yet.")
paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)
cols_collect = [("r", numpy.float32), ("M", numpy.float32)]
if args.ics is None or args.ics == -1:
nsims = paths.get_ics("csiborg")
else:
nsims = args.ics
# We loop over simulations. Here later optionally add MPI.
for i, nsim in enumerate(nsims):
if rank == 0:
now = datetime.now()
print(f"{now}: calculating {i}th simulation `{nsim}`.", flush=True)
nsnap = max(paths.get_snapshots(nsim, "csiborg"))
box = csiborgtools.read.CSiBORG1Box(nsnap, nsim, paths)
f = csiborgtools.read.read_h5(paths.particles(nsim, "csiborg"))
particles = f["particles"]
clump_map = f["clumpmap"]
clid2map = {clid: i for i, clid in enumerate(clump_map[:, 0])}
clumps_cat = csiborgtools.read.ClumpsCatalogue(nsim, paths, rawdata=True,
load_fitted=False)
ismain = clumps_cat.ismain
ntasks = len(clumps_cat)
# We loop over halos and add ther particle positions to this dictionary,
# which we will later save as an archive.
out = {}
for j in trange(ntasks) if nproc == 1 else range(ntasks):
# If we are fitting halos and this clump is not a main, then continue.
if not ismain[j]:
continue
clumpid = clumps_cat["index"][j]
parts = csiborgtools.read.load_parent_particles(
clumpid, particles, clump_map, clid2map, clumps_cat)
# If we have no particles, then do not save anything.
if parts is None:
continue
obj = csiborgtools.fits.Clump(parts, clumps_cat[j], box)
r200m, m200m = obj.spherical_overdensity_mass(200, npart_min=10,
kind="matter")
r = obj.r()
mask = r <= r200m
_out = csiborgtools.read.cols_to_structured(numpy.sum(mask),
cols_collect)
_out["r"] = r[mask]
_out["M"] = obj["M"][mask]
out[str(clumpid)] = _out
# Finished, so we save everything.
fout = paths.radpos_path(nsnap, nsim)
now = datetime.now()
print(f"{now}: saving radial profiles for simulation {nsim} to `{fout}`",
flush=True)
numpy.savez(fout, **out)
# Clean up the memory just to be sure.
del out
collect()