Procedural STL file generation in Python

generate_pinch_pot.py

import functools
import numpy
import struct
import svg.path   # pip install svg.path==2.2
from typing import List, Dict, Tuple, Callable  # kitchen sink?
array=numpy.array

# SVG <path> string describing the undistorted profile of the pot
PATH_DATA='''
m 0,50.916666 9.1801683,0 c 0.8686257,-1.2561 2.0631007,-1.983971 3.2272117,-2.612425 6.345522,-3.425674 6.113453,-5.441741 9.183746,-10.130426 3.920877,-5.987623 6.247456,-11.575049 11.615424,-15.913572 1.361948,-1.10076 2.477231,0.565914 1.236122,1.532279 -5.445873,4.240321 -7.508074,9.124891 -11.047477,15.345625 C 20.372917,44.45 19.05,46.566667 12.975328,49.970155 11.821211,50.616778 10.583333,51.59375 10.583333,52.916667 L 0,52.916667
'''
DOCUMENT_HEIGHT = 200  # Inkscape puts the origin in the bottom-left, but SVG puts it in the top-left. So I need to flip it vertically by the height of the original document.

# Parse the path string using a library
parsed_path = svg.path.parse_path(PATH_DATA)
'''
Becomes this (more or less):

Path(
    Line(start=50.916666j, end=(3.8885017+50.916666j)),
    CubicBezier(start=(3.8885017+50.916666j), control1=(4.7571276+49.660566j), control2=(5.951602599999999+48.932695j), end=(7.115713700000001+48.304241j)),
    CubicBezier(start=(7.115713700000001+48.304241j), control1=(13.461235+44.878567j), control2=(15.264759+44.373267j), end=(21.591126000000003+38.173815j)),
    CubicBezier(start=(21.591126000000003+38.173815j), control1=(26.703017000000003+33.164476j), control2=(27.838582000000002+26.598765999999998j), end=(33.20655000000001+22.260242999999996j)),
    CubicBezier(start=(33.20655000000001+22.260242999999996j), control1=(34.568498000000005+21.159482999999994j), control2=(35.68378100000001+22.826156999999995j), end=(34.44267200000001+23.792521999999995j)),
    CubicBezier(start=(34.44267200000001+23.792521999999995j), control1=(28.99679900000001+28.032842999999993j), control2=(28.507086000000008+34.12880799999999j), end=(23.395195000000008+39.138147j)),
    CubicBezier(start=(23.395195000000008+39.138147j), control1=(17.068828000000007+45.337599j), control2=(14.029183000000009+46.544481j), end=(7.683661400000009+49.970155j)),
    CubicBezier(start=(7.683661400000009+49.970155j), control1=(6.5195503000000095+50.598608999999996j), control2=(5.291666700000009+51.59375j), end=(5.291666700000009+52.916667j)),
    Line(start=(5.291666700000009+52.916667j), end=52.916667j),

    closed=False
)
'''

# Strip off the flat lines forming the bottom of the pot, which are the the "Line" elements of the Path above.
# We can add the bottom face later by hand.
parsed_path = parsed_path[1:-1]

# Extract the control points.
# I guess we could just use the CubicBezier objects later instead of indexing into here,
# but I didn't use an SVG path parser originally.
SPLINES: List[List[numpy.array]] = [
    [
        array([ pair.real, DOCUMENT_HEIGHT - pair.imag ])
        for pair in [
            bezier.start,
            bezier.control1,
            bezier.control2,
            bezier.end,
        ]
    ]
    for bezier in parsed_path
]


@functools.lru_cache()
def t_coefficients(t):
    t_sq=t*t
    t_cu=t_sq*t
    return array([1, t, t_sq, t_cu])


# Cache output of trig functions to potentially speed up calculations
sin = functools.lru_cache()(numpy.sin)
cos = functools.lru_cache()(numpy.cos)


def interp_cubic_spline(ctrl_points: List[numpy.array]) -> Callable[[float], numpy.array]:
    '''Interpolate a cubic spline along four control points'''
    a,b,c,d=ctrl_points

    coefficients = array([
           a,
        -3*a + 3*b,
         3*a - 6*b + 3*c,
        -1*a + 3*b - 3*c + d,
    ]).transpose()

    def interp(t):
        return coefficients.dot( t_coefficients(t) )

    return interp


def test_spline():
    '''Unused, ad-hoc test for interp_cubic_spline() to make sure it was behaving'''
    test_spline = [
        array((0,0)),
        array((1,0)),
        array((1,1)),
        array((0,1)),
    ]
    interp = interp_cubic_spline(test_spline)
    for t in numpy.linspace(0,1,11):
        print(f'{t}  {list(interp(t))}')


def print_passthru(*args):
    '''Log helper'''
    print(repr(args))
    return args


