HW · ME 595R

HW 8: Convolutional neural net for super resolution (with physics)

due 3/13/2025 before midnight via Learning Suite 25 possible points

We will use convolutional neural nets to perform super resolution of coarse MRI data using the methodology in this paper. We’ll just focus on the very first case (summarized in Figure 3). Read the Methodology section (section II) and Known Boundary Condition (III.B.1), which is the case we’ll focus on. The goal is to reproduce Fig 3c (CNN).

You can download the data in these two files: sr_lfdata.npy, and sr_hfdata.npy and load as follows:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm

# load low resolution data, which serves as input to our model
lfdata = np.load("sr_lfdata.npy")
lfx = lfdata[0, :, :]  # size 14 x 9  (height x width)
lfy = lfdata[1, :, :]
lfu = lfdata[4, :, :]
lfv = lfdata[5, :, :]

# plot the low resolution data (like fig 3a except we are using MRI noise here rather than Gaussian noise so it will look a bit different)
plt.figure()
plt.pcolormesh(lfx, lfy, np.sqrt(lfu**2 + lfv**2), cmap=cm.coolwarm, vmin=0.0, vmax=1.0)
plt.colorbar()

# load high resolution grids and mapping from low resolution to high resolution grid
hfdata = np.load("sr_hfdata.npy")
Jinv = hfdata[0, :, :]  # size 77 x 49 (height x width)
dxdxi = hfdata[1, :, :]
dxdeta = hfdata[2, :, :]
dydxi = hfdata[3, :, :]
dydeta = hfdata[4, :, :]
hfx = hfdata[5, :, :]
hfy = hfdata[6, :, :]

ny, nx = hfx.shape  #(77 x 49)
h = 0.01  # grid spacing in high fidelity (needed for derivatives)

plt.show()

I’ve also done (part of) the differentiation for you in the code below. The convolution filters shown in Appendix A aren’t actually used because they won’t work at the boundaries of the grid. We still use central finite differencing in the interior, but use one-sided finite differencing on the edges. The only difference from finite differencing you’ve done before is that we are using a larger stencil (5 pts for interior, 4 for edges). Note that the below code takes derivatives in the transformed Cartesian coordinates. You then need to apply the coordinate transformation using the formulas shown in eq 10a/b using the provided tensors above (Jinv which is written as 1/J in the text, dxdxi, dxdeta, dydxi, dydeta). Please note that the below assumes a grid given in height x width (not ‘x’ x ‘y’) just to match the ordering in pytorch convolution layers.

# see https://en.wikipedia.org/wiki/Finite_difference_coefficient
# or https://web.media.mit.edu/~crtaylor/calculator.html

# f should be a tensor of size: nbatch x nchannels x height (y or eta) x width (x or xi)
# This is written in a general way if one had more data, but for this case there is only 1 data sample, and there are only a few channels it might be clearer to you to separate the channels out into separate variables, in which case the below could be simplified (i.e., you remove the first two dimensions from everything so that input is just height x width if you desire).
def ddxi(f, h):
    # 5-pt stencil
    dfdx_central = (f[:, :, :, 0:-4] - 8*f[:, :, :, 1:-3] + 8*f[:, :, :, 3:-1] - f[:, :, :, 4:]) / (12*h)
    # 1-sided 4pt stencil
    dfdx_left = (-11*f[:, :, :, 0:2] + 18*f[:, :, :, 1:3] -9*f[:, :, :, 2:4] + 2*f[:, :, :, 3:5]) / (6*h)
    dfdx_right = (-2*f[:, :, :, -5:-3] + 9*f[:, :, :, -4:-2] -18*f[:, :, :, -3:-1] + 11*f[:, :, :, -2:]) / (6*h)

    return torch.cat((dfdx_left, dfdx_central, dfdx_right), dim=3)

def ddeta(f, h):
    # 5-pt stencil
    dfdy_central = (f[:, :, 0:-4, :] - 8*f[:, :, 1:-3, :] + 8*f[:, :, 3:-1, :] - f[:, :, 4:, :]) / (12*h)
    # 1-sided 4pt stencil
    dfdy_bot = (-11*f[:, :, 0:2, :] + 18*f[:, :, 1:3, :] -9*f[:, :, 2:4, :] + 2*f[:, :, 3:5, :]) / (6*h)
    dfdy_top = (-2*f[:, :, -5:-3, :] + 9*f[:, :, -4:-2, :] -18*f[:, :, -3:-1, :] + 11*f[:, :, -2:, :]) / (6*h)

    return torch.cat((dfdy_bot, dfdy_central, dfdy_top), dim=2)

