def rHalo(self, masscut=5e12, masstag='m200b'): """ Given the positions and masses of a set of halos, calculate the distance to the nearest halo above the mass cut specified Parameters ---------- pos: array_like An array of positions in 3D mass: array_like The masses of the halos masscut: float The mass above which to find the nearest halo """ sii = self.halos[masstag].argsort() mass = self.halos[masstag][sii] pos = self.halos[['x', 'y', 'z']][sii] del sii lii = mass.searchsorted(masscut) self.rhalo = np.ndarray(len(pos), dtype=np.float64) with fast3tree(pos[lii:]) as tree: for i, p in enumerate(pos): self.rhalo[i] = tree.query_nearest_distance(p)
def rankSigma5(self, z, magnitude, sigma5, zwindow, magwindow): dsigma5 = np.max(sigma5) - np.min(sigma5) ranksigma5 = np.zeros(len(z)) pos = np.zeros((len(z), 3)) pos[:, 0] = z pos[:, 1] = magnitude pos[:, 2] = sigma5 max_distances = np.array([zwindow, magwindow, dsigma5]) neg_max_distances = -1.0 * max_distances tree_dsidx = np.random.choice(np.arange(len(z)), size=len(z) // 10) tree_pos = pos[tree_dsidx] with fast3tree(tree_pos) as tree: for i, p in enumerate(pos): tpos = tree.query_box(p + neg_max_distances, p + max_distances, output='pos') tpos = tpos[:, 2] tpos.sort() ranksigma5[i] = tpos.searchsorted(pos[i, 2]) / (len(tpos) + 1) return ranksigma5
def get_dist_and_attrs(hosts, gals, nn, attrs, box_size=250.0): """ Selects the nn-th nearest neighbor in real space from gals in hosts. Accepts: hosts - list of objects to search for nearest neighbors gals - objects to find nearest neighbors of nn - specifies which nearest neighbor to find attrs - list of properties to grab from halos (i.e. ['mvir', 'c']) Returns: dnbr - distances to the nn-th neighbor res - array of shape (len(attrs), len(gals)) """ pos_tags = ['x', 'y', 'z'] N = len(gals[pos_tags[0]]) pos = make_pos(hosts, pos_tags) dnbr = np.zeros(N) res = [np.zeros(N) for attr in attrs] with fast3tree(pos) as tree: for i in xrange(N): if i % 10000 == 0: print i, N center = [gals[tag][i] for tag in pos_tags] if box_size > 0: r, ind = get_nearest_nbr_periodic(center, tree, box_size, num_neighbors=nn) else: r, ind = get_nearest_nbr(center, tree) dnbr[i] = np.log10(r) for j, attr in enumerate(attrs): res[j][i] = hosts[attr][ind] return dnbr, res
def count_pairs_JKindex(points, rbins, box_size): pairs = [[] for r in rbins] with fast3tree(points) as tree: tree.set_boundaries(0, box_size) for pidx, p in enumerate(points): pJKidx = point2JKindex(p, box_size = box_size) for i, r in enumerate(rbins): idx, pos = tree.query_radius(p,r, periodic=True, output='both') pos = pos[idx < pidx] pairs[i].extend([pJKidx, j] for j in point2JKindex(pos, box_size = box_size)) pairs = np.array([np.array(i) for i in pairs]) return np.array(pairs) #def count_pairs_JKindex_old(points, rbins, box_size): # pairs = np.zeros(len(rbins), dtype=int) # with fast3tree(points) as tree: # tree.set_boundaries(0, box_size) # pairs = [] # for i, r in enumerate(rbins): # loc = [] # for p in points: # pindex = point2JKindex(p) # #save the distinct pairs # ppairs = np.array([np.sort([pindex, point2JKindex(pfound)]) for pfound in tree.query_radius(p,r, periodic=True, output='pos') if np.sum(np.abs(pfound - p)) > 0.0]) # if len(ppairs) : # loc.append(ppairs) # #Since every pair is counted twice # if len(loc) > 0: # loc = np.sort(np.concatenate(loc), axis = 0)[::2] # else: # loc = np.array([]) # pairs.append(loc) # return np.array(pairs)
def projected_correlation(points, rbins, zmax , box_size): points = np.asarray(points) s = points.shape if len(s) != 2 or s[1] != 3: raise ValueError('`points` must be a 2-d array with last dim=3') N = s[0] rbins = np.asarray(rbins) rbins_sq = rbins*rbins dcorner2 = np.array([rbins[-1], rbins[-1], zmax]) dcorner1 = -dcorner2 if np.any(dcorner2*2 > box_size): print "[Warning] box too small!" pairs_rand = float(N*N) / box_size**3 \ * (rbins[1:]**2-rbins[:-1]**2)*np.pi*zmax*2.0 dcorner1[2] = 0 #save some time pairs = np.zeros(len(rbins)-1, dtype=int) with fast3tree(points) as tree: for p in points: for pp in _yield_periodic_points(p,dcorner1,dcorner2,box_size): x,y=tree.query_box(pp+dcorner1,pp+dcorner2,output='p').T[:2] x -= pp[0]; x *= x y -= pp[1]; y *= y x += y; x.sort() pairs += np.ediff1d(np.searchsorted(x, rbins_sq)) return pairs
def calculate_red_density_score(gals, box_size, r_max=1, red_cut=-11, color='ssfr'): """Returns count of red neighbors within r_max of each galaxy. Future versions will have a weighting score for neighbor galaxies. This could be a inverse distance weighting or one that incorporates color. Accepts: gals - catalog containing galaxy positions and colors box_size - periodicity of objects in pos r_max - radius around center whitn which to grab neighbors Returns: red_neighbors - array of red neighbor counts for each galaxy """ N = len(gals) pos = make_pos(gals) red_neighbors = np.zeros(N) with fast3tree(pos) as tree: tree.set_boundaries(0.0, box_size) for i in xrange(N): if i % 10000 == 0: print i, N indices = tree.query_radius(pos[i], r_max, periodic=box_size, output='index') neighbor_colors = gals[color][indices] red_neighbors[i] = len(np.where(neighbor_colors < red_cut)[0]) return red_neighbors
def count_pairs(points, rbins, box_size): pairs = np.zeros(len(rbins), dtype=int) with fast3tree(points) as tree: tree.set_boundaries(0, box_size) for p in points: pairs += np.array([tree.query_radius(p, r, periodic=True, \ output='c') for r in rbins]) return (pairs-points.shape[0])/2
def calNN(data,boxsize): if rank ==0: print "calculate NN" min = boxsize/306/2. max = boxsize/2.-1. npoint = len(data) rho = npoint/boxsize**3 count = np.zeros((N+1,),dtype=np.float64) xi = np.zeros((N+1,),dtype=np.float64) r = np.zeros(N+1) dx = (np.log10(max)-np.log10(min))/(N) start_n = rank*npoint/size stop_n = np.array([(rank+1)*npoint/size-1,npoint-1]).min() comm.Barrier() if(stop_n > start_n): with fast3tree.fast3tree(data) as tree: tree.set_boundaries(0.0,boxsize) tree.rebuild_boundaries() for j in range(start_n,stop_n+1): if rank == 0: if j%((stop_n-start_n)/10) == 0: print "process",j*100./(stop_n-start_n),"%" for i in range(N+1): upper_r = 10.**(np.log10(min)+i*dx) idx = tree.query_radius(data[j],upper_r, periodic=True,output='count') - 1 r[i] = upper_r count[i] += idx comm.Barrier() # the 'totals' array will hold the sum of each 'data' array if comm.rank==0: print "Reducing data" # only processor 0 will actually get the data totals = np.zeros_like(count) else: totals = None # use MPI to get the totals comm.Reduce( [count, MPI.DOUBLE], [totals, MPI.DOUBLE], op = MPI.SUM, root = 0 ) if rank == 0: xi[0] = totals[0]/(4./3*np.pi*r[0]**3)/rho/npoint for i in range(1,N+1): dV = (4./3*np.pi*r[i]**3-4./3*np.pi*r[i-1]**3) xi[i] = (totals[i]-totals[i-1])/rho/dV/npoint print xi #pylab.plot(r,xi) #pylab.xscale('log') #pylab.show() #xi = comm.bcast(xi, root=0) return (r,xi)
def reassign_colors_cam(self, px, py, pz, hpx, hpy, hpz, m200, mr, amag, mhalo=12.466, corr=0.749, alpham=0.0689): centrals = h[(h['HOST_HALOID'] == -1) & (h['M200B'] > 10**mhalo)] cpos = np.zeros((len(centrals), 3)) pos = np.zeros((len(g), 3)) pos[:, 0] = g['PX'] pos[:, 1] = g['PY'] pos[:, 2] = g['PZ'] cpos[:, 0] = centrals['PX'] cpos[:, 1] = centrals['PY'] cpos[:, 2] = centrals['PZ'] rhalo = np.zeros(len(pos)) with fast3tree(cpos) as tree: for i in range(len(pos)): d = tree.query_nearest_distance(pos[i, :]) rhalo[i] = d mr = copy(g['MAG_R_EVOL']) mr[mr < -22] = -22 mr[mr > -18] = -18 gr = amag[:, 0] - amag[:, 1] idx = np.argsort(rhalo) rhalo_sorted = rhalo[idx] rank_rhalo = np.arange(len(rhalo)) / len(rhalo) corr_coeff = corr * (mr + 22)**(alpham) corr_coeff[corr_coeff > 1] = 1. noisy_rank_rhalo = abunmatch.noisy_percentile(rank_rhalo, corr_coeff) g = g[idx] gr = g['AMAG'][:, 0] - g['AMAG'][:, 1] idx_swap = abunmatch.conditional_abunmatch(g['MAG_R_EVOL'], noisy_rank_rhalo, g['MAG_R_EVOL'], -gr, 99, return_indexes=True) temp_sedid = g['SEDID'][idx_swap] return temp_sedid
def find_pairs(halos, host_flag, D): pairs = [] points = halos[host_flag][list('xyz')].view(float).reshape(-1, 3) with fast3tree(points) as tree: for i, p in enumerate(points): idx = tree.query_radius(p, D, True, 'index') idx = idx[idx > i] pairs.extend(((i, j) if halos['mvir'][host_flag[i]] > halos['mvir'][host_flag[j]] else (j, i) for j in idx)) return pairs
def test_fast3tree(): c = np.array([0.5, 0.5, 0.5]) r = 0.1 with fast3tree(points) as tree: ind = tree.query_radius(c, r) ind.sort() ind_true = find_sphere(c, points, r) assert len(ind) == len(ind_true) assert (ind == ind_true).all()
def test_fast3tree_periodic(): c = np.array([0, 0, 0]) r = 0.2 with fast3tree(points) as tree: tree.set_boundaries(0, 1) ind = tree.query_radius(c, r, periodic=True) ind.sort() ind_true = find_sphere(c, points, r, box_size=1.0) assert len(ind) == len(ind_true) assert (ind == ind_true).all()
def test_fast3tree_index(): c = np.array([0.5, 0.5, 0.5]) r = 0.1 index = np.random.randint(0, 100000, size=len(points)) with fast3tree(points, index) as tree: ind = tree.query_radius(c, r) ind.sort() ind_true = index[find_sphere(c, points, r)] ind_true.sort() assert len(ind) == len(ind_true) assert (ind == ind_true).all()
def radial_profile_counts(gals, hosts, box_size, r, rbins, rmax, col='ssfr'): """ Calculates the distribution of gals around hosts as a function of r """ num_halos = len(hosts) results = [] pos = make_pos(gals) pos_tags = ['x', 'y', 'zr'] with fast3tree(pos) as tree: tree.set_boundaries(0.0, box_size) #mass_select = hosts[hosts['mbin_idx'] == i] num_reds, num_blues = [],[] num_pred_reds, num_pred_blues = [], [] diff_reds, diff_blues = [], [] for j in xrange(len(hosts)): num_red, num_blue = np.zeros(len(r)), np.zeros(len(r)) num_pred_red, num_pred_blue = np.zeros(len(r)), np.zeros(len(r)) center = [hosts[tag][j] for tag in pos_tags] idxs, pos = tree.query_radius(center, rmax, periodic=box_size, output='both') rs = get_3d_distance(center, pos, box_size=box_size) msk = rs > 0 rs = rs[msk] idxs = idxs[msk] for dist, sat_idx in zip(rs, idxs): rbin = np.digitize([dist], rbins) - 1 # -1 for the r vs rbin # indexing reflects return values from tree if gals['ssfr'][sat_idx] < -11: num_red[rbin] += 1 else: num_blue[rbin] += 1 if gals['pred'][sat_idx] < -11: num_pred_red[rbin] += 1 else: num_pred_blue[rbin] += 1 num_reds.append(num_red) num_blues.append(num_blue) num_pred_reds.append(num_pred_red) num_pred_blues.append(num_pred_blue) diff_reds.append(num_red - num_pred_red) diff_blues.append(num_blue - num_pred_blue) all_counts = map(np.array, [num_reds, num_blues, num_pred_reds, num_pred_blues]) means = map(lambda x: np.mean(x, axis=0), all_counts) stds = map(lambda x: np.std(x, axis=0), [diff_reds, diff_blues]) for counts in means: results.append(counts) for errs in stds: results.append(errs) return results
def calculate_density_score(gals, box_size, r_max=1): """ Calculates the number of neighbors within r_max of each galaxy. """ N = len(gals) pos = make_pos(gals) neighbors = np.zeros(N) with fast3tree(pos) as tree: tree.set_boundaries(0.0, box_size) for i in xrange(N): if i % 10000 == 0: print i, N count = tree.query_radius(pos[i], r_max, periodic=box_size, output='count') neighbors[i] = count return neighbors
def get_all_neighbors(pos, center, box_size): """Returns indices and positions of all neighbors within a half box_size of center in pos. Accepts: pos - list of positions (array-like) center - list of 3 coordinates box_size - periodicity of objects in pos Returns: idxs - indices of of neighbors in pos loc - positions of neighbors in pos """ with fast3tree(pos) as tree: tree.set_boundaries(0.0, box_size) rmax = box_size/2 - 1e-6 return tree.query_radius(center, rmax, periodic=box_size, output='both')
def computeSigma5(self, z, mag, pos_gals, dt=1.5): pos_bright_gals = pos_gals[mag < -19.8] max_distances = np.array([dt, dt, 1000]) neg_max_distances = -1.0 * max_distances sigma5 = np.zeros(len(pos_gals)) with fast3tree(pos_bright_gals) as tree: for i, p in enumerate(pos_gals): tpos = tree.query_box(p + neg_max_distances, p + max_distances, output='pos') dtheta = np.abs(tpos - p) dtheta = 2 * np.arcsin( np.sqrt( np.sin(dtheta[:, 0] / 2)**2 + np.cos(p[0] * np.cos(tpos[:, 0])) * np.sin(dtheta[:, 1] / 2)**2)) dtheta.sort() try: sigma5[i] = dtheta[4] except IndexError as e: sigma5[i] = -1 z_a = copy(z) z_a[z_a < 1e-6] = 1e-6 try: da = self.nbody.cosmo.angularDiameterDistance(z_a) except RuntimeError: print(np.min(z_a)) print(np.max(z_a)) print(np.isfinite(z_a).all()) print(z_a[~np.isfinite(z_a)]) print(self.nbody.domain.zmin, self.nbody.domain.zmax, self.nbody.domain.nest) raise RuntimeError sigma5 = sigma5 * self.nbody.cosmo.angularDiameterDistance(z_a) return sigma5
def get_projected_dist_and_attrs(hosts, gals, nn, attrs, box_size=250.0): """ Selects the nn-th nearest neighbor in redshift space from gals in hosts. Accepts: hosts - list of objects to search for nearest neighbors gals - objects to find nearest neighbors of nn - specifies which nearest neighbor to find attrs - list of properties to grab from halos (i.e. ['mvir', 'c']) Returns: dnbr - distances to the nn-th neighbor res - array of shape (len(attrs), len(gals)) """ width_by_2 = 0.01 pos_tags = ['x', 'y'] host_pos = make_pos(hosts, pos_tags) gal_pos = make_pos(gals, pos_tags) N = len(gals[pos_tags[0]]) gal_z = gals['zr'] host_z = hosts['zr'] dnbr = np.zeros(N) res = [np.zeros(N) for attr in attrs] for i in xrange(N): if i % 10000 == 0: print i, N sel = np.where(np.abs(gal_z[i] - host_z) < width_by_2)[0] center = gal_pos[i] with fast3tree(host_pos[sel]) as tree: if len(sel) <= nn: print "Insufficient number of neighbors in redshift bin" print "redshift: ", gal_z[i] assert False if box_size < 0: r, ind = get_nearest_nbr_periodic(center, tree, box_size, num_neighbors=nn) else: r, ind = get_nearest_nbr(center, tree) dnbr[i] = np.log10(r) for j, attr in enumerate(attrs): res[j][i] = hosts[attr][sel][ind] return dnbr, res
def remove_subhalos(halos, halos_nd, mcut = 1.e14): hosts = halos[halos['mvir']>mcut][['x','y','z']].view((float,3)) rvir = halos[halos['mvir']>mcut][['rvir']].view((float,1)) hostidx = halos[halos['mvir']>mcut][['id']].view((int,1)) hostidx_nd = np.array([np.where(halos_nd['id'].view((int,1))==idx)[0] for idx in hostidx]) #deal with the case that the host halos are not in the selected nd sample hostidx_nd = np.array([idx[0] for idx in hostidx_nd if len(idx)==1]) points = halos_nd[list('xyz')].view((float,3)) not_subhalo_bool = np.ones(len(points), dtype=int) with fast3tree(points) as tree: #tree.set_boundaries(0, box_size) #no periodic boundary condition for i, (p, rv, idx_nd) in enumerate(zip(hosts, rvir, hostidx_nd)): # rv is in unit Kpc, so divide by 1000 to get units in Mpc sublist = tree.query_radius(p, 2*rv/1000., periodic=False, output='index') for k in sublist: if k!= idx_nd: not_subhalo_bool[k] = 0 return halos_nd[np.where(not_subhalo_bool)]
def calc_mcf(mark, pos, rbins, box_size): """ mark : 1d ndarray, length N pos : 2d ndarray, shape (N, 3) rbins : 1d ndarray box_size : float """ pairs = [] rmax = rbins[-1] with fast3tree(pos) as tree: tree.set_boundaries(0, box_size) for i, c in enumerate(pos): j = tree.query_radius(c, rmax, True) j = j[j>i] pairs.extend((i, _j) for _j in j) pairs = np.array(pairs) d = pos[pairs[:,0]] - pos[pairs[:,1]] d[d > (0.5*box_size)] -= box_size d[d < (-0.5*box_size)] += box_size d *= d d = np.sqrt(d.sum(axis=-1)) s = d.argsort() k = np.searchsorted(d, rbins, sorter=s) del d # make mark_rank to span -1 to +1 mark_rank = rankdata(mark) mark_rank -= 1.0 mark_rank *= (2.0/float(len(mark)-1)) mark_rank -= 1.0 mcf = [] for i, j in zip(k[:-1], k[1:]): if j==i: mcf.append(np.nan) else: ii, jj = pairs[s[i:j]].T mcf.append((mark_rank[ii]*mark_rank[jj]).mean()) return np.array(mcf)
def compute_distances(px, py, pz, hpx, hpy, hpz, hmass, mcut): idx = (hmass > 10**mcut) cpos = np.zeros((np.sum(idx), 3)) pos = np.zeros((len(px), 3)) pos[:, 0] = px pos[:, 1] = py pos[:, 2] = pz cpos[:, 0] = hpx[idx] cpos[:, 1] = hpy[idx] cpos[:, 2] = hpz[idx] rhalo = np.zeros(len(pos)) with fast3tree(cpos) as tree: for i in range(len(pos)): d = tree.query_nearest_distance(pos[i, :]) rhalo[i] = d return rhalo
def radial_conformity(centrals, neighbors, msmin, msmax, box_size, rbins, satellites=False, red_cut=-11, col='ssfr'): """ Calculates quenched fraction of satellites of quenched/star-forming centrals binned by radius. """ rmin, rmax = np.min(rbins), np.max(rbins) nrbins = len(rbins) - 1 all_central_nbr_counts = [[] for _ in xrange(nrbins)] q_central_nbr_counts = [[] for _ in xrange(nrbins)] sf_central_nbr_counts = [[] for _ in xrange(nrbins)] n_pos = make_pos(neighbors) with fast3tree(n_pos) as tree: for c_pos, c_color, c_id in zip(centrals[list('xyz')], centrals[col], centrals['id']): idx, pos = tree.query_radius(list(c_pos), rmax, periodic=box_size, output='both') distances = get_3d_distance(c_pos, pos, box_size) for ii, dist in zip(idx, distances): #print dist if satellites: if neighbors[ii]['upid'] is not c_id: continue if dist < rmin or dist > rmax: continue rbin = np.digitize([dist], rbins, right=True)[0] - 1 nbr_red = centrals[ii][col] < red_cut if c_color < red_cut: q_central_nbr_counts[rbin].append(nbr_red) else: sf_central_nbr_counts[rbin].append(nbr_red) all_central_nbr_counts[rbin].append(nbr_red) def quenched_neighbor_fraction(nbr_counts): return np.array([np.mean(count) for count in nbr_counts]) return quenched_neighbor_fraction(q_central_nbr_counts), \ quenched_neighbor_fraction(sf_central_nbr_counts), \ quenched_neighbor_fraction(all_central_nbr_counts)
def get_dist_and_attrs(dp, d, nn, attrs): """ dp - parent set of halos d - galaxy set nn - num neighbors attrs - list of attributes (i.e. ['mvir','vmax'] """ pos = np.zeros((len(dp), 3)) for i, tag in enumerate(['x', 'y', 'z']): pos[:, i] = dp[tag][:] dnbr = np.zeros(len(d)) res = [np.zeros(len(d)) for attr in attrs] with fast3tree(pos) as tree: for i in xrange(len(d)): if i % 10000 == 0: print i, len(d) center = [d['x'].values[i], d['y'].values[i], d['z'].values[i]] r, ind = get_nearest_nbr_periodic(center, tree, box_size, num_neighbors=nn, exclude_self=True) dnbr[i] = np.log10(r) for j, attr in enumerate(attrs): res[j][i] = dp[attr].values[ind] return dnbr, res
def correlation3d(points, rbins, box_size): """ Calculate the 3D correlation function xi(r) for a periodic box. Parameters ---------- points : array_like Must be a 2-d array whose last dimension is 3 (i.e. has 3 columns). rbins : array_like A 1-d array that has the edges of the rp bins. Must be sorted. box_size : float The side length of the periodic box. Returns ------- xi : ndarray A 1-d array that has wp. The length of this retured array would be len(rbins) - 1. """ points = np.asarray(points) s = points.shape if len(s) != 2 or s[1] != 3: raise ValueError('`points` must be a 2-d array with last dim=3') N = s[0] rbins = np.asarray(rbins) pairs_rand = float(N*N) / box_size**3 \ * (rbins[1:]**3-rbins[:-1]**3)*(np.pi*4.0/3.0) pairs = np.zeros(len(rbins)-1, dtype=int) with fast3tree(points) as tree: tree.set_boundaries(0, box_size) for p in points: pairs += np.ediff1d([tree.query_radius(p, r, periodic=True, \ output='c') for r in rbins]) return pairs.astype(float)/pairs_rand - 1.0
def correlation3d(points, rbins, box_size): """ Calculate the 3D correlation function xi(r) for a periodic box. Parameters ---------- points : array_like Must be a 2-d array whose last dimension is 3 (i.e. has 3 columns). rbins : array_like A 1-d array that has the edges of the rp bins. Must be sorted. box_size : float The side length of the periodic box. Returns ------- xi : ndarray A 1-d array that has wp. The length of this retured array would be len(rbins) - 1. """ points = np.asarray(points) s = points.shape if len(s) != 2 or s[1] != 3: raise ValueError('`points` must be a 2-d array with last dim=3') N = s[0] rbins = np.asarray(rbins) pairs_rand = float(N*N) / box_size**3 \ * (rbins[1:]**3-rbins[:-1]**3)*(np.pi*4.0/3.0) pairs = np.zeros(len(rbins) - 1, dtype=int) with fast3tree(points) as tree: tree.set_boundaries(0, box_size) for p in points: pairs += np.ediff1d([tree.query_radius(p, r, periodic=True, \ output='c') for r in rbins]) return pairs.astype(float) / pairs_rand - 1.0
def calculate_clustering_score(gals, box_size, pos_tags=['x','y','zp'], rbins=[0,1,10]): """ Counts the number of pairs in different bins of cylindrical annuli """ N = len(gals) pos = make_pos(gals, pos_tags) neighbors = np.zeros((N, len(rbins) - 1)) r_max = rbins[-1] with fast3tree(pos) as tree: if box_size > 0: tree.set_boundaries(0.0, box_size) for i in xrange(N): if i % 10000 == 0: print i, N idxs, loc = tree.query_radius(pos[i], r_max, periodic=box_size, output='both') distances = get_cylinder_distance(pos[i], loc, box_size) for j in xrange(len(rbins) - 1): binned_neighbors = len(np.where((distances > rbins[j]) & (distances < rbins[j+1]))[0]) neighbors[i][j] = binned_neighbors return neighbors
def isolated(pairs, halos, host_flag, D_iso, D_M33, box_size, vmax_cut=None, periodic=True): print 'before: ', len(pairs) larger_host = [] larger_host_sub = [] smaller_host = [] smaller_host_sub = [] with fast3tree(halos[list('xyz')].view(float).reshape(-1, 3)) as tree: for pair in pairs: pair_idx = host_flag[list(pair)] h1, h2 = halos[pair_idx] #h1 is the larger halo p1 = np.fromiter((h1[ax] for ax in 'xyz'), float) p2 = np.fromiter((h2[ax] for ax in 'xyz'), float) mid = find_periodic_midpoint(p1, p2, box_size, periodic) idx = tree.query_radius(mid, D_iso, periodic, 'index') biggest = halos['mvir'][idx].argmax() idx = np.delete(idx, biggest) biggest = halos['mvir'][idx].argmax() #biggest is now second-biggest # the second index has a smaller mass if h2['mvir'] == halos['mvir'][idx[biggest]]: larger_host.append(pair[0]) smaller_host.append(pair[1]) else: continue #no need to search M33 #find M33 M33_ind, M33_mass = find_largest_sub(tree, p1, D_M33, halos, pair_idx, vmax_cut, periodic) LMC_ind, LMC_mass = find_largest_sub(tree, p2, D_M33, halos, pair_idx, vmax_cut, periodic) larger_host_sub.append(M33_ind) smaller_host_sub.append(LMC_ind) print 'after: ', len(larger_host) return larger_host, larger_host_sub, smaller_host, smaller_host_sub
def projected_correlation(points, rbins, zmax, box_size, jackknife_nside=0): """ Calculate the projected correlation function wp(rp) and its covariance matrix for a periodic box, with the plane-parallel approximation and the Jackknife method. Parameters ---------- points : array_like Must be a 2-d array whose last dimension is 3 (i.e. has 3 columns) The last column will be used as the redshift distance. rbins : array_like A 1-d array that has the edges of the rp bins. Must be sorted. zmax : float The integral of \pi goes from -zmax to zmax (redshift distance). box_size : float The side length of the periodic box. jackknife_nside : int, optional (Default: 0) If <= 1 , it will not do Jackknife. Returns ------- wp : ndarray A 1-d array that has wp. The length of this retured array would be len(rbins) - 1. wp_cov : ndarray (returned if jackknife_nside > 1) The len(wp) by len(wp) covariance matrix of wp. """ points = np.asarray(points) s = points.shape if len(s) != 2 or s[1] != 3: raise ValueError('`points` must be a 2-d array with last dim=3') N = s[0] rbins = np.asarray(rbins) rbins_sq = rbins*rbins dcorner2 = np.array([rbins[-1], rbins[-1], zmax]) dcorner1 = -dcorner2 if np.any(dcorner2*2 > box_size): print "[Warning] box too small!" pairs_rand = float(N*N) / box_size**3 \ * (rbins[1:]**2-rbins[:-1]**2)*np.pi*zmax*2.0 jackknife_nside = int(jackknife_nside) if jackknife_nside <= 1: #no jackknife dcorner1[2] = 0 #save some time pairs = np.zeros(len(rbins)-1, dtype=int) with fast3tree(points) as tree: for p in points: for pp in _yield_periodic_points(p,dcorner1,dcorner2,box_size): x,y=tree.query_box(pp+dcorner1,pp+dcorner2,output='p').T[:2] x -= pp[0]; x *= x y -= pp[1]; y *= y x += y; x.sort() pairs += np.ediff1d(np.searchsorted(x, rbins_sq)) return (pairs.astype(float)*2.0/pairs_rand - 1.0) * zmax*2.0 else: #do jackknife jack_ids = np.floor(np.remainder(points[:,0], box_size)\ / box_size*jackknife_nside).astype(int) jack_ids += np.floor(np.remainder(points[:,1], box_size)\ / box_size*jackknife_nside).astype(int) * jackknife_nside n_jack = jackknife_nside*jackknife_nside pairs = np.zeros((n_jack, len(rbins)-1), dtype=int) auto_pairs = np.zeros_like(pairs) with fast3tree(points) as tree: for p, jid in izip(points, jack_ids): for pp in _yield_periodic_points(p,dcorner1,dcorner2,box_size): idx,pos = tree.query_box(pp+dcorner1,pp+dcorner2,output='b') x, y = pos.T[:2] x -= pp[0]; x *= x y -= pp[1]; y *= y x += y y = x[jack_ids[idx]==jid] y.sort(); x.sort() pairs[jid] += np.ediff1d(np.searchsorted(x, rbins_sq)) auto_pairs[jid] += np.ediff1d(np.searchsorted(y, rbins_sq)) idx = pos = x = y = None pairs_sum = pairs.sum(axis=0) pairs = pairs_sum - pairs*2 + auto_pairs wp_jack = (pairs.astype(float) \ / pairs_rand \ / _jackknife_2d_random(rbins, box_size, jackknife_nside)\ - 1.0) * zmax*2.0 wp_full = (pairs_sum.astype(float)/pairs_rand - 1.0) * zmax*2.0 wp = wp_full*n_jack - wp_jack.mean(axis=0)*(n_jack-1) wp_cov = np.cov(wp_jack, rowvar=0, bias=1)*(n_jack-1) return wp, wp_cov
def plot_density_profile(d0, d_gals, test_gals, name): """ Binning by mass and radius, this function makes use of the count_neighbors_within_r helper to show radial profiles of parent galaxies. """ rmin, rmax, Nrp = 0.1, 5.0, 10 rbins = np.logspace(np.log10(rpmin), np.log10(rpmax), Nrp+1) r = np.sqrt(rbins[1:]*rbins[:-1]) dp = d0[d0['upid'] == -1] # Set up a tree with the galaxy set dp = dp.reset_index(drop=True) d_gals = d_gals.reset_index(drop=True) pos = np.zeros((len(d_gals), 3)) for i, tag in enumerate(['x', 'y', 'z']): pos[:, i] = d_gals[tag][:] mvir = dp['mvir'].values mmin, mmax = np.min(mvir), np.max(mvir) nmbins = 3 mbins = np.logspace(np.log10(mmin), np.log10(mmax), nmbins+1) num_halos, _ = np.histogram(mvir, mbins) dp['mbin_idx'] = np.digitize(mvir, mbins) with fast3tree(pos) as tree: tree.set_boundaries(0.0, box_size) for i in xrange(nmbins): mass_select = dp[dp['mbin_idx'] == i] num_blue_actual = np.zeros(len(rbins)) num_blue_pred = np.zeros(len(rbins)) num_red_actual = np.zeros(len(rbins)) num_red_pred = np.zeros(len(rbins)) # change the order of the loop after querying the radius for j, halo in mass_select.iterrows(): center = halo['x'], halo['y'], halo['z'] idxs, pos = tree.query_radius(center, rmax, periodic=box_size, output='both') dx = get_distance(center[0], pos[:, 0], box_size=box_size) dy = get_distance(center[1], pos[:, 1], box_size=box_size) dz = get_distance(center[2], pos[:, 2], box_size=box_size) r2 = dx*dx + dy*dy + dz*dz msk = r2 > 0.0 q = np.argsort(r2[msk]) rs = np.sqrt(r2[msk][q]) idxs = idxs[msk][q] for dist, sat_idx in zip(rs, idxs): #print dist #print rbins rbin = np.digitize([dist], rbins) # query for large radius and then do processing in here if d_gals['ssfr'].values[sat_idx] < -11: num_red_actual[rbin] += 1 else: num_blue_actual[rbin] += 1 test_gal = test_gals[test_gals['id'] == d_gals['id'].values[sat_idx]] if len(test_gal): if test_gal['pred'].values[0] < -11: num_red_pred[rbin] += 1 else: num_blue_pred[rbin] += 1 volumes = [4./3 * np.pi * r**3 for r in rbins] num_red_actual /= num_halos[i] num_blue_actual /= num_halos[i] num_red_pred /= num_halos[i] num_blue_pred /= num_halos[i] num_red_actual /= volumes num_blue_actual /= volumes num_red_pred /= volumes num_blue_pred /= volumes plt.figure(i) plt.loglog(rbins, num_red_actual, color=red_col, lw=4, label='input') plt.loglog(rbins, num_red_pred, color=red_col, label='pred', alpha=0.5) plt.loglog(rbins, num_blue_actual, color=blue_col, lw=4, label='input') plt.loglog(rbins, num_blue_pred, color=blue_col, label='pred', alpha=0.5) plt.legend(loc='best') plt.xlabel('distance') plt.ylabel('<centrals + satellites>') plt.title('Radial density {:.2f}'.format(mbins[i]) + ' < mvir < {:.2f}'.format(mbins[i+1])) return
# select the points inside the JK box xr, yr, zr = JKindex2JKbox(JKidx, box_size, n) xbool = (xr[0] - halos[:,0])*(xr[1] - halos[:,0]) < 0 ybool = (yr[0] - halos[:,1])*(yr[1] - halos[:,1]) < 0 zbool = (zr[0] - halos[:,2])*(zr[1] - halos[:,2]) < 0 halos_JK = halos[np.all([xbool, ybool, zbool],axis=0)] #determine the neighboring 27 JK indices (including itself) center = np.mean([xr, yr, zr], axis=-1) xx, yy, zz = [grid.flatten() for grid in np.meshgrid([-1.,0,1.],[-1.,0,1.],[-1.,0,1.],indexing='ij')] dx, dy, dz = float(box_size)/n*np.identity(3) shifts = np.outer(xx, dx) + np.outer(yy, dy) + np.outer(zz,dz) neighbor_JKindex = point2JKindex(np.mod(center + shifts , box_size), box_size, n) # calculate the number of halo partners that are inside each of the 27 neighboring boxes # crosspairslist[r, 13] is the number of autopairs # start the autopairs with negative so that we don't count the halo pairing with itself crosspairslist = np.zeros((len(rbins), 27)) crosspairslist[:,13] = -len(halos_JK)*np.ones_like(rbins) with fast3tree(halos) as tree: tree.set_boundaries(0, box_size) for p in halos_JK: for i, r in enumerate(rbins): pos = tree.query_radius(p,r, periodic=True, output='pos') JKpos = point2JKindex(pos, box_size, n) for k, neighborJK in enumerate(neighbor_JKindex): crosspairslist[i,k] += np.count_nonzero(JKpos == neighborJK) JKcrosspairs[JKidx] = crosspairslist fname = 'pairs_{}/boundaries/{}_nd{:.1f}_nsave{}'.format(proxy, case, nd_log, n) np.save(fname, JKcrosspairs)
def calculate_r_hill(halo_set, massive_halos, rmax=5): """ calculate_r_hill iterates over all halos in the halo set. For each halo, calculate the rhill generated by the 10 more massive nearest neighbors in massive_halos, save the r_hill min and corresponding distance and mass of the determining neighbor halo. """ half_box_size = box_size/2 massive_halos = massive_halos.reset_index(drop=True) halo_set = halo_set.reset_index(drop=True) r_hills = [] halo_dists = [] halo_masses = [] r_hills = [] # Iterate over the halos and compare for i, halo in halo_set.iterrows(): if i % 5000 == 0: print i m_sec = halo['mvir'] center = [halo['x'], halo['y'], halo['z']] larger_halos = massive_halos[massive_halos['mvir'] > m_sec] #.values? pos = np.zeros((len(larger_halos), 3)) for i, tag in enumerate(['x', 'y', 'z']): pos[:, i] = larger_halos[tag][:] num_tries = 0 with fast3tree(pos) as tree: tree.set_boundaries(0.0, box_size) # First find the nearest neighbor to get a sense of the density rmax = tree.query_nearest_distance(center) + 5 if rmax > half_box_size: rmax = half_box_size - 1e-6 while True: if num_tries > 3: break idxs, pos = tree.query_radius(center, rmax, periodic=box_size, output='both') # repeat with a larger radius if there are fewer than 10 neighbors if len(idxs) < 10: rmax *= 3.0 num_tries += 1 else: break dx = get_distance(center[0], pos[:, 0], box_size=box_size) dy = get_distance(center[1], pos[:, 1], box_size=box_size) dz = get_distance(center[2], pos[:, 2], box_size=box_size) r2 = dx*dx + dy*dy + dz*dz msk = r2 > 0.0 rs = np.sqrt(r2[msk]) idxs = idxs[msk] if len(rs) > 0: rhill_candidates = [r * (m_sec/(3 * massive_halos['mvir'][idx]))**(1./3) for r, idx in zip(rs, idxs)] rhill = min(rhill_candidates) idx_idx = np.argmin(rhill_candidates) halo_dist = rs[idx_idx] halo_mass = massive_halos['mvir'][idxs[idx_idx]] else: rhill = half_box_size halo_dist = np.nan halo_mass = np.nan r_hills.append(rhill) halo_dists.append(halo_dist) halo_masses.append(halo_mass) return r_hills, halo_dists, halo_masses
def main(): for isnap in range(NSnaps): print "snap", isnap folder = outputfolder + "/snap_%03d/" % (isnap) ahf_halos = [] rho = overdensities[0] desc_folder = outputfolder + "/snap_%03d/multidenshalos" % (isnap) outfile_prefix = folder + "/" + prefix_template + "_rho_%04d" % ( long(rho + 0.5)) z = get_z(outfile_prefix) #get_nhalos(outfile_prefix,z) halos = [] pids = [] print "read halo catalogue" for rho in overdensities: outfile_prefix = folder + "/" + prefix_template + "_rho_%04d" % ( long(rho + 0.5)) (halo, pid) = read_ahf_halos_snap(outfile_prefix, z) halos.append(halo) pids.append(pid) for i in range(len(overdensities) - 1): print "doing rho =", overdensities[i] lo_index = numpy.where(halos[i][:, 1] == 0)[0] lo_halos = halos[i][lo_index] hi_index = numpy.where(halos[i + 1][:, 1] == 0)[0] hi_halos = halos[i + 1][hi_index] lo_halos[:, 0:2] = -1 hi_halos[:, 0:2] = -1 for ih in range(len(lo_halos)): lo_halos[ih, 0] = ih for ih in range(len(hi_halos)): hi_halos[ih, 0] = ih os.system("mkdir -p " + folder + "/multilevels/") outfile = folder + "/multilevels/" + str( overdensities[i]) + "_to_" + str(overdensities[i + 1]) + ".txt" f = open(outfile, "w") if len(hi_halos) & len(lo_halos): pos = hi_halos[:, 5:8] with fast3tree.fast3tree(pos) as tree: tree.set_boundaries(0.0, boxsize_kpc) tree.rebuild_boundaries() for ii in range(len(lo_halos)): h = lo_halos[ii] idx = tree.query_radius(h[5:8], h[11], periodic=True, output='index') if len(idx) > 0: string = "\t".join([str(id) for id in idx]) print >> f, "%d\t%s" % (ii, string) hi_halos[idx, 1] = ii else: print >> f, ii del (tree) # pos = lo_halos[:,5:8] # with fast3tree.fast3tree(pos) as tree: # tree.set_boundaries(0.0,boxsize_kpc) # tree.rebuild_boundaries() # for ii in range(len(hi_halos)): # h = hi_halos[ii] # idx = tree.query_radius(h[5:8],h[11], periodic=True,output='index') # if len(idx)>0: # if(idx[0]>=len(lo_halos)): # print "something is wrong, idx = ",idx,len(lo_halos) # lo_halos[idx[0],0] = ii # del(tree) f.close() outfile = folder + "/multilevels/" + prefix_template + str( overdensities[i + 1]) + "halo.txt" numpy.savetxt(outfile, hi_halos) outfile = folder + "/multilevels/" + prefix_template + str( overdensities[i + 1]) + "particle.txt" f = open(outfile, "w+") print >> f, len(hi_index) for ihalo in range(len(hi_index)): id_halo = hi_index[ihalo] print >> f, len(pids[i + 1][id_halo]), ihalo string = "\n".join( [str(id) + "\t1" for id in pids[i + 1][id_halo]]) print >> f, string f.close() if i == 0: outfile = folder + "/multilevels/" + prefix_template + str( overdensities[i]) + "halo.txt" numpy.savetxt(outfile, lo_halos) outfile = folder + "/multilevels/" + prefix_template + str( overdensities[i]) + "particle.txt" f = open(outfile, "w+") print >> f, len(lo_index) for ihalo in range(len(lo_index)): id_halo = lo_index[ihalo] print >> f, len(pids[i][id_halo]), ihalo string = "\n".join( [str(id) + "\t1" for id in pids[i][id_halo]]) print >> f, string f.close() for isnap in range(NSnaps - 1): folder_1 = outputfolder + "/snap_%03d/" % (isnap) folder_2 = outputfolder + "/snap_%03d/" % (isnap + 1) for i in range(len(overdensities)): infile = folder_1 + "/multilevels/" + prefix_template + str( overdensities[i]) + "particle.txt" outfile = folder_2 + "/multilevels/" + prefix_template + str( overdensities[i]) + "particle.txt" mtree_file_fw = folder_1 + "/multilevels/" + prefix_template + str( overdensities[i]) + "_fw" mtree_file_rw = folder_2 + "/multilevels/" + prefix_template + str( overdensities[i]) + "_rw" qsub_com = "qsub " + submission_script_single + " \"%s %d %s %s %s\"" % ( mergertree_exec, 2, infile, outfile, mtree_file_fw) os.system(qsub_com) qsub_com = "qsub " + submission_script_single + " \"%s %d %s %s %s\"" % ( mergertree_exec, 2, outfile, infile, mtree_file_rw) os.system(qsub_com)
def projected_correlation(points, rbins, zmax, box_size, jackknife_nside=0): """ Calculate the projected correlation function wp(rp) and its covariance matrix for a periodic box, with the plane-parallel approximation and the Jackknife method. Parameters ---------- points : array_like Must be a 2-d array whose last dimension is 3 (i.e. has 3 columns) The last column will be used as the redshift distance. rbins : array_like A 1-d array that has the edges of the rp bins. Must be sorted. zmax : float The integral of \pi goes from -zmax to zmax (redshift distance). box_size : float The side length of the periodic box. jackknife_nside : int, optional (Default: 0) If <= 1 , it will not do Jackknife. Returns ------- wp : ndarray A 1-d array that has wp. The length of this retured array would be len(rbins) - 1. wp_cov : ndarray (returned if jackknife_nside > 1) The len(wp) by len(wp) covariance matrix of wp. """ points = np.asarray(points) s = points.shape if len(s) != 2 or s[1] != 3: raise ValueError('`points` must be a 2-d array with last dim=3') N = s[0] rbins = np.asarray(rbins) rbins_sq = rbins * rbins dcorner2 = np.array([rbins[-1], rbins[-1], zmax]) dcorner1 = -dcorner2 if np.any(dcorner2 * 2 > box_size): print "[Warning] box too small!" pairs_rand = float(N*N) / box_size**3 \ * (rbins[1:]**2-rbins[:-1]**2)*np.pi*zmax*2.0 jackknife_nside = int(jackknife_nside) if jackknife_nside <= 1: #no jackknife dcorner1[2] = 0 #save some time pairs = np.zeros(len(rbins) - 1, dtype=int) with fast3tree(points) as tree: for p in points: for pp in _yield_periodic_points(p, dcorner1, dcorner2, box_size): x, y = tree.query_box(pp + dcorner1, pp + dcorner2, output='p').T[:2] x -= pp[0] x *= x y -= pp[1] y *= y x += y x.sort() pairs += np.ediff1d(np.searchsorted(x, rbins_sq)) return (pairs.astype(float) * 2.0 / pairs_rand - 1.0) * zmax * 2.0 else: #do jackknife jack_ids = np.floor(np.remainder(points[:,0], box_size)\ / box_size*jackknife_nside).astype(int) jack_ids += np.floor(np.remainder(points[:,1], box_size)\ / box_size*jackknife_nside).astype(int) * jackknife_nside n_jack = jackknife_nside * jackknife_nside pairs = np.zeros((n_jack, len(rbins) - 1), dtype=int) auto_pairs = np.zeros_like(pairs) with fast3tree(points) as tree: for p, jid in izip(points, jack_ids): for pp in _yield_periodic_points(p, dcorner1, dcorner2, box_size): idx, pos = tree.query_box(pp + dcorner1, pp + dcorner2, output='b') x, y = pos.T[:2] x -= pp[0] x *= x y -= pp[1] y *= y x += y y = x[jack_ids[idx] == jid] y.sort() x.sort() pairs[jid] += np.ediff1d(np.searchsorted(x, rbins_sq)) auto_pairs[jid] += np.ediff1d(np.searchsorted(y, rbins_sq)) idx = pos = x = y = None pairs_sum = pairs.sum(axis=0) pairs = pairs_sum - pairs * 2 + auto_pairs wp_jack = (pairs.astype(float) \ / pairs_rand \ / _jackknife_2d_random(rbins, box_size, jackknife_nside)\ - 1.0) * zmax*2.0 wp_full = (pairs_sum.astype(float) / pairs_rand - 1.0) * zmax * 2.0 wp = wp_full * n_jack - wp_jack.mean(axis=0) * (n_jack - 1) wp_cov = np.cov(wp_jack, rowvar=0, bias=1) * (n_jack - 1) return wp, wp_cov
#!/usr/bin/env python import numpy as np from fast3tree import fast3tree data = np.random.rand(10000, 3) with fast3tree(data) as tree: idx = tree.query_radius([0.5, 0.5, 0.5], 0.2) print("idx:", idx, flush=True)
def matchTrainingSet(self, mag, ranksigma5, redfraction, dm=0.1, ds=0.05): mag_train = self.trainingSet['ABSMAG'][:, 2] if self.match_magonly: ranksigma5_train = np.random.rand(len(self.trainingSet['PSIGMA5'])) else: ranksigma5_train = self.trainingSet['PSIGMA5'] isred_train = self.trainingSet['ISRED'] pos = np.zeros((mag.size, 3)) pos_train = np.zeros((mag_train.size, 3)) n_gal = mag.size rand = np.random.rand(n_gal) pos[:, 0] = mag pos[:, 1] = ranksigma5 pos[:, 2] = redfraction pos[pos[:, 0] < np.min(mag_train), 0] = np.min(mag_train) pos[pos[:, 0] > np.max(mag_train), 0] = np.max(mag_train) pos_train[:, 0] = mag_train pos_train[:, 1] = ranksigma5_train pos_train[:, 2] = isred_train # max distance in isred direction large enough to select all # make search distance in mag direction larger as we go fainter # as there are fewer galaxies in the training set there def max_distances(m): return np.array( [np.min([np.max([np.abs((22.5 + m)) * dm, 0.1]), 5]), ds, 1.1]) def neg_max_distances(m): return -1. * \ np.array( [np.min([np.max([np.abs((22.5 + m)) * dm, 0.1]), 5]), ds, 1.1]) sed_idx = np.zeros(n_gal, dtype=np.int) bad = np.zeros(n_gal, dtype=np.bool) with fast3tree(pos_train) as tree: for i, p in enumerate(pos): idx, tpos = tree.query_box(p + neg_max_distances(p[0]), p + max_distances(p[0]), output='both') rf = np.sum(tpos[:, 2]) / len(tpos) isred = rand[i] < (rf * redfraction[i]) idx = idx[tpos[:, 2] == int(isred)] tpos = tpos[tpos[:, 2] == int(isred)] tpos -= p if self.match_magonly: dt = np.abs(tpos[:, 0]) else: dt = np.abs(np.sum(tpos**2, axis=1)) try: sed_idx[i] = idx[np.argmin(dt)] except Exception as e: bad[i] = True isbad = np.where(bad)[0] print('Number of bad SED assignments: {}'.format(len(isbad))) if self.match_magonly: def max_distances(m): return np.array( [np.max([(22.5 + m)**2 * dm, dm]), ds * 10, 1.1]) def neg_max_distances(m): return -1. * \ np.array([np.max([(22.5 + m)**2 * dm, dm]), ds * 10, 1.1]) else: def max_distances(m): return np.array( [np.max([10 * (22.5 + m)**2 * dm, dm]), ds, 1.1]) def neg_max_distances(m): return -1. * \ np.array([np.max([10 * (22.5 + m)**2 * dm, dm]), ds, 1.1]) for i, p in enumerate(pos[bad]): idx, tpos = tree.query_box(p + neg_max_distances(p[0]), p + max_distances(p[0]), output='both') rf = np.sum(tpos[:, 2]) / len(tpos) isred = rand[i] < (rf * redfraction[i]) idx = idx[tpos[:, 2] == int(isred)] tpos = tpos[tpos[:, 2] == int(isred)] tpos -= p if self.match_magonly: dt = np.abs(tpos[:, 0]) else: dt = np.abs(np.sum(tpos**2, axis=1)) try: sed_idx[isbad[i]] = idx[np.argmin(dt)] except Exception as e: bad[i] = True return sed_idx, bad