<commit_before>import numpy as np
from scipy.special import erf
from scipy.weave import inline

from cbamf.util import Tile, cdd, amin, amax

#=============================================================================
# Forms of the platonic sphere interpolation function
#=============================================================================
def sphere_lerp(r, r0, alpha):
    """ Linearly interpolate the pixels for the platonic object """
    return (1-np.clip((r-(r0-alpha)) / (2*alpha), 0, 1))

def sphere_logistic(r, r0, alpha):
    """ Classic logistic interpolation """
    return 1.0/(1.0 + np.exp(alpha*(r - r0)))

def sphere_triangle_cdf(r, r0, alpha):
    """ Cumulative distribution function for the traingle distribution """
    p0 = (r-r0+alpha)**2/(2*alpha**2)*(r0 > r)*(r>r0-alpha)
    p1 = 1*(r>r0)-(r0+alpha-r)**2/(2*alpha**2)*(r0<r)*(r<r0+alpha)
    return (1-np.clip(p0+p1, 0, 1))

def sphere_analytical_gaussian_exact(r, r0, alpha=0.2765):
    """
    Analytically calculate the sphere's functional form by convolving the
    Heavyside function with first order approximation to the sinc, a Gaussian.
    The alpha parameters controls the width of the approximation -- should be
    1, but is fit to be roughly 0.2765
    """
    term1 = 0.5*(erf((r0+r)/(alpha*np.sqrt(2))) + erf((r0-r)/(alpha*np.sqrt(2))))
    term2 = np.sqrt(0.5/np.pi)*(alpha/(r+1e-10)) * (
                np.exp(-0.5*(r-r0)**2/alpha**2) - np.exp(-0.5*(r+r0)**2/alpha**2)
            )
    return term1 - term2

def sphere_analytical_gaussian_trim(r, r0, alpha=0.2765, cut=1.6):
    """
    See sphere_analytical_gaussian_exact.

    I trimmed to terms from the functional form that are essentially zero (1e-8)
    for r0 > cut (~1.5), a fine approximation for these platonic anyway.
    """
    dr = r-r0
    m = np.abs(dr) <= cut

    # only compute on the relevant scales
    rr = r[m]
    t = (r0-rr)/(alpha*np.sqrt(2))
    q = 0.5*(1 + erf(t)) - np.sqrt(0.5/np.pi)*(alpha/(rr+1e-10)) * np.exp(-t*t)

    # fill in the grid, inside the interpolation and outside where values are constant
    ans = 0*r
    ans[m] = q
    ans[dr >  cut] = 0
    ans[dr < -cut] = 1
    return ans

def sphere_analytical_gaussian_fast(r, r0, alpha=0.2765, cut=1.20):
    """
    See sphere_analytical_gaussian_trim, but implemented in C with
    fast erf and exp approximations found at
        Abramowitz and Stegun: Handbook of Mathematical Functions
        A Fast, Compact Approximation of the Exponential Function

    The default cut 1.25 was chosen based on the accuracy of fast_erf
    """

    functions = """
    double fast_erf(double x){
        double sgn = 1.0;

        if (x < 0){
            sgn = -1.0;
            x = -x;
        }

        double p = 0.47047;
        double a1 =  0.3480242;
        double a2 = -0.0958798;
        double a3 =  0.7478556;
        double t1 = 1.0/(1 + p*x);
        double t2 = t1*t1;
        double t3 = t1*t2;
        return sgn*(1 - (a1*t1 + a2*t2 + a3*t3)*exp(-x*x));
    }

    static union
    {
        double d;
        struct
        {
            #ifdef LITTLE_ENDIAN
                int j, i;
            #else
                int i, j;
            #endif
        } n;
    } _eco;

    #define EXP_A (1048576 /M_LN2)
    #define EXP_C 60801
    #define fast_exp(y) (_eco.n.i = EXP_A*(y) + (1072693248 - EXP_C), _eco.d)
    """

    code = """
    double coeff1 = 1.0/(alpha*sqrt(2.0));
    double coeff2 = sqrt(0.5/pi)*alpha;

    for (int i=0; i<N; i++){
        double dr = r[i]-r0;
        if (dr < cut && dr > -cut){
            double t = (r0 - r[i])*coeff1;
            ans[i] = 0.5*(1+fast_erf(t)) - coeff2/(r[i]+1e-10) * fast_exp(-t*t);
        } else {
            ans[i] = 0.0*(dr > cut) + 1.0*(dr < -cut);
        }
    }
    """

    shape = r.shape
    r = r.flatten()
    N = r.shape[0]
    ans = r*0
    pi = np.pi

    inline(code, arg_names=['r', 'r0', 'alpha', 'cut', 'ans', 'pi', 'N'],
            support_code=functions, verbose=0)
    return ans.reshape(shape)

