Extract from CNC slicer project for SDF for a mesh

bih3.go

package cnc_slicer

import (
	"gonum.org/v1/gonum/spatial/r3"
)

import (
	"fmt"
	"math"
	"sort"
	"unsafe" // used to store floats and ints in the same value
)

func max(a, b float64) float64 {
	if a > b {
		return a
	} else {
		return b
	}
}

func min(a, b float64) float64 {
	if a < b {
		return a
	} else {
		return b
	}
}

func dist_sq(r r3.Vec) float64 {
	return r.X*r.X + r.Y*r.Y + r.Z*r.Z
}

func dist(r r3.Vec) float64 {
	return math.Sqrt(r.X*r.X + r.Y*r.Y + r.Z*r.Z)
}

// pick the smaller magnitude number, effectively the distance from
// the closest triangle
func minAbs(a float64, b float64) float64 {
	if math.Abs(a) < math.Abs(b) {
		return a
	} else {
		return b
	}
}

func minVec(a r3.Vec, b r3.Vec) r3.Vec {
	return r3.Vec{min(a.X, b.X), min(a.Y, b.Y), min(a.Z, b.Z)}
}

func maxVec(a r3.Vec, b r3.Vec) r3.Vec {
	return r3.Vec{max(a.X, b.X), max(a.Y, b.Y), max(a.Z, b.Z)}
}

type Mesh struct {
	Edge_adj map[[2]int]int // edge to triangles mapping for adjacency
	Indices  [][3]int       // indices into the vertex list
	Vertices []r3.Vec
	Bb       r3.Box
}

func (m *Mesh) Tri(i int) r3.Triangle {
	idxs := m.Indices[i]
	return r3.Triangle{
		m.Vertices[idxs[0]],
		m.Vertices[idxs[1]],
		m.Vertices[idxs[2]],
	}
}

func NewMesh(model []r3.Triangle, vertTol float64) Mesh {
	vertices := make([]r3.Vec, 0)
	mesh := make([][3]int, len(model))
	bb := r3.Box{
		Min: r3.Vec{math.MaxFloat64, math.MaxFloat64, math.MaxFloat64},
		Max: r3.Vec{-math.MaxFloat64, -math.MaxFloat64, -math.MaxFloat64},
	}

	vert_cache := make(map[[3]int64]int)
	// copy the model into the vertex and mesh arrays
	hvertTol := 0.5 * vertTol
	for i, tri := range model {
		for j, v := range tri {
			// look for a vertex within tolerance
			index := [3]int64{
				int64((v.X + hvertTol) / vertTol),
				int64((v.Y + hvertTol) / vertTol),
				int64((v.Z + hvertTol) / vertTol),
			}
			if vidx, ok := vert_cache[index]; !ok {
				vertices = append(vertices, v)
				bb.Min = minVec(bb.Min, v)
				bb.Max = maxVec(bb.Max, v)
				mesh[i][j] = len(vertices) - 1
				vert_cache[index] = len(vertices) - 1
			} else {
				mesh[i][j] = vidx
			}
		}
	}

	// generate edge adjacency list and calculate pseudonormals
	edges := make(map[[2]int]int)
	for i, tri := range mesh {
		id := [2]int{tri[0], tri[1]}
		edges[id] = i
		id[0] = tri[1]
		id[1] = tri[2]
		edges[id] = i
		id[0] = tri[2]
		id[1] = tri[0]
		edges[id] = i
	}

	return Mesh{
		Edge_adj: edges,
		Indices:  mesh,
		Vertices: vertices,
		Bb:       bb,
	}
}

const (
	LEAF = iota
	X_CLIP
	Y_CLIP
	Z_CLIP
)

type bihInternal struct {
	flags       int // offset to children is stored in the upper 30 bits, lower two bits are used as flags
	left, right int64
	// either:
	// - left and right clipping plane values (reinterpret them as float64)
	// - or offset intwo the mesh list and the number of elements that belong
	//   to this leaf
}

func (b *bihInternal) is_leaf() bool {
	return (b.flags & 3) == LEAF
}

func (b *bihInternal) is_x() bool {
	return (b.flags & 3) == X_CLIP
}

func (b *bihInternal) is_y() bool {
	return (b.flags & 3) == Y_CLIP
}

func (b *bihInternal) is_z() bool {
	return (b.flags & 3) == Z_CLIP
}

