change input of NS2dSpectral to initialized class

Author: scaomath

fno/data_gen/data_gen_Kolmogorov2d.py

@@ -129,7 +129,7 @@ def main(args):
         drag=0.1,
         smooth=True,
         forcing_fn=forcing_fn,
-        solver=RK4CrankNicolsonStepper,
+        solver=RK4CrankNicolsonStepper(),
     ).to(device)
 
     num_batches = total_samples // batch_size

fno/data_gen/data_gen_McWilliams2d.py

@@ -113,7 +113,7 @@ def main(args):
         drag=0,
         smooth=True,
         forcing_fn=None,
-        solver=RK4CrankNicolsonStepper,
+        solver=RK4CrankNicolsonStepper(),
     ).to(device)
 
     num_batches = total_samples // batch_size

fno/data_gen/data_gen_fno.py

@@ -168,14 +168,12 @@ def main(args):
         device=device,
         dtype=torch.float64,
     )
-
+    step_fn = IMEXStepper(order=2)
     ns2d = NavierStokes2DSpectral(
         viscosity=visc,
         grid=grid,
         smooth=True,
         forcing_fn=forcing_fn,
-        solver=IMEXStepper,
-        order=2,
     ).to(device)
 
     if os.path.exists(data_filepath) and not force_rerun:

torch_cfd/equations.py

@@ -379,17 +379,16 @@ def __init__(
         drag: float = 0.0,
         smooth: bool = True,
         forcing_fn: Optional[Callable] = None,
-        solver: Optional[Callable] = RK4CrankNicolsonStepper,
-        requires_grad: bool = False,
-        **solver_kwargs,
+        solver: IMEXStepper = None,
+        **kwargs,
     ):
         super().__init__()
         self.viscosity = viscosity
         self.grid = grid
         self.drag = drag
         self.smooth = smooth
         self.forcing_fn = forcing_fn
-        self.solver = solver(requires_grad=requires_grad, **solver_kwargs)
+        self.solver = solver
         self._initialize()
 
     def _initialize(self):