def sphere_constrained_cubic(r, r0, alpha):
    """
    Sphere generated by a cubic interpolant constrained to be (1,0) on
    (r0-sqrt(3)/2, r0+sqrt(3)/2), the size of the cube in the (111) direction.
    """
    sqrt3 = np.sqrt(3)

    b_coeff = r0*0.5/sqrt3*(1 - 0.6*sqrt3*alpha)/(0.15 + r0*r0)
    rscl = np.clip(r-r0, -0.5*sqrt3, 0.5*sqrt3)

    a, d = rscl + 0.5*sqrt3, rscl - 0.5*sqrt3
    return alpha*d*a*rscl + b_coeff*d*a - d/sqrt3

try:
    sphere_analytical_gaussian_fast(np.linspace(0,10,10), 5.0)
except Exception as e:
    sphere_analytical_gaussian_fast = sphere_analytical_gaussian_trim

#=============================================================================
# Actual sphere collection (and slab)
#=============================================================================
class SphereCollectionRealSpace(object):
    def __init__(self, pos, rad, shape, support_size=4, typ=None, pad=None,
                 method='exact-gaussian-fast', alpha=None, method_function=None):
        """
        method can be one of:
            [
                'lerp', 'logistic', 'triangle', 'constrained-cubic',
                'exact-gaussian', 'exact-gaussian-trim', 'exact-gaussian-fast'
            ]
        """
        self.support_size = support_size
        self.pos = pos.astype('float')
        self.rad = rad.astype('float')
        self.N = rad.shape[0]

        # set the aliasing method and coefficient
        # FIXME -- check if function for method and set to that instead
        # method='user-defined' and pass method_function
        self._setup_sphere_functions(method=method, alpha=alpha)

        if typ is None:
            self.typ = np.ones(self.N)
            if pad is not None and pad <= self.N:
                self.typ[-pad:] = 0
        else:
            self.typ = typ.astype('float')

        self.shape = shape
        self._setup()

    def _setup_sphere_functions(self, method, alpha=None):
        self.methods = [
            'lerp', 'logistic', 'triangle', 'constrained-cubic',
            'exact-gaussian', 'exact-gaussian-trim', 'exact-gaussian-fast'
        ]

        self.sphere_functions = {
            'lerp': sphere_lerp,
            'logistic': sphere_logistic,
            'triangle': sphere_triangle_cdf,
            'exact-gaussian': sphere_analytical_gaussian_exact,
            'exact-gaussian-trim': sphere_analytical_gaussian_trim,
            'exact-gaussian-fast': sphere_analytical_gaussian_fast,
            'constrained-cubic': sphere_constrained_cubic
        }

        self.alpha_defaults = {
            'lerp': 0.4539,
            'logistic': 6.5,
            'triangle': 0.6618,
            'exact-gaussian': 0.27595,
            'exact-gaussian-trim': 0.27595,
            'exact-gaussian-fast': 0.27595,
            'constrained-cubic': 0.84990,
        }

        self.method = method
        self.alpha = alpha if alpha is not None else self.alpha_defaults[self.method]

    def _setup(self):
        z,y,x = Tile(self.shape).coords()
        self.rvecs = np.rollaxis(np.array(np.broadcast_arrays(z,y,x)), 0, 4)
        self.particles = np.zeros(self.shape)
        self._diff_field = np.zeros(self.shape)

    def _particle(self, pos, rad, zscale, sign=1, dodiff=False):
        p = np.round(pos)
        r = np.round(np.array([1.0/zscale,1,1])*np.ceil(rad)+self.support_size)

        tile = Tile(p-r, p+r, 0, self.shape)
        subr = self.rvecs[tile.slicer + (np.s_[:],)]
        rvec = (subr - pos)

        # apply the zscale and find the distances to make a ellipsoid
        rvec[...,0] *= zscale
        rdist = np.sqrt((rvec**2).sum(axis=-1))

        # keep backwards compatibility for save files with older logistic coefficient
        # FIXME -- this is a nasty hack which should be fazed out when older save
        # files are deleted
        if not hasattr(self, 'method'):
            self._setup_sphere_functions('logistic', 5.0)
        if not hasattr(self, 'alpha'):
            self._setup_sphere_functions(self.method)

        # calculate the anti-aliasing according to the interpolation type
        t = sign*self.sphere_functions[self.method](rdist, rad, self.alpha)
        self.particles[tile.slicer] += t

        if dodiff:
            self._diff_field[tile.slicer] += t

    def _update_particle(self, n, p, r, t, zscale, dodiff=True):
        if self.typ[n] == 1:
            self._particle(self.pos[n], self.rad[n], zscale, -1, dodiff=dodiff)

        self.pos[n] = p
        self.rad[n] = r
        self.typ[n] = t

        if self.typ[n] == 1:
            self._particle(self.pos[n], self.rad[n], zscale, +1, dodiff=dodiff)

    def initialize(self, zscale):
        if len(self.pos.shape) != 2:
            raise AttributeError("Position array needs to be (-1,3) shaped, (z,y,x) order")

        self.particles = np.zeros(self.shape)
        for p0, r0, t0 in zip(self.pos, self.rad, self.typ):
            if t0 == 1:
                self._particle(p0, r0, zscale)

    def set_tile(self, tile):
        self.tile = tile

    def update(self, ns, pos, rad, typ, zscale, difference=True):
        for n, p, r, t in zip(ns, pos, rad, typ):
            self._update_particle(n, p, r, t, zscale, dodiff=difference)

    def get_field(self):
        return self.particles[self.tile.slicer]

    def get_diff_field(self):
        c = self._diff_field[self.tile.slicer].copy()
        self._diff_field[self.tile.slicer] *= 0
        return c

    def get_support_size(self, p0, r0, t0, p1, r1, t1, zscale):
        rsc = self.support_size

        zsc = np.array([1.0/zscale, 1, 1])
        r0, r1 = zsc*r0, zsc*r1

        off0 = r0 + rsc
        off1 = r1 + rsc

        if t0[0] == 1 and t1[0] == 1:
            pl = amin(p0-off0-1, p1-off1-1)
            pr = amax(p0+off0+1, p1+off1+1)
        if t0[0] != 1 and t1[0] == 1:
            pl = (p1-off1-1)
            pr = (p1+off1+1)
        if t0[0] == 1 and t1[0] != 1:
            pl = (p0-off0-1)
            pr = (p0+off0+1)
        if t0[0] != 1 and t1[0] != 1:
            c = np.array(self.shape)
            pl = c/2 - c/8
            pr = c/2 + c/8

        if len(pl.shape) > 1:
            pl = pl[0]
            pr = pr[0]
        return pl, pr

    def get_params(self):
        return np.hstack([self.pos.ravel(), self.rad])

    def get_params_pos(self):
        return self.pos.ravel()

    def get_params_rad(self):
        return self.rad

    def get_params_typ(self):
        return self.typ

    def __getstate__(self):
        odict = self.__dict__.copy()
        cdd(odict, ['rvecs', 'particles', '_diff_field'])
        return odict

    def __setstate__(self, idict):
        self.__dict__.update(idict)
        self._setup()