func (b *bihInternal) leftClip() float64 {
	return *(*float64)(unsafe.Pointer(&b.left))
}
func (b *bihInternal) rightClip() float64 {
	return *(*float64)(unsafe.Pointer(&b.right))
}

// TODO figure out what to make public
type BIH struct {
	Mesh
	bih          []bihInternal
	face_normals []r3.Vec // len(mesh), stores all the face normals
	edge_normals []r3.Vec // 3*len(mesh), stores all the edge pseudonormals
	vert_normals []r3.Vec // vertex pseudonormals
}

func (b *BIH) findMinDistHelper(dist_box func(r3.Box) float64, dist_tri func(r3.Triangle) float64, idx int, bb r3.Box, dist float64, min_tri int) (float64, int) {
	bih := &b.bih[idx]
	if bih.is_leaf() {
		offset := bih.left
		end := bih.right

		for i := offset; i < end; i++ {
			cur_dist_tri := dist_tri(b.Tri(int(i)))
			//fmt.Printf("%v %d %f\n", target, i, dist_tri)
			if cur_dist_tri < dist {
				dist = cur_dist_tri
				min_tri = int(i)
			}
		}
	} else {
		// see which bounding box is closer to the target and
		// start with that one
		left_bb := bb
		right_bb := bb
		if bih.is_x() {
			left_bb.Max.X = bih.leftClip()
			right_bb.Min.X = bih.rightClip()
		} else if bih.is_y() {
			left_bb.Max.Y = bih.leftClip()
			right_bb.Min.Y = bih.rightClip()
		} else if bih.is_z() {
			left_bb.Max.Z = bih.leftClip()
			right_bb.Min.Z = bih.rightClip()
		}

		left_dist := dist_box(left_bb)
		right_dist := dist_box(right_bb)

		left_idx := bih.flags >> 2
		right_idx := left_idx + 1
		if left_dist < right_dist {
			if left_dist < dist {
				dist, min_tri = b.findMinDistHelper(dist_box, dist_tri, left_idx, left_bb, dist, min_tri)
			}
			if right_dist < dist {
				dist, min_tri = b.findMinDistHelper(dist_box, dist_tri, right_idx, right_bb, dist, min_tri)
			}
		} else {
			if left_dist < dist {
				dist, min_tri = b.findMinDistHelper(dist_box, dist_tri, left_idx, left_bb, dist, min_tri)
			}
			if right_dist < dist {
				dist, min_tri = b.findMinDistHelper(dist_box, dist_tri, right_idx, right_bb, dist, min_tri)
			}
		}
	}

	return dist, min_tri
}

func (b *BIH) FindMinDist(dist_box func(r3.Box) float64, dist_tri func(r3.Triangle) float64) (float64, int) {
	return b.findMinDistHelper(dist_box, dist_tri, 0, b.Bb, math.MaxFloat64, -1)
}

const (
	FACE   = 0
	EDGE   = 1 << 30
	VERTEX = 2 << 30
)

