Addressable RGB LED 2020 - Worldsemi WS2812B-2020

#!/usr/bin/env python3
"""Rebuild the WS2812B-2020 page's composite LED scene GLBs (WS2812B-2020.glb, WS2812B-2020.insertion.glb) for the
corrected 2020 package and land pattern. Same scene contract as before (node names/order, the 'insert'
animation on the LED and LED_lit nodes, the LED_light point light, the board mesh), so the page's viewer
(WS2812B-2020-viewer.html / the README timeline iframe, which look nodes up by name) keeps working:

- fr4_board: unchanged.
- pad_top: the four 0.70 x 0.70 mm pads of the new WS2812B-2020.kicad_mod (datasheet p.2 PCB Solder Pad),
  copper 0.04 mm; KiCad Y-down flipped to Y-up (same frame as the STEP and the hero overlay).
- silk: 0.12 mm ribbons from the new footprint's F.SilkS lines (pin-1 corner bracket).
- solder_top: the 300 um jet-paste solder balls, 2 x 2 per pad (scene_solderballs.py, Adom solder-ball standard).
- WS2812B-2020 (LED node): geometry from `step2glb convert` of the coloured, laser-etched STEP, merged per
  material (body, window, terminal, polarity mark, R/G/B dies, IC die); real height (0.84 mm), no Z scale.
  The laser etch is a child node Adom.LaserEtch.MPN (rides the insertion animation).
- LED_lit glow: a shell just outside the new clear moulding (1.45 x 2.00 mm, z 0.28-0.84 mm).
Usage: rebuild_ws_scene.py OLD.glb S2G.glb KICAD_MOD OUT.glb"""
import sys, re, json, hashlib, math
import numpy as np
sys.path.insert(0, '/home/adom/project/page-work/led-pass/tools')
import gltfio

old, s2g, kmod, out = sys.argv[1:5]
MM = 1e-3
def box(x0, y0, z0, x1, y1, z1):
    """axis-aligned box, flat normals: returns P, N, I (float64 metres)."""
    P, N, I = [], [], []
    faces = [((1, 0, 0), [(x1, y0, z0), (x1, y1, z0), (x1, y1, z1), (x1, y0, z1)]),
             ((-1, 0, 0), [(x0, y1, z0), (x0, y0, z0), (x0, y0, z1), (x0, y1, z1)]),
             ((0, 1, 0), [(x1, y1, z0), (x0, y1, z0), (x0, y1, z1), (x1, y1, z1)]),
             ((0, -1, 0), [(x0, y0, z0), (x1, y0, z0), (x1, y0, z1), (x0, y0, z1)]),
             ((0, 0, 1), [(x0, y0, z1), (x1, y0, z1), (x1, y1, z1), (x0, y1, z1)]),
             ((0, 0, -1), [(x0, y1, z0), (x1, y1, z0), (x1, y0, z0), (x0, y0, z0)])]
    for n, q in faces:
        b = len(P); P += q; N += [n] * 4; I += [b, b + 1, b + 2, b, b + 2, b + 3]
    return np.array(P), np.array(N, float), np.array(I)
def prism(profile_xz, y0, y1):
    """extrude a convex XZ polygon (CCW seen from -Y) along Y from y0 to y1."""
    P, N, I = [], [], []
    n = len(profile_xz)
    for yy, ny in ((y0, -1), (y1, 1)):          # caps
        b = len(P); P += [(x, yy, z) for x, z in profile_xz]; N += [(0, ny, 0)] * n
        for k in range(1, n - 1): I += ([b, b + k, b + k + 1] if ny < 0 else [b, b + k + 1, b + k])
    for k in range(n):                            # sides
        (xa, za), (xb, zb) = profile_xz[k], profile_xz[(k + 1) % n]
        nx, nz = (zb - za), -(xb - xa); L = math.hypot(nx, nz) or 1; nx, nz = nx / L, nz / L
        b = len(P); P += [(xa, y0, za), (xb, y0, zb), (xb, y1, zb), (xa, y1, za)]; N += [(nx, 0, nz)] * 4
        I += [b, b + 1, b + 2, b, b + 2, b + 3]
    return np.array(P), np.array(N, float), np.array(I)
def merge(parts):
    P, N, I, o = [], [], [], 0
    for p, n, i in parts: P.append(p); N.append(n); I.append(i + o); o += len(p)
    return np.concatenate(P), np.concatenate(N), np.concatenate(I)
def fix_winding(P, N, I):
    """make every triangle's winding agree with its vertex normals (outward), whatever order the builders used."""
    T = I.reshape(-1, 3).copy()
    a, b, c = P[T[:, 0]], P[T[:, 1]], P[T[:, 2]]
    fn = np.cross(b - a, c - a); flip = (fn * N[T[:, 0]]).sum(1) < 0
    T[flip] = T[flip][:, [0, 2, 1]]; return P, N, T.reshape(-1)

# ---- footprint (KiCad mm, Y down -> metres, Y up) ----
src = open(kmod).read()
pads = [(m.group(1), float(m.group(2)), -float(m.group(3)), float(m.group(4)), float(m.group(5)))
        for m in re.finditer(r'\(pad "(\d+)" smd \w+ \(at ([-\d.]+) ([-\d.]+)\) \(size ([\d.]+) ([\d.]+)\)', src)]
assert len(pads) == 4, pads
silk = [(float(a), -float(b), float(c), -float(d)) for a, b, c, d, lay in
        re.findall(r'\(fp_line \(start ([-\d.]+) ([-\d.]+)\) \(end ([-\d.]+) ([-\d.]+)\)\s*\(stroke \(width [\d.]+\) \(type \w+\)\) \(layer "([^"]+)"\)', src) if lay == 'F.SilkS']
assert silk
CU = 0.04
pad_geo = merge([box((x - w / 2) * MM, (y - h / 2) * MM, 0, (x + w / 2) * MM, (y + h / 2) * MM, CU * MM) for _, x, y, w, h in pads])
SW = 0.12
silk_parts = []
for x0, y0, x1, y1 in silk:
    if abs(y0 - y1) < 1e-9: silk_parts.append(box((min(x0, x1) - SW / 2) * MM, (y0 - SW / 2) * MM, 0.03 * MM, (max(x0, x1) + SW / 2) * MM, (y0 + SW / 2) * MM, 0.06 * MM))
    else: silk_parts.append(box((x0 - SW / 2) * MM, (min(y0, y1) - SW / 2) * MM, 0.03 * MM, (x0 + SW / 2) * MM, (max(y0, y1) + SW / 2) * MM, 0.06 * MM))
silk_geo = merge(silk_parts)
# solder: the 300 um jet-paste balls (Adom solder-ball standard), the same ones the hero shows (scene_solderballs.py)
import os as _os
sys.path.insert(0, _os.path.dirname(_os.path.abspath(__file__)))
import scene_solderballs as _sb
_P, _N, _I = _sb.hemispheres(_sb.balls(kmod), CU)
sol_geo = (np.array(_P) * MM, np.array(_N, float), np.array(_I))
glow_geo = box(-0.74 * MM, -1.01 * MM, 0.29 * MM, 0.74 * MM, 1.01 * MM, 0.85 * MM)

# ---- LED from step2glb ----
s, sb = gltfio.load(s2g)
ROLE = {}                       # material index -> role, by the STEP colours (linear, as step2glb writes them)
KNOWN = {'body': (0.8900, 0.8900, 0.8481), 'window': (0.6383, 0.6739, 0.7106), 'terminal': (0.5705, 0.5705, 0.5382),
         'polarity_mark': (0.0134, 0.0134, 0.0134), 'die_R': (0.6035, 0.0903, 0.0624), 'die_G': (0.0953, 0.4999, 0.2543),
         'die_B': (0.0361, 0.1300, 0.5034), 'ic_die': (0.1005, 0.1065, 0.1193), 'laser_etch': (0.012, 0.012, 0.012)}
for mi, m in enumerate(s['materials']):
    c = m['pbrMetallicRoughness']['baseColorFactor'][:3]
    ROLE[mi] = min(KNOWN, key=lambda r: max(abs(a - b) for a, b in zip(c, KNOWN[r])))
    assert max(abs(a - b) for a, b in zip(c, KNOWN[ROLE[mi]])) < 5e-4, (c, ROLE[mi])
G = {}
for mi, P, N, I in gltfio.flatten(s, sb):
    G.setdefault(ROLE[mi], []).append((P, N, I))
G = {r: merge(v) for r, v in G.items()}
assert set(G) == set(KNOWN), set(G)

d, binc = gltfio.load(old); arrays = gltfio.explode(d, binc)
names = [n.get('name') for n in d['nodes']]
def mesh_of(name): return d['nodes'][names.index(name)]['mesh']
def put(mesh_idx, geo, key, mat):
    P, N, I = fix_winding(*geo)
    arrays[key + '_P'] = (P.astype(np.float32), 34962); arrays[key + '_N'] = (N.astype(np.float32), 34962); arrays[key + '_I'] = (I, 34963)
    return {'attributes': {'POSITION': key + '_P', 'NORMAL': key + '_N'}, 'indices': key + '_I', 'material': mat, 'mode': 4}
matn = {m['name']: i for i, m in enumerate(d['materials'])}
d['meshes'][mesh_of('pad_top')]['primitives'] = [put(0, pad_geo, 'pads', matn['pad'])]
d['meshes'][mesh_of('silk')]['primitives'] = [put(0, silk_geo, 'silk', matn['silk'])]
d['meshes'][mesh_of('solder_top')]['primitives'] = [put(0, sol_geo, 'solder', matn['solder'])]
d['meshes'][mesh_of('LED_lit')]['primitives'] = [put(0, glow_geo, 'glow', matn['led_glow'])]
# LED materials: reuse the three old slots, add the rest
def srgb2lin(c): return c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4
mats = d['materials']
mats[matn['body_white']].update(name='body_white', pbrMetallicRoughness={'baseColorFactor': list(KNOWN['body']) + [1], 'metallicFactor': 0, 'roughnessFactor': 0.55})
mats[matn['die_window']].update(name='window_milky', alphaMode='BLEND', pbrMetallicRoughness={'baseColorFactor': list(KNOWN['window']) + [0.55], 'metallicFactor': 0, 'roughnessFactor': 0.2},
                                extensions={'KHR_materials_clearcoat': {'clearcoatFactor': 1.0, 'clearcoatRoughnessFactor': 0.08}})
mats[matn['terminal']].update(name='terminal_tinned', pbrMetallicRoughness={'baseColorFactor': list(KNOWN['terminal']) + [1], 'metallicFactor': 0.7, 'roughnessFactor': 0.35})
if 'KHR_materials_clearcoat' not in d.setdefault('extensionsUsed', []): d['extensionsUsed'].append('KHR_materials_clearcoat')
for role, rough in (('polarity_mark', 0.6), ('die_R', 0.45), ('die_G', 0.45), ('die_B', 0.45), ('ic_die', 0.4)):
    mats.append({'name': role, 'pbrMetallicRoughness': {'baseColorFactor': list(KNOWN[role]) + [1], 'metallicFactor': 0, 'roughnessFactor': rough}}); matn[role] = len(mats) - 1
slot = {'body': matn['body_white'], 'window': matn['die_window'] if 'die_window' in matn else None, 'terminal': matn['terminal']}
slot['window'] = [i for i, m in enumerate(mats) if m['name'] == 'window_milky'][0]
slot['terminal'] = [i for i, m in enumerate(mats) if m['name'] == 'terminal_tinned'][0]
slot.update({r: matn[r] for r in ('polarity_mark', 'die_R', 'die_G', 'die_B', 'ic_die')})
led = d['nodes'][names.index('WS2812B-2020')]
# opaque parts first, the translucent window last (draw order)
order = ['body', 'terminal', 'polarity_mark', 'ic_die', 'die_R', 'die_G', 'die_B', 'window']
d['meshes'][led['mesh']]['primitives'] = [put(0, G[r], 'led_' + r, slot[r]) for r in order]
led['scale'] = [1, 1, 1]
# laser etch child
mats.append({'name': 'laser_etch', 'pbrMetallicRoughness': {'baseColorFactor': list(KNOWN['laser_etch']) + [1], 'metallicFactor': 0, 'roughnessFactor': 0.6}, 'doubleSided': True})
d['meshes'].append({'name': 'laser_etch', 'primitives': [put(0, G['laser_etch'], 'etch', len(mats) - 1)]})
d['nodes'].append({'name': 'Adom.LaserEtch.MPN', 'mesh': len(d['meshes']) - 1,
                   'extras': {'adomRole': 'laser_etch', 'virtualMark': True, 'text': 'WS2812B-2020', 'face': 'side (-Y, white body band)',
                              'tooltip': 'Virtual laser-etched MPN (WS2812B-2020)\nidentification aid, not the manufacturer marking'}})
led.setdefault('children', []).append(len(d['nodes']) - 1)
d.setdefault('asset', {})['generator'] = 'adom ws2812b-2020-fix rebuild_ws_scene.py (scene from ' + old.split('/')[-1] + ', LED from step2glb of the datasheet-built STEP)'
n = gltfio.save(out, d, arrays)
tris = {r: int(len(G[r][2]) // 3) for r in G}
print(json.dumps(dict(out=out.split('/')[-1], bytes=n, sha256=hashlib.sha256(open(out, 'rb').read()).hexdigest(), led_tris=tris,
                      pads_m=[[p[0], round(p[1], 3), round(p[2], 3), p[3], p[4]] for p in pads], silk_lines=len(silk),
                      led_bbox_mm=[[round(float(min(G[r][0][:, k].min() for r in G)) * 1000, 3), round(float(max(G[r][0][:, k].max() for r in G)) * 1000, 3)] for k in range(3)])))