Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Simplify geo_factors_codegen w/ skip_dir_nesting #295

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Simplify geo_factors_codegen w/ skip_dir_nesting #295

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 13 additions & 56 deletions symforce/codegen/geo_factors_codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@

import symforce.symbolic as sf
from symforce import ops
from symforce import python_util
from symforce import typing as T
from symforce.codegen import Codegen
from symforce.codegen import CppConfig
Expand Down Expand Up @@ -84,38 +83,10 @@ def prior_factor(
return residual


def get_function_code(codegen: Codegen, cleanup: bool = True) -> str:
def generate_between_factors(types: T.Sequence[T.Type], output_dir: T.Openable) -> None:
"""
Return just the function code from a Codegen object.
Generates between factors for each type in types into output_dir.
"""
# Codegen
data = codegen.generate_function()

# Read
assert codegen.name is not None
filename = "{}.h".format(codegen.name)
func_code = (data.function_dir / filename).read_text()

# Cleanup
if cleanup:
python_util.remove_if_exists(data.output_dir)

return func_code


def get_filename(codegen: Codegen) -> str:
"""
Helper to get appropriate filename
"""
assert codegen.name is not None
return codegen.name + ".h"


def get_between_factors(types: T.Sequence[T.Type]) -> T.Dict[str, str]:
"""
Compute
"""
files_dict: T.Dict[str, str] = {}
for cls in types:
tangent_dim = ops.LieGroupOps.tangent_dim(cls)
between_codegen = Codegen.function(
Expand All @@ -125,7 +96,7 @@ def get_between_factors(types: T.Sequence[T.Type]) -> T.Dict[str, str]:
config=CppConfig(),
docstring=get_between_factor_docstring("a_T_b"),
).with_linearization(name=f"between_factor_{cls.__name__.lower()}", which_args=["a", "b"])
files_dict[get_filename(between_codegen)] = get_function_code(between_codegen)
between_codegen.generate_function(output_dir, skip_directory_nesting=True)

prior_codegen = Codegen.function(
func=prior_factor,
Expand All @@ -134,14 +105,12 @@ def get_between_factors(types: T.Sequence[T.Type]) -> T.Dict[str, str]:
config=CppConfig(),
docstring=get_prior_docstring(),
).with_linearization(name=f"prior_factor_{cls.__name__.lower()}", which_args=["value"])
files_dict[get_filename(prior_codegen)] = get_function_code(prior_codegen)

return files_dict
prior_codegen.generate_function(output_dir, skip_directory_nesting=True)


def get_pose3_extra_factors(files_dict: T.Dict[str, str]) -> None:
def generate_pose3_extra_factors(output_dir: T.Openable) -> None:
"""
Generates factors specific to Poses which penalize individual components
Generates factors specific to Poses which penalize individual components into output_dir.

This includes factors for only the position or rotation components of a Pose. This can't be
done by wrapping the other generated functions because we need jacobians with respect to the
Expand Down Expand Up @@ -210,53 +179,41 @@ def prior_factor_pose3_position(
config=CppConfig(),
docstring=get_between_factor_docstring("a_R_b"),
).with_linearization(name="between_factor_pose3_rotation", which_args=["a", "b"])
between_rotation_codegen.generate_function(output_dir, skip_directory_nesting=True)

between_position_codegen = Codegen.function(
func=between_factor_pose3_position,
output_names=["res"],
config=CppConfig(),
docstring=get_between_factor_docstring("a_t_b"),
).with_linearization(name="between_factor_pose3_position", which_args=["a", "b"])
between_position_codegen.generate_function(output_dir, skip_directory_nesting=True)

between_translation_norm_codegen = Codegen.function(
func=between_factor_pose3_translation_norm, output_names=["res"], config=CppConfig()
).with_linearization(name="between_factor_pose3_translation_norm", which_args=["a", "b"])
between_translation_norm_codegen.generate_function(output_dir, skip_directory_nesting=True)

prior_rotation_codegen = Codegen.function(
func=prior_factor_pose3_rotation,
output_names=["res"],
config=CppConfig(),
docstring=get_prior_docstring(),
).with_linearization(name="prior_factor_pose3_rotation", which_args=["value"])
prior_rotation_codegen.generate_function(output_dir, skip_directory_nesting=True)

prior_position_codegen = Codegen.function(
func=prior_factor_pose3_position,
output_names=["res"],
config=CppConfig(),
docstring=get_prior_docstring(),
).with_linearization(name="prior_factor_pose3_position", which_args=["value"])

files_dict[get_filename(between_rotation_codegen)] = get_function_code(between_rotation_codegen)
files_dict[get_filename(between_position_codegen)] = get_function_code(between_position_codegen)
files_dict[get_filename(between_translation_norm_codegen)] = get_function_code(
between_translation_norm_codegen
)
files_dict[get_filename(prior_rotation_codegen)] = get_function_code(prior_rotation_codegen)
files_dict[get_filename(prior_position_codegen)] = get_function_code(prior_position_codegen)
prior_position_codegen.generate_function(output_dir, skip_directory_nesting=True)


def generate(output_dir: Path) -> None:
"""
Prior factors and between factors for C++.
"""
# Compute code
files_dict = get_between_factors(types=TYPES)
get_pose3_extra_factors(files_dict)

# Create output dir
factors_dir = output_dir / "factors"
factors_dir.mkdir(parents=True)

# Write out
for filename, code in files_dict.items():
(factors_dir / filename).write_text(code)
generate_between_factors(types=TYPES, output_dir=output_dir / "factors")
generate_pose3_extra_factors(output_dir / "factors")