def generate_pvi_vs_vi(mdp, init): print('\nRunning PVI vs VI') lr = 0.01 # pvi core_init = random_parameterised_matrix(2, 2, 32, 4) core_init = approximate(init, core_init) params = utils.solve(parameterised_value_iteration(mdp, lr/len(core_init)), core_init) vs = np.vstack([np.max(build(c), axis=1) for c in params]) m = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='r', label='pvi') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(m-2), cmap='autumn', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='r', marker='x') # vi qs = utils.solve(value_iteration(mdp, lr), init) vs = np.vstack([np.max(q, axis=1) for q in qs]) n = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='g', label='vi') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(n-2), cmap='spring', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='g', marker='x') plt.title('VI: {}, PVI {}'.format(n, m)) plt.legend() # plt.colorbar() plt.savefig('figs/vi-vs-pvi.png', dpi=300) plt.close()
def generate_mpvi_vs_mvi(mdp, init): print('\nRunning MPVI vs MVI') lr = 1e-2 # mpvi core_init = random_parameterised_matrix(2, 2, 32, 8) core_init = approximate(init, core_init) c_init = (core_init, [np.zeros_like(c) for c in core_init]) params = utils.solve(momentum_bundler(parameterised_value_iteration(mdp, lr), 0.9), c_init) vs = np.vstack([np.max(build(c[0]), axis=-1) for c in params]) m = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='r', label='mpvi') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(m-2), cmap='autumn', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='r', marker='x') # mvi init = (init, np.zeros_like(init)) qs = utils.solve(momentum_bundler(value_iteration(mdp, lr), 0.9), init) vs = np.vstack([np.max(q[0], axis=-1) for q in qs]) n = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='g', label='mvi') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(n-2), cmap='spring', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='g', marker='x') plt.title('MVI: {}, MPVI {}'.format(n, m)) plt.legend() plt.savefig('figs/mpvi-vs-pvi.png') plt.close()
def act(self, state): while self.T.num_leaves() < self.target_num_leaves: state_node = self.T.softmax_sample_leaves() state = state_node.state noise_state = state_node.noise_state noise_state, action = dampedSpringNoiseStep(state.noise_state) state_est = state + self.dynamics.predict(np.vstack(state, action)) reward_est = self.reward.predict(np.vstack(state_est, action)) value_est = self.value.predict(state_est) new_state_node = nodeData( action_inbound=action, state=state_est, noise_state=noise_state, value_est=value_est, step_reward=reward_est, ) self.T.add_state( parent=state_node, child=new_state_node, ) target_node = self.T.softmax_sample_leaf_states() next_node = self.T.step_toward_and_reroot(target_node) action = next_node.action
def test_multiply(): x = np.linspace(-2.5, 15, 5)[:, np.newaxis].astype(np.float32) y = randn(x.size)[:, np.newaxis].astype(np.float32) gk_x = GaussianKernel(0.1) x_e1 = FiniteVec.construct_RKHS_Elem(gk_x, x) x_e2 = FiniteVec.construct_RKHS_Elem(gk_x, y) x_fv = FiniteVec(gk_x, np.vstack([x, y]), prefactors=np.hstack([x_e1.prefactors] * 2), points_per_split=x.size) oper_feat_vec = FiniteVec(gk_x, x) oper = FiniteOp(oper_feat_vec, oper_feat_vec, np.eye(len(x))) res_e1 = multiply(oper, x_e1) res_e2 = multiply(oper, x_e2) res_v = multiply(oper, x_fv) assert np.allclose( res_e1.prefactors, (oper.matr @ oper.inp_feat.inner(x_e1) ).flatten()), "Application of operator to RKHS element failed." assert np.allclose( res_v.inspace_points, np.vstack([res_e1.inspace_points, res_e2.inspace_points]) ), "Application of operator to all vectors in RKHS vector failed at inspace points." assert np.allclose( res_v.prefactors, np.hstack([ res_e1.prefactors, res_e2.prefactors ])), "Application of operator to all vectors in RKHS vector failed." assert np.allclose( multiply(oper, oper).matr, oper.inp_feat.inner( oper.outp_feat)), "Application of operator to operator failed."
def make_n_step_data_from_DDQN(S, A, C, n, stdizer, params, key=randKey()): S_0_th = angle_normalize(S[0, :-n]) S_0_thdot = S[1, 0:-n] S_n_th = angle_normalize(S[0, n:]) S_n_thdot = S[1, n:] A_0 = A[:-n] A_n = A[n:] stdizer.observe_reward_vec(C) R = jnp.vstack( [stdizer.standardize_reward(C[m:-(n - m)]) for m in range(n)]) episode = jnp.vstack([ jnp.cos(S_0_th), jnp.sin(S_0_th), S_0_thdot, A_0, jnp.cos(S_n_th), jnp.sin(S_n_th), S_n_thdot, A_n, R, ]) return episode
def generate_nested_circles(key, n_samples, inner_radius=2, outer_radius=4, noise=0.15): k1, k2, k3, k4 = random.split(key, 4) # Generate the circles inner_t = random.uniform(k1, shape=(n_samples // 2, )) * 2 * jnp.pi inner_circle = inner_radius * jnp.vstack( [jnp.cos(inner_t), jnp.sin(inner_t)]) outer_t = random.uniform(k2, shape=(n_samples // 2, )) * 2 * jnp.pi outer_circle = outer_radius * jnp.vstack( [jnp.cos(outer_t), jnp.sin(outer_t)]) data = jnp.vstack([inner_circle.T, outer_circle.T]) # Keep track of the labels y = jnp.hstack([jnp.zeros(n_samples // 2), jnp.ones(n_samples // 2)]) # Shuffle the data idx = jnp.arange(n_samples) idx = random.permutation(k3, idx) data = data[idx] y = y[idx] data += random.normal(k4, data.shape) * noise return data, y
def bo(key, x_obs, y_obs, obj, params_init, params_bounds, search_space, num_bo_rounds, acquisition_fn=thompson_sampling, batch_size=1, num_points=2000, num_steps=1000, method='tfp'): """Run seq-batch Bayesian Optimization (BO).""" gp_util = gp.GPUtils() additional_info_dict = {} for i in range(num_bo_rounds//batch_size): key_loop = jax.random.fold_in(key, i) gaussian_process = gp_util.fit_gp( x_obs, y_obs, params_init, params_bounds, steps=num_steps) x_new = acquisition_fn(key_loop, x_obs, y_obs, gaussian_process, gp_util, search_space, batch_size=batch_size, num_points=num_points, method=method) y_new, additional_info = obj(x_new) additional_info_dict[i] = additional_info x_obs = jnp.vstack((x_obs, x_new)) y_obs = jnp.vstack((y_obs, y_new)) return x_obs, y_obs, additional_info_dict
def _logG(Xp, X, parameter_dict, t): xp, x = Xp['x'], X['x'] #compute importance_weight (there is 1 more index in the IW parameters than in everything else...) mu_t, mu_tm1 = self.IW_parameters['mu'][ t + 1], self.IW_parameters['mu'][t] cov_force_t, cov_force_tm1 = jnp.exp( self.IW_parameters['lcov_force'][t + 1]), jnp.exp( self.IW_parameters['lcov_force'][t]) IWs = -self._IW_energy_fn(x, jnp.vstack( (mu_t, cov_force_t))) + self._IW_energy_fn( xp, jnp.vstack((mu_tm1, cov_force_tm1))) #compute M_kernel_logp: we don't need to recompute this k_t_logps = normalized_Gaussian_logp(x, X['forward_mu'], X['forward_cov']) #compute l_kernel_logp Lparams = parameter_dict['L_parameters'] potential_params = jnp.vstack( (Lparams['mu'][t], jnp.exp(Lparams['lcov_force'][t]))) l_mu, l_cov = EL_mu_sigma(x, self._kernel_energy_fn, parameter_dict['ldt'], potential_params) l_tm1_logps = normalized_Gaussian_logp(xp, l_mu, l_cov) #finalize and return lws = IWs + l_tm1_logps - k_t_logps return lws
def flux_through_all_loops(r_surf, ll, dl, I_arr): """ returns the magnetic flux through all poloidal and toroidal loops of the surface array. Arguments: *r_surf*: (n, m, 3) array of cartesian coordinates on the surface. [:,0,:] must be the first poloidal loop, and [0,:,:] the first toroidal loop. *dl*: ( n_coils, nsegments, 3)-array along the coil segments. other coil line segment *l* ( n_coils, nsegments, 3)-array of the position of each coil segment returns: *polint*: n-array of poloidal fluxes *torint*: m-array of toroidal fluxes """ mapped_vector_potential = vmap( vmap(vector_potential, (0, None, None, None), 0), (1, None, None, None), 1) vec_surf = mapped_vector_potential(r_surf, ll, dl, I_arr) dl_pol = np.hstack((r_surf[:, 1:, :] - r_surf[:, :-1, :], (r_surf[:, 0, :] - r_surf[:, -1, :])[:, None, :])) A_midpol = 0.5 * np.hstack( (vec_surf[:, 1:, :] + vec_surf[:, :-1, :], (vec_surf[:, 0, :] + vec_surf[:, -1, :])[:, None, :])) # sum over axis 0 and 2 only, the product of the two. Calculates the inner product, and then "integrates" poloidally. polint = np.einsum('ijk, ijk->i', dl_pol, A_midpol) dl_tor = np.vstack((r_surf[1:, :, :] - r_surf[:-1, :, :], (r_surf[0, :, :] - r_surf[-1, :, :])[None, :, :])) A_midtor = 0.5 * np.vstack( (vec_surf[1:, :, :] + vec_surf[:-1, :, :], (vec_surf[0, :, :] + vec_surf[-1, :, :])[None, :, :])) torint = np.einsum('ijk, ijk->j', dl_tor, A_midtor) return polint, torint
def get_labels_and_idxs_bbox(self, outputs, targets, indices, num_boxes): """ Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss. Targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]. The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size. """ bbx_idx = self._get_src_permutation_idx(indices) B, N, _ = outputs['pred_boxes'].shape target_boxes = jnp.concatenate( [jnp.zeros((B, N, 2)), jnp.ones((B, N, 2))], -1) # [[0,0,1,1]] box_loss_mask = jnp.zeros(outputs['pred_boxes'].shape[:2]) target_boxes_o = jnp.concatenate( [t["boxes"][np.array(i)] for t, (_, i) in zip(targets, indices)], axis=0) box_loss_indices = jnp.ones(len(target_boxes_o)) target_boxes = self.scatter_nd(target_boxes, jnp.vstack(bbx_idx).astype(jnp.int32).T, target_boxes_o) box_loss_mask = self.scatter_nd( box_loss_mask, jnp.vstack(bbx_idx).astype(jnp.int32).T, box_loss_indices) return target_boxes, box_loss_mask
def generate_PG_vs_VI(mdp, init): print('\nRunning PG vs VI') lr = 0.001 # PG logits = utils.solve(policy_gradient_iteration_logits(mdp, lr), init) vs = np.vstack([value_functional(mdp.P, mdp.r, softmax(logit), mdp.discount).T for logit in logits]) n = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='g', label='PG') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(n-2), cmap='spring', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='g', marker='x') # VI v = value_functional(mdp.P, mdp.r, softmax(init), mdp.discount) init = np.einsum('ijk,jl->jk', mdp.P, v) # V->Q qs = utils.solve(value_iteration(mdp, lr), init) vs = np.vstack([np.max(q, axis=1) for q in qs]) m = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='r', label='VI') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(m-2), cmap='autumn', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='r', marker='x') plt.legend() plt.title('PG: {}, VI {}'.format(n, m)) # plt.colorbar() plt.savefig('figs/pg-vs-vi.png') plt.close()
def dynamics_fwd(cls, s, v_m): v_m = np.clip(v_m, -cls.v_mx, cls.v_mx) x, theta, x_dot, theta_dot = s # Compute force acting on the cart: F = ((cls.eta_g * cls.Kg * cls.eta_m * cls.Kt) / (cls.Rm * cls.r_mp) * (-cls.Kg * cls.Km * x_dot / cls.r_mp + cls.eta_m * v_m)) # Compute acceleration: A11 = cls.mp + cls.Jeq A12 = cls.mp * cls.pl * np.cos(theta) A21 = cls.mp * cls.pl * np.cos(theta) A22 = cls.Jp + cls.mp * cls.pl**2 b11 = F - cls.Beq * x_dot - cls.mp * cls.pl * np.sin( theta) * theta_dot**2 b21 = 0. - cls.Bp * theta_dot - cls.mp * cls.pl * cls.g * np.sin(theta) A = np.vstack((np.hstack((A11, A12)), np.hstack((A21, A22)))) b = np.vstack((b11, b21)) Ainv = np.linalg.inv(A) s_ddot = np.dot(Ainv, b).squeeze() s_vel = np.hstack((x_dot, theta_dot)) + s_ddot * cls.dt s_pos = np.hstack((x, theta)) + s_vel * cls.dt s_next = np.hstack((s_pos, s_vel)) return s_next
def generate_PG_vs_PPG(mdp, init): print('\nRunning PG vs PPG') lr = 0.1 # PPG core_init = random_parameterised_matrix(2, 2, 32, 8) core_init = approximate(init, core_init) all_params = utils.solve(parameterised_policy_gradient_iteration(mdp, lr/len(core_init)), core_init) vs = np.vstack([np.max(value_functional(mdp.P, mdp.r, softmax(build(params), axis=-1), mdp.discount), axis=1) for params in all_params]) m = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='g', label='PPG') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(m-2), cmap='spring', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='g', marker='x') # PG logits = utils.solve(policy_gradient_iteration_logits(mdp, lr), init) vs = np.vstack([np.max(value_functional(mdp.P, mdp.r, softmax(logit, axis=-1), mdp.discount), axis=1) for logit in logits]) n = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='r', label='PG') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(n-2), cmap='autumn', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='r', marker='x') plt.title('PG: {}, PPG {}'.format(n, m)) plt.legend() # plt.colorbar() plt.savefig('figs/pg-vs-ppg.png', dpi=300) plt.close()
def generate_vi_sgd_vs_mom(mdp, init, lr=0.01): print('\nRunning VI SGD vs Mom') # sgd qs = utils.solve(value_iteration(mdp, lr), init) qs = clipped_stack(qs,1000) vs = np.vstack([np.max(q, axis=1) for q in qs]) n = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='m', label='gd') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(n-2), cmap='spring', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='m', marker='x') # momentum init = (init, np.zeros_like(init)) qs = utils.solve(momentum_bundler(value_iteration(mdp, lr), 0.99), init) qs = clipped_stack(qs,1000) vs = np.vstack([np.max(q[0], axis=1) for q in qs]) m = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='r', label='momentum') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(m-2), cmap='autumn', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='r', marker='x') plt.title('SGD: {}, Mom {}, Lr: {}'.format(n, m, lr)) plt.legend() plt.savefig('figs/vi_sgd-vs-vi_mom_{}.png'.format(lr)) plt.close()
def generate_sample_grid(theta_mean, theta_std, n): """ Create a meshgrid of n ** n_dim samples, tiling [theta_mean[i] - 5 * theta_std[i], theta_mean[i] + 5 * theta_std] into n portions. Also returns the volume element. Parameters ---------- theta_mean, theta_std : ndarray (n_dim) Returns ------- theta_samples : ndarray (nobj, n_dim) vol_element: scalar Volume element """ n_components = theta_mean.size xs = [ np.linspace( theta_mean[i] - 5 * theta_std[i], theta_mean[i] + 5 * theta_std[i], n, ) for i in range(n_components) ] mxs = np.meshgrid(*xs) orshape = mxs[0].shape mxsf = np.vstack([i.ravel() for i in mxs]).T dxs = np.vstack([np.diff(xs[i])[i] for i in range(n_components)]) vol_element = np.prod(dxs) theta_samples = np.vstack(mxsf) return theta_samples, vol_element
def generate_mppg_vs_mpg(mdp, init): print('\nRunning MPG vs MPPG') lr = 0.001 # MPPG core_init = random_parameterised_matrix(2, 2, 32, 8) core_init = approximate(init, core_init) core_init = (core_init, [np.zeros_like(c) for c in core_init]) all_params = utils.solve(momentum_bundler(parameterised_policy_gradient_iteration(mdp, lr), 0.9), core_init) vs = np.vstack([np.max(value_functional(mdp.P, mdp.r, softmax(build(params), axis=-1), mdp.discount), axis=1) for params, mom in all_params]) m = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='g', label='MPPG') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(m-2), cmap='spring', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='g', marker='x') # MPG init = (init, np.zeros_like(init)) logits = utils.solve(momentum_bundler(policy_gradient_iteration_logits(mdp, lr), 0.9), init) vs = np.vstack([np.max(value_functional(mdp.P, mdp.r, softmax(logit, axis=-1), mdp.discount), axis=1) for logit, mom in logits]) n = vs.shape[0] plt.scatter(vs[0, 0], vs[0, 1], c='r', label='MPG') plt.scatter(vs[1:-1, 0], vs[1:-1, 1], c=range(n-2), cmap='autumn', s=10) plt.scatter(vs[-1, 0], vs[-1, 1], c='r', marker='x') plt.title('PG: {}, PPG {}'.format(n, m)) plt.legend() # plt.colorbar() plt.savefig('figs/mpg-vs-mppg.png', dpi=300) plt.close()
def fill_cov(S, dim): m, m = S.shape S_eye = jnp.identity(dim - m) S_fill = jnp.zeros((m, dim - m)) S_fill_left = jnp.vstack((S_fill, S_eye)) S_final = jnp.vstack((S, S_fill.T)) S_final = jnp.hstack((S_final, S_fill_left)) return S_final
def hits(adjMatrix, p: int = 100): """ Calculate the hub and authority score in a net Parameters ------------------- adjMatrix: a numpy array NxN. p: int. Default 100. The max iteration. Returns ------------------- hub: dict. The hub score for each node in the net. authority: dict. The authority score for each node in the net. h: numpy array Nxp. (Optional). The hub score for each node in the net for each algorithm's step. a: numpy array Nxp. (Optional). The authority score for each node in the net for each algorithm's step. Example ------------------- >> import numpy as np >> from influencer.lazy_centrality import hits Create an adjiacency matrix with numpy >> adjM = np.random.rand(10, 10) >> adjM[adjM>0.5]=1 >> adjM[adjM<=0.5]=0 >> hub, aut, _, _ = hits(adjMatrix = adjM) """ n = adjMatrix.shape[0] a = np.ones([1,n]) h = np.ones([1,n]) pa=a #ph=h authority = {} hub = {} for k in range(1,p): h1 = np.dot(adjMatrix, pa.T)/np.linalg.norm(np.dot(adjMatrix, pa.T)) a1 = np.dot(adjMatrix.T, h1)/np.linalg.norm(np.dot(adjMatrix.T , h1)) h = np.vstack((h,np.dot(adjMatrix, a[k-1,:].T)/np.linalg.norm(np.dot(adjMatrix, a[k-1,:].T)))) a = np.vstack((a,np.dot(adjMatrix.T, h[k,:].T)/np.linalg.norm(np.dot(adjMatrix.T, h[k,:].T)))) pa = a1.T #ph = h1.T for i in range(n): authority[str(i)] = a[-1,i] hub[str(i)] = h[-1,i] return hub, authority, h, a
def distinguish_samples(a_samples, b_samples, classifier_params, logit_d): a_labels = np.zeros((a_samples.shape[0], 1)) b_labels = np.ones((b_samples.shape[0], 1)) samples = np.vstack([a_samples, b_samples]) true_labels = np.vstack([a_labels, b_labels]) parallel_logit_d = jax.vmap(logit_d, in_axes=(0, None, None)) pred_labels = parallel_logit_d(classifier_params, samples, samples) pred_labels = jax.nn.sigmoid(pred_labels).mean(0) return true_labels, pred_labels
def generate_dataset(n_b, n_c, n_d): x_b = domain * 1.0 u_b = u_fn(x_b) x_c = jnp.linspace(*domain[:, 0], n_c).reshape((-1, 1)) x_d = jnp.linspace(*domain[:, 0], n_d).reshape((-1, 1)) u_d = u_fn(x_d) dirichlet = dataset_Dirichlet(jnp.vstack([x_b, x_d]), jnp.vstack([u_b, u_d])) collocation = dataset_Collocation(jnp.vstack([x_b, x_c])) return dirichlet, collocation
def create_spatiotemporal_grid(X, Y): """ create a grid of data sized [T, R1, R2] note that this function removes full duplicates (i.e. where all dimensions match) TODO: generalise to >5D """ if Y.ndim < 2: Y = Y[:, None] num_spatial_dims = X.shape[1] - 1 if num_spatial_dims == 4: sort_ind = nnp.lexsort( (X[:, 4], X[:, 3], X[:, 2], X[:, 1], X[:, 0])) # sort by 0, 1, 2, 4 elif num_spatial_dims == 3: sort_ind = nnp.lexsort( (X[:, 3], X[:, 2], X[:, 1], X[:, 0])) # sort by 0, 1, 2, 3 elif num_spatial_dims == 2: sort_ind = nnp.lexsort((X[:, 2], X[:, 1], X[:, 0])) # sort by 0, 1, 2 elif num_spatial_dims == 1: sort_ind = nnp.lexsort((X[:, 1], X[:, 0])) # sort by 0, 1 else: raise NotImplementedError X = X[sort_ind] Y = Y[sort_ind] unique_time = np.unique(X[:, 0]) unique_space = nnp.unique(X[:, 1:], axis=0) N_t = unique_time.shape[0] N_r = unique_space.shape[0] if num_spatial_dims == 4: R = np.tile(unique_space, [N_t, 1, 1, 1, 1]) elif num_spatial_dims == 3: R = np.tile(unique_space, [N_t, 1, 1, 1]) elif num_spatial_dims == 2: R = np.tile(unique_space, [N_t, 1, 1]) elif num_spatial_dims == 1: R = np.tile(unique_space, [N_t, 1]) else: raise NotImplementedError R_flat = R.reshape(-1, num_spatial_dims) Y_dummy = np.nan * np.zeros([N_t * N_r, 1]) time_duplicate = np.tile(unique_time, [N_r, 1]).T.flatten() X_dummy = np.block([time_duplicate[:, None], R_flat]) X_all = np.vstack([X, X_dummy]) Y_all = np.vstack([Y, Y_dummy]) X_unique, ind = nnp.unique(X_all, axis=0, return_index=True) Y_unique = Y_all[ind] grid_shape = (unique_time.shape[0], ) + unique_space.shape R_grid = X_unique[:, 1:].reshape(grid_shape) Y_grid = Y_unique.reshape(grid_shape[:-1] + (1, )) return unique_time[:, None], R_grid, Y_grid
def realfftbasis(nx): """ Basis of sines+cosines for nn-point discrete fourier transform (DFT). Ported from MatLab code: https://github.com/leaduncker/SimpleEvidenceOpt/blob/master/util/realfftbasis.m """ import numpy as np nn = nx ncos = np.ceil((nn + 1) / 2) nsin = np.floor((nn - 1) / 2) wvec = np.hstack( [np.arange(start=0., stop=ncos), np.arange(start=-nsin, stop=0.)]) wcos = wvec[wvec >= 0] wsin = wvec[wvec < 0] x = np.arange(nx) t0 = np.cos(np.outer(wcos * 2 * np.pi / nn, x)) t1 = np.sin(np.outer(wsin * 2 * np.pi / nn, x)) B = np.vstack([t0, t1]) / np.sqrt(nn * 0.5) return B, wvec
def conway_graph(size) -> jraph.GraphsTuple: """Returns a graph representing the game field of conway's game of life.""" # Creates nodes: each node represents a cell in the game. n_node = size**2 nodes = np.zeros((n_node, 1)) node_indices = jnp.arange(n_node) # Creates edges, senders and receivers: # the senders represent the connections to the 8 neighboring fields. n_edge = 8 * n_node edges = jnp.zeros((n_edge, 1)) senders = jnp.vstack([ node_indices - size - 1, node_indices - size, node_indices - size + 1, node_indices - 1, node_indices + 1, node_indices + size - 1, node_indices + size, node_indices + size + 1 ]) senders = senders.T.reshape(-1) senders = (senders + size**2) % size**2 receivers = jnp.repeat(node_indices, 8) # Adds a glider to the game nodes[0, 0] = 1.0 nodes[1, 0] = 1.0 nodes[2, 0] = 1.0 nodes[2 + size, 0] = 1.0 nodes[1 + 2 * size, 0] = 1.0 return jraph.GraphsTuple(n_node=jnp.array([n_node]), n_edge=jnp.array([n_edge]), nodes=jnp.asarray(nodes), edges=edges, globals=None, senders=senders, receivers=receivers)
def run_mlstm1900_example(): # Set up an example sequence = "MRKGEELFTGVVPILVELDGDVNGHKFSVRGEGEGDATNGKLTLKFICTTGKLPVPWPTLVTTLTYGVQCFARYPDHMKQHDFFKSAMPEGYVQERTISFKDDGTYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNFNSHNVYITADKQKNGIKANFKIRHNVEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSVLSKDPNEKRDHMVLLEFVTAAGITHGMDELYK" print("sequence length: ", len(sequence)) sequence = aa_seq_to_int(sequence)[:-1] embeddings = np.load("embed_matrix:0.npy") x = np.vstack([embeddings[i] for i in sequence]) print("embedding shape: ", x.shape) # x = sliding_window(sequence, size=10) params = dict() params["gh"] = np.load("rnn_mlstm_mlstm_gh:0.npy") params["gmh"] = np.load("rnn_mlstm_mlstm_gmh:0.npy") params["gmx"] = np.load("rnn_mlstm_mlstm_gmx:0.npy") params["gx"] = np.load("rnn_mlstm_mlstm_gx:0.npy") params["wh"] = np.load("rnn_mlstm_mlstm_wh:0.npy") params["wmh"] = np.load("rnn_mlstm_mlstm_wmh:0.npy") params["wmx"] = np.load("rnn_mlstm_mlstm_wmx:0.npy") params["wx"] = np.load("rnn_mlstm_mlstm_wx:0.npy") params["b"] = np.load("rnn_mlstm_mlstm_b:0.npy") # Pass through mLSTM1900 out = mlstm1900(params, x) print("output: ", out) print("reps: ", out.mean(axis=0)) print("output shape: ", out.shape) assert out.shape == (x.shape[0], 1900)
def explicit_sigma(aug_y, t, args): y, y_adj, arg_adj = unpack(aug_y) gval = g(y, t, args) jac_g, jac_args = jacobian(g, argnums=(0, 2))(y, t, args) sigma_adj = jac_g * -y_adj sigma_theta = jac_args.T * -y_adj return np.vstack([np.diag(gval), sigma_adj, sigma_theta])
def _M(X, parameter_dict, t): """ proposal for a singular particle x """ #grab the parameters x = X['x'] seed = X['seed'] Dx = x.shape ldt = parameter_dict['ldt'] Mparams = parameter_dict['M_parameters'] potential_params = jnp.vstack( (Mparams['mu'][t], jnp.exp(Mparams['lcov_force'][t]))) #split the random seed run_seed, x_seed = random.split(seed) #call EL mu, cov = EL_mu_sigma(x, self._kernel_energy_fn, ldt, potential_params) #ula move new_x = ULA_move(x, self._kernel_energy_fn, ldt, run_seed, potential_params) return { 'x': new_x, 'seed': x_seed, 'forward_mu': mu, 'forward_cov': cov }
def test_scan_decorated(self): class SimpleScan(nn.Module): @partial(nn.scan, variable_broadcast='params', split_rngs={'params': False}) @nn.compact def __call__(self, c, xs): return nn.LSTMCell(name="lstm_cell")(c, xs) key1, key2 = random.split(random.PRNGKey(0), 2) xs = random.uniform(key1, (3, 2)) dummy_rng = random.PRNGKey(0) init_carry = nn.LSTMCell.initialize_carry(dummy_rng, xs.shape[:1], xs.shape[-1]) model = SimpleScan() init_variables = model.init(key2, init_carry, xs) # simulate scan in python for comparison: c = init_carry ys = [] lstmcell_variables = freeze( {'params': init_variables['params']['lstm_cell']}) for i in range(xs.shape[0]): c, y = nn.LSTMCell().apply(lstmcell_variables, c, xs[i]) ys.append(y[None, ...]) y1 = jnp.vstack(ys) c2, y2 = model.apply(init_variables, init_carry, xs) np.testing.assert_allclose(y1, y2, atol=1e-7) np.testing.assert_allclose(c[0], c2[0], atol=1e-7) np.testing.assert_allclose(c[1], c2[1], atol=1e-7)
def controlled_gate(gate, c_qubit): input_qubits = [c_qubit] + gate.input_qubits if gate.func is None: U = gate.tensor sh, _ = U.shape I = np.eye(sh) zero0 = np.zeros((sh, sh)) zero1 = np.zeros((sh, sh)) A = np.hstack([I, zero0]) B = np.hstack([zero1, U]) gate = qtc.Gate(input_qubits, tensor=np.vstack([A, B])) return gate else: def func(params): U = gate.func(params) sh, _ = U.shape I = np.eye(sh) zero0 = np.zeros((sh, sh)) zero1 = np.zeros((sh, sh)) A = np.hstack([I, zero0]) B = np.hstack([zero1, U]) return np.vstack([A, B]) func_jit = jit(func) res = qtc.Gate(input_qubits, params=gate.params, func=func_jit) return res
def get_2d_array(self): nb_full_blocks = int(self.nb_rows / self.nb_columns) block_list = [] rng = self.key for _ in range(nb_full_blocks): rng, rng_input = jax.random.split(rng) unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_columns)) q, _ = jnp.linalg.qr(unstructured_block) q = jnp.transpose(q) block_list.append(q) remaining_rows = self.nb_rows - nb_full_blocks * self.nb_columns if remaining_rows > 0: rng, rng_input = jax.random.split(rng) unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_columns)) q, _ = jnp.linalg.qr(unstructured_block) q = jnp.transpose(q) block_list.append(q[0:remaining_rows]) final_matrix = jnp.vstack(block_list) if self.scaling == 0: multiplier = jnp.linalg.norm( random.normal(self.key, (self.nb_rows, self.nb_columns)), axis=1) elif self.scaling == 1: multiplier = jnp.sqrt(float(self.nb_columns)) * jnp.ones((self.nb_rows)) else: raise ValueError('Scaling must be one of {0, 1}. Was %s' % self._scaling) return jnp.matmul(jnp.diag(multiplier), final_matrix)
def curve_loss(t_, params_): # t_ should be (m,) vector. params_ = jnp.vstack([w1, params_, w2]) c = curve(t_, params_) loss_sum = 0 for c_ in c: loss_sum += loss_fn(c_) return loss_sum / float(len(t_))