Refactor equations' __call__ to nn.Module

Author: scaomath

fno/data_gen/data_gen_Kolmogorov2d.py

@@ -51,22 +51,25 @@ def main(args):
     log_filename = os.path.join(LOG_PATH, f"{current_time}_{log_name}.log")
     logger = get_logger(log_filename)
 
+    logger.info(f"Using the following arguments: ")
+    all_args = {k: v for k, v in vars(args).items() if not callable(v)}
+    logger.info(" | ".join(f"{k}={v}" for k, v in all_args.items()))
+
     total_samples = args.num_samples
     batch_size = args.batch_size  # 128
     n = args.grid_size  # 256
     scale = args.scale
     viscosity = args.visc
     dt = args.dt  # 1e-3
     T = args.time  # 10
-    subsample = args.subsample  # 4
-    ns = n // subsample
     T_warmup = args.time_warmup  # 4.5
     num_snapshots = args.num_steps  # 100
+    subsample = args.subsample  # 4
+    ns = n // subsample
     random_state = args.seed
     peak_wavenumber = args.peak_wavenumber  # 4
-    diam = (
-        eval(args.diam) if isinstance(args.diam, str) else args.diam
-    )  # "2 * torch.pi"
+    diam = args.diam  # "2 * torch.pi" default
+    diam = eval(diam) if isinstance(diam, str) else diam  # 
     force_rerun = args.force_rerun
 
     logger = logging.getLogger()
@@ -89,21 +92,29 @@ def main(args):
         )
         args.filename = filename
     data_filepath = os.path.join(DATA_PATH, filename)
-    if os.path.exists(data_filepath) and not force_rerun:
-        logger.info(f"Data already exists at {data_filepath}")
-        return
-    elif os.path.exists(data_filepath) and force_rerun:
-        logger.info(f"Force rerun and save data to {data_filepath}")
-        os.remove(data_filepath)
+    data_exist = os.path.exists(data_filepath)
+    if data_exist and not force_rerun:
+        logger.info(f"File {filename} exists with current data as follows:")
+        data = torch.load(data_filepath)
+        
+        for key, v in data.items():
+            if isinstance(v, torch.Tensor):
+                logger.info(f"{key:<12} | {v.shape} | {v.dtype}")
+            else:
+                logger.info(f"{key:<12} | {v.dtype}")
+        if len(data[key]) == total_samples:
+            return
+        elif len(data[key]) < total_samples:
+            total_samples -= len(data[key])
     else:
-        logger.info(f"Save data to {data_filepath}")
+        logger.info(f"Generating data and saving in {filename}")
 
     cuda = not args.no_cuda and torch.cuda.is_available()
     no_tqdm = args.no_tqdm
     device = torch.device("cuda:0" if cuda else "cpu")
 
-    torch.set_default_dtype(dtype)
-    logger.info(f"Using device: {device} | dtype: {dtype}")
+    torch.set_default_dtype(torch.float64)
+    logger.info(f"Using device: {device} | save dtype: {dtype} | computge dtype: {torch.get_default_dtype()}")
 
     grid = Grid(shape=(n, n), domain=((0, diam), (0, diam)), device=device)
 
@@ -120,7 +131,7 @@ def main(args):
         drag=0.1,
         smooth=True,
         forcing_fn=forcing_fn,
-        solver=RK4CrankNicholson,
+        solver=RK4CrankNicolsonStepper,
     ).to(device)
 
     num_batches = total_samples // batch_size
@@ -156,8 +167,8 @@ def main(args):
                     )
                     pbar.set_description(desc)
                     pbar.update(100)
