Skip to content

Commit

Permalink
Simplify geo_factors_codegen w/ skip_dir_nesting
Browse files Browse the repository at this point in the history
Previously, `geo_factors_codegen.py` was using an ad-hoc system of
generating all the C++ factors into a `factors` folder. It did this by
calling `Codegen.generate_function` to generate the code into a
temporary file, read the file into a string, then re-wrote the string
into the desired location. I assume this was to avoid all the fluff
that's generated by default with `Codegen.generate_function`.

However, there is the `skip_directory_nesting` optional argument for
`Codegen.generate_function` which does precisely that (I think the code
in this file might have been written before that option was added).

So, to reduce confusion (such as the confusion I faced when I first
started looking at this file) and complexity, I rewrote the code to
instead use the `skip_directory_nesting` argument.

Topic: geo_factors_use_skip_directory_nesting
  • Loading branch information
bradley-solliday-skydio committed Jan 18, 2023
1 parent 4c36881 commit f9a6f2d
Showing 1 changed file with 13 additions and 56 deletions.
69 changes: 13 additions & 56 deletions symforce/codegen/geo_factors_codegen.py
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")

0 comments on commit f9a6f2d

Please sign in to comment.