9f2dd4f6-0775-47e9-b700-af647027ebfa

Committed 22 Apr 2024 12:51PM UTC coverage: 74.53% (+0.3%) from 74.242%

Build # 9f2dd4f6-0775-47e9-b700-af647027ebfa

Build Type

push

circleci

Committed by

web-flow

Commit Message

Merge pull request #118 from njzifjoiez/fixed-size-vectorised-slice-sampler

vectorised slice sampler of fixed batch size

Run Details

1329 of 2026 branches covered (65.6%)

Branch coverage included in aggregate %.

79 of 80 new or added lines in 1 file covered. (98.75%)

1 existing line in 1 file now uncovered.

4026 of 5159 relevant lines covered (78.04%)

0.78 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

74.91

/ultranest/samplingpath.py

"""Sparsely sampled, virtual sampling path.

Supports reflections at unit cube boundaries, and regions.
"""


import numpy as np
from numpy.linalg import norm
import matplotlib.pyplot as plt


def nearest_box_intersection_line(ray_origin, ray_direction, fwd=True):
    r"""Compute intersection of a line (ray) and a unit box (0:1 in all axes).

    Based on
    http://www.iquilezles.org/www/articles/intersectors/intersectors.htm

    To continue forward traversing at the reflection point use::

        while True:
            # update current point x
            x, _, i = box_line_intersection(x, v)
            # change direction
            v[i] *= -1

    Parameters
    -----------
    ray_origin: vector
        starting point of line
    ray_direction: vector
        line direction vector

    Returns
    --------
    p: vector
        intersection point
    t: float
        intersection point distance from ray\_origin in units in ray\_direction
    i: int
        axes which change direction at pN

    """
    # make sure ray starts inside the box
    assert (ray_origin >= 0).all(), ray_origin
    assert (ray_origin <= 1).all(), ray_origin
    assert ((ray_direction**2).sum()**0.5 > 1e-200).all(), ray_direction

    # step size
    with np.errstate(divide='ignore', invalid='ignore'):
        m = 1. / ray_direction
        n = m * (ray_origin - 0.5)
        k = np.abs(m) * 0.5
        # line coordinates of intersection
        # find first intersecting coordinate
        if fwd:
            t2 = -n + k
            tF = np.nanmin(t2)
            iF = np.where(t2 == tF)[0]
        else:
            t1 = -n - k
            tF = np.nanmax(t1)
            iF = np.where(t1 == tF)[0]

    pF = ray_origin + ray_direction * tF
    eps = 1e-6
    assert (pF >= -eps).all(), (pF, ray_origin, ray_direction)
    assert (pF <= 1 + eps).all(), (pF, ray_origin, ray_direction)
    pF[pF < 0] = 0
    pF[pF > 1] = 1
    return pF, tF, iF


def box_line_intersection(ray_origin, ray_direction):
    """Find intersections of a line with the unit cube, in both sides.

    Parameters
    -----------
    ray_origin: vector
        starting point of line
    ray_direction: vector
        line direction vector

    Returns
    --------
    left: nearest_box_intersection_line return value
        from negative direction
    right: nearest_box_intersection_line return value
        from positive direction

    """
    pF, tF, iF = nearest_box_intersection_line(ray_origin, ray_direction, fwd=True)
    pN, tN, iN = nearest_box_intersection_line(ray_origin, ray_direction, fwd=False)
    if tN > tF or tF < 0:
        assert False, "no intersection"
    return (pN, tN, iN), (pF, tF, iF)


def linear_steps_with_reflection(ray_origin, ray_direction, t, wrapped_dims=None):
    """Go `t` steps in direction `ray_direction` from `ray_origin`.

    Reflect off the unit cube if encountered, respecting wrapped dimensions.
    In any case, the distance should be ``t * ray_direction``.

    Returns
    --------
    new_point: vector
        end point
    new_direction: vector
        new direction.

    """
    if t == 0:
        return ray_origin, ray_direction
    if t < 0:
        new_point, new_direction = linear_steps_with_reflection(ray_origin, -ray_direction, -t)
        return new_point, -new_direction

    if wrapped_dims is not None:
        reflected = np.zeros(len(ray_origin), dtype=bool)

    tleft = 1.0 * t
    while True:
        p, t, i = nearest_box_intersection_line(ray_origin, ray_direction, fwd=True)
        # print(p, t, i, ray_origin, ray_direction)
        assert np.isfinite(p).all()
        assert t >= 0, t
        if tleft <= t:  # stopping before reaching any border
            assert np.all(ray_origin + tleft * ray_direction >= 0), (ray_origin, tleft, ray_direction)
            assert np.all(ray_origin + tleft * ray_direction <= 1), (ray_origin, tleft, ray_direction)
            return ray_origin + tleft * ray_direction, ray_direction
        # go to reflection point
        ray_origin = p
        assert np.isfinite(ray_origin).all(), ray_origin
        # reflect
        ray_direction = ray_direction.copy()
        if wrapped_dims is None:
            ray_direction[i] *= -1
        else:
            # if we already once bumped into that (wrapped) axis,
            # do not continue but return this as end point
            if np.logical_and(reflected[i], wrapped_dims[i]).any():
                return ray_origin, ray_direction

            # note which axes we already flipped
            reflected[i] = True

            # in wrapped axes, we can keep going. Otherwise, reflects
            ray_direction[i] *= np.where(wrapped_dims[i], 1, -1)

            # in the i axes, we should wrap the coordinates
            assert np.logical_or(np.isclose(ray_origin[i], 1), np.isclose(ray_origin[i], 0)).all(), ray_origin[i]
            ray_origin[i] = np.where(wrapped_dims[i], 1 - ray_origin[i], ray_origin[i])

        assert np.isfinite(ray_direction).all(), ray_direction
        # reduce remaining distance
        tleft -= t


def get_sphere_tangent(sphere_center, edge_point):
    """Compute tangent at sphere surface point.

    Assume a sphere centered at sphere_center with radius
    so that edge_point is on the surface. At edge_point, in
    which direction does the normal vector point?

    Parameters
    -----------
    sphere_center: vector
        center of sphere
    edge_point: vector
        point at the surface

    Returns
    --------
    tangent: vector
        vector pointing to the sphere center.

    """
    arrow = sphere_center - edge_point
    return arrow / norm(arrow)


def get_sphere_tangents(sphere_center, edge_point):
    """Compute tangent at sphere surface point.

    Assume a sphere centered at sphere_center with radius
    so that edge_point is on the surface. At edge_point, in
    which direction does the normal vector point?

    This function is vectorized and handles arrays of arguments.

    Parameters
    -----------
    sphere_center: array
        centers of spheres
    edge_point: array
        points at the surface

    Returns
    --------
    tangent: array
        vectors pointing to the sphere center.

    """
    arrow = sphere_center - edge_point
    return arrow / norm(arrow, axis=1).reshape((-1, 1))


def reflect(v, normal):
    """Reflect vector ``v`` off a ``normal`` vector, return new direction vector."""
    return v - 2 * (normal * v).sum() * normal


def distances(direction, center, r=1):
    """Compute sphere-line intersection.

    Parameters
    -----------
    direction: vector
        direction vector (line starts at 0)
    center: vector
        center of sphere (coordinate vector)
    r: float
        radius of sphere

    Returns
    --------
    tpos, tneg: floats
        the positive and negative coordinate along the `l` vector where `r` is intersected.
        If no intersection, throws AssertError.

    """
    loc = (direction * center).sum()
    osqrnorm = (center**2).sum()
    # print(loc.shape, loc.shape, osqrnorm.shape)
    rootterm = loc**2 - osqrnorm + r**2
    # make sure we are crossing the sphere
    assert (rootterm > 0).all(), rootterm
    return -loc + rootterm**0.5, -loc - rootterm**0.5


def isunitlength(vec):
    """Verify that `vec` is of unit length."""
    assert np.isclose(norm(vec), 1), norm(vec)


def angle(a, b):
    """Compute dot product between vectors `a` and `b`.

    The arccos of the return value would give an actual angle.
    """
    # anorm = (a**2).sum()**0.5
    # bnorm = (b**2).sum()**0.5
    return (a * b).sum()  # / anorm / bnorm


def extrapolate_ahead(dj, xj, vj, contourpath=None):
    """Make `di` steps of size `vj` from `xj`.

    Reflect off unit cube if necessary.
    """
    assert dj == int(dj)

    # optimistically try to go there directly

    xk, vk = linear_steps_with_reflection(xj, vj, dj)

    return xk, vk  # deactivate feature below

    if contourpath is None:
        return xk, vk

    # check if we can already tell that the point will be outside
    region = contourpath.region
    if contourpath.region.inside(xk.reshape((1, -1))):
        return xk, vk

    # find first point outside region
    sign = 1 if dj > 0 else -1
    d = np.arange(0, dj, sign)
    first_point_outside = dj, xk, vk
    for di in d:
        xi, vi = linear_steps_with_reflection(xj, vj, di)
        if not region.inside(xk.reshape((1, -1))):
            first_point_outside = di, xi, vi
            break

    # reflect at this point (first outside)
    dout, reflpoint, v = first_point_outside

    if dout == 0:
        # already the starting point is outside.
        # return extrapolation and hope caller handles it
        return xk, vk

    # reversing:
    normal = contourpath.gradient(reflpoint)  # , v * sign)
    if normal is None:
        vnew = -v
    else:
        vnew = (v - 2 * angle(normal, v) * normal) * sign
    assert vnew.shape == v.shape, (vnew.shape, v.shape)
    assert np.isclose(norm(vnew), norm(v)), (vnew, v, norm(vnew), norm(v))

    # make one step (xl replaces first_point_outside/reflpoint)
    xl, vl = linear_steps_with_reflection(reflpoint, vnew, sign)
    # how many steps are still to do?
    dleft = dj - dout

    # make that many step in that direction, by recursing.

    # it is possible that this point is also outside. The next iteration
    # will catch that case.

    return extrapolate_ahead(dleft, xl, vl, contourpath=contourpath)


def interpolate(i, points, fwd_possible, rwd_possible, contourpath=None):
    """Interpolate a point on the path indicated by `points`.

    Given a sparsely sampled track (stored in .points),
    potentially encountering reflections,
    extract the corrdinates of the point with index `i`.
    That point may not have been evaluated yet.

    Parameters
    -----------
    i: int
        position on track to return.
    points: list of tuples (index, coordinate, direction, loglike)
        points on the path
    fwd_possible: bool
        whether the path could be extended in the positive direction.
    rwd_possible: bool
        whether the path could be extended in the negative direction.
    contourpath: ContourPath
        Use region to reflect. Not used at the moment.

    """
    points_before = [(j, xj, vj, Lj) for j, xj, vj, Lj in points if j <= i]
    points_after = [(j, xj, vj, Lj) for j, xj, vj, Lj in points if j >= i]

    # check if the point after is really after i
    if len(points_after) == 0 and not fwd_possible:
        # the path cannot continue, and i does not exist.
        # print("    interpolate_point %d: the path cannot continue fwd, and i does not exist." % i)
        j, xj, vj, Lj = max(points_before)
        return xj, vj, Lj, False

    # check if the point before is really before i
    if len(points_before) == 0 and not rwd_possible:
        # the path cannot continue, and i does not exist.
        k, xk, vk, Lk = min(points_after)
        # print("    interpolate_point %d: the path cannot continue rwd, and i does not exist." % i)
        return xk, vk, Lk, False

    if len(points_before) == 0 or len(points_after) == 0:
        # return None, None, None, False
        raise KeyError("cannot extrapolate outside path")

    j, xj, vj, Lj = max(points_before)
    k, xk, vk, Lk = min(points_after)

    # print("    interpolate_point %d between %d-%d" % (i, j, k))
    if j == i:  # we have this exact point in the chain
        return xj, vj, Lj, True

    assert not k == i  # otherwise the above would be true too

    # expand_to_step explores each reflection in detail, so
    # any points with change in v should have j == i
    # therefore we can assume:
    # assert (vj == vk).all()
    # this ^ is not true, because reflections on the unit cube can
    # occur, and change v without requiring a intermediate point.

    # j....i....k
    xl1, vj1 = extrapolate_ahead(i - j, xj, vj, contourpath=contourpath)
    xl2, vj2 = extrapolate_ahead(i - k, xk, vk, contourpath=contourpath)
    assert np.allclose(xl1, xl2), (xl1, xl2, i, j, k, xj, vj, xk, vk)
    assert np.allclose(vj1, vj2), (xl1, vj1, xl2, vj2, i, j, k, xj, vj, xk, vk)
    xl = xl1

    # xl = interpolate_between_two_points(i, xj, j, xk, k)
    # the new point is then just a linear interpolation
    # w = (i - k) * 1. / (j - k)
    # xl = xj * w + (1 - w) * xk

    return xl, vj, None, True


class SamplingPath(object):
    """Path described by a (potentially sparse) sequence of points.

    Convention of the stored point tuple ``(i, x, v, L)``:
    `i`: index (0 is starting point)
    `x`: point
    `v`: direction
    `L`: loglikelihood value
    """

    def __init__(self, x0, v0, L0):
        """Initialise with path starting point.

        Starting point (`x0`), direction (`v0`) and
        loglikelihood value (`L0`) of the path. Is given index 0.
        """
        self.reset(x0, v0, L0)

    def add(self, i, xi, vi, Li):
        """Add point `xi`, direction `vi` and value `Li` with index `i` to the path."""
        assert Li is not None
        assert len(xi.shape) == 1, (xi, xi.shape)
        assert len(vi.shape) == 1, (vi, vi.shape)
        assert len(np.shape(Li)) == 0, (Li, Li.shape)
        self.points.append((i, xi, vi, Li))

    def reset(self, x0, v0, L0):
        """Reset path, start from ``x0, v0, L0``."""
        self.points = []
        self.add(0, x0, v0, L0)
        self.fwd_possible = True
        self.rwd_possible = True

    def plot(self, **kwargs):
        """Plot the current path.

        Only uses first two dimensions.
        """
        x = np.array([x for i, x, v, L in sorted(self.points)])
        p, = plt.plot(x[:,0], x[:,1], 'o ', **kwargs)
        ilo, _, _, _ = min(self.points)
        ihi, _, _, _ = max(self.points)
        x = np.array([self.interpolate(i)[0] for i in range(ilo, ihi + 1)])
        kwargs['color'] = p.get_color()
        plt.plot(x[:,0], x[:,1], 'o-', ms=4, mfc='None', **kwargs)

    def interpolate(self, i):
        """Interpolate point with index `i` on path."""
        return interpolate(i, self.points, fwd_possible=self.fwd_possible, rwd_possible=self.rwd_possible)

    def extrapolate(self, i):
        """Advance beyond the current path, extrapolate from the end point.

        Parameters
        -----------
        i: int
            index on path.

        returns
        --------
        coords: vector
            coordinates of the new point.

        """
        if i >= 0:
            j, xj, vj, Lj = max(self.points)
            deltai = i - j
            assert deltai > 0, ("should be extrapolating", i, j)
        else:
            j, xj, vj, Lj = min(self.points)
            deltai = i - j
            assert deltai < 0, ("should be extrapolating", i, j)

        newpoint = extrapolate_ahead(deltai, xj, vj)
        return newpoint


