Added utility to automatically check jacobians of a given cost function.

[luisenp]

Closes #321.

Not integrated with TheseusLayer yet, but I think we can set this up so that the check is (optionally) performed once for each cost function right after the objective's version is fixed.

cc @xphsean12 as related to #461

[xphsean12]

I copy this function and run it in my local Jupyter, with 0.1.4 version of theseus. An error occurs in Line 153, showing "property' object is not iterable"

Besides, should Line 154 be aux_var = list(cf.aux_vars)?

[luisenp]

I copy this function and run it in my local Jupyter, with 0.1.4 version of theseus. An error occurs in Line 153, showing "property' object is not iterable"

Besides, should Line 154 be aux_var = list(cf.aux_vars)?

Good catch, fixed this in 1a0047b. Can you post a screen shot of the error you are getting? I wrote a unit test for this code and it seems to be working will with our existing cost functions.

[xphsean12]

I copy this function and run it in my local Jupyter, with 0.1.4 version of theseus. An error occurs in Line 153, showing "property' object is not iterable"
Besides, should Line 154 be aux_var = list(cf.aux_vars)?

Good catch, fixed this in 1a0047b. Can you post a screen shot of the error you are getting? I wrote a unit test for this code and it seems to be working will with our existing cost functions.

class toyCost(CostFunction):
    def __init__(
        self,
        Pos_x: Vector,
        Pos_y: Vector,

        Range: Vector, 
        BasePos_x: Vector,
        BasePos_y: Vector,


        cost_weight: CostWeight,
        
        name: Optional[str] = None,
    ):
        super().__init__(cost_weight, name=name)
     

        self.Pos_x = Pos_x
        self.Pos_y = Pos_y
        
        self.Range = as_variable(Range)
        self.BasePos_x = as_variable(BasePos_x)
        self.BasePos_y = as_variable(BasePos_y)

        self.register_optim_vars(["Pos_x", "Pos_y"])
       
        self.register_aux_vars(["Range", "BasePos_x", "BasePos_y"])
        self.weight = cost_weight


    # Add this function for autochecking Jocobian. What you need to do is to warp the inputs to two types: optim_vars and aux_vars    
    def err_func_numeric(self, optim_vars, aux_vars):
        Pos_x,Pos_y = optim_vars
        Pos_x = Pos_x.tensor
        Pos_y = Pos_y.tensor
        
        Range, BasePos_x, BasePos_y = aux_vars
        Range = Range.tensor
        BasePos_x = BasePos_x.tensor
        BasePos_y = BasePos_y.tensor
        
        # Copy it from the self-defined cost function _err_func
        distance_0 = ((Pos_x - BasePos_x).square()+ (Pos_y - BasePos_y).square()).sqrt()


        err = distance_0 - Range

        return err
    
    
    def _err_func(self):
        
        Pos_x = self.Pos_x.tensor
        Pos_y = self.Pos_y.tensor
        Range = self.Range.tensor
        BasePos_x = self.BasePos_x.tensor
        BasePos_y = self.BasePos_y.tensor
      
        distance_0 = ((Pos_x - BasePos_x).square()+ (Pos_y - BasePos_y).square()).sqrt()


        err = distance_0 - Range
        print(err.shape)
       
        #print(err_0.shape)
        return err
    
    
    def dim(self):
        return BasePos_x.shape[1]

    def error(self) -> torch.Tensor:
        return self._err_func()

    def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
        
        batch_size = self.Range.shape[0]
        
        
        dtype = self.Range.dtype
        device = self.Range.device
        

        
        Pos_x = self.Pos_x.tensor
        Pos_y = self.Pos_y.tensor
        

        Range = self.Range.tensor
        BasePos_x = self.BasePos_x.tensor
        BasePos_y = self.BasePos_y.tensor
        dof = BasePos_x.shape[1]
        
        interTerm = (Pos_x - BasePos_x).square()+ (Pos_y - BasePos_y).square()
        
        
        Jacob_Pos_x = torch.zeros(batch_size, dof, 1, dtype=dtype, device=device)
        Jacob_Pos_y = torch.zeros(batch_size, dof, 1, dtype=dtype, device=device)
        
        dErr_x = torch.pow(interTerm, -0.5) * (Pos_x - BasePos_x)

        dErr_y = torch.pow(interTerm, -0.5) * (Pos_y - BasePos_y)
       
        Jacob_Pos_x[:,:,0] = dErr_x
        Jacob_Pos_y[:,:,0] = dErr_y
        
        error = self.error()
        
        return [Jacob_Pos_x, Jacob_Pos_y], error

    def _copy_impl(self, new_name: Optional[str] = None) -> "toyCost":
        return toyCost(
            self.Pos_x.copy(),
            self.Pos_y.copy(),
            self.Range.copy(),
            self.BasePos_x.copy(),
            self.BasePos_y.copy(),
            self.weight.copy(),
            name=new_name,
        )

[luisenp]

I copy this function and run it in my local Jupyter, with 0.1.4 version of theseus. An error occurs in Line 153, showing "property' object is not iterable"
Besides, should Line 154 be aux_var = list(cf.aux_vars)?

