Add function for applying model

1 files changed, 10 insertions, 4 deletions

diff --git a/hdist.py b/hdist.py
index ca9dba6..d5d7aec 100644
--- a/hdist.py
+++ b/hdist.py
@@ -1,10 +1,16 @@
 from theano import tensor
+import theano
 import numpy
 
+def const(v):
+    if theano.config.floatX == 'float32':
+        return numpy.float32(v)
+    else:
+        return numpy.float64(v)
 
 def hdist(a, b):
-    rearth = numpy.float32(6371)
-    deg2rad = numpy.float32(3.14159265358979 / 180)
+    rearth = const(6371)
+    deg2rad = const(3.141592653589793 / 180)
 
     lat1 = a[:, 1] * deg2rad
     lon1 = a[:, 0] * deg2rad
@@ -15,9 +21,9 @@ def hdist(a, b):
     dlon = abs(lon1-lon2)
 
     al = tensor.sin(dlat/2)**2  + tensor.cos(lat1) * tensor.cos(lat2) * (tensor.sin(dlon/2)**2)
-    d = tensor.arctan2(tensor.sqrt(al), tensor.sqrt(numpy.float32(1)-al))
+    d = tensor.arctan2(tensor.sqrt(al), tensor.sqrt(const(1)-al))
 
-    hd = 2 * rearth * d
+    hd = const(2) * rearth * d
 
     return tensor.switch(tensor.eq(hd, float('nan')), (a-b).norm(2, axis=1), hd)