def nearest_point_distance(points, x_pixels=60, y_pixels=60): """Computes the distances from grid points to the nearest input point. Accepts a (batch of) list of xy points as input. The output is a (batch of) matrix of distances, where each entry of the matrix represents the distance from the corresponding point in the unit square to the nearest input point. Args: points: A batch_size x n x 2 matrix of points. x_pixels: How many x coordinates the output should have. y_pixels: How many y coordinates the output should have. Returns: A batch of matrices of distances. It is a tensor of dimension batch_size x x_pixels x y_pixels. """ x_coords = np.linspace(0.0, 1.0, x_pixels)[np.newaxis, np.newaxis, :] y_coords = np.linspace(0.0, 1.0, y_pixels)[np.newaxis, np.newaxis, :] x_points = points[:, :, 0:1] y_points = points[:, :, 1:2] x_diff_squared = jnp.square(x_points - x_coords) y_diff_squared = jnp.square(y_points - y_coords) distance_squared = jnp.amin( # distance to *nearest* point only x_diff_squared[:, :, :, np.newaxis] + y_diff_squared[:, :, np.newaxis, :], axis=1) return jnp.sqrt(distance_squared)
def _get_1d_latent_grid(self, paths_x): num_points_grid = 20 max_x = np.amax(paths_x) min_x = np.amin(paths_x) x_array = np.linspace(min_x, max_x, 20)[:, None] xx, tt = np.meshgrid(np.linspace(min_x, max_x, 20), np.linspace(0, 1, paths_x.shape[1])) txpairs = np.transpose(np.vstack([tt.reshape(-1), xx.reshape(-1)])) def scan_fn(carry, paths): drift, diffusion, index = carry time = index * self.config["delta_t"] - 0.5 gp_matrices, temp_drift_function, temp_diffusion_function = self.model.build( self.model.model_vars()) temp_drift = temp_drift_function(x_array, time) temp_diffusion = np.linalg.det( temp_diffusion_function(x_array, time)) drift = ops.index_add(drift, ops.index[index], temp_drift) diffusion = ops.index_add(diffusion, ops.index[index], temp_diffusion) index += 1 return (drift, diffusion, index), np.array([0.]) drift_grid = np.zeros((paths_x.shape[1], num_points_grid, 1)) diffusion_grid = np.zeros((paths_x.shape[1], num_points_grid)) (drift_grid, diffusion_grid, index), _ = lax.scan(scan_fn, (drift_grid, diffusion_grid, 0), np.transpose(paths_x, (1, 0, 2))) return txpairs, drift_grid.reshape(-1), diffusion_grid.reshape(-1)
def rescale_points(points, margin=0.1): """Rescales and translates points to lie inside unit square. Accepts a list of xy points as input, and applies a scalar multiplier and a translation so that the bounding box of the points becomes either [margin, 1 - margin] x (smaller interval centered on 0.5) or (smaller interval centered on 0.5) x [margin, 1 - margin]. Args: points: The points to transform. A batch_size x n x 2 matrix. margin: How much empty space to leave on each side of the unit square. Returns: A new array consisting of the translated and rotates points. """ min_corner = jnp.amin(points, axis=1, keepdims=True) # num_images x 1 x 2 points = points - min_corner # bounding boxes: [0, x_size] x [0, y_size] rectangular_size = jnp.amax(points, axis=1, keepdims=True) # num_images x 1 x 2 points = points - 0.5 * rectangular_size # bounding boxes: [-x_size/2, -y_size/2] x [x_size/2, y_size/2] square_size = jnp.amax(rectangular_size, axis=2, keepdims=True) # num_images x 1 x 1 points = points * (1 - 2 * margin) / jnp.maximum(square_size, 1e-5) points = points + 0.5 # center on (0.5, 0.5) return points
def uniform_stochastic_quantize(v: jnp.ndarray, num_levels: int, rng: PRNGKey, v_min: Optional[float] = None, v_max: Optional[float] = None) -> jnp.ndarray: """Uniform stochastic algorithm in https://arxiv.org/pdf/1611.00429.pdf. Args: v: vector to be quantized. num_levels: Number of levels of quantization. rng: jax random key. v_min: minimum threshold for quantization. If None, sets it to jnp.amin(v). v_max: maximum threshold for quantization. If None, sets it to jnp.amax(v). Returns: Quantized array. """ # Rescale the vector to be between zero to one. if v_min is None: v_min = jnp.amin(v) if v_max is None: v_max = jnp.amax(v) v = jnp.nan_to_num((v - v_min) / (v_max - v_min)) v = jnp.maximum(0., jnp.minimum(v, 1.)) # Compute the upper and lower boundary of each value. v_ceil = jnp.ceil(v * (num_levels - 1)) / (num_levels - 1) v_floor = jnp.floor(v * (num_levels - 1)) / (num_levels - 1) # uniformly quantize between v_ceil and v_floor. rand = jax.random.uniform(key=rng, shape=v.shape) threshold = jnp.nan_to_num((v - v_floor) / (v_ceil - v_floor)) quantized = jnp.where(rand > threshold, v_floor, v_ceil) # Rescale the values and return it. return v_min + quantized * (v_max - v_min)
def test_force(a_tensor, b_tensor): # a_eval, a_evec = np.linalg.eigh(a_tensor) # b_eval, b_evec = np.linalg.eigh(b_tensor) a_eval, a_evec = evp.dsyevv3(a_tensor) b_eval, b_evec = evp.dsyevv3(b_tensor) # print("ref w", a_eval) # print("test w", evp.dsyevc3(a_tensor)) # print("ref v", a_evec) # print("test v", evp.dsyevv3(a_tensor)) # assert 0 r = np.matmul(np.transpose(a_evec), b_evec) I = np.eye(3) rI = r * I # 3x3 -> 3x3 pos = np.sum(rI, axis=-1) # 3x3 -> 3 neg = -np.sum(rI, axis=-1) # 3x3 -> 3 acos_pos = np.arccos(pos) # 3 -> 3 acos_neg = np.arccos(neg) # 3 -> 3 a = np.amin([acos_pos, acos_neg], axis=0) # 2x3 -> 3 a2 = a * a # 3->3 l = np.sum(a2) # 3->1 # derivatives, start backprop dl_da2 = np.ones(3) # 1 x 3 da2_da = 2 * a * np.eye(3) # 3 x 3 da_darg = np.stack( [np.eye(3) * (acos_pos < acos_neg), np.eye(3) * (acos_neg < acos_pos)]) darg_dpn = np.stack([ np.eye(3) * (-1 / np.sqrt(1 - pos * pos)), np.eye(3) * (-1 / np.sqrt(1 - neg * neg)) ]) dl_darg = np.matmul(np.matmul(dl_da2, da2_da), da_darg) dpos = dl_darg[0] * (-1 / np.sqrt(1 - pos * pos)) dneg = dl_darg[1] * (-1 / np.sqrt(1 - neg * neg)) dneg = -dneg dpn_dr = np.array([ [[1, 1, 1], [0, 0, 0], [0, 0, 0]], [[0, 0, 0], [1, 1, 1], [0, 0, 0]], [[0, 0, 0], [0, 0, 0], [1, 1, 1]], ]) # element wise dr = (np.matmul(dpos, dpn_dr) + np.matmul(dneg, dpn_dr)) * np.eye(3) dr_daevec = np.matmul(b_evec, dr.T) dr_dbevec = np.matmul(a_evec, dr.T) dl_datensor = grad_eigh(a_eval, a_evec, np.zeros_like(a_eval), dr_daevec) dl_dbtensor = grad_eigh(b_eval, b_evec, np.zeros_like(b_eval), dr_dbevec) return dl_datensor, dl_dbtensor
def set_minmax(self, nn): try: self.xmin = nn.model.xmin self.xmax = nn.model.xmax except (NameError,AttributeError): self.xmin = np.amin(nn.model.x.data.numpy(),axis=0) self.xmax = np.amax(nn.model.x.data.numpy(),axis=0) self.range = (self.xmax - self.xmin)
def encode(self,x): try: self.xmin self.xmax except (NameError,AttributeError): self.xmin = np.amin(x.data.numpy(),axis=0) self.xmax = np.amax(x.data.numpy(),axis=0) x = (x.data.numpy()-self.xmin)/(self.xmax-self.xmin) return Variable(torch.from_numpy(x.copy()).type(dtype))
def sample(self, key, sample_shape=()): key_dirichlet, key_multinom = random.split(key) probs = self._dirichlet.sample(key_dirichlet, sample_shape) total_count = jnp.amax(self.total_count) if not_jax_tracer(total_count): # NB: the error can't be raised if inhomogeneous issue happens when tracing if jnp.amin(self.total_count) != total_count: raise NotImplementedError( "Inhomogeneous total count not supported" " by `sample`.") return Multinomial(total_count, probs).sample(key_multinom)
def pmi_u(r): I = np.eye(3) loss = [] for v, e in zip(r, I): a_pos = np.arccos(np.sum(v * e)) # norm is always 1 a_neg = np.arccos(np.sum(-v * e)) # norm is always 1 a = np.amin([a_pos, a_neg]) loss.append(a * a) return np.sum(loss)
def enumerate_support(self, expand=True): total_count = jnp.amax(self.total_count) if not_jax_tracer(total_count): # NB: the error can't be raised if inhomogeneous issue happens when tracing if jnp.amin(self.total_count) != total_count: raise NotImplementedError("Inhomogeneous total count not supported" " by `enumerate_support`.") values = jnp.arange(total_count + 1).reshape((-1,) + (1,) * len(self.batch_shape)) if expand: values = jnp.broadcast_to(values, values.shape[:1] + self.batch_shape) return values
def scan_fn(carry, paths): x, drift, diffusion, index = carry time = index * self.config["delta_t"] - 0.5 max_x = np.amax(paths_x[:, :, 0]) min_x = np.amin(paths_x[:, :, 0]) max_y = np.amax(paths_x[:, :, 1]) min_y = np.amin(paths_x[:, :, 1]) xx, yy = np.meshgrid(np.linspace(min_x, max_x, num_points_grid), np.linspace(min_y, max_y, num_points_grid)) temp = np.transpose(np.vstack([xx.reshape(-1), yy.reshape(-1)])) gp_matrices, temp_drift_function, temp_diffusion_function = self.model.build( self.model.model_vars()) temp_drift = temp_drift_function(temp, time) temp_diffusion = np.linalg.det(temp_diffusion_function(temp, time)) x = ops.index_add(x, ops.index[index], temp) drift = ops.index_add(drift, ops.index[index], temp_drift) diffusion = ops.index_add(diffusion, ops.index[index], temp_diffusion) index += 1 return (x, drift, diffusion, index), np.array([0.])
def __init__(self, data, leafsize=10): self.data = np.asarray(data) # if self.data.dtype.kind == 'c': # raise TypeError("KDTree does not work with complex data") self.n, self.m = np.shape(self.data) self.leafsize = int(leafsize) # if self.leafsize < 1: # raise ValueError("leafsize must be at least 1") self.maxes = np.amax(self.data,axis=0) self.mins = np.amin(self.data,axis=0) print('building tree') self.tree = self.__build(np.arange(self.n), self.maxes, self.mins)
def get_dx(grids): """Gets the grid spacing from grids array. Args: grids: Float numpy array with shape (num_grids,). Returns: Float, grid spacing. Raises: ValueError: If grids.ndim is not 1. """ if grids.ndim != 1: raise ValueError('grids.ndim is expected to be 1 but got %d' % grids.ndim) return (jnp.amax(grids) - jnp.amin(grids)) / (grids.size - 1)
def simplified_u(a_tensor, b_tensor): a_eval, a_evec = np.linalg.eigh(a_tensor) b_eval, b_evec = np.linalg.eigh(b_tensor) r = np.matmul(np.transpose(a_evec), b_evec) I = np.eye(3) rI = r * I # 3x3 -> 3x3 pos = np.sum(rI, axis=-1) neg = np.sum(-rI, axis=-1) acos_pos = np.arccos(pos) acos_neg = np.arccos(neg) # [a,b,c] # [d,e,f] # ------- # [min(a,d), min(b,e), min(c,f)] a = np.amin([acos_pos, acos_neg], axis=0) return np.sum(a * a)
def __build(self, idx, maxes, mins): if len(idx) <= self.leafsize: return KDTree.leafnode(idx) else: data = self.data[idx] # maxes = np.amax(data,axis=0) # mins = np.amin(data,axis=0) d = np.argmax(maxes-mins) maxval = maxes[d] minval = mins[d] if maxval == minval: # all points are identical; warn user? return KDTree.leafnode(idx) data = data[:,d] # sliding midpoint rule; see Maneewongvatana and Mount 1999 # for arguments that this is a good idea. split = (maxval+minval)/2 less_idx = np.nonzero(data <= split)[0] greater_idx = np.nonzero(data > split)[0] if len(less_idx) == 0: split = np.amin(data) less_idx = np.nonzero(data <= split)[0] greater_idx = np.nonzero(data > split)[0] if len(greater_idx) == 0: split = np.amax(data) less_idx = np.nonzero(data < split)[0] greater_idx = np.nonzero(data >= split)[0] if len(less_idx) == 0: # _still_ zero? all must have the same value if not np.all(data == data[0]): raise ValueError("Troublesome data array: %s" % data) split = data[0] less_idx = np.arange(len(data)-1) greater_idx = np.array([len(data)-1]) # lessmaxes = maxes.copy() # lessmaxes = index_update(lessmaxes, index[d], split) lessmaxes = np.asarray([x if (ii != d) else split for ii,x in enumerate(maxes)]) # lessmaxes[d] = split # greatermins = mins.copy() # greatermins = index_update(greatermins, index[d], split) greatermins = np.asarray([x if (ii != d) else split for ii,x in enumerate(mins)]) # greatermins[d] = split return KDTree.innernode(d, split, self.__build(idx[less_idx],lessmaxes,mins), self.__build(idx[greater_idx],maxes,greatermins))
def fprop(self, inputs, paddings=None): """Apply the stacking to inputs along the time axis. Args: inputs: The inputs tensor. It is expected to be of shape [batch, time, feature]. paddings: The paddings tensor. It is expected to be of shape [batch, time, 1], where all but the last dimension match inputs. Each value is 0 or 1 indicating whether a time step of a sequence is padded in the inputs to reach the max length in the batch. Returns: (outputs, out_paddings) pair. outputs is of shape [batch, ceil(time / stride), feature * stacking]. out_paddings is of shape [batch, ceil(time / stride), 1]. out_paddings will be 0 if any of the corresponding input padding is 0. """ p = self.params if paddings is None: paddings = jnp.zeros( jnp.concatenate([jnp.array(inputs.shape[:-1]), jnp.array([1])]), dtype=inputs.dtype) # Checks the inputs shape, paddings has 3 dimensions. base_layer.assert_has_shape(inputs, [-1, -1, -1]) base_layer.assert_has_shape(paddings, [-1, -1, 1]) # Trivial case. if 0 == p.left_context == p.right_context and 1 == p.stride: return inputs, paddings outputs = self._applystack(inputs) # Stack the padding values with the same context and stride parameters. # Then take the minimum padding values within each stacking window, since # an output time step becomes a padded one only if all of the underlying # stacked steps are padded ones. out_paddings = self._applystack(paddings, pad_value=1) if p.padding_reduce_option == 'reduce_min': out_paddings = jnp.amin(out_paddings, axis=2, keepdims=True) else: out_paddings = jnp.amax(out_paddings, axis=2, keepdims=True) return outputs, out_paddings
def pmi_restraints(conf, params, box, lamb, a_idxs, b_idxs, masses, angle_force, com_force): a_com, a_tensor = inertia_tensor(conf[a_idxs], masses[a_idxs]) b_com, b_tensor = inertia_tensor(conf[b_idxs], masses[b_idxs]) a_eval, a_evec = np.linalg.eigh(a_tensor) # already sorted b_eval, b_evec = np.linalg.eigh(b_tensor) # already sorted r = np.matmul(np.transpose(a_evec), b_evec) I = np.eye(3) loss = [] for v, e in zip(r, I): a_pos = np.arccos(np.sum(v * e)) # norm is always 1 a_neg = np.arccos(np.sum(-v * e)) # norm is always 1 a = np.amin([a_pos, a_neg]) loss.append(a * a)
def __fitCensoredMRR(self): ''' Median rank regression method (50%) to estimate parameters. Better for low number of failures. this method does not provides confidence bonds for parameters, however it is possible to estimate confidence bound in time using median rank for specified limit (other than 50%). :return: ''' iks = jnp.log(self.est_data['time'].to_numpy()) igrek = jnp.log( jnp.log(1.0 / (1.0 - util.median_rank( self.N, self.est_data['new_order_num'], 0.5)))) if self.method == Method.MRRCensored2p: slope, intercept, r, _, _ = self.__lineregress(iks, igrek) elif self.method == Method.MRRCensored3p: locs = jnp.linspace(0, jnp.amin(iks), 10) arr = np.empty((0, 4), float) for l in locs: l_iks = iks - l l_slope, l_intercept, l_r, _, _ = self.__lineregress( iks, igrek) arr = np.append(arr, np.array([[l, l_slope, l_intercept, l_r]]), axis=0) g = arr[np.argmax(arr[:, 3])] slope = g[1] intercept = g[2] r = g[3] self.loc = g[0] # assigning estimated parameters self.shape = slope self.scale = jnp.exp(-intercept / slope) self.est_data['ub'] = (jnp.log(1.0 / (1.0 - util.median_rank( self.N, self.est_data['new_order_num'], self.CL))) **(1.0 / self.shape) * self.scale) self.est_data['lb'] = (jnp.log(1.0 / (1.0 - util.median_rank( self.N, self.est_data['new_order_num'], 1.0 - self.CL))) **(1.0 / self.shape) * self.scale) self.r2 = r**2 self.est_data = self.est_data[['time', 'cdf', 'lb', 'ub']] self.method = Method.MRRCensored2p self.converged = True
def logit_transformer(logits, temp=1.0, confidence_quantile_threshold=1.0, self_supervised_label_transformation='soft', logit_indices=None): """Transforms logits into labels used as targets in a loss functions. Args: logits: jnp float array; Prediction of a model. temp: float; Softmax temp. confidence_quantile_threshold: float; Training examples are weighted based on this. self_supervised_label_transformation: str; Type of labels to produce (soft or sharp). logit_indices: list(int); Usable Indices for logits (list of indices to use). Returns: """ # Compute confidence for each prediction: confidence = jnp.amax(logits, axis=-1) - jnp.amin(logits, axis=-1) # Compute confidence threshold: alpha = jnp.quantile(confidence, confidence_quantile_threshold) # Only train on confident outputs: weights = jnp.float32(confidence >= alpha) if self_supervised_label_transformation == 'sharp': if logit_indices: logits = logits[Ellipsis, logit_indices] new_labels = jnp.argmax(logits, axis=-1) elif self_supervised_label_transformation == 'soft': new_labels = nn.softmax(logits / (temp or 1.0), axis=-1) else: new_labels = logits return new_labels, weights
def binary_stochastic_quantize(v: jnp.ndarray, rng: PRNGKey, v_min: Optional[float] = None, v_max: Optional[float] = None) -> jnp.ndarray: """Binary stochastic algorithm in https://arxiv.org/pdf/1611.00429.pdf. Args: v: vector to be quantized. rng: jax random key. v_min: minimum threshold for quantization. If None, sets it to jnp.amin(v). v_max: maximum threshold for quantization. If None, sets it to jnp.amax(v). Returns: Quantized array. """ if v_min is None: v_min = jnp.amin(v) if v_max is None: v_max = jnp.amax(v) v = jnp.nan_to_num((v - v_min) / (v_max - v_min)) v = jnp.maximum(0., jnp.minimum(v, 1.)) rand = jax.random.uniform(key=rng, shape=v.shape) return jnp.where(rand > v, v_min, v_max)
def get_groups(self, rng, state): if self.split_factor is None: # if no factor is given by default, we use prior infection rate. if np.size(state.prior_infection_rate) > 1: raise ValueError( 'Dorfman Splitting cannot be used with individual infection rates.' + ' Consider using Informative Dorfman instead.') # set group size to value defined by Dorfman testing group_size = 1 + np.ceil( 1 / np.sqrt(np.squeeze(state.prior_infection_rate))) # adjust to take into account testing limits group_size = np.amin((group_size, state.max_group_size)) split_factor = -(-state.num_patients // group_size) else: # ensure the split factor does not produce groups that are too large min_splits = -(-state.num_patients // state.max_group_size) split_factor = np.maximum(self.split_factor, min_splits) indices = onp.array_split(np.arange(state.num_patients), split_factor) new_groups = onp.zeros((len(indices), state.num_patients)) for i in range(len(indices)): new_groups[i, indices[i]] = True return np.array(new_groups, dtype=bool)
def subtract_min(x): y = x - np.amin(x) return (y)
def _amin(x, dim, keepdims=False): return np.amin(x, axis=dim, keepdims=keepdims)
lamb = 1e-10 print("len of spectrum") print(len(spectrum)) pvals = npo.logspace(np.log10(10), np.log10(num_pca-1), 500) sols = solve_implicit_z(npo.array(spectrum), pvals, lamb) mode_errs = theory_learning_curves(spectrum, vecs, pvals, lamb, y) sort_inds = np.argsort(spectrum)[::-1] theory0 = npo.sum(mode_errs, axis = 1) theory_adj = npo.sum(mode_errs, axis = 1) * num_pca / (num_pca - pvals + 1e-3) plt.loglog(pvals, theory0, label = 'original theory') plt.loglog(pvals, theory_adj, label = 'rescaled theory') plt.legend() plt.ylim([np.amin(theory_adj), np.amax(theory_adj)]) plt.xlabel(r'$p$', fontsize = 20) plt.ylabel(r'$E_g$', fontsize=20) plt.tight_layout() plt.savefig('rescale_risk.pdf') plt.show() inds = [10, 100, 1000] for i, j in enumerate(sort_inds[inds]): if inds[i]==0: plt.loglog(pvals, mode_errs[:,j] / mode_errs[0,j], label = r'$k=0$') else: plt.loglog(pvals, mode_errs[:,j] / mode_errs[0,j], label = r'$k = 10^{%d}$' % int(np.log10(inds[i])+0.01) )