useShakeController—apply a shake effect to any Object3D

use-shake-controller.ts

import { MutableRefObject, useEffect, useMemo, useRef, useState } from 'react';
import { Clock, Euler, MathUtils, Object3D, Vector3 } from 'three';
import { SimplexNoise } from 'three-stdlib';

import { ShakeController } from '@react-three/drei';
import { useFrame } from '@react-three/fiber';

export interface ShakeControllerParam {
  max: number;
  frequency: number;
}

export interface ShakeControllerVectorConfig {
  x?: ShakeControllerParam;
  y?: ShakeControllerParam;
  z?: ShakeControllerParam;
}

export interface ShakeControllerConfig {
  intensity?: number;
  position?: ShakeControllerVectorConfig;
  rotation?: ShakeControllerVectorConfig;
}

/**
 * Spread a single-dimension shake config over all three dimensions.
 */
export const uniformShakeConfig = ({
  max,
  frequency,
}: ShakeControllerParam): ShakeControllerVectorConfig =>
  ['x', 'y', 'z'].reduce(
    (rot, v) => ({
      ...rot,
      [v]: { max, frequency },
    }),
    {}
  );

/**
 * Attaches a shake effect to the referenced Object3D.
 *
 * The shake alters the position or rotation vectors of the object using
 * simplex noise.
 *
 * This hook returns a controller interface that sets the intensity of the
 * shake.
 */
export function useShakeController(
  target: MutableRefObject<Object3D | null>,
  { intensity = 0, position, rotation }: ShakeControllerConfig
): ShakeController {
  const intensityRef = useRef<number>(intensity);
  const initialPosition = useRef<Vector3>();
  const initialRotation = useRef<Euler>();

  const [xNoise] = useState(() => new SimplexNoise());
  const [yNoise] = useState(() => new SimplexNoise());
  const [zNoise] = useState(() => new SimplexNoise());

  useEffect(() => {
    initialPosition.current = target.current.position.clone();
    initialRotation.current = target.current.rotation.clone();
  }, [target]);

  useFrame(({ clock }) => {
    if (!target.current) return;

    const shake = Math.pow(intensityRef.current, 2);

    if (position) {
      const x = noiseForDimension(xNoise, clock, shake, position.x);
      const y = noiseForDimension(yNoise, clock, shake, position.y);
      const z = noiseForDimension(zNoise, clock, shake, position.z);

      target.current.position.set(
        initialPosition.current.x + x,
        initialPosition.current.y + y,
        initialPosition.current.z + z
      );
    }

    if (rotation) {
      const pitch = noiseForDimension(xNoise, clock, shake, rotation.x);
      const yaw = noiseForDimension(yNoise, clock, shake, rotation.y);
      const roll = noiseForDimension(zNoise, clock, shake, rotation.z);

      target.current.rotation.set(
        initialRotation.current.x + pitch,
        initialRotation.current.y + yaw,
        initialRotation.current.z + roll
      );
    }
  });

  return useMemo(
    () => ({
      getIntensity: (): number => intensityRef.current,
      setIntensity: (val: number): void => {
        intensityRef.current = MathUtils.clamp(val, 0, 1);
      },
    }),
    [intensityRef]
  );
}

function noiseForDimension(
  { noise }: SimplexNoise,
  clock: Clock,
  strength: number,
  { max, frequency }: ShakeControllerParam = { max: 0, frequency: 0 }
) {
  return strength * max * noise(clock.elapsedTime * frequency, 1);
}