Tips:

Although they used 3 separate networks that each produce 1 output channel (as shown in Fig 1), I just used 1 network with 3 output channels. I can be effective to learn shared features if output quantities are related. That would in general require adjusting the layer dimensions since we’ve reduced the available parameters with only 1 network. In this particular case I found that the default hyperparameters were still sufficient for the 1-network-3-channel-output version.
To better understand the two coordinate systems compare figure 2 and figure 9. Figure 2 is the actual physical space (x and y), and figure 9 is the same grid but mapped to a Cartesian space (\(\xi\) and \(\eta\)). We do most of our work in the Cartesian space where it is easier to compute finite difference derivatives, impose boundary conditions, and apply convolution filters. We then just need to map back to the regular space to impose our physics loss function (and plot). Mapping the velocities and pressures back is automatic as the grid connectivity is the same, and mapping derivatives back is what equations 10 a/b do (referred to earlier).
The paper only states some of the boundary conditions, and it does not explicitly state how they are imposed (because they are well known for the audience of this journal). So, I expand on them here. To impose boundary conditions, we simply set values for the velocities and pressures on the edges (bottom, left, top, right). The bottom edge (\(\eta\) = 0) is inflow with conditions: \(u = 0, v = 1, dp/d\eta = 0\). So if u, v, p were tensors of size neta x nxi we would set: u[0, :] = 0; v[0, :] = 1; p[0, :] = p[1, :]. The latter forces the pressure gradient to be zero at the inlet (which just means it is at some unknown constant pressure). The left and right edges are walls with conditions: \(u = 0, v = 0, dp/d\xi = 0\) (the latter is a result from incompressible boundary layer theory). At the top (outlet) we set \(du/d\eta = 0, dv/d\eta = 0, p = 0\) (the outlet pressure is unknown, but pressure is only defined relative to a reference point so we arbitrarily choose the outlet as a zero reference).
There are two ways you could do this: 1) predict results on a 77x49 grid then overwrite all the boundary values predicted by the network, or 2) predict results on a 75x47 grid then zero pad to create the boundaries (see nn.ConstantPad2d). The size would now be 77x49 and you can then modify those added boundary values as needed (many of the boundary conditions are zero and so won’t require further modification).
In your neural network you first step should be to upsample the coarse input using bicubic interpolation like shown in the paper. That is your best estimate of the flow field on the high dimensional grid, then apply convolutional layers from there. Using your best estimate will allow the neural net portion to learn a simpler function. You can upscale using nn.Upsample with mode=’bicubic`. The size you upscale to depends on whether you add boundaries or modify boundaries as discussed in the prior bullet.
In the case we are analyzing, the only loss comes from the physics residuals shown in eq (4). I’ve expanded them in Cartesian coordinates for you below (and I’ve also substituted in the density and viscosity for this particular problem).
\[\frac{du}{dx} + \frac{dv}{dy} = 0\\ u\frac{du}{dx} + v\frac{du}{dy} + \frac{dp}{dx} - 0.01 \left(\frac{d^2u}{dx^2} + \frac{d^2u}{dy^2}\right) = 0\\ u\frac{dv}{dx} + v\frac{dv}{dy} + \frac{dp}{dy} - 0.01 \left(\frac{d^2v}{dx^2} + \frac{d^2v}{dy^2}\right) = 0\]

To plot your final results you can use the following where u and v are your predicted flow fields from your network.

  plt.figure()
  plt.pcolormesh(hfx, hfy, np.sqrt(u**2 + v**2), cmap=cm.coolwarm, vmin=0.0, vmax=1.3)
  plt.colorbar()

After finishing, it would be worth skimming through section III.C (at least III.C.1). We just learned to super resolve 1 data instance through an optimization process. But if we instead parameterized the inflow, and trained a model with multiple datasets, we could instead learn a model that could super resolve any MRI input without further training (just model evaluation).