Move functions out of uvdata for better reusability

Author: bhazelton

src/pyuvdata/utils/frequency.py

@@ -8,6 +8,7 @@
 
 from . import tools
 from .pol import jstr2num, polstr2num
+from .types import FloatArray, IntArray
 
 
 def _check_flex_spw_contiguous(*, spw_array, flex_spw_id_array, strict=True):
@@ -549,3 +550,41 @@ def _select_freq_helper(
         freq_inds = freq_inds.tolist()
 
     return freq_inds, spw_inds, selections
+
+
+def _add_freq_order(spw_id: IntArray, freq_arr: FloatArray) -> IntArray:
+    """
+    Get the sorting order for the frequency axis after an add.
+
+    Sort first by spw then by channel, but don't reorder channels if they are
+    changing monotonically (all ascending or descending) within the spw.
+
+    Parameters
+    ----------
+    spw_id : np.ndarray of int
+        SPW id array of combined data to be sorted.
+    freq_arr : np.ndarray of float
+        Frequency array of combined data to be sorted.
+
+    Returns
+    -------
+    f_order : np.ndarray of int
+        index array giving the sort order.
+
+    """
+    spws = np.unique(spw_id)
+    f_order = np.concatenate([np.where(spw_id == spw)[0] for spw in np.unique(spw_id)])
+
+    # With spectral windows sorted, check and see if channels within
+    # windows need sorting. If they are ordered in ascending or descending
+    # fashion, leave them be. If not, sort in ascending order
+    for spw in spws:
+        select_mask = spw_id[f_order] == spw
+        check_freqs = freq_arr[f_order[select_mask]]
+        if not np.all(np.diff(check_freqs) > 0) and not np.all(
+            np.diff(check_freqs) < 0
+        ):
+            subsort_order = f_order[select_mask]
+            f_order[select_mask] = subsort_order[np.argsort(check_freqs)]
+
+    return f_order

src/pyuvdata/utils/tools.py

@@ -9,6 +9,8 @@
 
 import numpy as np
 
+from .types import FloatArray, IntArray, StrArray
+
 
 def _get_iterable(x):
     """Return iterable version of input."""
@@ -687,3 +689,43 @@ def _ntimes_to_nblts(uvd):
         inds.append(np.where(unique_t == i)[0][0])
 
     return np.asarray(inds)
+
+
+def flt_ind_str_arr(
+    *,
+    fltarr: FloatArray,
+    intarr: IntArray,
+    flt_tols: tuple[float, float],
+    flt_first: bool = True,
+) -> StrArray:
+    """
+    Create a string array built from float and integer arrays for matching.
+
+    Parameters
+    ----------
+    fltarr : np.ndarray of float
+        float array to be used in output string array
+    intarr : np.ndarray of int
+        integer array to be used in output string array
+    flt_tols : 2-tuple of float
+        Tolerances (relative, absolute) to use in formatting the floats as strings.
+    flt_first : bool
+        Whether to put the float first in the out put string or not (if False
+        the int comes first.)
+
+    Returns
+    -------
+    np.ndarray of str
+        String array that combines the float and integer values, useful for matching.
+
+    """
+    prec_flt = -2 * np.floor(np.log10(flt_tols[-1])).astype(int)
+    prec_int = 8
+    flt_str_list = ["{1:.{0}f}".format(prec_flt, flt) for flt in fltarr]
+    int_str_list = [str(intv).zfill(prec_int) for intv in intarr]
+    list_of_lists = []
+    if flt_first:
+        list_of_lists = [flt_str_list, int_str_list]
+    else:
+        list_of_lists = [int_str_list, flt_str_list]
+    return np.array(["_".join(zpval) for zpval in zip(*list_of_lists, strict=True)])

src/pyuvdata/uvbase.py

@@ -16,6 +16,7 @@
 
 from . import __version__, parameter as uvp
 from .utils.tools import _get_iterable, slicify
+from .utils.types import IntArray
 
 __all__ = ["UVBase"]
 
@@ -865,3 +866,135 @@ def _select_along_param_axis(self, param_dict: dict):
                     # here in the case of a repeated param_name in the form.
                     attr.value = attr.get_from_form(slice_dict)
                     attr.setter(self)
+
+    def _axis_add_helper(
+        self,
+        other,
+        axis_name: str,
+        other_inds: IntArray,
+        final_order: IntArray | None = None,
+    ):
+        """
+        Combine UVParameter objects with a single axis along an axis.
+
+        Parameters
+        ----------
+        other : UVBase
+            The UVBase object to be added.
+        axis_name : str
+            The axis name (e.g. "Nblts", "Npols").
+        other_inds : np.ndarray of int
+            Indices into the other object along this axis to include.
+        final_order : np.ndarray of int
+            Final ordering array giving the sort order after concatenation.
+
+        """
+        update_params = self._get_param_axis(axis_name, single_named_axis=True)
+        other_form_dict = {axis_name: other_inds}
+        for param, axis_list in update_params.items():
+            axis = axis_list[0]
+            new_array = np.concatenate(
+                [
+                    getattr(self, param),
+                    getattr(other, "_" + param).get_from_form(other_form_dict),
+                ],
+                axis=axis,
+            )
+            if final_order is not None:
+                new_array = np.take(new_array, final_order, axis=axis)
+
+            setattr(self, param, new_array)
+
+    def _axis_pad_helper(self, axis_name: str, add_len: int):
+        """
+        Pad out UVParameter objects with multiple dimensions along an axis.
+
+        Parameters
+        ----------
+        axis_name : str
+            The axis name (e.g. "Nblts", "Npols").
+        add_len : int
+            The extra length to be padded on for this axis.
+
+        """
+        update_params = self._get_param_axis(axis_name)
+        multi_axis_params = self._get_multi_axis_params()
+        for param, axis_list in update_params.items():
+            if param not in multi_axis_params:
+                continue
+            this_param_shape = getattr(self, param).shape
+            this_param_type = getattr(self, "_" + param).expected_type
+            bool_type = this_param_type is bool or bool in this_param_type
+            pad_shape = list(this_param_shape)
+            for ax in axis_list:
+                pad_shape[ax] = add_len
+                if bool_type:
+                    pad_array = np.ones(tuple(pad_shape), dtype=bool)
+                else:
+                    pad_array = np.zeros(tuple(pad_shape))
+                new_array = np.concatenate([getattr(self, param), pad_array], axis=ax)
+                if bool_type:
+                    new_array = new_array.astype(np.bool_)
+                setattr(self, param, new_array)
+
+    def _fill_multi_helper(self, other, t2o_dict: dict, order_dict: dict):
+        """
+        Fill UVParameter objects with multiple dimensions from the right side object.
+
+        Parameters
+        ----------
+        other : UVBase
+            The UVBase object to be added.
+        t2o_dict : dict
+            dict giving the indices in the left object to be filled from the right
+            object for each axis (keys are axes, values are index arrays).
+        order_dict : dict
+            dict giving the final sort indices for each axis (keys are axes, values
+            are index arrays for sorting).
+
+        """
+        multi_axis_params = self._get_multi_axis_params()
+        for param in multi_axis_params:
+            form = getattr(self, "_" + param).form
+            index_list = []
+            for axis in form:
+                index_list.append(t2o_dict[axis])
+            new_arr = getattr(self, param)
+            new_arr[np.ix_(*index_list)] = getattr(other, param)
+            setattr(self, param, new_arr)
+
+            # Fix ordering
+            for axis_ind, axis in enumerate(form):
+                if order_dict[axis] is not None:
+                    unique_order_diffs = np.unique(np.diff(order_dict[axis]))
+                    if np.array_equal(unique_order_diffs, np.array([1])):
+                        # everything is already in order
+                        continue
+                    setattr(
+                        self,
+                        param,
+                        np.take(getattr(self, param), order_dict[axis], axis=axis_ind),
+                    )
+
+    def _axis_fast_concat_helper(self, other, axis_name: str):
+        """
+        Concatenate UVParameter objects along an axis assuming no overlap.
+
+        Parameters
+        ----------
+        other : UVBase
+            The UVBase object to be added.
+        axis_name : str
+            The axis name (e.g. "Nblts", "Npols").
+        """
+        update_params = self._get_param_axis(axis_name)
+        for param, axis_list in update_params.items():
+            axis = axis_list[0]
+            setattr(
+                self,
+                param,
+                np.concatenate(
+                    [getattr(self, param)] + [getattr(obj, param) for obj in other],
+                    axis=axis,
+                ),
+            )