class Slab(object):
    def __init__(self, zpos, shape, normal=(1,0,0)):
        self.zpos = float(zpos)
        self.normal = np.array(normal).astype('float')
        self.normal /= np.sqrt(self.normal.dot(self.normal))

        self.shape = shape
        self._setup()

    def _setup(self):
        z,y,x = Tile(self.shape).coords()
        self.rvecs = np.rollaxis(np.array(np.broadcast_arrays(z,y,x)), 0, 4)
        self.image = np.zeros(self.shape)

    def _slab(self, zpos, norm, sign=1):
        # for the position at zpos, and the center in the x-y plane
        pos = np.array([zpos, self.shape[1]/2, self.shape[2]/2])

        p = (self.rvecs - pos).dot(norm)
        self.image = 1.0/(1.0 + np.exp(7*p))

    def initialize(self):
        self.image = np.zeros(self.shape)
        self._slab(self.zpos, self.normal)

    def set_tile(self, tile):
        self.tile = tile

    def update(self, params):
        zpos, norm = params[0], params[1:]
        norm = norm / np.sqrt(norm.dot(norm))

        self.zpos = zpos
        self.normal = norm

        self._slab(self.zpos, self.normal, +1)

    def get_field(self):
        return self.image[self.tile.slicer]

    def get_support_size(self, p=None):
        return np.zeros(3), np.array(self.shape)

    def get_params(self):
        return np.hstack([self.zpos, self.normal.ravel()])

    def __getstate__(self):
        odict = self.__dict__.copy()
        cdd(odict, ['rvecs', 'image'])
        return odict

    def __setstate__(self, idict):
        self.__dict__.update(idict)
        self._setup()

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return "{} <{}, {}>".format(str(self.__class__.__name__), self.zpos, list(self.normal))
<commit_msg>changing name of analytical gaussian exact for compatibility<commit_after>import numpy as np
from scipy.special import erf
from scipy.weave import inline

