Ejemplo n.º 1
0
def test_residual(x, gt, bc):
    y = utils.fd_step(x, bc, None)
    residual = y - x
    e = gt - x
    # Ae = -r
    A = utils.loss_kernel.view(1, 1, 3, 3)
    Ae = F.conv2d(e.unsqueeze(1), A).squeeze(1)
    # z should be all zeros
    z = Ae + residual[:, 1:-1, 1:-1]
    print(z)

    # Solve Ae = -r iteratively
    e = residual
    f = -residual
    for i in range(400):
        e = utils.fd_step(e, torch.zeros(1, 4), f)

    Ae = F.conv2d(e.unsqueeze(1), A).squeeze(1)
    # z should be all zeros
    z = Ae + residual[:, 1:-1, 1:-1]
    #print(z)

    # e = Se
    e = utils.fd_step(e, torch.zeros(1, 4), None)
    final = y + e
    print(torch.abs(gt - final))
    def multigrid_step(self, x, bc, f, step):
        '''
    One layer of multigrid. Recursive function.
    Find solution x to Ax + b = 0.
    '''
        batch_size, image_size, _ = x.size()
        # Pre smoothing
        for i in range(self.pre_smoothing):
            x = utils.fd_step(x, bc, f)

        if step > 1:
            # Downsample
            if f is not None:
                f_sub = 4 * utils.subsample(f)
            else:
                f_sub = None

            if self.is_bc_mask:
                # Subsample geometry
                bc_sub = utils.subsample(
                    bc.view(batch_size * 2, image_size, image_size))
                bc_sub = bc_sub.view(batch_size, 2, *bc_sub.size()[-2:])
            else:
                bc_sub = bc

            x_sub = utils.restriction(x, bc_sub)
            # Refine x_sub recursively
            x_sub = self.multigrid_step(x_sub, bc_sub, f_sub, step - 1)
            # Upsample
            x = utils.interpolation(x_sub, bc)

        # Post smoothing
        for i in range(self.post_smoothing):
            x = utils.fd_step(x, bc, f)
        return x
Ejemplo n.º 3
0
def get_solution(x, bc, f):
    '''
  Iterate until error is below a threshold.
  '''
    frames = [x]

    error_threshold = 0.00001
    max_iters = 8000

    error = utils.fd_error(x, bc, f)
    largest_error = error.max().item()
    print('largest error {}'.format(largest_error))
    if largest_error >= error_threshold:
        # Iterate with Jacobi until ground truth
        for i in range(max_iters):
            x = utils.fd_step(x, bc, f)
            error = utils.fd_error(x, bc, f)
            if (i + 1) % 100 == 0:
                largest_error = error.max().item(
                )  # largest error in the batch
                print('Iter {}: largest error {}'.format(i + 1, largest_error))
                if largest_error < error_threshold:
                    break

    # Add ground truth to frames
    y = x.cpu().numpy()
    frames.append(y)
    frames = np.stack(
        frames,
        axis=1)  # batch_size x (n_frames + 1) x image_size x image_size
    return frames
Ejemplo n.º 4
0
def test_heat():
    image_size = 65
    scale = np.random.uniform(350, 450) / (image_size**2)
    f = -gaussian(image_size) * scale
    f = torch.Tensor(f).unsqueeze(0)
    f = utils.pad_boundary(f, torch.zeros(1, 4))

    bc = torch.Tensor(np.random.rand(1, 4) * 80)
    x = torch.zeros(1, image_size + 2, image_size + 2)
    x = utils.set_boundary(x, bc)
    x = utils.initialize(x, bc, 'avg')

    y = x.clone()
    for i in range(2000):
        y = utils.fd_step(y, bc, None)

    z = x.clone()
    for i in range(4000):
        z = utils.fd_step(z, bc, f)

    # Au = 0
    A = utils.loss_kernel.view(1, 1, 3, 3)
    r = F.conv2d(y.unsqueeze(1), A).squeeze(1)
    error = torch.abs(r).max().item()
    print(error)

    # Au = f
    A = utils.loss_kernel.view(1, 1, 3, 3)
    r = F.conv2d(z.unsqueeze(1), A).squeeze(1) - f[:, 1:-1, 1:-1]
    error = torch.abs(r).max().item()
    print(error)

    y = (y / 100).numpy().squeeze(0)
    z = (z / 100).numpy().squeeze(0)

    plt.imshow(y)
    plt.colorbar()
    plt.show()

    plt.imshow(z)
    plt.colorbar()
    plt.show()
    def multigrid_step(self, x, bc, f, step):
        '''
    One layer of multigrid. Recursive function.
    Find solution x to Ax + b = f.
    Algorithm:
      - Update rule: u^{k+1} = S u^{k} + b - f
      - Residual r^{k} = u^{k+1} - u^{k} = A u^{k} + b - f
      - Solve A e^{k} = - r^{k} recursively.
      - u' = u^{k} + e^{k}
    '''
        if step == 0:
            return None

        # Pre smoothing
        x = utils.set_boundary(x, bc)
        for i in range(self.pre_smoothing):
            x = utils.fd_step(x, bc, f)
        # Calculate residual
        y = utils.fd_step(x, bc, f)
        r = y - x

        # Solve e: A e = -r
        # Restriction: downsample by 2
        zeros_bc = torch.zeros(1, 4)
        r_sub = utils.restriction(r, zeros_bc)

        # Recursive
        ek_sub = self.multigrid_step(r_sub, zeros_bc, -r_sub, step - 1)

        # Upsample
        if ek_sub is not None:
            ek = utils.interpolation(ek_sub, zeros_bc)
            # Add to x
            x = x + ek

        # Post smoothing
        x = utils.set_boundary(x, bc)
        for i in range(self.post_smoothing):
            x = utils.fd_step(x, bc, f)
        return x
Ejemplo n.º 6
0
def test_upsampling_poisson(x, gt, bc, f):
    print('Upsampling multigrid')
    f_sub = utils.subsample(f)
    x_sub = utils.restriction(x, bc)
    for i in range(1000):
        x_sub = utils.fd_step(x_sub, bc, f_sub)

    # Upsample
    x = utils.interpolation(x_sub, bc)

    A = utils.loss_kernel.view(1, 1, 3, 3)
    r = F.conv2d(x.unsqueeze(1), A).squeeze(1)
    r = utils.pad_boundary(r, torch.zeros(1, 4)) - f
    r = r.cpu().numpy()
    print(r.max())
Ejemplo n.º 7
0
def test_subsampling_poisson(x, gt, bc, f):
    print('Subsampling multigrid')
    for i in range(2000):
        x = utils.fd_step(x, bc, f)

    A = utils.loss_kernel.view(1, 1, 3, 3)
    r = F.conv2d(x.unsqueeze(1), A).squeeze(1)
    r = utils.pad_boundary(r, torch.zeros(1, 4)) - f
    print(np.abs(r.cpu().numpy()).max())

    # Subsample
    x_sub = x
    f_sub = f
    for i in range(3):
        f_sub = 4 * utils.subsample(f)
        x_sub = utils.restriction(x, bc)
        r_sub = F.conv2d(x_sub.unsqueeze(1), A).squeeze(1)
        r_sub = utils.pad_boundary(r_sub, torch.zeros(1, 4)) - f_sub
        print(x_sub.size())
        print(np.abs(r_sub.cpu().numpy()).max())
 def forward(self, x, bc, f):
     '''
 x: size (batch_size x image_size x image_size)
 return: same size
 '''
     return utils.fd_step(x, bc, f)