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