3D Dual Contouring

DualContouring3D.py

bl_info = {
    "name": "3D Implicit Function Contouring",
    "description": "Contour 3D implicit function",
    "author": "TheBusyTypist",
    "location": "View3D > Add > Mesh",
    "category": "Add Mesh"
}

import bpy
import bmesh
import numpy as np
import math

import traceback

class Config:
    def __init__(self, begin=-4.0, end=4.0, ticks=20):
        self.begin = begin
        self.end = end
        self.ticks = ticks

def Sample(f, config):
    begin = config.begin
    end = config.end
    ticks = config.ticks

    x = np.linspace(begin, end, ticks)
    y = np.linspace(begin, end, ticks)
    z = np.linspace(begin, end, ticks)
    X, Y, Z = np.meshgrid(x, y, z)

    vF = np.vectorize(f)
    s = vF(X, Y, Z)
    return s

def SolveIntersection(f, p0, p1, config):
    EPS = 1e-12
    m = (p0 + p1) * 0.5
    v = f(m[0], m[1], m[2])

    while abs(v) >= EPS:
        if v > 0:
            p1 = m
        else:
            p0 = m
        m = (p0 + p1) * 0.5
        v = f(m[0], m[1], m[2])
    return m

def HaveOppositeSigns(a, b):
    return (
        a < 0 and b >= 0
        or
        a >= 0 and b < 0
    )

def ComputeHermiteData(f, samples, config):
    def D(f, p):
        h = 1e-8
        fx1 = f(p[0] + h, p[1], p[2])
        fx0 = f(p[0] - h, p[1], p[2])
        fy1 = f(p[0], p[1] + h, p[2])
        fy0 = f(p[0], p[1] - h, p[2])
        fz0 = f(p[0], p[1], p[2] - h)
        fz1 = f(p[0], p[1], p[2] + h)
        n = np.array([
            (fx1 - fx0) * 0.5 / h,
            (fy1 - fy0) * 0.5 / h,
            (fz1 - fz0) * 0.5 / h])
        return n

    n = config.ticks
    g = [None for i in range(n) for j in range(n) for k in range(n)]
    begin = config.begin
    end = config.end
    step = (end - begin) / (n - 1)

    for z in range(n):
        for y in range(n):
            for x in range(n):
                o = np.array([
                    begin + x * step,
                    begin + y * step,
                    begin + z * step])

                u = np.array([
                    begin + (x + 1) * step,
                    begin + y * step,
                    begin + z * step])

                v = np.array([
                    begin + x * step,
                    begin + (y + 1) * step,
                    begin + z * step])

                w = np.array([
                    begin + x * step,
                    begin + y * step,
                    begin + (z + 1) * step])

                fo = samples[y, x, z]

                if x != n - 1:
                    fu = samples[y, x + 1, z]
                else:
                    fu = f(u[0], u[1], u[2])

                if y != n - 1:
                    fv = samples[y + 1, x, z]
                else:
                    fv = f(v[0], v[1], v[2])

                if z != n - 1:
                    fw = samples[y, x, z + 1]
                else:
                    fw = f(w[0], w[1], w[2])

                a = None
                if HaveOppositeSigns(fo, fu):
                    if fo < fu:
                        a = SolveIntersection(f, o, u, config)
                    else:
                        a = SolveIntersection(f, u, o, config)

                b = None
                if HaveOppositeSigns(fo, fv):
                    if fo < fv:
                        b = SolveIntersection(f, o, v, config)
                    else:
                        b = SolveIntersection(f, v, o, config)

                c = None
                if HaveOppositeSigns(fo, fw):
                    if fo < fw:
                        c = SolveIntersection(f, o, w, config)
                    else:
                        c = SolveIntersection(f, w, o, config)

                if a is not None:
                    na = D(f, a)

                if b is not None:
                    nb = D(f, b)

                if c is not None:
                    nc = D(f, c)

                h = (
                    None if a is None else (a, na),
                    None if b is None else (b, nb),
                    None if c is None else (c, nc)
                )

                g[z * n * n + y * n + x] = h

    return g

def ConstructQEF(a):
    A = np.array([])
    b = []
    px = []
    py = []
    pz = []
    g = None

    if len(a) >= 2:
        for p, n in a:
            px.append(p[0])
            py.append(p[1])
            pz.append(p[2])
            b.append(n.dot(p))
            if A.shape[0] == 0:
                A = n
            else:
                A = np.vstack((A, n))

        g = np.array([sum(px) / len(px), sum(py) / len(py), sum(pz) / len(pz)])

    return A, np.array(b), g
    
