gsplat 1.0.0

data/lib/gsplat/backend/ruby/projection_input_vjp.rb

@@ -0,0 +1,252 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Backend
+    # Input-side VJP helpers for RubyProjectionBackward.
+    module RubyProjectionBackward
+      module_function
+
+      # rubocop:disable Metrics/ParameterLists
+      def camera_mean_vjp(means, intrinsics, grad_means2d, grad_depths, grad_jacobians, width, height,
+                          options)
+        # rubocop:enable Metrics/ParameterLists
+        output = means.class.zeros(*means.shape)
+        # rubocop:disable Metrics/BlockLength
+        means.shape[0].times do |camera_index|
+          camera_means = means[camera_index, true, true]
+          camera_intrinsics = intrinsics[camera_index, true, true]
+          output_gradient = grad_means2d[camera_index, true, true]
+          camera_model = options.fetch(:camera_model)
+          gradient = if camera_model.to_s == "ortho"
+                       ortho_mean_vjp(camera_means, camera_intrinsics, output_gradient)
+                     elsif extended_camera?(options)
+                       distorted_mean_vjp(
+                         camera_means,
+                         camera_intrinsics,
+                         output_gradient,
+                         grad_jacobians[camera_index, true, true, true],
+                         options,
+                         camera_index
+                       )
+                     else
+                       projected_mean_vjp(camera_means, camera_intrinsics, output_gradient) +
+                         pinhole_jacobian_vjp(
+                           camera_means,
+                           camera_intrinsics,
+                           grad_jacobians[camera_index, true, true, true],
+                           width,
+                           height
+                         )
+                     end
+          gradient = add_depth_vjp(
+            gradient,
+            camera_means,
+            grad_depths[camera_index, true],
+            options.fetch(:global_z_order, true)
+          )
+          output[camera_index, true, true] = gradient
+        end
+        # rubocop:enable Metrics/BlockLength
+        output
+      end
+      private_class_method :camera_mean_vjp
+
+      def add_depth_vjp(gradient, means, grad_depth, global_z_order)
+        if global_z_order
+          gradient[true, 2] += grad_depth
+          return gradient
+        end
+
+        distance = (means**2).sum(axis: 1)**0.5
+        safe_distance = distance.dup
+        safe_distance[safe_distance.lt(1e-12)] = 1 if safe_distance.lt(1e-12).any?
+        gradient + (means * (grad_depth / safe_distance).reshape(means.shape[0], 1))
+      end
+      private_class_method :add_depth_vjp
+
+      def extended_camera?(options)
+        options.fetch(:camera_model).to_s == "fisheye" ||
+          %i[radial_coeffs tangential_coeffs thin_prism_coeffs].any? { |key| options[key] }
+      end
+      private_class_method :extended_camera?
+
+      # rubocop:disable Metrics/AbcSize, Metrics/ParameterLists
+      def distorted_mean_vjp(means, intrinsics, grad_projected, grad_jacobians, options, camera_index)
+        # rubocop:enable Metrics/AbcSize, Metrics/ParameterLists
+        output = means.class.zeros(*means.shape)
+        epsilon = means.instance_of?(Numo::DFloat) ? 1e-4 : 2e-3
+        means.shape[0].times do |gaussian_index|
+          point = means[gaussian_index, true].to_a
+          3.times do |axis|
+            positive = point.dup
+            negative = point.dup
+            positive[axis] += epsilon
+            negative[axis] -= epsilon
+            positive_outputs = distorted_point_outputs(positive, intrinsics, options, camera_index, means.class)
+            negative_outputs = distorted_point_outputs(negative, intrinsics, options, camera_index, means.class)
+            projected_derivative = (positive_outputs[0] - negative_outputs[0]) / (2 * epsilon)
+            jacobian_derivative = (positive_outputs[1] - negative_outputs[1]) / (2 * epsilon)
+            output[gaussian_index, axis] =
+              (projected_derivative * grad_projected[gaussian_index, true]).sum +
+              (jacobian_derivative * grad_jacobians[gaussian_index, true, true]).sum
+          end
+        end
+        output
+      end
+      private_class_method :distorted_mean_vjp
+
+      def distorted_point_outputs(point, intrinsics, options, camera_index, type)
+        projected, jacobians = Math::CameraDistortion.project_camera(
+          type.cast([point]),
+          intrinsics,
+          options.fetch(:camera_model),
+          radial: coefficient_row(options[:radial_coeffs], camera_index),
+          tangential: coefficient_row(options[:tangential_coeffs], camera_index),
+          thin_prism: coefficient_row(options[:thin_prism_coeffs], camera_index)
+        )
+        [projected.flatten, jacobians.reshape(2, 3)]
+      end
+      private_class_method :distorted_point_outputs
+
+      # rubocop:disable Metrics/AbcSize
+      def projected_mean_vjp(means, intrinsics, gradient)
+        x_coord = means[true, 0]
+        y_coord = means[true, 1]
+        z_coord = means[true, 2]
+        homogeneous_x = (intrinsics[0, 0] * x_coord) + (intrinsics[0, 1] * y_coord) +
+                        (intrinsics[0, 2] * z_coord)
+        homogeneous_y = (intrinsics[1, 0] * x_coord) + (intrinsics[1, 1] * y_coord) +
+                        (intrinsics[1, 2] * z_coord)
+        grad_x = gradient[true, 0]
+        grad_y = gradient[true, 1]
+        output = means.class.zeros(*means.shape)
+        output[true, 0] = (
+          (grad_x * intrinsics[0, 0]) + (grad_y * intrinsics[1, 0])
+        ) / z_coord
+        output[true, 1] = (
+          (grad_x * intrinsics[0, 1]) + (grad_y * intrinsics[1, 1])
+        ) / z_coord
+        output[true, 2] = (
+          (grad_x * ((intrinsics[0, 2] * z_coord) - homogeneous_x)) +
+          (grad_y * ((intrinsics[1, 2] * z_coord) - homogeneous_y))
+        ) / (z_coord**2)
+        output
+      end
+      # rubocop:enable Metrics/AbcSize
+      private_class_method :projected_mean_vjp
+
+      def ortho_mean_vjp(means, intrinsics, gradient)
+        output = means.class.zeros(*means.shape)
+        output[true, 0] = gradient[true, 0] * intrinsics[0, 0]
+        output[true, 1] = gradient[true, 1] * intrinsics[1, 1]
+        output
+      end
+      private_class_method :ortho_mean_vjp
+
+      # rubocop:disable Metrics/AbcSize
+      def pinhole_jacobian_vjp(means, intrinsics, gradient, width, height)
+        x_coord = means[true, 0]
+        y_coord = means[true, 1]
+        z_coord = means[true, 2]
+        focal_x = intrinsics[0, 0]
+        focal_y = intrinsics[1, 1]
+        ratio_x = x_coord / z_coord
+        ratio_y = y_coord / z_coord
+        min_x = -((intrinsics[0, 2] / focal_x) + (0.15 * width / focal_x))
+        max_x = ((width - intrinsics[0, 2]) / focal_x) + (0.15 * width / focal_x)
+        min_y = -((intrinsics[1, 2] / focal_y) + (0.15 * height / focal_y))
+        max_y = ((height - intrinsics[1, 2]) / focal_y) + (0.15 * height / focal_y)
+        clipped_x, inside_x = clipped_ratio_and_mask(ratio_x, min_x, max_x)
+        clipped_y, inside_y = clipped_ratio_and_mask(ratio_y, min_y, max_y)
+        inside_x = means.class.cast(inside_x)
+        inside_y = means.class.cast(inside_y)
+        output = means.class.zeros(*means.shape)
+        output[true, 0] = gradient[true, 0, 2] * (-focal_x / (z_coord**2)) * inside_x
+        output[true, 1] = gradient[true, 1, 2] * (-focal_y / (z_coord**2)) * inside_y
+        derivative_x_z = (
+          (inside_x * 2 * focal_x * x_coord / (z_coord**3)) +
+          ((1 - inside_x) * focal_x * clipped_x / (z_coord**2))
+        )
+        derivative_y_z = (
+          (inside_y * 2 * focal_y * y_coord / (z_coord**3)) +
+          ((1 - inside_y) * focal_y * clipped_y / (z_coord**2))
+        )
+        output[true, 2] = (
+          (gradient[true, 0, 0] * (-focal_x / (z_coord**2))) +
+          (gradient[true, 1, 1] * (-focal_y / (z_coord**2))) +
+          (gradient[true, 0, 2] * derivative_x_z) +
+          (gradient[true, 1, 2] * derivative_y_z)
+        )
+        output
+      end
+      # rubocop:enable Metrics/AbcSize
+      private_class_method :pinhole_jacobian_vjp
+
+      def clipped_ratio_and_mask(ratio, minimum, maximum)
+        clipped = ratio.dup
+        below = ratio.lt(minimum)
+        above = ratio.gt(maximum)
+        clipped[below] = minimum if below.any?
+        clipped[above] = maximum if above.any?
+        [clipped, below.eq(0) & above.eq(0)]
+      end
+      private_class_method :clipped_ratio_and_mask
+
+      def world_to_cam_vjp(grad_means, grad_covars, viewmats)
+        gaussian_count = grad_means.shape[1]
+        world_means = grad_means.class.zeros(gaussian_count, 3)
+        world_covars = grad_means.class.zeros(gaussian_count, 3, 3)
+        grad_means.shape[0].times do |camera_index|
+          rotation = viewmats[camera_index, 0...3, 0...3]
+          rotation_batch = grad_means.class.zeros(gaussian_count, 3, 3)
+          rotation_batch[true, true, true] = rotation
+          world_means += grad_means[camera_index, true, true].dot(rotation)
+          world_covars += Math::Mat.matmul_batch(
+            Math::Mat.matmul_batch(rotation_batch.transpose(0, 2, 1),
+                                   grad_covars[camera_index, true, true, true]),
+            rotation_batch
+          )
+        end
+        [world_means, world_covars]
+      end
+      private_class_method :world_to_cam_vjp
+
+      # rubocop:disable Metrics/ParameterLists
+      def input_covariance_gradients(means, covars, quaternions, scales, grad_means, grad_covars)
+        # rubocop:enable Metrics/ParameterLists
+        if covars
+          return [
+            grad_means,
+            format_covariance_gradient(grad_covars, covars.shape, means.shape[0]),
+            nil,
+            nil
+          ]
+        end
+
+        grad_quaternions, grad_scales = RubyQuatScaleToCovarPreci.backward(
+          means.class.cast(quaternions),
+          means.class.cast(scales),
+          grad_covars,
+          nil,
+          triu: false
+        )
+        [grad_means, nil, grad_quaternions, grad_scales]
+      end
+      private_class_method :input_covariance_gradients
+
+      def format_covariance_gradient(gradient, input_shape, gaussian_count)
+        return gradient if input_shape == [gaussian_count, 3, 3]
+
+        output = gradient.class.zeros(gaussian_count, 6)
+        output[true, 0] = gradient[true, 0, 0]
+        output[true, 1] = gradient[true, 0, 1] + gradient[true, 1, 0]
+        output[true, 2] = gradient[true, 0, 2] + gradient[true, 2, 0]
+        output[true, 3] = gradient[true, 1, 1]
+        output[true, 4] = gradient[true, 1, 2] + gradient[true, 2, 1]
+        output[true, 5] = gradient[true, 2, 2]
+        output
+      end
+      private_class_method :format_covariance_gradient
+    end
+  end
+end

