TYP: remove string types

olivierverdier · Sep 27, 2024 · 843573f · 843573f
1 parent 3805203
commit 843573f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/diffeopt/cometric/laplace.py b/diffeopt/cometric/laplace.py
@@ -1,7 +1,7 @@
-import numpy as np
 import torch
+from ..group.base import BaseDiffeoGroup
 
-def get_fourier_cometric(group: "BaseDiffeoGroup", s: int):
+def get_fourier_cometric(group: BaseDiffeoGroup, s: int):
     shape = group.shape
     idx, idy = group.get_raw_identity()
     lap = 4. - 2.*(torch.cos(2.*torch.pi*idx/shape[0]) + torch.cos(2.*torch.pi*idy/shape[1]))
@@ -11,7 +11,7 @@ def get_fourier_cometric(group: "BaseDiffeoGroup", s: int):
     lapinv[0,0] = 1.
     return lapinv
 
-def get_laplace_cometric(group: "BaseDiffeoGroup", s:int=1):
+def get_laplace_cometric(group: BaseDiffeoGroup, s:int=1):
     lapinv = get_fourier_cometric(group, s)
     def cometric(momentum: torch.Tensor) -> torch.Tensor:
         fx = torch.fft.fftn(momentum[0])

diff --git a/diffeopt/utils.py b/diffeopt/utils.py
@@ -1,5 +1,6 @@
 import numpy as np
 import torch
+from .group.base import BaseDiffeoGroup
 
 
 def get_volume(shape):
@@ -18,7 +19,7 @@ def normalize(I):
 
 from diffeopt.cometric import laplace
 
-def get_random_diffeo(group: "BaseDiffeoGroup", nb_steps:int=10, scale:float=1.) -> torch.Tensor:
+def get_random_diffeo(group: BaseDiffeoGroup, nb_steps:int=10, scale:float=1.) -> torch.Tensor:
     cometric = laplace.get_laplace_cometric(group, s=2)
     rm = torch.randn(*group.zero().shape)
     rv = cometric(rm)