support cuda in fix for numeric overflow

UDST · May 4, 2018 · 9794d9e · 9794d9e
1 parent d84469d
commit 9794d9e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/urbansim/urbanchoice/mnl.py b/urbansim/urbanchoice/mnl.py
@@ -36,8 +36,7 @@ def mnl_probs(data, beta, numalts):
     utilities.reshape(numalts, utilities.size() // numalts)
 
     # https://stats.stackexchange.com/questions/304758/softmax-overflow
-    if clamp:
-        utilities.mat -= utilities.mat.max(0)
+    utilities = utilities.subtract(utilities.max(0))
 
     exponentiated_utility = utilities.exp(inplace=True)
     if clamp:

diff --git a/urbansim/urbanchoice/pmat.py b/urbansim/urbanchoice/pmat.py
@@ -79,6 +79,12 @@ def cumsum(self, axis):
         # elif self.typ == 'cuda':
         #  return PMAT(misc.cumsum(self.mat,axis=axis))
 
+    def max(self, axis):
+        if self.typ == 'numpy':
+            return PMAT(np.max(self.mat, axis=axis))
+        elif self.typ == 'cuda':
+            return PMAT(self.mat.max(axis=axis))
+
     def argmax(self, axis):
         if self.typ == 'numpy':
             return PMAT(np.argmax(self.mat, axis=axis))