Merge pull request #23 from LCOGT/feature/django-test

util & operation tests, stack arrays bugfix, comments

datalab/datalab_session/analysis/line_profile.py

@@ -8,6 +8,15 @@
 def line_profile(input: dict):
   """
     Creates an array of luminosity values and the length of the line in arcseconds
+    input = {
+      basename (str): The name of the file to analyze
+      height (int): The height of the image
+      width (int): The width of the image
+      x1 (int): The x coordinate of the starting point
+      y1 (int): The y coordinate of the starting point
+      x2 (int): The x coordinate of the ending point
+      y2 (int): The y coordinate of the ending point
+    }
   """
   sci_hdu = get_hdu(input['basename'], 'SCI')
 
@@ -21,20 +30,26 @@ def line_profile(input: dict):
   try: 
     wcs = WCS(sci_hdu.header)
 
+    if(wcs.get_axis_types()[0].get('coordinate_type') == None):
+      raise WcsError("No valid WCS solution")
+
     start_sky_coord = wcs.pixel_to_world(x_points[0], y_points[0])
     end_sky_coord = wcs.pixel_to_world(x_points[1], y_points[1])
 
     arcsec_angle = start_sky_coord.separation(end_sky_coord).arcsecond
 
     start_coords = [start_sky_coord.ra.deg, start_sky_coord.dec.deg]
     end_coords = [end_sky_coord.ra.deg, end_sky_coord.dec.deg]
-
-  except WcsError as e:
-    error = f'{input["basename"]} does not have a valid WCS header'
-    print(error, e)
-    # if theres no valid WCS solution then we default to using pixscale to calculate the angle, and no coordinates
+  except WcsError:
+    # no valid WCS solution
     start_coords = None
     end_coords = None
-    arcsec_angle = len(line_profile) * sci_hdu.header["PIXSCALE"]
+
+    try:
+      # attempt using pixscale to calculate the angle
+      arcsec_angle = len(line_profile) * sci_hdu.header["PIXSCALE"]
+    except KeyError as e:
+      # no valid WCS solution, and no pixscale
+      arcsec_angle = None
 
   return {"line_profile": line_profile, "arcsec": arcsec_angle, "start_coords": start_coords, "end_coords": end_coords}

datalab/datalab_session/file_utils.py

@@ -81,15 +81,15 @@ def create_jpgs(cache_key, fits_paths: str, color=False) -> list:
 
 def stack_arrays(array_list: list):
   """
-  Takes a list of numpy arrays, crops them to an equal shape, and stacks them to be a 3d numpy array
-
+  Takes a list of numpy arrays from fits images and stacks them to be a 3d numpy array
+  cropped since fits images can be different sizes
   """
-  min_shape = min(arr.shape for arr in array_list)
-  cropped_data_list = [arr[:min_shape[0], :min_shape[1]] for arr in array_list]
+  min_x = min(arr.shape[0] for arr in array_list)
+  min_y = min(arr.shape[1] for arr in array_list)
 
-  stacked = np.stack(cropped_data_list, axis=2)
+  cropped_data_list = [arr[:min_x, :min_y] for arr in array_list]
 
