From b4b255eab8d7ee4bdeb7f83a7ad67d36c9ee7735 Mon Sep 17 00:00:00 2001
From: Chen Siyuan <siyuan.chen@cnrs-orleans.fr>
Date: Sat, 23 Dec 2023 21:06:06 +0800
Subject: [PATCH 1/3] add new EPTA and general selection functions

---
 enterprise/signals/selections.py | 86 ++++++++++++++++++++++++++++++++
 1 file changed, 86 insertions(+)

diff --git a/enterprise/signals/selections.py b/enterprise/signals/selections.py
index c701e7b0..9f2fba35 100644
--- a/enterprise/signals/selections.py
+++ b/enterprise/signals/selections.py
@@ -103,6 +103,24 @@ def by_band(flags):
     return {val: flags["B"] == val for val in flagvals}
 
 
+def by_freq_band(bands = None):
+    def backends(freqs):
+        """Selection function to split by EPTA and custom frequency band values
+        bands: dict of bands and ranges
+        default recovers EPTA freq bands
+        """
+        nonlocal bands
+        if isinstance(bands, dict):
+            pass
+        else:
+            bands = {"Band.1":[0,1000], "Band.2":[1000,2000],
+                     "Band.3":[2000,3000], "Band.4":[3000,10000]}
+        return {val: (freqs>=fl) & (freqs<fh) for val,
+                (fl, fh) in bands.items() if any((freqs>=fl) & (freqs<fh))}
+
+    return backends
+
+
 def by_frontend(flags):
     """Selection function to split by frontend under -fe flag"""
     flagvals = np.unique(flags["fe"])
@@ -129,6 +147,74 @@ def by_telescope(telescope):
     return {t: (telescope == t) for t in telescopes}
 
 
+def by_index(name, idx):
+    def indexvals(toas):
+        """Selection function to split by ToA index values"""
+        return {name: np.isin(np.arange(len(toas)), idx)}
+    return indexvals
+
+
 def no_selection(toas):
     """Default selection with no splitting."""
     return {"": np.ones_like(toas, dtype=bool)}
+
+
+def custom_backends(cb):
+    def backends(backend_flags):
+        """Selection function to split by custom backend flags only.
+        cb : list of str of the backends
+        use None to recover by_backend
+        use ["ASP", "GASP", "GUPPI", "PUPPI", "YUPPI"] to recover nanograv_backends
+        """
+        nonlocal cb
+        flagvals = np.unique(backend_flags)
+        if cb is not None:
+            cb = list(np.atleast_1d(cb))
+            flagvals = [val for val in flagvals if any([b in val for b in cb])]
+        else:
+            pass
+        return {val: backend_flags == val for val in flagvals}
+
+    return backends
+
+
+def custom_backends_dict(cb):
+    def backends(backend_flags, flags, toas):
+        """Selection function to split by custom flags dictionary only.
+        cb : str, list or dict of flags and names
+        use None to recover no_selection
+        use {"B":None} to recover by_band
+        use {"fe":None} to recover by_frontend
+        use {"backend":None} to recover by_backend
+        use {"backend":["ASP", "GASP", "GUPPI", "PUPPI", "YUPPI"]} to recover nanograv_backends
+        """
+        nonlocal cb
+        if isinstance(cb, str) or isinstance(cb, list):
+            flagvals = np.unique(backend_flags)
+            cb = list(np.atleast_1d(cb))
+            flagvals = [val for val in flagvals if any([b in val for b in cb])]
+            return {val: backend_flags == val for val in flagvals}
+        elif isinstance(cb, dict):
+            flagdict = {}
+            for flagname in cb.keys():
+                if flagname == "backend":
+                    flagvals = np.unique(backend_flags)
+                    if cb["backend"] is not None:
+                        cb_key = list(np.atleast_1d(cb["backend"]))
+                        flagvals = [val for val in flagvals if any([b in val for b in cb_key])]
+                    else:
+                        pass
+                    flagdict.update({val: backend_flags == val for val in flagvals})
+                else:
+                    flagvals = np.unique(flags[flagname])
+                    if cb[flagname] is not None:
+                        cb_key = list(np.atleast_1d(cb[flagname]))
+                        flagvals = [val for val in flagvals if any([b in val for b in cb_key])]
+                    else:
+                        pass
+                    flagdict.update({val: flags[flagname] == val for val in flagvals})
+            return flagdict
+        else:
+            return {"": np.ones_like(toas, dtype=bool)}
+
+    return backends

