Merge pull request #96 from MagriLab/feature/solver_protocol

Provided Solver Protocol

.github/workflows/build_and_publish.yaml

@@ -18,7 +18,7 @@ jobs:
  - name: Install Python
  uses: actions/[email protected]
  with:
- python-version: '3.6.8'
+ python-version: '3.8.10'
 
  - name: Upgrade pip / Install poetry
  run: |

.github/workflows/linting.yml

@@ -24,7 +24,7 @@ jobs:
  - name: Install Python 3
  uses: actions/[email protected]
  with:
- python-version: '3.6.8'
+ python-version: '3.8.10'
 
  - name: Upgrade Pip / Install Poetry
  run: |

kolsol/base/base_solver.py

@@ -1,9 +1,18 @@
-from typing import List, Optional
+from typing import Protocol
+
+import numpy as np
+import torch
 
 from ..utils.types import TypeTensor
 
 
-class BaseKolSol:
+class Solver(Protocol[TypeTensor]):
+
+ def dynamics(self, u_hat: TypeTensor) -> TypeTensor:
+ ...
+
+
+class BaseNumpyKolSol:
 
  def __init__(self, nk: int, nf: int, re: float, ndim: int = 2) -> None:
 
@@ -15,62 +24,39 @@ def __init__(self, nk: int, nf: int, re: float, ndim: int = 2) -> None:
  self.nk_grid = 2 * self.nk + 1
  self.mk_grid = 4 * self.nk + 1
 
- def dynamics(self, u_hat: TypeTensor) -> TypeTensor:
+ def dynamics(self, u_hat: np.ndarray) -> np.ndarray:
 
  r"""Computes \partial u_hat / \partial t
 
  Parameters
  ----------
- u_hat: torch.Tensor
+ u_hat: np.ndarray
  Velocity field in the Fourier domain.
 
  Returns
  -------
- TypeTensor
+ np.ndarray
  Computed time-derivative.
  """
 
- raise NotImplementedError('BaseKolSol::dynamics()')
-
- def aap(self, f1: TypeTensor, f2: TypeTensor) -> TypeTensor:
-
- """Computes anti-aliased product between two tensors.
-
- Parameters
- ----------
- f1: torch.Tensor
- Tensor one.
- f2: torch.Tensor
- Tensor two.
-
- Returns
- -------
- TypeTensor
- Anti-aliased product of tensor one, tensor two.
- """
+ raise NotImplementedError('BaseNumpyKolSol::dynamics()')
 
- raise NotImplementedError('BaseKolSol::aap()')
 
- def dissip(self, u_hat: TypeTensor) -> TypeTensor:
+class BaseTorchKolSol:
 
- """Computes dissipation from a given flow-field.
-
- Parameters
- ----------
- u_hat: torch.Tensor
- Velocity field in the Fourier domain.
+ def __init__(self, nk: int, nf: int, re: float, ndim: int = 2) -> None:
 
- Returns
- -------
- TypeTensor
- Computed dissipation.
- """
+ self.nk = nk
+ self.nf = nf
+ self.re = re
+ self.ndim = ndim
 
- raise NotImplementedError('BaseKolSol::dissip()')
+ self.nk_grid = 2 * self.nk + 1
+ self.mk_grid = 4 * self.nk + 1
 
- def vorticity(self, u_hat: TypeTensor) -> TypeTensor:
+ def dynamics(self, u_hat: torch.Tensor) -> torch.Tensor:
 
- """Computes vorticity of the given velocity field.
+ r"""Computes \partial u_hat / \partial t
 
  Parameters
  ----------
@@ -79,63 +65,9 @@ def vorticity(self, u_hat: TypeTensor) -> TypeTensor:
 
  Returns
  -------
- TypeTensor
- Computed vorticity field in the Fourier domain.
- """
-
- raise NotImplementedError('BaseKolSol::vorticity()')
-
- def fourier_to_phys(self, t_hat: TypeTensor) -> TypeTensor:
-
- """Functional mapping from Fourier domain to physical domain.
-
- Parameters
- ----------
- t_hat: torch.Tensor
- Tensor in the Fourier domain to map to the physical domain.
-
- Returns
- -------
- TypeTensor
- Tensor in the physical domain.
- """
-
- raise NotImplementedError('BaseKolSol::fourier_to_phys()')
-
- def phys_to_fourier(self, t: TypeTensor) -> TypeTensor:
-
- """Functional mapping from the physical domain to the Fourier domain.
-
- Parameters
- ----------
- t: torch.Tensor
- Tensor in the physical domain to map to the Fourier domain.
-
- Returns
- -------
- TypeTensor
- Tensor in the Fourier domain.
+ torch.Tensor
+ Computed time-derivative.
  """
 
- raise NotImplementedError('BaseKolSol::phys_to_fourier()')
-
- def random_field(self, magnitude: float, sigma: float, k_offset: Optional[List[int]] = None) -> TypeTensor:
-
- """Generate random field in the Fourier domain based on k distribution.
-
- Parameters
- ----------
- magnitude: float
- Magnitude to use to compute distribution.
- sigma: float
- Sigma to use to compute distribution.
- k_offset: Optional[List[int]]
- Magnitude of wavenumber offset in each dimension.
-
- Returns
- -------
- TypeTensor
- Generated random field in the Fourier domain.
- """
+ raise NotImplementedError('BaseTorchKolSol::dynamics()')
 
- raise NotImplementedError('BaseKolSol::random_field()')

kolsol/numpy/solver.py

@@ -6,11 +6,11 @@
 import numpy as np
 import opt_einsum as oe
 
-from ..base.base_solver import BaseKolSol
+from ..base.base_solver import BaseNumpyKolSol
 from ..utils.enums import eDirection
 
 
-class KolSol(BaseKolSol):
+class KolSol(BaseNumpyKolSol):
 
  def __init__(self, nk: int, nf: int, re: float, ndim: int = 2) -> None:
 
@@ -281,7 +281,7 @@ def phys_to_fourier(self, t: np.ndarray) -> np.ndarray:
 
  return t_hat
 
- def random_field(self, magnitude: float, sigma: float, k_offset: Optional[List[int]] = None) -> np.array:
+ def random_field(self, magnitude: float, sigma: float, k_offset: Optional[List[int]] = None) -> np.ndarray:
 
  """Generate random velocity field in the Fourier domain.
 

kolsol/torch/solver.py

@@ -7,11 +7,11 @@
 import opt_einsum as oe
 import torch
 
-from ..base.base_solver import BaseKolSol
+from ..base.base_solver import BaseTorchKolSol
 from ..utils.enums import eDirection
 
 
-class KolSol(BaseKolSol):
+class KolSol(BaseTorchKolSol):
 
  def __init__(self,
  nk: int,

kolsol/utils/types.py

@@ -1,6 +1,7 @@
-from typing import Union
+from typing import TypeVar
 
 import numpy as np
 import torch
 
-TypeTensor = Union[np.array, torch.Tensor]
+
+TypeTensor = TypeVar('TypeTensor', np.ndarray, torch.Tensor)