def export_tris(tris):
    '''Binary STL exporter'''

    # Header is 80 bytes, but usually ignored for some reason
    header = b'\0'*80 + struct.pack('<I', len(tris))

    LOOP_FORMAT = '<' + 'f'*12 + 'H'

    return b''.join(
        [header] + [
            struct.pack(LOOP_FORMAT,
                *numpy.cross(tri[1] - tri[0], tri[2] - tri[0]),
                *tri[0],
                *tri[1],
                *tri[2],
                0
            )
            for tri in tris
        ]
    )


@functools.lru_cache()
def fade_in(v):
    '''
    Piecewise "fade-in" function from 0 to 1, where the slope = 0 at both endpoints.
    Made of two mirrored quadratic functions
    '''
    if v < 0.5:
        return (v*2)**2 / 2
    else:
        return 1-((1-v)*2)**2/2


def tex_map(θ, r, z, v):
    # Distort the shape of the pot near the top.
    # This bit was as much art as science
    MAX_AMPLITUDE = 2
    BUMPS = 5

    # Displacements
    Δz = 3.0 * fade_in(v) * sin(θ * BUMPS - 0.5 * sin(θ * BUMPS))
    Δr = -2.0 * fade_in(v) * cos(θ * BUMPS + 1)
    Δθ = -0.1 * fade_in(v) * sin(θ * BUMPS)

    # Technically, the delta-z and delta-r are getting applied to both of them,
    # but at right-angles -- the displacement is rotated 45 degrees to reflect
    # the approximate angle of the pot.
    return θ + Δθ, r + Δz + Δr, z + Δz - Δr

def tex_map_inv(θ, r, z, v):
    # Invert the surface position parameter ("v") for the inside of the pot
    return tex_map(θ, r, z, 1-v)

def tex_map_brim(θ, r, z, v):
    # Maximum, constant distortion on the rim
    return tex_map(θ, r, z, 1)

def tex_map_passthru(θ, r, z, v):
    # No distortion
    return θ, r, z

# Yes, I'm using non-latin characters in code. It's just throwaway work; I get to have fun once in a while.

# Apply our distortion functions to our splines
mesh_tex_maps: List[Callable[[float, float, float], Tuple[float, float, float]]] = [
    tex_map_passthru,   # Outer base/bottom rim
    tex_map_passthru,   # Outer lower
    
    tex_map,            # Outer upper
    tex_map_brim,       # Rim
    tex_map_inv,        # Inner upper

    tex_map_passthru,   # Inner lower/mid
    tex_map_passthru,   # Inner lowest
]

# How many points to sample along the length of each spline
radial_resolutions: List[int] = [
    8,
    8,
    16,
    8,
    16,
    8,
    8,
]

# How many samples we should take rotating the splines around the z-axis
NUMBER_ANGLES = 120

# Pre-calculate input arrays since they don't really change
thetas = numpy.linspace(0, 2*numpy.pi, NUMBER_ANGLES + 1, endpoint=True)

# Extrude each spline into a mesh (2D array of points) by rotating it around the z-axis.
# Then, deform each mesh by a previously-defined function.
# Do this for each mesh.
meshes: List[List[List[numpy.array]]] = [
    [
        [
            array((_r * cos(_θ), _r * sin(_θ), _z))
            for θ in thetas
            for _θ, _r, _z in [mesh_tex_map(θ, r, z, v)]
        ]
        for v in numpy.linspace(0, 1, radial_resolution + 1, endpoint=True)
        for r, z in [spline_interp(v)]
    ]
    for spline, mesh_tex_map, radial_resolution in zip(SPLINES, mesh_tex_maps, radial_resolutions)
    for spline_interp in [interp_cubic_spline(spline)]
]

# Turn the meshes into triangles
tris: List[Tuple[Tuple[float, float, float], Tuple[float, float, float], Tuple[float, float, float]]] = [
    tri

    for mesh_i, mesh in enumerate(meshes)
    for _log_mesh in [print(f'Mesh {mesh_i} -- len {len(mesh)}')]
    for u_i in range(NUMBER_ANGLES)
    for v_i in range(len(mesh)-1)

    for quad_points in ((  # fancy variable declaration
        mesh[v_i][u_i],
        mesh[v_i][u_i+1],
        mesh[v_i+1][u_i+1],
        mesh[v_i+1][u_i],
    ),)
    for tri in (
        (quad_points[0], quad_points[1], quad_points[2]),
        (quad_points[0], quad_points[2], quad_points[3]),
    )
]

print(f'First tri:')
print(repr(tris[0]))

# Export!
tri_data = export_tris(tris)
with open('pinch_pot.stl', 'wb') as outfile:
    outfile.write(tri_data)