-  return stacked
+  return np.stack(cropped_data_list, axis=2)
 
 def scale_points(height_1: int, width_1: int, height_2: int, width_2: int, x_points=[], y_points=[], flip_y = False, flip_x = False):
   """

datalab/datalab_session/tests/test_analysis.py

@@ -0,0 +1,48 @@
+from unittest import mock
+import json
+
+from django.test import TestCase
+from numpy.testing import assert_almost_equal
+
+from datalab.datalab_session.analysis import line_profile, source_catalog
+
+class TestAnalysis(TestCase):
+    analysis_test_path = 'datalab/datalab_session/tests/test_files/analysis/'
+    analysis_fits_1_path = f'{analysis_test_path}fits_1.fits.fz'
+
+    def setUp(self):
+        with open(f'{self.analysis_test_path}test_line_profile.json') as f:
+            self.test_line_profile_data = json.load(f)['test_line_profile']
+
+        with open(f'{self.analysis_test_path}test_source_catalog.json') as f:
+            self.test_source_catalog_data = json.load(f)['test_source_catalog']
+
+    @mock.patch('datalab.datalab_session.file_utils.get_fits')
+    def test_line_profile(self, mock_get_fits):
+
+        mock_get_fits.return_value = self.analysis_fits_1_path
+
+        output = line_profile.line_profile({
+            'basename': 'fits_1',
+            'height': 100,
+            'width': 100,
+            'x1': 25,
+            'y1': 25,
+            'x2': 75,
+            'y2': 75
+        })
+
+        assert_almost_equal(output.get('line_profile').tolist(), self.test_line_profile_data, decimal=3)
+
+    @mock.patch('datalab.datalab_session.file_utils.get_fits')
+    def test_source_catalog(self, mock_get_fits):
+
+        mock_get_fits.return_value = self.analysis_fits_1_path
+
+        output = source_catalog.source_catalog({
+            'basename': 'fits_1',
+            'height': 100,
+            'width': 100,
+            })
+
+        self.assertEqual(output, self.test_source_catalog_data)

datalab/datalab_session/tests/test_files/analysis/test_line_profile.json

@@ -0,0 +1,14 @@
+{
+  "test_line_profile": [135.66614, 187.97342, 186.27275, 152.81543, 200.08119, 190.92934,
+       190.35522, 200.84515, 185.5646 , 182.38306, 182.62105, 195.83347,
+       181.14099, 170.32465, 193.30614, 189.06567, 172.27455, 194.88974,
+       184.23103, 183.86188, 177.11359, 184.32747, 185.9386 , 177.15134,
+       171.93048, 183.24039, 188.6636 , 171.44247, 175.03983, 169.98991,
+       201.46959, 173.0296 , 172.73636, 169.93402, 170.76671, 177.4004 ,
+       182.55852, 190.3085 , 181.39314, 181.369  , 171.61351, 158.37466,
+       178.958  , 168.84396, 170.62454, 186.83702, 173.42389, 165.36671,
+       186.99101, 181.13121, 153.41772, 185.963  , 165.45833, 175.77798,
+       180.02434, 177.85484, 182.0603 , 178.34627, 178.07858, 180.45233,
+       170.2989 , 176.22253, 169.75783, 171.57878, 180.5748 , 180.7788 ,
+       185.37752, 180.1012 , 157.70186, 181.05095, 168.38719, 159.55258]
+}

datalab/datalab_session/tests/test_files/analysis/test_source_catalog.json

@@ -0,0 +1,54 @@
+{
+  "test_source_catalog": [
+    {"x": 792, "y": 623, "flux": 3164455, "ra": "260.9333", "dec": "-33.0225"},
+    {"x": 919, "y": 120, "flux": 807813, "ra": "260.9667", "dec": "-32.9137"}, 
+    {"x": 996, "y": 627, "flux": 383099, "ra": "260.9861", "dec": "-33.0235"}, 
+    {"x": 997, "y": 641, "flux": 335359, "ra": "260.9862", "dec": "-33.0265"}, 
+    {"x": 587, "y": 803, "flux": 304833, "ra": "260.8802", "dec": "-33.0612"}, 
+    {"x": 1015, "y": 554, "flux": 291847, "ra": "260.9910", "dec": "-33.0076"}, 
+    {"x": 490, "y": 568, "flux": 226494, "ra": "260.8555", "dec": "-33.0101"}, 
+    {"x": 526, "y": 147, "flux": 209307, "ra": "260.8654", "dec": "-32.9191"}, 
+    {"x": 740, "y": 301, "flux": 177543, "ra": "260.9203", "dec": "-32.9526"}, 
+    {"x": 334, "y": 231, "flux": 168666, "ra": "260.8157", "dec": "-32.9370"}, 
+    {"x": 843, "y": 277, "flux": 132858, "ra": "260.9468", "dec": "-32.9476"}, 
+    {"x": 776, "y": 805, "flux": 115180, "ra": "260.9290", "dec": "-33.0618"}, 
+    {"x": 13, "y": 341, "flux": 101317, "ra": "260.7327", "dec": "-32.9604"}, 
+    {"x": 384, "y": 332, "flux": 93054, "ra": "260.8284", "dec": "-32.9590"}, 
+    {"x": 451, "y": 307, "flux": 82848, "ra": "260.8457", "dec": "-32.9536"}, 
+    {"x": 830, "y": 815, "flux": 76200, "ra": "260.9427", "dec": "-33.0641"}, 
+    {"x": 846, "y": 138, "flux": 70788, "ra": "260.9478", "dec": "-32.9176"}, 
+    {"x": 569, "y": 179, "flux": 69761, "ra": "260.8763", "dec": "-32.9261"}, 
+    {"x": 818, "y": 146, "flux": 69329, "ra": "260.9406", "dec": "-32.9191"}, 
+    {"x": 843, "y": 912, "flux": 63748, "ra": "260.9462", "dec": "-33.0851"}, 
+    {"x": 395, "y": 151, "flux": 60681, "ra": "260.8315", "dec": "-32.9198"}, 
+    {"x": 542, "y": 122, "flux": 60653, "ra": "260.8696", "dec": "-32.9136"}, 
+    {"x": 214, "y": 390, "flux": 60057, "ra": "260.7845", "dec": "-32.9712"}, 
+    {"x": 188, "y": 691, "flux": 59835, "ra": "260.7774", "dec": "-33.0363"}, 
+    {"x": 173, "y": 438, "flux": 59768, "ra": "260.7738", "dec": "-32.9816"}, 
+    {"x": 1022, "y": 188, "flux": 56811, "ra": "260.9931", "dec": "-32.9285"}, 
+    {"x": 431, "y": 936, "flux": 56575, "ra": "260.8396", "dec": "-33.0898"}, 
+    {"x": 333, "y": 112, "flux": 51136, "ra": "260.8157", "dec": "-32.9112"}, 
+    {"x": 211, "y": 958, "flux": 50722, "ra": "260.7828", "dec": "-33.0941"}, 
+    {"x": 664, "y": 535, "flux": 45352, "ra": "260.9004", "dec": "-33.0033"}, 
+    {"x": 989, "y": 975, "flux": 44414, "ra": "260.9838", "dec": "-33.0988"}, 
+    {"x": 1003, "y": 463, "flux": 38238, "ra": "260.9878", "dec": "-32.9879"}, 
+    {"x": 336, "y": 242, "flux": 37970, "ra": "260.8161", "dec": "-32.9393"}, 
+    {"x": 294, "y": 565, "flux": 37707, "ra": "260.8050", "dec": "-33.0093"}, 
+    {"x": 535, "y": 916, "flux": 37575, "ra": "260.8666", "dec": "-33.0854"}, 
+    {"x": 66, "y": 415, "flux": 37365, "ra": "260.7463", "dec": "-32.9765"}, 
+    {"x": 533, "y": 91, "flux": 36539, "ra": "260.8673", "dec": "-32.9069"}, 
+    {"x": 879, "y": 505, "flux": 35138, "ra": "260.9559", "dec": "-32.9970"}, 
+    {"x": 293, "y": 638, "flux": 34703, "ra": "260.8045", "dec": "-33.0250"}, 
+    {"x": 444, "y": 635, "flux": 33406, "ra": "260.8434", "dec": "-33.0245"}, 
+    {"x": 400, "y": 82, "flux": 33366, "ra": "260.8330", "dec": "-32.9048"}, 
+    {"x": 465, "y": 59, "flux": 32920, "ra": "260.8497", "dec": "-32.8998"}, 
+    {"x": 874, "y": 50, "flux": 32280, "ra": "260.9553", "dec": "-32.8984"}, 
+    {"x": 1007, "y": 416, "flux": 29422, "ra": "260.9892", "dec": "-32.9779"}, 
+    {"x": 281, "y": 294, "flux": 28626, "ra": "260.8019", "dec": "-32.9506"}, 
+    {"x": 682, "y": 39, "flux": 28508, "ra": "260.9057", "dec": "-32.8959"}, 
+    {"x": 972, "y": 426, "flux": 27457, "ra": "260.9801", "dec": "-32.9799"}, 
+    {"x": 480, "y": 557, "flux": 26744, "ra": "260.8529", "dec": "-33.0077"}, 
+    {"x": 812, "y": 291, "flux": 25618, "ra": "260.9389", "dec": "-32.9506"}, 
+    {"x": 443, "y": 206, "flux": 25483, "ra": "260.8438", "dec": "-32.9317"}
+  ]
+}

datalab/datalab_session/tests/test_files/file_extended_test_case.py

@@ -0,0 +1,23 @@
+import pathlib as pl
+from hashlib import md5
+from os import path, remove, listdir
+
+from django.test import TestCase
+
+# extending the TestCase class to include a custom assertions for file operations
+class FileExtendedTestCase(TestCase):
+    def assertIsFile(self, path):
+        if not pl.Path(path).resolve().is_file():
+            raise AssertionError("File does not exist: %s" % str(path))
+
+    def assertFilesEqual(self, image_1: str, image_2: str):
+        with open(image_1, 'rb') as file_1, open(image_2, 'rb') as file_2:
+            self.assertEqual(md5(file_1.read()).hexdigest(), md5(file_2.read()).hexdigest())
+
+    def clean_test_dir(self):
+        test_files_dir = 'datalab/datalab_session/tests/test_files'
+        for file_name in listdir(test_files_dir):
+            if file_name.startswith('temp_'):
+                file_path = path.join(test_files_dir, file_name)
+                if path.isfile(file_path):
+                    remove(file_path)