gsplat 1.0.0

data/lib/gsplat/ops/spherical_harmonics.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require_relative "../backend/ruby/spherical_harmonics"
+
+# Differentiable real spherical harmonics operation.
+module Gsplat
+  module Ops
+    # Evaluates real SH bases through degree four.
+    class SphericalHarmonics < Autograd::Function
+      class << self
+        # Evaluates SH and stores inputs for its analytic VJP.
+        # @api private
+        def forward(context, degree, directions, coefficients, masks: nil)
+          context.save(degree, directions, coefficients, masks)
+          Backend.dispatch(
+            :spherical_harmonics_forward,
+            degree,
+            directions,
+            coefficients,
+            masks: masks
+          )
+        end
+
+        # Propagates an SH value gradient to directions and coefficients.
+        # @api private
+        def backward(context, grad_output)
+          degree, directions, coefficients, masks = context.saved_values
+          grad_directions, grad_coefficients = Backend.dispatch(
+            :spherical_harmonics_backward,
+            degree,
+            directions,
+            coefficients,
+            grad_output,
+            masks: masks,
+            grad_dirs: context.needs_input_grad.fetch(1),
+            grad_coeffs: context.needs_input_grad.fetch(2)
+          )
+          [nil, grad_directions, grad_coefficients]
+        end
+      end
+    end
+  end
+
+  class << self
+    # Evaluates real spherical harmonics.
+    #
+    # @param degree [Integer] degree 0..4
+    # @param directions [Autograd::Variable, Numo::NArray] [...,3]
+    # @param coefficients [Autograd::Variable, Numo::NArray] [...,K,D], K >= (degree+1)^2
+    # @param masks [Numo::Bit, nil] [...]
+    # @return [Autograd::Variable, Numo::NArray] [...,D]
+    def spherical_harmonics(degree, directions, coefficients, masks: nil)
+      if [directions, coefficients].any?(Autograd::Variable)
+        return Ops::SphericalHarmonics.apply(degree, directions, coefficients, masks: masks)
+      end
+
+      Backend.dispatch(:spherical_harmonics_forward, degree, directions, coefficients, masks: masks)
+    end
+  end
+
+  Backend.register(
+    :spherical_harmonics_forward,
+    :ruby,
+    Backend::RubySphericalHarmonics.method(:forward)
+  )
+  Backend.register(
+    :spherical_harmonics_backward,
+    :ruby,
+    Backend::RubySphericalHarmonics.method(:backward)
+  )
+end

data/lib/gsplat/ops/tensor_shape_ops.rb

