Merge branch 'add-brigtness-calc-to-comp-all-orbits' of https://githu…

…b.com/sblunt/orbitize into add-brigtness-calc-to-comp-all-orbits

orbitize/read_input.py

@@ -103,8 +103,9 @@ def read_file(filename):
             "quant12_corr",
             "quant_type",
             "instrument",
+            "brightness",
         ),
-        dtype=(float, int, float, float, float, float, float, "S5", "S5"),
+        dtype=(float, int, float, float, float, float, float, "S5", "S5", float),
     )
 
     # read file
@@ -179,6 +180,10 @@ def read_file(filename):
                 have_seppacorr = np.zeros(
                     num_measurements, dtype=bool
                 )  # zeros are False
+            if "brightness" in input_table.columns:
+                have_brightness = ~input_table["brightness"].mask
+            else:
+                have_brightness = np.zeros(num_measurements, dtype=bool)
             if "rv" in input_table.columns:
                 have_rv = ~input_table["rv"].mask
             else:
@@ -231,18 +236,27 @@ def read_file(filename):
             else:
                 have_rv = np.zeros(num_measurements, dtype=bool)  # zeros are False
 
-            # Rob: not sure if we need this but adding just in case
             if "instrument" in input_table.columns:
                 have_inst = np.ones(num_measurements, dtype=bool)
             else:
                 have_inst = np.zeros(num_measurements, dtype=bool)
+            if "brightness" in input_table.columns:
+                have_brightness = np.ones(num_measurements, dtype=bool)
+            else:
+                have_brightness = np.zeros(num_measurements, dtype=bool)
 
     # orbitize! backwards compatability since we added new columns, some old data formats may not have them
     # fill in with default values
     if orbitize_style:
         if "quant12_corr" not in input_table.keys():
             default_corrs = np.nan * np.ones(len(input_table))
             input_table.add_column(default_corrs, name="quant12_corr")
+        if "brightness" not in input_table.keys():
+            default_brightness = np.nan * np.ones(len(input_table))
+            input_table.add_column(default_brightness, name="brightness")
+            have_brightness = np.zeros(num_measurements, dtype=bool)
+        else:
+            have_brightness = np.ones(num_measurements, dtype=bool)
         if "instrument" not in input_table.keys():
             default_insts = []
             for this_quant_type in input_table["quant_type"]:
@@ -275,6 +289,10 @@ def read_file(filename):
             MJD = row["epoch"] - 2400000.5
         else:
             MJD = row["epoch"]
+        if have_brightness[index]:
+            brightness = row["brightness"]
+        else:
+            brightness = None
 
         # check that "object" is an integer (instead of ABC/bcd)
         if not isinstance(row["object"], (int, np.int32, np.int64)):
@@ -294,6 +312,7 @@ def read_file(filename):
                         None,
                         row["quant_type"],
                         row["instrument"],
+                        None,
                     ]
                 )
 
@@ -322,6 +341,7 @@ def read_file(filename):
                         quant12_corr,
                         row["quant_type"],
                         row["instrument"],
+                        brightness,
                     ]
                 )
             else:  # catch wrong formats
@@ -332,6 +352,7 @@ def read_file(filename):
                 )
 
         else:  # When not in orbitize style
+
             if have_ra[index] and have_dec[index]:
                 # check if there's a covariance term
                 if have_radeccorr[index]:
@@ -362,6 +383,7 @@ def read_file(filename):
                         quant12_corr,
                         "radec",
                         this_inst,
+                        brightness,
                     ]
                 )
 
@@ -395,6 +417,7 @@ def read_file(filename):
                         quant12_corr,
                         "seppa",
                         this_inst,
+                        brightness,
                     ]
                 )
 
@@ -411,6 +434,7 @@ def read_file(filename):
                             None,
                             "rv",
                             row["instrument"],
+                            brightness,
                         ]
                     )
                 else:
@@ -426,6 +450,7 @@ def read_file(filename):
                             None,
                             "rv",
                             "defrv",
+                            brightness,
                         ]
                     )
 

tests/test_brightness.py

@@ -11,34 +11,45 @@
 def test_brightness_calculation():
 
     num_secondary_bodies = 1
-    input_file = os.path.join(DATADIR, "GJ504.csv")
+
+    # input_file = os.path.join(DATADIR, "GJ504.csv")
+    input_file = os.path.join(DATADIR, "betaPic.csv")
     data_table = read_input.read_file(input_file)
+
+    times = data_table["epoch"].value
+
     system_mass = 1.47
     plx = 24.30
 
     test_system = system.System(num_secondary_bodies, data_table, system_mass, plx)
 
+    print(test_system.param_idx)
+
     params = np.array(
         [
-            7.2774010e01,
-            4.1116819e-02,
-            5.6322372e-01,
-            3.5251172e00,
-            4.2904768e00,
-            9.4234377e-02,
-            4.5418411e01,
-            1.4317369e-03,
-            5.6322372e-01,
-            3.1016846e00,
-            4.2904768e00,
-            3.4033456e-01,
-            2.4589758e01,
-            1.4849439e00,
+            10.0,
+            0.1,
+            np.radians(89),
+            np.radians(21),
+            np.radians(31),
+            0.0,  # note: I didn't convert tau here, just picked random number
+            51.5,
+            1.75,
         ]
     )
 
     ra, dec, vz, brightness = test_system.compute_all_orbits(params)
 
+    # TODO (farrah): make plot of brightness vs time
+
+
+def test_read_input_with_brightness():
+
+    # TODO (farrah): use code above as inspiration to read in a csv file with a brightness column
+
+    print("hello! :D ")
+
 
 if __name__ == "__main__":
     test_brightness_calculation()
+    test_read_input_with_brightness()