Reorganising AdlerData, small fixes.

src/adler/dataclasses/AdlerData.py

@@ -7,41 +7,136 @@ class AdlerData:
     """
     Class for storing Adler-calculated values.
 
-    Note that for all per-filter attributes, the type is an array in order u, g, r, i, z, y.
-
     Attributes:
     -----------
+    filter_list : list
+        List of filters under investigation.
+
+    model_lists : list
+        List of lists of models per-filter. Length = len(filter_list). For example, if the filters are ['u', 'g', 'i'],
+        and HG has been calculated for all three and HG1G2 for 'u' only, model_lists = [ ['HG', 'HG1G2'], ['HG'], ['HG'] ].
+        This list of lists is used to track the parameters in the H and phase_parameter attributes, which will take an identical shape.
+
     phaseAngle_min_adler : array_like
-        Minimum phase angle of observations used in fitting model (degrees)
+        Minimum phase angle of observations used in fitting model (degrees). Size = len(filter_list).
 
     phaseAngle_range_adler : array_like
-        Max minus min phase angle range of observations used in fitting model (degrees)
+        Max minus min phase angle range of observations used in fitting model (degrees). Size = len(filter_list).
 
     nobs_adler : array_like
-        Number of observations used in fitting model
+        Number of observations used in fitting model. Size = len(filter_list).
 
     arc_adler : array_like
-        Observational arc used to fit model (days)
+        Observational arc used to fit model (days). Size = len(filter_list).
+
+    H_adler : list
+        Absolute magnitude. List of lists arranged identically to model_lists, with per-filter and per-model values.
+
+    H_err_adler : list
+        Error in absolute magnitude. List of lists arranged identically to model_lists, with per-filter and per-model values.
 
-    H_P16_adler : array_like
-        Absolute magnitude of the Penttila et al. 2016 model (mag)
+    phase_parameter_1 : list
+        First phase parameter of the model in question. List of lists arranged identically to model_lists, with per-filter and per-model values.
 
-    G12_P16_adler : array_like
-        Phase parameter of the Penttila et al. 2016 model
+    phase_parameter_err : list
+        Error in the first phase parameter. List of lists arranged identically to model_lists, with per-filter and per-model values.
 
-    HErr_P16_adler : array_like
-        Uncertainty in absolute magnitude (mag)
+    phase_parameter_2 : list
+        Second phase parameter of the model in question. List of lists arranged identically to model_lists, with per-filter and per-model values.
 
-    G12Err_P16_adler : array_like
-        Uncertainty in phase parameter
+    phase_parameter_err : list
+        Error in the second phase parameter. List of lists arranged identically to model_lists, with per-filter and per-model values.
 
     """
 
-    phaseAngle_min_adler: field(default_factory=np.zeros(6))
-    phaseAngle_range_adler: field(default_factory=np.zeros(6))
-    nobs_adler: field(default_factory=np.zeros(6))
-    arc_adler: field(default_factory=np.zeros(6))
-    H_P16_adler: field(default_factory=np.zeros(6))
-    G12_P16_adler: field(default_factory=np.zeros(6))
-    HErr_P16_adler: field(default_factory=np.zeros(6))
-    G12Err_P16_adler: field(default_factory=np.zeros(6))
+    filter_list: list
+    model_lists: list = field(default_factory=list)
+
+    phaseAngle_min_adler: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=float))
+    phaseAngle_range_adler: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=float))
+    nobs_adler: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=int))
+    arc_adler: np.ndarray = field(default_factory=lambda: np.zeros(0, dtype=float))
+
+    # Note that these are lists and not arrays because appending to Numpy arrays (i.e. to add parameters for a different model) is not recommended.
+    # They can be cast to Numpy arrays later if needed.
+    H_adler: list = field(default_factory=list, init=False)
+    H_err_adler: list = field(default_factory=list, init=False)
+    phase_parameter_1: list = field(default_factory=list, init=False)
+    phase_parameter_2: list = field(default_factory=list, init=False)
+    phase_parameter_1_err: list = field(default_factory=list, init=False)
+    phase_parameter_2_err: list = field(default_factory=list, init=False)
+
+    def __post_init__(self):
+        """This runs post-initialisation and creates all of the class attributes where one dimension is "filters" to ensure the arrays
+        and lists have the correct size. This makes population a little easier.
+        """
+        filter_length = len(self.filter_list)
+
+        self.model_lists = [[] for a in range(0, filter_length)]
+
+        self.phaseAngle_min_adler = np.zeros(filter_length, dtype=float)
+        self.phaseAngle_range_adler = np.zeros(filter_length, dtype=float)
+        self.nobs_adler = np.zeros(filter_length, dtype=int)
+        self.arc_adler = np.zeros(filter_length, dtype=float)
+
+        self.H_adler = [[] for a in range(0, filter_length)]
+        self.H_err_adler = [[] for a in range(0, filter_length)]
+        self.phase_parameter_1 = [[] for a in range(0, filter_length)]
+        self.phase_parameter_2 = [[] for a in range(0, filter_length)]
+        self.phase_parameter_1_err = [[] for a in range(0, filter_length)]
+        self.phase_parameter_2_err = [[] for a in range(0, filter_length)]
+
+    def populate_phase_parameters(
+        self,
+        filter_name,
+        model_name,
+        phaseAngle_min,
+        phaseAngle_range,
+        nobs,
+        arc,
+        H,
+        H_err,
+        parameter_1,
+        parameter_1_err,
+        parameter_2=None,
+        parameter_2_err=None,
+    ):
+        """Convenience method to correctly populate phase curve arrays/lists."""
+
+        # Raise an exception if only one of parameter_2 and parameter_2_err is given.
+        if (parameter_2 is None) != (parameter_2_err is None):
+            raise Exception(
+                "If using a model with 2 phase parameters, both parameter_2 and parameter_2_err must be supplied."
+            )
+
+        # Make sure the supplied filter is in the filter list.
+        try:
+            filter_index = self.filter_list.index(filter_name)
+        except ValueError:
+            raise Exception("Filter {} is not in supplied filter list.".format(filter_name))
+
+        self.phaseAngle_min_adler[filter_index] = phaseAngle_min
+        self.phaseAngle_range_adler[filter_index] = phaseAngle_range
+        self.nobs_adler[filter_index] = nobs
+        self.arc_adler[filter_index] = arc
+
+        # Check and see if the model has already been calculated for this filter.
+        if model_name not in self.model_lists[filter_index]:
+            self.model_lists[filter_index].append(model_name)
+
+            self.H_adler[filter_index].append(H)
+            self.H_err_adler[filter_index].append(H_err)
+            self.phase_parameter_1[filter_index].append(parameter_1)
+            self.phase_parameter_1_err[filter_index].append(parameter_1_err)
+            self.phase_parameter_2[filter_index].append(parameter_2)
+            self.phase_parameter_2_err[filter_index].append(parameter_2_err)
+
+        else:
+            model_index = self.model_lists[filter_index].index(model_name)
+
+            self.H_adler[filter_index][model_index] = H
+            self.H_err_adler[filter_index][model_index] = H_err
+            self.phase_parameter_1[filter_index][model_index] = parameter_1
+            self.phase_parameter_1_err[filter_index][model_index] = parameter_1_err
+            self.phase_parameter_2[filter_index][model_index] = parameter_2
+            self.phase_parameter_2_err[filter_index][model_index] = parameter_2_err

src/adler/dataclasses/DataSchema.py

@@ -1,5 +1,4 @@
 import numpy as np
-import sys
 from dataclasses import dataclass, field
 from lsst.rsp import get_tap_service
 
@@ -28,7 +27,7 @@ def populate(self, population_location, sql_query, sql_filename):
         elif population_location == "SQL":
             data_table = self.get_SQL_table(sql_query, sql_filename)
         else:
-            sys.exit(
+            raise Exception(
                 "Population source not recognised. Please supply either 'RSP' or 'SQL' for population_location argument."
             )
 
@@ -46,7 +45,7 @@ def get_RSP_table(self, sql_query):  # pragma: no cover
 
         self.service = get_tap_service("ssotap")
 
-        return self.service.search(sql_query).to_table()
+        return self.service.search(sql_query)
 
     def get_SQL_table(self, sql_query, sql_filename):
         pass
@@ -121,6 +120,9 @@ class Observations(DataSchema):
     reduced_mag: array_like
         The reduced magnitude.
 
+    num_obs : int
+        The number of observations contained in this structure.
+
     """
 
     ssObjectId: str = ""
@@ -133,6 +135,7 @@ class Observations(DataSchema):
     topocentricDist: np.ndarray = field(default_factory=np.zeros(0))
     heliocentricDist: np.ndarray = field(default_factory=np.zeros(0))
     reduced_mag: np.ndarray = field(default_factory=np.zeros(0))
+    num_obs: int = 0
 
     def __init__(self, ssObjectId, population_location, sql_query, sql_filename=None):
         """Initialises the Observations object.
@@ -154,6 +157,7 @@ def __init__(self, ssObjectId, population_location, sql_query, sql_filename=None
         self.ssObjectId = ssObjectId
         self.population_location = population_location
         self.populate(self.population_location, sql_query, sql_filename)
+        self.num_obs = len(midpointMjdTai)
         self.calculate_reduced_mag()
 
     def populate_from_table(self, data_table):
@@ -334,7 +338,7 @@ class SSObject(DataSchema):
     r_H: float = 0.0
     r_G12: float = 0
     r_Herr: float = 0.0
-    r_G12_err: float = 0.0
+    r_G12err: float = 0.0
     r_nData: int = 0
     maxExtendedness: float = 0.0
     minExtendedness: float = 0.0

tests/adler/dataclasses/test_AdlerData.py

@@ -0,0 +1,49 @@
+from numpy.testing import assert_array_equal
+import pytest
+
+
+def test_populate_phase_parameters():
+    from adler.dataclasses.AdlerData import AdlerData
+
+    test_object = AdlerData(["u", "g", "r"])
+
+    test_object.populate_phase_parameters("u", "model_1", 11.0, 12.0, 13, 14.0, 15.0, 16.0, 17.0, 18.0)
+    test_object.populate_phase_parameters(
+        "u", "model_2", 21.0, 22.0, 23, 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0
+    )
+    test_object.populate_phase_parameters("g", "model_1", 31.0, 32.0, 33, 34.0, 35.0, 36.0, 37.0, 38.0)
+    test_object.populate_phase_parameters(
+        "r", "model_2", 41.0, 42.0, 43, 44.0, 45.0, 46.0, 47.0, 48.0, 49.0, 50.0
+    )
+
+    assert test_object.filter_list == ["u", "g", "r"]
+    assert test_object.model_lists == [["model_1", "model_2"], ["model_1"], ["model_2"]]
+
+    assert_array_equal(test_object.phaseAngle_min_adler, [21.0, 31.0, 41.0])
+    assert_array_equal(test_object.phaseAngle_range_adler, [22.0, 32.0, 42.0])
+    assert_array_equal(test_object.nobs_adler, [23, 33, 43])
+    assert_array_equal(test_object.arc_adler, [24.0, 34.0, 44.0])
+
+    # test to make sure the object is correctly populated
+    assert test_object.H_adler == [[15.0, 25.0], [35.0], [45.0]]
+    assert test_object.H_err_adler == [[16.0, 26.0], [36.0], [46.0]]
+    assert test_object.phase_parameter_1 == [[17.0, 27.0], [37.0], [47.0]]
+    assert test_object.phase_parameter_1_err == [[18.0, 28.0], [38.0], [48.0]]
+    assert test_object.phase_parameter_2 == [[None, 29.0], [None], [49.0]]
+    assert test_object.phase_parameter_2_err == [[None, 30.0], [None], [50.0]]
+
+    # testing to make sure the correct error messages trigger
+    with pytest.raises(Exception) as error_info_1:
+        test_object.populate_phase_parameters(
+            "u", "model_1", 11.0, 12.0, 13, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0
+        )
+
+    assert (
+        error_info_1.value.args[0]
+        == "If using a model with 2 phase parameters, both parameter_2 and parameter_2_err must be supplied."
+    )
+
+    with pytest.raises(Exception) as error_info_2:
+        test_object.populate_phase_parameters("y", "model_1", 11.0, 12.0, 13, 14.0, 15.0, 16.0, 17.0, 18.0)
+
+    assert error_info_2.value.args[0] == "Filter y is not in supplied filter list."