Github action: auto-update.

Author: github-actions[bot]

dev/_downloads/02a7a82499cf14d2702c9234035c44c4/plot_diffusion_advection_solver.zip

@@ -22,7 +22,7 @@
       },
       "outputs": [],
       "source": [
-        "import torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\n\nfrom neuralop.losses.finite_diff import central_diff_2d  \n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
+        "import torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\n\nfrom neuralop.losses.differentiation import FiniteDiff  \n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
       ]
     },
     {
@@ -40,7 +40,7 @@
       },
       "outputs": [],
       "source": [
-        "## Simulation parameters\nLx, Ly = 2.0, 2.0   # Domain lengths\nnx, ny = 64, 64   # Grid resolution\nT = 1.6    # Total simulation time\ndt = 0.001  # Time step\nnu = 0.02   # diffusion coefficient\ncx, cy = 1.0, 0.6  # advection speeds\n\n## Create grid\nX = torch.linspace(0, Lx, nx, device=device).repeat(ny, 1).T \nY = torch.linspace(0, Ly, ny, device=device).repeat(nx, 1)  \ndx = Lx / (nx - 1)\ndy = Ly / (ny - 1)\nnt = int(T / dt)\n\n\n## Initial condition and source term\nu = (-torch.sin(2 * np.pi * Y) * torch.cos(2 * np.pi * X)\n        + 0.3 * torch.exp(-((X - 0.75)**2 + (Y - 0.5)**2) / 0.02)\n        - 0.3 * torch.exp(-((X - 1.25)**2 + (Y - 1.5)**2) / 0.02)).to(device)\n\ndef source_term(X, Y, t):\n    return 0.2 * torch.sin(3 * np.pi * X) * torch.cos(3 * np.pi * Y) * torch.cos(4 * np.pi * t)"
+        "## Simulation parameters\nLx, Ly = 2.0, 2.0   # Domain lengths\nnx, ny = 64, 64   # Grid resolution\nT = 1.6    # Total simulation time\ndt = 0.001  # Time step\nnu = 0.02   # diffusion coefficient\ncx, cy = 1.0, 0.6  # advection speeds\n\n## Create grid\nX = torch.linspace(0, Lx, nx, device=device).repeat(ny, 1).T \nY = torch.linspace(0, Ly, ny, device=device).repeat(nx, 1)  \ndx = Lx / (nx - 1)\ndy = Ly / (ny - 1)\nnt = int(T / dt)\n\n## Initialize finite difference operator\nfd = FiniteDiff(dim=2, h=(dx, dy))\n\n\n## Initial condition and source term\nu = (-torch.sin(2 * np.pi * Y) * torch.cos(2 * np.pi * X)\n        + 0.3 * torch.exp(-((X - 0.75)**2 + (Y - 0.5)**2) / 0.02)\n        - 0.3 * torch.exp(-((X - 1.25)**2 + (Y - 1.5)**2) / 0.02)).to(device)\n\ndef source_term(X, Y, t):\n    return 0.2 * torch.sin(3 * np.pi * X) * torch.cos(3 * np.pi * Y) * torch.cos(4 * np.pi * t)"
       ]
     },
     {
@@ -58,7 +58,7 @@
       },
       "outputs": [],
       "source": [
-        "u_evolution = [u.clone()]\n\nt = torch.tensor(0.0)\nfor _ in range(nt):\n    \n    # Compute derivatives\n    u_x, u_y = central_diff_2d(u, [dx, dy])\n    u_xx, _ = central_diff_2d(u_x, [dx, dy])\n    _, u_yy = central_diff_2d(u_y, [dx, dy])\n\n    # Evolve one step in time using Euler's method\n    u = u + dt * (-cx * u_x - cy * u_y + nu * (u_xx + u_yy) + source_term(X, Y, t))\n    t += dt\n    u_evolution.append(u.clone())\n\nu_evolution = torch.stack(u_evolution).cpu().numpy()"
+        "u_evolution = [u.clone()]\n\nt = torch.tensor(0.0)\nfor _ in range(nt):\n    \n    # Compute derivatives\n    u_x = fd.dx(u)\n    u_y = fd.dy(u)\n    u_xx = fd.dx(u_x)\n    u_yy = fd.dy(u_y)\n\n    # Evolve one step in time using Euler's method\n    u = u + dt * (-cx * u_x - cy * u_y + nu * (u_xx + u_yy) + source_term(X, Y, t))\n    t += dt\n    u_evolution.append(u.clone())\n\nu_evolution = torch.stack(u_evolution).cpu().numpy()"
       ]
     },
     {

dev/_downloads/07795e165bda9881ff980becb31905ab/plot_diffusion_advection_solver.ipynb

@@ -3,7 +3,7 @@
 
 Fourier Differentiation
 ========================================================
-An example of usage of our Fourier Differentiation Function on 1d data.
+An example of usage of our Fourier Differentiation Function
 """
 
 # %%
@@ -13,33 +13,33 @@
 import torch
 import numpy as np
 import matplotlib.pyplot as plt
-from neuralop.losses.fourier_diff import fourier_derivative_1d
+from neuralop.losses.differentiation import FourierDiff
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 
 # %%
 # Creating an example of periodic 1D curve
-# --------------------
-# Here we consider sin(x) and cos(x), which are periodic on the interval [0,2pi]
-L = 2*torch.pi
+# ----------------------------------------
+# Here we consider sin(x) and cos(x), which are periodic on the interval [0, 2π]
+L = 2 * torch.pi
 x = torch.linspace(0, L, 101)[:-1]
 f = torch.stack([torch.sin(x), torch.cos(x)], dim=0)
 x_np = x.cpu().numpy()
 
 # %%
 # Differentiate the signal
 # -----------------------------------------
-# We use the Fourier differentiation function to differentiate the signal
-dfdx = fourier_derivative_1d(f, order=1, L=L)
-df2dx2 = fourier_derivative_1d(f, order=2, L=L)
-df3dx3 = fourier_derivative_1d(f, order=3, L=L)
+# We use the FourierDiff class to differentiate the signal
+fd1d = FourierDiff(dim=1, L=L, use_fc=False)
+derivatives = fd1d.compute_multiple_derivatives(f, [1, 2, 3])
+dfdx, df2dx2, df3dx3 = derivatives
 
 
 # %%
 # Plot the results for sin(x)
-# ----------------------
+# ---------------------------
 plt.figure()
 plt.plot(x_np, dfdx[0].squeeze().cpu().numpy(), label='Fourier dfdx')
 plt.plot(x_np, np.cos(x_np), '--', label='dfdx')
@@ -53,7 +53,7 @@
 
 # %%
 # Plot the results for cos(x)
-# ----------------------
+# ---------------------------
 plt.figure()
 plt.plot(x_np, dfdx[1].squeeze().cpu().numpy(), label='Fourier dfdx')
 plt.plot(x_np, -np.sin(x_np), '--', label='dfdx')
@@ -69,24 +69,25 @@
 
 # %%
 # Creating an example of non-periodic 1D curve
-# --------------------
-# Here we consider sin(16x)-cos(8x) and exp(-0.8x)+sin(x)
-L = 2*torch.pi
-x = torch.linspace(0, L, 101)[:-1]    
-f = torch.stack([torch.sin(3*x) - torch.cos(x), torch.exp(-0.8*x)+torch.sin(x)], dim=0)
+# -------------------------------------------
+# Here we consider sin(3x)-cos(x) and exp(-0.8x)+sin(x)
+L = 2 * torch.pi
+x = torch.linspace(0, L, 101)[:-1]
+f = torch.stack([torch.sin(3*x) - torch.cos(x), torch.exp(-0.8*x) + torch.sin(x)], dim=0)
 x_np = x.cpu().numpy()
 
 # %%
 # Differentiate the signal
 # -----------------------------------------
-# We use the Fourier differentiation function with Fourier continuation to differentiate the signal
-dfdx = fourier_derivative_1d(f, order=1, L=L, use_FC='Legendre', FC_d=4, FC_n_additional_pts=30)
-df2dx2 = fourier_derivative_1d(f, order=2, L=L, use_FC='Legendre', FC_d=4, FC_n_additional_pts=30)
+# We use the FourierDiff class with Fourier continuation to differentiate the signal
+fd1d = FourierDiff(dim=1, L=L, use_fc='Legendre', fc_degree=4, fc_n_additional_pts=50)
+derivatives = fd1d.compute_multiple_derivatives(f, [1, 2])
+dfdx, df2dx2 = derivatives
 
 
 # %%
-# Plot the results for sin(16x)-cos(8x)
-# ----------------------
+# Plot the results for sin(3x)-cos(x)
+# --------------------------------------
 plt.figure()
 plt.plot(x_np, dfdx[0].squeeze().cpu().numpy(), label='Fourier dfdx')
 plt.plot(x_np, 3*torch.cos(3*x) + torch.sin(x), '--', label='dfdx')
@@ -98,12 +99,186 @@
 
 # %%
 # Plot the results for exp(-0.8x)+sin(x)
-# ----------------------
+# ---------------------------------------
 plt.figure()
 plt.plot(x_np, dfdx[1].squeeze().cpu().numpy(), label='Fourier dfdx')
 plt.plot(x_np, -0.8*torch.exp(-0.8*x)+torch.cos(x), '--', label='dfdx')
 plt.plot(x_np, df2dx2[1].squeeze().cpu().numpy(), label='Fourier df2dx2')
 plt.plot(x_np, 0.64*torch.exp(-0.8*x)-torch.sin(x), '--', label='df2dx2')
 plt.xlabel('x')
 plt.legend()
+plt.show()
+
+
+# %%
+# 2D Fourier Differentiation Examples
+# ===================================
+# Here we demonstrate the FourierDiff class for 2D functions
+
+# %%
+# Creating an example of periodic 2D function
+# -----------------------------------------
+# Here we consider f(x,y) = sin(x) * cos(y), which is periodic on the interval [0, 2π] × [0, 2π]
+L_x, L_y = 2 * torch.pi, 2 * torch.pi
+nx, ny = 180, 186
+x = torch.linspace(0, L_x, nx, dtype=torch.float64)
+y = torch.linspace(0, L_y, ny, dtype=torch.float64)
+X, Y = torch.meshgrid(x, y, indexing='ij')
+
+# Test function: f(x,y) = sin(x) * cos(y)
+f_2d = torch.sin(X) * torch.cos(Y)
+
+# %%
+# Differentiate the 2D signal
+# -----------------------------------------
+# We use the FourierDiff class to compute derivatives
+fd2d = FourierDiff(dim=2, L=(L_x, L_y))
+
+# Compute derivatives
+df_dx = fd2d.dx(f_2d)
+df_dy = fd2d.dy(f_2d)
+laplacian = fd2d.laplacian(f_2d)
+
+# Expected analytical results for f(x,y) = sin(x) * cos(y)
+df_dx_expected = torch.cos(X) * torch.cos(Y)  
+df_dy_expected = -torch.sin(X) * torch.sin(Y)  
+laplacian_expected = -2 * torch.sin(X) * torch.cos(Y) 
+
+# %%
+# Plot the 2D results
+# ----------------------
+fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+fig.suptitle('2D Fourier Differentiation Results: f(x,y) = sin(x) * cos(y)')
+
+# Compute consistent colorbar limits for each derivative pair
+df_dx_min = min(df_dx.min().item(), df_dx_expected.min().item())
+df_dx_max = max(df_dx.max().item(), df_dx_expected.max().item())
+df_dy_min = min(df_dy.min().item(), df_dy_expected.min().item())
+df_dy_max = max(df_dy.max().item(), df_dy_expected.max().item())
+
+# Original function
+im0 = axes[0, 0].imshow(f_2d.cpu().numpy())
+axes[0, 0].set_title('Original: sin(x) * cos(y)')
+plt.colorbar(im0, ax=axes[0, 0], shrink=0.57)
+
+# ∂f/∂x computed
+im1 = axes[0, 1].imshow(df_dx.cpu().numpy(), vmin=df_dx_min, vmax=df_dx_max)
+axes[0, 1].set_title('∂f/∂x (computed)')
+plt.colorbar(im1, ax=axes[0, 1], shrink=0.57)
+
+# ∂f/∂x expected
+im2 = axes[0, 2].imshow(df_dx_expected.cpu().numpy(), vmin=df_dx_min, vmax=df_dx_max)
+axes[0, 2].set_title('∂f/∂x (expected: cos(x) * cos(y))')
+plt.colorbar(im2, ax=axes[0, 2], shrink=0.57)
+
+# ∂f/∂y computed
+im3 = axes[1, 0].imshow(df_dy.cpu().numpy(), vmin=df_dy_min, vmax=df_dy_max)
+axes[1, 0].set_title('∂f/∂y (computed)')
+plt.colorbar(im3, ax=axes[1, 0], shrink=0.57)
+
+# ∂f/∂y expected
+im4 = axes[1, 1].imshow(df_dy_expected.cpu().numpy(), vmin=df_dy_min, vmax=df_dy_max)
+axes[1, 1].set_title('∂f/∂y (expected: -sin(x) * sin(y))')
+plt.colorbar(im4, ax=axes[1, 1], shrink=0.57)
+
+# Laplacian
+im5 = axes[1, 2].imshow(laplacian.cpu().numpy())
+axes[1, 2].set_title('∇²f (computed)')
+plt.colorbar(im5, ax=axes[1, 2], shrink=0.57)
+
+plt.tight_layout()
+plt.show()
+
+
+
+
+# %%
+# 3D Fourier Differentiation Examples
+# ===================================
+# Here we demonstrate the FourierDiff class for 3D functions
+
+# %%
+# Creating an example of periodic 3D function
+# -----------------------------------------
+# Here we consider f(x,y,z) = sin(x) * cos(y) * sin(z), which is periodic on [0, 2π]³
+L_x, L_y, L_z = 2 * torch.pi, 2 * torch.pi, 2 * torch.pi
+nx, ny, nz = 176, 180, 192
+x = torch.linspace(0, L_x, nx, dtype=torch.float64)
+y = torch.linspace(0, L_y, ny, dtype=torch.float64)
+z = torch.linspace(0, L_z, nz, dtype=torch.float64)
+X, Y, Z = torch.meshgrid(x, y, z, indexing='ij')
+
+# Test function: f(x,y,z) = sin(x) * cos(y) * sin(z)
+f_3d = torch.sin(X) * torch.cos(Y) * torch.sin(Z)
+
+# Alternative: create tensor directly like in the test
+f_3d_alt = torch.randn(nx, ny, nz, dtype=torch.float64)
+
+# %%
+# Differentiate the 3D signal
+# -----------------------------------------
+# We use the FourierDiff class to compute derivatives
+fd3d = FourierDiff(dim=3, L=(L_x, L_y, L_z))
+
+# Compute derivatives
+df_dx_3d = fd3d.dx(f_3d)
+df_dy_3d = fd3d.dy(f_3d)
+df_dz_3d = fd3d.dz(f_3d)
+laplacian_3d = fd3d.laplacian(f_3d)
+
+# Expected analytical results for f(x,y,z) = sin(x) * cos(y) * sin(z)
+df_dx_expected_3d = torch.cos(X) * torch.cos(Y) * torch.sin(Z) 
+df_dy_expected_3d = -torch.sin(X) * torch.sin(Y) * torch.sin(Z) 
+df_dz_expected_3d = torch.sin(X) * torch.cos(Y) * torch.cos(Z) 
+laplacian_expected_3d = -3 * torch.sin(X) * torch.cos(Y) * torch.sin(Z)  
+
+# %%
+# Plot a slice of the 3D results (z=0 plane)
+# ------------------------------------------
+z_slice_idx = nz // 2
+fig, axes = plt.subplots(2, 3, figsize=(18, 12))
+fig.suptitle('3D Fourier Differentiation Results (z=0 slice): f(x,y,z) = sin(x) * cos(y) * sin(z)')
+
+# Compute consistent colorbar limits for each derivative pair at the z-slice
+df_dx_3d_slice = df_dx_3d[:, :, z_slice_idx]
+df_dx_expected_3d_slice = df_dx_expected_3d[:, :, z_slice_idx]
+df_dy_3d_slice = df_dy_3d[:, :, z_slice_idx]
+df_dy_expected_3d_slice = df_dy_expected_3d[:, :, z_slice_idx]
+
+df_dx_3d_min = min(df_dx_3d_slice.min().item(), df_dx_expected_3d_slice.min().item())
+df_dx_3d_max = max(df_dx_3d_slice.max().item(), df_dx_expected_3d_slice.max().item())
+df_dy_3d_min = min(df_dy_3d_slice.min().item(), df_dy_expected_3d_slice.min().item())
+df_dy_3d_max = max(df_dy_3d_slice.max().item(), df_dy_expected_3d_slice.max().item())
+
+# Original function slice
+im0 = axes[0, 0].imshow(f_3d[:, :, z_slice_idx].cpu().numpy())
+axes[0, 0].set_title('Original: sin(x) * cos(y) * sin(z)')
+plt.colorbar(im0, ax=axes[0, 0], shrink=0.57)
+
+# ∂f/∂x slice
+im1 = axes[0, 1].imshow(df_dx_3d_slice.cpu().numpy(), vmin=df_dx_3d_min, vmax=df_dx_3d_max)
+axes[0, 1].set_title('∂f/∂x (computed)')
+plt.colorbar(im1, ax=axes[0, 1], shrink=0.57)
+
+# ∂f/∂x expected slice
+im2 = axes[0, 2].imshow(df_dx_expected_3d_slice.cpu().numpy(), vmin=df_dx_3d_min, vmax=df_dx_3d_max)
+axes[0, 2].set_title('∂f/∂x (expected: cos(x) * cos(y) * sin(z))')
+plt.colorbar(im2, ax=axes[0, 2], shrink=0.57)
+
+# ∂f/∂y slice
+im3 = axes[1, 0].imshow(df_dy_3d_slice.cpu().numpy(), vmin=df_dy_3d_min, vmax=df_dy_3d_max)
+axes[1, 0].set_title('∂f/∂y (computed)')
+plt.colorbar(im3, ax=axes[1, 0], shrink=0.57)
+
+# ∂f/∂y expected slice
+im4 = axes[1, 1].imshow(df_dy_expected_3d_slice.cpu().numpy(), vmin=df_dy_3d_min, vmax=df_dy_3d_max)
+axes[1, 1].set_title('∂f/∂y (expected: -sin(x) * sin(y) * sin(z))')
+plt.colorbar(im4, ax=axes[1, 1], shrink=0.57)
+
+# ∂f/∂z slice
+im5 = axes[1, 2].imshow(df_dz_3d[:, :, z_slice_idx].cpu().numpy())
+axes[1, 2].set_title('∂f/∂z (computed)')
+plt.colorbar(im5, ax=axes[1, 2], shrink=0.57)
+
+plt.tight_layout()
 plt.show()

dev/_downloads/07fcc19ba03226cd3d83d4e40ec44385/auto_examples_python.zip

@@ -22,7 +22,7 @@
       },
       "outputs": [],
       "source": [
-        "import torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\nfrom neuralop.losses.finite_diff import central_diff_2d  \n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
+        "import torch\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.animation as animation\nfrom neuralop.losses.differentiation import FiniteDiff  \n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
       ]
     },
     {
@@ -40,7 +40,7 @@
       },
       "outputs": [],
       "source": [
-        "## Simulation parameters\nLx, Ly = 2.0, 2.0   # Domain lengths\nnx, ny = 64, 64   # Grid resolution\nT = 1    # Total simulation time\ndt = 0.001  # Time step\nnu = 0.04   # Viscosity\n\n## Create grid\nX = torch.linspace(0, Lx, nx, device=device).repeat(ny, 1).T \nY = torch.linspace(0, Ly, ny, device=device).repeat(nx, 1)  \ndx = Lx / (nx-1)\ndy = Ly / (ny-1)\nnt = int(T / dt)\n\n## Initial condition \nu = -torch.sin(2 * np.pi * Y).to(device)\nv =  torch.cos(2 * np.pi * X).to(device)"
+        "## Simulation parameters\nLx, Ly = 2.0, 2.0   # Domain lengths\nnx, ny = 64, 64   # Grid resolution\nT = 1    # Total simulation time\ndt = 0.001  # Time step\nnu = 0.04   # Viscosity\n\n## Create grid\nX = torch.linspace(0, Lx, nx, device=device).repeat(ny, 1).T \nY = torch.linspace(0, Ly, ny, device=device).repeat(nx, 1)  \ndx = Lx / (nx-1)\ndy = Ly / (ny-1)\nnt = int(T / dt)\n\n## Initialize finite difference operator\nfd = FiniteDiff(dim=2, h=(dx, dy))\n\n## Initial condition \nu = -torch.sin(2 * np.pi * Y).to(device)\nv =  torch.cos(2 * np.pi * X).to(device)"
       ]
     },
     {
@@ -58,7 +58,7 @@
       },
       "outputs": [],
       "source": [
-        "u_evolution = [u.clone()]\nv_evolution = [v.clone()]\n\nfor _ in range(nt):\n    \n    # Compute first-order derivatives\n    u_x, u_y = central_diff_2d(u, [dx, dy])\n    v_x, v_y = central_diff_2d(v, [dx, dy])\n\n    # Compute second-order derivatives\n    u_xx, _ = central_diff_2d(u_x, [dx, dy])\n    _, u_yy = central_diff_2d(u_y, [dx, dy])\n    v_xx, _ = central_diff_2d(v_x, [dx, dy])\n    _, v_yy = central_diff_2d(v_y, [dx, dy])\n\n    # Evolve in time using Euler's method\n    u_next = u + dt * (-u * u_x - v * u_y + nu * (u_xx + u_yy))\n    v_next = v + dt * (-u * v_x - v * v_y + nu * (v_xx + v_yy))\n    \n    u, v = u_next.clone(), v_next.clone()\n    u_evolution.append(u.clone())\n    v_evolution.append(v.clone())\n\nu_evolution = torch.stack(u_evolution).cpu().numpy()\nv_evolution = torch.stack(v_evolution).cpu().numpy()"
+        "u_evolution = [u.clone()]\nv_evolution = [v.clone()]\n\nfor _ in range(nt):\n    \n    # Compute first-order derivatives\n    u_x = fd.dx(u)\n    u_y = fd.dy(u)\n    v_x = fd.dx(v)\n    v_y = fd.dy(v)\n\n    # Compute second-order derivatives\n    u_xx = fd.dx(u_x)\n    u_yy = fd.dy(u_y)\n    v_xx = fd.dx(v_x)\n    v_yy = fd.dy(v_y)\n\n    # Evolve in time using Euler's method\n    u_next = u + dt * (-u * u_x - v * u_y + nu * (u_xx + u_yy))\n    v_next = v + dt * (-u * v_x - v * v_y + nu * (v_xx + v_yy))\n    \n    u, v = u_next.clone(), v_next.clone()\n    u_evolution.append(u.clone())\n    v_evolution.append(v.clone())\n\nu_evolution = torch.stack(u_evolution).cpu().numpy()\nv_evolution = torch.stack(v_evolution).cpu().numpy()"
       ]
     },
     {

dev/_downloads/1bff93ede268f57f2123a573685b3499/plot_darcy_flow_spectrum.zip

@@ -15,7 +15,7 @@
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 
-from neuralop.losses.finite_diff import central_diff_2d  
+from neuralop.losses.differentiation import FiniteDiff  
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -45,6 +45,9 @@
 dy = Ly / (ny - 1)
 nt = int(T / dt)
 
+## Initialize finite difference operator
+fd = FiniteDiff(dim=2, h=(dx, dy))
+
 
 ## Initial condition and source term
 u = (-torch.sin(2 * np.pi * Y) * torch.cos(2 * np.pi * X)
@@ -63,9 +66,10 @@ def source_term(X, Y, t):
 for _ in range(nt):
 
     # Compute derivatives
-    u_x, u_y = central_diff_2d(u, [dx, dy])
-    u_xx, _ = central_diff_2d(u_x, [dx, dy])
-    _, u_yy = central_diff_2d(u_y, [dx, dy])
+    u_x = fd.dx(u)
+    u_y = fd.dy(u)
+    u_xx = fd.dx(u_x)
+    u_yy = fd.dy(u_y)
 
     # Evolve one step in time using Euler's method
     u = u + dt * (-cx * u_x - cy * u_y + nu * (u_xx + u_yy) + source_term(X, Y, t))