-
-        result = get_trajectory_rk4(
+        logger.info(f"generate data from {T_warmup} to {T}")
+        result = get_trajectory_imex(
             ns2d,
             vort_hat,
             dt,
@@ -172,9 +183,9 @@ def main(args):
                 f"variable: {field} | shape: {value.shape} | dtype: {value.dtype}"
             )
             if subsample > 1:
-                result[field] = F.interpolate(value, size=(ns, ns), mode="bilinear")
-            else:
-                result[field] = value
+                assert value.ndim == 4, f"Subsampling only works for (b, c, h, w) tensors, current shape: {value.shape}"
+                value = F.interpolate(value, size=(ns, ns), mode="bilinear")
+            result[field] = value
 
         result["random_states"] = torch.tensor(
             [random_state + idx + k for k in range(batch_size)], dtype=torch.int32
@@ -198,5 +209,5 @@ def main(args):
 
 
 if __name__ == "__main__":
-    args = get_args("Params Kolmogorov 2d flow data generation")
+    args = get_args_ns2d("Params Kolmogorov 2d flow data generation")
     main(args)

fno/data_gen/data_gen_McWilliams2d.py

@@ -117,7 +117,7 @@ def main(args):
         drag=0,
         smooth=True,
         forcing_fn=None,
-        solver=RK4CrankNicholson,
+        solver=RK4CrankNicolsonStepper,
     ).to(device)
 
     num_batches = total_samples // batch_size
@@ -145,7 +145,7 @@ def main(args):
                     pbar.set_description(desc)
                     pbar.update(100)
 
-        result = get_trajectory_rk4(
+        result = get_trajectory_imex(
             ns2d,
             vort_hat,
             dt,
@@ -191,5 +191,5 @@ def main(args):
 
 
 if __name__ == "__main__":
-    args = get_args("Meta parameters for generating NSE 2d with McWilliams IV")
+    args = get_args_ns2d("Meta parameters for generating NSE 2d with McWilliams IV")
     main(args)

fno/data_gen/data_gen_fno.py

@@ -1,3 +1,12 @@
+# The MIT License (MIT)
+# Copyright © 2025 Shuhao Cao
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
 import argparse
 import math
 import os
@@ -10,6 +19,9 @@
 from grf import GRF2d
 from solvers import *
 from data_utils import *
+from torch_cfd.grids import *
+from torch_cfd.equations import *
+from torch_cfd.forcings import *
 from fno.pipeline import DATA_PATH, LOG_PATH
 
 def main(args):
@@ -25,13 +37,16 @@ def main(args):
     Sample usage:
 
     - Training data for Spectral-Refiner ICLR 2025 paper 'fnodata_extra_64x64_N1280_v1e-3_T50_steps100_alpha2.5_tau7.pt'
-    >>> python data_gen_fno.py --num-samples 1280 --batch-size 256 --grid-size 256 --subsample 4 --extra-vars --time 50 --time-warmup 30 --num-steps 100 --dt 1e-3 --visc 1e-3
+    >>> python data_gen_fno.py --num-samples 1280 --batch-size 256 --grid-size 256 --subsample 4 --extra-vars --time 50 --time-warmup 30 --num-steps 100 --dt 1e-3 --visc 1e-3 --scale 0.1
 
     - Test data
-    >>> python data_gen_fno.py --num-samples 16 --batch-size 8 --grid-size 256 --subsample 1 --double --extra-vars --time 50 --time-warmup 30 --num-steps 100 --dt 1e-3 --replicable-init --seed 42
+    >>> python data_gen_fno.py --num-samples 16 --batch-size 8 --grid-size 256 --subsample 1 --double --extra-vars --time 50 --time-warmup 30 --num-steps 100 --dt 1e-3 --scale 0.1 --replicable-init --seed 42
 
     - Test data fine
-    >>> python data_gen_fno.py --num-samples 2 --batch-size 1 --grid-size 512 --subsample 1 --double --extra-vars --time 50 --time-warmup 30 --num-steps 200 --dt 5e-4 --replicable-init --seed 42
+    >>> python data_gen_fno.py --num-samples 2 --batch-size 1 --grid-size 512 --subsample 1 --double --extra-vars --time 50 --time-warmup 30 --num-steps 200 --dt 5e-4 --scale 0.1 --replicable-init --seed 42
+
+    - Testing if the code works
+    >>> python data_gen/data_gen_fno.py --num-samples 4 --batch-size 2 --grid-size 128 --subsample 1 --double --extra-vars --time 2 --time-warmup 1 --num-steps 10 --dt 1e-3 --scale 0.1 --replicable-init --seed 42
 
     """
 
@@ -43,9 +58,6 @@ def main(args):
     log_filename = os.path.join(logpath, f"{current_time}_{log_name}.log")
     logger = get_logger(log_filename)
 
-    cuda = not args.no_cuda and torch.cuda.is_available()
-    device = torch.device("cuda" if cuda else "cpu")
-    logger.info(f"Using device: {device}")
     logger.info(f"Using the following arguments: ")
     all_args = {k: v for k, v in vars(args).items() if not callable(v)}
     logger.info(" | ".join(f"{k}={v}" for k, v in all_args.items()))
@@ -60,40 +72,37 @@ def main(args):
         raise ValueError(
             f"Grid size {n} is larger than the maximum allowed {n_grid_max}"
         )
+    scale = args.scale
     visc = args.visc if args.Re is None else 1/args.Re # 1e-3
     T = args.time  # 50
     T_warmup = args.time_warmup  # 30
     T_new = T - T_warmup
-    delta_t = args.dt  # 1e-4
+    record_steps = args.num_steps
+    dt = args.dt  # 1e-4
+    logger.info(f"Using dt = {dt}")
+
+    warmup_steps = int(T_warmup / dt)
+    total_steps = int(T_new / dt)
+    record_every_iters = int(total_steps / record_steps)
 
     alpha = args.alpha  # 2.5
     tau = args.tau  # 7
-    f = args.forcing  # FNO's default sin+cos
+    peak_wavenumber = args.peak_wavenumber
+    
     dtype = torch.float64 if args.double else torch.float32
     normalize = args.normalize
     filename = args.filename
     force_rerun = args.force_rerun
     replicate_init = args.replicable_init
     dealias = not args.no_dealias
     pbar = not args.no_tqdm
-    torch.set_default_dtype(torch.float64)
-    logger.info(f"Using device: {device} | save dtype: {dtype} | computge dtype: {torch.get_default_dtype()}")
 
     # Number of solutions to generate
     total_samples = args.num_samples  # 8
 
-    # Number of snapshots from solution
-    record_steps = args.num_steps
-
     # Batch size
     batch_size = args.batch_size  # 8
 
-    solver_kws = dict(visc=visc,
-            delta_t=delta_t,
-            diam=diam,
-            dealias=dealias,
-            dtype=torch.float64)
-
     extra = "_extra" if args.extra_vars else ""
     dtype_str = "_fp64" if args.double else ""
     if filename is None:
@@ -127,9 +136,27 @@ def main(args):
     else:
         logger.info(f"Generating data and saving in {filename}")
 
+    cuda = not args.no_cuda and torch.cuda.is_available()
+    no_tqdm = args.no_tqdm
+    device = torch.device("cuda:0" if cuda else "cpu")
+
+    torch.set_default_dtype(torch.float64)
+    logger.info(f"Using device: {device} | save dtype: {dtype} | computge dtype: {torch.get_default_dtype()}")
     # Set up 2d GRF with covariance parameters
     # Parameters of covariance C = tau^0.5*(2*alpha-2)*(-Laplacian + tau^2 I)^(-alpha)
     # Note that we need alpha > d/2 (here d= 2)
+
+    grid = Grid(shape=(n, n), domain=((0, diam), (0, diam)), device=device)
+    
+    forcing_fn = SinCosForcing(
+        grid=grid,
+        scale=scale,
+        diam=diam,
+        k=peak_wavenumber,
+        vorticity=True,
+    )
+    # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y)))
+    
     grf = GRF2d(
         n=n,
         alpha=alpha,
@@ -139,14 +166,14 @@ def main(args):
         dtype=torch.float64,
     )
 
-    # Forcing function: 0.1*(sin(2pi(x+y)) + cos(2pi(x+y)))
-    grid = torch.linspace(0, 1, n + 1, device=device)
-    grid = grid[0:-1]
-
-    X, Y = torch.meshgrid(grid, grid, indexing="ij")
-    # FNO's original implementation
-    # fh = 0.1 * (torch.sin(2 * math.pi * (X + Y)) + torch.cos(2 * math.pi * (X + Y)))
-    fh = f(X, Y)
+    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:
         logger.info(f"Data already exists at {data_filepath}")
@@ -165,46 +192,48 @@ def main(args):
         # Sample random fields
         seeds = [args.seed + idx + k for k in range(batch_size)]
         n0 = n_grid_max if replicate_init else n
-        w0 = [grf.sample(1, n0, random_state=s) for _, s in zip(range(batch_size), seeds)]
-        w0 = torch.stack(w0)
+        vort_init = [grf.sample(1, n0, random_state=s) for _, s in zip(range(batch_size), seeds)]
+        vort_init = torch.stack(vort_init)
         if n != n0:
-            w0 = F.interpolate(w0, size=(n, n), mode="nearest")
-        w0 = w0.squeeze(1)
-        
-        logger.info(f"initial condition {w0.shape}")
+            vort_init = F.interpolate(vort_init, size=(n, n), mode="nearest")
+        vort_init = vort_init.squeeze(1)
+        vort_hat = fft.rfft2(vort_init).to(device)
+
+        logger.info(f"initial condition {vort_init.shape}")
 
         if T_warmup > 0:
-            logger.info(f"warm up till {T_warmup}")
-            tmp = get_trajectory_imex_crank_nicolson(
-                w0,
-                fh,
-                T=T_warmup,
-                record_steps=record_steps,
-                subsample=1,
-                pbar=pbar,
-                **solver_kws,
-            )
-            w0 = tmp["vorticity"][:, -1].to(device)
-            del tmp
-            logger.info(f"warmup initial condition {w0.shape}")
+            with tqdm(total=warmup_steps, disable=no_tqdm) as pbar:
+                for j in range(warmup_steps):
+                    vort_hat, _ = ns2d.step(vort_hat, dt)
+                    if j % 100 == 0:
+                        vort_norm = torch.linalg.norm(fft.irfft2(vort_hat)).item() / n
+                        desc = (
+                            datetime.now().strftime("%d-%b-%Y %H:%M:%S")
+                            + f" - Warmup | vort_hat ell2 norm {vort_norm:.4e}"
+                        )
+                        pbar.set_description(desc)
+                        pbar.update(100)
 
         logger.info(f"generate data from {T_warmup} to {T}")
-        result = get_trajectory_imex_crank_nicolson(
-            w0,
-            fh,
-            T=T_new,
-            record_steps=record_steps,
-            subsample=subsample,
-            pbar=pbar,
-            **solver_kws,
+        result = get_trajectory_imex(
+            ns2d,
+            vort_hat,
+            dt,
+            num_steps=total_steps,
+            record_every_steps=record_every_iters,
+            pbar=not no_tqdm,
         )
 
         for field, value in result.items():
-            if subsample > 1 and value.ndim == 4:
+            value = fft.irfft2(value).real.cpu().to(dtype)
+            logger.info(
+                f"variable: {field} | shape: {value.shape} | dtype: {value.dtype}"
+            )
+            if subsample > 1:
+                assert value.ndim == 4, f"Subsampling only works for (b, c, h, w) tensors, current shape: {value.shape}"
                 value = F.interpolate(value, size=(ns, ns), mode="bilinear")
-            result[field] = value.cpu().to(dtype)
+            result[field] = value
             logger.info(f"{field:<15} | {value.shape} | {value.dtype}")
-            
 
         if not extra:
             for key in ["vort_t", "stream", "residual"]:
@@ -233,5 +262,5 @@ def main(args):
 
 
 if __name__ == "__main__":
-    args = get_args("Generate the original FNO data for NSE in 2D")
+    args = get_args_ns2d("Generate the original FNO data for NSE in 2D")
     main(args)