Good catch, fixed this in 1a0047b. Can you post a screen shot of the error you are getting? I wrote a unit test for this code and it seems to be working will with our existing cost functions.

class toyCost(CostFunction):
    def __init__(
        self,
        Pos_x: Vector,
        Pos_y: Vector,

        Range: Vector, 
        BasePos_x: Vector,
        BasePos_y: Vector,


        cost_weight: CostWeight,
        
        name: Optional[str] = None,
    ):
        super().__init__(cost_weight, name=name)
     

        self.Pos_x = Pos_x
        self.Pos_y = Pos_y
        
        self.Range = as_variable(Range)
        self.BasePos_x = as_variable(BasePos_x)
        self.BasePos_y = as_variable(BasePos_y)

        self.register_optim_vars(["Pos_x", "Pos_y"])
       
        self.register_aux_vars(["Range", "BasePos_x", "BasePos_y"])
        self.weight = cost_weight


    # Add this function for autochecking Jocobian. What you need to do is to warp the inputs to two types: optim_vars and aux_vars    
    def err_func_numeric(self, optim_vars, aux_vars):
        Pos_x,Pos_y = optim_vars
        Pos_x = Pos_x.tensor
        Pos_y = Pos_y.tensor
        
        Range, BasePos_x, BasePos_y = aux_vars
        Range = Range.tensor
        BasePos_x = BasePos_x.tensor
        BasePos_y = BasePos_y.tensor
        
        # Copy it from the self-defined cost function _err_func
        distance_0 = ((Pos_x - BasePos_x).square()+ (Pos_y - BasePos_y).square()).sqrt()


        err = distance_0 - Range

        return err
    
    
    def _err_func(self):
        
        Pos_x = self.Pos_x.tensor
        Pos_y = self.Pos_y.tensor
        Range = self.Range.tensor
        BasePos_x = self.BasePos_x.tensor
        BasePos_y = self.BasePos_y.tensor
      
        distance_0 = ((Pos_x - BasePos_x).square()+ (Pos_y - BasePos_y).square()).sqrt()


        err = distance_0 - Range
        print(err.shape)
       
        #print(err_0.shape)
        return err
    
    
    def dim(self):
        return BasePos_x.shape[1]

    def error(self) -> torch.Tensor:
        return self._err_func()

    def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
        
        batch_size = self.Range.shape[0]
        
        
        dtype = self.Range.dtype
        device = self.Range.device
        

        
        Pos_x = self.Pos_x.tensor
        Pos_y = self.Pos_y.tensor
        

        Range = self.Range.tensor
        BasePos_x = self.BasePos_x.tensor
        BasePos_y = self.BasePos_y.tensor
        dof = BasePos_x.shape[1]
        
        interTerm = (Pos_x - BasePos_x).square()+ (Pos_y - BasePos_y).square()
        
        
        Jacob_Pos_x = torch.zeros(batch_size, dof, 1, dtype=dtype, device=device)
        Jacob_Pos_y = torch.zeros(batch_size, dof, 1, dtype=dtype, device=device)
        
        dErr_x = torch.pow(interTerm, -0.5) * (Pos_x - BasePos_x)

        dErr_y = torch.pow(interTerm, -0.5) * (Pos_y - BasePos_y)
       
        Jacob_Pos_x[:,:,0] = dErr_x
        Jacob_Pos_y[:,:,0] = dErr_y
        
        error = self.error()
        
        return [Jacob_Pos_x, Jacob_Pos_y], error

    def _copy_impl(self, new_name: Optional[str] = None) -> "toyCost":
        return toyCost(
            self.Pos_x.copy(),
            self.Pos_y.copy(),
            self.Range.copy(),
            self.BasePos_x.copy(),
            self.BasePos_y.copy(),
            self.weight.copy(),
            name=new_name,
        )

Ah, I see. The error is that you are directly passing the class, but this method requires passing a class instance. I do think passing the class would be more intuitive, but it will unfortunately make the code much more complicated, because there is currently no easy way to get the information about the variable types just from the class.

I thought for a bit about adding this, but I think it would this PR more large and complicated, so decided to keep as is for now.

[mhmukadam]

LGTM! If we want to add improvements like you mention then we can flag an issue to switch to passing class name instead of an instance.

I am not sure we should enable some check for the full objective since it would slow down the construction with unnecessary checks especially if there are many cost functions of the same type in the objective or check cost functions in the library that are already verified in unit tests. Could we assume a user would just use this function adhoc when they need it for their custom cost functions?

[luisenp]

LGTM! If we want to add improvements like you mention then we can flag an issue to switch to passing class name instead of an instance.

I am not sure we should enable some check for the full objective since it would slow down the construction with unnecessary checks especially if there are many cost functions of the same type in the objective or check cost functions in the library that are already verified in unit tests. Could we assume a user would just use this function adhoc when they need it for their custom cost functions?

I'll at least add a method one can run on the full objective on-demand, even if it's slow. I think this would be the most convenient API for this feature.

[mhmukadam]

LGTM! On demand checks is a good compromise than always silently running the checks under the hood.