Merge pull request #34 from liamtoney/pass-rij

Allow users to directly pass an `rij` array to `ltsva()`

Author: jwbishop

lts_array/classes/lts_data_class.py

@@ -17,14 +17,21 @@ def __init__(self, window_length, window_overlap, alpha):
         self.window_overlap = window_overlap
         self.alpha = alpha
 
-    def build_data_arrays(self, st, latlist, lonlist, remove_elements=None):
+    def build_data_arrays(self, st, latlist, lonlist, remove_elements=None, rij=None):
         """ Collect basic data from stream file. Project lat/lon into r_ij coordinates.
 
+        Alternatively, if ``rij`` is provided, ``latlist`` and ``lonlist`` are not used
+        and the provided Cartesion coordinates ``rij`` are used instead.
+
         Args:
             st (stream): An obspy stream object.
             latlist (list): A list of latitude points.
             lonlist (list): A list of longitude points.
             remove_elements (list): A list of elements to remove before processing. Python numbering is used, so "[0]" removes the first element.
+            rij (array or None): A NumPy array with the first row corresponding to
+                cartesian "X" - coordinates and the second row corresponding to
+                cartesian "Y" - coordinates, in units of km. If this is provided then
+                ``latlist`` and ``lonlist`` are ignored.
         """
         # Check that all traces have the same length
         if len(set([len(tr) for tr in st])) != 1:
@@ -35,8 +42,14 @@ def build_data_arrays(self, st, latlist, lonlist, remove_elements=None):
             remove_elements = np.sort(remove_elements)
             for jj in range(0, len(remove_elements)):
                 st.remove(st[remove_elements[jj]])
-                latlist.remove(latlist[remove_elements[jj]])
-                lonlist.remove(lonlist[remove_elements[jj]])
+                if rij is None:
+                    latlist.remove(latlist[remove_elements[jj]])
+                    lonlist.remove(lonlist[remove_elements[jj]])
+                else:
+                    _x, _y = [list(_) for _ in rij]
+                    _x.remove(_x[remove_elements[jj]])
+                    _y.remove(_y[remove_elements[jj]])
+                    rij = np.array([_x, _y])
                 remove_elements -= 1
 
         # Save the station name
@@ -69,7 +82,11 @@ def build_data_arrays(self, st, latlist, lonlist, remove_elements=None):
         for ii, tr in enumerate(st):
             self.data[:, ii] = tr.data
         # Set the array coordinates
-        self.rij = self.getrij(latlist, lonlist)
+        if rij is None:
+            self.rij = self.getrij(latlist, lonlist)
+        else:
+            print('Ignoring `lat_list` and `lon_list` since `rij` array was provided!')
+            self.rij = rij
         # Make sure the least squares problem is well-posed
         # rij must have at least 3 elements
         if np.shape(self.rij)[1] < 3:
@@ -92,10 +109,9 @@ def getrij(self, latlist, lonlist):
 
         Returns:
             (array):
-            ``rij``: A numpy array with the first row corresponding to
+            ``rij``: A NumPy array with the first row corresponding to
             cartesian "X" - coordinates and the second row
-            corresponding to cartesian "Y" - coordinates.
-
+            corresponding to cartesian "Y" - coordinates, in units of km.
         """
 
         # Check that the lat-lon arrays are the same size.

lts_array/ltsva.py

@@ -6,7 +6,7 @@
 warnings.simplefilter(action='ignore', category=FutureWarning)
 
 
-def ltsva(st, lat_list, lon_list, window_length, window_overlap, alpha=1.0, plot_array_coordinates=False, remove_elements=None):
+def ltsva(st, lat_list, lon_list, window_length, window_overlap, alpha=1.0, plot_array_coordinates=False, remove_elements=None, rij=None):
     r""" Process infrasound or seismic array data with least trimmed squares (LTS).
 
     Args:
@@ -19,6 +19,10 @@ def ltsva(st, lat_list, lon_list, window_length, window_overlap, alpha=1.0, plot
             Choose 1.0 for ordinary least squares (default).
         plot_array_coordinates (bool): Plot array coordinates? Defaults to False.
         remove_elements (list): (Optional) Remove element number(s) from ``st``, ``lat_list``, and ``lon_list`` before processing. Here numbering refers to the Python index (e.g. [0] = remove 1st element in stream).
+        rij (array or None): A NumPy array with the first row corresponding to cartesian
+            "X" - coordinates and the second row corresponding to cartesian "Y" -
+            coordinates, in units of km. If this is provided then ``lat_list`` and
+            ``lon_list`` are ignored.
 
     Returns:
         (tuple):
@@ -36,7 +40,7 @@ def ltsva(st, lat_list, lon_list, window_length, window_overlap, alpha=1.0, plot
 
     # Build data object
     data = DataBin(window_length, window_overlap, alpha)
-    data.build_data_arrays(st, lat_list, lon_list, remove_elements)
+    data.build_data_arrays(st, lat_list, lon_list, remove_elements, rij)
 
     # Plot array coordinates as a check
     if plot_array_coordinates: