@takram/three-atmosphere
Version:
A Three.js and R3F implementation of Precomputed Atmospheric Scattering
419 lines (371 loc) • 15.4 kB
text/typescript
// Based on Intel's Outdoor Light Scattering Sample: https://github.com/GameTechDev/OutdoorLightScattering
/**
* Copyright 2017 Intel Corporation.
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
* use this file except in compliance with the License. You may obtain a copy of
* the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations under
* the License.
*
* Modified from the original source code.
*/
import {
HalfFloatType,
LinearFilter,
RenderTarget,
RGFormat,
type Camera,
type Vector2,
type Vector4
} from 'three'
import { float, Fn, max, min, uniform, uv, vec2, vec4 } from 'three/tsl'
import {
NodeMaterial,
NodeUpdateType,
QuadMesh,
RendererUtils,
type NodeBuilder,
type NodeFrame,
type TextureNode,
type UniformNode
} from 'three/webgpu'
import {
bvecAnd,
bvecNot,
Node,
outputTexture,
textureGather
} from '@takram/three-geospatial/webgpu'
import {
getCameraZUnit,
getOutermostScreenPixelCoords,
transformUVToNDC
} from './common'
const { resetRendererState, restoreRendererState } = RendererUtils
export class UnwarpEpipolarNode extends Node {
static override get type(): string {
return 'UnwarpEpipolarNode'
}
sliceEndpointsNode!: TextureNode
coordinateNode!: TextureNode
epipolarShadowLengthNode!: TextureNode
viewZUnitNode!: TextureNode // Must be filterable
camera!: Camera
epipolarSliceCount!: UniformNode<number> // float
maxSliceSampleCount!: UniformNode<number> // float
screenSize!: UniformNode<Vector2> // vec2
lightScreenPosition!: UniformNode<Vector4> // vec4
refinementThreshold = uniform(0.03)
private readonly textureNode: TextureNode
private readonly renderTarget: RenderTarget
private readonly material = new NodeMaterial()
private readonly mesh = new QuadMesh(this.material)
private rendererState?: RendererUtils.RendererState
constructor() {
super()
this.updateType = NodeUpdateType.FRAME // After CSM's updateBefore
this.material.name = 'UnwarpEpipolar'
this.mesh.name = 'UnwarpEpipolar'
const renderTarget = new RenderTarget(1, 1, {
depthBuffer: false,
type: HalfFloatType,
format: RGFormat
})
const texture = renderTarget.texture
texture.name = 'UnwarpEpipolar'
texture.minFilter = LinearFilter
texture.magFilter = LinearFilter
texture.generateMipmaps = false
this.renderTarget = renderTarget
this.textureNode = outputTexture(this, renderTarget.texture)
}
getTextureNode(): TextureNode {
return this.textureNode
}
override update({ renderer }: NodeFrame): void {
if (renderer == null) {
return
}
const { width, height } = this.screenSize.value
this.renderTarget.setSize(width, height)
this.rendererState = resetRendererState(renderer, this.rendererState)
renderer.setRenderTarget(this.renderTarget)
this.mesh.render(renderer)
restoreRendererState(renderer, this.rendererState)
}
private setupFragmentNode(builder: NodeBuilder): Node<'vec2'> {
const {
sliceEndpointsNode,
coordinateNode,
epipolarShadowLengthNode,
refinementThreshold,
viewZUnitNode,
camera,
maxSliceSampleCount,
epipolarSliceCount,
screenSize,
lightScreenPosition
} = this
return Fn(() => {
const uvNode = uv().toConst()
const positionNDC = transformUVToNDC(uvNode).toConst()
const cameraZ = getCameraZUnit(camera, uvNode, viewZUnitNode).toConst()
// Compute direction of the ray going from the light through the pixel.
const rayDirection = positionNDC
.sub(lightScreenPosition.xy)
.normalize()
.toConst()
// Find, which boundary the ray intersects. For this, we will find which
// two of four half spaces the rayDirection belongs to.
// Each of four half spaces is produced by the line connecting one of four
// screen corners and the current pixel:
// ________________ _______'________ ________________
// |' . '| | ' | | |
// | ' . ' | | ' | . | |
// | ' . ' | | ' | '|. hs1 |
// | *. | | * hs0 | | '*. |
// | ' ' . | | ' | | ' . |
// | ' ' . | | ' | | ' . |
// |'____________ '_| |'_______________| | ____________ '_.
// ' '
// ________________ . '________________
// | . '| |' |
// | hs2 . ' | | ' |
// | . ' | | ' |
// | . * | | * |
// . ' | | ' |
// | | | hs3 ' |
// |________________| |______'_________|
// '
// Note that in fact the outermost visible screen pixels do not lie
// exactly on the boundary (+1 or -1), but are biased by 0.5 screen pixel
// size inwards. Using these adjusted boundaries improves precision and
// results in smaller number of pixels which require correction.
const boundaries = getOutermostScreenPixelCoords(screenSize).toConst() // left, bottom, right, top
const halfSpaceEquationTerms = positionNDC.xxyy
.sub(boundaries.xzyw)
.mul(rayDirection.yyxx)
.toConst()
const halfSpaceFlags = halfSpaceEquationTerms.xyyx
.lessThan(halfSpaceEquationTerms.zzww)
.toConst()
// Now compute mask indicating which of four sectors the rayDirection
// belongs to and consequently which border the ray intersects:
// ________________
// |' . '| 0 : hs3 && !hs0
// | ' 3 . ' | 1 : hs0 && !hs1
// | ' . ' | 2 : hs1 && !hs2
// |0 *. 2 | 3 : hs2 && !hs3
// | ' ' . |
// | ' 1 ' . |
// |'____________ '_|
//
// Note that sectorFlags now contains true (1) for the exit boundary and
// false (0) for 3 other.
const sectorFlags = bvecAnd(
halfSpaceFlags.wxyz,
bvecNot(halfSpaceFlags.xyzw)
).toConst()
// Compute distances to boundaries:
const distanceToBoundaries = boundaries
.sub(lightScreenPosition.xyxy)
.div(
rayDirection.xyxy.add(vec4(rayDirection.xyxy.abs().lessThan(1e-6)))
)
.toConst()
// Select distance to the exit boundary:
const distanceToExitBoundary = vec4(sectorFlags)
.dot(distanceToBoundaries)
.toConst()
// Compute exit point on the boundary:
const exitPoint = lightScreenPosition.xy
.add(rayDirection.mul(distanceToExitBoundary))
.toConst()
// Compute epipolar slice for each boundary:
const epipolarSlice = vec4(0, 0.25, 0.5, 0.75)
.add(
exitPoint.yxyx
.sub(boundaries.wxyz)
.mul(vec4(-1, 1, 1, -1))
.div(boundaries.wzwz.sub(boundaries.yxyx))
.saturate()
.div(4)
)
.toConst()
// Select the right value:
const epipolarSliceValue = vec4(sectorFlags).dot(epipolarSlice).toConst()
// Now find two closest epipolar slices, from which we will interpolate.
// First, find index of the slice which precedes our slice.
// Note that 0 <= epipolarSlice <= 1, and both 0 and 1 refer to the first
// slice.
const precedingSliceIndex = min(
epipolarSliceValue.mul(epipolarSliceCount).floor(),
epipolarSliceCount.sub(1)
).toConst()
// Compute EXACT texture coordinates of preceding and succeeding slices
// and their weights.
// Note that slice 0 is stored in the first texel which has exact texture
// coordinate 0.5 / epipolarSliceCount.
const sourceSliceV0 = precedingSliceIndex
.div(epipolarSliceCount)
.add(float(0.5).div(epipolarSliceCount))
.toConst()
const sourceSliceV1 = sourceSliceV0
.add(float(1).div(epipolarSliceCount))
.fract()
.toConst()
const sourceSliceV = [sourceSliceV0, sourceSliceV1]
// Compute slice weights.
const sliceWeight1 = epipolarSliceValue
.mul(epipolarSliceCount)
.sub(precedingSliceIndex)
.toConst()
const sliceWeight0 = sliceWeight1.oneMinus().toConst()
const sliceWeights = [sliceWeight0, sliceWeight1]
const shadowLength = vec2(0).toVar()
const totalWeight = float(0).toVar()
// Unrolled loop for 2 slices:
for (let i = 0; i < 2; ++i) {
// Load epipolar line endpoints.
const sliceEndpoints = sliceEndpointsNode
.sample(vec2(sourceSliceV[i], 0.5))
.toConst()
// Compute line direction on the screen.
const sliceDirection = sliceEndpoints.zw
.sub(sliceEndpoints.xy)
.toConst()
const sliceLengthSquare = sliceDirection.dot(sliceDirection).toConst()
// Project current pixel onto the epipolar line.
const samplePositionOnLine = positionNDC
.sub(sliceEndpoints.xy)
.dot(sliceDirection)
.div(sliceLengthSquare.max(1e-8))
.toConst()
// Compute index of the slice on the line.
// Note that the first sample on the line (samplePositionOnLine==0) is
// exactly the Entry Point, while the last sample
// (samplePositionOnLine==1) is exactly the exit point.
const sampleIndex = samplePositionOnLine
.mul(maxSliceSampleCount.sub(1))
.toConst()
// We have to manually perform bilateral filtering of the texture to
// eliminate artifacts at depth discontinuities.
const precedingSampleIndex = sampleIndex.floor().toConst()
// Get bilinear filtering weight.
const uWeight = sampleIndex.sub(precedingSampleIndex).toConst()
// Get texture coordinate of the left source texel. Again, offset by 0.5
// is essential to align with the texel center.
const precedingSampleU = precedingSampleIndex
.add(0.5)
.div(maxSliceSampleCount)
.toConst()
const shadowLengthUV = vec2(precedingSampleU, sourceSliceV[i]).toConst()
// Gather 4 camera space z values
// Note that we need to bias sourceColorUV by 0.5 texel size to refer
// the location between all four texels and get the required values for
// sure.
// The values in vec4, which Gather() returns are arranged as follows:
// _______ _______
// | | |
// | x | y |
// |_______o_______| o gather location
// | | |
// | *w | z | * sourceColorUV
// |_______|_______|
// |<----->|
// 1/shadowLengthTextureSize.x
const shadowLengthTextureSize = vec2(
maxSliceSampleCount,
epipolarSliceCount
).toConst()
const sourceLocationsCameraZ = textureGather(
coordinateNode,
shadowLengthUV.add(vec2(0.5).div(shadowLengthTextureSize)),
2 // Z component
).wz
// Compute depth weights in a way that if the difference is less than
// the threshold, the weight is 1 and the weights fade out to 0 as the
// difference becomes larger than the threshold:
const maxZ = max(sourceLocationsCameraZ, max(cameraZ, 1)).toConst()
const depthWeights = refinementThreshold
.div(
cameraZ
.sub(sourceLocationsCameraZ)
.abs()
.div(maxZ)
.max(refinementThreshold)
)
.saturate()
.toVar()
// Note that if the sample is located outside the [-1,1] × [-1,1] area,
// the sample is invalid and currentCameraZ == invalidCoordinate.
// Depth weight computed for such sample will be zero.
depthWeights.assign(depthWeights.pow4())
// Multiply bilinear weights with the depth weights:
const bilateralUWeight = vec2(uWeight.oneMinus(), uWeight)
.mul(depthWeights)
.mul(sliceWeights[i])
.toVar()
// If the sample projection is behind [0,1], we have to discard this
// slice.
// We however must take into account the fact that if at least one
// sample from the two bilinear sources is correct, the sample can still
// be properly computed.
//
// -1 0 1 N-2 N-1 N Sample index
// | X | X | X | X | ...... | X | X | X | X |
// 1-1/(N-1) 0 1/(N-1) 1 1+1/(N-1) samplePositionOnLine
// | |
// |<-------------------Clamp range------------------->|
bilateralUWeight.mulAssign(
vec2(
samplePositionOnLine
.sub(0.5)
.abs()
.lessThan(maxSliceSampleCount.sub(1).reciprocal().add(0.5))
)
)
// We now need to compute the following weighted sum:
// We will use hardware to perform bilinear filtering and get this value
// using single bilinear fetch:
const subpixelUOffset = bilateralUWeight.y
.div(bilateralUWeight.x.add(bilateralUWeight.y).max(0.001))
.toVar()
subpixelUOffset.divAssign(shadowLengthTextureSize.x)
const filteredShadowLength = bilateralUWeight.x
.add(bilateralUWeight.y)
.mul(
epipolarShadowLengthNode.sample(
shadowLengthUV.add(vec2(subpixelUOffset, 0))
).xy
)
.toConst()
shadowLength.addAssign(filteredShadowLength)
// Update total weight.
totalWeight.addAssign(bilateralUWeight.dot(vec2(1)))
}
return totalWeight
.greaterThan(1e-6)
.select(shadowLength.div(totalWeight), 0)
})()
}
override setup(builder: NodeBuilder): unknown {
const { material } = this
material.fragmentNode = this.setupFragmentNode(builder)
material.needsUpdate = true
return this.textureNode
}
override dispose(): void {
this.renderTarget.dispose()
this.material.dispose()
this.mesh.geometry.dispose()
super.dispose()
}
}