def get_interf_res(atomlist, reftree, model):
if len(atomlist) == 0 :
return([], [], [])
array = np.array([ a.get_coord() for a in atomlist])
tree = KDTree(array)
#list of list of interfacing nucl-atom per aa-atom
interf = tree.query_ball_tree(reftree, args.cutoff)
if max([len(i) for i in interf]) == 0:
return([], [], [])
# mask on aa-atoms: interface / non-interface
mask = [len(j) > 0 for j in interf ]
interf = array[np.array(mask)]
# list of interfacing aa-atoms
interf_at = set([id(a) for anr,a in enumerate(atomlist) if mask[anr]])
interf_res = []
for residue in model.get_residues():
for atom in residue:
if id(atom) in interf_at:
interf_res.append(residue)
break
interftree = KDTree(interf)
nainterf = natree.sparse_distance_matrix(interftree, args.cutoff)
A = nainterf.toarray()
A[A==0] = 99
mininterf = np.min(A, axis=1)
return(mininterf, interf, interf_res)
def get_N_ngbs(positions, radii, N=12, maxdist=3.0, edge = None):
"""N first neighbours, with a maximum relative distances, such that $r_{ij} < maxdist (R_i + R_j)$.
If a potential neighbour is further away than the distance to the edge of the field of view,
the current particle of interest is considered as "on the edge" and the neighbour not taken into account.
Returns neighbours, inside"""
assert len(positions)==len(radii)
if edge is None:
edge = (positions.min(0), positions.max(0))
#initialize the geometry of each particle
to_edge = np.minimum((positions - edge[0]).min(-1), (edge[0] - positions).min(-1))**2
inside = np.full(len(positions), True, dtype=bool)
neighbours = np.full([len(positions), N], -1, dtype=int)
tree = KDTree(positions)
rmax = radii.max()
for i, js in enumerate(tree.query_ball_tree(tree, 2*rmax*maxdist)):
disq = np.sum((positions[js] - positions[i])**2, -1)
ags = np.argsort(disq)[:N]
if disq[ags[-1]] < to_edge[i]:
neighbours[i, :len(js)] = np.array(js)[ags]
else:
inside[i] = False
N2 = np.where(disq[ags] < to_edge[i])[0][0]+1
neighbours[i, :N2] = np.array(js)[ags[:N2]]
return neighbours, inside
def find_pairs(cutoff, X, Y=None):
"""
Find pairs with euclidean distance below C{cutoff}. Either between
C{X} and C{Y}, or within C{X} if C{Y} is C{None}.
Uses a KDTree and thus is memory efficient and reasonable fast.
@type cutoff: float
@type X: (m,n) numpy.array
@type Y: (k,n) numpy.array
@return: set of index tuples
@rtype: iterable
"""
try:
from scipy.spatial import cKDTree as KDTree
KDTree.query_pairs
KDTree.query_ball_tree
except (ImportError, AttributeError):
from scipy.spatial import KDTree
tree = KDTree(X, len(X))
if Y is None:
return tree.query_pairs(cutoff)
other = KDTree(Y, len(Y))
contacts = tree.query_ball_tree(other, cutoff)
return ((i, j) for (i, js) in enumerate(contacts) for j in js)
def binder(positions, orientations, bl, m=4, method='ball', margin=0):
"""Calculate the binder cumulant, given positions and orientations.
bl: the binder length scale, such that
B(bl) = 1 - .333 * S4 / S2^2
where SN are <phibl^N> averaged over each block/cluster of size bl in frame.
"""
if margin:
if margin < ss:
margin *= ss
center = 0.5*(positions.max(0) + positions.min(0))
dmask = d < d.max() - margin
positions = positions[dmask]
orientations = orientations[dmask]
if 'neigh' in method or 'ball' in method:
tree = KDTree(positions)
balls = tree.query_ball_tree(tree, bl)
balls, ball_mask = helpy.pad_uneven(balls, 0, True, int)
ball_orient = orientations[balls]
ball_orient[ball_mask] = np.nan
phis = np.nanmean(np.exp(m*ball_orient*1j), 1)
phi2 = np.dot(phis, phis) / len(phis)
phiphi = phis*phis
phi4 = np.dot(phiphi, phiphi) / len(phiphi)
return 1 - phi4 / (3*phi2*phi2)
else: # elif method=='block':
raise ValueError("method {} not implemented".format(method))
def get_N_ngbs(positions, radii, N=12, maxdist=3.0, edge = None):
"""N first neighbours, with a maximum relative distances, such that $r_{ij} < maxdist (R_i + R_j)$.
If a potential neighbour is further away than the distance to the edge of the field of view,
the current particle of interest is considered as "on the edge" and the neighbour not taken into account.
Returns neighbours, inside"""
assert len(positions)==len(radii)
if edge is None:
edge = (positions.min(0), positions.max(0))
#initialize the geometry of each particle
to_edge = np.minimum((positions - edge[0]).min(-1), (edge[0] - positions).min(-1))**2
inside = np.full(len(positions), True, dtype=bool)
neighbours = np.full([len(positions), N], -1, dtype=int)
tree = KDTree(positions)
rmax = radii.max()
for i, js in enumerate(tree.query_ball_tree(tree, 2*rmax*maxdist)):
disq = np.sum((positions[js] - positions[i])**2, -1)
ags = np.argsort(disq)[:N]
if disq[ags[-1]] < to_edge[i]:
neighbours[i, :len(js)] = np.array(js)[ags]
else:
inside[i] = False
N2 = np.where(disq[ags] < to_edge[i])[0][0]+1
neighbours[i, :N2] = np.array(js)[ags[:N2]]
return neighbours, inside
def binder(positions, orientations, bl, m=4, method='ball', margin=0):
"""Calculate the binder cumulant, given positions and orientations.
bl: the binder length scale, such that
B(bl) = 1 - .333 * S4 / S2^2
where SN are <phibl^N> averaged over each block/cluster of size bl in frame.
"""
if margin:
if margin < ss:
margin *= ss
center = 0.5 * (positions.max(0) + positions.min(0))
dmask = d < d.max() - margin
positions = positions[dmask]
orientations = orientations[dmask]
if 'neigh' in method or 'ball' in method:
tree = KDTree(positions)
balls = tree.query_ball_tree(tree, bl)
balls, ball_mask = helpy.pad_uneven(balls, 0, True, int)
ball_orient = orientations[balls]
ball_orient[ball_mask] = np.nan
phis = np.nanmean(np.exp(m * ball_orient * 1j), 1)
phi2 = np.dot(phis, phis) / len(phis)
phiphi = phis * phis
phi4 = np.dot(phiphi, phiphi) / len(phiphi)
return 1 - phi4 / (3 * phi2 * phi2)
else: # elif method=='block':
raise ValueError("method {} not implemented".format(method))
def find_pairs(cutoff, X, Y=None):
"""
Find pairs with euclidean distance below C{cutoff}. Either between
C{X} and C{Y}, or within C{X} if C{Y} is C{None}.
Uses a KDTree and thus is memory efficient and reasonable fast.
@type cutoff: float
@type X: (m,n) numpy.array
@type Y: (k,n) numpy.array
@return: set of index tuples
@rtype: iterable
"""
try:
from scipy.spatial import cKDTree as KDTree
KDTree.query_pairs
KDTree.query_ball_tree
except (ImportError, AttributeError):
from scipy.spatial import KDTree
tree = KDTree(X, len(X))
if Y is None:
return tree.query_pairs(cutoff)
other = KDTree(Y, len(Y))
contacts = tree.query_ball_tree(other, cutoff)
return ((i, j) for (i, js) in enumerate(contacts) for j in js)
def spatialCorelation(points,
fields,
vectorColumns=None,
Nbins=200,
maxDist=50.0):
"""Compute the spatial corellation of each field
points -- 2D array of points coordinates. Shape is (N,d) with d the number of spatial dimensions.
fields -- 2D array of scalar field or of coordinates of vector fields. Shape is (N, F)
with F the sum of the dimensions of each field.
vectorColumns -- 1D array indexing the columns of fields into vector fields.
for example [0, 1, 1, 1] means that the first column of fields is the scalar field 0 and
the next 3 columns are the coordinates of a 3D vector field.
Nbins -- The number of bins of the histogram
maxLength -- The maximum distance between a pair of points taken into account in the histogram
"""
#parameters parsing
if len(points) != len(fields):
raise ValueError('You must have exactly one field value per point\n' +
'Here points id %i and fieds is %i' %
(len(points), len(fields)))
if vectorColumns == None:
vectorColumns = np.arange(fields.shape[1])
if len(vectorColumns) != fields.shape[1]:
vectorColumns = np.concatenate(
(vectorColumns, np.arange(vectorColumns.max() + 1,
fields.shape[1])))
slices = [
np.where(vectorColumns == v)[0] for v in range(vectorColumns.max() + 1)
]
#spatial query
lowerBound = points.min(axis=0) + maxDist / 2
upperBound = points.max(axis=0) - maxDist / 2
inside_id = [
i for i, p in enumerate(points)
if (p >= lowerBound).all() and (p <= upperBound).all()
]
tree = KDTree(points)
inside_tree = KDTree(points[inside_id])
pairs = inside_tree.query_ball_tree(tree, maxDist)
#binning
coord_bins = np.zeros((Nbins, fields.shape[1]))
nb_bins = np.zeros((Nbins), dtype=int)
for p, qs in zip(inside_id, pairs):
qs.remove(p)
rs = np.asarray(np.sqrt(
((points[qs] - points[p])**2).sum(axis=1)) * Nbins / maxDist,
dtype=int)
nb_bins[rs] += 1
coord_bins[rs] += fields[qs] * fields[p]
bins = np.column_stack(
[coord_bins[:, cols].sum(axis=1) for cols in slices])
bins[np.nonzero(nb_bins)] /= nb_bins[np.nonzero(nb_bins)][:, np.newaxis]
return np.column_stack((np.arange(Nbins, dtype=float) / maxDist, bins))
def test_1D_regular():
grid = RegularGrid([0], [1], [3])
pos = np.arange(-1, 4, 0.1)[:, None]
tree = KDTree(pos)
edges = np.array([[i, j]
for i, js in enumerate(tree.query_ball_tree(tree, 0.15))
for j in sorted(js) if i < j])
sumw, count = bin_texture(pos, edges, grid)
assert_array_almost_equal(sumw, np.full((3, 1), 0.3))
assert_array_equal(count, np.full((3), 30))
assert_array_almost_equal(sumw / count[:, None], np.full((3, 1), 0.01))
def match(x1, y1, x2, y2, radius):
"""
"""
# Build the trees
p1 = np.array([*zip(x1, y1)])
t1 = KDTree(p1)
p2 = np.array([*zip(x2, y2)])
t2 = KDTree(p2)
# Do the matching
m = t1.query_ball_tree(t2, radius)
return m
def _sparse_mdm(self, big_tree: cKDTree, small_tree: cKDTree, iter_over: str) -> sparse.csr_matrix:
neighbours = small_tree.query_ball_tree(big_tree, self.max_distance, p=2.0, eps=0)
with job.Parallel(backend=self.joblib_backend, n_jobs=self.n_jobs, verbose=self.verbose) as parallel:
results = parallel(job.delayed(self._apply_function)
(small_tree.data[i], big_tree.data[neighbours[i]], i, neighbours[i], iter_over)
for i in range(small_tree.data.shape[0]))
coo_data = np.concatenate([results[i][0] for i in range(len(results))])
coo_i = np.concatenate([results[i][1] for i in range(len(results))])
coo_j = np.concatenate([results[i][2] for i in range(len(results))])
sparse_mdm = sparse.coo_matrix((coo_data, (coo_i, coo_j)),
shape=(self.samples1.shape[0], self.samples2.shape[0]),
dtype=self.dtype).tocsr()
return sparse_mdm
def spatialCorelation(points, fields, vectorColumns=None, Nbins=200, maxDist=50.0):
"""Compute the spatial corellation of each field
points -- 2D array of points coordinates. Shape is (N,d) with d the number of spatial dimensions.
fields -- 2D array of scalar field or of coordinates of vector fields. Shape is (N, F)
with F the sum of the dimensions of each field.
vectorColumns -- 1D array indexing the columns of fields into vector fields.
for example [0, 1, 1, 1] means that the first column of fields is the scalar field 0 and
the next 3 columns are the coordinates of a 3D vector field.
Nbins -- The number of bins of the histogram
maxLength -- The maximum distance between a pair of points taken into account in the histogram
"""
#parameters parsing
if len(points) != len(fields):
raise ValueError(
'You must have exactly one field value per point\n'
+ 'Here points id %i and fieds is %i'%(len(points), len(fields))
)
if vectorColumns==None:
vectorColumns = np.arange(fields.shape[1])
if len(vectorColumns) != fields.shape[1]:
vectorColumns = np.concatenate((
vectorColumns,
np.arange(vectorColumns.max()+1,fields.shape[1])
))
slices = [np.where(vectorColumns==v)[0] for v in range(vectorColumns.max()+1)]
#spatial query
lowerBound = points.min(axis=0) + maxDist/2
upperBound = points.max(axis=0) - maxDist/2
inside_id = [
i for i, p in enumerate(points)
if (p>= lowerBound).all() and (p <= upperBound).all()
]
tree = KDTree(points)
inside_tree = KDTree(points[inside_id])
pairs = inside_tree.query_ball_tree(tree, maxDist)
#binning
coord_bins = np.zeros((Nbins, fields.shape[1]))
nb_bins = np.zeros((Nbins), dtype=int)
for p, qs in zip(inside_id, pairs):
qs.remove(p)
rs = np.asarray(
np.sqrt(
((points[qs] - points[p])**2).sum(axis=1)
) * Nbins / maxDist,
dtype=int)
nb_bins[rs] += 1
coord_bins[rs] += fields[qs]*fields[p]
bins = np.column_stack([coord_bins[:,cols].sum(axis=1) for cols in slices])
bins[np.nonzero(nb_bins)] /= nb_bins[np.nonzero(nb_bins)][:,np.newaxis]
return np.column_stack((np.arange(Nbins, dtype=float)/maxDist,bins))
def remove_close_set(points_fixed, points_reduce, radius):
'''
Given two sets of points and a radius, return a set of points
that is the subset of points_reduce where no point is within
radius of any point in points_fixed
'''
tree_fixed = KDTree(points_fixed)
tree_reduce = KDTree(points_reduce)
reduce_duplicates = tree_fixed.query_ball_tree(tree_reduce, r = radius)
reduce_duplicates = np.unique(np.hstack(reduce_duplicates).astype(int))
reduce_mask = np.ones(len(points_reduce), dtype=np.bool)
reduce_mask[reduce_duplicates] = False
points_clean = points_reduce[reduce_mask]
return points_clean
def runMaxFlowMinCost(data, d_th, k1, tagList, n_threads, dth_max, prior):
print("Generate Graph Model..." + " (" +
datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + ") " + "\n")
num_hyb = np.arange(1, int(np.amax(data.hyb)) + 1)
data.sort_values('hyb', inplace=True)
data = data.reset_index(drop=True)
## Graphical Model Data Structures
print("[INFO] Add D vars" + " (" +
datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + ") ")
G = nx.Graph()
G.add_nodes_from(data.index.values)
print("[INFO] Add T vars" + " (" +
datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + ") ")
for h1 in tqdm(num_hyb):
KDTree_h1 = KDTree(data[data.hyb == h1][['x', 'y', 'z']])
for h2 in num_hyb[h1:]:
if h1 != h2:
KDTree_h2 = KDTree(data[data.hyb == h2][['x', 'y', 'z']])
query = KDTree_h1.query_ball_tree(KDTree_h2, d_th, p=2)
E = []
offset1 = data.index[data.hyb == h1].min()
offset2 = data.index[data.hyb == h2].min()
for i, e1 in enumerate(query):
if e1:
for e2 in e1:
E.append((i + offset1, e2 + offset2))
G.add_edges_from(E)
print("[INFO] Generate model" + " (" +
datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + ") ")
conn_comps = [list(c) for c in nx.connected_components(G)]
for c in tqdm(range(len(conn_comps))):
data.loc[conn_comps[c], 'conn_comp'] = c
# Drop conn components with less than n_hybs elements
gr = data.groupby('conn_comp')
for i, group in gr:
if len(group) < len(num_hyb):
data = data.drop(group.index)
labels = np.unique(data.conn_comp)
if labels.size > 0:
res = []
for l in tqdm(np.nditer(labels), total=len(labels)):
res.append(
runComponent(data, int(l), d_th, k1, tagList, num_hyb, dth_max,
prior))
#return maxFlowMinCost
return [x for x in res if x is not None]
def remove_close_set(points_fixed, points_reduce, radius):
'''
Given two sets of points and a radius, return a set of points
that is the subset of points_reduce where no point is within
radius of any point in points_fixed
'''
from scipy.spatial import cKDTree as KDTree
tree_fixed = KDTree(points_fixed)
tree_reduce = KDTree(points_reduce)
reduce_duplicates = tree_fixed.query_ball_tree(tree_reduce, r=radius)
reduce_duplicates = np.unique(np.hstack(reduce_duplicates).astype(int))
reduce_mask = np.ones(len(points_reduce), dtype=np.bool)
reduce_mask[reduce_duplicates] = False
points_clean = points_reduce[reduce_mask]
return points_clean
def remove_ps1_source(ps1_cat, ns_cat):
ps1 = fits.open(ps1_cat)[1].data
ps1_ra = ps1['f0']
ps1_dec = ps1['f1']
ra = ns_cat['X_WORLD']
dec = ns_cat['Y_WORLD']
ns_tree = KDTree(zip(ra, dec))
ps1_tree = KDTree(zip(ps1_ra, ps1_dec))
match = ns_tree.query_ball_tree(ps1_tree, r=.5 / 3600.)
#print len(match), len(ra), len(ps1)
ps1_match = []
for i in match:
if i != []:
ps1_match.append(True)
else:
ps1_match.append(False)
return array(ps1_match)
def get_bonds(positions, radii, maxdist=3.0):
"""Bonds by relative distances, such that $r_{ij} < maxdist (R_i + R_j)$.
Returns pairs, distances. Pairs are sorted and unique."""
assert len(positions)==len(radii)
tree = KDTree(positions)
rmax = radii.max()
#fetch all potential pairs, already sorted
pairs = np.array([
[i,j]
for i, js in enumerate(tree.query_ball_tree(tree, 2*rmax*maxdist))
for j in sorted(js)
if i<j
])
#compute all pair's square distances via numpy
dists = np.sum((positions[pairs[:,0]] - positions[pairs[:,1]])**2, -1)
#filter out the pairs that are too far
good = dists < maxdist**2 * radii[pairs].sum(-1)**2
return pairs[good], np.sqrt(dists[good])
def get_bonds(positions, radii, maxdist=3.0):
"""Bonds by relative distances, such that $r_{ij} < maxdist (R_i + R_j)$.
Returns pairs, distances. Pairs are sorted and unique."""
assert len(positions)==len(radii)
tree = KDTree(positions)
rmax = radii.max()
#fetch all potential pairs, already sorted
pairs = np.array([
[i,j]
for i, js in enumerate(tree.query_ball_tree(tree, 2*rmax*maxdist))
for j in sorted(js)
if i<j
])
if len(pairs) == 0:
return np.zeros((0,2), int), np.zeros(0)
#compute all pair's square distances via numpy
dists = np.sum((positions[pairs[:,0]] - positions[pairs[:,1]])**2, -1)
#filter out the pairs that are too far
good = dists < maxdist**2 * radii[pairs].sum(-1)**2
return pairs[good], np.sqrt(dists[good])
class kdtree():
def __init__(self, nodes, get_latlon=id, transformer=None):
if transformer:
self.transformer = transformer
else:
from pyproj import CRS, Transformer
self.transformer = Transformer.from_crs(CRS("WGS84"),
CRS("EPSG:28992"))
if nodes:
from scipy.spatial import cKDTree as KDTree
self.nodes = list(nodes)
self.tree = KDTree(
list(
map(lambda x: self.transformer.transform(*get_latlon(x)),
self.nodes)))
else:
self.nodes = self.tree = None
self.get_latlon = get_latlon
def query(self, node, around, get_latlon=lambda x: x):
if not self.tree:
return []
latlon = get_latlon(node)
projected = self.transformer.transform(*latlon)
indices = self.tree.query_ball_point(projected, r=around)
results = list(map(lambda x: self.nodes[x], indices))
return results
def get_neighbours(self, around):
if not self.tree:
return []
indicess = self.tree.query_ball_tree(self.tree, r=around)
results = list(
map(lambda indices: list(map(lambda x: self.nodes[x], indices)),
indicess))
return indicess, results
def keypoints_match_geometry(src_frame, dest_frame):
"""
Finds keypoint matches based on scene geometry
@return pairs, angles
`pairs[n] = (src_pt_n, dest_pt_n)`
`angles[n]` = angle btw rays of src_pt_n and dest_pt_n
"""
if hasattr(src_frame, 'kpt_proj_cloud'):
helper_cloud = src_frame.kpt_proj_cloud
else:
helper_cloud = build_point_cloud_for_projection(src_frame)
# spatial and perspective projection src_frame -> dest_frame
spatial_mat = dest_frame.world_to_camera @ src_frame.camera_to_world
full_projection_mat = dest_frame.intrinsic_mat @ spatial_mat[:3, :]
# project and retrieve keypoints
helper_projected = projection_apply_rowvec(full_projection_mat,
helper_cloud)
proj_pts, proj_sizes, proj_orientations = kptproj_interpret_projected_vectors(
helper_projected)
# find pairs of neighbours in radius of MATCH_DISTANCE
tree_proj = KDTree(proj_pts)
tree_dest = KDTree(dest_frame.kpt_locs)
match_suggestions = tree_dest.query_ball_tree(tree_proj, r=MATCH_DISTANCE)
matches = []
# kpt_matched_id[n] = id of point in src_frame which matches n
#frame.kpt_matched_id = np.zeros(frame.pt_count, dtype=np.int32)
#frame.kpt_matched_id[:] = -1 # no match
#print(src_frame.pt_count, dest_frame.pt_count, len(match_suggestions))
for dest_pt_idx, suggestions in enumerate(match_suggestions):
dest_pt_size = dest_frame.kpt_sizes[dest_pt_idx]
for src_pt_idx in suggestions:
proj_size = proj_sizes[src_pt_idx]
if ((max(dest_pt_size / proj_size, proj_size / dest_pt_size) <
MATCH_SIZE_DIFF_RELATIVE) and (angular_distance_abs(
dest_frame.kpt_orientations[dest_pt_idx],
proj_orientations[src_pt_idx]) < MATCH_ANGLE_DIFF)):
matches.append((src_pt_idx, dest_pt_idx))
#frame.kpt_matched_id[pt_idx] = src_pt_idx
break
if len(matches) == 0:
return AttrDict(
pairs=np.zeros((0, 2), dtype=np.int32),
angles=np.zeros(0, dtype=np.float32),
)
else:
# store matched pairs
match_pairs = np.array(matches, dtype=np.int32)
# sort by src_point_id
match_pairs = match_pairs[np.argsort(match_pairs[:, 0]), :]
view_angle_changes = derive_keypoint_view_angle_change(
src_frame, dest_frame, match_pairs)
return AttrDict(
pairs=match_pairs,
angles=view_angle_changes,
)
if result: break
nstruc += 1
coor = collectlib.collect_all_coor()
pdbsizes2 = np.cumsum([0] + pdbsizes)
coors = [coor[pdbsizes2[n]:pdbsizes2[n+1]] for n in range(len(pdbs))]
for n1 in range(len(pdbs)):
c1 = coors[n1]
tree1 = KDTree(c1)
resids1 = resids[n1]
rescounts1 = rescounts[n1]
for n2 in range(n1+1, len(pdbs)):
c2 = coors[n2]
tree2 = KDTree(c2)
resids2 = resids[n2]
rescounts2 = rescounts[n2]
pairs = tree1.query_ball_tree(tree2, cutoff)
for p1,pp in enumerate(pairs):
if not len(pp): continue
resid1 = resids1[p1]
if resid1 not in rescounts1:
rescounts1[resid1] = 0
rescounts1[resid1] += len(pp)
for p2 in pp:
resid2 = resids2[p2]
if resid2 not in rescounts2:
rescounts2[resid2] = 0
rescounts2[resid2] += 1
avg_contacts = [sum(c.values())/float(nstruc) for c in rescounts]
#from pprint import pprint
#pprint(rescounts)
#print(sum(rescounts[0].values()), sum(rescounts[1].values()), avg_contacts)
dtype=np.float64))
tid = np.array(tycho2["tycho2_id"])
print("Loading the Kepler catalog...")
kepler = load_kepler()
kepler_xyz = convert_to_cartesian(np.array(kepler[["ra", "dec"]],
dtype=np.float64))
# Building KD-tree.
print("Building KD trees...")
tycho2_tree = KDTree(tycho2_xyz)
kepler_tree = KDTree(kepler_xyz)
# Cross match.
print("Cross matching trees...")
match = kepler_tree.query_ball_tree(tycho2_tree, np.sqrt(2-2*np.cos(tol)))
match_flag = np.zeros(len(kepler), dtype=bool)
match_id = np.zeros(len(kepler), dtype=np.uint64)
distances = np.nan + np.zeros(len(kepler))
tycho_id = []
for i, m in enumerate(match):
if len(m):
# Compute the angular distance.
d = np.arccos(np.dot(kepler_xyz[i], tycho2_xyz[m].T))
distances[i] = np.min(d)
if distances[i] <= tol:
match_id[i] = m[np.argmin(d)]
match_flag[i] = True
print(len(match_id))
TGAS_PATH = "/export/bbq2/angusr/gaia/Tgas.csv"
def test_2D_3branches():
grid = RegularGrid([-1.2, -1.2], [0.8, 0.8], [3, 3])
pos0 = np.array([[0, 0], [-1, 0], [0, -1], [0, 1]], dtype=float)
tree = KDTree(pos0)
edges = np.array([[i, j]
for i, js in enumerate(tree.query_ball_tree(tree, 1.5))
for j in sorted(js) if i < j])
sumw0, count0 = bin_texture(pos0, edges, grid)
assert_array_equal(count0, [[1, 4, 1], [3, 3, 3], [0, 0, 0]])
assert_array_almost_equal(sumw0, [[[1, -1, 1], [4, 0, 2], [1, 1, 1]],
[[1, -1, 3], [1, 0, 2], [1, 1, 3]],
[[0, 0, 0], [0, 0, 0], [0, 0, 0]]])
#move 1 point within the same grid element
pos1 = np.copy(pos0)
pos1[-1] += [0, 0.1]
sumw1, count1 = bin_texture(pos1, edges, grid)
assert_array_equal(count0, count1)
assert_array_almost_equal(sumw1,
[[[1, -1, 1], [4, 0.1, 2.21], [1, 1.1, 1.21]],
[[1, -1, 3], [1, 0, 2.21], [1, 1.1, 3.63]],
[[0, 0, 0], [0, 0, 0], [0, 0, 0]]])
sumw, count = bin_geometrical_changes(pos0, pos1, edges, grid)
assert_array_equal(count0, count)
sumB = C2B(sumw)
assert_array_almost_equal(sumB, sumw1 - sumw0)
## there should be no topology change
sumC, countc, sumT, counta, countd = bin_changes(pos0, pos1, edges, edges,
grid)
assert_array_equal(count, countc)
assert_array_almost_equal(sumw, sumC)
assert_array_equal(counta, np.zeros_like(count))
assert_array_equal(countd, np.zeros_like(count))
assert_array_almost_equal(sumT, np.zeros_like(sumB))
#flip the motif along x
pos1 = np.array([[0, 0], [1, 0], [0, -1], [0, 1]], dtype=float)
sumw1, count1 = bin_texture(pos1, edges, grid)
assert_array_equal(count0, [[1, 4, 1], [3, 3, 3], [0, 0, 0]])
assert_array_almost_equal(
sumw1,
[[[0, 0, 0], [0, 0, 0], [0, 0, 0]], [[1, 1, 3], [1, 0, 2], [1, -1, 3]],
[[1, 1, 1], [4, 0, 2], [1, -1, 1]]])
sumw, count = bin_geometrical_changes(pos0, pos1, edges, grid)
assert_array_equal(count, [[0, 0, 0], [3, 3, 3], [0, 0, 0]])
sumB = C2B(sumw)
assert_array_almost_equal(
sumB,
[[[0, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 2, 0], [0, 0, 0], [0, -2, 0]],
[[0, 0, 0], [0, 0, 0], [0, 0, 0]]])
## there should be no topology change
sumC, countc, sumT, counta, countd = bin_changes(pos0, pos1, edges, edges,
grid)
assert_array_equal(count, countc)
assert_array_almost_equal(sumw, sumC)
assert_array_equal(counta, np.zeros_like(count))
assert_array_equal(countd, np.zeros_like(count))
assert_array_almost_equal(sumT, np.zeros_like(sumB))
#rotate the motif by 90°
pos1 = np.array([[0, 0], [0, -1], [1, 0], [-1, 0]], dtype=float)
sumw1, count1 = bin_texture(pos1, edges, grid)
assert_array_equal(count1, count0.T)
assert_array_almost_equal(
sumw1,
[[[1, -1, 1], [3, -1, 1], [0, 0, 0]],
[[2, 0, 4], [2, 0, 1], [0, 0, 0]], [[1, 1, 1], [3, 1, 1], [0, 0, 0]]])
sumw, count = bin_geometrical_changes(pos0, pos1, edges, grid)
assert_array_equal(count, [[0, 0, 0], [0, 3, 0], [0, 0, 0]])
sumB = C2B(sumw)
assert_array_almost_equal(
sumB,
[[[0, 0, 0], [0, 0, 0], [0, 0, 0]], [[0, 0, 0], [1, 0, -1], [0, 0, 0]],
[[0, 0, 0], [0, 0, 0], [0, 0, 0]]])
## there should be no topology change
sumC, countc, sumT, counta, countd = bin_changes(pos0, pos1, edges, edges,
grid)
assert_array_equal(count, countc)
assert_array_almost_equal(sumw, sumC)
assert_array_equal(counta, np.zeros_like(count))
assert_array_equal(countd, np.zeros_like(count))
assert_array_almost_equal(sumT, np.zeros_like(sumB))
#change topology only
pos1 = np.copy(pos0)
edges1 = np.array([[i, j]
for i, js in enumerate(tree.query_ball_tree(tree, 1.1))
for j in sorted(js) if i < j] + [[2, 3]])
sumB, countc, sumT, counta, countd = bin_changes(pos0, pos1, edges, edges1,
grid)
assert_array_equal(countc, [[0, 2, 0], [2, 3, 2], [0, 0, 0]])
assert_array_almost_equal(sumB, np.zeros_like(sumw))
assert_array_equal(counta, [[0, 0, 0], [1, 1, 1], [0, 0, 0]])
assert_array_equal(countd, [[1, 2, 1], [1, 0, 1], [0, 0, 0]])
assert_array_almost_equal(sumT, [[[-1, 1, -1], [-2, 0, -2], [-1, -1, -1]],
[[-1, 1, 3], [0, 0, 4], [-1, -1, 3]],
[[0, 0, 0], [0, 0, 0], [0, 0, 0]]])
def test_1D_notregular():
grid = RegularGrid([0], [1], [3])
pos0 = np.array([0.1, 0.5, 0.9, 1.5, 1.9, 2.2, 2.3])[:, None]
tree = KDTree(pos0)
edges = np.array([[i, j]
for i, js in enumerate(tree.query_ball_tree(tree, 0.5))
for j in sorted(js) if i < j])
assert_array_equal(
edges, np.array([[0, 1], [1, 2], [3, 4], [4, 5], [4, 6], [5, 6]]))
sumw0, count0 = bin_texture(pos0, edges, grid)
assert_array_equal(count0, [6, 5, 7])
assert_array_almost_equal(
sumw0,
np.array([
6 * 0.16, 3 * 0.16 + 1 * 0.09 + 1 * 0.16,
2 * 0.09 + 2 * 0.16 + 3 * 0.01
])[:, None])
#move 2 points within the same grid element
pos1 = np.copy(pos0)
pos1[:2] += 0.2
sumw1, count1 = bin_texture(pos1, edges, grid)
assert_array_equal(count0, count1)
assert_array_almost_equal(
sumw1,
np.array([
3 * 0.16 + 3 * 0.04, 3 * 0.16 + 1 * 0.09 + 1 * 0.16,
2 * 0.09 + 2 * 0.16 + 3 * 0.01
])[:, None])
sumw, count = bin_geometrical_changes(pos0, pos1, edges, grid)
assert_array_equal(count0, count)
sumB = C2B(sumw)
assert_array_almost_equal(
sumB,
np.array([2 * 3 * 0.3 * (0.2 - 0.4), 0, 0])[:, None])
## there should be no topology change
sumC, countc, sumT, counta, countd = bin_changes(pos0, pos1, edges, edges,
grid)
assert_array_equal(count, countc)
assert_array_almost_equal(sumw, sumC)
assert_array_equal(counta, np.zeros_like(count))
assert_array_equal(countd, np.zeros_like(count))
assert_array_almost_equal(sumT, np.zeros_like(sumB))
#move one point to a different grid element, without straining any existing bond
pos1 = np.copy(pos0)
pos1[:3] += 0.2
sumw1, count1 = bin_texture(pos1, edges, grid)
assert_array_equal(count1, [5, 6, 7])
assert_array_almost_equal(
sumw1,
np.array([
5 * 0.16, 4 * 0.16 + 1 * 0.09 + 1 * 0.16,
2 * 0.09 + 2 * 0.16 + 3 * 0.01
])[:, None])
sumw, count = bin_geometrical_changes(pos0, pos1, edges, grid)
assert_array_equal(count, [5, 5, 7])
sumB = C2B(sumw)
assert_array_almost_equal(sumB, np.array([0, 0, 0])[:, None])
## there should be no topology change
sumC, countc, sumT, counta, countd = bin_changes(pos0, pos1, edges, edges,
grid)
assert_array_equal(count, countc)
assert_array_almost_equal(sumw, sumC)
assert_array_equal(counta, np.zeros_like(count))
assert_array_equal(countd, np.zeros_like(count))
assert_array_almost_equal(sumT, np.zeros_like(sumB))
##update topology
tree1 = KDTree(pos1)
edges1 = np.array(
[[i, j] for i, js in enumerate(tree1.query_ball_tree(tree1, 0.5))
for j in sorted(js) if i < j])
assert_array_equal(
edges1,
np.array([[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [4, 6], [5, 6]]))
sumC, countc, sumT, counta, countd = bin_changes(pos0, pos1, edges, edges1,
grid)
assert_array_equal(countc, np.array([5, 5, 7]))
assert_array_almost_equal(sumw, sumC)
assert_array_equal(counta, np.array([0, 3, 0]))
assert_array_equal(countd, np.zeros_like(count))
assert_array_almost_equal(sumT, np.array([0, 3 * 0.16, 0])[:, None])
def runComponent(data, l, d_th, k1, tagList, num_hyb, dth_max, prior):
if len(data[data.conn_comp == l]):
if len(np.unique(data[data.conn_comp == l].hyb)) == len(num_hyb):
data_tmp = data[data.conn_comp == l].sort_values(['hyb']).copy()
data_tmp.reset_index(inplace=True, drop=True)
Dvar_tmp = pd.DataFrame(data={
'x': [],
'y': [],
'z': [],
'hyb': [],
'ch': [],
'X_idx': []
})
for h in num_hyb:
for i, row in data_tmp[data_tmp.hyb == h].iterrows():
# add detection variables
Dvar_tmp = Dvar_tmp.append(pd.Series([
row.x, row.y, row.z, h, row.ch, i,
prob2Eng(row.p0),
prob2Eng(row.p1)
],
index=[
'x', 'y', 'z',
'hyb', 'ch',
'X_idx', 'E_0',
'E_1'
]),
ignore_index=True)
X_idx_tmp = len(Dvar_tmp)
Tvar_tmp = pd.DataFrame(
data={
'x_idx': [],
'anchestor_x_idx': [],
'descendant_x_idx': [],
'E_0': [],
'E_1': []
})
for h1 in num_hyb[:-1]:
h2 = h1 + 1
if len(data_tmp[data_tmp.hyb == h1]) == 0:
continue
KDTree_h1 = KDTree(
data_tmp[data_tmp.hyb == h1][['x', 'y', 'z']])
if len(data_tmp[data_tmp.hyb == h2]) == 0:
continue
KDTree_h2 = KDTree(
data_tmp[data_tmp.hyb == h2][['x', 'y', 'z']])
query = KDTree_h1.query_ball_tree(KDTree_h2, dth_max, p=2)
for i in range(len(query)):
if len(query[i]):
data_i_idx = data_tmp.index[data_tmp.hyb == h1][i]
row_i = data_tmp.loc[[data_i_idx]]
for j in range(len(query[i])):
data_j_idx = data_tmp.index[data_tmp.hyb == h2][
query[i][j]]
row_j = data_tmp.loc[[data_j_idx]]
d = np.linalg.norm(
np.array(row_i[['x', 'y', 'z']]) -
np.array(row_j[['x', 'y', 'z']]))
# create transition var only if two signals are close enough
fI = np.absolute(
np.float64(row_i.Imax_gf) -
np.float64(row_j.Imax_gf))
mu_d = 1 / (1 + k1 * d)
Tvar_tmp = Tvar_tmp.append(pd.Series(
[
X_idx_tmp, data_i_idx, data_j_idx,
prob2Eng(1 - mu_d),
prob2Eng(mu_d), mu_d
],
index=[
'x_idx', 'anchestor_x_idx',
'descendant_x_idx', 'E_0', 'E_1', 'mu_D'
]),
ignore_index=True)
X_idx_tmp = X_idx_tmp + 1
if prior == "prior":
# Exclude paths not in taglist
Tvar_tmp.x_idx = np.ravel_multi_index([
Tvar_tmp.anchestor_x_idx.astype(np.int),
Tvar_tmp.descendant_x_idx.astype(np.int)
], (len(Dvar_tmp), len(Dvar_tmp)))
hyb_sets = []
for h in num_hyb:
hyb_sets.append(
np.unique(np.array(Dvar_tmp[Dvar_tmp.hyb == h].ch)))
allPaths = list(itertools.product(*hyb_sets))
allowedPaths_D_list = []
allowedPaths_T_list = []
# Remove not existing paths
if len(allPaths):
for tag in range(len(tagList)):
if (allPaths == np.array(
tagList.iloc[tag, :]).astype(int)
).all(1).any():
allowedDvar = pd.DataFrame(
data={
'x': [],
'y': [],
'z': [],
'hyb': [],
'ch': [],
'X_idx': [],
'E_0': [],
'E_1': []
})
for h in num_hyb:
allowedDvar = allowedDvar.append(
Dvar_tmp[(Dvar_tmp.hyb == h)
& (Dvar_tmp.ch == int(
tagList.iloc[tag,
int(h - 1)]))],
ignore_index=True)
DGraph = nx.Graph()
DGraph.add_nodes_from(allowedDvar.X_idx)
DGraph.add_edges_from(
Tvar_tmp[(Tvar_tmp.anchestor_x_idx.isin(
allowedDvar.X_idx))
& (Tvar_tmp.descendant_x_idx.isin(
allowedDvar.X_idx))][[
'anchestor_x_idx',
'descendant_x_idx'
]].values.tolist())
for c in nx.connected_components(DGraph):
if len(allowedDvar[allowedDvar.X_idx.isin(
c)].hyb.unique()) == len(num_hyb):
allowedPaths_D_list.append(
np.array(list(c)))
tmp_dvar = np.array(Tvar_tmp[
(Tvar_tmp.anchestor_x_idx.isin(c)) &
(Tvar_tmp.descendant_x_idx.isin(c))][[
'anchestor_x_idx',
'descendant_x_idx'
]])
tmp_path = []
for v in range(len(tmp_dvar)):
tmp_path.append(
np.ravel_multi_index([
int(tmp_dvar[v, 0]),
int(tmp_dvar[v, 1])
], (len(Dvar_tmp), len(Dvar_tmp))))
allowedPaths_T_list.append(
np.array(tmp_path))
Dvar_tmp = Dvar_tmp[Dvar_tmp.X_idx.isin(
np.concatenate(allowedPaths_D_list))]
Tvar_tmp = Tvar_tmp[Tvar_tmp.x_idx.isin(
np.concatenate(allowedPaths_T_list))]
d = dict(zip(Dvar_tmp.X_idx, np.arange(len(Dvar_tmp))))
Dvar_tmp.X_idx = [d[x] for x in Dvar_tmp.X_idx]
Tvar_tmp.anchestor_x_idx = [
d[x] for x in Tvar_tmp.anchestor_x_idx
]
Tvar_tmp.descendant_x_idx = [
d[x] for x in Tvar_tmp.descendant_x_idx
]
Dvar_tmp.X_idx = Dvar_tmp.X_idx + 1
Tvar_tmp.anchestor_x_idx = Tvar_tmp.anchestor_x_idx + 1
Tvar_tmp.descendant_x_idx = Tvar_tmp.descendant_x_idx + 1
Dvar_tmp['X'] = np.arange(1, len(Dvar_tmp) + 1)
sink = Dvar_tmp.X.max() + 1
# Inizialize graph
G = nx.DiGraph()
E = [] # Edges
n = sink + 1
for h in num_hyb:
for idx, row in Dvar_tmp[(Dvar_tmp.hyb == h)].iterrows():
if h == 1:
E.append((0, row.X, {'capacity': 1, 'weight': 0}))
# Add detection edges
E.append((row.X, n, {
'capacity': 1,
'weight': np.round(row.E_1 * 1000000).astype(int)
}))
n = n + 1
G.add_edges_from(E)
E = []
for idx, row in Tvar_tmp[(Tvar_tmp.anchestor_x_idx.isin(
Dvar_tmp[Dvar_tmp.hyb == h].X_idx))].iterrows():
# Add transition edges
E.append((list(G.successors(row.anchestor_x_idx))[0],
row.descendant_x_idx, {
'capacity': 1,
'weight':
np.round(row.E_1 * 1000000).astype(int)
}))
G.add_edges_from(E)
# For each D of last cycle connect to sink
E = []
for idx, row in Dvar_tmp[(
Dvar_tmp.hyb == num_hyb.max())].iterrows():
E.append((list(G.successors(row.X_idx))[0], sink, {
'capacity': 1,
'weight': 0
}))
G.add_edges_from(E)
# Prune graph removing leaf nodes
remove_nodes = []
for n in G.nodes:
n_set = nx.algorithms.descendants(G, n)
if not sink in n_set:
remove_nodes.append(n)
if n == 0: # source and sink are not connected
return
remove_nodes.remove(sink)
G.remove_nodes_from(remove_nodes)
MaxFlowMinCost = nx.max_flow_min_cost(G, 0, sink)
# Decode sequence
E = []
for n1 in MaxFlowMinCost:
for n2 in MaxFlowMinCost[n1]:
if MaxFlowMinCost[n1][n2] == 1:
E.append((n1, n2))
G = nx.Graph()
G.add_edges_from(E)
G.remove_node(0)
G.remove_node(sink)
return {'G': G, 'Dvar': Dvar_tmp, 'Tvar': Tvar_tmp}
sys.stdout.flush()
if name is not None:
newargs = initargs + ['--imodes','flexm-'+str(nstruc+1)+name+'.dat']
if not os.path.exists('flexm-'+str(nstruc+1)+name+'.dat'):
break
collectlib.collect_iattract(newargs)
result = collectlib.collect_next()
if result: break
nstruc += 1
coor = collectlib.collect_all_coor()
pdbsizes2 = np.cumsum([0] + pdbsizes)
coors = [coor[pdbsizes2[n]:pdbsizes2[n+1]] for n in range(len(pdbs))]
energy = 0
eblock = 0
for n1 in range(len(pdbs)):
c1 = coors[n1]
tree1 = KDTree(c1)
for n2 in range(n1+1, len(pdbs)):
c2 = coors[n2]
tree2 = KDTree(c2)
energyblock = energyblocks[eblock]
pairs = tree1.query_ball_tree(tree2, 10)
ene = sum([sum(e[p]) for e,p in zip(energyblock,pairs) if len(p)])
energy += ene
eblock += 1
f1.write("%.3f\n" % energy)
def _compute_global_cell_graph_features(
centroids,
neighbor_distances,
neighbor_counts,
):
"""Internal support for compute_global_cell_graph_features that
returns its result in a nested nametuple structure instead of a
pandas DataFrame.
"""
vor = Voronoi(centroids)
centroids = vor.points
vertices = vor.vertices
regions = [r for r in vor.regions if r and -1 not in r]
areas = np.stack(_poly_area(vertices[r]) for r in regions)
peris = np.stack(_poly_peri(vertices[r]) for r in regions)
max_dists = np.stack(pdist(vertices[r]).max() for r in regions)
poly_props = PolyProps._make(map(_pop_stats, (areas, peris, max_dists)))
de = Delaunay(centroids)
# From the docs: "Coplanar points are input points which were not
# included in the triangulation due to numerical precision
# issues." I don't know how this would affect the results if
# present, and it doesn't appear to happen, so it's excluded here.
assert not de.coplanar.size
indptr, indices = de.vertex_neighbor_vertices
bin_connectivity = sparse.csr_matrix(
(np.ones(len(indices), dtype=bool), indices, indptr),
(len(centroids), ) * 2)
ridge_points = sparse.triu(bin_connectivity, format='coo')
ridge_points = np.stack((ridge_points.row, ridge_points.col), axis=-1)
# This isn't exactly the collection of sides, since if they should
# be counted per-triangle then we weight border ridges wrong
# relative to ridges that are part of two triangles.
ridge_lengths = _dist(*np.swapaxes(centroids[ridge_points], 0, 1))
sides = ridge_lengths
areas = np.stack(_poly_area(centroids[t]) for t in de.simplices)
tri_props = TriProps._make(map(_pop_stats, (sides, areas)))
graph = sparse.coo_matrix((ridge_lengths, ridge_points.T),
(len(centroids), len(centroids)))
mst = minimum_spanning_tree(graph)
# Without looking into exactly how minimum_spanning_tree
# constructs its output, elimate any explicit zeros to be on the
# safe side.
mst_branches = _pop_stats(mst.data[mst.data != 0])
tree = KDTree(centroids)
neigbors_in_distance = {
# Yes, we just throw away the actual points
r: _pop_stats(np.stack(map(len, tree.query_ball_tree(tree, r))) - 1)
for r in neighbor_distances
}
distance_for_neighbors = dict(
zip(
neighbor_counts,
map(_pop_stats,
tree.query(centroids, [c + 1 for c in neighbor_counts])[0].T),
))
density_props = DensityProps(neigbors_in_distance, distance_for_neighbors)
return Props(poly_props, tri_props, mst_branches, density_props)
def _compute_global_cell_graph_features(
centroids,
neighbor_distances,
neighbor_counts,
):
"""Internal support for compute_global_cell_graph_features that
returns its result in a nested nametuple structure instead of a
pandas DataFrame.
"""
vor = Voronoi(centroids)
centroids = vor.points
vertices = vor.vertices
regions = [r for r in vor.regions if r and -1 not in r]
areas = np.stack(_poly_area(vertices[r]) for r in regions)
peris = np.stack(_poly_peri(vertices[r]) for r in regions)
max_dists = np.stack(pdist(vertices[r]).max() for r in regions)
poly_props = PolyProps._make(map(_pop_stats, (areas, peris, max_dists)))
# Assume that each Voronoi vertex is on exactly three ridges.
ridge_points = vor.ridge_points
# This isn't exactly the collection of sides, since if they should
# be counted per-triangle then we weight border ridges wrong
# relative to ridges that are part of two triangles.
ridge_lengths = _dist(*np.swapaxes(centroids[ridge_points], 0, 1))
sides = ridge_lengths
# Point indices of each triangle
tris = [[] for _ in vertices]
for p, ri in enumerate(vor.point_region):
for vi in vor.regions[ri]:
if vi != -1:
tris[vi].append(p)
# This will only fail in particular symmetrical cases where a
# Voronoi vertex is associated with more than three centroids.
# Since this should be unlikely in practice, we don't handle it
# beyond throwing an AssertionError
assert all(len(t) == 3 for t in tris)
tris = np.asarray(tris)
areas = np.stack(_poly_area(centroids[t]) for t in tris)
tri_props = TriProps._make(map(_pop_stats, (sides, areas)))
graph = sparse.coo_matrix((ridge_lengths, np.sort(ridge_points).T),
(len(centroids), len(centroids)))
mst = minimum_spanning_tree(graph)
# Without looking into exactly how minimum_spanning_tree
# constructs its output, elimate any explicit zeros to be on the
# safe side.
mst_branches = _pop_stats(mst.data[mst.data != 0])
tree = KDTree(centroids)
neigbors_in_distance = {
# Yes, we just throw away the actual points
r: _pop_stats(np.stack(map(len, tree.query_ball_tree(tree, r))) - 1)
for r in neighbor_distances
}
distance_for_neighbors = dict(zip(
neighbor_counts,
map(_pop_stats, tree.query(centroids, [c + 1 for c in neighbor_counts])[0].T),
))
density_props = DensityProps(neigbors_in_distance, distance_for_neighbors)
return Props(poly_props, tri_props, mst_branches, density_props)
def fast_fit_big_image(im,
centers_zxy,
radius_fit=4,
avoid_neigbors=True,
recenter=False,
verbose=True,
better_fit=False,
troubleshoot=False):
ps = []
if len(centers_zxy) > 0:
if avoid_neigbors:
centers_tree = KDTree(centers_zxy)
list_inters = centers_tree.query_ball_tree(centers_tree,
radius_fit * 2)
centers_ = centers_zxy
if verbose:
from tqdm import tqdm_notebook as tqdm
centers_ = tqdm(centers_zxy)
zb, xb, yb = np.reshape(
np.indices([radius_fit * 2] * 3) - radius_fit, [3, -1]).astype(int)
keep = zb * zb + xb * xb + yb * yb <= radius_fit**2
zb, xb, yb = zb[keep], xb[keep], yb[keep]
sz, sx, sy = im.shape
X_c = np.array([zb, xb, yb]).T
for ic, (zc, xc, yc) in enumerate(centers_):
if avoid_neigbors:
common_pts_ids = list_inters[ic]
index_ic = list_inters[ic].index(ic)
centers__ = centers_zxy[common_pts_ids] - [zc, xc, yc]
nns_ = np.argmin(cdist(centers__, X_c), 0)
zb_, xb_, yb_ = X_c[nns_ == index_ic].T
else:
zb_, xb_, yb_ = zb, xb, yb
z_keep, x_keep, y_keep = int(zc) + zb_, int(xc) + xb_, int(
yc) + yb_
z_keep, x_keep, y_keep = in_dim(z_keep, x_keep, y_keep, sz, sx, sy)
X_ = np.array([z_keep, x_keep, y_keep]).T
im_ = im[z_keep, x_keep, y_keep]
#recenter to max in the imge - generally indicated if using seeds from a different image
if recenter:
if len(im_) > 0:
#while True:
im_c = np.argmax(im_)
zcn, xcn, ycn = z_keep[im_c], x_keep[im_c], y_keep[im_c]
#if zcn==zc and xcn==xc and ycn==yc:
# break
zc, xc, yc = zcn, xcn, ycn
z_keep, x_keep, y_keep = int(zc) + zb_, int(xc) + xb_, int(
yc) + yb_
z_keep, x_keep, y_keep = in_dim(z_keep, x_keep, y_keep, sz,
sx, sy)
X_ = np.array([z_keep, x_keep, y_keep]).T
im_ = im[z_keep, x_keep, y_keep]
if not better_fit:
reconstruct, plt_val, compare_with_fitting = False, False, False
if troubleshoot:
reconstruct, plt_val, compare_with_fitting = True, True, True
p_ = gfit_fast(im_,
X_.T,
bk_f=0.1,
reconstruct=reconstruct,
plt_val=plt_val,
compare_with_fitting=compare_with_fitting)
else:
p_ = np.array([np.nan] * 11)
if len(im_) > 0:
center = X_[np.argmax(im_)]
obj = GaussianFit(im_,
X_.T,
center=center,
delta_center=2.5)
obj.fit()
p_ = obj.p
ps.append(p_)
ps = np.array(ps)
return ps
def mean_neighbors(tree, radius):
if isinstance(tree, np.ndarray):
tree = KDTree(tree)
indices = tree.query_ball_tree(tree, radius)
num_neighbors = [len(i) for i in indices]
return np.mean(num_neighbors)
def spherematch(lon1, lat1, lon2=None, lat2=None, matchrad=None,
nnearest=0, maxmatches=-1):
"""
---------------------------------------------------------------------------
Finds matches in one catalog to another.
Parameters
lon1 : array-like
Longitude-like (RA, etc.) coordinate in degrees of the first catalog
lat1 : array-like
Latitude-like (Dec, etc.) coordinate in degrees of the first catalog (shape of array must match `lon1`)
lon2 : array-like
Latitude-like (RA, etc.) coordinate in degrees of the second catalog
lat2 : array-like
Latitude-like (Dec, etc.) in degrees of the second catalog (shape of array must match `ra2`)
matchrad : float or None, optional
How close (in degrees) a match has to be to count as a match. If None,
all nearest neighbors for the first catalog will be returned.
nnearest : int, optional
The nth neighbor to find.
maxmatches : int, optional
Maximum number of matches to find. If maxmatches > 0, the code finds
all matches up to maxmatches satisfying matchrad and nnearest is
ignored.
Returns
-------
m1 : int array
Indices into the first catalog of the matches.
m2 : int array
Indices into the second catalog of the matches.
d12 : float array
Distance (in degrees) between the matches
-------------------------------------------------------
"""
try:
lon1, lat1
except NameError:
raise NameError('lon1 and/or lat1 not defined. Aborting spherematch()')
if not isinstance(lon1, (list, NP.ndarray)):
lon1 = NP.asarray(lon1)
if not isinstance(lat1, (list, NP.ndarray)):
lat1 = NP.asarray(lat1)
try:
nnearest = int(nnearest)
except TypeError:
raise TypeError('nnearest should be a non-negative integer.')
try:
maxmatches = int(maxmatches)
except TypeError:
raise TypeError('maxmatches should be a non-negative integer.')
if matchrad is None:
if maxmatches > 0:
nnearest = 0
print 'No matchrad specified. Will determine all the {0} nearest neighbours.'.format(maxmatches)
else:
maxmatches = -1
if nnearest <= 0:
nnearest = 1
print 'No matchrad specified. Will determine the nearest neighbour # {0}.'.format(nnearest)
elif not isinstance(matchrad, (int,float)):
raise TypeError('matchrad should be a scalar number.')
elif matchrad > 0.0:
matchrad_cartesian = 2.0*NP.sin(0.5*matchrad*NP.pi/180.0)
if maxmatches >= 0:
nnearest = 0
else:
if nnearest <= 0:
nnearest = 1
print 'maxmatches is negative. Will determine the nearest neighbour # {0}.'.format(nnearest)
else:
raise ValueError('matchrad is not positive.')
self_match = False
if (lon2 is None) and (lat2 is None):
self_match = True
elif lon2 is None:
lon2 = lon1
elif lat2 is None:
lat2 = lat2
if lon1.shape != lat1.shape:
raise ValueError('lon1 and lat1 should be of same length')
if not self_match:
if lon2.shape != lat2.shape:
raise ValueError('lon2 and lat2 should be of same length')
if lon1.size == 1:
lon1 = NP.asarray(lon1).reshape(1)
lat1 = NP.asarray(lat1).reshape(1)
if lon2.size == 1:
lon2 = NP.asarray(lon2).reshape(1)
lat2 = NP.asarray(lat2).reshape(1)
x1, y1, z1 = sph2xyz(lon1.ravel(), lat1.ravel())
# this is equivalent to, but faster than just doing NP.array([x1, y1, z1])
coords1 = NP.empty((x1.size, 3))
coords1[:, 0] = x1
coords1[:, 1] = y1
coords1[:, 2] = z1
if (lon2 is not None) and (lat2 is not None):
x2, y2, z2 = sph2xyz(lon2.ravel(), lat2.ravel())
# this is equivalent to, but faster than just doing NP.array([x1, y1, z1])
coords2 = NP.empty((x2.size, 3))
coords2[:, 0] = x2
coords2[:, 1] = y2
coords2[:, 2] = z2
if maxmatches == 0:
kdt1 = KDT(coords1)
if not self_match:
kdt2 = KDT(coords2)
ngbr_of_first_in_second = kdt1.query_ball_tree(kdt2, matchrad_cartesian)
m1 = [i for i in xrange(len(ngbr_of_first_in_second)) for j in ngbr_of_first_in_second[i] if ngbr_of_first_in_second[i] != []]
m2 = [j for i in xrange(len(ngbr_of_first_in_second)) for j in ngbr_of_first_in_second[i] if ngbr_of_first_in_second[i] != []]
d12 = sphdist(lon1[m1], lat1[m1], lon2[m2], lat2[m2])
else:
ngbr_of_first_in_itself = kdt1.query_ball_tree(kdt1, matchrad_cartesian)
m1 = [i for i in xrange(len(ngbr_of_first_in_itself)) for j in ngbr_of_first_in_itself[i] if ngbr_of_first_in_itself[i] != [] and i != j]
m2 = [j for i in xrange(len(ngbr_of_first_in_itself)) for j in ngbr_of_first_in_itself[i] if ngbr_of_first_in_itself[i] != [] and i != j]
d12 = sphdist(lon1[m1], lat1[m1], lon1[m2], lat1[m2])
else:
if not self_match:
kdt2 = KDT(coords2)
if matchrad is None:
m2 = kdt2.query(coords1, max(maxmatches,nnearest))[1]
else:
m2 = kdt2.query(coords1, max(maxmatches,nnearest), distance_upper_bound=matchrad_cartesian)[1]
else:
kdt1 = KDT(coords1)
if matchrad is None:
m2 = kdt1.query(coords1, max(maxmatches,nnearest)+1)[1]
else:
m2 = kdt2.query(coords1, max(maxmatches,nnearest)+1, distance_upper_bound=matchrad_cartesian)[1]
if nnearest > 0:
m1 = NP.arange(lon1.size)
if nnearest > 1:
m2 = m2[:,-1]
else:
m1 = NP.repeat(NP.arange(lon1.size).reshape(lon1.size,1), maxmatches, axis=1).flatten()
m2 = m2[:,-maxmatches:].flatten() # Extract the last maxmatches columns
if not self_match:
msk = m2 < lon2.size
else:
msk = m1 < lon1.size
m1 = m1[msk]
m2 = m2[msk]
if not self_match:
d12 = sphdist(lon1[m1], lat1[m1], lon2[m2], lat2[m2])
else:
d12 = sphdist(lon1[m1], lat1[m1], lon1[m2], lat1[m2])
if matchrad is not None:
msk = d12 <= matchrad
m1 = m1[msk]
m2 = m2[msk]
d12 = d12[msk]
return m1, m2, d12