func subdivide(b []bihInternal, bih_idx int, mesh_idx int, mesh [][3]int, verts []r3.Vec, centroids map[[3]int][3]float64, bb r3.Box) []bihInternal {
	//fmt.Printf("%d %d %d %v \n", bih_idx, mesh_idx, len(mesh), bb)
	if len(mesh) <= 4 {
		b[bih_idx] = bihInternal{
			flags: LEAF,
			left:  int64(mesh_idx),
			right: int64(len(mesh) + mesh_idx),
		}
		return b
	} else {

		// TODO maybe swap out with a more modern partition method
		// classical heuristic, the longest axis
		// using the median as the pivot point
		dims := r3.Sub(bb.Max, bb.Min)
		clipping_plane := X_CLIP

		if dims.X >= dims.Y && dims.X >= dims.Z {
			clipping_plane = X_CLIP
		} else if dims.Y >= dims.X && dims.Y >= dims.Z {
			clipping_plane = Y_CLIP
		} else {
			clipping_plane = Z_CLIP
		}

		sort.Slice(mesh, func(i, j int) bool {
			return centroids[mesh[i]][clipping_plane-1] < centroids[mesh[j]][clipping_plane-1]
		})

		left_half := len(mesh) / 2
		left_bb := r3.Box{
			Min: r3.Vec{math.MaxFloat64, math.MaxFloat64, math.MaxFloat64},
			Max: r3.Vec{-math.MaxFloat64, -math.MaxFloat64, -math.MaxFloat64},
		}
		right_bb := r3.Box{
			Min: r3.Vec{math.MaxFloat64, math.MaxFloat64, math.MaxFloat64},
			Max: r3.Vec{-math.MaxFloat64, -math.MaxFloat64, -math.MaxFloat64},
		}

		for _, tri := range mesh[0:left_half] {
			for _, idx := range tri {
				left_bb.Min = minVec(left_bb.Min, verts[idx])
				left_bb.Max = maxVec(left_bb.Max, verts[idx])
			}
		}
		for _, tri := range mesh[left_half:len(mesh)] {
			for _, idx := range tri {
				right_bb.Min = minVec(right_bb.Min, verts[idx])
				right_bb.Max = maxVec(right_bb.Max, verts[idx])
			}
		}
		//fmt.Printf("sub %v %v\n", left_bb, right_bb)

		// append two new nodes to store the children
		children_idx := len(b)
		b = append(b, bihInternal{})
		b = append(b, bihInternal{})
		b = subdivide(b, children_idx, mesh_idx, mesh[0:left_half], verts, centroids, left_bb)
		b = subdivide(b, children_idx+1, mesh_idx+left_half, mesh[left_half:len(mesh)], verts, centroids, right_bb)

		b[bih_idx].flags = (children_idx << 2) | clipping_plane
		var left_plane float64
		var right_plane float64
		switch clipping_plane {
		case X_CLIP:
			left_plane = left_bb.Max.X
			right_plane = right_bb.Min.X
		case Y_CLIP:
			left_plane = left_bb.Max.Y
			right_plane = right_bb.Min.Y
		case Z_CLIP:
			left_plane = left_bb.Max.Z
			right_plane = right_bb.Min.Z
		}
		*(*float64)(unsafe.Pointer(&b[bih_idx].left)) = left_plane
		*(*float64)(unsafe.Pointer(&b[bih_idx].right)) = right_plane
		return b
	}
}

func calc_alpha_wnormal(tri [3]int, vert_idx int, vertices []r3.Vec) (float64, r3.Vec) {
	idx := -1
	for i, tri_idx := range tri {
		if tri_idx == vert_idx {
			idx = i
			break
		}
	}

	nidx := []int{1, 2, 0}[idx]
	bidx := 3 - idx - nidx

	a := vertices[vert_idx]
	b := vertices[tri[nidx]]
	c := vertices[tri[bidx]]

	ba := r3.Sub(b, a)
	ca := r3.Sub(c, a)
	norm := r3.Cross(ba, ca)

	cosalpha := r3.Dot(ba, ca) / (dist(ba) * dist(ca))
	alpha := math.Acos(cosalpha)
	return alpha, r3.Scale(alpha, r3.Unit(norm))
}

