Source code for refl1d.util

__all__ = ["merge_ends"]

import sys

import numpy as np

# CRUFT: 2.7 support
if sys.version_info[0] > 2:
    def asbytes(s):
        return s.encode('utf-8')
else:
    def asbytes(s):
        return s



def merge_ends(w, p, tol=1e-3):
    """
    join the leading and trailing ends of the profile together so fewer
    slabs are required and so that gaussian roughness can be used.
    """
    # TODO: accept rho, rhoi pairs as well
    # TODO: make sure we apply an interface to the right as well as the left
    try:
        # Assuming there p[0] != p[-1] within tolerance, we are guaranteed
        # that we will have a first value not equivalent to p[0] on the
        # left, with index > 0 and a last value not equivalent to p[-1] on
        # the right, with index < -1.  We are going to put the first value
        # at left index - 1 and the last value at right index + 1, accumulating
        # the widths of the identical layers.
        lidx = np.where(abs(p-p[0]) > tol)[0][0]-1
        ridx = len(p) - np.where(abs(p[::-1]-p[-1]) > tol)[0][0]
        w[lidx], p[lidx] = np.sum(w[:lidx+1]), p[0]
        w[ridx], p[ridx] = np.sum(w[ridx:]), p[-1]
        return w[lidx:ridx+1], p[lidx:ridx+1]
    except Exception:
        if len(w):
            # All one big layer
            w[0] = np.sum(w)
            return w[0:1], p[0:1]
        else:
            return w, p