from cbamf.util import Tile, cdd, amin, amax

#=============================================================================
# Forms of the platonic sphere interpolation function
#=============================================================================
def sphere_lerp(r, r0, alpha):
    """ Linearly interpolate the pixels for the platonic object """
    return (1-np.clip((r-(r0-alpha)) / (2*alpha), 0, 1))

def sphere_logistic(r, r0, alpha):
    """ Classic logistic interpolation """
    return 1.0/(1.0 + np.exp(alpha*(r - r0)))

def sphere_triangle_cdf(r, r0, alpha):
    """ Cumulative distribution function for the traingle distribution """
    p0 = (r-r0+alpha)**2/(2*alpha**2)*(r0 > r)*(r>r0-alpha)
    p1 = 1*(r>r0)-(r0+alpha-r)**2/(2*alpha**2)*(r0<r)*(r<r0+alpha)
    return (1-np.clip(p0+p1, 0, 1))

def sphere_analytical_gaussian(r, r0, alpha=0.2765):
    """
    Analytically calculate the sphere's functional form by convolving the
    Heavyside function with first order approximation to the sinc, a Gaussian.
    The alpha parameters controls the width of the approximation -- should be
    1, but is fit to be roughly 0.2765
    """
    term1 = 0.5*(erf((r0+r)/(alpha*np.sqrt(2))) + erf((r0-r)/(alpha*np.sqrt(2))))
    term2 = np.sqrt(0.5/np.pi)*(alpha/(r+1e-10)) * (
                np.exp(-0.5*(r-r0)**2/alpha**2) - np.exp(-0.5*(r+r0)**2/alpha**2)
            )
    return term1 - term2

