globe/src/atmosphere.ts

import {
  BackSide,
  DirectionalLight,
  Mesh,
  Object3D,
  ShaderMaterial,
  SphereGeometry,
  Vector3,
} from 'three';

import { plainText as vertexShader } from './shaders/atmosphere.vert';
import { plainText as fragmentShader } from './shaders/atmosphere.frag';

export class Atmosphere extends Object3D {
  public Rayleigh = 0.0025;
  public Mie = 0.0005;
  public Exposure = 15.0;
  public Scale = 1;
  public ScaleDepth = 0.25;
  public G = -0.95;
  public Wavelength = new Vector3(0.65, 0.57, 0.475);

  public geom?: SphereGeometry;
  public mesh?: Mesh;
  public shader = new ShaderMaterial({
    vertexShader,
    fragmentShader,
    side: BackSide,
    transparent: true,
  });

  constructor(public innerRadius: number, public outerRadius: number) {
    super();
    this.Scale = 1 / (this.outerRadius - this.innerRadius);
    this.setUniforms(this.shader);
    this.initialize();
  }

  set planetRadius(radius: number) {
    const thickness = this.thickness;
    this.innerRadius = radius;
    this.outerRadius = radius + thickness;
    this.refreshScale();
  }
  get planetRadius() {
    return this.innerRadius;
  }

  set thickness(thickness: number) {
    this.outerRadius = thickness + this.innerRadius;
    this.refreshScale();
  }
  get thickness() {
    return this.outerRadius - this.innerRadius;
  }

  setUniforms(shader: ShaderMaterial) {
    shader.uniforms.invWavelength = {
      value: [
        1 / Math.pow(this.Wavelength.x, 4),
        1 / Math.pow(this.Wavelength.y, 4),
        1 / Math.pow(this.Wavelength.z, 4),
      ],
    };
    shader.uniforms.outerRadius = { value: this.outerRadius };
    shader.uniforms.innerRadius = { value: this.innerRadius };
    shader.uniforms.Kr = { value: this.Rayleigh };
    shader.uniforms.Km = { value: this.Mie };
    shader.uniforms.ESun = { value: this.Exposure };
    shader.uniforms.scale = { value: this.Scale };
    shader.uniforms.scaleDepth = { value: this.ScaleDepth };
    shader.uniforms.g = { value: this.G };
  }

  setLight(light: DirectionalLight) {
    this.shader.uniforms.lightDirection = {
      value: light.position.clone().normalize(),
    };
  }

  initialize() {
    if (this.geom) {
      this.geom.dispose();
      this.mesh && this.remove(this.mesh);
    }
    this.geom = new SphereGeometry(this.outerRadius, 128, 128);
    this.mesh = new Mesh(this.geom, this.shader);
    this.add(this.mesh);
  }

  private refreshScale() {
    this.Scale = 1 / (this.outerRadius - this.innerRadius);
    this.setUniforms(this.shader);
    this.initialize();
  }
}