# -*- coding: utf-8 -*-
"""
Created on Thu Jul 22 09:26:54 2021
@author: Rasmus Vest Nielsen
"""
## Load libraries
import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import spsolve # import with from to use its short name
import matplotlib.pyplot as plt # combines namespace of numpy and pyplot
## Setup of the model by specifying its dimensions
# Specify a rectangular grid
x = np.arange(-1000., 1000., 25.)
y = np.arange(-1000., 1000., 25.) # backward, i.e. first row grid line has highest y
z = np.arange(-100., 0., 20.) # backward, i.e. from top to bottom
# As well as the number of cells along the three axes
Nx = len(x)-1
Ny = len(y)-1
Nz = len(z)-1
# Grid
sz = (Nz,Ny,Nx) # the shape of the model
Nod = np.prod(sz) # total number of cells in the model
# from this we have the width of columns, rows and layers
dx = np.diff(x).reshape(1, 1, Nx)
dy = np.diff(y).reshape(1, Ny, 1)
dz = np.abs(np.diff(z)).reshape(Nz, 1,1)
IBOUND = np.ones(sz)
IBOUND[:,-1,:] = -1 # last row of model heads are prescribed
IBOUND[:, 40:45, 20:70]=0 # these cells are inactive
active = (IBOUND>0).reshape(Nod) # active is now a vector of booleans of length Nod
inact = (IBOUND==0).reshape(Nod) # dito for inact
fxhd = (IBOUND<0).reshape(Nod) # dito for fxhd
k = 10.0 # m/d uniform conductivity
kx = k * np.ones(sz) # [L/T] 3D kx array
ky = k * np.ones(sz) # [L/T] 3D ky array with same values as kx
kz = k * np.ones(sz) # [L/T] 3D kz array with same values as kx
# half cell flow resistances
Rx = 0.5 * dx / (dy * dz) / kx # [T/L2], flow resistance half cell in x-direction
Ry = 0.5 * dy / (dz * dx) / ky # same in y-direction
Rz = 0.5 * dz / (dx * dy) / kz # same in z-direction
Rx = Rx.reshape(Nod,); Rx[inact] = np.Inf; Rx=Rx.reshape(sz)
Ry = Ry.reshape(Nod,); Ry[inact] = np.Inf; Ry=Ry.reshape(sz)
Rz = Rz.reshape(Nod,); Rz[inact] = np.Inf; Rz=Rz.reshape(sz)
# Conductances between adjacent cells
Cx = 1 / (Rx[:, :, :-1] + Rx[:, :,1:]) # [L2/T] in x-direction
Cy = 1 / (Ry[:, :-1,:] + Ry[:, 1:,:]) # idem in y-direction
Cz = 1 / (Rz[:-1,:,:] + Rz[1:,:,:]) # idem in z-direction
NOD = np.arange(Nod).reshape(sz) # this is a full 3D array of node numbers (cell numbers)
IE = NOD[:, :, 1:] # numbers of the eastern neighbors of each cell
IW = NOD[:, :, :-1] # same western neighbors
IN = NOD[:, :-1,:] # same northern neighbors
IS = NOD[:, 1:,:] # southern neighbors
IT = NOD[:-1,:,:] # top neighbors
IB = NOD[1:,:,:] # bottom neighbors
R = lambda x : x.ravel() # define short hand for x.ravel()
# notice the call signature:
# csc_matrix( (data, (row_index, col_index) ), (M,N)); This is a tuple within tuple.
A = sp.csc_matrix(( np.concatenate(( R(Cx), R(Cx), R(Cy), R(Cy), R(Cz), R(Cz)) ),\
(np.concatenate(( R(IE), R(IW), R(IN), R(IS), R(IB), R(IT)) ),\
np.concatenate(( R(IW), R(IE), R(IS), R(IN), R(IT), R(IB)) ),\
)),(Nod,Nod))
# to use the vector of diagonal values int a call of sp.diags() we need to have it aa a
# standard nondimensional numpy vector.
# To get this:
# - first turn the matrix obtained by A.sum(axis=1) into a np.array by np.array( .. )
# - then take the whole column to loose the array orientation (to get a dimensionless numpy vector)
adiag = np.array(-A.sum(axis=1))[:,0]
Adiag = sp.diags(adiag)
## Boundary conditions
FQ = np.zeros(sz) # all flows zero. Note sz is the shape of the model grid
FQ[2, 30, 25] = -1200 # [m3/d] extraction in this cell
RHS = FQ.reshape(Nod)
# Fixed heads
HI = np.zeros(sz)
RHS = FQ.reshape(Nod,1) - A[:,fxhd].dot(HI.reshape(Nod,1)[fxhd])
# Solving the matrix equation for the unknown heads
Phi = HI.flatten()
Phi[active] = spsolve( (A+Adiag)[active][:,active] ,RHS[active] )
Phi[inact] = np.NaN
Phi=Phi.reshape(sz) # reshape vector Phi to 3D shape of the grid
# Plotting the heads as contours
xm = 0.5 * (x[:-1] + x[1:]) # [L] coordinates of column centers
ym = 0.5 * (y[:-1] + y[1:]) # [L] coordinates of row centers
layer = 2 # contours for this layer
nc = 50 # number of contours in total
plt.xlabel('x [m]')
plt.ylabel('y [m]')
plt.title("Contours (%d in total) of the head in layer %d with inactive section" % (nc, layer))
plt.contour(xm, ym, Phi[layer], nc)