class ContourSamplingPath(object):
    """Region-aware form of the sampling path.

    Uses region points to guess a likelihood contour gradient.
    """

    def __init__(self, samplingpath, region):
        """Initialise with `samplingpath` and `region`."""
        self.samplingpath = samplingpath
        self.points = self.samplingpath.points
        self.region = region

    def add(self, i, x, v, L):
        """Add point `xi`, direction `vi` and value `Li` with index `i` to the path."""
        self.samplingpath.add(i, x, v, L)

    def interpolate(self, i):
        """Interpolate point with index `i` on path."""
        return interpolate(
            i, self.samplingpath.points,
            fwd_possible=self.samplingpath.fwd_possible,
            rwd_possible=self.samplingpath.rwd_possible,
            contourpath=self)

    def extrapolate(self, i):
        """Advance beyond the current path, extrapolate from the end point.

        Parameters
        -----------
        i: int
            index on path.

        returns
        --------
        coords: vector
            coordinates of the new point.

        """
        if i >= 0:
            j, xj, vj, Lj = max(self.samplingpath.points)
            deltai = i - j
            assert deltai > 0, ("should be extrapolating", i, j)
        else:
            j, xj, vj, Lj = min(self.samplingpath.points)
            deltai = i - j
            assert deltai < 0, ("should be extrapolating", i, j)

        newpoint = extrapolate_ahead(deltai, xj, vj, contourpath=self)
        return newpoint

    def gradient(self, reflpoint, plot=False):
        """Compute gradient approximation.

        Finds spheres enclosing the `reflpoint`, and chooses their mean
        as the direction to go towards. If no spheres enclose the
        reflpoint, use nearest sphere.

        v is not used, because that would break detailed balance.

        Considerations:
           - in low-d, we want to focus on nearby live point spheres
             The border traced out is fairly accurate, at least in the
             normal away from the inside.

           - in high-d, reflpoint is contained by all live points,
             and none of them point to the center well. Because the
             sampling is poor, the "region center" position
             will be very stochastic.

        Parameters
        -----------
        reflpoint: vector
            point outside the likelihood contour, reflect there
        v: vector
            previous direction vector

        Returns
        ---------
        gradient: vector
            normal of ellipsoid

        """
        if plot:
            plt.plot(reflpoint[0], reflpoint[1], '+ ', color='k', ms=10)

        # check which the reflections the ellipses would make
        region = self.region
        bpts = region.transformLayer.transform(reflpoint.reshape((1,-1)))
        dist = ((bpts - region.unormed)**2).sum(axis=1)
        assert dist.shape == (len(region.unormed),), (dist.shape, len(region.unormed))
        nearby = dist < region.maxradiussq
        assert nearby.shape == (len(region.unormed),), (nearby.shape, len(region.unormed))
        if not nearby.any():
            nearby = dist == dist.min()
        sphere_centers = region.u[nearby,:]

        tsphere_centers = region.unormed[nearby,:]
        nlive, ndim = region.unormed.shape
        assert tsphere_centers.shape[1] == ndim, (tsphere_centers.shape, ndim)

        # choose mean among those points
        tsphere_center = tsphere_centers.mean(axis=0)
        assert tsphere_center.shape == (ndim,), (tsphere_center.shape, ndim)
        tt = get_sphere_tangent(tsphere_center, bpts.flatten())
        assert tt.shape == tsphere_center.shape, (tt.shape, tsphere_center.shape)

        # convert back to u space
        sphere_center = region.transformLayer.untransform(tsphere_center)
        t = region.transformLayer.untransform(tt * 1e-3 + tsphere_center) - sphere_center

        if plot:
            tt_all = get_sphere_tangent(tsphere_centers, bpts)
            t_all = region.transformLayer.untransform(tt_all * 1e-3 + tsphere_centers) - sphere_centers

            """
            # plot in transformed space too:
            origax = plt.gca()
            ax = plt.axes([0.7, 0.7, 0.27, 0.27])
            plt.plot(bpts[:,0], bpts[:,1], '+ ', color='k', ms=10)
            plt.plot(region.unormed[:,0], region.unormed[:,1], 'x ', color='k', ms=4)
            plt.plot(tsphere_centers[:,0], tsphere_centers[:,1], 'o ', mfc='None', mec='b', ms=10, mew=1)
            for si, ti in zip(tsphere_centers, tt_all):
                plt.plot([si[0], ti[0] + si[0]], [si[1], ti[1] + si[1]], '--', lw=2, color='gray', alpha=0.5)
            plt.plot(tsphere_center[0], tsphere_center[1], '^ ', mfc='None', mec='g', ms=10, mew=3)
            plt.plot([tsphere_center[0], tt[0] + tsphere_center[0]],
                [tsphere_center[1], tt[1] + tsphere_center[1]], lw=1, color='gray')
            plt.sca(origax)
            """

            plt.plot(sphere_centers[:,0], sphere_centers[:,1], 'o ', mfc='None', mec='b', ms=10, mew=1)
            if not (dist < region.maxradiussq).any():
                plt.plot(sphere_centers[:,0], sphere_centers[:,1], 's ', mfc='None', mec='b', ms=10, mew=1)
            for si, ti in zip(sphere_centers, t_all):
                plt.plot([si[0], ti[0] * 1000 + si[0]], [si[1], ti[1] * 1000 + si[1]], '--', lw=2, color='gray', alpha=0.5)
            plt.plot(sphere_center[0], sphere_center[1], '^ ', mfc='None', mec='g', ms=10, mew=3)
            plt.plot([sphere_center[0], t[0] * 1000 + sphere_center[0]], [sphere_center[1], t[1] * 1000 + sphere_center[1]], color='gray')

        # compute new vector
        normal = t / norm(t)
        isunitlength(normal)
        assert normal.shape == t.shape, (normal.shape, t.shape)

        return normal

