Module: Gsplat::Strategy::Ops

Defined in:
lib/gsplat/strategy/ops.rb,
lib/gsplat/strategy/mcmc_ops.rb

Overview

MCMC-specific parameter edits synchronized with Adam state.

Class Method Summary collapse

Class Method Details

.duplicate!(params, optimizers, mask) ⇒ Integer

Appends exact copies of selected Gaussian rows and zeroed moments.

Returns:

  • (Integer) —

    number of rows appended



13
14
15
16
17
18
19
20
21
22
23
24
 
# File 'lib/gsplat/strategy/ops.rb', line 13

def duplicate!(params, optimizers, mask)
  indices = selected_indices(mask, parameter_count(params))
  return 0 if indices.empty?

  prepare_optimizers!(optimizers)
  params.each do |key, variable|
    appended = select_rows(variable.data, indices)
    variable.replace_data!(append_rows(variable.data, appended))
    optimizers.fetch(key).append!(indices.size)
  end
  indices.size
end

.inject_position_noise!(params, scaler:, rng: Gsplat.rng) ⇒ Object

rubocop:disable Metrics/AbcSize



54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/gsplat/strategy/mcmc_ops.rb', line 54

def inject_position_noise!(params, scaler:, rng: Gsplat.rng)
  return params if scaler.zero?

  opacities = sigmoid_mcmc(params.fetch(:opacities).data)
  scales = Numo::NMath.exp(params.fetch(:scales).data)
  covariances, = Gsplat.quat_scale_to_covar_preci(
    params.fetch(:quats).data,
    scales,
    compute_preci: false
  )
  noise = params.fetch(:means).data.class.zeros(*params.fetch(:means).data.shape)
  noise.shape[0].times do |index|
    gate = 1.0 / (1.0 + ::Math.exp(-100.0 * ((1.0 - opacities[index].to_f) - 0.995)))
    standard = noise.class.cast(Array.new(3) { normal_mcmc_sample(rng) })
    noise[index, true] = covariances[index, true, true].dot(standard) * gate * scaler
  end
  params.fetch(:means).data[] = params.fetch(:means).data + noise
  params
end

.relocate!(params, optimizers, mask, binoms:, **options) ⇒ Object

rubocop:disable Metrics/AbcSize

Raises:



10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
 
# File 'lib/gsplat/strategy/mcmc_ops.rb', line 10

def relocate!(params, optimizers, mask, binoms:, **options)
  min_opacity = options.fetch(:min_opacity, 0.005)
  rng = options.fetch(:rng, Gsplat.rng)
  count = params.fetch(:means).data.shape[0]
  validate_mcmc_mask!(mask, count)
  dead = Numo::Bit.cast(mask).where.to_a
  return 0 if dead.empty?

  alive = Numo::Bit.cast(mask).eq(0).where.to_a
  raise Gsplat::Error, "cannot relocate when every Gaussian is dead" if alive.empty?

  weights = sigmoid_mcmc(params.fetch(:opacities).data[alive])
  sampled = weighted_indices(weights, dead.size, rng).map { |index| alive.fetch(index) }
  update_relocation_sources!(params, sampled, binoms, min_opacity)
  prepare_mcmc_optimizers!(optimizers)
  params.each do |key, variable|
    variable.data[*([dead] + Array.new(variable.data.ndim - 1, true))] =
      variable.data[*([sampled] + Array.new(variable.data.ndim - 1, true))]
    optimizers.fetch(key).zero_state_at!(sampled)
    optimizers.fetch(key).zero_state_at!(dead)
  end
  dead.size
end

.remove!(params, optimizers, mask) ⇒ Integer

Removes selected Gaussian rows and matching optimizer moments.

Returns:

  • (Integer) —

    number of rows removed



29
30
31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/gsplat/strategy/ops.rb', line 29

def remove!(params, optimizers, mask)
  validate_mask!(mask, parameter_count(params))
  keep = Numo::Bit.cast(mask).eq(0)
  removed = mask.count_true
  return 0 if removed.zero?

  prepare_optimizers!(optimizers)
  params.each do |key, variable|
    variable.replace_data!(select_rows(variable.data, keep))
    optimizers.fetch(key).select!(keep)
  end
  removed
end

.reset_opacity!(params, optimizers, maximum:) ⇒ Object

Raises:

  • (ArgumentError) 


64
65
66
67
68
69
70
71
72
73
74
75
76
77
 
# File 'lib/gsplat/strategy/ops.rb', line 64

def reset_opacity!(params, optimizers, maximum:)
  raise ArgumentError, "maximum opacity must be in (0,1)" unless maximum.positive? && maximum < 1

  variable = params.fetch(:opacities)
  cap = ::Math.log(maximum / (1 - maximum))
  values = variable.data.dup
  above = values.gt(cap)
  values[above] = cap if above.any?
  variable.replace_data!(values)
  optimizer = optimizers.fetch(:opacities)
  optimizer.state
  optimizer.zero_state_at!(0...values.shape[0])
  above.count_true
end

.sample_add!(params, optimizers, count, binoms:, **options) ⇒ Object

rubocop:enable Metrics/AbcSize

Raises:

  • (ArgumentError) 


35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/gsplat/strategy/mcmc_ops.rb', line 35

def sample_add!(params, optimizers, count, binoms:, **options)
  min_opacity = options.fetch(:min_opacity, 0.005)
  rng = options.fetch(:rng, Gsplat.rng)
  raise ArgumentError, "count must be a non-negative integer" unless count.is_a?(Integer) && !count.negative?
  return 0 if count.zero?

  weights = sigmoid_mcmc(params.fetch(:opacities).data)
  sampled = weighted_indices(weights, count, rng)
  update_relocation_sources!(params, sampled, binoms, min_opacity)
  prepare_mcmc_optimizers!(optimizers)
  params.each do |key, variable|
    rows = variable.data[*([sampled] + Array.new(variable.data.ndim - 1, true))].dup
    variable.replace_data!(append_mcmc_rows(variable.data, rows))
    optimizers.fetch(key).append!(count)
  end
  count
end

.split!(params, optimizers, mask, rng: Gsplat.rng) ⇒ Integer

Replaces selected Gaussians with two sampled children.

Returns:

  • (Integer) —

    number of parents split



46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
 
# File 'lib/gsplat/strategy/ops.rb', line 46

def split!(params, optimizers, mask, rng: Gsplat.rng)
  count = parameter_count(params)
  indices = selected_indices(mask, count)
  return 0 if indices.empty?

  children = split_children(params, indices, rng)
  keep = Numo::Bit.cast(mask).eq(0)
  prepare_optimizers!(optimizers)
  params.each do |key, variable|
    retained = select_rows(variable.data, keep)
    variable.replace_data!(append_rows(retained, children.fetch(key)))
    optimizer = optimizers.fetch(key)
    optimizer.select!(keep)
    optimizer.append!(indices.size * 2)
  end
  indices.size
end