Skip to content

Commit

Permalink
attempting to speed up a few more things with cython
Browse files Browse the repository at this point in the history
  • Loading branch information
@JohannesBuchner
JohannesBuchner committed Nov 30, 2020
1 parent 373e6fc commit cb16331
Show file tree
Hide file tree
Showing 2 changed files with 108 additions and 49 deletions.
21 changes: 14 additions & 7 deletions ultranest/integrator.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1998,7 +1998,8 @@ def _should_node_be_expanded(
# we do not need to expand this one
# expand_node = len(node.children) == 0
# prefer 1 child, or the number required, if specified
expand_node = len(node.children) < target_min_num_children.get(node.id, 1)
nmin = target_min_num_children.get(node.id, 1) if target_min_num_children else 1
expand_node = len(node.children) < nmin
# print("not expanding, because we are quite wide", nlive, minimal_width, minimal_widths_sequence)
# but we have to expand the first iteration,
# otherwise the integrator never sets H
Expand All @@ -2023,7 +2024,8 @@ def run(
max_num_improvement_loops=-1,
min_num_live_points=400,
cluster_num_live_points=40,
order_test_window=10,
insertion_test_window=10,
insertion_test_zscore_threshold=2,
):
"""Run until target convergence criteria are fulfilled.
Expand Down Expand Up @@ -2085,7 +2087,11 @@ def run(
cluster_num_live_points: int
require at least this many live points per detected cluster
order_test_window: float
insertion_test_zscore_threshold: float
z-score used as a threshold for the insertion order test.
Set to infinity to disable.
insertion_test_window: float
Number of iterations after which the insertion order test is reset.
"""
Expand All @@ -2102,7 +2108,8 @@ def run(
cluster_num_live_points=cluster_num_live_points,
show_status=show_status,
viz_callback=viz_callback,
order_test_window=order_test_window,
insertion_test_window=insertion_test_window,
insertion_test_zscore_threshold=insertion_test_zscore_threshold,
):
if self.log:
self.logger.debug("did a run_iter pass!")
Expand All @@ -2129,7 +2136,7 @@ def run_iter(
cluster_num_live_points=40,
show_status=True,
viz_callback='auto',
order_test_window=10,
insertion_test_window=10,
insertion_test_zscore_threshold=2,
):
"""Iterate towards convergence.
Expand Down Expand Up @@ -2338,14 +2345,14 @@ def run_iter(
u, p, L = self._create_point(Lmin=Lmin, ndraw=ndraw, active_u=active_u, active_values=active_values)
child = self.pointpile.make_node(L, u, p)
main_iterator.Lmax = max(main_iterator.Lmax, L)
if len(np.unique(active_values)) == nlive:
if np.isfinite(insertion_test_zscore_threshold) and nlive > 1:
insertion_test.add((active_values < L).sum(), nlive)
if abs(insertion_test.zscore) > insertion_test_zscore_threshold:
insertion_test_runs.append(insertion_test.N)
insertion_test_quality = insertion_test.N
insertion_test_direction = np.sign(insertion_test.zscore)
insertion_test.reset()
elif insertion_test.N > nlive * order_test_window:
elif insertion_test.N > nlive * insertion_test_window:
insertion_test_quality = np.inf
insertion_test_direction = 0
insertion_test.reset()
Expand Down
136 changes: 94 additions & 42 deletions ultranest/mlfriends.pyx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def count_nearby(np.ndarray[np.float_t, ndim=2] apts,
cdef count_nearby(np.ndarray[np.float_t, ndim=2] apts,
np.ndarray[np.float_t, ndim=2] bpts,
np.float_t radiussq,
np.ndarray[np.int_t, ndim=1] nnearby
Expand Down Expand Up @@ -149,10 +149,14 @@ def compute_mean_pair_distance(
return total_dist / Npairs





def update_clusters(upoints, tpoints, maxradiussq, clusterids=None):
@cython.boundscheck(False)
@cython.wraparound(False)
cdef _update_clusters(
np.ndarray[np.float_t, ndim=2] upoints,
np.ndarray[np.float_t, ndim=2] tpoints,
np.float_t maxradiussq,
np.ndarray[np.int_t, ndim=1] clusterids,
):
"""Clusters `upoints`, so that clusters are distinct if no member pair is within a radius of sqrt(`maxradiussq`)
clusterids are the cluster indices of each point
Expand All @@ -172,18 +176,16 @@ def update_clusters(upoints, tpoints, maxradiussq, clusterids=None):
"""
#print("clustering with maxradiussq %f..." % maxradiussq)
assert upoints.shape == tpoints.shape
assert upoints.shape[0] == tpoints.shape[0], ('different number of points', upoints.shape[0], tpoints.shape[0])
assert upoints.shape[1] == tpoints.shape[1], ('different dimensionality of points', upoints.shape[1], tpoints.shape[1])
clusteridxs = np.zeros(len(tpoints), dtype=int)
currentclusterid = 1
i = 0
if clusterids is None:
clusterids = np.zeros(len(tpoints), dtype=int)
else:
# avoid issues when old clusterids are from a longer array
clusterids = clusterids[:len(tpoints)]
existing = clusterids == currentclusterid
if existing.any():
i = np.where(existing)[0][0]
# avoid issues when old clusterids are from a longer array
clusterids = clusterids[:len(tpoints)]
existing = clusterids == currentclusterid
if existing.any():
i = np.where(existing)[0][0]

clusteridxs[i] = currentclusterid
while True:
Expand Down Expand Up @@ -234,15 +236,46 @@ def update_clusters(upoints, tpoints, maxradiussq, clusterids=None):
# so use the mean of the entire point population instead
group_mean = upoints.mean(axis=0).reshape((1,-1))
overlapped_upoints[clusteridxs == idx,:] = group_upoints - group_mean
#print("clustering done, %d clusters" % nclusters)
#if nclusters > 1:
# np.savetxt("clusters%d.txt" % nclusters, upoints)
# np.savetxt("clusters%d_radius.txt" % nclusters, [maxradiussq])

return nclusters, clusteridxs, overlapped_upoints

@cython.boundscheck(False)
@cython.wraparound(False)
def update_clusters(
np.ndarray[np.float_t, ndim=2] upoints,
np.ndarray[np.float_t, ndim=2] tpoints,
np.float_t maxradiussq,
clusterids = None,
):
"""Clusters `upoints`, so that clusters are distinct if no member pair is within a radius of sqrt(`maxradiussq`)
clusterids are the cluster indices of each point
clusterids re-uses the existing ids to assign new cluster ids
clustering is performed on a transformed coordinate space (`tpoints`).
Returned values are based on upoints.
Returns
---------
nclusters: int
the number of clusters found, which is also clusterids.max()
new_clusterids: array of int
the new clusterids for each point
overlapped_points:
upoints with their cluster centers subtracted.
"""
if clusterids is None:
clusterids = np.zeros(len(tpoints), dtype=int)
return _update_clusters(upoints, tpoints, maxradiussq, clusterids)


def make_eigvals_positive(a, targetprod):
@cython.boundscheck(False)
@cython.wraparound(False)
def make_eigvals_positive(
np.ndarray[np.float_t, ndim=2] a,
np.float_t targetprod
):
"""For the symmetric square matrix ``a``, increase any zero eigenvalues
to fulfill the given target product of eigenvalues.
Expand All @@ -263,7 +296,12 @@ def make_eigvals_positive(a, targetprod):

return a

def bounding_ellipsoid(x, minvol=0.):
@cython.boundscheck(False)
@cython.wraparound(False)
def bounding_ellipsoid(
np.ndarray[np.float_t, ndim=2] x,
np.float_t minvol=0.
):
"""Calculate bounding ellipsoid containing a set of points x.
Parameters
Expand All @@ -281,7 +319,8 @@ def bounding_ellipsoid(x, minvol=0.):
"""
# Function taken from nestle, MIT licensed, (C) kbarbary

npoints, ndim = x.shape
npoints = x.shape[0]
ndim = x.shape[1]

# Calculate covariance of points
ctr = np.mean(x, axis=0)
Expand Down Expand Up @@ -523,7 +562,7 @@ class AffineLayer(ScalingLayer):
return u


def vol_prefactor(n):
def vol_prefactor(np.int_t n):
"""Volume constant for an `n`-dimensional sphere.
for `n` even: $$ (2pi)^(n /2) / (2 * 4 * ... * n)$$
Expand All @@ -542,6 +581,37 @@ def vol_prefactor(n):

return f

def _inside_ellipsoid(
np.ndarray[np.float_t, ndim=2] points,
np.ndarray[np.float_t, ndim=1] ellipsoid_center,
np.ndarray[np.float_t, ndim=2] ellipsoid_invcov,
np.float_t square_radius
):
"""Check if inside ellipsoid
Parameters
----------
points: array of vectors
Points to check
ellipsoid_center: vector
center of ellipsoid
ellipsoid_invcov: matrix
inverse covariance matrix
square_radius: float
square radius
Returns
---------
is_inside: array of bools
True if inside wrapping for each point in `points`.
"""
# compute distance vector to center
d = points - ellipsoid_center
# distance in normalised coordates: vector . matrix . vector
# where the matrix is the ellipsoid inverse covariance
r = np.einsum('ij,jk,ik->i', d, ellipsoid_invcov, d)
# (r <= 1) means inside
return r <= square_radius

class MLFriends(object):
"""MLFriends region.
Expand Down Expand Up @@ -816,16 +886,7 @@ class MLFriends(object):
True if inside wrapping ellipsoid, for each point in `pts`.
"""
# to disable wrapping ellipsoid
#return np.ones(len(u), dtype=bool)

# compute distance vector to center
d = u - self.ellipsoid_center
# distance in normalised coordates: vector . matrix . vector
# where the matrix is the ellipsoid inverse covariance
r = np.einsum('ij,jk,ik->i', d, self.ellipsoid_invcov, d)
# (r <= 1) means inside
return r <= self.enlarge
return _inside_ellipsoid(u, self.ellipsoid_center, self.ellipsoid_invcov, self.enlarge)

def compute_mean_pair_distance(self):
return compute_mean_pair_distance(self.unormed, self.transformLayer.clusterids)
Expand Down Expand Up @@ -903,13 +964,4 @@ class WrappingEllipsoid(object):
True if inside wrapping ellipsoid, for each point in `pts`.
"""
# to disable wrapping ellipsoid
#return np.ones(len(u), dtype=bool)

# compute distance vector to center
d = u - self.ellipsoid_center
# distance in normalised coordates: vector . matrix . vector
# where the matrix is the ellipsoid inverse covariance
r = np.einsum('ij,jk,ik->i', d, self.ellipsoid_invcov, d)
# (r <= 1) means inside
return r <= self.enlarge
return _inside_ellipsoid(u, self.ellipsoid_center, self.ellipsoid_invcov, self.enlarge)

0 comments on commit cb16331

Please sign in to comment.