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 applyDiseaseAnimation(self, axS, axI, animationLength=10, **kvargs): artistlist = [] fps = 24 numFrames = animationLength * fps numTimes = self.UList.shape[0] step = numTimes // numFrames max = np.nanmax(self.UList) for timeIndex in range(0, len(self.times), step): S, I = self.UList[timeIndex, :self.domainSize // 2], self.UList[timeIndex, self.domainSize // 2:] artistS = Shape.plotImage2d(S, self.shapeObject, ax=axS, geometry=self.geometryS, title=f"Susceptible", animated=True, colorbar=False, **kvargs) artistI = Shape.plotImage2d(I, self.shapeObject, ax=axI, geometry=self.geometryI, title=f"Infected", animated=True, colorbar=False, **kvargs) artistlist.append([artistS, artistI]) if timeIndex == 0: axI.figure.colorbar(artistI, ax=[axS, axI]) return animation.ArtistAnimation(axI.figure, artistlist, blit=True, interval=1000 / fps)
def plotImage(self, timeIndex=0, ax=None, group="I", show=False, animated=False, **kwargs): if group == "S": U = self.UList[int(timeIndex), :self.domainSize // 2] elif group == "I": U = self.UList[int(timeIndex), self.domainSize // 2:] else: print(f"Group: {group} not found") return None max = np.nanmax(U) return Shape.plotImage2d(U, self.shapeObject, ax=ax, show=show, geometry=self.geometryS, animated=animated, **kwargs)
def plot(self, timeIndex=0, ax=None, group="I", show=False, view=None, **kwargs): if group == "S": U = self.UList[int(timeIndex), :self.domainSize // 2] title = f"Susceptible at t = {timeIndex * self.k}" elif group == "I": U = self.UList[int(timeIndex), self.domainSize // 2:] title = f"Infected at t = {timeIndex * self.k}" else: print(f"Group: {group} not found") return None return Shape.plotOnShape(U, self.shapeObject, ax=ax, show=show, view=view, geometry=self.geometryS, title=title, **kwargs)
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
# Starting conditions for number of infected def getU_0(x): return np.exp(0 * a) * (x) def fDeriv(u): return a * u fig, (ax1, ax2) = plt.subplots(1,2, figsize=(13,4)) for k in kListSmall: diffList = [] for N in NList: shapeObject = Shape(dim=dim, N=N) # Inserting into MCN object that solves the problem. Note neumann conditions are UList, times = ModifiedCrankNicolson.modifiedCrankNicolson(getU0=getU_0, fDeriv=fDeriv, scheme=schemes.makeLaplacian1D(1), shape=shapeObject, T=T, k=k, g=U_exact) t, x = np.ogrid[0:T+k:k,0:1+1/N-0.0001:1/N] Uex = U_exact(x,t) diffList.append(np.nanmax(abs(UList-Uex))) ax1.plot(NList,diffList,label=f"$k={k}$") ax1.set_xlabel("$N$") ax1.set_ylabel("Error") ax1.legend() for k in kListLarge: diffList = []
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)