def sphere_analytical_gaussian_trim(r, r0, alpha=0.2765, cut=1.6):
    """
    See sphere_analytical_gaussian_exact.

    I trimmed to terms from the functional form that are essentially zero (1e-8)
    for r0 > cut (~1.5), a fine approximation for these platonic anyway.
    """
    dr = r-r0
    m = np.abs(dr) <= cut

    # only compute on the relevant scales
    rr = r[m]
    t = (r0-rr)/(alpha*np.sqrt(2))
    q = 0.5*(1 + erf(t)) - np.sqrt(0.5/np.pi)*(alpha/(rr+1e-10)) * np.exp(-t*t)

    # fill in the grid, inside the interpolation and outside where values are constant
    ans = 0*r
    ans[m] = q
    ans[dr >  cut] = 0
    ans[dr < -cut] = 1
    return ans

def sphere_analytical_gaussian_fast(r, r0, alpha=0.2765, cut=1.20):
    """
    See sphere_analytical_gaussian_trim, but implemented in C with
    fast erf and exp approximations found at
        Abramowitz and Stegun: Handbook of Mathematical Functions
        A Fast, Compact Approximation of the Exponential Function

    The default cut 1.25 was chosen based on the accuracy of fast_erf
    """

    functions = """
    double fast_erf(double x){
        double sgn = 1.0;

        if (x < 0){
            sgn = -1.0;
            x = -x;
        }

        double p = 0.47047;
        double a1 =  0.3480242;
        double a2 = -0.0958798;
        double a3 =  0.7478556;
        double t1 = 1.0/(1 + p*x);
        double t2 = t1*t1;
        double t3 = t1*t2;
        return sgn*(1 - (a1*t1 + a2*t2 + a3*t3)*exp(-x*x));
    }

    static union
    {
        double d;
        struct
        {
            #ifdef LITTLE_ENDIAN
                int j, i;
            #else
                int i, j;
            #endif
        } n;
    } _eco;

    #define EXP_A (1048576 /M_LN2)
    #define EXP_C 60801
    #define fast_exp(y) (_eco.n.i = EXP_A*(y) + (1072693248 - EXP_C), _eco.d)
    """

    code = """
    double coeff1 = 1.0/(alpha*sqrt(2.0));
    double coeff2 = sqrt(0.5/pi)*alpha;

    for (int i=0; i<N; i++){
        double dr = r[i]-r0;
        if (dr < cut && dr > -cut){
            double t = (r0 - r[i])*coeff1;
            ans[i] = 0.5*(1+fast_erf(t)) - coeff2/(r[i]+1e-10) * fast_exp(-t*t);
        } else {
            ans[i] = 0.0*(dr > cut) + 1.0*(dr < -cut);
        }
    }
    """

    shape = r.shape
    r = r.flatten()
    N = r.shape[0]
    ans = r*0
    pi = np.pi

    inline(code, arg_names=['r', 'r0', 'alpha', 'cut', 'ans', 'pi', 'N'],
            support_code=functions, verbose=0)
    return ans.reshape(shape)

def sphere_constrained_cubic(r, r0, alpha):
    """
    Sphere generated by a cubic interpolant constrained to be (1,0) on
    (r0-sqrt(3)/2, r0+sqrt(3)/2), the size of the cube in the (111) direction.
    """
    sqrt3 = np.sqrt(3)

    b_coeff = r0*0.5/sqrt3*(1 - 0.6*sqrt3*alpha)/(0.15 + r0*r0)
    rscl = np.clip(r-r0, -0.5*sqrt3, 0.5*sqrt3)

    a, d = rscl + 0.5*sqrt3, rscl - 0.5*sqrt3
    return alpha*d*a*rscl + b_coeff*d*a - d/sqrt3

try:
    sphere_analytical_gaussian_fast(np.linspace(0,10,10), 5.0)
except Exception as e:
    sphere_analytical_gaussian_fast = sphere_analytical_gaussian_trim

