DensityIntegrationUncertaintyQuantification.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Apr  3 15:51:35 2017

@author: jiachengzhang
"""

import numpy as np
import sys
import scipy.sparse as scysparse
import scipy.sparse.linalg as splinalg
import scipy.linalg as linalg


def Density_integration_Poisson_uncertainty(Xn,Yn,fluid_mask,grad_x,grad_y,dirichlet_label,dirichlet_value,
    uncertainty_quantification=True, sigma_grad_x=None, sigma_grad_y=None, sigma_dirichlet=None):
  # Evaluate the density field from the gradient fields by solving the Poisson equation.
  # The uncertainty of the density field can be quantified.
  """
  Inputs:
    Xn,Yn: 2d array of mesh grid. 
    fluid_mask: 2d array of binary mask of flow field. Boundary points are considered in the flow (mask should be True)
    grad_x, grad_y: 2d array of gradient field.
    dirichlet_label: 2d array of binary mask indicating the Dirichlet BC locations. Ohterwise Neumann. At least one point should be dirichlet.
    dirichlet_value: 2d array of diriclet BC values.
    uncertainty_quantification: True to perform the uncertainty quantification, False only perform integration.
    sigma_grad_x, sigma_grad_y, sigma_dirichlet: the uncertainty given as std of the input fields. 2d array of fields.
  Returns: 
    Pn: the integrated density field.
    sigma_Pn: the uncertainty (std) of the integrated density field. 
  """

  Ny, Nx = np.shape(Xn)
  dx = Xn[1,1] - Xn[0,0]
  dy = Yn[1,1] - Yn[0,0]
  invdx = 1.0/dx
  invdy = 1.0/dy
  invdx2 = invdx**2
  invdy2 = invdy**2

  fluid_mask_ex = np.zeros((Ny+2,Nx+2)).astype('bool')
  fluid_mask_ex[1:-1,1:-1] = fluid_mask

  grad_x_ex = np.zeros((Ny+2,Nx+2))
  grad_x_ex[1:-1,1:-1] = grad_x
  grad_y_ex = np.zeros((Ny+2,Nx+2))
  grad_y_ex[1:-1,1:-1] = grad_y
  dirichlet_label_ex = np.zeros((Ny+2,Nx+2)).astype('bool')
  dirichlet_label_ex[1:-1,1:-1] = dirichlet_label
  dirichlet_value_ex = np.zeros((Ny+2,Nx+2))
  dirichlet_value_ex[1:-1,1:-1] = dirichlet_value
  
  sigma_grad_x_ex = np.zeros((Ny+2,Nx+2))
  sigma_grad_y_ex = np.zeros((Ny+2,Nx+2))
  sigma_dirichlet_ex = np.zeros((Ny+2,Nx+2))
  if uncertainty_quantification==True:
    sigma_grad_x_ex[1:-1,1:-1] = sigma_grad_x
    sigma_grad_y_ex[1:-1,1:-1] = sigma_grad_y
    sigma_dirichlet_ex[1:-1,1:-1] = sigma_dirichlet

  # Generate the linear operator.
  j,i = np.where(fluid_mask_ex==True)
  Npts = len(j)
  fluid_index = -np.ones(fluid_mask_ex.shape).astype('int')
  fluid_index[j,i] = range(Npts)
  iC = fluid_index[j,i]
  iC_label = dirichlet_label_ex[j,i]
  iE = fluid_index[j,i+1]
  iW = fluid_index[j,i-1]
  iN = fluid_index[j+1,i]
  iS = fluid_index[j-1,i]

  LaplacianOperator = scysparse.csr_matrix((Npts,Npts),dtype=np.float)
  # RHS = np.zeros(Npts)
  # Var_RHS = np.zeros(Npts) # variance

  # Also generate the linear operator which maps from the grad_x, grad_y, dirichlet val to the rhs.
  Map_grad_x = scysparse.csr_matrix((Npts,Npts),dtype=np.float)
  Map_grad_y = scysparse.csr_matrix((Npts,Npts),dtype=np.float)
  Map_dirichlet_val = scysparse.csr_matrix((Npts,Npts),dtype=np.float)

  # First, construct the linear operator and RHS as if all they are all Nuemanan Bc

  # if the east and west nodes are inside domain
  loc = (iE!=-1)*(iW!=-1)
  LaplacianOperator[iC[loc],iC[loc]] += -2.0*invdx2
  LaplacianOperator[iC[loc],iE[loc]] += +1.0*invdx2
  LaplacianOperator[iC[loc],iW[loc]] += +1.0*invdx2
  # RHS[iC[loc]] += (grad_x_ex[j[loc],i[loc]+1] - grad_x_ex[j[loc],i[loc]-1])*(invdx*0.5)
  # Var_RHS[iC[loc]] += (invdx*0.5)**2 * (sigma_grad_x_ex[j[loc],i[loc]+1]**2 + sigma_grad_x_ex[j[loc],i[loc]-1]**2)
  Map_grad_x[iC[loc],iE[loc]] += invdx*0.5
  Map_grad_x[iC[loc],iW[loc]] += -invdx*0.5

  # if the east node is ouside domian
  loc = (iE==-1)
  LaplacianOperator[iC[loc],iC[loc]] += -2.0*invdx2
  LaplacianOperator[iC[loc],iW[loc]] += +2.0*invdx2
  # RHS[iC[loc]] += -(grad_x_ex[j[loc],i[loc]] + grad_x_ex[j[loc],i[loc]-1])*invdx
  # Var_RHS[iC[loc]] += invdx**2 * (sigma_grad_x_ex[j[loc],i[loc]]**2 + sigma_grad_x_ex[j[loc],i[loc]-1]**2)
  Map_grad_x[iC[loc],iC[loc]] += -invdx
  Map_grad_x[iC[loc],iW[loc]] += -invdx

  # if the west node is ouside domian
  loc = (iW==-1)
  LaplacianOperator[iC[loc],iC[loc]] += -2.0*invdx2
  LaplacianOperator[iC[loc],iE[loc]] += +2.0*invdx2
  # RHS[iC[loc]] += (grad_x_ex[j[loc],i[loc]] + grad_x_ex[j[loc],i[loc]+1])*invdx
  # Var_RHS[iC[loc]] += invdx**2 * (sigma_grad_x_ex[j[loc],i[loc]]**2 + sigma_grad_x_ex[j[loc],i[loc]+1]**2)
  Map_grad_x[iC[loc],iC[loc]] += invdx
  Map_grad_x[iC[loc],iE[loc]] += invdx

  # if the north and south nodes are inside domain
  loc = (iN!=-1)*(iS!=-1)
  LaplacianOperator[iC[loc],iC[loc]] += -2.0*invdy2
  LaplacianOperator[iC[loc],iN[loc]] += +1.0*invdy2
  LaplacianOperator[iC[loc],iS[loc]] += +1.0*invdy2
  # RHS[iC[loc]] += (grad_y_ex[j[loc]+1,i[loc]] - grad_y_ex[j[loc]-1,i[loc]])*(invdy*0.5)
  # Var_RHS[iC[loc]] += (invdy*0.5)**2 * (sigma_grad_y_ex[j[loc]+1,i[loc]]**2 + sigma_grad_y_ex[j[loc]-1,i[loc]]**2)
  Map_grad_y[iC[loc],iN[loc]] += invdy*0.5
  Map_grad_y[iC[loc],iS[loc]] += -invdy*0.5

  # if the north node is ouside domian
  loc = (iN==-1)
  LaplacianOperator[iC[loc],iC[loc]] += -2.0*invdy2
  LaplacianOperator[iC[loc],iS[loc]] += +2.0*invdy2
  # RHS[iC[loc]] += -(grad_y_ex[j[loc],i[loc]] + grad_y_ex[j[loc]-1,i[loc]])*invdy
  # Var_RHS[iC[loc]] += invdy**2 * (sigma_grad_y_ex[j[loc],i[loc]]**2 + sigma_grad_y_ex[j[loc]-1,i[loc]]**2)
  Map_grad_y[iC[loc],iC[loc]] += -invdy
  Map_grad_y[iC[loc],iS[loc]] += -invdy

  # if the south node is ouside domian
  loc = (iS==-1)
  LaplacianOperator[iC[loc],iC[loc]] += -2.0*invdy2
  LaplacianOperator[iC[loc],iN[loc]] += +2.0*invdy2
  # RHS[iC[loc]] += (grad_y_ex[j[loc],i[loc]] + grad_y_ex[j[loc]+1,i[loc]])*invdy
  # Var_RHS[iC[loc]] += invdy**2 * (sigma_grad_y_ex[j[loc],i[loc]]**2 + sigma_grad_y_ex[j[loc]+1,i[loc]]**2)
  Map_grad_y[iC[loc],iC[loc]] += invdy
  Map_grad_y[iC[loc],iN[loc]] += invdy

  # Then change the boundary conidtion at locatiosn of Dirichlet.
  loc = (iC_label==True)
  LaplacianOperator[iC[loc],:] = 0.0
  LaplacianOperator[iC[loc],iC[loc]] = 1.0*invdx2
  # RHS[iC[loc]] = dirichlet_value_ex[j[loc],i[loc]] * invdx2
  # Var_RHS[iC[loc]] = sigma_dirichlet_ex[j[loc],i[loc]]**2 * invdx2**2
  Map_grad_x[iC[loc],:] = 0.0
  Map_grad_y[iC[loc],:] = 0.0
  Map_dirichlet_val[iC[loc],iC[loc]] = 1.0*invdx2

  LaplacianOperator.eliminate_zeros()
  Map_grad_x.eliminate_zeros()
  Map_grad_y.eliminate_zeros()
  Map_dirichlet_val.eliminate_zeros()

  # Solve for the field.
  grad_x_vect = grad_x_ex[j,i]
  grad_y_vect = grad_y_ex[j,i]
  dirichlet_val_vect = dirichlet_value_ex[j,i]
  RHS = Map_grad_x.dot(grad_x_vect) + Map_grad_y.dot(grad_y_vect) + Map_dirichlet_val.dot(dirichlet_val_vect)
  # Solve the linear system
  Pn_ex = np.zeros((Ny+2,Nx+2))
  p_vect = splinalg.spsolve(LaplacianOperator, RHS)
  Pn_ex[j,i] = p_vect
  
  # Uncertainty propagation
  sigma_Pn_ex = np.zeros((Ny+2,Nx+2))
  if uncertainty_quantification==True:
    # Propagate to get the covariance matrix for RHS
    Cov_grad_x = scysparse.diags(sigma_grad_x_ex[j,i]**2,shape=(Npts,Npts),format='csr')
    Cov_grad_y = scysparse.diags(sigma_grad_y_ex[j,i]**2,shape=(Npts,Npts),format='csr')
    Cov_dirichlet_val = scysparse.diags(sigma_dirichlet_ex[j,i]**2,shape=(Npts,Npts),format='csr')
    Cov_RHS = Map_grad_x*Cov_grad_x*Map_grad_x.transpose() + Map_grad_y*Cov_grad_y*Map_grad_y.transpose() + \
              Map_dirichlet_val*Cov_dirichlet_val*Map_dirichlet_val.transpose()

    Laplacian_inv = linalg.inv(LaplacianOperator.A)
    Cov_p = np.matmul(np.matmul(Laplacian_inv, Cov_RHS.A), Laplacian_inv.T)
    Var_p_vect = np.diag(Cov_p)

    sigma_Pn_ex[j,i] = Var_p_vect**0.5

  return Pn_ex[1:-1,1:-1], sigma_Pn_ex[1:-1,1:-1]