Add diagonal & isotropic sqrt_infos in geo factors

<NOTE>
A core premise of this PR has turned out to be flawed. The idea
was that we can overload the factors by the type of sqrt_info, but that
only works if we actually call the factors with normal arguments, which
isn't the normal use case. Rather, what the user would ordinarily do
is pass the function pointer to a `sym::Factor` object, at which point
it is impossible to disambiguate which overload is desired. Instead you
would have to use a static_cast or something (which would be unpleasant
as typing out the exact signature would be a pain).

Now, the motivation to reduce op counts by taking advantage of simpler
sqrt_infos is still present and valid, but will need to rethink where
these variants will be located and what they will be named
</NOTE>

Often times the square root information matrix is a diagonal or
isotropic matrix. However, in our generated geo factors we always assume
they are full dense matrices leading to many more operations than
needed.

This commit modifies `geo_factors_codegen.py` to create 3 variants for
each geo factor, one which assumes `sqrt_info` is a square matrix (the
existing behavior), one which assumes it is a diagonal matrix, and one
which assumes it is an isotropic matrix.

These variants are all generated with the same name & are distinguished
only by their signatures (thus this only works for C++, which is all we
currently generate factors for anyway). The square version takes a square
eigen matrix, the diagonal version a eigen vector whose entries are
those of the diagonal of `sqrt_info`, and a single scalar for the
isotropic version.

The header that the user imports remains the same, except now rather
than containing the implementation directly, it re-imports the headers
for the 3 different versions (which are stores in sub-directories).

Topic: sqrt_info_variants
Relative: geo_factors_use_skip_directory_nesting

symforce/codegen/geo_factors_codegen.py

@@ -3,13 +3,17 @@
 # This source code is under the Apache 2.0 license found in the LICENSE file.
 # ----------------------------------------------------------------------------
 
+import copy
+import functools
+import inspect
 from pathlib import Path
 
 import symforce.symbolic as sf
 from symforce import ops
 from symforce import typing as T
 from symforce.codegen import Codegen
 from symforce.codegen import CppConfig
+from symforce.typing_util import get_type
 
 TYPES = (sf.Rot2, sf.Rot3, sf.V3, sf.Pose2, sf.Pose3)
 
@@ -28,7 +32,8 @@ def get_between_factor_docstring(between_argument_name: str) -> str:
     Args:
         sqrt_info: Square root information matrix to whiten residual. This can be computed from
                    a covariance matrix as the cholesky decomposition of the inverse. In the case
-                   of a diagonal it will contain 1/sigma values. Must match the tangent dim.
+                   of a diagonal it will contain 1/sigma values. Pass in a single scalar if
+                   isotropic, a vector if diagonal, and a square matrix otherwise.
     """.format(
         a_T_b=between_argument_name
     )
@@ -47,7 +52,8 @@ def get_prior_docstring() -> str:
     Args:
         sqrt_info: Square root information matrix to whiten residual. This can be computed from
                    a covariance matrix as the cholesky decomposition of the inverse. In the case
-                   of a diagonal it will contain 1/sigma values. Must match the tangent dim.
+                   of a diagonal it will contain 1/sigma values. Pass in a single scalar if
+                   isotropic, a vector if diagonal, and a square matrix otherwise.
     """
 
 
