generate.py

import symforce

symforce.set_epsilon_to_symbol()


from pathlib import Path
import symforce.symbolic as sf
from symforce.codegen import (
    CppConfig,
    Codegen,
    values_codegen,
    CodegenConfig,
    template_util,
)
from symforce.codegen.codegen_config import RenderTemplateConfig
import numpy as np
import sym
from symforce import typing as T
from symforce.values import Values
from utils import build_cube_values

import re
import shutil
import textwrap

NUM_POINTS_PER_FACE = 1000
NUM_CORRESPONDENCES = NUM_POINTS_PER_FACE * 6


def point_to_plane_residual(
    world_T_lidar: sf.Pose3,
    point_lidar: sf.V3,
    centroid_world: sf.V3,
    normal_world: sf.V3,
) -> sf.V3:

    estimated_position = world_T_lidar * point_lidar
    v = estimated_position - centroid_world

    d_pred = normal_world.T * v

    return d_pred


def point_to_plane_sum_residual(
    world_T_lidar: sf.Pose3,
    points: T.List[sf.V3],
    centroids: T.List[sf.V3],
    normals: T.List[sf.V3],
) -> sf.V1:
    residuals = []
    for i in range(points.shape[0]):
        residuals.append(
            point_to_plane_residual(
                world_T_lidar,
                points[i, :].T,
                centroids[i, :].T,
                normals[i, :].T,
            )
        )
    return sf.Matrix(residuals)


def build_sum_residual(values: Values) -> Values:
    return sf.Matrix(
        point_to_plane_sum_residual(
            values.attr.world_T_lidar,
            values.attr.points,
            values.attr.centroids,
            values.attr.normals,
        )
    )


def build_codegen_object(config: T.Optional[CodegenConfig] = None) -> Codegen:
    """
    Create Codegen object for the linearization function
    """
    if config is None:
        config = CppConfig()

    values = build_cube_values(NUM_POINTS_PER_FACE)

    def symbolic(k: str, v: T.Any) -> T.Any:
        if isinstance(v, sym.Pose3):
            return sf.Pose3.symbolic(k)
        elif isinstance(v, float):
            return sf.Symbol(k)
        elif isinstance(v, np.ndarray):
            if len(v.shape) == 1:
                return sf.Matrix(v.shape[0], 1).symbolic(k)
            else:
                return sf.Matrix(*v.shape).symbolic(k)
        else:
            assert False, k

    values = Values(**{key: symbolic(key, v) for key, v in values.items_recursive()})

    residual = build_sum_residual(values)

    flat_keys = {key: re.sub(r"[\.\[\]]+", "_", key) for key in values.keys_recursive()}

    inputs = Values(
        **{flat_keys[key]: value for key, value in values.items_recursive()}
    )
    outputs = Values(residual=residual)

    optimized_keys = ["world_T_lidar"]

    linearization_func = Codegen(
        inputs=inputs,
        outputs=outputs,
        config=config,
        docstring=textwrap.dedent(
            """
            This function was autogenerated. Do not modify by hand.

            Computes the linearization of the residual around the given state,
            and returns the relevant information about the resulting linear system.

            Input args: The state to linearize around

            Output args:
                residual (Eigen::Matrix*): The residual vector
            """
        ),
    ).with_linearization(
        name="linearization",
        which_args=[flat_keys[key] for key in optimized_keys],
        sparse_linearization=True,
    )

    return linearization_func


def generate_point_to_plane_residual_code(
    output_dir: T.Optional[Path] = None, print_code: bool = False
) -> None:
    """
    Generate C++ code for the point-to-plane residual function. A C++ Factor can then be
    constructed and optimized from this function without any Python dependency.
    """
    # Create a Codegen object for the symbolic residual function, targeted at C++
    codegen = Codegen.function(point_to_plane_residual, config=CppConfig())

    # Create a Codegen object that computes a linearization from the residual Codegen object,
    # by introspecting and symbolically differentiating the given arguments
    codegen_with_linearization = codegen.with_linearization(
        which_args=["world_T_lidar"]
    )

    # Generate the function and print the code
    generated_paths = codegen_with_linearization.generate_function(
        output_dir=output_dir,
        namespace="ICP",
        skip_directory_nesting=True,
    )
    if print_code:
        print(generated_paths.generated_files[0].read_text())

    # generate_function writes to a new temporary directory if output_dir is None. Delete the
    # temporary directory.
    if output_dir is None:
        shutil.rmtree(generated_paths.output_dir)


def generate(output_dir: Path) -> None:
    generate_point_to_plane_residual_code(output_dir)
    values = build_cube_values(NUM_POINTS_PER_FACE)
    values_codegen.generate_values_keys(
        values,
        output_dir,
        config=CppConfig(),
        skip_directory_nesting=True,
    )
    build_codegen_object(config=CppConfig()).generate_function(
        output_dir,
        namespace="ICP",
        skip_directory_nesting=True,
    )


if __name__ == "__main__":
    generate(Path("gen"))