UIKA / submodules /pytorch3d /tests /test_harmonic_embedding.py

Add files using upload-large-folder tool

e2b9934 verified 3 months ago

6.03 kB

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import unittest

import torch
from pytorch3d.renderer.implicit import HarmonicEmbedding
from torch.distributions import MultivariateNormal

from .common_testing import TestCaseMixin


class TestHarmonicEmbedding(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        super().setUp()
        torch.manual_seed(1)

    def test_correct_output_dim(self):
        embed_fun = HarmonicEmbedding(n_harmonic_functions=2, append_input=False)
        # input_dims * (2 * n_harmonic_functions + int(append_input))
        output_dim = 3 * (2 * 2 + int(False))
        self.assertEqual(
            output_dim,
            embed_fun.get_output_dim_static(
                input_dims=3, n_harmonic_functions=2, append_input=False
            ),
        )
        self.assertEqual(output_dim, embed_fun.get_output_dim())

    def test_correct_frequency_range(self):
        embed_fun_log = HarmonicEmbedding(n_harmonic_functions=3)
        embed_fun_lin = HarmonicEmbedding(n_harmonic_functions=3, logspace=False)
        self.assertClose(embed_fun_log._frequencies, torch.FloatTensor((1.0, 2.0, 4.0)))
        self.assertClose(embed_fun_lin._frequencies, torch.FloatTensor((1.0, 2.5, 4.0)))

    def test_correct_embed_out(self):
        n_harmonic_functions = 2
        x = torch.randn((1, 5))
        D = 5 * n_harmonic_functions * 2  # sin + cos

        embed_fun = HarmonicEmbedding(
            n_harmonic_functions=n_harmonic_functions, append_input=False
        )
        embed_out = embed_fun(x)

        self.assertEqual(embed_out.shape, (1, D))
        # Sum the squares of the respective frequencies
        # cos^2(x) + sin^2(x) = 1
        sum_squares = embed_out[0, : D // 2] ** 2 + embed_out[0, D // 2 :] ** 2
        self.assertClose(sum_squares, torch.ones((D // 2)))

        # Test append input
        embed_fun = HarmonicEmbedding(
            n_harmonic_functions=n_harmonic_functions, append_input=True
        )
        embed_out_appended_input = embed_fun(x)
        self.assertClose(
            embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
        )
        # Last plane in output is the input
        self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
        self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)

    def test_correct_embed_out_with_diag_cov(self):
        n_harmonic_functions = 2
        x = torch.randn((1, 3))
        diag_cov = torch.randn((1, 3))
        D = 3 * n_harmonic_functions * 2  # sin + cos

        embed_fun = HarmonicEmbedding(
            n_harmonic_functions=n_harmonic_functions, append_input=False
        )
        embed_out = embed_fun(x, diag_cov=diag_cov)

        self.assertEqual(embed_out.shape, (1, D))

        # Compute the scaling factor introduce in MipNerf
        scale_factor = (
            -0.5 * diag_cov[..., None] * torch.pow(embed_fun._frequencies[None, :], 2)
        )
        scale_factor = torch.exp(scale_factor).reshape(1, -1).tile((1, 2))
        # If we remove this scaling factor, we should go back to the
        # classical harmonic embedding:
        # Sum the squares of the respective frequencies
        # cos^2(x) + sin^2(x) = 1
        embed_out_without_cov = embed_out / scale_factor
        sum_squares = (
            embed_out_without_cov[0, : D // 2] ** 2
            + embed_out_without_cov[0, D // 2 :] ** 2
        )
        self.assertClose(sum_squares, torch.ones((D // 2)))

        # Test append input
        embed_fun = HarmonicEmbedding(
            n_harmonic_functions=n_harmonic_functions, append_input=True
        )
        embed_out_appended_input = embed_fun(x, diag_cov=diag_cov)
        self.assertClose(
            embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
        )
        # Last plane in output is the input
        self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
        self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)

    def test_correct_behavior_between_ipe_and_its_estimation_from_harmonic_embedding(
        self,
    ):
        """
        Check that the HarmonicEmbedding with integrated_position_encoding (IPE) set to
        True is coherent with the HarmonicEmbedding.
        What is the idea behind this test?
        We wish to produce an IPE that is the expectation
        of our lifted multivariate gaussian, modulated by the sine and cosine of
        the coordinates. These expectation has a closed-form
        (see equations 11, 12, 13, 14 of [1]).
        We sample N elements from the multivariate gaussian defined by its mean and covariance
        and compute the HarmonicEmbedding. The expected value of those embeddings should be
        equal to our IPE.
        Inspired from:
        https://github.com/google/mipnerf/blob/84c969e0a623edd183b75693aed72a7e7c22902d/internal/mip_test.py#L359
        References:
            [1] `MIP-NeRF <https://arxiv.org/abs/2103.13415>`_.
        """
        num_dims = 3
        n_harmonic_functions = 6
        mean = torch.randn(num_dims)
        diag_cov = torch.rand(num_dims)

        he_fun = HarmonicEmbedding(
            n_harmonic_functions=n_harmonic_functions, logspace=True, append_input=False
        )
        ipe_fun = HarmonicEmbedding(
            n_harmonic_functions=n_harmonic_functions,
            append_input=False,
        )

        embedding_ipe = ipe_fun(mean, diag_cov=diag_cov)

        rand_mvn = MultivariateNormal(mean, torch.eye(num_dims) * diag_cov)

        # Providing a large enough number of samples
        # we should obtain an estimation close to our IPE
        num_samples = 100000
        embedding_he = he_fun(rand_mvn.sample_n(num_samples))
        self.assertClose(embedding_he.mean(0), embedding_ipe, rtol=1e-2, atol=1e-2)

	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the BSD-style license found in the
	# LICENSE file in the root directory of this source tree.

	import unittest

	import torch
	from pytorch3d.renderer.implicit import HarmonicEmbedding
	from torch.distributions import MultivariateNormal

	from .common_testing import TestCaseMixin


	class TestHarmonicEmbedding(TestCaseMixin, unittest.TestCase):
	def setUp(self) -> None:
	super().setUp()
	torch.manual_seed(1)

	def test_correct_output_dim(self):
	embed_fun = HarmonicEmbedding(n_harmonic_functions=2, append_input=False)
	# input_dims * (2 * n_harmonic_functions + int(append_input))
	output_dim = 3 * (2 * 2 + int(False))
	self.assertEqual(
	output_dim,
	embed_fun.get_output_dim_static(
	input_dims=3, n_harmonic_functions=2, append_input=False
	),
	)
	self.assertEqual(output_dim, embed_fun.get_output_dim())

	def test_correct_frequency_range(self):
	embed_fun_log = HarmonicEmbedding(n_harmonic_functions=3)
	embed_fun_lin = HarmonicEmbedding(n_harmonic_functions=3, logspace=False)
	self.assertClose(embed_fun_log._frequencies, torch.FloatTensor((1.0, 2.0, 4.0)))
	self.assertClose(embed_fun_lin._frequencies, torch.FloatTensor((1.0, 2.5, 4.0)))

	def test_correct_embed_out(self):
	n_harmonic_functions = 2
	x = torch.randn((1, 5))
	D = 5 * n_harmonic_functions * 2 # sin + cos

	embed_fun = HarmonicEmbedding(
	n_harmonic_functions=n_harmonic_functions, append_input=False
	)
	embed_out = embed_fun(x)

	self.assertEqual(embed_out.shape, (1, D))
	# Sum the squares of the respective frequencies
	# cos^2(x) + sin^2(x) = 1
	sum_squares = embed_out[0, : D // 2] 2 + embed_out[0, D // 2 :] 2
	self.assertClose(sum_squares, torch.ones((D // 2)))

	# Test append input
	embed_fun = HarmonicEmbedding(
	n_harmonic_functions=n_harmonic_functions, append_input=True
	)
	embed_out_appended_input = embed_fun(x)
	self.assertClose(
	embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
	)
	# Last plane in output is the input
	self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
	self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)

	def test_correct_embed_out_with_diag_cov(self):
	n_harmonic_functions = 2
	x = torch.randn((1, 3))
	diag_cov = torch.randn((1, 3))
	D = 3 * n_harmonic_functions * 2 # sin + cos

	embed_fun = HarmonicEmbedding(
	n_harmonic_functions=n_harmonic_functions, append_input=False
	)
	embed_out = embed_fun(x, diag_cov=diag_cov)

	self.assertEqual(embed_out.shape, (1, D))

	# Compute the scaling factor introduce in MipNerf
	scale_factor = (
	-0.5 * diag_cov[..., None] * torch.pow(embed_fun._frequencies[None, :], 2)
	)
	scale_factor = torch.exp(scale_factor).reshape(1, -1).tile((1, 2))
	# If we remove this scaling factor, we should go back to the
	# classical harmonic embedding:
	# Sum the squares of the respective frequencies
	# cos^2(x) + sin^2(x) = 1
	embed_out_without_cov = embed_out / scale_factor
	sum_squares = (
	embed_out_without_cov[0, : D // 2] ** 2
	+ embed_out_without_cov[0, D // 2 :] ** 2
	)
	self.assertClose(sum_squares, torch.ones((D // 2)))

	# Test append input
	embed_fun = HarmonicEmbedding(
	n_harmonic_functions=n_harmonic_functions, append_input=True
	)
	embed_out_appended_input = embed_fun(x, diag_cov=diag_cov)
	self.assertClose(
	embed_out_appended_input.shape, torch.tensor((1, D + x.shape[-1]))
	)
	# Last plane in output is the input
	self.assertClose(embed_out_appended_input[..., -x.shape[-1] :], x)
	self.assertClose(embed_out_appended_input[..., : -x.shape[-1]], embed_out)

	def test_correct_behavior_between_ipe_and_its_estimation_from_harmonic_embedding(
	self,
	):
	"""
	Check that the HarmonicEmbedding with integrated_position_encoding (IPE) set to
	True is coherent with the HarmonicEmbedding.

	What is the idea behind this test?

	We wish to produce an IPE that is the expectation
	of our lifted multivariate gaussian, modulated by the sine and cosine of
	the coordinates. These expectation has a closed-form
	(see equations 11, 12, 13, 14 of [1]).

	We sample N elements from the multivariate gaussian defined by its mean and covariance
	and compute the HarmonicEmbedding. The expected value of those embeddings should be
	equal to our IPE.

	Inspired from:
	https://github.com/google/mipnerf/blob/84c969e0a623edd183b75693aed72a7e7c22902d/internal/mip_test.py#L359

	References:
	[1] `MIP-NeRF <https://arxiv.org/abs/2103.13415>`_.
	"""
	num_dims = 3
	n_harmonic_functions = 6
	mean = torch.randn(num_dims)
	diag_cov = torch.rand(num_dims)

	he_fun = HarmonicEmbedding(
	n_harmonic_functions=n_harmonic_functions, logspace=True, append_input=False
	)
	ipe_fun = HarmonicEmbedding(
	n_harmonic_functions=n_harmonic_functions,
	append_input=False,
	)

	embedding_ipe = ipe_fun(mean, diag_cov=diag_cov)

	rand_mvn = MultivariateNormal(mean, torch.eye(num_dims) * diag_cov)

	# Providing a large enough number of samples
	# we should obtain an estimation close to our IPE
	num_samples = 100000
	embedding_he = he_fun(rand_mvn.sample_n(num_samples))
	self.assertClose(embedding_he.mean(0), embedding_ipe, rtol=1e-2, atol=1e-2)