@@ -83,29 +89,194 @@ def prior_factor(
     return residual
 
 
+def modify_argument(
+    core_func: T.Callable, arg_to_modify: str, new_arg_type: T.Type, modification: T.Callable
+) -> T.Callable:
+    """
+    Returns a wrapper which applies modification to the arg_to_modify parameter before forwarding
+    it and any other arguments on to core_func. Also, sets the type annotation of arg_to_modify
+    to new_arg_type in the returned function.
+
+    If arg_to_modify is not a parameter of core_func, then a Value Error will be raised.
+    """
+    try:
+        arg_index = list(inspect.signature(core_func).parameters).index(arg_to_modify)
+    except ValueError as error:
+        raise ValueError(f"{arg_to_modify} is not an argument of {core_func}") from error
+
+    @functools.wraps(core_func)
+    def wrapper(*args: T.Any, **kwargs: T.Any) -> T.Any:
+        args_list = list(args)
+        if arg_index < len(args):
+            # Then arg_to_modify was passed in args
+            args_list[arg_index] = modification(args[arg_index])
+        else:
+            # arg_to_modify should have been passed in kwargs
+            try:
+                kwargs[arg_to_modify] = modification(kwargs[arg_to_modify])
+            except KeyError as error:
+                raise TypeError(f"{wrapper} missing required argument {arg_to_modify}") from error
+
+        return core_func(*args_list, **kwargs)
+
+    wrapper.__annotations__ = dict(wrapper.__annotations__, **{arg_to_modify: new_arg_type})
+
+    return wrapper
+
+
+def is_not_fixed_size_square_matrix(type_t: T.Type) -> bool:
+    return (
+        not issubclass(type_t, sf.Matrix)
+        or type_t == sf.Matrix
+        or type_t.SHAPE[0] != type_t.SHAPE[1]
+    )
+
+
+def _get_sqrt_info_dim(func: T.Callable) -> int:
+    """
+    Raises ValueError if func does not have a parameter named sqrt_info annotated as a fixed
+    sized square symbolic matrix.
+    Returns the matrix dimension of sqrt_info type annotation of func.
+    """
+    if "sqrt_info" not in func.__annotations__:
+        raise ValueError(
+            "sqrt_info missing annotation. Either add one or explicitly pass in expected number of dimensions"
+        )
+    sqrt_info_type = func.__annotations__["sqrt_info"]
+    if is_not_fixed_size_square_matrix(sqrt_info_type):
+        raise ValueError(
+            f"""Expected sqrt_info to be annotated as a fixed size square matrix. Instead
+            found {sqrt_info_type}. Either fix annotation or explicitly pass in expected number
+            of dimensions of sqrt_info."""
+        )
+    return sqrt_info_type.SHAPE[0]
+
+
+def isotropic_sqrt_info_wrapper(func: T.Callable) -> T.Callable:
+    """
+    Wraps func, except instead of taking a square matrix for the sqrt_info argument, it takes a
+    scalar and passes that scalar times the identity matrix in for the value of sqrt_info.
+
+    Raises ValueError if func does not have an argument named sqrt_info annotated as a fixed size
+    symbolic square matrix.
+    """
+    sqrt_info_dim = _get_sqrt_info_dim(func)
+
+    return modify_argument(
+        func,
+        arg_to_modify="sqrt_info",
+        new_arg_type=T.Scalar,
+        modification=lambda sqrt_info: sqrt_info * sf.M.eye(sqrt_info_dim, sqrt_info_dim),
+    )
+
+
+def diagonal_sqrt_info_wrapper(func: T.Callable) -> T.Callable:
+    """
+    Wraps func, except instead of taking a square matrix for the sqrt_info argument, it takes a
+    vector representing the diagonal and passes the corresponding diagonal matrix in for the
+    value of sqrt_info.
+
+    Raises ValueError if func does not have an argument named sqrt_info annotated as a fixed size
+    symbolic square matrix.
+    """
+    sqrt_info_dim = _get_sqrt_info_dim(func)
+
+    return modify_argument(
+        func,
+        arg_to_modify="sqrt_info",
+        new_arg_type=type(sf.M(sqrt_info_dim, 1)),
+        modification=sf.M.diag,
+    )
+
+
+def override_annotations(func: T.Callable, input_types: T.Sequence[T.ElementOrType]) -> T.Callable:
+    """
+    Returns copy of func which is the same except its parameters are annotated with input_types.
+
+    Raises a ValueError if the length of input_types does not match parameter count of func.
+    """
+    parameters = inspect.signature(func).parameters
+    if len(parameters) != len(input_types):
+        raise ValueError(
+            f"{func} has {len(parameters)} inputs, but input_types has length {len(input_types)}"
+        )
+    new_func = copy.copy(func)
+    new_func.__annotations__ = dict(
+        func.__annotations__, **{param: get_type(tp) for param, tp in zip(parameters, input_types)}
+    )
+    return new_func
+
+
+def generate_with_alternate_sqrt_infos(
+    output_dir: T.Openable,
+    func: T.Callable,
+    name: str,
+    which_args: T.Sequence[str],
+    input_types: T.Sequence[T.ElementOrType] = None,
+    output_names: T.Sequence[str] = None,
+    docstring: str = None,
+) -> None:
+    """
+    Generates func with linearization into output_dir / name, along with two overloads of func:
+    one which takes a single scalar representing an isotropic matrix for the sqrt_info argument,
+    and another which instead takes a vector representing a diagonal matrix for the same argument.
+
+    A common header located at output_dir / name.h will re-export each of these overloads.
+
+    As usual, if func does not have concrete type annotations, then input_types must be passed in
+    (the same as you would if calling Codegen.function direction on func).
+    """
+
+    common_header = Path(output_dir, name + ".h")
+    common_header.parent.mkdir(exist_ok=True, parents=True)
+
+    annotated_func = func if input_types is None else override_annotations(func, input_types)
+
+    for func_variant, variant_name in [
+        (diagonal_sqrt_info_wrapper(annotated_func), f"{name}_diagonal"),
+        (isotropic_sqrt_info_wrapper(annotated_func), f"{name}_isotropic"),
+        (annotated_func, f"{name}_square"),
+    ]:
+        Codegen.function(
+            func=func_variant,
+            output_names=output_names,
+            config=CppConfig(),
+            docstring=docstring,
+        ).with_linearization(name=name, which_args=which_args).generate_function(
+            Path(output_dir, name),
+            skip_directory_nesting=True,
+            generated_file_name=variant_name,
+        )
+
+        with common_header.open("a") as f:
+            f.write(f'#include "./{name}/{variant_name}.h"\n')
+
+
 def generate_between_factors(types: T.Sequence[T.Type], output_dir: T.Openable) -> None:
     """
     Generates between factors for each type in types into output_dir.
     """
     for cls in types:
         tangent_dim = ops.LieGroupOps.tangent_dim(cls)
-        between_codegen = Codegen.function(
+        generate_with_alternate_sqrt_infos(
+            output_dir,
             func=between_factor,
+            name=f"between_factor_{cls.__name__.lower()}",
+            which_args=["a", "b"],
             input_types=[cls, cls, cls, sf.M(tangent_dim, tangent_dim), sf.Symbol],
             output_names=["res"],
-            config=CppConfig(),
             docstring=get_between_factor_docstring("a_T_b"),
-        ).with_linearization(name=f"between_factor_{cls.__name__.lower()}", which_args=["a", "b"])
-        between_codegen.generate_function(output_dir, skip_directory_nesting=True)
+        )
 
-        prior_codegen = Codegen.function(
+        generate_with_alternate_sqrt_infos(
+            output_dir,
             func=prior_factor,
+            name=f"prior_factor_{cls.__name__.lower()}",
+            which_args=["value"],
             input_types=[cls, cls, sf.M(tangent_dim, tangent_dim), sf.Symbol],
             output_names=["res"],
-            config=CppConfig(),
             docstring=get_prior_docstring(),
-        ).with_linearization(name=f"prior_factor_{cls.__name__.lower()}", which_args=["value"])
-        prior_codegen.generate_function(output_dir, skip_directory_nesting=True)
+        )
 
 
 def generate_pose3_extra_factors(output_dir: T.Openable) -> None:
@@ -173,42 +344,46 @@ def prior_factor_pose3_position(
     ) -> sf.Matrix:
         return prior_factor(value.t, prior, sqrt_info, epsilon)
 
-    between_rotation_codegen = Codegen.function(
+    generate_with_alternate_sqrt_infos(
+        output_dir,
         func=between_factor_pose3_rotation,
+        name="between_factor_pose3_rotation",
+        which_args=["a", "b"],
         output_names=["res"],
-        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(
+    generate_with_alternate_sqrt_infos(
+        output_dir,
         func=between_factor_pose3_position,
+        name="between_factor_pose3_position",
+        which_args=["a", "b"],
         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(
+    generate_with_alternate_sqrt_infos(
+        output_dir,
         func=prior_factor_pose3_rotation,
+        name="prior_factor_pose3_rotation",
         output_names=["res"],
-        config=CppConfig(),
+        which_args=["value"],
         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(
+    generate_with_alternate_sqrt_infos(
+        output_dir,
         func=prior_factor_pose3_position,
+        name="prior_factor_pose3_position",
         output_names=["res"],
-        config=CppConfig(),
+        which_args=["value"],
         docstring=get_prior_docstring(),
-    ).with_linearization(name="prior_factor_pose3_position", which_args=["value"])
-    prior_position_codegen.generate_function(output_dir, skip_directory_nesting=True)
+    )
 
 
 def generate(output_dir: Path) -> None: