Update Duration-priors.py to reduce support

IntervalsModel/network-data/POLYMOD/Duration-priors.py

@@ -92,24 +92,22 @@ for x, y, z in list(product(Ages,Ages, Locales)):
     Durlist = [C[C['duration_multi']==i]["duration_multi"].count() for i in range(1,6)]
     DurFreqAAL[(x,y,z)] = Durlist/np.sum(Durlist)
 
-       
-       
 # Define error functions for poisson random variables
-durcutoff = 6*16
-def PoisArray(lam):
+durcutoff = {'home': 6*16, 'workschool': 6*8, 'rest': 6*8}
+def PoisArray(lam, loc):
     arr = [math.exp(-lam + k*math.log(lam) - np.sum([math.log(n) for n in range(1,k+1)]) ) 
-           for k in range(durcutoff)]
-    return [math.exp(-lam)] + arr
+           for k in range(durcutoff[loc])]
+    return arr
 
-def PoisBin(lam):
-    Arr = PoisArray(lam)
+def PoisBin(lam,loc):
+    Arr = PoisArray(lam,loc)
     out = [Arr[0], Arr[1], np.sum(Arr[2:6]), np.sum(Arr[6:24]), np.sum(Arr[24:])]
     return out/np.sum(out)
 
 def PoisErr(lam, label): # label is SymAge-Location
     agesrc, agetar, loc = label
-    err = np.sum((PoisBin(lam) - DurFreqAAL[label])**2)
-    errop = np.sum((PoisBin(lam) - DurFreqAAL[(agetar, agesrc, loc)])**2)
+    err = np.sum((PoisBin(lam, loc) - DurFreqAAL[label])**2)
+    errop = np.sum((PoisBin(lam, loc) - DurFreqAAL[(agetar, agesrc, loc)])**2)
     return err+errop
 PoisErr = np.vectorize(PoisErr, excluded=[1])
 
@@ -117,9 +115,11 @@ PoisPrior={}
 SymAge = [('Y','Y'), ('Y', 'M'), ('Y', 'O'), ('M', 'M'),('M', 'O'), ('O','O')]
 for symage, loc in list(product(SymAge, Locales)):
     AAL = (symage[0],symage[1],loc)
-    arr = [1/PoisErr(i,AAL) for i in range(1,durcutoff)]
-    PoisPrior[AAL] = arr/np.sum(arr)
-    
+    arr = [1/PoisErr(i,AAL) for i in range(1,durcutoff[loc])]
+    zeropad = [0 for i in range(durcutoff['home'] - durcutoff[loc])]
+    PoisPrior[AAL] = (arr+zeropad)/np.sum(arr)
+
+
 # Save them to csv
 df = pd.DataFrame(list(PoisPrior.items()), columns=['col1','col2'])
 dfkeys = pd.DataFrame([pd.Series(x) for x in df.col1])