def SolveQEF(A, b, g):
    ATA = A.T.dot(A)
    ATb = A.T.dot(b)
    d = ATb - ATA.dot(g)
    c = np.linalg.pinv(ATA).dot(d)
    return c + g
    # return np.linalg.pinv(ATA).dot(ATb)

def IsInside(p, lx, ux, ly, uy, lz, uz):
    return (
        p[0] <= ux and p[0] >= lx
        and p[1] <= uy and p[1] >= ly
        and p[2] <= uz and p[2] >= lz)

def Contour(H, config):
    n = config.ticks - 1
    m = config.ticks
    v = [None for i in range(n) for j in range(n) for k in range(n)]

    def AppendIfNotNone(a, h):
        if h is not None:
            a.append(h)

    for z in range(n):
        for y in range(n):
            for x in range(n):
                a = []
                
                h = H[z * m * m + y * m + x]
                hx = H[z * m * m + y * m + x + 1]
                hxz = H[(z + 1) * m * m + y * m + x + 1]
                hz = H[(z + 1) * m * m + y * m + x]
                hxy = H[z * m * m + (y + 1) * m + x]
                hy = H[z * m * m + (y + 1) * m + x]
                hyz = H[(z + 1) * m * m + (y + 1) * m + x]

                AppendIfNotNone(a, h[0])
                AppendIfNotNone(a, h[1])
                AppendIfNotNone(a, h[2])

                AppendIfNotNone(a, hx[1])
                AppendIfNotNone(a, hx[2])

                AppendIfNotNone(a, hxz[1])

                AppendIfNotNone(a, hz[0])
                AppendIfNotNone(a, hz[1])

                AppendIfNotNone(a, hy[0])
                AppendIfNotNone(a, hy[2])

                AppendIfNotNone(a, hxy[2])

                AppendIfNotNone(a, hyz[0])

                A, b, g = ConstructQEF(a)

                if len(a) >= 2:
                    p = SolveQEF(A, b, g)

                    step = (config.end - config.begin) / n
                    x0 = config.begin + step * x
                    y0 = config.begin + step * y
                    z0 = config.begin + step * z
                    if not IsInside(p,
                        x0, x0 + step, y0, y0 + step, z0, z0 + step):
                        p = g
                        
                    v[z * n * n + y * n + x] = (p[0], p[1], p[2])

    return v

def GenerateMesh(v, config, context):
    n = config.ticks - 1

    indices = [None for i in range(n) for j in range(n) for k in range(n)]
    vertices = []
    for i, e in enumerate(v):
        if e is not None:
            indices[i] = len(vertices)
            vertices.append((e[0], e[1], e[2]))

    edges = []
    for z in range(n):
        for y in range(n):
            for x in range(n):
                ci = z * n * n + y * n + x
                c = indices[ci]

                r = None
                if x != n - 1:
                    ri = z * n * n + y * n + x + 1
                    r = indices[ri]

                u = None
                if y != n - 1:
                    ui = z * n * n + (y + 1) * n + x
                    u = indices[ui]

                f = None
                if z != n - 1:
                    fi = (z + 1) * n * n + y * n + x
                    f = indices[fi]

                if c is not None:
                    if r is not None:
                        edges.append((c, r))
                    if u is not None:
                        edges.append((c, u))
                    if f is not None:
                        edges.append((c, f))

    mesh = bpy.data.meshes.new("Mesh")
    mesh.from_pydata(vertices, edges, [])
    obj = bpy.data.objects.new("MeshObj", mesh)
    context.scene.objects.link(obj)

class DualContouring3D(bpy.types.Operator):
    """Contour 3D implicit function"""
    bl_idname = "object.contour_3d"
    bl_label = "Contour 3D implicit function"
    bl_options = {"REGISTER", "UNDO"}

    Ticks = bpy.props.IntProperty(name="Ticks", default=64)
    FunctionSource = bpy.props.StringProperty(name="Function",
        # default="x ** 4 + y ** 4 + z ** 4 - 1"
        default="(2 * x**2 + y**2 + z**2 - 1)**3 - x**2 * z**3 / 10 - y**2 * z**3"
    )

    def execute(self, context):
        f = eval("lambda x, y, z: " + self.FunctionSource)
        config = Config()
        config.ticks = self.Ticks
        samples = Sample(f, config)

        H = ComputeHermiteData(f, samples, config)

        v = Contour(H, config)

        GenerateMesh(v, config, context)
        
        return {"FINISHED"}

def register():
    bpy.utils.register_class(DualContouring3D)

def unregister():
    bpy.utils.unregister_class(DualContouring3D)

if __name__ == "__main__":
    register()