fixed mcwilliams data gen fp64 bug

Author: scaomath

README.md

@@ -4,7 +4,8 @@ This repository featuers two parts:
 - The first part is a native PyTorch port of [Google's Computational Fluid Dynamics package in Jax](https://github.com/google/jax-cfd). The main changes are documented in the `README.md` under the [`torch_cfd` directory](torch_cfd). Most significant changes in all routines include:
   - Routines that rely on the functional programming of Jax have been rewritten to be a more debugger-friendly PyTorch tensor-in-tensor-out style.
   - Functions and operators are in general implemented as `nn.Module` like a factory template.
-  - Extra fields computation and tracking are made easier, such as time derivatives and PDE residual $R(\boldsymbol{v}):=\boldsymbol{f}-\partial_t \boldsymbol{v}-(\boldsymbol{v}\cdot\nabla)\boldsymbol{v} + \nu \Delta \boldsymbol{v}$.
+  - Jax-cfd's `funcutils.trajectory` function supports to track only one field variable (vorticity or velocity), Extra fields computation and tracking are made easier, such as time derivatives and PDE residual $R(\boldsymbol{v}):=\boldsymbol{f}-\partial_t \boldsymbol{v}-(\boldsymbol{v}\cdot\nabla)\boldsymbol{v} + \nu \Delta \boldsymbol{v}$.
+  - All ops takes batch dimension of tensors into consideration, not a single trajectory.
 - Neural Operator-Assisted Navier-Stokes Equations solver.
   - The **Spatiotempoeral Fourier Neural Operator** (SFNO) that is a spacetime tensor-to-tensor learner (or trajectory-to-trajectory), inspiration drawn from the [3D FNO in Nvidia's Neural Operator repo](https://github.com/neuraloperator/neuraloperator).
   - Data generation for the meta-example of the isotropic turbulence with energy spectra matching the inverse cascade of Kolmogorov flow in a periodic box. Ref: McWilliams, J. C. (1984). The emergence of isolated coherent vortices in turbulent flow. *Journal of Fluid Mechanics*, 146, 21-43.
@@ -16,7 +17,7 @@ To install `torch-cfd`'s current release, simply do:
 ```bash
 pip install torch-cfd
 ```
-If one wants to play with the neural operator part, it is recommended cloning this repo and playing it local by creating a venv using `requirements.txt`. Note: it is recommended using PyTorch version >=2.0.0 for the broadcasting semantics.
+If one wants to play with the neural operator part, it is recommended to clone this repo and play it locally by creating a venv using `requirements.txt`. Note: using PyTorch version >=2.0.0 for the broadcasting semantics.
 
 ## Data
 The data are available at https://huggingface.co/datasets/scaomath/navier-stokes-dataset 
@@ -52,6 +53,6 @@ PR welcome. Currently, the port of `torch-cfd` currently includes:
 
 ## Acknowledgments
 The research of Brarda and Xi is supported by the National Science Foundation award DMS-2208412. 
-The work of  Li was performed under the auspices of
+The work of Li was performed under the auspices of
 the U.S. Department of Energy by Lawrence Livermore National Laboratory under Contract DEAC52-07NA27344 and was supported by the LLNL-LDRD program under Project No. 24ERD033. Cao is greatful for the support from [Long Chen (UC Irvine)](https://github.com/lyc102/ifem) and 
 [Ludmil Zikatanov (Penn State)](https://github.com/HAZmathTeam/hazmath) over the years, and their efforts in open-sourcing scientific computing codes. Cao also appreciates the support from the National Science Foundation DMS-2309778, and the free A6000 credits at the SSE ML cluster from the University of Missouri.

fno/data/data_gen_McWilliams2d.py

@@ -19,12 +19,6 @@
 from torch_cfd.initial_conditions import *
 from torch_cfd.finite_differences import *
 from torch_cfd.forcings import *
-
-from collections import defaultdict
-
-import matplotlib.pyplot as plt
-import seaborn as sns
-import xarray
 from tqdm import tqdm
 from .data_gen import *
 
@@ -79,10 +73,10 @@ def main(args):
     filename = args.filename
     if filename is None:
         filename = (
-            f"McWilliams2d{dtype_str}"
-            + f"_N{total_samples}_n{ns}"
-            + f"_v{viscosity:.0e}_T{T}.pt"
-        ).replace("e-0", "e-")
+            f"McWilliams2d{dtype_str}_{ns}x{ns}"
+            + f"_N{total_samples}_v{viscosity:.0e}"
+            + f"_T{num_snapshots}.pt".replace("e-0", "e-")
+        )
         args.filename = filename
     data_filepath = os.path.join(DATA_PATH, filename)
     if os.path.exists(data_filepath) and not force_rerun:
@@ -112,10 +106,10 @@ def main(args):
     ).to(device)
 
     for i, idx in enumerate(range(0, total_samples, batch_size)):
+        logger.info(f"Generate trajectory for {i+1}-th batch of {total_samples}")
         logger.info(
-            f"Generate trajectory for {i+1}-th batch of {total_samples}"
+            f"random state: {random_state + idx} to {random_state + idx + batch_size-1}"
         )
-        logger.info(f"random state: {random_state + idx} to {random_state + idx + batch_size-1}")
 
         vort_init = torch.stack(
             [
@@ -143,7 +137,10 @@ def main(args):
         )
 
         for field, value in result.items():
-            value = fft.irfft2(value).real.cpu().to(torch.float32)
+            value = fft.irfft2(value).real.cpu().to(dtype)
+            logger.info(
+                f"variable: {field} | shape: {value.shape} | dtype: {value.dtype}"
+            )
             if subsample > 1:
                 result[field] = F.interpolate(value, size=(ns, ns), mode="bilinear")
             else:
@@ -158,7 +155,9 @@ def main(args):
 
     pickle_to_pt(data_filepath)
 
-    verify_trajectories(data_filepath, dt=record_every_iters * dt, T_warmup=T_warmup)
+    verify_trajectories(
+        data_filepath, dt=record_every_iters * dt, T_warmup=T_warmup, n_samples=1
+    )
 
 
 if __name__ == "__main__":

fno/sfno.py

@@ -207,7 +207,7 @@ def forward(self, x, out_steps=None):
             nx,
             ny,
             nt // 2 + 1,
-            dtype=torch.cfloat,
+            dtype=x_ft.dtype,
             device=x.device,
         )
         slice_x = [slice(0, self.modes_x), slice(-self.modes_x, None)]
@@ -319,6 +319,13 @@ def hook(model, input, output):
         for k, block in enumerate(blocks):
             block.register_forward_hook(_get_latent_tensors(f"{layer_name}_{k}"))
 
+    def double(self):
+        for param in self.parameters():
+            if param.dtype == torch.float32:
+                param.data = param.data.to(torch.float64)
+            elif param.dtype == torch.complex64:
+                param.data = param.data.to(torch.complex128)
+
     def forward(self, x, out_steps=None):
         """
         if out_steps is None, it will try to use self.out_steps
@@ -327,13 +334,16 @@ def forward(self, x, out_steps=None):
         if out_steps is None:
             out_steps = self.out_steps if self.out_steps is not None else x.size(-1)
         x_res = x  # save skip connection
-        x = self.p(x.unsqueeze(1))  # [b, 1, n, n, T] -> [b, H, n, n, T]
+        x = self.p(x.unsqueeze(1) )  # [b, 1, n, n, T] -> [b, H, n, n, T]
+
+        if self.debug:
+            print(f"in proj: {x.size()}")
 
         x = F.pad(
             x,
             [0, 0, self.padding, self.padding, self.padding, self.padding],
             mode="circular",
-        )
+        )  # pad the domain if input is non-periodic
 
         for conv, mlp, w, nonlinear in zip(
             self.spectral_conv, self.mlp, self.w, self.activations

fno/visualizations.py

@@ -1,18 +1,17 @@
+import matplotlib.pyplot as plt
+import numpy as np
 import plotly.express as px
 import plotly.figure_factory as ff
 import plotly.graph_objects as go
-
-import xarray
-import numpy as np
+import seaborn as sns
 import torch
 import torch.fft as fft
-import matplotlib.pyplot as plt
-import seaborn as sns
+
+import xarray
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-def plot_contour(z, 
-                func=plt.imshow,
-                **kwargs):
+
+def plot_contour(z, func=plt.imshow, **kwargs):
     if isinstance(z, torch.Tensor):
         z = z.cpu().numpy()
     _, ax = plt.subplots(figsize=(3, 3))
@@ -23,9 +22,10 @@ def plot_contour(z,
     cax = divider.append_axes("right", size="7%", pad=0.1)
     cbar = plt.colorbar(f, ax=ax, cax=cax)
     cbar.ax.tick_params(labelsize=10)
-    cbar.ax.locator_params(nbins=9)   
+    cbar.ax.locator_params(nbins=9)
     cbar.update_ticks()
 
+
 def plot_contour_plotly(
     z,
     colorscale="RdYlBu",
@@ -83,7 +83,7 @@ def plot_contour_plotly(
         contour_kwargs["colorbar"] = dict(
             thickness=0.15 * layout_kwargs["height"],
             tickwidth=0.3,
-            exponentformat = 'e',
+            exponentformat="e",
         )
         layout_kwargs["width"] = 1.32 * layout_kwargs["height"]
     else:
@@ -111,122 +111,129 @@ def plot_contour_plotly(
     return fig
 
 
-
 def get_enstrophy_spectrum(vorticity, h):
     if isinstance(vorticity, np.ndarray):
         vorticity = torch.from_numpy(vorticity)
     n = vorticity.shape[0]
     kx = fft.fftfreq(n, d=h)
     ky = fft.fftfreq(n, d=h)
     kx, ky = torch.meshgrid([kx, ky], indexing="ij")
-    kmax = n//2
+    kmax = n // 2
     kx = kx[..., : kmax + 1]
     ky = ky[..., : kmax + 1]
-    k2 = (4*torch.pi**2)*(kx**2 + ky**2)
+    k2 = (4 * torch.pi**2) * (kx**2 + ky**2)
     k2[0, 0] = 1.0
 
     wh = fft.rfft2(vorticity)
 
-    tke = (0.5*wh*wh.conj()).real
+    tke = (0.5 * wh * wh.conj()).real
     kmod = torch.sqrt(k2)
-    k = torch.arange(1, kmax, dtype=torch.float64) # Nyquist limit for this grid
+    k = torch.arange(1, kmax, dtype=torch.float64)  # Nyquist limit for this grid
     Ens = torch.zeros_like(k)
-    dk = (torch.max(k)-torch.min(k))/(2*n)
+    dk = (torch.max(k) - torch.min(k)) / (2 * n)
     for i in range(len(k)):
-        Ens[i] += (tke[(kmod<k[i]+dk) & (kmod>=k[i]-dk)]).sum()
+        Ens[i] += (tke[(kmod < k[i] + dk) & (kmod >= k[i] - dk)]).sum()
 
-    Ens = Ens/Ens.sum()
+    Ens = Ens / Ens.sum()
     return Ens
 
 
-def plot_enstrophy_spectrum(fields:list,
-                            h=None, 
-                            slope=5,
-                            factor=None,
-                            cutoff=1e-15,
-                            plot_cutoff_factor=1/8,
-                            labels=None,
-                            title=None,
-                            legend_loc="upper right",
-                            fontsize=15,
-                            subplot_kw={"figsize": (5, 5), "dpi": 100, "facecolor": "w"},
-                            **kwargs):
+def plot_enstrophy_spectrum(
+    fields: list,
+    h=None,
+    slope=5,
+    factor=None,
+    cutoff=1e-15,
+    plot_cutoff_factor=1 / 8,
+    labels=None,
+    title=None,
+    legend_loc="upper right",
+    fontsize=15,
+    subplot_kw={"figsize": (5, 5), "dpi": 100, "facecolor": "w"},
+    **kwargs,
+):
     for k, field in enumerate(fields):
         if isinstance(field, np.ndarray):
             fields[k] = torch.from_numpy(field)
     if labels is None:
         labels = [f"Field {i}" for i in range(len(fields))]
     n = fields[0].shape[0]
-    if h is None: h = 1 / n
-    kmax = n//2
-    k = torch.arange(1, kmax, dtype=torch.float64) # Nyquist limit for this grid
+    if h is None:
+        h = 1 / n
+    kmax = n // 2
+    k = torch.arange(1, kmax, dtype=torch.float64)  # Nyquist limit for this grid
     Es = [get_enstrophy_spectrum(field, h) for field in fields]
     if factor is None:
-        factor = Es[-1].quantile(0.8)/(k[-1] ** (-slope))
+        factor = Es[-1].quantile(0.8) / (k[-1] ** (-slope))
         # print(factor)
-    
+
     fig, ax = plt.subplots(**subplot_kw)
-    plot_cutoff = int(n*plot_cutoff_factor)
+    plot_cutoff = int(n * plot_cutoff_factor)
     for i, E in enumerate(Es):
         if cutoff is not None:
             E[E < cutoff] = np.nan
         E[-plot_cutoff:] = np.nan
         plt.loglog(k, E, label=f"{labels[i]}")
 
-    plt.loglog(k[:-plot_cutoff], (factor*k ** (-slope))[:-plot_cutoff], "b--", 
-               label=f"$O(k^{{{-slope:.3g}}})$",)
+    plt.loglog(
+        k[:-plot_cutoff],
+        (factor * k ** (-slope))[:-plot_cutoff],
+        "b--",
+        label=f"$O(k^{{{-slope:.3g}}})$",
+    )
     plt.grid(True, which="both", ls="--", linewidth=0.4)
-    plt.autoscale(enable=True, axis='x', tight=True)
+    plt.autoscale(enable=True, axis="x", tight=True)
     plt.legend(fontsize=fontsize, loc=legend_loc)
     plt.title(title, fontsize=fontsize)
     plt.xlabel("Wavenumber", fontsize=fontsize)
     ax.xaxis.set_tick_params(labelsize=fontsize)
     ax.yaxis.set_tick_params(labelsize=fontsize)
 
 
-def plot_contour_trajectory(field, 
-                            num_snapshots=5, 
-                            contourf=False,
-                            T_start=4.5,
-                            dt = 1e-1,
-                            **kwargs):
+def plot_contour_trajectory(
+    field,
+    num_snapshots=5,
+    contourf=False,
+    T_start=4.5,
+    dt=1e-1,
+    cb_kws=dict(orientation="vertical", pad=0.01, aspect=10),
+    subplot_kws=dict(
+        xticks=[],
+        yticks=[],
+        ylabel="",
+        xlabel="",
+    ),
+    plot_kws=dict(
+        col_wrap=5,
+        cmap=sns.cm.icefire,
+        robust=True,
+        add_colorbar=True,
+        xticks=None,
+        yticks=None,
+        size=3,
+        aspect=1,
+    ),
+    **kwargs,
+):
     """
     plot trajectory using xarray's imshow or contourf wrapper
     """
     field = field.detach().cpu().numpy()
     *size, T = field.shape
-    grid = np.linspace(0, 1, size[0]+1)[:-1]
+    grid = np.linspace(0, 1, size[0] + 1)[:-1]
     time = np.arange(T) * dt + T_start
     coords = {
-        'x': grid,
-        'y': grid,
-        't':  time,
+        "x": grid,
+        "y": grid,
+        "t": time,
     }
-    ds = xarray.DataArray(
-        field,
-        dims=["x", "y", "t"],
-        coords=coords
-    )
-    t_steps = T//num_snapshots
+    ds = xarray.DataArray(field, dims=["x", "y", "t"], coords=coords)
+    t_steps = T // num_snapshots
     ds = ds.thin({"t": t_steps})
     plot_func = ds.plot.contourf if contourf else ds.plot.imshow
-    plot_func(col='t', 
-            col_wrap=5, 
-            cmap=sns.cm.icefire, 
-            robust=True,
-            add_colorbar=True,
-            xticks=None,
-            yticks=None,
-            size=3, 
-            aspect=1,
-            interpolation="hermite",
-            subplot_kws=dict(xticks=[], 
-                        yticks=[],
-                        ylabel= "",
-                        xlabel= "",
-            ),
-            cbar_kwargs=dict(orientation="vertical",
-                          pad=0.01,
-                          aspect= 10
-        ),
-        **kwargs)
+    plot_func(
+        col="t",
+        subplot_kws=subplot_kws,
+        cbar_kwargs=cb_kws,
+        **plot_kws,
+    )