reduce indentation levels, #247

casangi · Sep 13, 2024 · 4f841e7 · 4f841e7
1 parent 487b815
commit 4f841e7
Showing 1 changed file with 63 additions and 65 deletions.
diff --git a/src/xradio/vis/_vis_utils/_ms/create_antenna_xds.py b/src/xradio/vis/_vis_utils/_ms/create_antenna_xds.py
@@ -11,6 +11,7 @@
     load_generic_table,
     convert_casacore_time_to_mjd,
     make_taql_where_between_min_max,
+    table_exists,
 )
 from xradio._utils.schema import convert_generic_xds_to_xradio_schema
 from xradio.vis._vis_utils._ms.msv4_sub_xdss import interpolate_to_time
@@ -305,78 +306,75 @@ def create_gain_curve_xds(
     """
 
     gain_curve_xds = None
-    if os.path.exists(
-        os.path.join(in_file, "GAIN_CURVE")
-    ):  # Check if the table exists.
-        generic_gain_curve_xds = load_generic_table(
-            in_file,
-            "GAIN_CURVE",
-            taql_where=f" where (ANTENNA_ID IN [{','.join(map(str,ant_xds.antenna_id.values))}]) AND (SPECTRAL_WINDOW_ID = {spectral_window_id})",
-        )
+    if not table_exists(os.path.join(in_file, "GAIN_CURVE")):
+        return gain_curve_xds
 
-        if (
-            generic_gain_curve_xds.data_vars
-        ):  # Some times the gain_curve table is empty (this is the case with ngEHT simulation data we have).
-
-            assert (
-                len(generic_gain_curve_xds.SPECTRAL_WINDOW_ID) == 1
-            ), "Only one spectral window is supported."
-            generic_gain_curve_xds = generic_gain_curve_xds.isel(
-                SPECTRAL_WINDOW_ID=0, drop=True
-            )  # Drop the spectral window dimension as it is singleton.
-
-            assert (
-                len(generic_gain_curve_xds.TIME) == 1
-            ), "Only one gain curve measurement per antenna is supported."
-            generic_gain_curve_xds = generic_gain_curve_xds.isel(TIME=0, drop=True)
-
-            generic_gain_curve_xds = generic_gain_curve_xds.sel(
-                ANTENNA_ID=ant_xds.antenna_id, drop=False
-            )  # Make sure the antenna_id is in the same order as the xds .
-
-            gain_curve_xds = xr.Dataset(attrs={"type": "gain_curve"})
-
-            to_new_data_variables = {
-                "INTERVAL": ["GAIN_CURVE_INTERVAL", ["antenna_name"]],
-                "GAIN": [
-                    "GAIN_CURVE",
-                    ["antenna_name", "poly_term", "receptor_label"],
-                ],
-                "SENSITIVITY": [
-                    "GAIN_CURVE_SENSITIVITY",
-                    ["antenna_name", "receptor_label"],
-                ],
-            }
+    generic_gain_curve_xds = load_generic_table(
+        in_file,
+        "GAIN_CURVE",
+        taql_where=f" where (ANTENNA_ID IN [{','.join(map(str,ant_xds.antenna_id.values))}]) AND (SPECTRAL_WINDOW_ID = {spectral_window_id})",
+    )
 
-            to_new_coords = {
-                "TYPE": ["gain_curve_type", ["antenna_name"]],
-            }
+    if not generic_gain_curve_xds.data_vars:
+        # Some times the gain_curve table is empty (this is the case with ngEHT simulation data we have).
+        return gain_curve_xds
 
-            gain_curve_xds = convert_generic_xds_to_xradio_schema(
-                generic_gain_curve_xds,
-                gain_curve_xds,
-                to_new_data_variables,
-                to_new_coords,
-            )
+    assert (
+        len(generic_gain_curve_xds.SPECTRAL_WINDOW_ID) == 1
+    ), "Only one spectral window is supported."
+    generic_gain_curve_xds = generic_gain_curve_xds.isel(
+        SPECTRAL_WINDOW_ID=0, drop=True
+    )  # Drop the spectral window dimension as it is singleton.
 
-            ant_borrowed_coords = {
-                "antenna_name": ant_xds.coords["antenna_name"],
-                "station": ant_xds.coords["station"],
-                "mount": ant_xds.coords["mount"],
-                "telescope_name": ant_xds.coords["telescope_name"],
-                "receptor_label": ant_xds.coords["receptor_label"],
-                "polarization_type": ant_xds.coords["polarization_type"],
-            }
-            gain_curve_xds = gain_curve_xds.assign_coords(ant_borrowed_coords)
+    assert (
+        len(generic_gain_curve_xds.TIME) == 1
+    ), "Only one gain curve measurement per antenna is supported."
+    generic_gain_curve_xds = generic_gain_curve_xds.isel(TIME=0, drop=True)
 
-            gain_curve_xds["GAIN_CURVE"] = gain_curve_xds["GAIN_CURVE"].transpose(
-                "antenna_name", "receptor_label", "poly_term"
-            )
+    generic_gain_curve_xds = generic_gain_curve_xds.sel(
+        ANTENNA_ID=ant_xds.antenna_id, drop=False
+    )  # Make sure the antenna_id is in the same order as the xds .
 
-        return gain_curve_xds
+    gain_curve_xds = xr.Dataset(attrs={"type": "gain_curve"})
 
-    else:
-        return gain_curve_xds
+    to_new_data_variables = {
+        "INTERVAL": ["GAIN_CURVE_INTERVAL", ["antenna_name"]],
+        "GAIN": [
+            "GAIN_CURVE",
+            ["antenna_name", "poly_term", "receptor_label"],
+        ],
+        "SENSITIVITY": [
+            "GAIN_CURVE_SENSITIVITY",
+            ["antenna_name", "receptor_label"],
+        ],
+    }
+
+    to_new_coords = {
+        "TYPE": ["gain_curve_type", ["antenna_name"]],
+    }
+
+    gain_curve_xds = convert_generic_xds_to_xradio_schema(
+        generic_gain_curve_xds,
+        gain_curve_xds,
+        to_new_data_variables,
+        to_new_coords,
+    )
+
+    ant_borrowed_coords = {
+        "antenna_name": ant_xds.coords["antenna_name"],
+        "station": ant_xds.coords["station"],
+        "mount": ant_xds.coords["mount"],
+        "telescope_name": ant_xds.coords["telescope_name"],
+        "receptor_label": ant_xds.coords["receptor_label"],
+        "polarization_type": ant_xds.coords["polarization_type"],
+    }
+    gain_curve_xds = gain_curve_xds.assign_coords(ant_borrowed_coords)
+
+    gain_curve_xds["GAIN_CURVE"] = gain_curve_xds["GAIN_CURVE"].transpose(
+        "antenna_name", "receptor_label", "poly_term"
+    )
+
+    return gain_curve_xds
 
 
 def extract_phase_cal_info(