def get_meshes(seed, galaxies=False): mesh = {} mesh['s'] = tools.readbigfile(path + ftypefpm % (bs, nc, seed, step) + 'mesh/s/') partp = tools.readbigfile(path + ftypefpm % (bs, nc, seed, step) + 'dynamic/1/Position/') mesh['cic'] = tools.paintcic(partp, bs, nc) mesh['decic'] = tools.decic(mesh['cic'], kk, kny) mesh['R1'] = tools.fingauss(mesh['cic'], kk, R1, kny) mesh['R2'] = tools.fingauss(mesh['cic'], kk, R2, kny) mesh['GD'] = mesh['R1'] - mesh['R2'] hmesh = {} hposall = tools.readbigfile(path + ftype % (bs, ncf, seed, stepf) + 'FOF/PeakPosition/')[1:] massall = tools.readbigfile(path + ftype % (bs, ncf, seed, stepf) + 'FOF/Mass/')[1:].reshape(-1) * 1e10 hposd = hposall[:num].copy() massd = massall[:num].copy() hmesh['pcic'] = tools.paintcic(hposd, bs, nc) hmesh['pnn'] = tools.paintnn(hposd, bs, nc) hmesh['pnnsm'] = tools.fingauss(hmesh['pnn'], kk, R1, kny) hmesh['mnn'] = tools.paintnn(hposd, bs, nc, massd) return mesh, hmesh
def get_meshes(seed, galaxies=False): mesh = {} mesh['s'] = tools.readbigfile(path + ftypefpm % (bs, nc, seed, step) + 'mesh/s/') partp = tools.readbigfile(path + ftypefpm % (bs, nc, seed, step) + 'dynamic/1/Position/') mesh['cic'] = tools.paintcic(partp, bs, nc) mesh['cicovd'] = mesh['cic'] / mesh['cic'].mean() - 1 mesh['decic'] = tools.decic(mesh['cic'], kk, kny) mesh['R1'] = tools.fingauss(mesh['cic'], kk, R1, kny) mesh['R2'] = tools.fingauss(mesh['cic'], kk, R2, kny) mesh['GD'] = mesh['R1'] - mesh['R2'] hmesh = {} hposall = tools.readbigfile(path + ftype % (bs, ncf, seed, stepf) + 'FOF/PeakPosition/')[1:] massall = tools.readbigfile(path + ftype % (bs, ncf, seed, stepf) + 'FOF/Mass/')[1:].reshape(-1) * 1e10 hposd = hposall[:num].copy() massd = massall[:num].copy() print(massd[-1] / 1e10) hmesh['pcic'] = tools.paintcic(hposd, bs, nc) hmesh['pnn'] = tools.paintnn(hposd, bs, nc) hmesh['mnn'] = tools.paintnn(hposd, bs, nc, massd) hmesh['mcic'] = tools.paintcic(hposd, bs, nc, massd) hmesh['mcicovd'] = (hmesh['mcic'] - hmesh['mcic'].mean()) / hmesh['mcic'].mean() data = hmesh['mcicovd'] print(data.min(), data.max(), data.mean(), data.std()) return mesh, hmesh
def get_meshes(seed, pdict=defdict): for i in pdict.keys(): locals()[i] = pdict[i] mesh = {} mesh['s'] = tools.readbigfile(path + ftypefpm%(bs, nc, seed, step) + 'mesh/s/') partp = tools.readbigfile(path + ftypefpm%(bs, nc, seed, step) + 'dynamic/1/Position/') mesh['cic'] = tools.paintcic(partp, bs, ncp) #mesh['decic'] = tools.decic(mesh['cic'], kk, kny) mesh['R1'] = tools.fingauss(mesh['cic'], kk, R1, kny) mesh['R2'] = tools.fingauss(mesh['cic'], kk, R2, kny) mesh['GD'] = mesh['R1'] - mesh['R2'] hmesh = {} hpath = path + ftype%(bs, ncf, seed, stepf) + 'FOF/' hposd = tools.readbigfile(hpath + 'PeakPosition/') massd = tools.readbigfile(hpath + 'Mass/').reshape(-1)*1e10 #galtype = tools.readbigfile(hpath + 'gal_type/').reshape(-1).astype(bool) hposall = tools.readbigfile(path + ftype%(bs, ncf, seed, stepf) + 'FOF/PeakPosition/')[1:] massall = tools.readbigfile(path + ftype%(bs, ncf, seed, stepf) + 'FOF/Mass/')[1:].reshape(-1)*1e10 hposd = hposall[:num].copy() massd = massall[:num].copy() #hmesh['pcic'] = tools.paintcic(hposd, bs, nc) hmesh['pnn'] = tools.paintnn(hposd, bs, ncp) hmesh['mnn'] = tools.paintnn(hposd, bs, ncp, massd) #hmesh['pnnsat'] = tools.paintnn(hposd[galtype], bs, ncp) #hmesh['pnncen'] = tools.paintnn(hposd[~galtype], bs, ncp) return mesh, hmesh
def generate_training_data(): meshes = {} cube_features, cube_target = [[] for i in range(len(cube_sizes)) ], [[] for i in range(len(cube_sizes))] for seed in seeds: mesh = {} partp = tools.readbigfile(path + ftype % (bs, nc, seed, step) + 'dynamic/1/Position/') mesh['cic'] = tools.paintcic(partp, bs, ncp) #mesh['decic'] = tools.decic(mesh['cic'], kk, kny) mesh['R1'] = tools.fingauss(mesh['cic'], kk, R1, kny) #mesh['R2'] = tools.fingauss(mesh['cic'], kk, R2, kny) #mesh['GD'] = mesh['R1'] - mesh['R2'] hmesh = {} hpath = path + ftype % (bs, ncf, seed, stepf) + 'galaxies_n05/galcat/' hposd = tools.readbigfile(hpath + 'Position/') massd = tools.readbigfile(hpath + 'Mass/').reshape(-1) * 1e10 galtype = tools.readbigfile(hpath + 'gal_type/').reshape(-1).astype(bool) hmesh['pnn'] = tools.paintnn(hposd, bs, ncp) hmesh['mnn'] = tools.paintnn(hposd, bs, ncp, massd) hmesh['pnnsat'] = tools.paintnn(hposd[galtype], bs, ncp) hmesh['pnncen'] = tools.paintnn(hposd[~galtype], bs, ncp) meshes[seed] = [mesh, hmesh] print('All the mesh have been generated for seed = %d' % seed) #Create training voxels ftlist = [mesh[i].copy() for i in ftname] ftlistpad = [np.pad(i, pad, 'wrap') for i in ftlist] # targetmesh = hmesh['pnn'] targetmesh = [hmesh[i].copy() for i in tgname] for i, size in enumerate(cube_sizes): print('For size = ', size) if size == nc: features = [np.stack(ftlistpad, axis=-1)] target = [np.stack(targetmesh, axis=-1)] else: numcubes = int(num_cubes / size * 4) features, target = dtools.randomvoxels(ftlistpad, targetmesh, numcubes, max_offset[i], size, cube_sizesft[i], seed=seed, rprob=0) cube_features[i] = cube_features[i] + features cube_target[i] = cube_target[i] + target # # for i in range(cube_sizes.size): cube_target[i] = np.stack(cube_target[i], axis=0) cube_features[i] = np.stack(cube_features[i], axis=0) print(cube_features[i].shape, cube_target[i].shape) return meshes, cube_features, cube_target
def get_meshes(seed, galaxies=False, inverse=True): mesh = {} mesh['s'] = tools.readbigfile(path + ftypefpm % (bs, nc, seed, step) + 'mesh/s/') partp = tools.readbigfile(path + ftypefpm % (bs, nc, seed, step) + 'dynamic/1/Position/') mesh['cic'] = tools.paintcic(partp, bs, nc) mesh['ciclog'] = np.log(1e-3 + mesh['cic']) mesh['cicovd'] = mesh['cic'] / mesh['cic'].mean() - 1 mesh['decic'] = tools.decic(mesh['cic'], kk, kny) mesh['R1'] = tools.fingauss(mesh['cic'], kk, R1, kny) mesh['R2'] = tools.fingauss(mesh['cic'], kk, R2, kny) mesh['GD'] = mesh['R1'] - mesh['R2'] hmesh = {} hposall = tools.readbigfile(path + ftype % (bs, ncf, seed, stepf) + 'FOF/PeakPosition/')[1:] if stellar: massall = np.load(path + ftype % (bs, ncf, seed, stepf) + 'stellarmass.npy') else: massall = tools.readbigfile(path + ftype % (bs, ncf, seed, stepf) + 'FOF/Mass/')[1:].reshape(-1) * 1e10 hposd = hposall[:num].copy() massd = massall[:num].copy() print(massall.min() / 1e10, massall.max() / 1e10) print(massd.min() / 1e10, massd.max() / 1e10) hmesh['pcic'] = tools.paintcic(hposd, bs, nc) hmesh['pnn'] = tools.paintnn(hposd, bs, nc) hmesh['mnn'] = tools.paintnn(hposd, bs, nc, massd) hmesh['mcic'] = tools.paintcic(hposd, bs, nc, massd) hmesh['mcicnomean'] = (hmesh['mcic']) / hmesh['mcic'].mean() hmesh['mcicovd'] = (hmesh['mcic'] - hmesh['mcic'].mean()) / hmesh['mcic'].mean() hmesh['pcicovd'] = (hmesh['pcic'] - hmesh['pcic'].mean()) / hmesh['pcic'].mean() hmesh['pcicovdR3'] = tools.fingauss(hmesh['pcicovd'], kk, R1, kny) if inverse: return hmesh, mesh else: return mesh, hmesh
module = hub.Module('./../code/models/n%02d/%s/%s.hub'%(numd*1e4, suff, chkname)) xx = tf.placeholder(tf.float32, shape=[None, cube_sizeft, cube_sizeft, cube_sizeft, nchannels], name='input') yy = tf.placeholder(tf.float32, shape=[None, cube_size, cube_size, cube_size, 1], name='labels') output = module(dict(input=xx, label=yy, keepprob=1), as_dict=True)['prediction'] sess = tf.Session() sess.run(tf.initializers.global_variables()) # ############################# meshes = {} cube_features, cube_target = [], [] for seed in seeds: mesh = {} partp = tools.readbigfile(path + ftype%(bs, nc, seed, step) + 'dynamic/1/Position/') mesh['cic'] = tools.paintcic(partp, bs, ncp) mesh['decic'] = tools.decic(mesh['cic'], kk, kny) mesh['R1'] = tools.fingauss(mesh['cic'], kk, R1, kny) mesh['R2'] = tools.fingauss(mesh['cic'], kk, R2, kny) mesh['GD'] = mesh['R1'] - mesh['R2'] mesh['s'] = tools.readbigfile(path + ftype%(bs, nc, seed, step) + 'mesh/s/') hmesh = {} hposall = tools.readbigfile(path + ftype%(bs, ncf, seed, stepf) + 'FOF/PeakPosition/')[1:] hposd = hposall[:num].copy() hmesh['pcic'] = tools.paintcic(hposd, bs, nc) hmesh['pnn'] = tools.paintnn(hposd, bs, ncp) hmesh['target'] = hmesh['pnn'].copy() print('All the mesh have been generated for seed = %d'%seed) #Create training voxels ftlist = [mesh[i].copy() for i in ftname]
def main(_): infield = True dtype = tf.float32 mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape) nc, bs = FLAGS.nc, FLAGS.box_size a0, a, nsteps = FLAGS.a0, FLAGS.af, FLAGS.nsteps stages = np.linspace(a0, a, nsteps, endpoint=True) numd = 1e-3 ##Begin here klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) #pypath = '/global/cscratch1/sd/chmodi/cosmo4d/output/version2/L0400_N0128_05step-fof/lhd_S0100/n10/opt_s999_iM12-sm3v25off/meshes/' final = tools.readbigfile('../data//L0400_N0128_S0100_05step/mesh/d/') ic = tools.readbigfile('../data/L0400_N0128_S0100_05step/mesh/s/') fpos = tools.readbigfile( '../data/L0400_N0128_S0100_05step/dynamic/1/Position/') hpos = tools.readbigfile( '../data/L0400_N0512_S0100_40step/FOF/PeakPosition//')[1:int(bs**3 * numd)] hmass = tools.readbigfile( '../data/L0400_N0512_S0100_40step/FOF/Mass//')[1:int(bs**3 * numd)].flatten() meshpos = tools.paintcic(hpos, bs, nc) meshmass = tools.paintcic(hpos, bs, nc, hmass.flatten() * 1e10) data = meshmass data /= data.mean() data -= 1 kv = tools.fftk([nc, nc, nc], bs, symmetric=True, dtype=np.float32) datasm = tools.fingauss(data, kv, 3, np.pi * nc / bs) ic, data = np.expand_dims(ic, 0), np.expand_dims(data, 0).astype(np.float32) datasm = np.expand_dims(datasm, 0).astype(np.float32) print("Min in data : %0.4e" % datasm.min()) np.save(fpath + 'ic', ic) np.save(fpath + 'data', data) #################################################### # tf.reset_default_graph() tfic = tf.constant(ic.astype(np.float32)) state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) with tf.Session() as sess: state = sess.run(final_state) fpos = state[0, 0] * bs / nc bparams, bmodel = getbias(bs, nc, data[0] + 1, ic[0], fpos) #bmodel += 1 #np.expand_dims(bmodel, 0) + 1 errormesh = data - np.expand_dims(bmodel, 0) kerror, perror = tools.power(errormesh[0] + 1, boxsize=bs) kerror, perror = kerror[1:], perror[1:] print("Error power spectra", kerror, perror) print("\nkerror", kerror.min(), kerror.max(), "\n") print("\nperror", perror.min(), perror.max(), "\n") suff = "-error" dg.saveimfig(suff, [ic, errormesh], [ic, data], fpath + '/figs/') dg.save2ptfig(suff, [ic, errormesh], [ic, data], fpath + '/figs/', bs) ipkerror = iuspline(kerror, perror) #################################################### #stdinit = srecon.standardinit(bs, nc, meshpos, hpos, final, R=8) recon_estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir=fpath) def predict_input_fn(data=data, M0=0., w=3., R0=0., off=None, istd=None, x0=None): features = {} features['datasm'] = data features['R0'] = R0 features['x0'] = x0 features['bparams'] = bparams features['ipkerror'] = [kerror, perror] #ipkerror return features, None eval_results = recon_estimator.predict( input_fn=lambda: predict_input_fn(x0=ic), yield_single_examples=False) for i, pred in enumerate(eval_results): if i > 0: break suff = '-model' dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/') dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/', bs) np.save(fpath + '/reconmeshes/ic_true' + suff, pred['ic']) np.save(fpath + '/reconmeshes/fin_true' + suff, pred['final']) np.save(fpath + '/reconmeshes/model_true' + suff, pred['model']) # randominit = np.random.normal(size=data.size).reshape(data.shape) #eval_results = recon_estimator.predict(input_fn=lambda : predict_input_fn(x0 = np.expand_dims(stdinit, 0)), yield_single_examples=False) eval_results = recon_estimator.predict( input_fn=lambda: predict_input_fn(x0=randominit), yield_single_examples=False) for i, pred in enumerate(eval_results): if i > 0: break suff = '-init' dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/') dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/', bs) np.save(fpath + '/reconmeshes/ic_init' + suff, pred['ic']) np.save(fpath + '/reconmeshes/fin_init' + suff, pred['final']) np.save(fpath + '/reconmeshes/model_init' + suff, pred['model']) # # Train and evaluate model. RRs = [4., 2., 1., 0.5, 0.] niter = 100 iiter = 0 for R0 in RRs: print('\nFor iteration %d\n' % iiter) print('With R0=%0.2f \n' % (R0)) def train_input_fn(): features = {} features['datasm'] = data features['R0'] = R0 features['bparams'] = bparams features['ipkerror'] = [kerror, perror] #ipkerror #features['x0'] = np.expand_dims(stdinit, 0) features['x0'] = randominit features['lr'] = 0.01 return features, None recon_estimator.train(input_fn=train_input_fn, max_steps=iiter + niter) eval_results = recon_estimator.predict(input_fn=predict_input_fn, yield_single_examples=False) for i, pred in enumerate(eval_results): if i > 0: break iiter += niter # suff = '-%d-R%d' % (iiter, R0) dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/') dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/', bs) np.save(fpath + '/reconmeshes/ic' + suff, pred['ic']) np.save(fpath + '/reconmeshes/fin' + suff, pred['final']) np.save(fpath + '/reconmeshes/model' + suff, pred['model']) sys.exit(0) ## exit(0)
num = int(numd * bs**3) seed = 100 R1 = 3 R2 = 3 * 1.2 kny = np.pi * ncp / bs kk = tools.fftk((ncp, ncp, ncp), bs) suff = '2ftmdg256' ############################# ##Read data and generate meshes #mesh = tools.readbigfile(path + ftype%(bs, nc, seed, step) + 'mesh/d/') partp = tools.readbigfile(path + ftype % (bs, nc, seed, step) + 'dynamic/1/Position/') mesh = tools.paintcic(partp, bs, ncp) meshdecic = tools.decic(mesh, kk, kny) meshR1 = tools.fingauss(mesh, kk, R1, kny) meshR2 = tools.fingauss(mesh, kk, R2, kny) meshdg = meshR1 - meshR2 #ftlist = [meshdecic.copy(), meshR1.copy(), meshdg.copy()] hposall = tools.readbigfile(path + ftype % (bs, ncf, seed, stepf) + 'FOF/PeakPosition/')[1:] hposd = hposall[:num].copy() #hposall = hposall[:2*num] # hpmeshall = tools.paintcic(hposall, bs, nc) # hpmeshd = tools.paintcic(hposd, bs, nc) #hpmeshall = tools.paintnn(hposall, bs, ncp) hpmeshd = tools.paintnn(hposd, bs, ncp)
def main(_): infield = True dtype = tf.float32 mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape) nc, bs = FLAGS.nc, FLAGS.box_size a0, a, nsteps = FLAGS.a0, FLAGS.af, FLAGS.nsteps stages = np.linspace(a0, a, nsteps, endpoint=True) numd = 1e-3 startw = time.time() print(mesh_shape) #layout_rules = mtf.convert_to_layout_rules(FLAGS.layout) #mesh_shape = [("row", FLAGS.nx), ("col", FLAGS.ny)] layout_rules = [("nx_lr", "row"), ("ny_lr", "col"), ("nx", "row"), ("ny", "col"), ("ty", "row"), ("tz", "col"), ("ty_lr", "row"), ("tz_lr", "col"), ("nx_block", "row"), ("ny_block", "col")] # Resolve the cluster from SLURM environment cluster = tf.distribute.cluster_resolver.SlurmClusterResolver( {"mesh": mesh_shape.size // FLAGS.gpus_per_task}, port_base=8822, gpus_per_node=FLAGS.gpus_per_node, gpus_per_task=FLAGS.gpus_per_task, tasks_per_node=FLAGS.tasks_per_node) cluster_spec = cluster.cluster_spec() print(cluster_spec) # Create a server for all mesh members server = tf.distribute.Server(cluster_spec, "mesh", cluster.task_id) print(server) if cluster.task_id > 0: server.join() # Otherwise we are the main task, let's define the devices devices = [ "/job:mesh/task:%d/device:GPU:%d" % (i, j) for i in range(cluster_spec.num_tasks("mesh")) for j in range(FLAGS.gpus_per_task) ] print("List of devices", devices) mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl( mesh_shape, layout_rules, devices) ##Begin here klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) final = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0128_S0100_05step/mesh/d/' ) ic = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0128_S0100_05step/mesh/s/' ) pypath = '/global/cscratch1/sd/chmodi/cosmo4d/output/version2/L0400_N0128_05step-fof/lhd_S0100/n10/opt_s999_iM12-sm3v25off/meshes/' fin = tools.readbigfile(pypath + 'decic//') hpos = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0512_S0100_40step/FOF/PeakPosition//' )[1:int(bs**3 * numd)] hmass = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0512_S0100_40step/FOF/Mass//' )[1:int(bs**3 * numd)].flatten() #meshpos = tools.paintcic(hpos, bs, nc) meshmass = tools.paintcic(hpos, bs, nc, hmass.flatten() * 1e10) data = meshmass kv = tools.fftk([nc, nc, nc], bs, symmetric=True, dtype=np.float32) datasm = tools.fingauss(data, kv, 3, np.pi * nc / bs) ic, data = np.expand_dims(ic, 0), np.expand_dims(data, 0).astype(np.float32) datasm = np.expand_dims(datasm, 0).astype(np.float32) print("Min in data : %0.4e" % datasm.min()) ic, data = np.expand_dims(ic, 0), np.expand_dims(data, 0).astype(np.float32) np.save(fpath + 'ic', ic) np.save(fpath + 'data', data) #################################################### tf.reset_default_graph() print('ic constructed') graph = mtf.Graph() mesh = mtf.Mesh(graph, "my_mesh") initial_conditions, data_field, loss, var_grads, update_op, linear_op, input_field, lr, R0, M0, width, chisq, prior, tf_off, tf_istd = recon_prototype( mesh, datasm, nc=FLAGS.nc, batch_size=FLAGS.batch_size, dtype=dtype) # Lower mesh computation start = time.time() lowering = mtf.Lowering(graph, {mesh: mesh_impl}) restore_hook = mtf.MtfRestoreHook(lowering) end = time.time() print('\n Time for lowering : %f \n' % (end - start)) tf_initc = lowering.export_to_tf_tensor(initial_conditions) tf_data = lowering.export_to_tf_tensor(data_field) tf_chisq = lowering.export_to_tf_tensor(chisq) tf_prior = lowering.export_to_tf_tensor(prior) tf_grads = lowering.export_to_tf_tensor(var_grads[0]) #tf_lr = lowering.export_to_tf_tensor(lr) tf_linear_op = lowering.lowered_operation(linear_op) tf_update_ops = lowering.lowered_operation(update_op) n_block_x, n_block_y, n_block_z = FLAGS.nx, FLAGS.ny, 1 nc = FLAGS.nc with tf.Session(server.target) as sess: start = time.time() sess.run(tf_linear_op, feed_dict={input_field: ic}) ic_check, data_check = sess.run([tf_initc, tf_data], {width: 3}) dg.saveimfig('-check', [ic_check, data_check], [ic, data], fpath + '/figs/') dg.save2ptfig('-check', [ic_check, data_check], [ic, data], fpath + '/figs/', bs) print('Total time taken for mesh thingy is : ', time.time() - start) sess.run(tf_linear_op, feed_dict={ input_field: np.random.normal(size=ic.size).reshape(ic.shape) }) ic0, data0 = sess.run([tf_initc, tf_data], {width: 3}) dg.saveimfig('-init', [ic0, data0], [ic, data], fpath) start = time.time() titer = 20 niter = 101 iiter = 0 start0 = time.time() RRs = [4, 2, 1, 0.5, 0] wws = [1, 2, 3] lrs = np.array([0.1, 0.1, 0.1, 0.1, 0.1]) * 2 #lrs = [0.1, 0.05, 0.01, 0.005, 0.001] readin = True mm0, ww0, RR0 = 1e12, 3, 0.5 if readin: icread = np.load(fpath + '/figs-M%02d-R%02d-w%01d/ic_recon.npy' % (np.log10(mm0), 10 * RR0, ww0)) sess.run(tf_linear_op, feed_dict={input_field: icread}) for mm in [1e12, 1e11]: print('Fraction of points above 1 for mm = %0.2e: ' % mm, (datasm > mm).sum() / datasm.size) noisefile = '/project/projectdirs/m3058/chmodi/cosmo4d/train/L0400_N0128_05step-n10/width_3/Wts_30_10_1/r1rf1/hlim-13_nreg-43_batch-5/eluWts-10_5_1/blim-20_nreg-23_batch-100/hist_M%d_na.txt' % ( np.log10(mm) * 10) offset, ivar = setnoise(datasm, noisefile, noisevar=0.25) for iR, zlR in enumerate(zip(RRs, lrs)): RR, lR = zlR for ww in wws: for ff in [ fpath + '/figs-M%02d-R%02d-w%01d' % (np.log10(mm), 10 * RR, ww) ]: try: os.makedirs(ff) except Exception as e: print(e) if readin: if mm > mm0: continue elif mm == mm0 and RR > RR0: print(RR, RR0, RRs) continue elif RR == RR0 and ww <= ww0: print(ww, ww0, wws) continue else: print('Starting from %0.2e' % mm, RR, ww) print('Do for %0.2e' % mm, RR, ww) for i in range(niters[iR]): iiter += 1 sess.run( tf_update_ops, { lr: lR, M0: mm, R0: RR, width: ww, tf_off: offset, tf_istd: ivar**0.5 }) if (i % titer == 0): end = time.time() print('Iter : ', i) print('Time taken for %d iterations: ' % titer, end - start) start = end ## ic1, data1, cc, pp = sess.run( [tf_initc, tf_data, tf_chisq, tf_prior], { M0: mm, R0: RR, width: ww, tf_off: offset, tf_istd: ivar**0.5 }) print('Chisq and prior are : ', cc, pp) dg.saveimfig(i, [ic1, data1], [ic, data], ff) dg.save2ptfig(i, [ic1, data1], [ic, data], ff, bs) ic1, data1 = sess.run([tf_initc, tf_data], {width: ww}) np.save(ff + '/ic_recon', ic1) np.save(ff + '/data_recon', data1) dg.saveimfig(iiter, [ic1, data1], [ic, data], fpath + '/figs') dg.save2ptfig(iiter, [ic1, data1], [ic, data], fpath + '/figs', bs) wws = [3] RRs = [0] niters = [201, 101, 201] lrs = np.array([0.1, 0.1, 0.1]) ic1, data1 = sess.run([tf_initc, tf_data], {width: 3}) print('Total time taken for %d iterations is : ' % iiter, time.time() - start0) dg.saveimfig('', [ic1, data1], [ic, data], fpath) dg.save2ptfig('', [ic1, data1], [ic, data], fpath, bs) np.save(fpath + 'ic_recon', ic1) np.save(fpath + 'data_recon', data1) print('Total wallclock time is : ', time.time() - start0) ## exit(0)
def main(_): infield = True dtype = tf.float32 mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape) nc, bs = FLAGS.nc, FLAGS.box_size a0, a, nsteps = FLAGS.a0, FLAGS.af, FLAGS.nsteps stages = np.linspace(a0, a, nsteps, endpoint=True) numd = 1e-3 ##Begin here klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) #pypath = '/global/cscratch1/sd/chmodi/cosmo4d/output/version2/L0400_N0128_05step-fof/lhd_S0100/n10/opt_s999_iM12-sm3v25off/meshes/' final = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0128_S0100_05step/mesh/d/' ) ic = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0128_S0100_05step/mesh/s/' ) fpos = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0128_S0100_05step/dynamic/1/Position/' ) aa = 1 zz = 1 / aa - 1 rsdfactor = float(100 / (aa**2 * cosmo.H(zz).value**1)) print('\nRsdfactor used is : ', rsdfactor) hpos = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0512_S0100_40step/FOF/PeakPosition//' )[1:int(bs**3 * numd)] hvel = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0512_S0100_40step/FOF/CMVelocity//' )[1:int(bs**3 * numd)] rsdpos = hpos + hvel * rsdfactor * np.array([0, 0, 1]) print('Effective displacement : ', (hvel[:, -1] * rsdfactor).std()) hmass = tools.readbigfile( '/project/projectdirs/m3058/chmodi/cosmo4d/data/L0400_N0512_S0100_40step/FOF/Mass//' )[1:int(bs**3 * numd)].flatten() meshpos = tools.paintcic(rsdpos, bs, nc) meshmass = tools.paintcic(rsdpos, bs, nc, hmass.flatten() * 1e10) data = meshmass kv = tools.fftk([nc, nc, nc], bs, symmetric=True, dtype=np.float32) datasm = tools.fingauss(data, kv, 3, np.pi * nc / bs) ic, data = np.expand_dims(ic, 0), np.expand_dims(data, 0).astype(np.float32) datasm = np.expand_dims(datasm, 0).astype(np.float32) print("Min in data : %0.4e" % datasm.min()) # #################################################### stdinit = srecon.standardinit(bs, nc, meshpos, hpos, final, R=8) recon_estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir=fpath) def predict_input_fn(data=data, M0=0., w=3., R0=0., off=None, istd=None, x0=None): features = {} features['datasm'] = data features['rsdfactor'] = rsdfactor features['M0'] = M0 features['w'] = w features['R0'] = R0 features['off'] = off features['istd'] = istd features['x0'] = x0 return features, None eval_results = recon_estimator.predict( input_fn=lambda: predict_input_fn(x0=ic), yield_single_examples=False) for i, pred in enumerate(eval_results): if i > 0: break suff = '-model' dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/') dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/', bs) np.save(fpath + '/reconmeshes/ic_true' + suff, pred['ic']) np.save(fpath + '/reconmeshes/fin_true' + suff, pred['final']) np.save(fpath + '/reconmeshes/model_true' + suff, pred['model']) # randominit = np.random.normal(size=data.size).reshape(data.shape) #eval_results = recon_estimator.predict(input_fn=lambda : predict_input_fn(x0 = np.expand_dims(stdinit, 0)), yield_single_examples=False) eval_results = recon_estimator.predict( input_fn=lambda: predict_input_fn(x0=randominit), yield_single_examples=False) for i, pred in enumerate(eval_results): if i > 0: break suff = '-init' dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/') dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/', bs) np.save(fpath + '/reconmeshes/ic_init' + suff, pred['ic']) np.save(fpath + '/reconmeshes/fin_init' + suff, pred['final']) np.save(fpath + '/reconmeshes/model_init' + suff, pred['model']) # # Train and evaluate model. mms = [1e12, 1e11] wws = [1., 2., 3.] RRs = [4., 2., 1., 0.5, 0.] niter = 100 iiter = 0 for mm in mms: noisefile = '/project/projectdirs/m3058/chmodi/cosmo4d/train/L0400_N0128_05step-n10/width_3/Wts_30_10_1/r1rf1/hlim-13_nreg-43_batch-5/eluWts-10_5_1/blim-20_nreg-23_batch-100/hist_M%d_na.txt' % ( np.log10(mm) * 10) offset, ivar = setnoise(datasm, noisefile, noisevar=0.25) istd = ivar**0.5 if not FLAGS.offset: offset = None if not FLAGS.istd: istd = None for R0 in RRs: for ww in wws: print('\nFor iteration %d\n' % iiter) print('With mm=%0.2e, R0=%0.2f, ww=%d \n' % (mm, R0, ww)) def train_input_fn(): features = {} features['datasm'] = datasm features['rsdfactor'] = rsdfactor features['M0'] = mm features['w'] = ww features['R0'] = R0 features['off'] = offset features['istd'] = istd features['x0'] = np.expand_dims( stdinit, 0 ) #np.random.normal(size=datasm.size).reshape(datasm.shape) features['lr'] = 0.01 return features, None recon_estimator.train(input_fn=train_input_fn, max_steps=iiter + niter) eval_results = recon_estimator.predict( input_fn=predict_input_fn, yield_single_examples=False) for i, pred in enumerate(eval_results): if i > 0: break iiter += niter # suff = '-%d-M%d-R%d-w%d' % (iiter, np.log10(mm), R0, ww) dg.saveimfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/') dg.save2ptfig(suff, [pred['ic'], pred['model']], [ic, data], fpath + '/figs/', bs) suff = '-M%d-R%d-w%d' % (np.log10(mm), R0, ww) np.save(fpath + '/reconmeshes/ic' + suff, pred['ic']) np.save(fpath + '/reconmeshes/fin' + suff, pred['final']) np.save(fpath + '/reconmeshes/model' + suff, pred['model']) RRs = [1., 0.5, 0.] wws = [3.] niter = 200 sys.exit(0) ## exit(0)