Replace geodesic distance with haversine in WRF

To speed up covariance matrix computation

da/statevectors/statevector_wrfchem.py

@@ -336,11 +336,13 @@ class WRFChemStateVector(StateVector):
         dasystem rc file.
         """
 
+        logging.info("Creating covariance matrix")
         covariancematrix = np.zeros([self.nparams, self.nparams])
         
         # Fill flux covariances
         if self.spatial_correlations:
-            distances = self.get_distance_matrix()
+            if not hasattr(self, "dm"):
+                self.dm = self.get_distance_matrix()
         for nproc in range(self.n_emis_proc):
             i0 = self.nparams_proc*nproc
             i1 = self.nparams_proc*(nproc+1)
@@ -350,7 +352,7 @@ class WRFChemStateVector(StateVector):
                 sigma_nproc = float(dacycle.dasystem['sigma_scale_' + str(nproc+1)])
 
             if self.spatial_correlations:
-                covariancematrix[i0:i1, i0:i1] = sigma_nproc**2 * np.exp(-distances/self.correlation_length_km)
+                covariancematrix[i0:i1, i0:i1] = sigma_nproc**2 * np.exp(-self.dm/self.correlation_length_km)
             else:
                 covariancematrix[i0:i1, i0:i1] = sigma_nproc**2 * np.identity(i1-i0)
             
@@ -382,13 +384,66 @@ class WRFChemStateVector(StateVector):
             lats[nr] = self.lat[is_nr].mean()
             lons[nr] = self.lon[is_nr].mean()
         coords = np.vstack([lats, lons]).T
-        
-        # https://stackoverflow.com/a/49039973
-        from geopy import distance
-        from scipy.spatial.distance import pdist, squareform
         geodist=lambda p1, p2: distance.distance(p1, p2).km
-        return squareform(pdist(coords, geodist))
+        dm = self.cdist(coords, coords)
+
+        return dm
+
+    @staticmethod
+    def cdist(c1, c2):
+        """
+        Haversine distance between each combination of
+        points in input.
+        
+        Parameters
+        ----------
+        c1 : n by 2 array of (lat, lon) points
+        c2 : m by 2 array of (lat, lon) points
+        
+        Returns
+        -------
+        dm : n by m array of distances in km
+        """
 
+        def haversine(p1, p2):
+            """
+            Haversine distance between points in input
+            
+            Parameters
+            ----------
+            p1, p2 : n by 2 array of (lat, lon) points
+            
+            Returns
+            -------
+            d : n by 1 array of distances in km
+            """
+            EARTH_RADIUS = 6371.009
+            lat1 = np.radians(p1[:, 0])
+            lat2 = np.radians(p2[:, 0])
+            dlon = np.radians(p2[:, 1]) - np.radians(p1[:, 1])
+            h = np.sin((lat2-lat1)/2)**2 \
+                + np.cos(lat1)*np.cos(lat2)*np.sin(dlon/2)**2
+            h[h>1.0] = 1.0
+            d = 2*EARTH_RADIUS*np.arcsin(np.sqrt(h))
+            return d
+            
+    
+        # Ensure dimensions
+        c1 = np.array(c1, ndmin=2)
+        c2 = np.array(c2, ndmin=2)
+        
+        # Combine
+        p1 = np.repeat(c1,len(c2), 0)
+        p2 = np.tile(c2, [len(c1),1])
+        
+        # Compute
+        distances = haversine(p1, p2)
+        
+        # Put in matrix form
+        dm = np.reshape(distances, (len(c1), len(c2)))
+        
+        return dm
+    
 ################### End Class StateVector ###################
 
 if __name__ == "__main__":