From c47b04af2e58aa5b2c78f5feb1dd377823a0cef4 Mon Sep 17 00:00:00 2001
From: Chen Siyuan <siyuan.chen@cnrs-orleans.fr>
Date: Sun, 7 Apr 2024 21:43:44 +0800
Subject: [PATCH 2/3] liniting

---
 enterprise/signals/gp_bases.py | 1 -
 enterprise/signals/utils.py    | 2 --
 2 files changed, 3 deletions(-)

diff --git a/enterprise/signals/gp_bases.py b/enterprise/signals/gp_bases.py
index 1847e54e..9b621f4e 100644
--- a/enterprise/signals/gp_bases.py
+++ b/enterprise/signals/gp_bases.py
@@ -219,7 +219,6 @@ def createfourierdesignmatrix_eph(
 
 @function
 def createfourierdesignmatrix_chromatic(toas, freqs, nmodes=30, Tspan=None, logf=False, fmin=None, fmax=None, idx=4):
-
     """
     Construct Scattering-variation fourier design matrix.
 
diff --git a/enterprise/signals/utils.py b/enterprise/signals/utils.py
index eff440e1..e2986f55 100644
--- a/enterprise/signals/utils.py
+++ b/enterprise/signals/utils.py
@@ -322,7 +322,6 @@ def create_stabletimingdesignmatrix(designmat, fastDesign=True):
 
 
 def make_ecc_interpolant():
-
     """
     Make interpolation function from eccentricity file to
     determine number of harmonics to use for a given
@@ -339,7 +338,6 @@ def make_ecc_interpolant():
 
 
 def get_edot(F, mc, e):
-
     """
     Compute eccentricity derivative from Taylor et al. (2016)
 

From d105ddd509405ba8800bf4065080a43010c9aa6d Mon Sep 17 00:00:00 2001
From: Chen Siyuan <siyuan.chen@cnrs-orleans.fr>
Date: Sun, 7 Apr 2024 21:50:21 +0800
Subject: [PATCH 3/3] more linting

---
 enterprise/signals/selections.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/enterprise/signals/selections.py b/enterprise/signals/selections.py
index 9f2fba35..b0a570b8 100644
--- a/enterprise/signals/selections.py
+++ b/enterprise/signals/selections.py
@@ -103,7 +103,7 @@ def by_band(flags):
     return {val: flags["B"] == val for val in flagvals}
 
 
-def by_freq_band(bands = None):
+def by_freq_band(bands=None):
     def backends(freqs):
         """Selection function to split by EPTA and custom frequency band values
         bands: dict of bands and ranges
@@ -113,10 +113,10 @@ def backends(freqs):
         if isinstance(bands, dict):
             pass
         else:
-            bands = {"Band.1":[0,1000], "Band.2":[1000,2000],
-                     "Band.3":[2000,3000], "Band.4":[3000,10000]}
-        return {val: (freqs>=fl) & (freqs<fh) for val,
-                (fl, fh) in bands.items() if any((freqs>=fl) & (freqs<fh))}
+            bands = {"Band.1": [0, 1000], "Band.2": [1000, 2000], "Band.3": [2000, 3000], "Band.4": [3000, 10000]}
+        return {
+            val: (freqs >= fl) & (freqs < fh) for val, (fl, fh) in bands.items() if any((freqs >= fl) & (freqs < fh))
+        }
 
     return backends
 
@@ -151,6 +151,7 @@ def by_index(name, idx):
     def indexvals(toas):
         """Selection function to split by ToA index values"""
         return {name: np.isin(np.arange(len(toas)), idx)}
+
     return indexvals