@@ -0,0 +1,173 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Ops
+    # Internal helpers for composing coarse differentiable operations.
+    module TensorOps
+      module_function
+
+      def apply(function, *inputs, **)
+        return function.apply(*inputs, **) if inputs.any?(Autograd::Variable)
+
+        needs_grad = Array.new(inputs.length, false)
+        function.forward(Autograd::Context.new(needs_grad), *inputs, **)
+      end
+
+      def data(value)
+        value.is_a?(Autograd::Variable) ? value.data : value
+      end
+    end
+
+    # Adds a leading camera dimension to a tensor.
+    class CameraBroadcast < Autograd::Function
+      class << self
+        def forward(context, value, camera_count)
+          valid_count = camera_count.is_a?(Integer) && camera_count.positive?
+          raise ArgumentError, "camera_count must be positive" unless valid_count
+
+          context.save(value.shape)
+          output = value.class.zeros(*([camera_count] + value.shape))
+          camera_count.times do |camera_index|
+            output[*([camera_index] + Array.new(value.ndim, true))] = value
+          end
+          output
+        end
+
+        def backward(_context, gradient)
+          [gradient.sum(axis: 0), nil]
+        end
+      end
+    end
+
+    # Computes camera-to-Gaussian viewing directions from world-to-camera matrices.
+    class CameraDirections < Autograd::Function
+      class << self
+        def forward(context, means, viewmats)
+          unless means.ndim == 2 && means.shape[-1] == 3 &&
+                 viewmats.ndim == 3 && viewmats.shape[1..] == [4, 4]
+            raise ShapeError,
+                  "expected means [N,3] and viewmats [C,4,4], " \
+                  "got #{means.shape.inspect} and #{viewmats.shape.inspect}"
+          end
+
+          output = means.class.zeros(viewmats.shape[0], means.shape[0], 3)
+          viewmats.shape[0].times do |camera_index|
+            rotation = means.class.cast(viewmats[camera_index, 0...3, 0...3])
+            translation = means.class.cast(viewmats[camera_index, 0...3, 3])
+            camera_center = -translation.dot(rotation)
+            output[camera_index, true, true] = means - camera_center
+          end
+          context.save(means.shape)
+          output
+        end
+
+        def backward(_context, gradient)
+          [gradient.sum(axis: 0), nil]
+        end
+      end
+    end
+
+    # Adds a singleton feature axis to depths.
+    class DepthFeatures < Autograd::Function
+      class << self
+        def forward(_context, depths)
+          depths.reshape(*(depths.shape + [1]))
+        end
+
+        def backward(_context, gradient)
+          [gradient.reshape(*gradient.shape[0...-1])]
+        end
+      end
+    end
+
+    # Concatenates per-Gaussian features and depth.
+    class ConcatDepth < Autograd::Function
+      class << self
+        def forward(context, features, depth_features)
+          unless features.shape[0...-1] == depth_features.shape[0...-1] && depth_features.shape[-1] == 1
+            raise ShapeError,
+                  "expected features [...,D] and depths [...,1], " \
+                  "got #{features.shape.inspect} and #{depth_features.shape.inspect}"
+          end
+
+          feature_count = features.shape[-1]
+          output = features.class.zeros(*(features.shape[0...-1] + [feature_count + 1]))
+          output[*Array.new(features.ndim - 1, true), 0...feature_count] = features
+          output[*Array.new(features.ndim - 1, true), feature_count] = depth_features[
+            *Array.new(depth_features.ndim - 1, true), 0
+          ]
+          context.save(feature_count)
+          output
+        end
+
+        def backward(context, gradient)
+          feature_count = context.saved_values.fetch(0)
+          leading = Array.new(gradient.ndim - 1, true)
+          [
+            gradient[*leading, 0...feature_count].dup,
+            gradient[*leading, feature_count].reshape(*(gradient.shape[0...-1] + [1]))
+          ]
+        end
+      end
+    end
+
+    # Extracts a contiguous interval from the last tensor axis.
+    class FeatureSlice < Autograd::Function
+      class << self
+        def forward(context, value, range)
+          unless range.is_a?(Range) && range.begin.is_a?(Integer) && range.end.is_a?(Integer)
+            raise ArgumentError, "feature range must have integer bounds"
+          end
+
+          stop = range.exclude_end? ? range.end : range.end + 1
+          unless range.begin.between?(0, value.shape[-1]) && stop.between?(range.begin, value.shape[-1])
+            raise ShapeError, "feature range #{range.inspect} is outside #{value.shape.inspect}"
+          end
+
+          context.save(value.shape, range.begin, stop)
+          value[*Array.new(value.ndim - 1, true), range.begin...stop].dup
+        end
+
+        def backward(context, gradient)
+          shape, start, stop = context.saved_values
+          output = gradient.class.zeros(*shape)
+          output[*Array.new(output.ndim - 1, true), start...stop] = gradient
+          [output, nil]
+        end
+      end
+    end
+
+    # Concatenates any number of tensors along their final feature axis.
+    class ConcatFeatures < Autograd::Function
+      class << self
+        def forward(context, *values)
+          raise ArgumentError, "at least one feature tensor is required" if values.empty?
+
+          leading_shape = values.first.shape[0...-1]
+          unless values.all? { |value| value.shape[0...-1] == leading_shape }
+            raise ShapeError, "feature tensors must share leading dimensions"
+          end
+
+          sizes = values.map { |value| value.shape[-1] }
+          output = values.first.class.zeros(*(leading_shape + [sizes.sum]))
+          cursor = 0
+          values.each_with_index do |value, index|
+            output[*Array.new(output.ndim - 1, true), cursor...(cursor + sizes[index])] = value
+            cursor += sizes[index]
+          end
+          context.save(sizes)
+          output
+        end
+
+        def backward(context, gradient)
+          cursor = 0
+          context.saved_values.first.map do |size|
+            value = gradient[*Array.new(gradient.ndim - 1, true), cursor...(cursor + size)].dup
+            cursor += size
+            value
+          end
+        end
+      end
+    end
+  end
+end