#=============================================================================
# Actual sphere collection (and slab)
#=============================================================================
class SphereCollectionRealSpace(object):
    def __init__(self, pos, rad, shape, support_size=4, typ=None, pad=None,
                 method='exact-gaussian-fast', alpha=None, method_function=None):
        """
        method can be one of:
            [
                'lerp', 'logistic', 'triangle', 'constrained-cubic',
                'exact-gaussian', 'exact-gaussian-trim', 'exact-gaussian-fast'
            ]
        """
        self.support_size = support_size
        self.pos = pos.astype('float')
        self.rad = rad.astype('float')
        self.N = rad.shape[0]

        # set the aliasing method and coefficient
        # FIXME -- check if function for method and set to that instead
        # method='user-defined' and pass method_function
        self._setup_sphere_functions(method=method, alpha=alpha)

        if typ is None:
            self.typ = np.ones(self.N)
            if pad is not None and pad <= self.N:
                self.typ[-pad:] = 0
        else:
            self.typ = typ.astype('float')

        self.shape = shape
        self._setup()

    def _setup_sphere_functions(self, method, alpha=None):
        self.methods = [
            'lerp', 'logistic', 'triangle', 'constrained-cubic',
            'exact-gaussian', 'exact-gaussian-trim', 'exact-gaussian-fast'
        ]

        self.sphere_functions = {
            'lerp': sphere_lerp,
            'logistic': sphere_logistic,
            'triangle': sphere_triangle_cdf,
            'exact-gaussian': sphere_analytical_gaussian,
            'exact-gaussian-trim': sphere_analytical_gaussian_trim,
            'exact-gaussian-fast': sphere_analytical_gaussian_fast,
            'constrained-cubic': sphere_constrained_cubic
        }

        self.alpha_defaults = {
            'lerp': 0.4539,
            'logistic': 6.5,
            'triangle': 0.6618,
            'exact-gaussian': 0.27595,
            'exact-gaussian-trim': 0.27595,
            'exact-gaussian-fast': 0.27595,
            'constrained-cubic': 0.84990,
        }

        self.method = method
        self.alpha = alpha if alpha is not None else self.alpha_defaults[self.method]

    def _setup(self):
        z,y,x = Tile(self.shape).coords()
        self.rvecs = np.rollaxis(np.array(np.broadcast_arrays(z,y,x)), 0, 4)
        self.particles = np.zeros(self.shape)
        self._diff_field = np.zeros(self.shape)

    def _particle(self, pos, rad, zscale, sign=1, dodiff=False):
        p = np.round(pos)
        r = np.round(np.array([1.0/zscale,1,1])*np.ceil(rad)+self.support_size)

        tile = Tile(p-r, p+r, 0, self.shape)
        subr = self.rvecs[tile.slicer + (np.s_[:],)]
        rvec = (subr - pos)

        # apply the zscale and find the distances to make a ellipsoid
        rvec[...,0] *= zscale
        rdist = np.sqrt((rvec**2).sum(axis=-1))

        # keep backwards compatibility for save files with older logistic coefficient
        # FIXME -- this is a nasty hack which should be fazed out when older save
        # files are deleted
        if not hasattr(self, 'method'):
            self._setup_sphere_functions('logistic', 5.0)
        if not hasattr(self, 'alpha'):
            self._setup_sphere_functions(self.method)

        # calculate the anti-aliasing according to the interpolation type
        t = sign*self.sphere_functions[self.method](rdist, rad, self.alpha)
        self.particles[tile.slicer] += t

        if dodiff:
            self._diff_field[tile.slicer] += t

    def _update_particle(self, n, p, r, t, zscale, dodiff=True):
        if self.typ[n] == 1:
            self._particle(self.pos[n], self.rad[n], zscale, -1, dodiff=dodiff)

        self.pos[n] = p
        self.rad[n] = r
        self.typ[n] = t

        if self.typ[n] == 1:
            self._particle(self.pos[n], self.rad[n], zscale, +1, dodiff=dodiff)

    def initialize(self, zscale):
        if len(self.pos.shape) != 2:
            raise AttributeError("Position array needs to be (-1,3) shaped, (z,y,x) order")

        self.particles = np.zeros(self.shape)
        for p0, r0, t0 in zip(self.pos, self.rad, self.typ):
            if t0 == 1:
                self._particle(p0, r0, zscale)

    def set_tile(self, tile):
        self.tile = tile

    def update(self, ns, pos, rad, typ, zscale, difference=True):
        for n, p, r, t in zip(ns, pos, rad, typ):
            self._update_particle(n, p, r, t, zscale, dodiff=difference)

    def get_field(self):
        return self.particles[self.tile.slicer]

    def get_diff_field(self):
        c = self._diff_field[self.tile.slicer].copy()
        self._diff_field[self.tile.slicer] *= 0
        return c

    def get_support_size(self, p0, r0, t0, p1, r1, t1, zscale):
        rsc = self.support_size

        zsc = np.array([1.0/zscale, 1, 1])
        r0, r1 = zsc*r0, zsc*r1

        off0 = r0 + rsc
        off1 = r1 + rsc

        if t0[0] == 1 and t1[0] == 1:
            pl = amin(p0-off0-1, p1-off1-1)
            pr = amax(p0+off0+1, p1+off1+1)
        if t0[0] != 1 and t1[0] == 1:
            pl = (p1-off1-1)
            pr = (p1+off1+1)
        if t0[0] == 1 and t1[0] != 1:
            pl = (p0-off0-1)
            pr = (p0+off0+1)
        if t0[0] != 1 and t1[0] != 1:
            c = np.array(self.shape)
            pl = c/2 - c/8
            pr = c/2 + c/8

        if len(pl.shape) > 1:
            pl = pl[0]
            pr = pr[0]
        return pl, pr

    def get_params(self):
        return np.hstack([self.pos.ravel(), self.rad])

    def get_params_pos(self):
        return self.pos.ravel()

    def get_params_rad(self):
        return self.rad

    def get_params_typ(self):
        return self.typ

    def __getstate__(self):
        odict = self.__dict__.copy()
        cdd(odict, ['rvecs', 'particles', '_diff_field'])
        return odict

    def __setstate__(self, idict):
        self.__dict__.update(idict)
        self._setup()


