Poisson Equation Solution

main.py

from math import sqrt, fabs, ceil
import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np


class PoissonEquation(object):
    def __init__(self, step=0.1):
        try:
            self.step = float(step)
        except Exception:
            self.step = 0.1
        self.eps = self.step
        self.count_iterations = 0
        self.length_x = 1
        self.length_y = self.length_x
        self.size_x = ceil(self.length_x / self.step)
        self.size_y = ceil(self.length_y / self.step)

        self.curr_u = [[0 for x in range(self.size_y)]
                       for x in range(self.size_x)]
        self.prev_u = self.curr_u.copy()

    def f(self, i=0, j=0):
        x = i * self.step
        y = j * self.step
        return 2 * (x * x + y * y) - 2 * (x + y)

    @property
    def get_max_norm(self):
        _maxNorm = 0.0
        for i in range(self.size_y):
            for j in range(self.size_x):
                _maxNorm += fabs(self.curr_u[i][j] * self.curr_u[i][j])
        return sqrt(_maxNorm)

    def run(self):
        max_norm = 0.0
        self.count_iterations = 0

        while (fabs(max_norm - self.get_max_norm) > self.eps or
               self.count_iterations == 0):
            self.prev_u = self.curr_u.copy()

            for i in range(1, self.size_y - 1):
                for j in range(1, self.size_x - 1):
                    self.curr_u[i][j] = (self.prev_u[i - 1][j] +
                                         self.prev_u[i + 1][j] +
                                         self.prev_u[i][j - 1] +
                                         self.prev_u[i][j + 1] -
                                         self.step * self.step * self.f(i, j)
                                        ) / 4

            self.count_iterations += 1
            max_norm = self.get_max_norm

        return self.count_iterations

    def print_u(self):
        print('\n'.join([' '.join(['{:4}'.format(item) for item in row])
                         for row in self.curr_u]))

    def save_u(self, filename='data.txt'):
        file = open(filename, 'w+')
        for row in self.curr_u:
            file.write(' '.join(str(item) for item in row) + '\n')
        file.close()
        print('Data was written to file `%s`.' % filename)

    def plot_u(self):
        if self.step < 0.001:
            print('Graph will not be built. There will be saved to `data.txt`.')
            self.save_u()
            return

        X = np.arange(0, self.length_x, self.step)
        Y = np.arange(0, self.length_y, self.step)
        X, Y = np.meshgrid(X, Y)
        Z = np.array(self.curr_u)

        levels = MaxNLocator(nbins=15).tick_values(Z.min(), Z.max())
        cmap = plt.get_cmap('PiYG')

        plt.contourf(X, Y, Z, levels=levels, cmap=cmap)
        plt.colorbar()
        plt.title('Poisson Equation Solution (h = %.4f)' % self.step)
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.show()
        print('Graph was built.')


def main():
    step = input('Please, choose step for this method (by default: 0.1): ')
    poissonEquation = PoissonEquation(step=step)

    print('Count of iterations:', poissonEquation.run())
    print('Step:', poissonEquation.step)

    poissonEquation.plot_u()


if __name__ == '__main__':
    main()