func ImportModel(model []r3.Triangle, vertTol float64) BIH {
	vertices := make([]r3.Vec, 0)
	mesh := make([][3]int, len(model))
	bb := r3.Box{
		Min: r3.Vec{math.MaxFloat64, math.MaxFloat64, math.MaxFloat64},
		Max: r3.Vec{-math.MaxFloat64, -math.MaxFloat64, -math.MaxFloat64},
	}

	vert_cache := make(map[[3]int64]int)
	// copy the model into the vertex and mesh arrays
	hvertTol := 0.5 * vertTol
	for i, tri := range model {
		for j, v := range tri {
			// look for a vertex within tolerance
			index := [3]int64{
				int64((v.X + hvertTol) / vertTol),
				int64((v.Y + hvertTol) / vertTol),
				int64((v.Z + hvertTol) / vertTol),
			}
			if vidx, ok := vert_cache[index]; !ok {
				vertices = append(vertices, v)
				bb.Min = minVec(bb.Min, v)
				bb.Max = maxVec(bb.Max, v)
				mesh[i][j] = len(vertices) - 1
				vert_cache[index] = len(vertices) - 1
			} else {
				mesh[i][j] = vidx
			}
		}
	}
	centroids := make(map[[3]int][3]float64)
	for _, tri := range mesh {
		v := [3]float64{0, 0, 0}
		for _, idx := range tri {
			v[0] += vertices[idx].X
			v[1] += vertices[idx].Y
			v[2] += vertices[idx].Z
		}
		v[0] /= 3.
		v[1] /= 3.
		v[2] /= 3.
		centroids[tri] = v
	}

	b := make([]bihInternal, 1)
	b = subdivide(b, 0, 0, mesh, vertices, centroids, bb)

	// generate edge adjacency list and calculate pseudonormals
	edges := make(map[[2]int]int)
	face_normals := make([]r3.Vec, len(mesh))
	for i, tri := range mesh {
		id := [2]int{tri[0], tri[1]}
		edges[id] = i
		id[0] = tri[1]
		id[1] = tri[2]
		edges[id] = i
		id[0] = tri[2]
		id[1] = tri[0]
		edges[id] = i

		ba := r3.Sub(vertices[tri[1]], vertices[tri[0]])
		ca := r3.Sub(vertices[tri[2]], vertices[tri[0]])

		face_normals[i] = r3.Unit(r3.Cross(ba, ca))
	}

	next_j := []int{1, 2, 0}
	edge_pseudonormals := make([]r3.Vec, 3*len(mesh))
	first_nonclosed := true
	for i, tri := range mesh {
		cur_normal := face_normals[i]
		for j := range tri {
			// flip order to find the other edge
			other := [2]int{tri[next_j[j]], tri[j]}
			if other_face, ok := edges[other]; !ok {
				if first_nonclosed {
					fmt.Println("w: non closed edge detected")
					first_nonclosed = false
				}
				// this edge is only adjacent to this triangle,
				// store this triangle's face normal as its normal
				edge_pseudonormals[3*i+j] = cur_normal
			} else {
				other_normal := face_normals[other_face]
				edge_normal := r3.Add(r3.Scale(0.5, cur_normal), r3.Scale(0.5, other_normal))
				edge_pseudonormals[3*i+j] = edge_normal
			}
		}
	}

	vertex_pseudonormals := make([]r3.Vec, len(vertices))
	for i, tri := range mesh {
		// for each vertex check to see if it hasn't already been calculated
		// if it hasn't then calculate it
		for j, idx := range tri {
			if dist_sq(vertex_pseudonormals[idx]) < 0.5 {
				// calculate it by traversing along all the edges sharing that vertex
				alpha, wnormal := calc_alpha_wnormal(tri, idx, vertices)
				normal_tot := wnormal
				weights := alpha

				edge_vert := tri[next_j[j]]
				next_tri, ok := edges[[2]int{edge_vert, idx}]
				//fmt.Printf("next tri %v %v %v\n", mesh[next_tri], edge_vert, idx)
				for ok && next_tri != i {
					alpha, wnormal = calc_alpha_wnormal(mesh[next_tri], idx, vertices)
					normal_tot = r3.Add(normal_tot, wnormal)
					weights += alpha

					vert_idx := -1
					edge_idx := -1
					for k, jdx := range mesh[next_tri] {
						if jdx == idx {
							vert_idx = k
						}
						if jdx == edge_vert {
							edge_idx = k
						}
					}

					if vert_idx == -1 || edge_idx == -1 {
						fmt.Printf("triangle doesn't match vertex or edge %v %v %v\n", mesh[next_tri], edge_vert, idx)
						panic("unreachable state; reached a triangle not containing the proper adjacent edge")
					}
					// the indexes all sum up to 3, so subtracting the other two indices
					// returns the missing one
					edge_vert = mesh[next_tri][3-vert_idx-edge_idx]
					// and each matching triangle can be found using
					// its edge some edge-vertex
					next_tri, ok = edges[[2]int{edge_vert, idx}]
					//fmt.Printf("next tri2 %v %v %v\n", mesh[next_tri], edge_vert, idx)
				}
				if !ok {
					//fmt.Printf("w: incomplete vertex traversal, %v\n", vertices[idx])
				}
				vertex_pseudonormals[idx] = r3.Scale(1/alpha, normal_tot)
			}
		}
	}

	return BIH{
		Mesh: Mesh{
			Edge_adj: edges,
			Indices:  mesh,
			Vertices: vertices,
			Bb:       bb,
		},
		bih:          b,
		face_normals: face_normals,
		edge_normals: edge_pseudonormals,
		vert_normals: vertex_pseudonormals,
	}
}

func (b *BIH) Evaluate(p r3.Vec) float64 {
	dist := b.DistNearestTri(p)
	return dist
}

func (b *BIH) Bounds() r3.Box {
	return b.Bb
}