add MPI and do for all boxes

This commit is contained in:
rstiskalek 2022-11-14 15:52:52 +00:00
parent 17badf2652
commit 76349bab3f

View file

@ -17,8 +17,9 @@ Script to split particles into smaller files according to their clump
membership for faster manipulation. Running this will require a lot of memory. membership for faster manipulation. Running this will require a lot of memory.
""" """
from tqdm import tqdm
from os.path import join from os.path import join
from mpi4py import MPI
from datetime import datetime
try: try:
import csiborgtools import csiborgtools
except ModuleNotFoundError: except ModuleNotFoundError:
@ -27,12 +28,21 @@ except ModuleNotFoundError:
import csiborgtools import csiborgtools
import utils import utils
Nsims = [9844] # Get MPI things
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
nproc = comm.Get_size()
# Nsims = [9844]
Nsims = csiborgtools.io.get_csiborg_ids("/mnt/extraspace/hdesmond")
Nsnap = 1016 Nsnap = 1016
partcols = ["x", "y", "z", "vx", "vy", "vz", "M", "level"] partcols = ["x", "y", "z", "vx", "vy", "vz", "M", "level"]
dumpdir = join(utils.dumpdir, "temp") dumpdir = join(utils.dumpdir, "temp")
for Nsim in tqdm(Nsims): jobs = csiborgtools.fits.split_jobs(len(Nsims), nproc)[rank]
for icount, Nsim in enumerate(jobs):
print("{}: rank {} working {} / {} jobs.".format(datetime.now(), rank,
icount + 1, len(jobs)))
simpath = csiborgtools.io.get_sim_path(Nsim) simpath = csiborgtools.io.get_sim_path(Nsim)
# Load the clumps, particles' clump IDs and particles. # Load the clumps, particles' clump IDs and particles.
clumps = csiborgtools.io.read_clumps(Nsnap, simpath) clumps = csiborgtools.io.read_clumps(Nsnap, simpath)