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)