class Slab(object):
    def __init__(self, zpos, shape, normal=(1,0,0)):
        self.zpos = float(zpos)
        self.normal = np.array(normal).astype('float')
        self.normal /= np.sqrt(self.normal.dot(self.normal))

        self.shape = shape
        self._setup()

    def _setup(self):
        z,y,x = Tile(self.shape).coords()
        self.rvecs = np.rollaxis(np.array(np.broadcast_arrays(z,y,x)), 0, 4)
        self.image = np.zeros(self.shape)

    def _slab(self, zpos, norm, sign=1):
        # for the position at zpos, and the center in the x-y plane
        pos = np.array([zpos, self.shape[1]/2, self.shape[2]/2])

        p = (self.rvecs - pos).dot(norm)
        self.image = 1.0/(1.0 + np.exp(7*p))

    def initialize(self):
        self.image = np.zeros(self.shape)
        self._slab(self.zpos, self.normal)

    def set_tile(self, tile):
        self.tile = tile

    def update(self, params):
        zpos, norm = params[0], params[1:]
        norm = norm / np.sqrt(norm.dot(norm))

        self.zpos = zpos
        self.normal = norm

        self._slab(self.zpos, self.normal, +1)

    def get_field(self):
        return self.image[self.tile.slicer]

    def get_support_size(self, p=None):
        return np.zeros(3), np.array(self.shape)

    def get_params(self):
        return np.hstack([self.zpos, self.normal.ravel()])

    def __getstate__(self):
        odict = self.__dict__.copy()
        cdd(odict, ['rvecs', 'image'])
        return odict

    def __setstate__(self, idict):
        self.__dict__.update(idict)
        self._setup()

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return "{} <{}, {}>".format(str(self.__class__.__name__), self.zpos, list(self.normal))

{'fraction_non_alphanumeric': 0.16457261377687368, 'fraction_numerical': 0.033664393123529174, 'mean_word_length': 2.2325941206807633, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3781479', 'n_tokens_mistral': 9312, 'n_tokens_neox': 8700, 'n_words': 2025}