def train_step(inputs): x_true, y = inputs if args.stdinit: x_init = y[:, 1] / b1eul if args.diffps: x_init = x_init + linear_field( nc, bs, ipkdiff, batch_size=y.shape[0]) elif args.priorinit: x_init = linear_field(nc, bs, ipklin, batch_size=y.shape[0]) else: x_init = tf.random.normal(x_true.shape) y = y[:, 0] if len(rim.trainable_variables) == 0: #Hack since sometimes this si the first time RIM is called and so hasn't been inisitalized i = 0 a, b, c = x_init[i:i + 1], y[i:i + 1], x_true[i:i + 1] _ = rim(tf.constant(a), tf.constant(b), grad_fn, tf.constant(c), grad_params)[1] / args.batch_size # gradients = [0.] * len(rim.trainable_variables) #n = args.sims_in_loop for i in range(args.batch_size // world_size): with tf.GradientTape() as tape: a, b, c = x_init[i:i + 1], y[i:i + 1], x_true[i:i + 1] loss = rim(tf.constant(a), tf.constant(b), grad_fn, tf.constant(c), grad_params)[1] / args.batch_size grads = tape.gradient(loss, rim.trainable_variables) for j in range(len(grads)): gradients[j] = gradients[j] + grads[j] optimizer.apply_gradients(zip(gradients, rim.trainable_variables)) return loss
def test_step(inputs): x_true, y = inputs #x_init = tf.random.normal(x_true.shape) if args.stdinit: x_init = y[:, 1] / b1eul + linear_field( nc, bs, ipkdiff, batch_size=y.shape[0]) #x_init = y[:, 1] y = y[:, 0] elif args.priorinit: x_init = linear_field(nc, bs, ipklin, batch_size=y.shape[0]) else: x_init = tf.random.normal(x_true.shape) x_pred = rim(x_init, y, grad_fn, grad_params) return x_pred, x_init, x_true, y
def train_step(inputs): x_true, y = inputs if len(rim.trainable_variables) == 0: #Hack since sometimes this si the first time RIM is called and so hasn't been inisitalized i = 0 b, c = y[i:i+1], x_true[i:i+1] if args.stdinit: a = b[:, 1] / b1eul if args.diffps : a = a + linear_field(nc, bs, ipkdiff, batch_size=b.shape[0]) else: a = tf.random.normal(c.shape) b = b[:, 0] try: a, b, c = tf.constant(a), tf.constant(b), tf.constant(c) except: pass _ = rim(a, b, grad_fn, c, grad_params)[1] / args.batch_size # gradients = [0.]*len(rim.trainable_variables) for i in range(args.batch_size // world_size): with tf.GradientTape() as tape: b, c = y[i:i+1], x_true[i:i+1] a = tf.random.normal(c.shape) b = b[:, 0] try: a, b, c = tf.constant(a), tf.constant(b), tf.constant(c) except: pass loss = rim(a, b, grad_fn, c, grad_params)[1] / args.batch_size grads = tape.gradient(loss, rim.trainable_variables) for j in range(len(grads)): gradients[j] = gradients[j] + grads[j] optimizer.apply_gradients(zip(gradients, rim.trainable_variables)) return loss
def check_2pt(xx, yy, rim, grad_fn, compares, nrim=10, fname=None): truemesh = [xx[0], yy[0]] rimpreds = [] for it in range(nrim): x_init = flowpm.linear_field(nc, bs, ipklin, batch_size=xx.shape[0]).numpy() #x_init = np.random.normal(size=xx.size).reshape(xx.shape).astype(np.float32) #x_init = (yy - (yy.max() - yy.min())/2.)/yy.std() + np.random.normal(size=xx.size).reshape(xx.shape).astype(np.float32) pred = rim(tf.constant(x_init), tf.constant(yy), grad_fn)[-1] rimpreds.append([pred[0].numpy(), pm(pred)[0].numpy()]) fig, ax = plt.subplots(1, 2, figsize=(9, 4), sharex=True) for ip, preds in enumerate(rimpreds): k, pks = tools.get_ps(preds, truemesh, bs) for i in range(2): lbl = None if ip == 0 and i == 0: lbl = 'Linear' if ip == 0 and i == 1: lbl = 'Final' ax[0].plot(k, pks[i][2] / (pks[i][0] * pks[i][1])**0.5, 'C%d-' % i, alpha=0.4, label=lbl) ax[1].plot(k, (pks[i][0] / pks[i][1])**0.5, 'C%d-' % i, alpha=0.4) lss = ['-', '--', ':', '-.'] lws = [1, 1, 2, 2] lbls = ['Adam', 'Adam 10x', 'Best recon'] #for ip, preds in enumerate([pred_adam, pred_adam10]): for ip, preds in enumerate(compares): k, pks = tools.get_ps(preds, truemesh, bs) for i in range(2): lbl = None if i == 0: lbl = lbls[ip] ax[0].plot(k, pks[i][2] / (pks[i][0] * pks[i][1])**0.5, 'C%d' % i, ls=lss[ip + 1], lw=lws[ip + 1]) ax[1].plot(k, (pks[i][0] / pks[i][1])**0.5, 'C%d' % i, label=lbl, ls=lss[ip + 1], lw=lws[ip + 1]) for axis in ax: axis.semilogx() axis.grid(which='both') axis.legend(fontsize=12) axis.set_xlabel('k(h/Mpc)', fontsize=12) ax[0].set_ylim(-0.1, 1.2) ax[1].set_ylim(-0.5, 2.0) ax[0].set_ylabel('$r_c$', fontsize=12) ax[1].set_ylabel('$t_f$', fontsize=12) plt.tight_layout() if fname is not None: plt.savefig(fname) else: plt.savefig('rim-2pt.png') plt.close()
def test_step(inputs): x_true, y = inputs if args.stdinit: x_init = y[:, 1] / b1eul if args.diffps : x_init = x_init + linear_field(nc, bs, ipkdiff, batch_size=y.shape[0]) else: x_init = tf.random.normal(x_true.shape) y = y[:, 0] x_pred = rim(x_init, y, grad_fn, x_true, grad_params)[0] return x_pred, x_init, x_true, y
def test_save_state(): """ Tests the BigFile saving function """ klin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) a0 = 0.1 nc = [16, 16, 16] boxsize = [100., 100., 100.] cosmo = flowpm.cosmology.Planck15() initial_conditions = flowpm.linear_field( nc, # size of the cube boxsize, # Physical size of the cube ipklin, # Initial powerspectrum batch_size=2) # Sample particles state = flowpm.lpt_init(cosmo, initial_conditions, a0) with tempfile.TemporaryDirectory() as tmpdirname: filename = tmpdirname + '/testsave' save_state(cosmo, state, a0, nc, boxsize, filename) # Now try to reload the information using BigFile bf = bigfile.BigFile(filename) # Testing recovery of header header = bf['Header'] assert_allclose(np.array(header.attrs['NC']), np.array(nc)) assert_allclose(np.array(header.attrs['BoxSize']), np.array(boxsize)) assert_allclose(np.array(header.attrs['OmegaCDM']), np.array(cosmo.Omega_c)) assert_allclose(np.array(header.attrs['OmegaB']), np.array(cosmo.Omega_b)) assert_allclose(np.array(header.attrs['OmegaK']), np.array(cosmo.Omega_k)) assert_allclose(np.array(header.attrs['h']), np.array(cosmo.h)) assert_allclose(np.array(header.attrs['Sigma8']), np.array(cosmo.sigma8)) assert_allclose(np.array(header.attrs['w0']), np.array(cosmo.w0)) assert_allclose(np.array(header.attrs['wa']), np.array(cosmo.wa)) assert_allclose(np.array(header.attrs['Time']), np.array(a0)) # Testing recovery of data pos = bf['1/Position'] assert_allclose(pos[:], state[0, 1].numpy() / nc[0] * boxsize[0]) vel = bf['1/Velocity'] assert_allclose(vel[:], state[1, 1].numpy() / nc[0] * boxsize[0]) # Closing file bf.close()
def pm_data_test(self, dummy): args = self.args nc, bs = args.nc, args.bs print("PM graph") linear = flowpm.linear_field(nc, bs, args.ipklin, batch_size=args.world_size) base = self.pm(linear) sample = self.gal_sample(base) return linear, sample
def pm_data_test(dummy): print("PM graph") linear = flowpm.linear_field(nc, bs, ipklin, batch_size=world_size) if args.nbody: print('Nobdy sim') state = lpt_init(linear, a0=a0, order=args.lpt_order) final_state = nbody(state, stages, nc) else: print('ZA/2LPT sim') final_state = lpt_init(linear, a0=af, order=args.lpt_order) tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0]) return linear, tfinal_field
def test_step(inputs): x_true, y = inputs #x_init = tf.random.normal(x_true.shape) if args.stdinit: x_init = y[:, 1] / b1eul + linear_field( nc, bs, ipkdiff, batch_size=y.shape[0]) #x_init = y[:, 1] y = y[:, 0] else: print('data init') x_init = y x_pred = cnn(tf.expand_dims(x_init, -1))[..., 0] #x_pred = cnn(x_init) return x_pred, x_init, x_true, y
def train_step(inputs): x_true, y = inputs #x_init = tf.random.normal(x_true.shape) if args.stdinit: x_init = y[:, 1] / b1eul + linear_field( nc, bs, ipkdiff, batch_size=y.shape[0]) #x_init = y[:, 1] y = y[:, 0] elif args.priorinit: x_init = linear_field(nc, bs, ipklin, batch_size=y.shape[0]) else: x_init = tf.random.normal(x_true.shape) with tf.GradientTape() as tape: x_pred = rim(x_init, y, grad_fn, grad_params) res = (x_true - x_pred) #print(res.shape) loss = tf.reduce_mean( tf.square(res), axis=(0, 2, 3, 4)) ##This is not advised, come back to this loss = tf.reduce_sum(loss) / args.batch_size gradients = tape.gradient(loss, rim.trainable_variables) optimizer.apply_gradients(zip(gradients, rim.trainable_variables)) print(optimizer._decayed_lr(tf.float32)) return loss
def main(_): cosmology = flowpm.cosmology.Planck15() # Compute the k vectora that will be needed in the PGD fit k, _ = flowpm.power_spectrum(tf.zeros([1] + [FLAGS.nc] * 3), boxsize=np.array([FLAGS.box_size] * 3), kmin=np.pi / FLAGS.box_size, dk=2 * np.pi / FLAGS.box_size) # Create some initial conditions klin = tf.constant(np.logspace(-4, 1, 512), dtype=tf.float32) pk = linear_matter_power(cosmology, klin) pk_fun = lambda x: tf.cast( tf.reshape( interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]), klin, pk), x.shape), tf.complex64) initial_conditions = flowpm.linear_field( [FLAGS.nc, FLAGS.nc, FLAGS.nc], [FLAGS.box_size, FLAGS.box_size, FLAGS.box_size], pk_fun, batch_size=FLAGS.batch_size) initial_state = flowpm.lpt_init(cosmology, initial_conditions, FLAGS.a_init) stages = np.linspace(FLAGS.a_init, 1., FLAGS.nsteps, endpoint=True) print('Starting simulation') state, scale_factors, pgdparams = fit_nbody(cosmology, initial_state, stages, [FLAGS.nc, FLAGS.nc, FLAGS.nc], pm_nc_factor=FLAGS.B) print('Simulation done') pickle.dump( { 'B': FLAGS.B, 'nsteps': FLAGS.nsteps, 'params': pgdparams, 'scale_factors': scale_factors, 'cosmology': cosmology.to_dict(), 'boxsize': FLAGS.box_size, 'nc': FLAGS.nc }, open(FLAGS.filename, "wb"))
def pm_poisson(): print("PM graph") linear = flowpm.linear_field(nc, bs, ipklin, batch_size=batch_size) if args.nbody: print('Nobdy sim') state = lpt_init(linear, a0=a0, order=args.lpt_order) final_state = nbody(state, stages, nc) else: print('ZA/2LPT sim') final_state = lpt_init(linear, a0=af, order=args.lpt_order) tfinal_field = cic_paint(tf.zeros_like(linear), final_state[0]) base = tfinal_field if Rsm != 0: basek = r2c3d(tfinal_field, norm=nc**3) basek = tf.multiply(basek, tf.cast(smwts, tf.complex64)) base = c2r3d(basek, norm=nc**3) galmean = tfp.distributions.Poisson(rate = plambda * (1 + base)) result = galmean.sample() return linear, tfinal_field, result, base
def train_step(inputs): x_true, y = inputs #x_init = tf.random.normal(x_true.shape) if args.stdinit: print('std init') x_init = y[:, 1] / b1eul + linear_field( nc, bs, ipkdiff, batch_size=y.shape[0]) #x_init = y[:, 1] y = y[:, 0] else: print('data init') x_init = y with tf.GradientTape() as tape: x_pred = cnn(tf.expand_dims(x_init, -1))[..., 0] print(x_pred.shape) res = (x_true - x_pred) #print(res.shape) loss = tf.reduce_mean(tf.square(res)) gradients = tape.gradient(loss, cnn.trainable_variables) optimizer.apply_gradients(zip(gradients, cnn.trainable_variables)) print("learnign rate : ", optimizer._decayed_lr(tf.float32)) print('Len trainable : ', len(cnn.trainable_variables)) return loss
else: ofolder = './models/L%04d_N%03d/LPT%d%s/' % (bs, nc, args.lpt_order, suffpath) try: os.makedirs(ofolder) except Exception as e: print(e) ####################################### x_test, y_test = testdata[0:1, 0], testdata[0:1, 1:] x_test = tf.constant(x_test, dtype=tf.float32) if args.stdinit: x_init = tf.constant(y_test[:, 1] / b1eul, dtype=tf.float32) if args.diffps: x_init = x_init + linear_field( nc, bs, ipkdiff, batch_size=y_test.shape[0]) elif args.priorinit: x_init = linear_field(nc, bs, ipklin, batch_size=y_test.shape[0]) else: x_init = tf.random.normal(x_test.shape) y_test = tf.constant(y_test[:, 0]) minic, minfin = datamodel.reconstruct(tf.constant(y_test), bias, errormesh, RRs=RRs, niter=args.rim_iter * 20, lr=0.5, x_init=x_init, useprior=True) pred_adam, _ = datamodel.reconstruct(tf.constant(y_test), bias,
def main(_): mesh_shape = [("row", 2), ("col", 2)] layout_rules = [("nx_lr", "row"), ("ny_lr", "col"), ("nx", "row"), ("ny", "col"), ("ty_lr", "row"), ("tz_lr", "col"), ("nx_block", "row"), ("ny_block", "col")] mesh_hosts = ["localhost:%d" % (8222 + j) for j in range(4)] # Create a cluster from the mesh hosts. cluster = tf.train.ClusterSpec({ "mesh": mesh_hosts, "master": ["localhost:8488"] }) # Create a server for local mesh members server = tf.train.Server(cluster, job_name="master", task_index=0) mesh_devices = [ '/job:mesh/task:%d' % i for i in range(cluster.num_tasks("mesh")) ] print("List of devices", mesh_devices) mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl( mesh_shape, layout_rules, mesh_devices) # Build the model # Create computational graphs and some initializations graph = mtf.Graph() mesh = mtf.Mesh(graph, "nbody_mesh") # Compute a few things first, using simple tensorflow a0 = FLAGS.a0 a = FLAGS.af nsteps = FLAGS.nsteps bs, nc = FLAGS.box_size, FLAGS.nc klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) stages = np.linspace(a0, a, nsteps, endpoint=True) #pt = PerturbationGrowth(cosmology, a=[a], a_normalize=1.0) # Generate a batch of 3D initial conditions initial_conditions = flowpm.linear_field( FLAGS.nc, # size of the cube FLAGS.box_size, # Physical size of the cube ipklin, # Initial power spectrum batch_size=FLAGS.batch_size) state = lpt_init(initial_conditions, a0=a0, order=1) #final_state = state final_state = nbody(state, stages, nc) tfinal_field = cic_paint(tf.zeros_like(initial_conditions), final_state[0]) # Compute necessary Fourier kernels kvec = flowpm.kernels.fftk((nc, nc, nc), symmetric=False) from flowpm.kernels import laplace_kernel, gradient_kernel lap = tf.cast(laplace_kernel(kvec), tf.complex64) grad_x = gradient_kernel(kvec, 0) grad_y = gradient_kernel(kvec, 1) grad_z = gradient_kernel(kvec, 2) derivs = [lap, grad_x, grad_y, grad_z] mesh_final_field = lpt_prototype(mesh, initial_conditions, derivs, bs=FLAGS.box_size, nc=FLAGS.nc, batch_size=FLAGS.batch_size) # Lower mesh computation lowering = mtf.Lowering(graph, {mesh: mesh_impl}) # Retrieve output of computation result = lowering.export_to_tf_tensor(mesh_final_field) with tf.Session(server.target, config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=False)) as sess: a, b, c = sess.run([initial_conditions, tfinal_field, result]) np.save('init', a) np.save('reference_final', b) np.save('mesh_pyramid', c) plt.figure(figsize=(15, 3)) plt.subplot(141) plt.imshow(a[0].sum(axis=2)) plt.title('Initial Conditions') plt.subplot(142) plt.imshow(b[0].sum(axis=2)) plt.title('TensorFlow (single GPU)') plt.colorbar() plt.subplot(143) plt.imshow(c[0].sum(axis=2)) plt.title('Mesh TensorFlow Single') plt.colorbar() plt.subplot(144) plt.imshow((b[0] - c[0]).sum(axis=2)) plt.title('Residuals') plt.colorbar() plt.savefig("comparison-single.png") exit(0)
def test_convergence_Born(return_results=False): """ This function tests that given a set of density planes, both lenstools and flowpm recover the same convergence maps in angular coordinates. """ klin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) cosmo = flowpm.cosmology.Planck15() a0 = 0.9 # Create a state vector initial_conditions = flowpm.linear_field([nc, nc, 10 * nc], [bs, bs, 10 * bs], ipklin, batch_size=2) state = flowpm.lpt_init(cosmo, initial_conditions, a0) r = tf.linspace(0., 10 * bs, 11) r_center = 0.5 * (r[1:] + r[:-1]) a_center = flowpm.background.a_of_chi(cosmo, r_center) constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2 # To make it convenient to access simulation properties in lenstools # let's quicly export and reload the sim # TODO: remove the need for this! flowpm.io.save_state(cosmo, state, a0, [nc, nc, 10 * nc], [bs, bs, 10 * bs], 'snapshot_born_testing', attrs={'comoving_distance': r_center[0]}) # Reload the snapshot with lenstools snapshot = FlowPMSnapshot.open('snapshot_born_testing') # Get some density planes and create lenstool tracer lensplanes = [] tracer = lt.simulations.RayTracer(lens_type=lt.simulations.DensityPlane) for i in range(len(r_center)): plane = flowpm.raytracing.density_plane( state, [nc, nc, 10 * nc], r_center[i] / bs * nc, width=nc, plane_resolution=plane_resolution) r, a, p = r_center[i], a_center[i], plane[0] lensplanes.append((r, a, plane)) density_normalization = bs * r / a # We upsample the lensplanes before giving them to lenstools because # lentools is using a weird kind of interpolation when converting from # comoving coordinates to angular coords. with a larger p = tf.image.resize( tf.reshape(p, [1, plane_resolution, plane_resolution, 1]), [2048, 2048]) p = (p[0, :, :, 0] * constant_factor * density_normalization).numpy() p = p - np.mean(p) lt_plane = lt.simulations.DensityPlane( p, angle=snapshot.header["box_size"], redshift=1 / a - 1, cosmology=snapshot.cosmology) tracer.addLens(lt_plane) # Adding dummy lensplane at the end tracer.addLens( lt.simulations.DensityPlane(np.zeros((2048, 2048)), angle=snapshot.header["box_size"], redshift=0.99, cosmology=snapshot.cosmology)) tracer.addLens( lt.simulations.DensityPlane(np.zeros((2048, 2048)), angle=snapshot.header["box_size"], redshift=2, cosmology=snapshot.cosmology)) tracer.reorderLenses() # Create an array of coordinates at which to retrieve the convernge maps xgrid, ygrid = np.meshgrid( np.linspace(0, field, npix, endpoint=False), # range of X coordinates np.linspace(0, field, npix, endpoint=False)) # range of Y coordinates coords = np.stack([xgrid, ygrid], axis=0) * u.deg c = coords.reshape([2, -1]).T # Compute convergence map with lenstool lt_map = tracer.convergenceBorn(coords, z=1.0) # Compute convergemce map with flowpm fpm_map = flowpm.raytracing.convergenceBorn(cosmo, lensplanes, bs / nc, bs, c.to(u.rad), z_source=tf.ones([1]), field_npix=npix) # Comparing the final maps assert_allclose(lt_map, fpm_map[0].numpy().reshape([npix, npix, -1])[:, :, -1], atol=5e-4) if return_results: return lt_map, fpm_map
def main(): """ Model function for the CosmicRIM. """ if args.parallel: rim = build_rim_parallel(params) else: rim = build_rim_split(params) grad_fn = recon_dm_grad # # # train_dataset = tf.data.Dataset.range(args.batch_in_epoch) # train_dataset = train_dataset.map(pm_data) # # dset = dset.apply(tf.data.experimental.unbatch()) # train_dataset = train_dataset.prefetch(-1) # test_dataset = tf.data.Dataset.range(1).map(pm_data_test).prefetch(-1) # traindata, testdata = get_data() idx = np.random.randint(0, traindata.shape[0], 1) xx, yy = traindata[idx, 0].astype(np.float32), traindata[idx, 1].astype(np.float32), x_init = flowpm.linear_field(nc, bs, ipklin, batch_size=xx.shape[0]) #x_init = np.random.normal(size=xx.size).reshape(xx.shape).astype(np.float32) x_pred = rim(x_init, yy, grad_fn) # # @tf.function def rim_train(x_true, x_init, y): with tf.GradientTape() as tape: x_pred = rim(x_init, y, grad_fn) res = (x_true - x_pred) loss = tf.reduce_mean(tf.square(res)) gradients = tape.gradient(loss, rim.trainable_variables) return loss, gradients ##Train and save piter, testiter = 10, 50 losses = [] lrs = [0.0001, 0.0001] liters = [1001, 2001] trainiter = 0 start = time.time() x_test, y_test = None, None for il in range(len(lrs)): print('Learning rate = %0.3e' % lrs[il]) opt = tf.keras.optimizers.Adam(learning_rate=lrs[il]) for i in range(liters[il]): idx = np.random.randint(0, traindata.shape[0], args.batch_size) xx, yy = traindata[idx, 0].astype( np.float32), traindata[idx, 1].astype(np.float32), x_init = flowpm.linear_field(nc, bs, ipklin, batch_size=xx.shape[0]).numpy() #x_init = np.random.normal(size=xx.size).reshape(xx.shape).astype(np.float32) #x_init = (yy - (yy.max() - yy.min())/2.)/yy.std() + np.random.normal(size=xx.size).reshape(xx.shape).astype(np.float32) loss, gradients = rim_train(x_true=tf.constant(xx), x_init=tf.constant(x_init), y=tf.constant(yy)) losses.append(loss.numpy()) opt.apply_gradients(zip(gradients, rim.trainable_variables)) if i % piter == 0: print("Time taken for %d iterations : " % piter, time.time() - start) print("Loss at iteration %d : " % i, losses[-1]) start = time.time() if i % testiter == 0: plt.plot(losses) plt.savefig(ofolder + 'losses.png') plt.close() ##check 2pt and comapre to Adam #idx = np.random.randint(0, testdata.shape[0], 1) #xx, yy = testdata[idx, 0].astype(np.float32), testdata[idx, 1].astype(np.float32), if x_test is None: idx = np.random.randint(0, testdata.shape[0], 1) x_test, y_test = testdata[idx, 0].astype( np.float32), testdata[idx, 1].astype(np.float32), x_init = flowpm.linear_field( nc, bs, ipklin, batch_size=x_test.shape[0]).numpy() pred_adam = adam(tf.constant(x_init), tf.constant(y_test), grad_fn) pred_adam = [ pred_adam[0].numpy(), pm(pred_adam)[0].numpy() ] pred_adam10 = adam10(tf.constant(x_init), tf.constant(y_test), grad_fn) pred_adam10 = [ pred_adam10[0].numpy(), pm(pred_adam10)[0].numpy() ] minic, minfin = fid_recon.reconstruct(tf.constant(y_test), RRs=[1.0, 0.0], niter=args.rim_iter * 10, lr=0.1) compares = [pred_adam, pred_adam10, [minic[0], minfin[0]]] print('Test set generated') #x_init = np.random.normal(size=x_test.size).reshape(x_test.shape).astype(np.float32) x_init = flowpm.linear_field( nc, bs, ipklin, batch_size=x_test.shape[0]).numpy() #x_init = (y_test - (y_test.max() - y_test.min())/2.)/y_test.std() + np.random.normal(size=x_test.size).reshape(x_test.shape).astype(np.float32) pred = rim(tf.constant(x_init), tf.constant(y_test), grad_fn)[-1] check_im(x_test[0], x_init[0], pred.numpy()[0], fname=ofolder + 'rim-im-%04d.png' % trainiter) check_2pt(x_test, y_test, rim, grad_fn, compares, fname=ofolder + 'rim-2pt-%04d.png' % trainiter) rim.save_weights(ofolder + '/%d' % trainiter) trainiter += 1
def main(_): dtype=tf.float32 startw = time.time() tf.random.set_random_seed(100) np.random.seed(100) # Compute a few things first, using simple tensorflow a0=FLAGS.a0 a=FLAGS.af nsteps=FLAGS.nsteps bs, nc = FLAGS.box_size, FLAGS.nc klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) stages = np.linspace(a0, a, nsteps, endpoint=True) tf.reset_default_graph() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype) if FLAGS.nbody: state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) else: final_state = lpt_init(tfic, a0=stages[-1], order=1) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) with tf.Session() as sess: ic, fin = sess.run([tfic, tfinal_field]) np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) ################################################################ tf.reset_default_graph() print('ic constructed') noise = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype(np.float32) data_noised = fin + noise data = data_noised minimum = data.copy() start = noise.copy().flatten().astype(np.float32) Rsm = tf.placeholder(tf.float32, name='smoothing') def recon_prototype(linear): """ """ linear = tf.reshape(linear, minimum.shape) #loss = tf.reduce_sum(tf.square(linear - minimum)) state = lpt_init(linear, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) final_field = cic_paint(tf.zeros_like(linear), final_state[0]) residual = final_field - data.astype(np.float32) base = residual ## Rsmsq = tf.multiply(Rsm*bs/nc, Rsm*bs/nc) ## smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq)) ## basek = r2c3d(base, norm=nc**3) ## basek = tf.multiply(basek, tf.cast(smwts, tf.complex64)) ## base = c2r3d(basek, norm=nc**3) #### # chisq = tf.multiply(base, base) chisq = tf.reduce_sum(chisq) chisq = tf.multiply(chisq, 1/nc**3, name='chisq') #Prior lineark = r2c3d(linear, norm=nc**3) priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32)) prior = tf.reduce_sum(tf.multiply(priormesh, 1/priorwt)) prior = tf.multiply(prior, 1/nc**3, name='prior') # loss = chisq + prior return loss @tf.function def min_lbfgs(): return tfp.optimizer.lbfgs_minimize( make_val_and_grad_fn(recon_prototype), initial_position=tf.constant(start), tolerance=1e-5, max_iterations=50) with tf.Session() as sess: #results = sess.run(min_lbfgs(), {Rsm:4}) results = sess.run(min_lbfgs()) print(results) minimum = results.position print(minimum) tf.reset_default_graph() print('\nminimized\n') tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype)*0 + minimum.reshape(data_noised.shape) 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: minic, minfin = sess.run([tfic, tfinal_field]) dg.saveimfig(0, [minic, minfin], [ic, fin], fpath+'') dg.save2ptfig(0, [minic, minfin], [ic, fin], fpath+'', bs) ## exit(0)
def test_density_plane(return_results=False): """ Tests cutting density planes from snapshots against lenstools """ klin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt(data_path + '/flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) cosmo = flowpm.cosmology.Planck15() a0 = 0.9 r0 = flowpm.background.rad_comoving_distance(cosmo, a0) # Create a state vector initial_conditions = flowpm.linear_field(nc, bs, ipklin, batch_size=2) state = flowpm.lpt_init(cosmo, initial_conditions, a0) # Export the snapshot flowpm.io.save_state(cosmo, state, a0, [nc, nc, nc], [bs, bs, bs], 'snapshot_density_testing', attrs={'comoving_distance': r0}) # Reload the snapshot with lenstools snapshot = FlowPMSnapshot.open('snapshot_density_testing') # Cut a lensplane in the middle of the volume lt_plane, resolution, NumPart = snapshot.cutPlaneGaussianGrid( normal=2, plane_resolution=plane_resolution, center=(bs / 2) * snapshot.Mpc_over_h, thickness=(bs / 4) * snapshot.Mpc_over_h, left_corner=np.zeros(3) * snapshot.Mpc_over_h, smooth=None, kind='density') # Cut the same lensplane with flowpm fpm_plane = flowpm.raytracing.density_plane( state, nc, center=nc / 2, width=nc / 4, plane_resolution=plane_resolution) # Apply additional normalization terms to match lenstools definitions constant_factor = 3 / 2 * cosmo.Omega_m * (constants.H0 / constants.c)**2 density_normalization = bs / 4 * r0 / a0 fpm_plane = fpm_plane * density_normalization * constant_factor # Checking first the mean value, which accounts for any normalization # issues assert_allclose(np.mean(fpm_plane[0]), np.mean(lt_plane), rtol=1e-5) # To check pixelwise difference, we need to do some smoothing as lenstools and # flowpm use different painting kernels smooth_lt_plane = np.fft.ifft2(fourier_gaussian(np.fft.fft2(lt_plane), 3)).real smooth_fpm_plane = np.fft.ifft2( fourier_gaussian(np.fft.fft2(fpm_plane[0]), 3)).real assert_allclose(smooth_fpm_plane, smooth_lt_plane, rtol=2e-2) if return_results: return fpm_plane, lt_plane, smooth_fpm_plane, smooth_lt_plane
bs, nc = FLAGS.box_size, FLAGS.nc klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) stages = np.linspace(a0, a, nsteps, endpoint=True) #tf.reset_default_graph() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype) if FLAGS.nbody: state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) else: final_state = lpt_init(tfic, a0=stages[-1], order=1) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) with tf.Session() as sess: ic, fin = sess.run([tfic, tfinal_field]) np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) ################################################################ #tf.reset_default_graph()
def main(_): infield = True mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape) 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) #If initc, run normal flowpm to generate data tf.reset_default_graph() if infield: tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100) if FLAGS.nbody: state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) else: final_state = lpt_init(tfic, a0=stages[-1], order=1) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) start = time.time() with tf.Session(server.target) as sess: ic, fin = sess.run([tfic, tfinal_field]) print("\nTime taken for the vanilla flowpm thingy :\n ", time.time() - start) else: ic = None tf.reset_default_graph() print('ic constructed') graph = mtf.Graph() mesh = mtf.Mesh(graph, "my_mesh") initial_conditions, final_field, input_field = nbody_prototype( mesh, infield, nc=FLAGS.nc, batch_size=FLAGS.batch_size) # 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_final = lowering.export_to_tf_tensor(final_field) nc = FLAGS.nc with tf.Session(server.target) as sess: start = time.time() if infield: ic_check, fin_check = sess.run([tf_initc, tf_final], feed_dict={input_field: ic}) else: ic_check, fin_check = sess.run([tf_initc, tf_final]) ic, fin = ic_check, fin_check print('\n Time for the mesh run : %f \n' % (time.time() - start)) plt.figure(figsize=(15, 3)) plt.subplot(141) plt.imshow(ic_check[0].sum(axis=2)) plt.title('Initial Conditions') plt.subplot(142) plt.imshow(fin[0].sum(axis=2)) plt.title('TensorFlow (single GPU)') plt.colorbar() plt.subplot(143) plt.imshow(fin_check[0].sum(axis=2)) plt.title('Mesh TensorFlow') plt.colorbar() plt.subplot(144) plt.imshow((fin_check[0] - fin[0]).sum(axis=2)) plt.title('Residuals') plt.colorbar() plt.savefig("comparison_mesh.png") exit(0) ## exit(0)
def main(_): cosmology = flowpm.cosmology.Planck15() # Create a simple Planck15 cosmology without neutrinos, and makes sure sigma8 # is matched nbdykit_cosmo = Cosmology.from_astropy(Planck15.clone(m_nu=0 * u.eV)) nbdykit_cosmo = nbdykit_cosmo.match(sigma8=cosmology.sigma8.numpy()) # Compute the k vectora that will be needed in the PGD fit k, _ = flowpm.power_spectrum(tf.zeros([1] + [FLAGS.nc] * 3), boxsize=np.array([FLAGS.box_size] * 3), kmin=np.pi / FLAGS.box_size, dk=2 * np.pi / FLAGS.box_size) # Create some initial conditions klin = tf.constant(np.logspace(-4, 1, 512), dtype=tf.float32) pk = linear_matter_power(cosmology, klin) pk_fun = lambda x: tf.cast( tf.reshape( interpolate.interp_tf(tf.reshape(tf.cast(x, tf.float32), [-1]), klin, pk), x.shape), tf.complex64) initial_conditions = flowpm.linear_field( [FLAGS.nc, FLAGS.nc, FLAGS.nc], [FLAGS.box_size, FLAGS.box_size, FLAGS.box_size], pk_fun, batch_size=FLAGS.batch_size) initial_state = flowpm.lpt_init(cosmology, initial_conditions, FLAGS.a_init) stages = np.linspace(FLAGS.a_init, 1., FLAGS.nsteps, endpoint=True) print('Starting simulation') # Run the Nbody states = flowpm.nbody(cosmology, initial_state, stages, [FLAGS.nc, FLAGS.nc, FLAGS.nc], pm_nc_factor=FLAGS.B, return_intermediate_states=True) print('Simulation done') # Initialize PGD params alpha = tf.Variable([FLAGS.alpha0], dtype=tf.float32) scales = tf.Variable([FLAGS.kl0, FLAGS.ks0], dtype=tf.float32) optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.learning_rate) params = [] scale_factors = [] # We begin by fitting the last time step for j, (a, state) in enumerate(states[::-1]): # Let's compute the target power spectrum at that scale factor target_pk = HalofitPower(nbdykit_cosmo, 1. / a - 1.)(k).astype('float32') for i in range(FLAGS.niter if j == 0 else FLAGS.niter_refine): optimizer.minimize( partial(pgd_loss, alpha, scales, state, target_pk), [alpha] if (FLAGS.fix_scales and j > 0) else [alpha, scales]) if i % 10 == 0: loss, pk = pgd_loss(alpha, scales, state, target_pk, return_pk=True) if i == 0: pk0 = pk print("step %d, loss:" % i, loss) params.append(np.concatenate([alpha.numpy(), scales.numpy()])) scale_factors.append(a) print("Fitted params (alpha, kl, ks)", params[-1]) plt.loglog(k, target_pk, "k") plt.loglog(k, pk0, ':', label='starting') plt.loglog(k, pk, '--', label='after n steps') plt.grid(which='both') plt.savefig('PGD_fit_%0.2f.png' % a) plt.close() pickle.dump( { 'B': FLAGS.B, 'nsteps': FLAGS.nsteps, 'params': params, 'scale_factors': scale_factors, 'cosmology': cosmology.to_dict(), 'boxsize': FLAGS.box_size, 'nc': FLAGS.nc }, open(FLAGS.filename, "wb"))
def benchmark_model(mesh): """ Initializes a 3D volume with random noise, and execute a forward FFT """ # Setup parameters bs = FLAGS.box_size nc = FLAGS.cube_size batch_size = FLAGS.batch_size a0 = FLAGS.a0 a = 1.0 nsteps = FLAGS.pm_steps # Compute a few things first, using simple tensorflow klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) stages = np.linspace(a0, a, nsteps, endpoint=True) # Initialize the integration steps stages = np.linspace(FLAGS.a0, 1.0, FLAGS.pm_steps, endpoint=True) # Generate a batch of 3D initial conditions initial_conditions = flowpm.linear_field( nc, # size of the cube bs, # Physical size of the cube ipklin, # Initial power spectrum batch_size=batch_size) # Compute necessary Fourier kernels kvec = flowpm.kernels.fftk((nc, nc, nc), symmetric=False) from flowpm.kernels import laplace_kernel, gradient_kernel lap = tf.cast(laplace_kernel(kvec), tf.complex64) grad_x = gradient_kernel(kvec, 0) grad_y = gradient_kernel(kvec, 1) grad_z = gradient_kernel(kvec, 2) # Define the named dimensions # Parameters of the small scales decomposition n_block_x = 8 n_block_y = 4 n_block_z = 1 halo_size = 4 # Parameters of the large scales decomposition downsampling_factor = 2 lnc = nc // 2**downsampling_factor fx_dim = mtf.Dimension("nx", nc) fy_dim = mtf.Dimension("ny", nc) fz_dim = mtf.Dimension("nz", nc) tfx_dim = mtf.Dimension("tx", nc) tfy_dim = mtf.Dimension("ty", nc) tfz_dim = mtf.Dimension("tz", nc) # Dimensions of the low resolution grid tx_dim = mtf.Dimension("tx_lr", nc) ty_dim = mtf.Dimension("ty_lr", nc) tz_dim = mtf.Dimension("tz_lr", nc) nx_dim = mtf.Dimension('nx_block', n_block_x) ny_dim = mtf.Dimension('ny_block', n_block_y) nz_dim = mtf.Dimension('nz_block', n_block_z) sx_dim = mtf.Dimension('sx_block', nc // n_block_x) sy_dim = mtf.Dimension('sy_block', nc // n_block_y) sz_dim = mtf.Dimension('sz_block', nc // n_block_z) batch_dim = mtf.Dimension("batch", batch_size) pk_dim = mtf.Dimension("npk", len(plin)) pk = mtf.import_tf_tensor(mesh, plin.astype('float32'), shape=[pk_dim]) # Compute necessary Fourier kernels kvec = flowpm.kernels.fftk((nc, nc, nc), symmetric=False) kx = mtf.import_tf_tensor(mesh, kvec[0].squeeze().astype('float32'), shape=[tfx_dim]) ky = mtf.import_tf_tensor(mesh, kvec[1].squeeze().astype('float32'), shape=[tfy_dim]) kz = mtf.import_tf_tensor(mesh, kvec[2].squeeze().astype('float32'), shape=[tfz_dim]) kv = [ky, kz, kx] kvec_lr = flowpm.kernels.fftk([nc, nc, nc], symmetric=False) kx_lr = mtf.import_tf_tensor(mesh, kvec_lr[0].squeeze().astype('float32'), shape=[tx_dim]) ky_lr = mtf.import_tf_tensor(mesh, kvec_lr[1].squeeze().astype('float32'), shape=[ty_dim]) kz_lr = mtf.import_tf_tensor(mesh, kvec_lr[2].squeeze().astype('float32'), shape=[tz_dim]) kv_lr = [ky_lr, kz_lr, kx_lr] # kvec for high resolution blocks shape = [batch_dim, fx_dim, fy_dim, fz_dim] lr_shape = [batch_dim, fx_dim, fy_dim, fz_dim] hr_shape = [batch_dim, nx_dim, ny_dim, nz_dim, sx_dim, sy_dim, sz_dim] part_shape = [batch_dim, fx_dim, fy_dim, fz_dim] initc = mtfpm.linear_field(mesh, shape, bs, nc, pk, kv) state = mtfpm.lpt_init_single( initc, a0, kv_lr, halo_size, lr_shape, hr_shape, part_shape[1:], antialias=True, ) #state = mtfpm.lpt_init(low, high, 0.1, kv_lr, kv_hr, halo_size, hr_shape, lr_shape, # part_shape[1:], downsampling_factor=downsampling_factor, antialias=True,) # Here we can run our nbody final_state = state #mtfpm.nbody(state, stages, lr_shape, hr_shape, kv_lr, kv_hr, halo_size, downsampling_factor=downsampling_factor) # paint the field final_field = mtf.zeros(mesh, shape=hr_shape) for block_size_dim in hr_shape[-3:]: final_field = mtf.pad(final_field, [halo_size, halo_size], block_size_dim.name) final_field = mesh_utils.cic_paint(final_field, final_state[0], halo_size) # Halo exchange for blocks_dim, block_size_dim in zip(hr_shape[1:4], final_field.shape[-3:]): final_field = mpm.halo_reduce(final_field, blocks_dim, block_size_dim, halo_size) # Remove borders for block_size_dim in hr_shape[-3:]: final_field = mtf.slice(final_field, halo_size, block_size_dim.size, block_size_dim.name) #final_field = mtf.reshape(final_field, [batch_dim, fx_dim, fy_dim, fz_dim]) # Hack usisng custom reshape because mesh is pretty dumb final_field = mtf.slicewise(lambda x: x[:, 0, 0, 0], [final_field], output_dtype=tf.float32, output_shape=[batch_dim, fx_dim, fy_dim, fz_dim], name='my_dumb_reshape', splittable_dims=part_shape[:-1] + hr_shape[:4]) return mtf.reduce_sum(final_field)
def main(_): dtype = tf.float32 startw = time.time() tf.random.set_random_seed(100) np.random.seed(100) # Compute a few things first, using simple tensorflow a0 = FLAGS.a0 a = FLAGS.af nsteps = FLAGS.nsteps bs, nc = FLAGS.box_size, FLAGS.nc klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) stages = np.linspace(a0, a, nsteps, endpoint=True) tf.reset_default_graph() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype) if FLAGS.nbody: state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) else: final_state = lpt_init(tfic, a0=stages[-1], order=1) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) with tf.Session() as sess: ic, fin = sess.run([tfic, tfinal_field]) np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) tf.reset_default_graph() print('ic constructed') linear, final_field, update_ops, loss, chisq, prior, Rsm = recon_prototype( fin) #initial_conditions = recon_prototype(mesh, fin, nc=FLAGS.nc, batch_size=FLAGS.batch_size, dtype=dtype) # Lower mesh computation with tf.Session() as sess: #ic_check, fin_check = sess.run([tf_initc, tf_final]) #sess.run(tf_linear_op, feed_dict={input_field:ic}) #ic_check, fin_check = sess.run([linear, final_field]) #dg.saveimfig('-check', [ic_check, fin_check], [ic, fin], fpath) #dg.save2ptfig('-check', [ic_check, fin_check], [ic, fin], fpath, bs) #sess.run(tf_linear_op, feed_dict={input_field:np.random.normal(size=ic.size).reshape(ic.shape)}) sess.run(tf.global_variables_initializer()) ic0, fin0 = sess.run([linear, final_field]) dg.saveimfig('-init', [ic0, fin0], [ic, fin], fpath) start = time.time() titer = 20 niter = 201 iiter = 0 start0 = time.time() RRs = [4, 2, 1, 0.5, 0] 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] for iR, zlR in enumerate(zip(RRs, lrs)): RR, lR = zlR for ff in [fpath + '/figs-R%02d' % (10 * RR)]: try: os.makedirs(ff) except Exception as e: print(e) for i in range(niter): iiter += 1 sess.run(update_ops, {Rsm: RR}) print(sess.run([loss, chisq, prior], {Rsm: RR})) if (i % titer == 0): end = time.time() print('Iter : ', i) print('Time taken for %d iterations: ' % titer, end - start) start = end ## #ic1, fin1, cc, pp = sess.run([tf_initc, tf_final, tf_chisq, tf_prior], {R0:RR}) #ic1, fin1, cc, pp = sess.run([tf_initc, tf_final, tf_chisq, tf_prior], {R0:RR}) ic1, fin1 = sess.run([linear, final_field]) #print('Chisq and prior are : ', cc, pp) dg.saveimfig(i, [ic1, fin1], [ic, fin], fpath + '/figs-R%02d' % (10 * RR)) dg.save2ptfig(i, [ic1, fin1], [ic, fin], fpath + '/figs-R%02d' % (10 * RR), bs) dg.saveimfig(i * (iR + 1), [ic1, fin1], [ic, fin], fpath + '/figs') dg.save2ptfig(i * (iR + 1), [ic1, fin1], [ic, fin], fpath + '/figs', bs) ic1, fin1 = sess.run([linear, final_field]) print('Total time taken for %d iterations is : ' % iiter, time.time() - start0) dg.saveimfig(i, [ic1, fin1], [ic, fin], fpath) dg.save2ptfig(i, [ic1, fin1], [ic, fin], fpath, bs) np.save(fpath + 'ic_recon', ic1) np.save(fpath + 'final_recon', fin1) print('Total wallclock time is : ', time.time() - start0) ## exit(0)
def main(_): mesh_shape = [("row", FLAGS.nx), ("col", FLAGS.ny)] layout_rules = [("nx_lr", "row"), ("ny_lr", "col"), ("nx", "row"), ("ny", "col"), ("ty_lr", "row"), ("tz_lr", "col"), ("nx_block", "row"), ("ny_block", "col")] mesh_impl = HvdSimdMeshImpl(mtf.convert_to_shape(mesh_shape), mtf.convert_to_layout_rules(layout_rules)) # Build the model # Create computational graphs and some initializations graph = mtf.Graph() mesh = mtf.Mesh(graph, "nbody_mesh") # Compute a few things first, using simple tensorflow a0 = FLAGS.a0 a = FLAGS.af nsteps = FLAGS.nsteps bs, nc = FLAGS.box_size, FLAGS.nc klin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../flowpm/data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) stages = np.linspace(a0, a, nsteps, endpoint=True) #pt = PerturbationGrowth(cosmology, a=[a], a_normalize=1.0) # Generate a batch of 3D initial conditions initial_conditions = flowpm.linear_field( FLAGS.nc, # size of the cube FLAGS.box_size, # Physical size of the cube ipklin, # Initial power spectrum batch_size=FLAGS.batch_size) cosmo = flowpm.cosmology.Planck15() state = lpt_init(cosmo, initial_conditions, a, order=1) final_state = state #final_state = nbody(cosmo, state, stages, nc) tfinal_field = cic_paint(tf.zeros_like(initial_conditions), final_state[0]) # Compute necessary Fourier kernels kvec = flowpm.kernels.fftk((nc, nc, nc), symmetric=False) from flowpm.kernels import laplace_kernel, gradient_kernel lap = tf.cast(laplace_kernel(kvec), tf.complex64) grad_x = gradient_kernel(kvec, 0) grad_y = gradient_kernel(kvec, 1) grad_z = gradient_kernel(kvec, 2) derivs = [lap, grad_x, grad_y, grad_z] mesh_final_field = lpt_prototype(mesh, initial_conditions, derivs, bs=FLAGS.box_size, nc=FLAGS.nc, batch_size=FLAGS.batch_size) # Lower mesh computation lowering = mtf.Lowering(graph, {mesh: mesh_impl}) # Retrieve output of computation result = lowering.export_to_tf_tensor(mesh_final_field) with tf.Session() as sess: a, b, c = sess.run([initial_conditions, tfinal_field, result]) if comm.rank == 0: np.save('init', a) np.save('reference_final', b) np.save('mesh_pyramid', c) plt.figure(figsize=(15, 3)) plt.subplot(141) plt.imshow(a[0].sum(axis=2)) plt.title('Initial Conditions') plt.subplot(142) plt.imshow(b[0].sum(axis=2)) plt.title('TensorFlow (single GPU)') plt.colorbar() plt.subplot(143) plt.imshow(c[0].sum(axis=2)) plt.title('Mesh TensorFlow Single') plt.colorbar() plt.subplot(144) plt.imshow((b[0] - c[0]).sum(axis=2)) plt.title('Residuals') plt.colorbar() plt.savefig("comparison-single.png") exit(0)
def main(_): dtype = tf.float32 mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape) 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) tf.reset_default_graph() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype) if FLAGS.nbody: state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) else: final_state = lpt_init(tfic, a0=stages[-1], order=1) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) with tf.Session(server.target) as sess: ic, fin = sess.run([tfic, tfinal_field]) np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) tf.reset_default_graph() print('ic constructed') graph = mtf.Graph() mesh = mtf.Mesh(graph, "my_mesh") initial_conditions, final_field, loss, var_grads, update_op, linear_op, input_field, lr, R0 = recon_prototype( mesh, fin, 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_final = lowering.export_to_tf_tensor(final_field) tf_grads = lowering.export_to_tf_tensor(var_grads[0]) 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 ic_hrshape = ic.reshape([ FLAGS.batch_size, n_block_x, nc // n_block_x, n_block_y, nc // n_block_y, n_block_z, nc // n_block_z ]) ic_hrshape = np.transpose(ic_hrshape, [0, 1, 3, 5, 2, 4, 6]) with tf.Session(server.target) as sess: #ic_check, fin_check = sess.run([tf_initc, tf_final]) sess.run(tf_linear_op, feed_dict={input_field: ic_hrshape}) ic_check, fin_check = sess.run([tf_initc, tf_final]) dg.saveimfig('-check', [ic_check, fin_check], [ic, fin], fpath) dg.save2ptfig('-check', [ic_check, fin_check], [ic, fin], fpath, bs) sess.run(tf_linear_op, feed_dict={ input_field: np.random.normal(size=ic.size).reshape(ic_hrshape.shape) }) ic0, fin0 = sess.run([tf_initc, tf_final]) dg.saveimfig('-init', [ic0, fin0], [ic, fin], fpath) start = time.time() niter = 5 iiter = 0 start0 = time.time() RRs = [4, 2, 1, 0.5, 0] lrs = np.array([0.2, 0.15, 0.1, 0.1, 0.1]) #lrs = [0.1, 0.05, 0.01, 0.005, 0.001] for iR, zlR in enumerate(zip(RRs, lrs)): RR, lR = zlR #for ff in [fpath + '/figs-R%02d'%(10*RR)]: for ff in [fpath + '/figsiter']: try: os.makedirs(ff) except Exception as e: print(e) for i in range(301): if (i % niter == 0): end = time.time() print('Iter : ', i) print('Time taken for %d iterations: ' % niter, end - start) start = end ## ic1, fin1 = sess.run([tf_initc, tf_final]) #dg.saveimfig(i, [ic1, fin1], [ic, fin], fpath+'/figs-R%02d'%(10*RR)) #dg.save2ptfig(i, [ic1, fin1], [ic, fin], fpath+'/figs-R%02d'%(10*RR), bs) dg.saveimfig2x2(iiter, [ic1, fin1], [ic, fin], fpath + '/figsiter') # sess.run(tf_update_ops, {lr: lR, R0: RR}) iiter += 1 dg.saveimfig(i * (iR + 1), [ic1, fin1], [ic, fin], fpath + '/figs') dg.save2ptfig(i * (iR + 1), [ic1, fin1], [ic, fin], fpath + '/figs', bs) ic1, fin1 = sess.run([tf_initc, tf_final]) print('Total time taken for %d iterations is : ' % iiter, time.time() - start0) dg.saveimfig(i, [ic1, fin1], [ic, fin], fpath) dg.save2ptfig(i, [ic1, fin1], [ic, fin], fpath, bs) np.save(fpath + 'ic_recon', ic1) np.save(fpath + 'final_recon', fin1) print('Total wallclock time is : ', time.time() - start0) ## exit(0)
def main(_): dtype = tf.float32 mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape) startw = time.time() print(mesh_shape) ## ## ##Begin here klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('..//data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) tf.reset_default_graph() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype) if FLAGS.nbody: state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) else: final_state = lpt_init(tfic, a0=stages[-1], order=1) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) with tf.Session() as sess: ic, fin = sess.run([tfic, tfinal_field]) np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) print(ic.shape, fin.shape) ######################################################## print(ic.shape, fin.shape) recon_estimator = tf.estimator.Estimator(model_fn=model_fn, model_dir=fpath) def eval_input_fn(): features = {} features['data'] = fin features['R0'] = 0 features['x0'] = None features['lr'] = 0 return features, None # Train and evaluate model. RRs = [4., 2., 1., 0.5, 0.] niter = 200 iiter = 0 for R0 in RRs: print('\nFor iteration %d and R=%0.1f\n' % (iiter, R0)) def train_input_fn(): features = {} features['data'] = fin features['R0'] = R0 features['x0'] = np.random.normal(size=fin.size).reshape(fin.shape) features['lr'] = 0.01 return features, None for _ in range(1): recon_estimator.train(input_fn=train_input_fn, max_steps=iiter + niter) eval_results = recon_estimator.predict(input_fn=eval_input_fn, yield_single_examples=False) for i, pred in enumerate(eval_results): if i > 0: break iiter += niter # dg.saveimfig(iiter, [pred['ic'], pred['data']], [ic, fin], fpath + '/figs/') dg.save2ptfig(iiter, [pred['ic'], pred['data']], [ic, fin], fpath + '/figs/', bs) sys.exit(0)
def main(_): dtype = tf.float32 startw = time.time() tf.random.set_random_seed(100) np.random.seed(100) # Compute a few things first, using simple tensorflow a0 = FLAGS.a0 a = FLAGS.af nsteps = FLAGS.nsteps bs, nc = FLAGS.box_size, FLAGS.nc klin = np.loadtxt('../data/Planck15_a1p00.txt').T[0] plin = np.loadtxt('../data/Planck15_a1p00.txt').T[1] ipklin = iuspline(klin, plin) stages = np.linspace(a0, a, nsteps, endpoint=True) tf.reset_default_graph() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype) if FLAGS.nbody: state = lpt_init(tfic, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) else: final_state = lpt_init(tfic, a0=stages[-1], order=1) tfinal_field = cic_paint(tf.zeros_like(tfic), final_state[0]) with tf.Session() as sess: ic, fin = sess.run([tfic, tfinal_field]) np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) k, pic = tools.power(ic[0] + 1, boxsize=bs) k, pfin = tools.power(fin[0], boxsize=bs) plt.plot(k, pic) plt.plot(k, pfin) plt.loglog() plt.grid(which='both') plt.savefig('pklin.png') plt.close() print(pic) print(pfin) #sys.exit(-1) ################################################################ tf.reset_default_graph() print('ic constructed') noise = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype( np.float32) * 1 data_noised = fin + noise data = data_noised startpos = noise.copy().flatten().astype(np.float32) x0 = tf.placeholder(dtype=tf.float32, shape=data.flatten().shape, name='initlin') Rsm = tf.placeholder(tf.float32, name='smoothing') def recon_prototype(linearflat): """ """ linear = tf.reshape(linearflat, data.shape) # #loss = tf.reduce_sum(tf.square(linear - minimum)) state = lpt_init(linear, a0=0.1, order=1) final_state = nbody(state, stages, FLAGS.nc) final_field = cic_paint(tf.zeros_like(linear), final_state[0]) residual = final_field - data.astype(np.float32) base = residual Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc) smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq)) basek = r2c3d(base, norm=nc**3) basek = tf.multiply(basek, tf.cast(smwts, tf.complex64)) base = c2r3d(basek, norm=nc**3) # chisq = tf.multiply(base, base) chisq = tf.reduce_sum(chisq) chisq = tf.multiply(chisq, 1 / nc**3, name='chisq') #Prior lineark = r2c3d(linear, norm=nc**3) priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32)) prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt)) prior = tf.multiply(prior, 1 / nc**3, name='prior') # loss = chisq + prior grad = tf.gradients(loss, linearflat) print(grad) return loss, grad[0] @tf.function def min_lbfgs(): return tfp.optimizer.lbfgs_minimize( #make_val_and_grad_fn(recon_prototype), recon_prototype, initial_position=x0, tolerance=1e-10, max_iterations=100) with tf.Session() as sess: start = time.time() results = sess.run(min_lbfgs(), {Rsm: 2, x0: startpos}) print("\n") print(results) print("\n") minimum = results.position print(minimum) print("\nTime taken : ", time.time() - start) start = time.time() results = sess.run(min_lbfgs(), {Rsm: 1, x0: minimum}) print("\n") print(results) minimum = results.position print("\n") print(minimum) print("\nTime taken : ", time.time() - start) start = time.time() results = sess.run(min_lbfgs(), {Rsm: 0, x0: minimum}) print("\n") print(results) minimum = results.position print("\n") print(minimum) print("\nTime taken : ", time.time() - start) tf.reset_default_graph() print("\n") print('\nminimized\n') tfic = linear_field( FLAGS.nc, FLAGS.box_size, ipklin, batch_size=1, seed=100, dtype=dtype) * 0 + minimum.reshape(data_noised.shape) 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: minic, minfin = sess.run([tfic, tfinal_field]) dg.saveimfig(0, [minic, minfin], [ic, fin], fpath + '') dg.save2ptfig(0, [minic, minfin], [ic, fin], fpath + '', bs) np.save(fpath + 'recon0ic', minic) np.save(fpath + 'recon-final', minfin) ## exit(0)
def main(): startw = time.time() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(nc, bs, ipklin, batch_size=1, seed=100, dtype=dtype) tfinal_field = pm(tfic) ic, fin = tfic.numpy(), tfinal_field.numpy() np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) print('\ndata constructed\n') noise = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype(np.float32) data_noised = fin + noise data = data_noised @tf.function def recon_prototype(linear, Rsm): """ """ linear = tf.reshape(linear, data.shape) #loss = tf.reduce_sum(tf.square(linear - minimum)) final_field = pm(linear) residual = final_field - data.astype(np.float32) base = residual if anneal: print("\nAdd annealing section to graph\n") Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc) smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq)) basek = r2c3d(base, norm=nc**3) basek = tf.multiply(basek, tf.cast(smwts, tf.complex64)) base = c2r3d(basek, norm=nc**3) chisq = tf.multiply(base, base) chisq = tf.reduce_sum(chisq) chisq = tf.multiply(chisq, 1 / nc**3, name='chisq') #Prior lineark = r2c3d(linear, norm=nc**3) priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32)) prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt)) prior = tf.multiply(prior, 1 / nc**3, name='prior') # loss = chisq + prior return loss @tf.function def val_and_grad(x, Rsm): print("val and grad : ", x.shape) with tf.GradientTape() as tape: tape.watch(x) loss = recon_prototype(x, tf.constant(Rsm, dtype=tf.float32)) grad = tape.gradient(loss, x) return loss, grad @tf.function def val_and_grad(x, Rsm): print("val and grad : ", x.shape) with tf.GradientTape() as tape: tape.watch(x) loss = recon_prototype(x, Rsm) grad = tape.gradient(loss, x) return loss, grad @tf.function def grad(x, Rsm): with tf.GradientTape() as tape: tape.watch(x) loss = recon_prototype(x, Rsm) grad = tape.gradient(loss, x) return grad #Function for LBFSG def func(x, RR): return [ vv.numpy().astype(np.float64) for vv in val_and_grad(x=tf.constant(x, dtype=tf.float32), Rsm=tf.constant(RR, dtype=tf.float32)) ] # # Create an optimizer for Adam. opt = tf.keras.optimizers.Adam(learning_rate=lr) #Loop it Reconstruction ##Reconstruction x0 = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype(np.float32) linear = tf.Variable(name='linmesh', shape=(1, nc, nc, nc), dtype=tf.float32, initial_value=x0, trainable=True) for iR, RR in enumerate(RRs): if optimizer == 'lbfgs': results = sopt.minimize(fun=func, x0=x0, args=RR, jac=True, method='L-BFGS-B', tol=1e-10, options={ 'maxiter': niter, 'ftol': 1e-12, 'gtol': 1e-12, 'eps': 1e-12 }) #results = sopt.minimize(fun=func, x0=x0, args = RR, jac=True, method='L-BFGS-B', # options={'maxiter':niter}) print(results) minic = results.x.reshape(data.shape) elif optimizer == 'adam': for i in range(niter): grads = grad([linear], tf.constant(RR, dtype=tf.float32)) opt.apply_gradients(zip(grads, [linear])) minic = linear.numpy().reshape(data.shape) # print('\nminimized\n') minfin = pm(tf.constant(minic, dtype=tf.float32)).numpy() dg.saveimfig("-R%d" % RR, [minic, minfin], [ic, fin], fpath + '') dg.save2ptfig("-R%d" % RR, [minic, minfin], [ic, fin], fpath + '', bs) ### x0 = minic np.save(fpath + 'ic-%d' % iR, minic) np.save(fpath + 'final-%d' % iR, minfin) exit(0)
def main(): startw = time.time() # Run normal flowpm to generate data try: ic, fin = np.load(fpath + 'ic.npy'), np.load(fpath + 'final.npy') print('Data loaded') except Exception as e: print('Exception occured', e) tfic = linear_field(nc, bs, ipklin, batch_size=1, seed=100, dtype=dtype) tfinal_field = pm(tfic) ic, fin = tfic.numpy(), tfinal_field.numpy() np.save(fpath + 'ic', ic) np.save(fpath + 'final', fin) print('\ndata constructed\n') noise = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype(np.float32) data_noised = fin + noise data = data_noised @tf.function def recon_prototype(linear, Rsm=0): """ """ linear = tf.reshape(linear, data.shape) #loss = tf.reduce_sum(tf.square(linear - minimum)) final_field = pm(linear) residual = final_field - data.astype(np.float32) base = residual if anneal: print("\nAdd annealing section to graph\n") Rsmsq = tf.multiply(Rsm * bs / nc, Rsm * bs / nc) smwts = tf.exp(tf.multiply(-kmesh**2, Rsmsq)) basek = r2c3d(base, norm=nc**3) basek = tf.multiply(basek, tf.cast(smwts, tf.complex64)) base = c2r3d(basek, norm=nc**3) ### chisq = tf.multiply(base, base) chisq = tf.reduce_sum(chisq) chisq = tf.multiply(chisq, 1 / nc**3, name='chisq') #Prior lineark = r2c3d(linear, norm=nc**3) priormesh = tf.square(tf.cast(tf.abs(lineark), tf.float32)) prior = tf.reduce_sum(tf.multiply(priormesh, 1 / priorwt)) prior = tf.multiply(prior, 1 / nc**3, name='prior') # loss = chisq + prior return loss # @tf.function def val_and_grad(x, RR): with tf.GradientTape() as tape: tape.watch(x) loss = recon_prototype(x, RR) grad = tape.gradient(loss, x) return loss, grad @tf.function def min_lbfgs(x0, RR): return tfp.optimizer.lbfgs_minimize( lambda x: val_and_grad(x, tf.constant(RR, dtype=tf.float32)), initial_position=x0, tolerance=1e-10, max_iterations=200) # # def make_val_and_grad_fn(value_fn, R): # @functools.wraps(value_fn) # def val_and_grad(x): # return tfp.math.value_and_gradient(value_fn, x) # return val_and_grad # # @tf.function # def min_lbfgs(x0, RR): # return tfp.optimizer.lbfgs_minimize( # make_val_and_grad_fn(recon_prototype), # recon_prototype, # initial_position=x0, # tolerance=1e-10, # max_iterations=100) # ##Reconstruction x0 = np.random.normal(0, 1, nc**3).reshape(fin.shape).astype( np.float32).flatten() RRs = [2, 1, 0.5, 0] for iR, RR in enumerate(RRs): results = min_lbfgs(x0, RR) #results = sopt.minimize(fun=func, x0=x0, args = RR, jac=True, method='L-BFGS-B', # options={'maxiter':200, 'ftol': 2.220446049250313e-09, 'gtol': 1e-10,}) print(results) ### minic = results.position.numpy().reshape(data.shape) print(minic.shape) print('\nminimized\n') minfin = pm(tf.constant(minic, dtype=tf.float32)).numpy() dg.saveimfig("-R%d" % RR, [minic, minfin], [ic, fin], fpath + '') dg.save2ptfig("-R%d" % RR, [minic, minfin], [ic, fin], fpath + '', bs)