@takram/three-geospatial 0.7.1 → 0.9.0

package/src/webgpu/FnLayout.ts

@@ -1,17 +1,17 @@
 import type {
   ProxiedTuple,
   ShaderCallNodeInternal,
-  ShaderNodeFn,
-  Struct
+  ShaderNodeFn
 } from 'three/src/nodes/TSL.js'
 import { Fn } from 'three/tsl'
-import type { NodeBuilder, Texture3DNode, TextureNode } from 'three/webgpu'
+import { StructTypeNode, type NodeBuilder } from 'three/webgpu'
 
 import type { Node, NodeType } from './node'
 
-// Note that "texture" and "texture3D" are just placeholders until TSL supports
-// texture types.
-type FnLayoutType = NodeType | Struct | 'texture' | 'texture3D'
+type FnLayoutType =
+  | NodeType
+  | (new (...args: any[]) => any)
+  | ((...args: any[]) => any)
 
 export interface FnLayoutInput<T extends FnLayoutType = FnLayoutType> {
   name: string
@@ -30,13 +30,11 @@ export interface FnLayout<
 
 type InferNodeObject<T extends FnLayoutType> = T extends NodeType
   ? Node<T>
-  : T extends Struct
-    ? ReturnType<T>
-    : T extends 'texture'
-      ? TextureNode
-      : T extends 'texture3D'
-        ? Texture3DNode
-        : never
+  : T extends new (...args: any[]) => any
+    ? InstanceType<T>
+    : T extends (...args: any[]) => any
+      ? ReturnType<T>
+      : never
 
 type InferNodeObjects<Inputs extends readonly FnLayoutInput[]> = {
   [K in keyof Inputs]: Inputs[K] extends FnLayoutInput<infer T>
@@ -58,10 +56,13 @@ function transformType(type: FnLayoutType): string {
   if (typeof type === 'string') {
     return type
   }
-  if (type.layout.name == null) {
-    throw new Error('Struct name is required.')
+  if ('layout' in type && type.layout instanceof StructTypeNode) {
+    if (type.layout.name == null) {
+      throw new Error('Struct name is required.')
+    }
+    return type.layout.name
   }
-  return type.layout.name
+  throw new Error(`Unsupported layout type: ${type}`)
 }
 
 export function FnLayout<

package/src/webgpu/HighpVelocityNode.ts

@@ -20,11 +20,11 @@ export class HighpVelocityNode extends TempNode {
 
   projectionMatrix?: Matrix4 | null
 
-  private readonly currentProjectionMatrix = uniform(new Matrix4())
-  private readonly previousProjectionMatrix = uniform('mat4' as const)
+  private readonly currentProjectionMatrix = uniform('mat4')
+  private readonly previousProjectionMatrix = uniform('mat4')
 
-  private readonly currentModelViewMatrix = uniform(new Matrix4())
-  private readonly previousModelViewMatrix = uniform('mat4' as const)
+  private readonly currentModelViewMatrix = uniform('mat4')
+  private readonly previousModelViewMatrix = uniform('mat4')
   private readonly objectModelViewMatrices = new WeakMap<Object3D, Matrix4>()
 
   constructor() {
@@ -42,6 +42,11 @@ export class HighpVelocityNode extends TempNode {
     this.updateAfterType = NodeUpdateType.OBJECT
   }
 
+  setProjectionMatrix(value: Matrix4 | null): this {
+    this.projectionMatrix = value
+    return this
+  }
+
   // Executed once per frame:
   override update({ camera }: NodeFrame): void {
     if (camera == null) {

package/src/webgpu/LensFlareNode.ts

@@ -1,5 +1,10 @@
-import { add, Fn, uniform } from 'three/tsl'
-import { TempNode, type NodeBuilder, type TextureNode } from 'three/webgpu'
+import { add, convertToTexture, Fn, rtt, uniform } from 'three/tsl'
+import {
+  TempNode,
+  type NodeBuilder,
+  type RTTNode,
+  type TextureNode
+} from 'three/webgpu'
 import invariant from 'tiny-invariant'
 
 import { DownsampleThresholdNode } from './DownsampleThresholdNode'
@@ -9,11 +14,6 @@ import { LensGlareNode } from './LensGlareNode'
 import { LensHaloNode } from './LensHaloNode'
 import { MipmapSurfaceBlurNode } from './MipmapSurfaceBlurNode'
 import type { Node } from './node'
-import {
-  convertToTexture,
-  rtTexture,
-  type RTTextureNode
-} from './RTTextureNode'
 import { isWebGPU } from './utils'
 
 export class LensFlareNode extends TempNode {
@@ -31,7 +31,7 @@ export class LensFlareNode extends TempNode {
 
   bloomIntensity = uniform(0.05)
 
-  featuresNode: RTTextureNode
+  featuresNode: RTTNode
 
   constructor(inputNode?: TextureNode | null) {
     super('vec4')
@@ -44,9 +44,9 @@ export class LensFlareNode extends TempNode {
     this.bloomNode = new MipmapSurfaceBlurNode(null, 8)
     this.glareNode = new LensGlareNode()
 
-    this.featuresNode = rtTexture(add(this.ghostNode, this.haloNode))
-    this.featuresNode.value.name = 'LensFlareNode.Features'
-    this.featuresNode.resolutionScale = 0.5
+    this.featuresNode = rtt(add(this.ghostNode, this.haloNode))
+    this.featuresNode.value.name = 'LensFlare_features'
+    this.featuresNode.pixelRatio = 0.5
 
     // Use the full resolution because the thresholdNode already downsamples the
     // input texture.
@@ -112,8 +112,4 @@ export class LensFlareNode extends TempNode {
 }
 
 export const lensFlare = (inputNode: Node | null): LensFlareNode =>
-  new LensFlareNode(
-    inputNode != null
-      ? convertToTexture(inputNode, 'LensFlareNode.Input')
-      : null
-  )
+  new LensFlareNode(inputNode != null ? convertToTexture(inputNode) : null)

package/src/webgpu/LensGlareNode.ts

@@ -10,6 +10,7 @@ import {
 import { hash } from 'three/src/nodes/core/NodeUtils.js'
 import {
   atomicAdd,
+  convertToTexture,
   Fn,
   globalId,
   If,
@@ -28,9 +29,9 @@ import {
   vec4
 } from 'three/tsl'
 import {
+  IndirectStorageBufferAttribute,
   MeshBasicNodeMaterial,
   RendererUtils,
-  StorageBufferAttribute,
   type ComputeNode,
   type NodeBuilder,
   type NodeFrame,
@@ -40,7 +41,6 @@ import invariant from 'tiny-invariant'
 
 import { FilterNode } from './FilterNode'
 import type { Node } from './node'
-import { convertToTexture } from './RTTextureNode'
 
 const { resetRendererState, restoreRendererState } = RendererUtils
 
@@ -92,7 +92,12 @@ export class LensGlareNode extends FilterNode {
 
   private computeNode?: ComputeNode
 
-  private readonly counterBuffer = new StorageBufferAttribute(1, 1)
+  // drawIndexedIndirect format:
+  // [indexCount, instanceCount, firstIndex, baseVertex, firstInstance]
+  private readonly indirectBuffer = new IndirectStorageBufferAttribute(
+    new Uint32Array([6, 0, 0, 0, 0]),
+    1
+  )
   private instanceBuffer = instancedArray(1, instanceStruct)
 
   private readonly renderTarget = this.createRenderTarget()
@@ -106,16 +111,19 @@ export class LensGlareNode extends FilterNode {
   private readonly camera = new PerspectiveCamera()
   private rendererState?: RendererUtils.RendererState
 
-  private readonly inputTexelSize = uniform(new Vector2())
-  private readonly outputTexelSize = uniform(new Vector2())
-  private readonly geometryRatio = uniform(new Vector2())
+  private readonly inputTexelSize = uniform('vec2')
+  private readonly outputTexelSize = uniform('vec2')
+  private readonly geometryRatio = uniform('vec2')
 
   constructor(inputNode?: TextureNode | null) {
     super(inputNode)
+    this.material.name = 'LensGlare'
+
     this.inputNode = inputNode
     this.resolutionScale = 0.5
 
     this.outputTexture = this.renderTarget.texture
+    this.mesh.geometry.indirect = this.indirectBuffer
   }
 
   override customCacheKey(): number {
@@ -151,7 +159,7 @@ export class LensGlareNode extends FilterNode {
     const { width: inputWidth, height: inputHeight } = inputNode.value
     this.setSize(inputWidth, inputHeight) // Compute node is initialized here.
 
-    const { computeNode, counterBuffer, renderTarget } = this
+    const { computeNode, indirectBuffer, renderTarget } = this
     invariant(computeNode != null)
 
     this.inputTexelSize.value.set(1 / inputWidth, 1 / inputHeight)
@@ -165,24 +173,12 @@ export class LensGlareNode extends FilterNode {
     const { width: outputWidth, height: outputHeight } = renderTarget
     this.outputTexelSize.value.set(1 / outputWidth, 1 / outputHeight)
 
-    // Reset the counter:
-    counterBuffer.array[0] = 0
-    counterBuffer.needsUpdate = true
+    // Reset instanceCount in the indirect buffer:
+    indirectBuffer.array[1] = 0
+    indirectBuffer.needsUpdate = true
 
     void renderer.compute(computeNode)
 
-    renderer
-      .getArrayBufferAsync(counterBuffer)
-      .then(arrayBuffer => {
-        // TODO: This is indeed a couple of frames behind, thus the number of
-        // computed instances above and the number of instances to be drawn by
-        // the mesh differ.
-        this.mesh.count = new Uint32Array(arrayBuffer)[0]
-      })
-      .catch((error: unknown) => {
-        console.error(error)
-      })
-
     this.rendererState = resetRendererState(renderer, this.rendererState)
 
     renderer.setRenderTarget(renderTarget)
@@ -195,16 +191,16 @@ export class LensGlareNode extends FilterNode {
     const {
       spikePairCount,
       inputNode,
-      counterBuffer,
+      indirectBuffer,
       instanceBuffer,
       outputTexelSize
     } = this
     invariant(inputNode != null)
 
-    const counterStorage = storage(
-      counterBuffer,
+    const indirectStorage = storage(
+      indirectBuffer,
       'uint',
-      counterBuffer.count
+      indirectBuffer.count
     ).toAtomic()
 
     this.computeNode = Fn(() => {
@@ -218,7 +214,11 @@ export class LensGlareNode extends FilterNode {
       const inputLuminance = inputColor.a // Alpha channel stores luminance
 
       If(inputLuminance.greaterThan(0.1), () => {
-        const countBefore = atomicAdd(counterStorage.element(0), spikePairCount)
+        // The first element is instanceCount in the drawIndexedIndirect buffer.
+        const countBefore = atomicAdd(
+          indirectStorage.element(1),
+          spikePairCount
+        )
         for (let i = 0; i < spikePairCount; ++i) {
           const instance = instanceBuffer.element(countBefore.add(i))
           instance.get('color').assign(inputColor.rgb)
@@ -308,8 +308,4 @@ export class LensGlareNode extends FilterNode {
 }
 
 export const lensGlare = (inputNode: Node | null): LensGlareNode =>
-  new LensGlareNode(
-    inputNode != null
-      ? convertToTexture(inputNode, 'LensGlareNode.Input')
-      : null
-  )
+  new LensGlareNode(inputNode != null ? convertToTexture(inputNode) : null)

package/src/webgpu/LensHaloNode.ts

@@ -38,8 +38,9 @@ export class LensHaloNode extends TempNode {
 
   constructor(inputNode?: TextureNode | null) {
     super('vec3')
-    this.inputNode = inputNode
     this.updateBeforeType = NodeUpdateType.FRAME
+
+    this.inputNode = inputNode
   }
 
   override updateBefore({ renderer }: NodeFrame): void {

package/src/webgpu/OutputTexture3DNode.ts

@@ -1,6 +1,8 @@
 import type { Texture } from 'three'
 import { Texture3DNode, type Node, type NodeBuilder } from 'three/webgpu'
 
+import { reinterpretType } from '../types'
+
 export class OutputTexture3DNode extends Texture3DNode {
   static override get type(): string {
     return 'OutputTexture3DNode'
@@ -11,6 +13,14 @@ export class OutputTexture3DNode extends Texture3DNode {
   constructor(owner: Node, texture: Texture) {
     super(texture)
     this.owner = owner
+
+    // WORKAROUND: Missing method as of r182. Adding these in the module
+    // augmentation breaks VSCode's auto completion.
+    reinterpretType<
+      typeof this & {
+        setUpdateMatrix: (value: boolean) => void
+      }
+    >(this)
     this.setUpdateMatrix(false)
   }
 

package/src/webgpu/OutputTextureNode.ts

@@ -1,6 +1,8 @@
 import type { Texture } from 'three'
 import { TextureNode, type Node, type NodeBuilder } from 'three/webgpu'
 
+import { reinterpretType } from '../types'
+
 export class OutputTextureNode extends TextureNode {
   static override get type(): string {
     return 'OutputTextureNode'
@@ -11,6 +13,14 @@ export class OutputTextureNode extends TextureNode {
   constructor(owner: Node, texture: Texture) {
     super(texture)
     this.owner = owner
+
+    // WORKAROUND: Missing method as of r182. Adding these in the module
+    // augmentation breaks VSCode's auto completion.
+    reinterpretType<
+      typeof this & {
+        setUpdateMatrix: (value: boolean) => void
+      }
+    >(this)
     this.setUpdateMatrix(false)
   }
 

package/src/webgpu/STBNTextureNode.ts

@@ -0,0 +1,58 @@
+import { hashString } from 'three/src/nodes/core/NodeUtils.js'
+import { Fn, frameId, nodeImmutable, screenCoordinate, vec3 } from 'three/tsl'
+import {
+  Data3DTexture,
+  NearestFilter,
+  RedFormat,
+  RepeatWrapping,
+  Texture3DNode,
+  type NodeBuilder
+} from 'three/webgpu'
+
+import { DEFAULT_STBN_URL } from '../constants'
+import { STBNLoader } from '../STBNLoader'
+
+const emptyTexture3D = /*#__PURE__*/ (() => {
+  const texture = new Data3DTexture(new Uint8Array(1))
+  texture.format = RedFormat
+  // BUG: TextureNode doesn't update these when the texture is swapped.
+  texture.minFilter = NearestFilter
+  texture.magFilter = NearestFilter
+  texture.wrapS = RepeatWrapping
+  texture.wrapT = RepeatWrapping
+  texture.wrapR = RepeatWrapping
+  texture.needsUpdate = true
+  return texture
+})()
+
+export class STBNTextureNode extends Texture3DNode {
+  url = DEFAULT_STBN_URL
+
+  constructor() {
+    super(emptyTexture3D)
+  }
+
+  override customCacheKey(): number {
+    return hashString(this.url)
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    new STBNLoader()
+      .loadAsync(this.url)
+      .then(texture => {
+        this.value = texture
+      })
+      .catch((error: unknown) => {
+        console.error(error)
+      })
+    return super.setup(builder)
+  }
+}
+
+export const stbnTexture = /*#__PURE__*/ nodeImmutable(STBNTextureNode)
+
+export const stbn = /*#__PURE__*/ Fn(() => {
+  return stbnTexture
+    .sample(vec3(screenCoordinate.xy, frameId.mod(64)).div(vec3(128, 128, 64)))
+    .r.toConst('stbn')
+}).once()()