Module: Gsplat::Math::Mat

Extended by:
SmallMatrixPrimitives
Defined in:
lib/gsplat/math/mat.rb

Overview

Closed-form batched operations for small matrices.

Class Method Summary collapse

Class Method Details

.det2x2(matrices, dtype: nil) ⇒ Numo::NArray

Determinant of batched 2x2 matrices.

Parameters:

  • matrices (Numo::NArray) —

    [...,2,2]

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

    optional float32/float64 calculation type

Returns:

  • (Numo::NArray) —

    [...]



19
20
21
22
23
 
# File 'lib/gsplat/math/mat.rb', line 19

def det2x2(matrices, dtype: nil)
  _, leading_shape, flat = prepare_matrix(matrices, 2, dtype)
  determinant = (flat[true, 0, 0] * flat[true, 1, 1]) - (flat[true, 0, 1] * flat[true, 1, 0])
  reshape_vector(determinant, leading_shape)
end

.det3x3(matrices, dtype: nil) ⇒ Numo::NArray

Determinant of batched 3x3 matrices.

Parameters:

  • matrices (Numo::NArray) —

    [...,3,3]

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

    optional float32/float64 calculation type

Returns:

  • (Numo::NArray) —

    [...]



66
67
68
69
70
 
# File 'lib/gsplat/math/mat.rb', line 66

def det3x3(matrices, dtype: nil)
  _, leading_shape, flat = prepare_matrix(matrices, 3, dtype)
  determinant = determinant3(flat)
  reshape_vector(determinant, leading_shape)
end

.eigvals2x2(matrices, dtype: nil) ⇒ Numo::NArray

Ordered real eigenvalues of batched 2x2 matrices.

Parameters:

  • matrices (Numo::NArray) —

    [...,2,2]

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

    optional float32/float64 calculation type

Returns:

  • (Numo::NArray) —

    [...,2] in ascending order



48
49
50
51
52
53
54
55
56
57
58
59
 
# File 'lib/gsplat/math/mat.rb', line 48

def eigvals2x2(matrices, dtype: nil)
  typed, leading_shape, flat = prepare_matrix(matrices, 2, dtype)
  half_trace = (flat[true, 0, 0] + flat[true, 1, 1]) * 0.5
  half_difference = (flat[true, 0, 0] - flat[true, 1, 1]) * 0.5
  discriminant = (half_difference**2) + (flat[true, 0, 1] * flat[true, 1, 0])
  discriminant[discriminant.lt(0)] = 0
  root = discriminant**0.5
  output = typed.class.zeros(flat.shape[0], 2)
  output[true, 0] = half_trace - root
  output[true, 1] = half_trace + root
  reshape_matrix(output, leading_shape, 2)
end

.inv2x2(matrices, dtype: nil, epsilon: nil) ⇒ Numo::NArray

Inverse of batched 2x2 matrices.

Parameters:

  • matrices (Numo::NArray) —

    [...,2,2]

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

    optional float32/float64 calculation type

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

    singularity threshold

Returns:

  • (Numo::NArray) —

    [...,2,2]



31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/gsplat/math/mat.rb', line 31

def inv2x2(matrices, dtype: nil, epsilon: nil)
  typed, leading_shape, flat = prepare_matrix(matrices, 2, dtype)
  determinant = (flat[true, 0, 0] * flat[true, 1, 1]) - (flat[true, 0, 1] * flat[true, 1, 0])
  ensure_invertible!(determinant, epsilon || default_epsilon(typed.class))
  output = typed.class.zeros(flat.shape[0], 2, 2)
  output[true, 0, 0] = flat[true, 1, 1] / determinant
  output[true, 0, 1] = -flat[true, 0, 1] / determinant
  output[true, 1, 0] = -flat[true, 1, 0] / determinant
  output[true, 1, 1] = flat[true, 0, 0] / determinant
  reshape_matrix(output, leading_shape, 2, 2)
end

.inv3x3(matrices, dtype: nil, epsilon: nil) ⇒ Numo::NArray

Inverse of batched 3x3 matrices.

Parameters:

  • matrices (Numo::NArray) —

    [...,3,3]

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

    optional float32/float64 calculation type

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

    singularity threshold

Returns:

  • (Numo::NArray) —

    [...,3,3]



78
79
80
81
82
83
84
85
 
# File 'lib/gsplat/math/mat.rb', line 78

def inv3x3(matrices, dtype: nil, epsilon: nil)
  typed, leading_shape, flat = prepare_matrix(matrices, 3, dtype)
  determinant = determinant3(flat)
  ensure_invertible!(determinant, epsilon || default_epsilon(typed.class))
  output = typed.class.zeros(flat.shape[0], 3, 3)
  fill_inverse3!(output, flat, determinant)
  reshape_matrix(output, leading_shape, 3, 3)
end

.matmul_batch(left, right, dtype: nil) ⇒ Numo::NArray

Batched matrix multiplication for small matrices.

Parameters:

  • left (Numo::NArray) —

    [...,M,K]

  • right (Numo::NArray) —

    [...,K,N]

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

    optional float32/float64 calculation type

Returns:

  • (Numo::NArray) —

    [...,M,N]



93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
 
# File 'lib/gsplat/math/mat.rb', line 93

def matmul_batch(left, right, dtype: nil)
  left = cast_float(left, dtype)
  right = cast_float(right, left.class)
  validate_rank!(left, 2)
  validate_rank!(right, 2)
  leading_shape = left.shape[0...-2]
  unless right.shape[0...-2] == leading_shape && left.shape[-1] == right.shape[-2]
    raise ShapeError, "matmul shape mismatch: left #{left.shape.inspect}, right #{right.shape.inspect}"
  end

  rows = left.shape[-2]
  shared = left.shape[-1]
  columns = right.shape[-1]
  batch_size = leading_shape.empty? ? 1 : leading_shape.inject(:*)
  left_flat = left.reshape(batch_size, rows, shared)
  right_flat = right.reshape(batch_size, shared, columns)
  output = multiply_flat(left_flat, right_flat, rows, shared, columns)
  reshape_matrix(output, leading_shape, rows, columns)
end