data/lib/gsplat/backend/ruby/quat_scale_to_covar_preci.rb

@@ -0,0 +1,139 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Backend
+    # Numo implementation of quaternion/scale covariance conversion.
+    module RubyQuatScaleToCovarPreci
+      module_function
+
+      def forward(quaternions, scales, compute_covar:, compute_preci:, triu:)
+        quaternions, scales, leading_shape = validate_inputs(quaternions, scales, compute_preci)
+        count = quaternions.size / 4
+        rotations = Math::Quaternion.to_rotmat(quaternions).reshape(count, 3, 3)
+        scales_flat = scales.reshape(count, 3)
+        covariance = covariance_matrix(rotations, scales_flat) if compute_covar
+        precision = precision_matrix(rotations, scales_flat) if compute_preci
+        [
+          format_output(covariance, leading_shape, triu),
+          format_output(precision, leading_shape, triu)
+        ]
+      end
+
+      def backward(quaternions, scales, grad_covar, grad_preci, triu:)
+        quaternions, scales, leading_shape = validate_inputs(quaternions, scales, !grad_preci.nil?)
+        count = quaternions.size / 4
+        rotations = Math::Quaternion.to_rotmat(quaternions).reshape(count, 3, 3)
+        scales_flat = scales.reshape(count, 3)
+        grad_rotations = quaternions.class.zeros(count, 3, 3)
+        grad_scales = quaternions.class.zeros(count, 3)
+        if grad_covar
+          accumulate_covar_gradients!(
+            grad_rotations,
+            grad_scales,
+            rotations,
+            scales_flat,
+            expand_gradient(grad_covar, leading_shape, triu, quaternions.class)
+          )
+        end
+        if grad_preci
+          accumulate_preci_gradients!(
+            grad_rotations,
+            grad_scales,
+            rotations,
+            scales_flat,
+            expand_gradient(grad_preci, leading_shape, triu, quaternions.class)
+          )
+        end
+        grad_quaternions = Math::Quaternion.to_rotmat_vjp(
+          quaternions,
+          grad_rotations.reshape(*(leading_shape + [3, 3]))
+        )
+        [grad_quaternions, grad_scales.reshape(*(leading_shape + [3]))]
+      end
+
+      def validate_inputs(quaternions, scales, compute_preci)
+        unless quaternions.is_a?(Numo::NArray) && [Numo::SFloat, Numo::DFloat].include?(quaternions.class)
+          raise ArgumentError, "quaternions must be Numo::SFloat or Numo::DFloat"
+        end
+
+        expected_scales = quaternions.shape[0...-1] + [3]
+        unless quaternions.ndim >= 1 && quaternions.shape[-1] == 4 && scales.shape == expected_scales
+          raise ShapeError,
+                "expected quaternions [...,4] and scales #{expected_scales.inspect}, " \
+                "got #{quaternions.shape.inspect} and #{scales.shape.inspect}"
+        end
+
+        scales = quaternions.class.cast(scales)
+        raise Gsplat::Error, "precision is undefined for zero scales" if compute_preci && scales.eq(0).any?
+
+        [quaternions, scales, quaternions.shape[0...-1]]
+      end
+      private_class_method :validate_inputs
+
+      def covariance_matrix(rotations, scales)
+        transform = rotations * scales.reshape(scales.shape[0], 1, 3)
+        Math::Mat.matmul_batch(transform, transform.transpose(0, 2, 1))
+      end
+      private_class_method :covariance_matrix
+
+      def precision_matrix(rotations, scales)
+        transform = rotations * (1.0 / scales).reshape(scales.shape[0], 1, 3)
+        Math::Mat.matmul_batch(transform, transform.transpose(0, 2, 1))
+      end
+      private_class_method :precision_matrix
+
+      def format_output(matrix, leading_shape, triu)
+        return nil unless matrix
+        return matrix.reshape(*(leading_shape + [3, 3])) unless triu
+
+        output = matrix.class.zeros(matrix.shape[0], 6)
+        [[0, 0], [0, 1], [0, 2], [1, 1], [1, 2], [2, 2]].each_with_index do |(row, column), index|
+          output[true, index] = matrix[true, row, column]
+        end
+        output.reshape(*(leading_shape + [6]))
+      end
+      private_class_method :format_output
+
+      def expand_gradient(gradient, leading_shape, triu, type)
+        gradient = type.cast(gradient)
+        expected = leading_shape + [triu ? 6 : 3, triu ? nil : 3].compact
+        unless gradient.shape == expected
+          raise ShapeError, "expected gradient #{expected.inspect}, got #{gradient.shape.inspect}"
+        end
+        return gradient.reshape(gradient.size / 9, 3, 3) unless triu
+
+        flat = gradient.reshape(gradient.size / 6, 6)
+        output = type.zeros(flat.shape[0], 3, 3)
+        [[0, 0], [1, 1], [2, 2]].each_with_index do |(row, column), diagonal_index|
+          output[true, row, column] = flat[true, [0, 3, 5].fetch(diagonal_index)]
+        end
+        [[0, 1, 1], [0, 2, 2], [1, 2, 4]].each do |row, column, index|
+          half = flat[true, index] * 0.5
+          output[true, row, column] = half
+          output[true, column, row] = half
+        end
+        output
+      end
+      private_class_method :expand_gradient
+
+      def accumulate_covar_gradients!(grad_rotations, grad_scales, rotations, scales, gradient)
+        symmetric = gradient + gradient.transpose(0, 2, 1)
+        transform = rotations * scales.reshape(scales.shape[0], 1, 3)
+        grad_transform = Math::Mat.matmul_batch(symmetric, transform)
+        grad_rotations[] = grad_rotations + (grad_transform * scales.reshape(scales.shape[0], 1, 3))
+        grad_scales[] = grad_scales + (grad_transform * rotations).sum(axis: 1)
+      end
+      private_class_method :accumulate_covar_gradients!
+
+      def accumulate_preci_gradients!(grad_rotations, grad_scales, rotations, scales, gradient)
+        inverse_scales = 1.0 / scales
+        symmetric = gradient + gradient.transpose(0, 2, 1)
+        transform = rotations * inverse_scales.reshape(scales.shape[0], 1, 3)
+        grad_transform = Math::Mat.matmul_batch(symmetric, transform)
+        grad_rotations[] = grad_rotations + (grad_transform * inverse_scales.reshape(scales.shape[0], 1, 3))
+        grad_scales[] = grad_scales - ((grad_transform * rotations).sum(axis: 1) / (scales**2))
+      end
+      private_class_method :accumulate_preci_gradients!
+    end
+  end
+end

