def modifiedCrankNicolson(getU0=None,
fDeriv=fZero,
scheme=None,
T=10,
k=1,
shape=None,
isNeumannFunc=None,
schemeNeumannFunc=None,
g=gDefault):
Amatrix, Finterior, Fboundary, geometry = FiniteDifference.getSystem(
shape=shape,
f=fZero,
g=gOne,
scheme=scheme,
isNeumannFunc=isNeumannFunc,
schemeNeumannFunc=schemeNeumannFunc,
interior=True)
if shape.dim == 1:
def getBoundaryVals(t):
return Shape.getVectorLattice(g(Shape.getMeshGrid(shape), t),
shape)
U_0 = Shape.getVectorLattice(getU0(Shape.getMeshGrid(shape)),
shape)
else:
def getBoundaryVals(t):
return Shape.getVectorLattice(g(*Shape.getMeshGrid(shape), t),
shape)
U_0 = Shape.getVectorLattice(getU0(*Shape.getMeshGrid(shape)),
shape)
return ModifiedCrankNicolson.modifiedCrankNicolsonSolver(
U_0,
fDeriv,
T=T,
k=k,
diffOperator=Amatrix,
boundaryCoeffs=-Fboundary,
g=getBoundaryVals,
geometry=geometry)
def getDiseaseModelF(getBeta, getGamma, shape):
# currying
if shape.dim == 1:
betaLattice = getBeta(Shape.getMeshGrid(shape))
gammaLattice = getGamma(Shape.getMeshGrid(shape))
else:
betaLattice = getBeta(*Shape.getMeshGrid(shape))
gammaLattice = getGamma(*Shape.getMeshGrid(shape))
if np.size(betaLattice) == 1:
beta = betaLattice
else:
beta = Shape.getVectorLattice(betaLattice, shape)
if np.size(gammaLattice) == 1:
gamma = gammaLattice
else:
gamma = Shape.getVectorLattice(gammaLattice, shape)
def F(U):
S, I = np.split(U, 2)
return np.concatenate((-beta * S * I, beta * S * I - gamma * I))
return F
def getBoundaryVals(t):
return Shape.getVectorLattice(g(*Shape.getMeshGrid(shape), t),
shape)
def __init__(self,
getU0_S,
getU0_I,
muS,
muI,
schemeS,
getBeta,
getGamma,
T=10,
k=1,
N=4,
isBoundaryFunction=None,
dim=2,
length=1,
origin=0,
isNeumannFunc=None,
schemeNeumannFunc=None,
g=gDefault):
"""
:param scheme: function returning array of coefficients.
:param f: function returning time derivative.
:param g: function returning boundry conditions.
:param isBoundaryFunction: return true if point is on boundary
:param isBoundaryFunction:
:param length: length of sides
:param origin: for plotting
:param isNeumannFunc: Function Returning true if point has Neumann conditions
:param schemeNeumannFunc: Scheme for Neumann conditions on that point.
"""
self.T, self.k = T, k
self.shapeObject = Shape(N=N,
isBoundaryFunc=isBoundaryFunction,
dim=dim,
length=length,
origin=origin)
AmatrixS, FinternalS, FboundaryS, self.geometryS = FiniteDifference.getSystem(
shape=self.shapeObject,
f=fZero,
g=gOne,
scheme=schemeS,
isNeumannFunc=isNeumannFunc,
schemeNeumannFunc=schemeNeumannFunc,
interior=True)
self.geometryI = self.geometryS
self.Fboundary = np.concatenate((muS * FboundaryS, muI * FboundaryS))
self.diffOperator = sparse.bmat(
[[AmatrixS * muS, None], [None, AmatrixS * muI]], format="csc")
geometrySI = np.concatenate((self.geometryS, self.geometryI), axis=1)
domainGeometrySI = FiniteDifference.getDomainGeometry(geometrySI)
self.domainSize = len(domainGeometrySI[0])
# Assuming R = 0 at t = 0
if self.shapeObject.dim == 1:
I_0 = Shape.getVectorLattice(
getU0_I(Shape.getMeshGrid(self.shapeObject)), self.shapeObject)
S_0 = Shape.getVectorLattice(
getU0_S(Shape.getMeshGrid(self.shapeObject)), self.shapeObject)
else:
I_0 = Shape.getVectorLattice(
getU0_I(*Shape.getMeshGrid(self.shapeObject)),
self.shapeObject)
S_0 = Shape.getVectorLattice(
getU0_S(*Shape.getMeshGrid(self.shapeObject)),
self.shapeObject)
self.U_0 = np.concatenate(
(S_0, I_0))[np.logical_or(geometrySI[0], geometrySI[1])]
diseaseModelF = DiseaseModel.getDiseaseModelF(getBeta, getGamma,
self.shapeObject)
self.UList, self.times = ModifiedCrankNicolson.modifiedCrankNicolsonSolver(
self.U_0,
f=diseaseModelF,
T=T,
k=k,
diffOperator=self.diffOperator,
boundaryCoeffs=-self.Fboundary,
geometry=geometrySI,
g=gDefault)