diff --git a/theseus/geometry/lie_group.py b/theseus/geometry/lie_group.py old mode 100644 new mode 100755 index 99caecf9..0ca2fede --- a/theseus/geometry/lie_group.py +++ b/theseus/geometry/lie_group.py @@ -4,7 +4,7 @@ # LICENSE file in the root directory of this source tree. import abc -from typing import Any, List, Optional, Tuple, cast +from typing import Any, List, Optional, Tuple, Union, cast import torch @@ -59,6 +59,7 @@ def dof(self) -> int: def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, @@ -70,6 +71,7 @@ def rand( def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, diff --git a/theseus/geometry/point_types.py b/theseus/geometry/point_types.py old mode 100644 new mode 100755 index f6bf6905..d14b6f7a --- a/theseus/geometry/point_types.py +++ b/theseus/geometry/point_types.py @@ -42,12 +42,17 @@ def __init__( def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, ) -> "Point2": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (2,): + raise ValueError("The scale must be None, a float, or a 2 element 1D tensor.") + elif scale is None: + scale = 1.0 return Point2( tensor=torch.rand( size[0], @@ -56,19 +61,24 @@ def rand( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale ) @staticmethod def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, ) -> "Point2": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (2,): + raise ValueError("The scale must be None, a float, or a 2 element 1D tensor.") + elif scale is None: + scale = 1.0 return Point2( tensor=torch.randn( size[0], @@ -77,7 +87,7 @@ def randn( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale ) @staticmethod @@ -144,12 +154,17 @@ def __init__( def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, ) -> "Point3": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (3,): + raise ValueError("The scale must be None, a float, or a 3D vector.") + elif scale is None: + scale = 1.0 return Point3( tensor=torch.rand( size[0], @@ -158,19 +173,24 @@ def rand( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale ) @staticmethod def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, ) -> "Point3": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (3,): + raise ValueError("The scale must be None, a float, or a 3D vector.") + elif scale is None: + scale = 1.0 return Point3( tensor=torch.randn( size[0], @@ -179,7 +199,7 @@ def randn( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale ) @staticmethod diff --git a/theseus/geometry/se2.py b/theseus/geometry/se2.py old mode 100644 new mode 100755 index 8b26a856..53de7ef5 --- a/theseus/geometry/se2.py +++ b/theseus/geometry/se2.py @@ -46,12 +46,17 @@ def __init__( def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SE2": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (3,): + raise ValueError("The scale must be None, a float, or a 3 element 1D tensor.") + elif scale is None: + scale = 1.0 x_y_theta = torch.rand( size[0], 3, @@ -59,7 +64,7 @@ def rand( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale x_y_theta[:, 2] = 2 * theseus.constants.PI * (x_y_theta[:, 2] - 0.5) return SE2(x_y_theta=x_y_theta) @@ -68,12 +73,17 @@ def rand( def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SE2": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (3,): + raise ValueError("The scale must be None, a float, or a 3 element 1D tensor.") + elif scale is None: + scale = 1.0 x_y_theta = torch.randn( size[0], 3, @@ -81,7 +91,7 @@ def randn( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale x_y_theta[:, 2] *= theseus.constants.PI return SE2(x_y_theta=x_y_theta) diff --git a/theseus/geometry/se3.py b/theseus/geometry/se3.py old mode 100644 new mode 100755 index 0074d30a..b690576f --- a/theseus/geometry/se3.py +++ b/theseus/geometry/se3.py @@ -45,15 +45,21 @@ def __init__( def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SE3": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (4,): + raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.") + elif scale is None: + scale = 1.0 tensor = SE3_base.rand( *size, generator=generator, + scale=scale, dtype=dtype, device=device, requires_grad=requires_grad, @@ -64,15 +70,21 @@ def rand( def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SE3": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (4,): + raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.") + elif scale is None: + scale = 1.0 tensor = SE3_base.randn( *size, generator=generator, + scale=scale, dtype=dtype, device=device, requires_grad=requires_grad, diff --git a/theseus/geometry/so2.py b/theseus/geometry/so2.py old mode 100644 new mode 100755 index 67c1def7..239dc974 --- a/theseus/geometry/so2.py +++ b/theseus/geometry/so2.py @@ -50,12 +50,17 @@ def __init__( def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SO2": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (1,): + raise ValueError("The scale must be None, or a float, or a single element tensor.") + elif scale is None: + scale = 1.0 return SO2.exp_map( 2 * theseus.constants.PI @@ -66,7 +71,7 @@ def rand( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale - theseus.constants.PI ) @@ -74,12 +79,17 @@ def rand( def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SO2": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (1,): + raise ValueError("The scale must be None, or a float, or a single element tensor.") + elif scale is None: + scale = 1.0 return SO2.exp_map( theseus.constants.PI * torch.randn( @@ -89,7 +99,7 @@ def randn( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale ) @staticmethod diff --git a/theseus/geometry/so3.py b/theseus/geometry/so3.py old mode 100644 new mode 100755 index 4f71bff7..d927dfbf --- a/theseus/geometry/so3.py +++ b/theseus/geometry/so3.py @@ -45,15 +45,21 @@ def __init__( def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SO3": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (4,): + raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.") + elif scale is None: + scale = 1.0 tensor = SO3_base.rand( *size, generator=generator, + scale=scale, dtype=dtype, device=device, requires_grad=requires_grad, @@ -64,15 +70,21 @@ def rand( def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: theseus.constants.DeviceType = None, requires_grad: bool = False, ) -> "SO3": if len(size) != 1: raise ValueError("The size should be 1D.") + if isinstance(scale, torch.Tensor) and scale.shape != (4,): + raise ValueError("The scale must be None, or a float, or a 4 element 1D tensor.") + elif scale is None: + scale = 1.0 tensor = SO3_base.randn( *size, generator=generator, + scale=scale, dtype=dtype, device=device, requires_grad=requires_grad, diff --git a/theseus/geometry/vector.py b/theseus/geometry/vector.py old mode 100644 new mode 100755 index 05c7a069..580ea18d --- a/theseus/geometry/vector.py +++ b/theseus/geometry/vector.py @@ -46,12 +46,17 @@ def dof(self) -> int: def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, ) -> "Vector": if len(size) != 2: raise ValueError("The size should be 2D.") + if isinstance(scale, torch.Tensor) and (scale.shape[0] != size[0] or scale.shape[1] != size[1]): + raise ValueError(f"The scale must be None, or a float, or a tensor of shape {size}.") + elif scale is None: + scale = 1.0 return Vector( tensor=torch.rand( size, @@ -59,19 +64,24 @@ def rand( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale ) @staticmethod def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: DeviceType = None, requires_grad: bool = False, ) -> "Vector": if len(size) != 2: raise ValueError("The size should be 2D.") + if isinstance(scale, torch.Tensor) and (scale.shape[0] != size[0] or scale.shape[1] != size[1]): + raise ValueError(f"The scale must be None, or a float, or a tensor of shape {size}.") + elif scale is None: + scale = 1.0 return Vector( tensor=torch.randn( size, @@ -79,7 +89,7 @@ def randn( dtype=dtype, device=device, requires_grad=requires_grad, - ) + ) * scale ) def __repr__(self) -> str: diff --git a/torchlie/torchlie/functional/se3_impl.py b/torchlie/torchlie/functional/se3_impl.py old mode 100644 new mode 100755 index 7c51ac8f..5200cda2 --- a/torchlie/torchlie/functional/se3_impl.py +++ b/torchlie/torchlie/functional/se3_impl.py @@ -3,7 +3,7 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import List, Optional, Tuple, cast +from typing import List, Optional, Tuple, Union, cast import torch @@ -121,14 +121,17 @@ def get_transform_tensor_size(transform_tensor: torch.Tensor): def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: constants.DeviceType = None, requires_grad: bool = False, ) -> torch.Tensor: - rotation = SO3.rand(*size, generator=generator, dtype=dtype, device=device) + if scale is None: + scale = 1.0 + rotation = SO3.rand(*size, generator=generator, scale=scale, dtype=dtype, device=device) translation = torch.rand( *size, 3, 1, generator=generator, dtype=dtype, device=device - ) + ) * scale if not isinstance(scale, torch.Tensor) else scale[:-1] ret = torch.cat((rotation, translation), dim=-1) ret.requires_grad_(requires_grad) return ret @@ -140,14 +143,17 @@ def rand( def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: constants.DeviceType = None, requires_grad: bool = False, ) -> torch.Tensor: - rotation = SO3.randn(*size, generator=generator, dtype=dtype, device=device) + if scale is None: + scale = 1.0 + rotation = SO3.randn(*size, generator=generator, scale=scale, dtype=dtype, device=device) translation = torch.randn( *size, 3, 1, generator=generator, dtype=dtype, device=device - ) + ) * scale if not isinstance(scale, torch.Tensor) else scale[:-1] ret = torch.cat((rotation, translation), dim=-1) ret.requires_grad_(requires_grad) return ret diff --git a/torchlie/torchlie/functional/so3_impl.py b/torchlie/torchlie/functional/so3_impl.py old mode 100644 new mode 100755 index 538af863..7cceb4be --- a/torchlie/torchlie/functional/so3_impl.py +++ b/torchlie/torchlie/functional/so3_impl.py @@ -3,7 +3,7 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import Dict, List, Optional, Tuple, cast +from typing import Dict, List, Optional, Tuple, Union, cast import torch @@ -155,13 +155,16 @@ def get_transform_tensor_size(transform_tensor: torch.Tensor): def rand( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: constants.DeviceType = None, requires_grad: bool = False, ) -> torch.Tensor: # Reference: # https://web.archive.org/web/20211105205926/http://planning.cs.uiuc.edu/node198.html - u = torch.rand(3, *size, generator=generator, dtype=dtype, device=device) + if scale is None: + scale = 1.0 + u = torch.rand(3, *size, generator=generator, dtype=dtype, device=device) * scale if not isinstance(scale, torch.Tensor) else scale[-1] u1 = u[0] u2, u3 = u[1:3] * 2 * constants.PI @@ -188,13 +191,17 @@ def rand( def randn( *size: int, generator: Optional[torch.Generator] = None, + scale: Optional[Union[float, torch.Tensor]] = None, dtype: Optional[torch.dtype] = None, device: constants.DeviceType = None, requires_grad: bool = False, ) -> torch.Tensor: + if scale is None: + scale = 1.0 ret = _exp_autograd_fn( constants.PI * torch.randn(*size, 3, generator=generator, dtype=dtype, device=device) + * scale if not isinstance(scale, torch.Tensor) else scale[-1] ) ret.requires_grad_(requires_grad) return ret