def array(self, N=0, filename=None, component=None, root=0):
"""Dump data to numpy format on root processor."""
assert (N == 0 or N == 1)
is_root = comm.Get_rank() == root
size = self.get_total_number_probes() if is_root else len(self)
comp = self.value_size() if component is None else 1
z = zeros((size, comp))
# Retrieve all values
if len(self) > 0:
for k in range(comp):
if is_root:
ids = self.get_probe_ids()
z[ids, k] = self.get_probes_component_and_snapshot(k, N)
else:
z[:, k] = self.get_probes_component_and_snapshot(k, N)
# Collect on root
recvfrom = comm.gather(len(self), root=root)
if is_root:
for j, k in enumerate(recvfrom):
if comm.Get_rank() != j:
ids = comm.recv(source=j, tag=101)
z0 = comm.recv(source=j, tag=102)
z[ids, :] = z0[:, :]
else:
ids = self.get_probe_ids()
comm.send(ids, dest=root, tag=101)
comm.send(z, dest=root, tag=102)
if is_root:
if filename:
save(filename + "_statistics", z)
return squeeze(z)
def array(self, N=0, filename=None, component=None, root=0):
"""Dump data to numpy format on root processor."""
assert(N == 0 or N == 1)
is_root = comm.Get_rank() == root
size = self.get_total_number_probes() if is_root else len(self)
comp = self.value_size() if component is None else 1
z = zeros((size, comp))
# Retrieve all values
if len(self) > 0:
for k in range(comp):
if is_root:
ids = self.get_probe_ids()
z[ids, k] = self.get_probes_component_and_snapshot(k, N)
else:
z[:, k] = self.get_probes_component_and_snapshot(k, N)
# Collect on root
recvfrom = comm.gather(len(self), root=root)
if is_root:
for j, k in enumerate(recvfrom):
if comm.Get_rank() != j:
ids = comm.recv(source=j, tag=101)
z0 = comm.recv(source=j, tag=102)
z[ids, :] = z0[:, :]
else:
ids = self.get_probe_ids()
comm.send(ids, dest=root, tag=101)
comm.send(z, dest=root, tag=102)
if is_root:
if filename:
z.dump(filename+"_statistics.probes")
return squeeze(z)
def test_partition(self):
""" Make sure the x_ranges span the entire space without any gaps. """
shapes = ((200,30,10), (33,10,10), (130,5,5), (111,2,2))
for shape in shapes:
space.initialize_space(shape)
x = comm.gather(space.get_space_info()['x_range'])
if comm.Get_rank() == 0:
self.assertEqual(x[0][0], 0)
self.assertEqual(x[-1][-1], space.get_space_info()['shape'][0])
for k in range(len(x)-1):
self.assertEqual(x[k][1], x[k+1][0])
def test_partition(self):
""" Make sure the x_ranges span the entire space without any gaps. """
shapes = ((200, 30, 10), (33, 10, 10), (130, 5, 5), (111, 2, 2))
for shape in shapes:
space.initialize_space(shape)
x = comm.gather(space.get_space_info()['x_range'])
if comm.Get_rank() == 0:
self.assertEqual(x[0][0], 0)
self.assertEqual(x[-1][-1], space.get_space_info()['shape'][0])
for k in range(len(x) - 1):
self.assertEqual(x[k][1], x[k + 1][0])
def plot_fields( self, save_figure=False ):
"""
Plot the Ex and By field using matplotlib
If save_figure is True, the plots are saved as PNG files,
in a folder named `diagnostics`
"""
# PLOTTING: NEW LINES RELATED TO MPI
Ex_from_all_procs = mpi_comm.gather( self.Ex[1:-1] )
By_from_all_procs = mpi_comm.gather( self.By[1:-1] )
if mpi_comm.rank == 0:
print( 'Plotting the fields at iteration %d' %self.n )
global_Ex = np.concatenate( Ex_from_all_procs )
global_By = np.concatenate( By_from_all_procs )
plt.clf()
plt.suptitle('Fields at iteration %d' %self.n)
# Plot of Ex
plt.subplot(211)
z = self.dz*np.arange( self.Nz_global )
plt.plot( z, global_Ex, '-' )
plt.ylim(-1.1, 1.1)
plt.xlim(0, self.Lz)
plt.ylabel('$E_x^n$')
plt.xlabel('z')
# Plot of By
plt.subplot(212)
z = self.dz*np.arange( self.Nz_global ) + 0.5*self.dz
plt.plot( z, global_By, '-' )
plt.ylim(-1.1/c, 1.1/c)
plt.xlim(0, self.Lz)
plt.ylabel('$B_y^{n-1/2}$')
plt.xlabel('z')
if save_figure is True:
# Check that the diagnostics folder exists
if os.path.exists('diagnostics') is False:
os.mkdir('diagnostics')
plt.savefig( "diagnostics/iteration%03d.png" %self.n)
def array(self, N=None, filename=None, component=None, root=0):
"""Dump data to numpy format on root processor for all or one snapshot."""
is_root = comm.Get_rank() == root
size = self.get_total_number_probes() if is_root else len(self)
comp = self.value_size() if component is None else 1
if not N is None:
z = zeros((size, comp))
else:
z = zeros((size, comp, self.number_of_evaluations()))
# Get all values
if len(self) > 0:
if not N is None:
for k in range(comp):
if is_root:
ids = self.get_probe_ids()
z[ids,
k] = self.get_probes_component_and_snapshot(k, N)
else:
z[:, k] = self.get_probes_component_and_snapshot(k, N)
else:
for i, (index, probe) in enumerate(self):
j = index if is_root else i
if not N is None:
z[j, :] = probe.get_probe_at_snapshot(N)
else:
for k in range(self.value_size()):
z[j, k, :] = probe.get_probe_sub(k)
# Collect values on root
recvfrom = comm.gather(len(self), root=root)
if is_root:
for j, k in enumerate(recvfrom):
if comm.Get_rank() != j:
ids = comm.recv(source=j, tag=101)
z0 = comm.recv(source=j, tag=102)
z[ids, :] = z0[:, :]
else:
ids = self.get_probe_ids()
comm.send(ids, dest=root, tag=101)
comm.send(z, dest=root, tag=102)
if is_root:
if filename:
if not N is None:
save(filename + "_snapshot_" + str(N), z)
else:
save(filename + "_all", z)
return squeeze(z)
def array(self, N=None, filename=None, component=None, root=0):
"""Dump data to numpy format on root processor for all or one snapshot."""
is_root = comm.Get_rank() == root
size = self.get_total_number_probes() if is_root else len(self)
comp = self.value_size() if component is None else 1
if not N is None:
z = zeros((size, comp))
else:
z = zeros((size, comp, self.number_of_evaluations()))
# Get all values
if len(self) > 0:
if not N is None:
for k in range(comp):
if is_root:
ids = self.get_probe_ids()
z[ids, k] = self.get_probes_component_and_snapshot(k, N)
else:
z[:, k] = self.get_probes_component_and_snapshot(k, N)
else:
for i, (index, probe) in enumerate(self):
j = index if is_root else i
if not N is None:
z[j, :] = probe.get_probe_at_snapshot(N)
else:
for k in range(self.value_size()):
z[j, k, :] = probe.get_probe_sub(k)
# Collect values on root
recvfrom = comm.gather(len(self), root=root)
if is_root:
for j, k in enumerate(recvfrom):
if comm.Get_rank() != j:
ids = comm.recv(source=j, tag=101)
z0 = comm.recv(source=j, tag=102)
z[ids, :] = z0[:, :]
else:
ids = self.get_probe_ids()
comm.send(ids, dest=root, tag=101)
comm.send(z, dest=root, tag=102)
if is_root:
if filename:
if not N is None:
z.dump(filename+"_snapshot_"+str(N)+".probes")
else:
z.dump(filename+"_all.probes")
return squeeze(z)
def get(self):
""" Redefined so that we don't get overlap data. """
# Get our section of the grid (excluding overlap).
if self._xlap is 0:
data = self.data.get()
else:
data = self.data.get()[self._xlap:-self._xlap,:,:]
# return np.concatenate(comm.allgather(data), axis=0) # Super-simple.
result = comm.gather(data) # Gather all peices to root.
if comm.Get_rank() == 0:
# Root node glues everything together.
return np.concatenate(result, axis=0)
else:
return None
def main(args):
parser = argparse.ArgumentParser(
description=
'Generate a legacypipe-compatible CCDs file from a set of reduced imaging.'
)
parser.add_argument('--file_list',
help='List of zeropoint fits files to concatenate')
parser.add_argument('--nproc',
type=int,
default=1,
help='number mpi tasks')
parser.add_argument('--outname',
type=str,
default='combined_legacy_zpt.fits',
help='Output directory.')
parser.add_argument('--fix_hdu',
action='store_true',
default=False,
help='')
parser.add_argument('--remove-file-prefix',
help='Remove prefix from image_filename, if present')
parser.add_argument('--cut',
action='store_true',
default=False,
help='Cut to ccd_cuts==0')
parser.add_argument('--cut-expnum',
action='store_true',
default=False,
help='Cut out rows with expnum==0')
parser.add_argument('files',
nargs='*',
help='Zeropoint files to concatenate')
opt = parser.parse_args(args)
fns = np.array(opt.files)
if opt.file_list:
fns = np.hstack([fns, read_lines(opt.file_list)])
# RUN HERE
if opt.nproc > 1:
from mpi4py.MPI import COMM_WORLD as comm
print('rank %d' % comm.rank)
fn_split = np.array_split(fns, comm.size)
cats = []
for cnt, fn in enumerate(fn_split[comm.rank]):
if comm.rank == 0:
#if cnt % 100 == 0:
print('%d/%d' % (cnt, len(fn_split[comm.rank])))
try:
if opt.fix_hdu:
T = fits_table(fn)
fix_hdu_b4merge(T)
cats.append(T)
else:
cats.append(fits_table(fn))
except IOError:
print('File is bad, skipping: %s' % fn)
cats = merge_tables(cats, columns='fillzero')
#if opt.fix_hdu:
# print('fixing hdu')
# cats= fix_hdu(cats)
# Gather
all_cats = comm.gather(cats, root=0)
if comm.rank == 0:
all_cats = merge_tables(all_cats, columns='fillzero')
write_cat(all_cats, outname=opt.outname)
print("Done")
else:
cats = []
for cnt, fn in enumerate(fns):
if (cnt % 500) == 0:
print('Reading CCDs table %d/%d' % (cnt, len(fns)))
t = fits_table(fn)
cats.append(t)
cats = merge_tables(cats, columns='fillzero')
if opt.fix_hdu:
print('fixing hdu')
cats = fix_hdu(cats)
if opt.cut:
print(len(cats), 'CCDs')
cats.cut(cats.ccd_cuts == 0)
print(len(cats), 'CCDs pass ccd_cuts == 0')
if opt.cut_expnum:
print(len(cats), 'CCDs')
cats.cut(cats.expnum > 0)
print(len(cats), 'CCDs have expnum > 0')
if opt.remove_file_prefix:
fns = []
for fn in cats.image_filename:
if fn.startswith(opt.remove_file_prefix):
fn = fn.replace(opt.remove_file_prefix, '', 1)
fns.append(fn)
cats.image_filename = np.array(fns)
write_cat(cats, outname=opt.outname)
print("Done")
fns = read_lines(opt.file_list)
# RUN HERE
if opt.nproc > 1:
from mpi4py.MPI import COMM_WORLD as comm
fn_split = np.array_split(fns, comm.size)
cats = []
for fn in fn_split[comm.rank]:
try:
cats.append(fits_table(fn))
except IOError:
print('File is bad, skipping: %s' % fn)
cats = merge_tables(cats, columns='fillzero')
# Gather
all_cats = comm.gather(cats, root=0)
if comm.rank == 0:
all_cats = merge_tables(all_cats, columns='fillzero')
if os.path.exists(opt.outname):
os.remove(opt.outname)
all_cats.writeto(opt.outname)
print('Wrote %s' % opt.outname)
print("Done")
else:
cats = []
for cnt, fn in enumerate(fns):
print('Reading %d/%d' % (cnt, len(fns)))
cats.append(fits_table(fn))
cats = merge_tables(cats, columns='fillzero')
if os.path.exists(opt.outname):
os.remove(opt.outname)
#if comm.rank == 0:
# if not os.path.exists(os.path.dirname(outfn)):
# os.makedirs(os.path.dirname(outfn))
#comm.Barrier()
#
#with stdouterr_redirected(to=outfn, comm=None):
#print('rank %d running image %s logging to %s' % (comm.rank,projfn,logfn))
with stdouterr_redirected(to=logfn, comm=None):
#t0=ptime('before-legacy_main',t0)
try:
legacy_main(image_list=[projfn], args=args_copy)
except:
print('Failed: %s' % projfn)
#t0=ptime('after-legacy_main',t0)
else:
legacy_main(image_list=image_list, args=args)
if args.nproc > 1:
if False:
# Wait for all mpi tasks to finish
confirm_files = comm.gather(image_list, root=0)
if comm.rank == 0:
tnow = Time()
print("Done, total time = %s" % (tnow - tbegin, ))
else:
# Crash after smallest image_list is finished (conserve cpu time)
if comm.rank == 0:
print("Done")
else:
tnow = Time()
print("Done, total time = %s" % (tnow - tbegin, ))
CW.Barrier()
#You might have to do some testing of doing a task on one process and sending that to everyone else though. If you are sending big stuff it might take a while.
#Something that is a litle bit different to broadcasting but still shares information on one node with others is scatter, which can scatter elements of a list, or array to all the process. This can be good with those independant tasks. So lets try scattering the pickle files:
if rank == 0:
pickle_file = glob.glob("test_data_*.pkl")
else:
pickle_file = None
pickle_file = CW.scatter(pickle_file, root=0)
print("On rank", rank, "pickle_file =", pickle_file)
CW.Barrier()
#Now that each rank has a file, it can do it's own thing.
#If individual processes were doing their own tasks, but you need to collate everything, you can do that using gather, which puts everything from the different ranks in a list:
data = (rank + 1)**2
print("On rank", rank, "data =", data)
CW.Barrier()
data = CW.gather(data, root=0)
print("On rank", rank, "data =", data)
CW.Barrier()
#When dealing with reading and writing files, it is probbaly best to only let one rank do this because you may run into I/O issues.
best_sol = np.array([])
good_sol = np.array([])
m = 0
e = 0
cel = np.array([])
period = n // size
full_range = np.linspace(0, size, size + 1) * period
if full_range[-1] < n:
full_range[-1] = n
tempsol = sol[int(full_range[rank]):int(full_range[rank + 1])]
cel = goal(tempsol, w)
mpi.Barrier()
cel = mpi.gather(cel)
if rank == 0:
cel = concatenate(cel)
sort = np.argsort(cel)
cel = cel[sort[::-1]]
sol = sol[sort[::-1]]
mpi.Barrier()
sol = mpi.bcast(sol, root=0)
cel = mpi.bcast(cel, root=0)
e = int(n * best)
m = int(n * qual)
best_sol = sol[0:e - 1]
good_sol = sol[e:m - 1]
# Log to unique file
outfn=os.path.join(outdir,"log.rank%d_%s" % \
(comm.rank,\
datetime.datetime.now().strftime("hr%H_min%M")))
else:
fns = np.array_split(image_list, 1)[0]
for fn in fns:
imgfn = os.path.join(rootdir, fn)
wtfn = imgfn.replace('_ooi_', '_oow_')
dqfn = imgfn.replace('_ooi_', '_ood_')
if args.nproc > 1:
# Log to unique file
with stdouterr_redirected(to=outfn, comm=None):
t0 = ptime(
'before: %s' %
os.path.basename(imgfn).replace('.fits.fz', ''), t0)
_ = newWeightMap(wtfn=wtfn, imgfn=imgfn, dqfn=dqfn)
t0 = ptime('after', t0)
else:
_ = newWeightMap(wtfn=wtfn, imgfn=imgfn, dqfn=dqfn)
if args.nproc > 1:
# Wait for all mpi tasks to finish
confirm_files = comm.gather(fns[comm.rank], root=0)
if comm.rank == 0:
tnow = Time()
print("Done, total time = %s" % (tnow - tbegin, ))
else:
tnow = Time()
print("Done, total time = %s" % (tnow - tbegin, ))
if args.dowhat == 'cleanup':
cleanup = True
merge_bybrick(bricks,
outdir=args.outdir,
prefix=args.prefix,
cleanup=cleanup)
elif args.dowhat == 'check':
fns = glob(
os.path.join(
args.outdir,
'input_sample/bybrick/%ssample_*[0-9][0-9].fits' %
args.prefix))
if len(fns) == 0: raise ValueError
fns = np.array_split(fns, comm.size)[comm.rank]
ids = combine(fns)
all_ids = comm.gather(ids, root=0)
if comm.rank == 0:
print('number of unique ids=%d, total number ra,dec pts=%d' % \
(len(set(all_ids)),len(all_ids)))
#if comm.rank == 0:
# merge_draws(outdir=args.outdir,prefix=args.prefix)
#plotobj= PlotTable(outdir=args.outdir,prefix=args.prefix)
#images_split= np.array_split(images, comm.size)
# HACK, not sure if need to wait for all proc to finish
#confirm_files = comm.gather( images_split[comm.rank], root=0 )
#if comm.rank == 0:
# print('Rank 0 gathered the results:')
# print('len(images)=%d, len(gathered)=%d' % (len(images),len(confirm_files)))
# tnow= Time()
# print("TIMING:total %s" % (tnow-tbegin,))
# print("Done")
