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)