#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
This program solves a 2D Poisson problem on the unit square, using the Scipy linear solver.

Uses the module FDM

@author: drager
"""


import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from numpy import pi, sin , cos
from matplotlib import cm
import FDM


def myfun(x,y):
    return sin((x**2+y**2)*2*pi)
    
n=100
h = 1/(n+1)

uvect = FDM.solution_vect(n, myfun)
ugrid = uvect.reshape((n,n))


xx,yy = np.ogrid[h:1:n*1j, h:1:n*1j]

stride = 10

plt.ion()
fig = plt.figure()
ax = Axes3D(fig)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
ax.plot_surface(xx,yy, ugrid, cmap = cm.Blues
                , alpha=0.85, rstride=stride, cstride = stride)
ax.elev, ax.azmin = 60, 45
plt.show()