data/lib/gsplat/ops/tensor_value_ops.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Ops
+    # Adds an offset and applies a lower clamp.
+    class AddClampMin < Autograd::Function
+      class << self
+        def forward(context, value, offset:, minimum:)
+          shifted = value + offset
+          context.save(shifted.gt(minimum))
+          below = shifted.lt(minimum)
+          shifted[below] = minimum if below.any?
+          shifted
+        end
+
+        # rubocop:disable Metrics/AbcSize
+        def backward(context, gradient)
+          mask = gradient.class.cast(context.saved_values.fetch(0))
+          [gradient * mask]
+        end
+      end
+    end
+
+    # Elementwise multiplication with VJPs for both inputs.
+    class Multiply < Autograd::Function
+      class << self
+        def forward(context, left, right)
+          unless left.shape == right.shape
+            raise ShapeError, "multiply shape mismatch: #{left.shape.inspect} and #{right.shape.inspect}"
+          end
+
+          context.save(left, right)
+          left * right
+        end
+
+        def backward(context, gradient)
+          left, right = context.saved_values
+          [gradient * right, gradient * left]
+        end
+      end
+    end
+
+    # Elementwise exponential activation.
+    class Exp < Autograd::Function
+      class << self
+        def forward(context, value)
+          output = Numo::NMath.exp(value)
+          context.save(output)
+          output
+        end
+
+        def backward(context, gradient)
+          [gradient * context.saved_values.first]
+        end
+      end
+    end
+
+    # Elementwise logistic activation.
+    class Sigmoid < Autograd::Function
+      class << self
+        def forward(context, value)
+          output = 1.0 / (1.0 + Numo::NMath.exp(-value))
+          context.save(output)
+          output
+        end
+
+        def backward(context, gradient)
+          output = context.saved_values.first
+          [gradient * output * (1.0 - output)]
+        end
+      end
+    end
+
+    # Last-axis quaternion normalization.
+    class NormalizeQuaternion < Autograd::Function
+      class << self
+        def forward(context, value)
+          context.save(value)
+          Math::Quaternion.normalize(value)
+        end
+
+        def backward(context, gradient)
+          [Math::Quaternion.normalize_vjp(context.saved_values.first, gradient)]
+        end
+      end
+    end
+
+    # Joins the direct-current and remaining SH coefficient blocks.
+    class ConcatCoefficients < Autograd::Function
+      class << self
+        def forward(context, direct, rest)
+          unless direct.ndim == 3 && rest.ndim == 3 &&
+                 direct.shape[0] == rest.shape[0] && direct.shape[2] == rest.shape[2]
+            raise ShapeError, "cannot concatenate SH shapes #{direct.shape.inspect} and #{rest.shape.inspect}"
+          end
+
+          direct_count = direct.shape[1]
+          output = direct.class.zeros(direct.shape[0], direct_count + rest.shape[1], direct.shape[2])
+          output[true, 0...direct_count, true] = direct
+          output[true, direct_count...output.shape[1], true] = rest
+          context.save(direct_count)
+          output
+        end
+
+        def backward(context, gradient)
+          direct_count = context.saved_values.first
+          [
+            gradient[true, 0...direct_count, true].dup,
+            gradient[true, direct_count...gradient.shape[1], true].dup
+          ]
+        end
+      end
+    end
+
+    # Divides one selected rendered depth channel by accumulated alpha.
+    class NormalizeDepth < Autograd::Function
+      class << self
+        def forward(context, rendered, alphas, depth_index:)
+          expected = rendered.shape[0...-1] + [1]
+          unless alphas.shape == expected
+            raise ShapeError, "expected alphas #{expected.inspect}, got #{alphas.shape.inspect}"
+          end
+
+          alpha = alphas[*Array.new(alphas.ndim - 1, true).push(0)]
+          depth = rendered[*Array.new(rendered.ndim - 1, true).push(depth_index)]
+          valid = alpha.gt(0)
+          output = rendered.dup
+          normalized = rendered.class.zeros(*alpha.shape)
+          normalized[valid] = depth[valid] / alpha[valid] if valid.any?
+          output[*Array.new(output.ndim - 1, true), depth_index] = normalized
+          context.save(alpha, depth, valid, depth_index)
+          output
+        end
+
+        def backward(context, gradient)
+          alpha, depth, valid, depth_index = context.saved_values
+          grad_rendered = gradient.dup
+          grad_depth = gradient[*Array.new(gradient.ndim - 1, true).push(depth_index)]
+          normalized_grad = gradient.class.zeros(*alpha.shape)
+          normalized_grad[valid] = grad_depth[valid] / alpha[valid] if valid.any?
+          grad_rendered[*Array.new(gradient.ndim - 1, true), depth_index] = normalized_grad
+          grad_alpha = gradient.class.zeros(*alpha.shape)
+          grad_alpha[valid] = -grad_depth[valid] * depth[valid] / (alpha[valid]**2) if valid.any?
+          [grad_rendered, grad_alpha.reshape(*(alpha.shape + [1]))]
+        end
+        # rubocop:enable Metrics/AbcSize
+      end
+    end
+  end
+end

data/lib/gsplat/optim/adam.rb

