fix some typing hints

Author: scaomath

torch_cfd/equations.py

@@ -450,7 +450,7 @@ def implicit_solve(self, vort_hat, dt):
     def step(self, *args, **kwargs):
         return self.forward(*args, **kwargs)
 
-    def forward(self, vort_hat, dt, steps=1):
+    def forward(self, vort_hat, dt, steps=1) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         vort_hat: (B, kx, ky) or (n_t, kx, ky) or (kx, ky)
         - if rfft2 is used then the shape is (*, nx, ny//2+1)

torch_cfd/grids.py

@@ -285,7 +285,7 @@ def clone(self, *args, **kwargs):
     def to(self, *args, **kwargs):
         return GridArray(self.data.to(*args, **kwargs), self.offset, self.grid)
 
-    _HANDLED_TYPES = (numbers.Number, Array)
+    _HANDLED_TYPES = (numbers.Number, torch.Tensor)
 
     @classmethod
     def __torch_function__(self, ufunc, types, args=(), kwargs=None):
@@ -483,7 +483,7 @@ def enforce_edge_bc(self, *args) -> GridVariable:
 GridVariableVector = Tuple[GridVariable, ...]
 
 
-class GridArrayTensor(Array):
+class GridArrayTensor(torch.Tensor):
     """A numpy array of GridArrays, representing a physical tensor field.
 
     Packing tensor coordinates into a numpy array of dtype object is useful

torch_cfd/initial_conditions.py

@@ -93,7 +93,7 @@ def spectral_filter(
     v: torch.Tensor,
     grid: grids.Grid,
 ) -> torch.Tensor:
-    """Filter an Array with white noise to match a prescribed spectral density."""
+    """Filter a torch.Tensor with white noise to match a prescribed spectral density."""
     k = _angular_frequency_magnitude(grid)
     filters = torch.where(k > 0, spectral_density(k), 0.0)
     # The output signal can safely be assumed to be real if our input signal was