Module: Gsplat::Math::Quaternion

Defined in:
lib/gsplat/math/quaternion.rb

Overview

Batched wxyz quaternion normalization, rotation conversion, and VJPs.

Class Method Summary collapse

Class Method Details

.normalize(quaternions, dtype: nil, epsilon: nil) ⇒ Numo::NArray

Normalizes wxyz quaternions along the last dimension.

Parameters:

  • quaternions (Numo::NArray) —

    [...,4]

  • dtype (Class, Symbol, nil) (defaults to: nil) —

    optional float32/float64 calculation type

  • epsilon (Numeric, nil) (defaults to: nil) —

    minimum norm

Returns:

  • (Numo::NArray) —

    [...,4]



15
16
17
18
19
 
# File 'lib/gsplat/math/quaternion.rb', line 15

def normalize(quaternions, dtype: nil, epsilon: nil)
  typed, leading_shape, flat = prepare(quaternions, 4, dtype)
  denominator = norm_denominator(flat, epsilon || default_epsilon(typed.class))
  (flat / denominator.reshape(flat.shape[0], 1)).reshape(*(leading_shape + [4]))
end

.normalize_vjp(quaternions, grad_output, dtype: nil, epsilon: nil) ⇒ Numo::NArray

VJP of quaternion normalization.

Parameters:

  • quaternions (Numo::NArray) —

    [...,4]

  • grad_output (Numo::NArray) —

    [...,4]

  • dtype (Class, Symbol, nil) (defaults to: nil) —

    optional float32/float64 calculation type

  • epsilon (Numeric, nil) (defaults to: nil) —

    minimum norm

Returns:

  • (Numo::NArray) —

    [...,4]



28
29
30
31
32
33
34
35
36
 
# File 'lib/gsplat/math/quaternion.rb', line 28

def normalize_vjp(quaternions, grad_output, dtype: nil, epsilon: nil)
  typed, leading_shape, flat = prepare(quaternions, 4, dtype)
  gradient, = prepare_matching(grad_output, leading_shape, 4, typed.class)
  denominator = norm_denominator(flat, epsilon || default_epsilon(typed.class))
  normalized = flat / denominator.reshape(flat.shape[0], 1)
  dot = (normalized * gradient).sum(axis: 1)
  result = (gradient - (normalized * dot.reshape(flat.shape[0], 1))) / denominator.reshape(flat.shape[0], 1)
  result.reshape(*(leading_shape + [4]))
end

.to_rotmat(quaternions, dtype: nil) ⇒ Numo::NArray

Converts wxyz quaternions to rotation matrices.

Parameters:

  • quaternions (Numo::NArray) —

    [...,4]

  • dtype (Class, Symbol, nil) (defaults to: nil) —

    optional float32/float64 calculation type

Returns:

  • (Numo::NArray) —

    [...,3,3]



43
44
45
46
47
48
49
50
51
52
53
 
# File 'lib/gsplat/math/quaternion.rb', line 43

def to_rotmat(quaternions, dtype: nil)
  typed, leading_shape, = prepare(quaternions, 4, dtype)
  flat = normalize(typed).reshape(typed.size / 4, 4)
  quat_w = flat[true, 0]
  quat_x = flat[true, 1]
  quat_y = flat[true, 2]
  quat_z = flat[true, 3]
  output = typed.class.zeros(flat.shape[0], 3, 3)
  fill_rotation!(output, quat_w, quat_x, quat_y, quat_z)
  output.reshape(*(leading_shape + [3, 3]))
end

.to_rotmat_vjp(quaternions, grad_output, dtype: nil) ⇒ Numo::NArray

VJP of quaternion-to-rotation-matrix conversion.

Parameters:

  • quaternions (Numo::NArray) —

    [...,4]

  • grad_output (Numo::NArray) —

    [...,3,3]

  • dtype (Class, Symbol, nil) (defaults to: nil) —

    optional float32/float64 calculation type

Returns:

  • (Numo::NArray) —

    [...,4]



61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/gsplat/math/quaternion.rb', line 61

def to_rotmat_vjp(quaternions, grad_output, dtype: nil)
  typed, leading_shape, = prepare(quaternions, 4, dtype)
  gradient = cast_float(grad_output, typed.class)
  unless gradient.shape == leading_shape + [3, 3]
    raise ShapeError, "expected gradient #{(leading_shape + [3, 3]).inspect}, got #{gradient.shape.inspect}"
  end

  normalized = normalize(typed).reshape(typed.size / 4, 4)
  grad_matrix = gradient.reshape(normalized.shape[0], 3, 3)
  grad_normalized = rotation_vjp(normalized, grad_matrix)
  normalize_vjp(typed, grad_normalized.reshape(*(leading_shape + [4])))
end