@@ -0,0 +1,212 @@
+# frozen_string_literal: true
+
+module Gsplat
+  # Optimizers and learning-rate schedules for autograd Variables.
+  module Optim
+    # Dense Adam optimizer with editable first-axis state for densification.
+    class Adam
+      # Immutable parameter-group configuration.
+      Group = Data.define(:name, :variable, :lr, :eps)
+      # Adam step count and first/second moments.
+      State = Data.define(:step, :exp_avg, :exp_avg_sq)
+
+      attr_reader :beta1, :beta2, :groups
+
+      # rubocop:disable Naming/MethodParameterName
+      def initialize(groups = nil, lr: 1e-3, betas: [0.9, 0.999], eps: 1e-15, **named_groups)
+        # rubocop:enable Naming/MethodParameterName
+        raise ArgumentError, "provide positional or named groups, not both" if groups && !named_groups.empty?
+
+        groups ||= named_groups
+        @beta1, @beta2 = betas
+        validate_hyperparameters!(lr, eps)
+        @groups = normalize_groups(groups, lr, eps)
+        @states = {}
+      end
+
+      # rubocop:disable Metrics/AbcSize
+      def step
+        groups.each_value do |group|
+          gradient = group.variable.grad
+          next unless gradient
+
+          state = state_for(group.name)
+          next_step = state.step + 1
+          exp_avg = (beta1 * state.exp_avg) + ((1 - beta1) * gradient)
+          exp_avg_sq = (beta2 * state.exp_avg_sq) + ((1 - beta2) * (gradient**2))
+          corrected_mean = exp_avg / (1 - (beta1**next_step))
+          corrected_variance = exp_avg_sq / (1 - (beta2**next_step))
+          update = group.lr * corrected_mean / ((corrected_variance**0.5) + group.eps)
+          group.variable.data[] = group.variable.data - update
+          @states[group.name] = State.new(
+            step: next_step,
+            exp_avg: exp_avg,
+            exp_avg_sq: exp_avg_sq
+          )
+        end
+        self
+      end
+      # rubocop:enable Metrics/AbcSize
+
+      def zero_grad!
+        groups.each_value { |group| group.variable.zero_grad! }
+        self
+      end
+
+      # Returns lazily initialized moment state for one or all groups.
+      #
+      # @return [State, Hash{Symbol=>State}]
+      def state(name = nil)
+        return state_for(name.to_sym) if name
+        return state_for(groups.keys.first) if groups.length == 1
+
+        groups.keys.to_h { |key| [key, state_for(key)] }
+      end
+
+      # Returns the current learning rate for one or all groups.
+      #
+      # @return [Numeric, Hash{Symbol=>Numeric}]
+      def learning_rate(name = nil)
+        return groups.fetch(name.to_sym).lr if name
+        return groups.values.first.lr if groups.length == 1
+
+        groups.transform_values(&:lr)
+      end
+
+      def learning_rate=(value)
+        raise ArgumentError, "learning rate must be positive" unless value.positive?
+
+        @groups = groups.transform_values do |group|
+          Group.new(name: group.name, variable: group.variable, lr: value, eps: group.eps)
+        end
+      end
+
+      def set_learning_rate(name, value)
+        raise ArgumentError, "learning rate must be positive" unless value.positive?
+
+        key = name.to_sym
+        group = groups.fetch(key)
+        @groups[key] = Group.new(name: key, variable: group.variable, lr: value, eps: group.eps)
+      end
+
+      # Selects first-axis optimizer rows after pruning parameters.
+      #
+      # @param selection [Numo::Bit, Array<Integer>, Range]
+      # @return [self]
+      def select!(selection)
+        groups.each_key do |name|
+          current = state_for(name)
+          @states[name] = State.new(
+            step: current.step,
+            exp_avg: select_first_axis(current.exp_avg, selection),
+            exp_avg_sq: select_first_axis(current.exp_avg_sq, selection)
+          )
+        end
+        self
+      end
+
+      def append!(count)
+        raise ArgumentError, "append count must be non-negative" unless count.is_a?(Integer) && !count.negative?
+        return self if count.zero?
+
+        groups.each_key do |name|
+          current = state_for(name)
+          @states[name] = State.new(
+            step: current.step,
+            exp_avg: append_zeros(current.exp_avg, count),
+            exp_avg_sq: append_zeros(current.exp_avg_sq, count)
+          )
+        end
+        self
+      end
+
+      # Clears moments at selected first-axis rows.
+      #
+      # @param indices [Integer, Array<Integer>, Range]
+      # @return [self]
+      def zero_state_at!(indices)
+        groups.each_key do |name|
+          current = state_for(name)
+          index = [indices] + Array.new(current.exp_avg.ndim - 1, true)
+          current.exp_avg[*index] = 0
+          current.exp_avg_sq[*index] = 0
+        end
+        self
+      end
+
+      # Restores one parameter group's moments from a checkpoint.
+      def load_state!(name = nil, step:, exp_avg:, exp_avg_sq:)
+        key = name ? name.to_sym : single_group_name
+        group = groups.fetch(key)
+        expected = group.variable.data.shape
+        unless exp_avg.shape == expected && exp_avg_sq.shape == expected
+          raise ShapeError,
+                "Adam state for #{key} expected #{expected.inspect}, " \
+                "got #{exp_avg.shape.inspect} and #{exp_avg_sq.shape.inspect}"
+        end
+        raise ArgumentError, "Adam step must be a non-negative integer" unless step.is_a?(Integer) && !step.negative?
+
+        type = group.variable.data.class
+        @states[key] = State.new(
+          step: step,
+          exp_avg: type.cast(exp_avg).dup,
+          exp_avg_sq: type.cast(exp_avg_sq).dup
+        )
+        self
+      end
+
+      private
+
+      def single_group_name
+        return groups.keys.first if groups.length == 1
+
+        raise ArgumentError, "group name is required for a multi-group Adam optimizer"
+      end
+
+      def validate_hyperparameters!(learning_rate, epsilon)
+        valid_betas = [beta1, beta2].all? { |value| value >= 0 && value < 1 }
+        raise ArgumentError, "betas must be in [0,1)" unless valid_betas
+        raise ArgumentError, "learning rate must be positive" unless learning_rate.positive?
+        raise ArgumentError, "eps must be positive" unless epsilon.positive?
+      end
+
+      def normalize_groups(input, default_lr, default_eps)
+        raw_groups = input.is_a?(Autograd::Variable) ? { default: { variable: input } } : input
+        unless raw_groups.is_a?(Hash) && !raw_groups.empty?
+          raise ArgumentError, "groups must be a Variable or non-empty Hash"
+        end
+
+        raw_groups.to_h do |name, specification|
+          specification = { variable: specification } if specification.is_a?(Autograd::Variable)
+          variable = specification.fetch(:variable)
+          raise ArgumentError, "group variable must be an Autograd::Variable" unless variable.is_a?(Autograd::Variable)
+
+          group_lr = specification.fetch(:lr, default_lr)
+          group_eps = specification.fetch(:eps, default_eps)
+          validate_hyperparameters!(group_lr, group_eps)
+          key = name.to_sym
+          [key, Group.new(name: key, variable: variable, lr: group_lr, eps: group_eps)]
+        end
+      end
+
+      def state_for(name)
+        group = groups.fetch(name)
+        @states[name] ||= State.new(
+          step: 0,
+          exp_avg: group.variable.data.class.zeros(*group.variable.data.shape),
+          exp_avg_sq: group.variable.data.class.zeros(*group.variable.data.shape)
+        )
+      end
+
+      def select_first_axis(array, selection)
+        array[*([selection] + Array.new(array.ndim - 1, true))].dup
+      end
+
+      def append_zeros(array, count)
+        output = array.class.zeros(*([array.shape[0] + count] + array.shape[1..]))
+        output[*([0...array.shape[0]] + Array.new(array.ndim - 1, true))] = array
+        output
+      end
+    end
+  end
+end

data/lib/gsplat/optim/lr_scheduler.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Optim
+    # Closed-form exponential learning-rate schedule.
+    class ExponentialLR
+      attr_reader :max_steps, :step_count
+
+      def initialize(optimizer, lr_final:, max_steps:)
+        raise ArgumentError, "lr_final must be positive" unless lr_final.positive?
+        raise ArgumentError, "max_steps must be positive" unless max_steps.is_a?(Integer) && max_steps.positive?
+
+        @optimizer = optimizer
+        @lr_final = lr_final
+        @max_steps = max_steps
+        @step_count = 0
+        @initial_rates = optimizer.groups.transform_values(&:lr)
+      end
+
+      # Advances or sets the schedule and updates every optimizer group.
+      #
+      # @param step [Integer, nil] absolute step, or nil to advance once
+      # @return [Numeric, Hash{Symbol=>Numeric}] updated learning rates
+      def step(step = nil)
+        @step_count = step || (step_count + 1)
+        progress = [step_count.to_f / max_steps, 1.0].min
+        @optimizer.groups.each_key do |name|
+          initial = @initial_rates.fetch(name)
+          rate = initial * ((@lr_final / initial)**progress)
+          @optimizer.set_learning_rate(name, rate)
+        end
+        @optimizer.learning_rate
+      end
+    end
+  end
+end