1	"""Sparsely sampled, virtual sampling path.
2
3	Supports reflections at unit cube boundaries, and regions.
4	"""
5
6
7	import numpy as np	1✔
8	from numpy.linalg import norm	1✔
9	import matplotlib.pyplot as plt	1✔
10
11
12	def nearest_box_intersection_line(ray_origin, ray_direction, fwd=True):	1✔
13	r"""Compute intersection of a line (ray) and a unit box (0:1 in all axes).
14
15	Based on
16	http://www.iquilezles.org/www/articles/intersectors/intersectors.htm
17
18	To continue forward traversing at the reflection point use::
19
20	while True:
21	# update current point x
22	x, _, i = box_line_intersection(x, v)
23	# change direction
24	v[i] *= -1
25
26	Parameters
27	-----------
28	ray_origin: vector
29	starting point of line
30	ray_direction: vector
31	line direction vector
32
33	Returns
34	--------
35	p: vector
36	intersection point
37	t: float
38	intersection point distance from ray\_origin in units in ray\_direction
39	i: int
40	axes which change direction at pN
41
42	"""
43	# make sure ray starts inside the box
44	assert (ray_origin >= 0).all(), ray_origin	1✔
45	assert (ray_origin <= 1).all(), ray_origin	1✔
46	assert ((ray_direction2).sum()0.5 > 1e-200).all(), ray_direction	1✔
47
48	# step size
49	with np.errstate(divide='ignore', invalid='ignore'):	1✔
50	m = 1. / ray_direction	1✔
51	n = m * (ray_origin - 0.5)	1✔
52	k = np.abs(m) * 0.5	1✔
53	# line coordinates of intersection
54	# find first intersecting coordinate
55	if fwd:	1✔
56	t2 = -n + k	1✔
57	tF = np.nanmin(t2)	1✔
58	iF = np.where(t2 == tF)[0]	1✔
59	else:
60	t1 = -n - k	1✔
61	tF = np.nanmax(t1)	1✔
62	iF = np.where(t1 == tF)[0]	1✔
63
64	pF = ray_origin + ray_direction * tF	1✔
65	eps = 1e-6	1✔
66	assert (pF >= -eps).all(), (pF, ray_origin, ray_direction)	1✔
67	assert (pF <= 1 + eps).all(), (pF, ray_origin, ray_direction)	1✔
68	pF[pF < 0] = 0	1✔
69	pF[pF > 1] = 1	1✔
70	return pF, tF, iF	1✔
71
72
73	def box_line_intersection(ray_origin, ray_direction):	1✔
74	"""Find intersections of a line with the unit cube, in both sides.
75
76	Parameters
77	-----------
78	ray_origin: vector
79	starting point of line
80	ray_direction: vector
81	line direction vector
82
83	Returns
84	--------
85	left: nearest_box_intersection_line return value
86	from negative direction
87	right: nearest_box_intersection_line return value
88	from positive direction
89
90	"""
91	pF, tF, iF = nearest_box_intersection_line(ray_origin, ray_direction, fwd=True)	1✔
92	pN, tN, iN = nearest_box_intersection_line(ray_origin, ray_direction, fwd=False)	1✔
93	if tN > tF or tF < 0:	1!
94	assert False, "no intersection"	×
95	return (pN, tN, iN), (pF, tF, iF)	1✔
96
97
98	def linear_steps_with_reflection(ray_origin, ray_direction, t, wrapped_dims=None):	1✔
99	"""Go `t` steps in direction `ray_direction` from `ray_origin`.
100
101	Reflect off the unit cube if encountered, respecting wrapped dimensions.
102	In any case, the distance should be ``t * ray_direction``.
103
104	Returns
105	--------
106	new_point: vector
107	end point
108	new_direction: vector
109	new direction.
110
111	"""
112	if t == 0:	1✔
113	return ray_origin, ray_direction	1✔
114	if t < 0:	1✔
115	new_point, new_direction = linear_steps_with_reflection(ray_origin, -ray_direction, -t)	1✔
116	return new_point, -new_direction	1✔
117
118	if wrapped_dims is not None:	1✔
119	reflected = np.zeros(len(ray_origin), dtype=bool)	1✔
120
121	tleft = 1.0 * t	1✔
122	while True:
123	p, t, i = nearest_box_intersection_line(ray_origin, ray_direction, fwd=True)	1✔
124	# print(p, t, i, ray_origin, ray_direction)
125	assert np.isfinite(p).all()	1✔
126	assert t >= 0, t	1✔
127	if tleft <= t: # stopping before reaching any border	1✔
128	assert np.all(ray_origin + tleft * ray_direction >= 0), (ray_origin, tleft, ray_direction)	1✔
129	assert np.all(ray_origin + tleft * ray_direction <= 1), (ray_origin, tleft, ray_direction)	1✔
130	return ray_origin + tleft * ray_direction, ray_direction	1✔
131	# go to reflection point
132	ray_origin = p	1✔
133	assert np.isfinite(ray_origin).all(), ray_origin	1✔
134	# reflect
135	ray_direction = ray_direction.copy()	1✔
136	if wrapped_dims is None:	1✔
137	ray_direction[i] *= -1	1✔
138	else:
139	# if we already once bumped into that (wrapped) axis,
140	# do not continue but return this as end point
141	if np.logical_and(reflected[i], wrapped_dims[i]).any():	1✔
142	return ray_origin, ray_direction	1✔
143
144	# note which axes we already flipped
145	reflected[i] = True	1✔
146
147	# in wrapped axes, we can keep going. Otherwise, reflects
148	ray_direction[i] *= np.where(wrapped_dims[i], 1, -1)	1✔
149
150	# in the i axes, we should wrap the coordinates
151	assert np.logical_or(np.isclose(ray_origin[i], 1), np.isclose(ray_origin[i], 0)).all(), ray_origin[i]	1✔
152	ray_origin[i] = np.where(wrapped_dims[i], 1 - ray_origin[i], ray_origin[i])	1✔
153
154	assert np.isfinite(ray_direction).all(), ray_direction	1✔
155	# reduce remaining distance
156	tleft -= t	1✔
157
158
159	def get_sphere_tangent(sphere_center, edge_point):	1✔
160	"""Compute tangent at sphere surface point.
161
162	Assume a sphere centered at sphere_center with radius
163	so that edge_point is on the surface. At edge_point, in
164	which direction does the normal vector point?
165
166	Parameters
167	-----------
168	sphere_center: vector
169	center of sphere
170	edge_point: vector
171	point at the surface
172
173	Returns
174	--------
175	tangent: vector
176	vector pointing to the sphere center.
177
178	"""
179	arrow = sphere_center - edge_point	1✔
180	return arrow / norm(arrow)	1✔
181
182
183	def get_sphere_tangents(sphere_center, edge_point):	1✔
184	"""Compute tangent at sphere surface point.
185
186	Assume a sphere centered at sphere_center with radius
187	so that edge_point is on the surface. At edge_point, in
188	which direction does the normal vector point?
189
190	This function is vectorized and handles arrays of arguments.
191
192	Parameters
193	-----------
194	sphere_center: array
195	centers of spheres
196	edge_point: array
197	points at the surface
198
199	Returns
200	--------
201	tangent: array
202	vectors pointing to the sphere center.
203
204	"""
205	arrow = sphere_center - edge_point	1✔
206	return arrow / norm(arrow, axis=1).reshape((-1, 1))	1✔
207
208
209	def reflect(v, normal):	1✔
210	"""Reflect vector ``v`` off a ``normal`` vector, return new direction vector."""
211	return v - 2 * (normal * v).sum() * normal	×
212
213
214	def distances(direction, center, r=1):	1✔
215	"""Compute sphere-line intersection.
216
217	Parameters
218	-----------
219	direction: vector
220	direction vector (line starts at 0)
221	center: vector
222	center of sphere (coordinate vector)
223	r: float
224	radius of sphere
225
226	Returns
227	--------
228	tpos, tneg: floats
229	the positive and negative coordinate along the `l` vector where `r` is intersected.
230	If no intersection, throws AssertError.
231
232	"""
233	loc = (direction * center).sum()	×
234	osqrnorm = (center**2).sum()	×
235	# print(loc.shape, loc.shape, osqrnorm.shape)
236	rootterm = loc2 - osqrnorm + r2	×
237	# make sure we are crossing the sphere
238	assert (rootterm > 0).all(), rootterm	×
239	return -loc + rootterm0.5, -loc - rootterm0.5	×
240
241
242	def isunitlength(vec):	1✔
243	"""Verify that `vec` is of unit length."""
244	assert np.isclose(norm(vec), 1), norm(vec)	1✔
245
246
247	def angle(a, b):	1✔
248	"""Compute dot product between vectors `a` and `b`.
249
250	The arccos of the return value would give an actual angle.
251	"""
252	# anorm = (a2).sum()0.5
253	# bnorm = (b2).sum()0.5
254	return (a * b).sum() # / anorm / bnorm	1✔
255
256
257	def extrapolate_ahead(dj, xj, vj, contourpath=None):	1✔
258	"""Make `di` steps of size `vj` from `xj`.
259
260	Reflect off unit cube if necessary.
261	"""
262	assert dj == int(dj)	1✔
263
264	# optimistically try to go there directly
265
266	xk, vk = linear_steps_with_reflection(xj, vj, dj)	1✔
267
268	return xk, vk # deactivate feature below	1✔
269
270	if contourpath is None:
271	return xk, vk
272
273	# check if we can already tell that the point will be outside
274	region = contourpath.region	×
275	if contourpath.region.inside(xk.reshape((1, -1))):	×
276	return xk, vk	×
277
278	# find first point outside region
279	sign = 1 if dj > 0 else -1	×
280	d = np.arange(0, dj, sign)	×
281	first_point_outside = dj, xk, vk	×
282	for di in d:	×
283	xi, vi = linear_steps_with_reflection(xj, vj, di)	×
284	if not region.inside(xk.reshape((1, -1))):	×
285	first_point_outside = di, xi, vi	×
286	break	×
287
288	# reflect at this point (first outside)
289	dout, reflpoint, v = first_point_outside	×
290
291	if dout == 0:	×
292	# already the starting point is outside.
293	# return extrapolation and hope caller handles it
294	return xk, vk	×
295
296	# reversing:
297	normal = contourpath.gradient(reflpoint) # , v * sign)	×
298	if normal is None:	×
299	vnew = -v	×
300	else:
301	vnew = (v - 2 * angle(normal, v) * normal) * sign	×
302	assert vnew.shape == v.shape, (vnew.shape, v.shape)	×
303	assert np.isclose(norm(vnew), norm(v)), (vnew, v, norm(vnew), norm(v))	×
304
305	# make one step (xl replaces first_point_outside/reflpoint)
306	xl, vl = linear_steps_with_reflection(reflpoint, vnew, sign)	×
307	# how many steps are still to do?
308	dleft = dj - dout	×
309
310	# make that many step in that direction, by recursing.
311
312	# it is possible that this point is also outside. The next iteration
313	# will catch that case.
314
315	return extrapolate_ahead(dleft, xl, vl, contourpath=contourpath)	×
316
317
318	def interpolate(i, points, fwd_possible, rwd_possible, contourpath=None):	1✔
319	"""Interpolate a point on the path indicated by `points`.
320
321	Given a sparsely sampled track (stored in .points),
322	potentially encountering reflections,
323	extract the corrdinates of the point with index `i`.
324	That point may not have been evaluated yet.
325
326	Parameters
327	-----------
328	i: int
329	position on track to return.
330	points: list of tuples (index, coordinate, direction, loglike)
331	points on the path
332	fwd_possible: bool
333	whether the path could be extended in the positive direction.
334	rwd_possible: bool
335	whether the path could be extended in the negative direction.
336	contourpath: ContourPath
337	Use region to reflect. Not used at the moment.
338
339	"""
340	points_before = [(j, xj, vj, Lj) for j, xj, vj, Lj in points if j <= i]	1✔
341	points_after = [(j, xj, vj, Lj) for j, xj, vj, Lj in points if j >= i]	1✔
342
343	# check if the point after is really after i
344	if len(points_after) == 0 and not fwd_possible:	1✔
345	# the path cannot continue, and i does not exist.
346	# print(" interpolate_point %d: the path cannot continue fwd, and i does not exist." % i)
347	j, xj, vj, Lj = max(points_before)	1✔
348	return xj, vj, Lj, False	1✔
349
350	# check if the point before is really before i
351	if len(points_before) == 0 and not rwd_possible:	1✔
352	# the path cannot continue, and i does not exist.
353	k, xk, vk, Lk = min(points_after)	1✔
354	# print(" interpolate_point %d: the path cannot continue rwd, and i does not exist." % i)
355	return xk, vk, Lk, False	1✔
356
357	if len(points_before) == 0 or len(points_after) == 0:	1✔
358	# return None, None, None, False
359	raise KeyError("cannot extrapolate outside path")	1✔
360
361	j, xj, vj, Lj = max(points_before)	1✔
362	k, xk, vk, Lk = min(points_after)	1✔
363
364	# print(" interpolate_point %d between %d-%d" % (i, j, k))
365	if j == i: # we have this exact point in the chain	1✔
366	return xj, vj, Lj, True	1✔
367
368	assert not k == i # otherwise the above would be true too	1✔
369
370	# expand_to_step explores each reflection in detail, so
371	# any points with change in v should have j == i
372	# therefore we can assume:
373	# assert (vj == vk).all()
374	# this ^ is not true, because reflections on the unit cube can
375	# occur, and change v without requiring a intermediate point.
376
377	# j....i....k
378	xl1, vj1 = extrapolate_ahead(i - j, xj, vj, contourpath=contourpath)	1✔
379	xl2, vj2 = extrapolate_ahead(i - k, xk, vk, contourpath=contourpath)	1✔
380	assert np.allclose(xl1, xl2), (xl1, xl2, i, j, k, xj, vj, xk, vk)	1✔
381	assert np.allclose(vj1, vj2), (xl1, vj1, xl2, vj2, i, j, k, xj, vj, xk, vk)	1✔
382	xl = xl1	1✔
383
384	# xl = interpolate_between_two_points(i, xj, j, xk, k)
385	# the new point is then just a linear interpolation
386	# w = (i - k) * 1. / (j - k)
387	# xl = xj * w + (1 - w) * xk
388
389	return xl, vj, None, True	1✔
390
391
392	class SamplingPath(object):	1✔
393	"""Path described by a (potentially sparse) sequence of points.
394
395	Convention of the stored point tuple ``(i, x, v, L)``:
396	`i`: index (0 is starting point)
397	`x`: point
398	`v`: direction
399	`L`: loglikelihood value
400	"""
401
402	def __init__(self, x0, v0, L0):	1✔
403	"""Initialise with path starting point.
404
405	Starting point (`x0`), direction (`v0`) and
406	loglikelihood value (`L0`) of the path. Is given index 0.
407	"""
408	self.reset(x0, v0, L0)	1✔
409
410	def add(self, i, xi, vi, Li):	1✔
411	"""Add point `xi`, direction `vi` and value `Li` with index `i` to the path."""
412	assert Li is not None	1✔
413	assert len(xi.shape) == 1, (xi, xi.shape)	1✔
414	assert len(vi.shape) == 1, (vi, vi.shape)	1✔
415	assert len(np.shape(Li)) == 0, (Li, Li.shape)	1✔
416	self.points.append((i, xi, vi, Li))	1✔
417
418	def reset(self, x0, v0, L0):	1✔
419	"""Reset path, start from ``x0, v0, L0``."""
420	self.points = []	1✔
421	self.add(0, x0, v0, L0)	1✔
422	self.fwd_possible = True	1✔
423	self.rwd_possible = True	1✔
424
425	def plot(self, **kwargs):	1✔
426	"""Plot the current path.
427
428	Only uses first two dimensions.
429	"""
430	x = np.array([x for i, x, v, L in sorted(self.points)])	×
431	p, = plt.plot(x[:,0], x[:,1], 'o ', **kwargs)	×
432	ilo, _, _, _ = min(self.points)	×
433	ihi, _, _, _ = max(self.points)	×
434	x = np.array([self.interpolate(i)[0] for i in range(ilo, ihi + 1)])	×
435	kwargs['color'] = p.get_color()	×
436	plt.plot(x[:,0], x[:,1], 'o-', ms=4, mfc='None', **kwargs)	×
437
438	def interpolate(self, i):	1✔
439	"""Interpolate point with index `i` on path."""
440	return interpolate(i, self.points, fwd_possible=self.fwd_possible, rwd_possible=self.rwd_possible)	1✔
441
442	def extrapolate(self, i):	1✔
443	"""Advance beyond the current path, extrapolate from the end point.
444
445	Parameters
446	-----------
447	i: int
448	index on path.
449
450	returns
451	--------
452	coords: vector
453	coordinates of the new point.
454
455	"""
456	if i >= 0:	1✔
457	j, xj, vj, Lj = max(self.points)	1✔
458	deltai = i - j	1✔
459	assert deltai > 0, ("should be extrapolating", i, j)	1✔
460	else:
461	j, xj, vj, Lj = min(self.points)	1✔
462	deltai = i - j	1✔
463	assert deltai < 0, ("should be extrapolating", i, j)	1✔
464
465	newpoint = extrapolate_ahead(deltai, xj, vj)	1✔
466	return newpoint	1✔
467
468
469	class ContourSamplingPath(object):	1✔
470	"""Region-aware form of the sampling path.
471
472	Uses region points to guess a likelihood contour gradient.
473	"""
474
475	def __init__(self, samplingpath, region):	1✔
476	"""Initialise with `samplingpath` and `region`."""
477	self.samplingpath = samplingpath	1✔
478	self.points = self.samplingpath.points	1✔
479	self.region = region	1✔
480
481	def add(self, i, x, v, L):	1✔
482	"""Add point `xi`, direction `vi` and value `Li` with index `i` to the path."""
483	self.samplingpath.add(i, x, v, L)	1✔
484
485	def interpolate(self, i):	1✔
486	"""Interpolate point with index `i` on path."""
487	return interpolate(	1✔
488	i, self.samplingpath.points,
489	fwd_possible=self.samplingpath.fwd_possible,
490	rwd_possible=self.samplingpath.rwd_possible,
491	contourpath=self)
492
493	def extrapolate(self, i):	1✔
494	"""Advance beyond the current path, extrapolate from the end point.
495
496	Parameters
497	-----------
498	i: int
499	index on path.
500
501	returns
502	--------
503	coords: vector
504	coordinates of the new point.
505
506	"""
507	if i >= 0:	1✔
508	j, xj, vj, Lj = max(self.samplingpath.points)	1✔
509	deltai = i - j	1✔
510	assert deltai > 0, ("should be extrapolating", i, j)	1✔
511	else:
512	j, xj, vj, Lj = min(self.samplingpath.points)	1✔
513	deltai = i - j	1✔
514	assert deltai < 0, ("should be extrapolating", i, j)	1✔
515
516	newpoint = extrapolate_ahead(deltai, xj, vj, contourpath=self)	1✔
517	return newpoint	1✔
518
519	def gradient(self, reflpoint, plot=False):	1✔
520	"""Compute gradient approximation.
521
522	Finds spheres enclosing the `reflpoint`, and chooses their mean
523	as the direction to go towards. If no spheres enclose the
524	reflpoint, use nearest sphere.
525
526	v is not used, because that would break detailed balance.
527
528	Considerations:
529	- in low-d, we want to focus on nearby live point spheres
530	The border traced out is fairly accurate, at least in the
531	normal away from the inside.
532
533	- in high-d, reflpoint is contained by all live points,
534	and none of them point to the center well. Because the
535	sampling is poor, the "region center" position
536	will be very stochastic.
537
538	Parameters
539	-----------
540	reflpoint: vector
541	point outside the likelihood contour, reflect there
542	v: vector
543	previous direction vector
544
545	Returns
546	---------
547	gradient: vector
548	normal of ellipsoid
549
550	"""
551	if plot:	1!
552	plt.plot(reflpoint[0], reflpoint[1], '+ ', color='k', ms=10)	×
553
554	# check which the reflections the ellipses would make
555	region = self.region	1✔
556	bpts = region.transformLayer.transform(reflpoint.reshape((1,-1)))	1✔
557	dist = ((bpts - region.unormed)**2).sum(axis=1)	1✔
558	assert dist.shape == (len(region.unormed),), (dist.shape, len(region.unormed))	1✔
559	nearby = dist < region.maxradiussq	1✔
560	assert nearby.shape == (len(region.unormed),), (nearby.shape, len(region.unormed))	1✔
561	if not nearby.any():	1✔
562	nearby = dist == dist.min()	1✔
563	sphere_centers = region.u[nearby,:]	1✔
564
565	tsphere_centers = region.unormed[nearby,:]	1✔
566	nlive, ndim = region.unormed.shape	1✔
567	assert tsphere_centers.shape[1] == ndim, (tsphere_centers.shape, ndim)	1✔
568
569	# choose mean among those points
570	tsphere_center = tsphere_centers.mean(axis=0)	1✔
571	assert tsphere_center.shape == (ndim,), (tsphere_center.shape, ndim)	1✔
572	tt = get_sphere_tangent(tsphere_center, bpts.flatten())	1✔
573	assert tt.shape == tsphere_center.shape, (tt.shape, tsphere_center.shape)	1✔
574
575	# convert back to u space
576	sphere_center = region.transformLayer.untransform(tsphere_center)	1✔
577	t = region.transformLayer.untransform(tt * 1e-3 + tsphere_center) - sphere_center	1✔
578
579	if plot:	1!
580	tt_all = get_sphere_tangent(tsphere_centers, bpts)	×
581	t_all = region.transformLayer.untransform(tt_all * 1e-3 + tsphere_centers) - sphere_centers	×
582
583	"""
584	# plot in transformed space too:
585	origax = plt.gca()
586	ax = plt.axes([0.7, 0.7, 0.27, 0.27])
587	plt.plot(bpts[:,0], bpts[:,1], '+ ', color='k', ms=10)
588	plt.plot(region.unormed[:,0], region.unormed[:,1], 'x ', color='k', ms=4)
589	plt.plot(tsphere_centers[:,0], tsphere_centers[:,1], 'o ', mfc='None', mec='b', ms=10, mew=1)
590	for si, ti in zip(tsphere_centers, tt_all):
591	plt.plot([si[0], ti[0] + si[0]], [si[1], ti[1] + si[1]], '--', lw=2, color='gray', alpha=0.5)
592	plt.plot(tsphere_center[0], tsphere_center[1], '^ ', mfc='None', mec='g', ms=10, mew=3)
593	plt.plot([tsphere_center[0], tt[0] + tsphere_center[0]],
594	[tsphere_center[1], tt[1] + tsphere_center[1]], lw=1, color='gray')
595	plt.sca(origax)
596	"""
597
598	plt.plot(sphere_centers[:,0], sphere_centers[:,1], 'o ', mfc='None', mec='b', ms=10, mew=1)	×
599	if not (dist < region.maxradiussq).any():	×
600	plt.plot(sphere_centers[:,0], sphere_centers[:,1], 's ', mfc='None', mec='b', ms=10, mew=1)	×
601	for si, ti in zip(sphere_centers, t_all):	×
602	plt.plot([si[0], ti[0] * 1000 + si[0]], [si[1], ti[1] * 1000 + si[1]], '--', lw=2, color='gray', alpha=0.5)	×
603	plt.plot(sphere_center[0], sphere_center[1], '^ ', mfc='None', mec='g', ms=10, mew=3)	×
604	plt.plot([sphere_center[0], t[0] * 1000 + sphere_center[0]], [sphere_center[1], t[1] * 1000 + sphere_center[1]], color='gray')	×
605
606	# compute new vector
607	normal = t / norm(t)	1✔
608	isunitlength(normal)	1✔
609	assert normal.shape == t.shape, (normal.shape, t.shape)	1✔
610
611	return normal	1✔

JohannesBuchner / UltraNest / 9f2dd4f6-0775-47e9-b700-af647027ebfa

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous