@takram/three-geospatial 0.4.0 → 0.5.0

package/src/webgpu/FnLayout.ts

@@ -0,0 +1,86 @@
+import type {
+  ProxiedTuple,
+  ShaderCallNodeInternal,
+  ShaderNodeFn,
+  Struct
+} from 'three/src/nodes/TSL.js'
+import { Fn } from 'three/tsl'
+import type { NodeBuilder, Texture3DNode, TextureNode } from 'three/webgpu'
+
+import type { NodeObject, NodeType } from './node'
+
+// Note that "texture" and "texture3D" are just placeholders until TSL supports
+// texture types.
+type FnLayoutType = NodeType | Struct | 'texture' | 'texture3D'
+
+export interface FnLayoutInput<T extends FnLayoutType = FnLayoutType> {
+  name: string
+  type: T
+}
+
+export interface FnLayout<
+  T extends FnLayoutType,
+  Inputs extends readonly FnLayoutInput[] = []
+> {
+  typeOnly?: boolean
+  name: string
+  type: T
+  inputs?: Inputs
+}
+
+type InferNodeObject<T extends FnLayoutType> = T extends NodeType
+  ? NodeObject<T>
+  : T extends Struct
+    ? ReturnType<T>
+    : T extends 'texture'
+      ? NodeObject<TextureNode>
+      : T extends 'texture3D'
+        ? NodeObject<Texture3DNode>
+        : never
+
+type InferNodeObjects<Inputs extends readonly FnLayoutInput[]> = {
+  [K in keyof Inputs]: Inputs[K] extends FnLayoutInput<infer T>
+    ? InferNodeObject<T>
+    : never
+}
+
+export type FnLayoutResult<
+  T extends FnLayoutType,
+  Inputs extends readonly FnLayoutInput[],
+  Nodes extends readonly unknown[] = InferNodeObjects<Inputs>
+> = (
+  callback: (
+    ...args: [] extends Nodes ? [NodeBuilder] : [Nodes, NodeBuilder]
+  ) => InferNodeObject<T> | NodeObject<ShaderCallNodeInternal>
+) => ShaderNodeFn<ProxiedTuple<Nodes>>
+
+function transformType(type: FnLayoutType): string {
+  if (typeof type === 'string') {
+    return type
+  }
+  if (type.layout.name == null) {
+    throw new Error('Struct name is required.')
+  }
+  return type.layout.name
+}
+
+export function FnLayout<
+  T extends FnLayoutType,
+  const Inputs extends readonly FnLayoutInput[] = []
+>({
+  typeOnly = false,
+  ...layout
+}: FnLayout<T, Inputs>): FnLayoutResult<T, Inputs> {
+  return typeOnly
+    ? callback => Fn(callback as any)
+    : callback =>
+        Fn(callback as any).setLayout({
+          ...layout,
+          type: transformType(layout.type),
+          inputs:
+            layout.inputs?.map(input => ({
+              ...input,
+              type: transformType(input.type)
+            })) ?? []
+        })
+}

package/src/webgpu/FnVar.ts

@@ -0,0 +1,26 @@
+import type { ProxiedTuple, ShaderNodeFn } from 'three/src/nodes/TSL.js'
+import { Fn } from 'three/tsl'
+import type { NodeBuilder } from 'three/webgpu'
+
+type NonCallable<T> = T extends (...args: any[]) => any ? never : T
+
+export function FnVar<Args extends readonly unknown[], R>(
+  callback: (...args: Args) => (builder: NodeBuilder) => R
+): ShaderNodeFn<ProxiedTuple<Args>>
+
+export function FnVar<Args extends readonly unknown[], R>(
+  callback: (...args: Args) => NonCallable<R>
+): ShaderNodeFn<ProxiedTuple<Args>>
+
+export function FnVar<Args extends readonly unknown[], R>(
+  callback:
+    | ((...args: Args) => R)
+    | ((...args: Args) => (builder: NodeBuilder) => R)
+): ShaderNodeFn<ProxiedTuple<Args>> {
+  return Fn((args: Args, builder: NodeBuilder) => {
+    const result = callback(...args)
+    return typeof result === 'function'
+      ? (result as (builder: NodeBuilder) => R)(builder)
+      : result
+  })
+}

package/src/webgpu/GaussianBlurNode.ts

@@ -0,0 +1,129 @@
+import { add, Fn, nodeObject, uv } from 'three/tsl'
+import type { NodeBuilder, TextureNode } from 'three/webgpu'
+import invariant from 'tiny-invariant'
+
+import type { Node, NodeObject } from './node'
+import { SeparableFilterNode } from './SeparableFilterNode'
+
+// Reference: https://github.com/pmndrs/postprocessing/blob/v6.37.8/src/core/GaussKernel.js
+
+function getCoefficients(n: number): Float64Array {
+  invariant(n >= 0)
+  if (n === 0) {
+    return new Float64Array(0)
+  }
+  if (n === 1) {
+    return new Float64Array([1])
+  }
+
+  let row0 = new Float64Array(n)
+  let row1 = new Float64Array(n)
+  let result = row1
+
+  // Incrementally build Pascal's Triangle to get the desired row.
+  for (let y = 1; y <= n; ++y) {
+    for (let x = 0; x < y; ++x) {
+      row1[x] = x === 0 || x === y - 1 ? 1 : row0[x - 1] + row0[x]
+    }
+    result = row1
+    row1 = row0
+    row0 = result
+  }
+  return result
+}
+
+interface GaussianKernel {
+  weights: Float64Array
+  offsets: Float64Array
+}
+
+function createGaussianKernel(
+  kernelSize: number,
+  edgeBias = 2
+): GaussianKernel {
+  invariant(kernelSize >= 3)
+
+  const n = kernelSize + edgeBias * 2
+  const coefficients =
+    edgeBias > 0
+      ? getCoefficients(n).slice(edgeBias, -edgeBias)
+      : getCoefficients(n)
+
+  const mid = Math.floor((coefficients.length - 1) / 2)
+  const sum = coefficients.reduce((a, b) => a + b, 0)
+  const weights = coefficients.slice(mid)
+  const offsets = [...Array(mid + 1).keys()] // [0..mid+1]
+
+  const linearWeights = new Float64Array(Math.floor(offsets.length / 2))
+  const linearOffsets = new Float64Array(linearWeights.length)
+  linearWeights[0] = weights[0] / sum
+
+  for (let i = 1, j = 1; i < offsets.length - 1; i += 2, ++j) {
+    const offset0 = offsets[i]
+    const offset1 = offsets[i + 1]
+    const weight0 = weights[i]
+    const weight1 = weights[i + 1]
+
+    const weight = weight0 + weight1
+    const offset = (offset0 * weight0 + offset1 * weight1) / weight
+    linearWeights[j] = weight / sum
+    linearOffsets[j] = offset
+  }
+
+  // Ensure that the weights add up to 1.
+  const linearWeightSum =
+    (linearWeights.reduce((a, b) => a + b, 0) - linearWeights[0] * 0.5) * 2
+
+  if (linearWeightSum !== 0) {
+    const scale = 1 / linearWeightSum
+    for (let i = 0; i < linearWeights.length; ++i) {
+      linearWeights[i] *= scale
+    }
+  }
+
+  return {
+    offsets: linearOffsets,
+    weights: linearWeights
+  }
+}
+
+export class GaussianBlurNode extends SeparableFilterNode {
+  static override get type(): string {
+    return 'GaussianBlurNode'
+  }
+
+  private readonly kernelSize: number
+
+  constructor(inputNode?: TextureNode | null, kernelSize = 35) {
+    super(inputNode)
+    this.kernelSize = kernelSize
+  }
+
+  protected override setupOutputNode(builder: NodeBuilder): Node {
+    const { inputNode, inputTexelSize, direction } = this
+    invariant(inputNode != null)
+
+    const { offsets, weights } = createGaussianKernel(this.kernelSize)
+
+    return Fn(() => {
+      const center = uv()
+      const offsetSize = direction.mul(inputTexelSize).toVertexStage()
+
+      const output = inputNode.sample(center).mul(weights[0])
+      for (let i = 1; i < offsets.length; ++i) {
+        const offset = offsetSize.mul(offsets[i])
+        output.addAssign(
+          add(
+            inputNode.sample(center.add(offset)),
+            inputNode.sample(center.sub(offset))
+          ).mul(weights[i])
+        )
+      }
+      return output
+    })()
+  }
+}
+
+export const gaussianBlur = (
+  ...args: ConstructorParameters<typeof GaussianBlurNode>
+): NodeObject<GaussianBlurNode> => nodeObject(new GaussianBlurNode(...args))

package/src/webgpu/HighpVelocityNode.ts

@@ -0,0 +1,115 @@
+import { Matrix4, type Object3D } from 'three'
+import type { NodeFrame } from 'three/src/Three.WebGPU.js'
+import {
+  nodeImmutable,
+  positionLocal,
+  positionPrevious,
+  sub,
+  uniform
+} from 'three/tsl'
+import { NodeUpdateType, TempNode, type NodeBuilder } from 'three/webgpu'
+
+export class HighpVelocityNode extends TempNode {
+  static override get type(): string {
+    return 'HighpVelocityNode'
+  }
+
+  projectionMatrix?: Matrix4 | null
+
+  private readonly currentProjectionMatrix = uniform(new Matrix4())
+  private readonly previousProjectionMatrix = uniform('mat4')
+
+  private readonly currentModelViewMatrix = uniform(new Matrix4())
+  private readonly previousModelViewMatrix = uniform('mat4')
+  private readonly objectModelViewMatrices = new WeakMap<Object3D, Matrix4>()
+
+  constructor() {
+    super('vec3')
+
+    // Sequence:
+    // - updateBefore() for the first object
+    // - update() for the current frame
+    // - updateAfter() for the first object
+    // - updateBefore() for the next object
+    // - updateAfter() for the next object
+    // - ...
+    this.updateType = NodeUpdateType.FRAME
+    this.updateBeforeType = NodeUpdateType.OBJECT
+    this.updateAfterType = NodeUpdateType.OBJECT
+  }
+
+  // Executed once per frame:
+  override update({ camera }: NodeFrame): void {
+    if (camera == null) {
+      return
+    }
+    const {
+      currentProjectionMatrix: current,
+      previousProjectionMatrix: previous
+    } = this
+
+    const projectionMatrix = this.projectionMatrix ?? camera.projectionMatrix
+    if (previous.value == null) {
+      previous.value = new Matrix4().copy(projectionMatrix)
+    } else {
+      previous.value.copy(current.value)
+    }
+    current.value.copy(projectionMatrix)
+  }
+
+  // Executed once per object before rendering:
+  override updateBefore({ object, camera }: NodeFrame): void {
+    if (object == null || camera == null) {
+      return
+    }
+    const {
+      currentModelViewMatrix: current,
+      previousModelViewMatrix: previous,
+      objectModelViewMatrices: matrices
+    } = this
+
+    current.value.multiplyMatrices(
+      camera.matrixWorldInverse,
+      object.matrixWorld
+    )
+    previous.value = matrices.get(object) ?? current.value
+  }
+
+  // Executed once per object after rendering:
+  override updateAfter({ object }: NodeFrame): void {
+    if (object == null) {
+      return
+    }
+    const {
+      currentModelViewMatrix: current,
+      objectModelViewMatrices: matrices
+    } = this
+
+    let matrix = matrices.get(object)
+    if (matrix == null) {
+      matrix = new Matrix4()
+      matrices.set(object, matrix)
+    }
+    matrix.copy(current.value)
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    const currentClip = this.currentProjectionMatrix
+      .mul(this.currentModelViewMatrix)
+      .mul(positionLocal)
+      .toVertexStage()
+    const previousClip = this.previousProjectionMatrix
+      .mul(this.previousModelViewMatrix)
+      .mul(positionPrevious)
+      .toVertexStage()
+
+    // Perspective divisions cannot be performed in the vertex shader.
+    // See: http://john-chapman-graphics.blogspot.com/2013/01/per-object-motion-blur.html
+    const currentNDC = currentClip.xyz.div(currentClip.w)
+    const previousNDC = previousClip.xyz.div(previousClip.w)
+
+    return sub(currentNDC, previousNDC)
+  }
+}
+
+export const highpVelocity = /*#__PURE__*/ nodeImmutable(HighpVelocityNode)

package/src/webgpu/KawaseBlurNode.ts

@@ -0,0 +1,76 @@
+import { add, nodeObject, uv, vec2, vec4 } from 'three/tsl'
+import type { NodeBuilder, TextureNode } from 'three/webgpu'
+import invariant from 'tiny-invariant'
+
+import { DualMipmapFilterNode } from './DualMipmapFilterNode'
+import type { Node, NodeObject } from './node'
+
+export class KawaseBlurNode extends DualMipmapFilterNode {
+  static override get type(): string {
+    return 'KawaseBlurNode'
+  }
+
+  constructor(inputNode?: TextureNode | null, levels = 4) {
+    super(inputNode, levels)
+    this.resolutionScale = 0.5
+  }
+
+  protected override setupDownsampleNode(builder: NodeBuilder): Node {
+    const { inputNode, inputTexelSize } = this
+    invariant(inputNode != null)
+
+    const center = uv()
+    const offset = vec4(1, 1, -1, -1)
+      .mul(inputTexelSize.xyxy.mul(0.5))
+      .add(center.xyxy)
+    const uv1 = offset.zy.toVertexStage() // -0.5, 0.5
+    const uv2 = offset.xy.toVertexStage() // 0.5, 0.5
+    const uv3 = offset.xw.toVertexStage() // 0.5, -0.5
+    const uv4 = offset.zw.toVertexStage() // -0.5, -0.5
+
+    return add(
+      inputNode.sample(center).mul(4),
+      inputNode.sample(uv1),
+      inputNode.sample(uv2),
+      inputNode.sample(uv3),
+      inputNode.sample(uv4)
+    ).mul(1 / 8)
+  }
+
+  protected override setupUpsampleNode(builder: NodeBuilder): Node {
+    const { inputNode, inputTexelSize } = this
+    invariant(inputNode != null)
+
+    const center = uv()
+    const offset = vec4(1, 1, -1, -1)
+      .mul(inputTexelSize.xyxy.mul(0.5))
+      .add(center.xyxy)
+    const uv1 = offset.zy.toVertexStage() // -0.5, 0.5
+    const uv2 = offset.xy.toVertexStage() // 0.5, 0.5
+    const uv3 = offset.xw.toVertexStage() // 0.5, -0.5
+    const uv4 = offset.zw.toVertexStage() // -0.5, -0.5
+    const uv5 = vec2(offset.z, center.y).toVertexStage() // -0.5, 0
+    const uv6 = vec2(offset.x, center.y).toVertexStage() // 0.5, 0
+    const uv7 = vec2(center.x, offset.y).toVertexStage() // 0, 0.5
+    const uv8 = vec2(center.x, offset.w).toVertexStage() // 0, -0.5
+
+    return add(
+      add(
+        inputNode.sample(uv1),
+        inputNode.sample(uv2),
+        inputNode.sample(uv3),
+        inputNode.sample(uv4)
+      ).mul(1 / 12),
+      add(
+        inputNode.sample(uv5),
+        inputNode.sample(uv6),
+        inputNode.sample(uv7),
+        inputNode.sample(uv8)
+      ).mul(1 / 6)
+    )
+  }
+}
+
+export const kawaseBlur = (
+  ...args: ConstructorParameters<typeof KawaseBlurNode>
+): NodeObject<KawaseBlurNode> => nodeObject(new KawaseBlurNode(...args))

package/src/webgpu/LensFlareNode.ts

@@ -0,0 +1,128 @@
+import { add, Fn, nodeObject, uniform } from 'three/tsl'
+import {
+  TempNode,
+  type Node,
+  type NodeBuilder,
+  type TextureNode
+} from 'three/webgpu'
+import invariant from 'tiny-invariant'
+
+import { DownsampleThresholdNode } from './DownsampleThresholdNode'
+import { GaussianBlurNode } from './GaussianBlurNode'
+import { LensGhostNode } from './LensGhostNode'
+import { LensGlareNode } from './LensGlareNode'
+import { LensHaloNode } from './LensHaloNode'
+import { MipmapSurfaceBlurNode } from './MipmapSurfaceBlurNode'
+import type { NodeObject } from './node'
+import {
+  convertToTexture,
+  rtTexture,
+  type RTTextureNode
+} from './RTTextureNode'
+import { isWebGPU } from './utils'
+
+export class LensFlareNode extends TempNode {
+  static override get type(): string {
+    return 'LensFlareNode'
+  }
+
+  inputNode?: TextureNode | null
+  thresholdNode: DownsampleThresholdNode
+  blurNode: GaussianBlurNode
+  ghostNode: LensGhostNode
+  haloNode: LensHaloNode
+  bloomNode: MipmapSurfaceBlurNode
+  glareNode: LensGlareNode
+
+  bloomIntensity = uniform(0.05)
+
+  featuresNode: RTTextureNode
+
+  constructor(inputNode?: TextureNode | null) {
+    super('vec4')
+    this.inputNode = inputNode
+
+    this.thresholdNode = new DownsampleThresholdNode()
+    this.blurNode = new GaussianBlurNode()
+    this.ghostNode = new LensGhostNode()
+    this.haloNode = new LensHaloNode()
+    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
+
+    // Use the full resolution because the thresholdNode already downsamples the
+    // input texture.
+    this.blurNode.resolutionScale = 1
+    this.bloomNode.resolutionScale = 1
+    this.glareNode.resolutionScale = 1
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    const {
+      inputNode,
+      thresholdNode,
+      blurNode,
+      ghostNode,
+      haloNode,
+      bloomNode,
+      featuresNode,
+      glareNode
+    } = this
+    invariant(inputNode != null)
+
+    const threshold = thresholdNode.getTextureNode()
+    const blur = blurNode.getTextureNode()
+
+    // input → threshold → blur → ghost
+    // input → threshold → blur → halo
+    thresholdNode.inputNode = inputNode
+    blurNode.inputNode = threshold
+    ghostNode.inputNode = blur
+    haloNode.inputNode = blur
+
+    // input → threshold → bloom
+    bloomNode.inputNode = threshold
+
+    // input → threshold → glare
+    glareNode.inputNode = threshold
+
+    const bloom = nodeObject(bloomNode.getTextureNode()).mul(
+      this.bloomIntensity
+    )
+    const glare = glareNode.getTextureNode()
+
+    // TODO: Add an option to switch to mixing the bloom:
+    return Fn(() => {
+      // TODO: Prevent the output from becoming too bright.
+      const output = nodeObject(inputNode)
+      output.addAssign(bloom)
+      if (isWebGPU(builder)) {
+        output.addAssign(glare)
+      }
+      return output.add(featuresNode)
+    })()
+  }
+
+  override dispose(): void {
+    this.thresholdNode.dispose()
+    this.blurNode.dispose()
+    this.ghostNode.dispose()
+    this.haloNode.dispose()
+    this.bloomNode.dispose()
+    this.glareNode.dispose()
+    this.featuresNode.dispose()
+    super.dispose()
+  }
+}
+
+export const lensFlare = (inputNode: Node | null): NodeObject<LensFlareNode> =>
+  nodeObject(
+    new LensFlareNode(
+      inputNode != null
+        ? convertToTexture(inputNode, 'LensFlareNode.Input')
+        : null
+    )
+  )

package/src/webgpu/LensGhostNode.ts

@@ -0,0 +1,62 @@
+import { add, sub, uniform, uv, vec3 } from 'three/tsl'
+import { TempNode, type NodeBuilder, type TextureNode } from 'three/webgpu'
+import invariant from 'tiny-invariant'
+
+import { FnLayout } from './FnLayout'
+
+export class LensGhostNode extends TempNode {
+  static override get type(): string {
+    return 'LensGhostNode'
+  }
+
+  inputNode?: TextureNode | null
+
+  intensity = uniform(1e-5)
+
+  constructor(inputNode?: TextureNode | null) {
+    super('vec3')
+    this.inputNode = inputNode
+  }
+
+  override setup(builder: NodeBuilder): unknown {
+    const { inputNode, intensity } = this
+    invariant(inputNode != null)
+
+    const sampleGhost = FnLayout({
+      name: 'sampleGhost',
+      type: 'vec3',
+      inputs: [
+        { name: 'uv', type: 'vec2' },
+        { name: 'direction', type: 'vec2' },
+        { name: 'color', type: 'vec3' },
+        { name: 'offset', type: 'float' }
+      ]
+    })(([uv, direction, color, offset]) => {
+      const suv = direction.mul(offset).add(uv.oneMinus()).saturate()
+      const result = inputNode.sample(suv).rgb.mul(color)
+
+      // Falloff at the perimeter:
+      const d = sub(0.5, suv)
+        .length()
+        .mul(1 / (Math.SQRT2 / 4))
+        .saturate()
+      result.mulAssign(d.oneMinus().pow(3))
+      return result
+    })
+
+    const uvNode = uv()
+    const direction = uvNode.sub(0.5)
+    const color = add(
+      sampleGhost(uvNode, direction, vec3(0.8, 0.8, 1), -5.0),
+      sampleGhost(uvNode, direction, vec3(1, 0.8, 0.4), -1.5),
+      sampleGhost(uvNode, direction, vec3(0.9, 1, 0.8), -0.4),
+      sampleGhost(uvNode, direction, vec3(1, 0.8, 0.4), -0.2),
+      sampleGhost(uvNode, direction, vec3(0.9, 0.7, 0.7), -0.1),
+      sampleGhost(uvNode, direction, vec3(0.5, 1, 0.4), 0.7),
+      sampleGhost(uvNode, direction, vec3(0.5, 0.5, 0.5), 1),
+      sampleGhost(uvNode, direction, vec3(1, 1, 0.6), 2.5),
+      sampleGhost(uvNode, direction, vec3(0.5, 0.8, 1), 10)
+    )
+    return color.mul(intensity)
+  }
+}