def testContactRates2(self): undirected = True numVertices = 10 graph = HIVGraph(numVertices, undirected) maleVertex = graph.getVertex(0) maleVertex[HIVVertices.genderIndex] = HIVVertices.male femaleVertex = maleVertex.copy() femaleVertex[HIVVertices.genderIndex] = HIVVertices.female for i in range(5): graph.setVertex(i, maleVertex) graph.setVertex(i+5, femaleVertex) V = graph.getVertexList().getVertices() contactList = range(numVertices) #Test that the parameters alpha and C do the right thing hiddenDegSeq = self.gen.rvs(size=graph.getNumVertices()) rates = HIVRates(graph, hiddenDegSeq) t = 0.2 logging.debug("Rates with no existing contacts") contactRateInds, contactRates = rates.contactRates(range(numVertices), contactList, t) #When there are no contacts the choice is easy and some random new contacts #are chosen. #Now test differences in choice between existing and new contact. t = 0.3 for i in range(5): rates.contactEvent(i, i+5, t) rates.alpha = 1.0 logging.debug("Rates with default alpha=" + str(rates.alpha)) contactRateInds, contactRates = rates.contactRates(range(numVertices), contactList, 0.4) for i in range(5): self.assertTrue(contactRates[i] == rates.contactRate) self.assertTrue(contactRateInds[i] == i+5) #Now try changing alpha logging.debug("Rates with alpha=0.5") rates.setAlpha(0.5) contactRateInds, contactRates = rates.contactRates(range(numVertices), contactList, 0.4) #Observed probabilities change as expected #Now increase time and observe probabilities logging.debug("Rates with t=20") contactRateInds, contactRates = rates.contactRates(range(numVertices), contactList, 20) #Test we don't pick from removed graph.getVertexList().setInfected(0, t) graph.getVertexList().setInfected(4, t) graph.getVertexList().setInfected(7, t) graph.getVertexList().setInfected(8, t) #graph.getVertexList().setDetected(4, t, HIVVertices.randomDetect) #graph.getVertexList().setDetected(7, t, HIVVertices.randomDetect) rates.removeEvent(4, HIVVertices.randomDetect, t) rates.removeEvent(7, HIVVertices.randomDetect, t) infectedSet = graph.getInfectedSet() susceptibleSet = graph.getSusceptibleSet() removedSet = graph.getRemovedSet() contactSet = infectedSet.union(susceptibleSet) infectedList = list(infectedSet) removedList = list(removedSet) contactList = list(contactSet) contactRateInds, contactRates = rates.contactRates(infectedList, contactList, 20) #Contacts cannot be in removed set self.assertTrue(numpy.intersect1d(contactRateInds, numpy.array(removedList)).shape[0]==0)