treewide: move all modules into new submodule to fix relative imports

python, why are you like this?
2024-03-05 13:06:44 +01:00
parent 2ef8664fd9
commit 3769a27a4f
15 changed files with 2 additions and 2 deletions
@@ -0,0 +1,32 @@
+from astroquery.mast import Observations
+import numpy as np
+from enum import Enum
+
+class ObservationSource(Enum):
+    KEPLER = 1
+    K2 = 2
+    TESS = 3
+
+def getStarObservations(starName: str, sources: list[ObservationSource], sequences: list[list[int]] = []) -> str:
+    obs = Observations.query_object(objectname=starName, radius="0 deg")
+
+    obsWantedFilter = obs["dataproduct_type"] == "timeseries"
+
+    if(len(sequences) > 1):
+        obsWantedFilter &= ((obs['sequence_number'] >= max(sequences)) &
+                  (obs['sequence_number'] <= min(sequences)))
+    elif(len(sequences) == 1):
+        obsWantedFilter &= (obs['sequence_number'] == sequences[0])
+
+    obsSourceFilter = np.ma.MaskedArray(data=np.full(obsWantedFilter.shape, False),
+                                        mask=False, fill_value=True)
+    if(ObservationSource.KEPLER in sources):
+        obsSourceFilter |= (obs['obs_collection'] == "Kepler")
+    if(ObservationSource.K2 in sources):
+        obsSourceFilter |= (obs['obs_collection'] == "K2")
+    if(ObservationSource.TESS in sources):
+        obsSourceFilter |= (obs['obs_collection'] == "TESS")
+
+    obsWantedFilter &= obsSourceFilter
+
+    return obs[obsWantedFilter]