data/lib/gsplat/backend/ruby/rasterize_to_indices_in_range.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+module Gsplat
+  module Backend
+    # Reference implementation of contribution-index enumeration.
+    module RubyRasterizeToIndicesInRange
+      module_function
+
+      # rubocop:disable Metrics/ParameterLists
+      def forward(range_start, range_end, transmittances, means2d, conics, opacities,
+                  width, height, tile_size, isect_offsets, flatten_ids)
+        # rubocop:enable Metrics/ParameterLists
+        values = validate_inputs(
+          range_start, range_end, transmittances, means2d, conics, opacities,
+          width, height, tile_size, isect_offsets, flatten_ids
+        )
+        transmittances, means2d, conics, opacities, offsets, ids = values
+        gaussian_ids = []
+        pixel_ids = []
+        image_ids = []
+        means = means2d.reshape(means2d.shape[0] * means2d.shape[1], 2)
+        conic_values = conics.reshape(conics.shape[0] * conics.shape[1], 3)
+        opacity_values = opacities.flatten
+        means2d.shape[0].times do |image|
+          height.times do |row|
+            width.times do |column|
+              enumerate_pixel!(
+                gaussian_ids, pixel_ids, image_ids, range_start, range_end,
+                transmittances, means, conic_values, opacity_values, offsets, ids,
+                image, row, column, means2d.shape[1], width, tile_size
+              )
+            end
+          end
+        end
+        [Numo::Int64.cast(gaussian_ids), Numo::Int64.cast(pixel_ids), Numo::Int64.cast(image_ids)]
+      end
+
+      # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/ParameterLists, Metrics/PerceivedComplexity
+      def validate_inputs(range_start, range_end, transmittances, means2d, conics, opacities,
+                          width, height, tile_size, isect_offsets, flatten_ids)
+        # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/ParameterLists, Metrics/PerceivedComplexity
+        valid_range = range_start.is_a?(Integer) && range_end.is_a?(Integer) &&
+                      range_start >= 0 && range_end >= range_start
+        raise ArgumentError, "range must be nonnegative integers with start <= end" unless valid_range
+        unless means2d.is_a?(Numo::NArray) && [Numo::SFloat, Numo::DFloat].include?(means2d.class)
+          raise ArgumentError, "means2d must be Numo::SFloat or Numo::DFloat"
+        end
+
+        valid_means = means2d.ndim == 3 && means2d.shape[-1] == 2
+        raise ShapeError, "expected means2d [C,N,2]" unless valid_means
+
+        camera_count, gaussian_count = means2d.shape[0...2]
+        type = means2d.class
+        transmittances = type.cast(transmittances)
+        conics = type.cast(conics)
+        opacities = type.cast(opacities)
+        valid = transmittances.shape == [camera_count, height, width] &&
+                conics.shape == [camera_count, gaussian_count, 3] &&
+                opacities.shape == [camera_count, gaussian_count]
+        raise ShapeError, "invalid contribution-index tensor shapes" unless valid
+
+        offsets = Numo::Int32.cast(isect_offsets)
+        ids = Numo::Int32.cast(flatten_ids)
+        valid_offsets = offsets.ndim == 3 && offsets.shape[0] == camera_count
+        raise ShapeError, "expected isect_offsets [C,TH,TW]" unless valid_offsets
+        raise ShapeError, "expected flatten_ids [M]" unless ids.ndim == 1
+
+        RubyRasterizeToPixels.send(:validate_indices!, ids, camera_count * gaussian_count, offsets)
+        RubyRasterizeToPixels.send(
+          :validate_dimensions!, width, height, tile_size, offsets.shape[-1], offsets.shape[-2]
+        )
+        [transmittances, means2d, conics, opacities, offsets, ids]
+      end
+      private_class_method :validate_inputs
+
+      # rubocop:disable Metrics/AbcSize, Metrics/ParameterLists
+      def enumerate_pixel!(gaussian_ids, pixel_ids, image_ids, range_start, range_end,
+                           transmittances, means, conics, opacities, offsets, flatten_ids,
+                           image, row, column, gaussian_count, width, tile_size)
+        # rubocop:enable Metrics/AbcSize, Metrics/ParameterLists
+        range = intersection_batch(
+          offsets, flatten_ids, image, row, column, tile_size, range_start, range_end
+        )
+        transmittance = transmittances[image, row, column].to_f
+        range.each do |intersection|
+          flat_id = flatten_ids[intersection].to_i
+          delta_x = means[flat_id, 0].to_f - column - 0.5
+          delta_y = means[flat_id, 1].to_f - row - 0.5
+          conic = conics[flat_id, true]
+          sigma = (0.5 * ((conic[0] * delta_x * delta_x) + (conic[2] * delta_y * delta_y))) +
+                  (conic[1] * delta_x * delta_y)
+          alpha = [opacities[flat_id].to_f * ::Math.exp(-sigma), RubyAccumulate::ALPHA_CLAMP].min
+          next if sigma.negative? || alpha < RubyAccumulate::ALPHA_SKIP
+
+          next_transmittance = transmittance * (1 - alpha)
+          break if next_transmittance <= RubyAccumulate::TRANSMITTANCE_STOP
+
+          gaussian_ids << (flat_id % gaussian_count)
+          pixel_ids << ((row * width) + column)
+          image_ids << image
+          transmittance = next_transmittance
+        end
+      end
+      private_class_method :enumerate_pixel!
+
+      # rubocop:disable Metrics/ParameterLists
+      def intersection_batch(offsets, flatten_ids, image, row, column, tile_size, first_batch, last_batch)
+        # rubocop:enable Metrics/ParameterLists
+        tile_height, tile_width = offsets.shape[-2..]
+        flat_tile = (((image * tile_height) + (row / tile_size)) * tile_width) + (column / tile_size)
+        tile_start = offsets[flat_tile].to_i
+        tile_end = flat_tile + 1 < offsets.size ? offsets[flat_tile + 1].to_i : flatten_ids.size
+        batch_size = tile_size * tile_size
+        start_index = [tile_start + (first_batch * batch_size), tile_end].min
+        end_index = [tile_start + (last_batch * batch_size), tile_end].min
+        start_index...end_index
+      end
+      private_class_method :intersection_batch
+    end
+  end
+end