Feature/updated tests

src/deepbench/astro_object/galaxy_object.py

@@ -1,5 +1,5 @@
 from src.deepbench.astro_object.astro_object import AstroObject
-from src.deepbench.shape_generator.shape_generator import ShapeGenerator
+from src.deepbench.image.sky_image import SkyImage
 from astropy.modeling.models import Sersic2D
 
 import numpy as np
@@ -15,7 +15,7 @@ def __init__(
         amplitude=1,
         radius=25,
         n=1.0,
-        noise_level = 0.2,
+        noise_level=0.2,
         ellipse=random.uniform(0.1, 0.9),
         theta=random.uniform(-1.5, 1.5),
     ):
@@ -37,7 +37,8 @@ def __init__(
             image_dimensions=image_dimensions,
             radius=radius,
             amplitude=amplitude,
-            noise_level=noise_level)
+            noise_level=noise_level,
+        )
 
         self._n = n
         self._ellipse = ellipse
@@ -58,7 +59,7 @@ def create_Sersic_profile(self, center_x, center_y):
 
         return profile(x, y)
 
-    def create_object(self, center_x = 5.0, center_y = 5.0) -> np.ndarray:
+    def create_object(self, center_x=5.0, center_y=5.0) -> np.ndarray:
         """
         Create the star object from a Moffat distribution and Poisson and PSF noise.
 
@@ -76,14 +77,10 @@ def create_object(self, center_x = 5.0, center_y = 5.0) -> np.ndarray:
         """
 
         # Create the empty image shape.
-        image_shape = ShapeGenerator.create_empty_shape(self._image)
-
         # Create the Poisson noise profile specific to Galaxy objects.
         noise_profile = self.create_noise(galaxy=True)
 
-        image_shape = self.create_Sersic_profile(
-            center_x=center_x, center_y=center_y
-        )
+        image_shape = self.create_Sersic_profile(center_x=center_x, center_y=center_y)
 
         # Append the noise profiles to the object.
         image_shape += noise_profile

src/deepbench/astro_object/spiral_galaxy.py

@@ -3,19 +3,19 @@
 from numpy import log2, random
 from numpy import tan
 
+
 class SpiralGalaxyObject(GalaxyObject):
     def __init__(
         self,
         image_dimensions,
         amplitude=1,
         radius=25,
         n=1.0,
-        noise_level = 0.2,
+        noise_level=0.2,
         ellipse=random.uniform(0.1, 0.9),
         theta=random.uniform(-1.5, 1.5),
-        winding_number:int=2,
-        spiral_pitch:float=0.2
-
+        winding_number: int = 2,
+        spiral_pitch: float = 0.2,
     ):
 
         self.winding_number = winding_number
@@ -28,22 +28,24 @@ def __init__(
             n=n,
             ellipse=ellipse,
             theta=theta,
-            noise_level=noise_level
+            noise_level=noise_level,
         )
 
     def create_sprial_profile(self, center_x, center_y):
         "ref paper: https://doi.org/10.1111/j.1365-2966.2009.14950.x"
 
-        spiral = self._amplitude/log2(self.spiral_pitch*tan(self.theta/(2*self.winding_number) ))
+        spiral = self._amplitude / log2(
+            self.spiral_pitch * tan(self._theta / (2 * self.winding_number))
+        )
         # TO BE IMPLEMENTED.
         # WHERE IS THE CODE FOR THIS????????
-        return spiral
+        return random.default_rng().uniform(size=self._image.shape) * spiral
 
     def create_object(self, center_x, center_y):
+        image_shape = self._image.copy()
 
-        image_shape = self.create_empty_shape(self._image)
         # Add the spiral to the image
-        spiral = self.create_sprial_profile(self, center_x, center_y)
+        spiral = self.create_sprial_profile(center_x, center_y)
         image_shape += spiral
 
         # Create the Poisson noise profile.

src/deepbench/image/image.py

@@ -13,6 +13,9 @@ def __init__(
         object_noise_type: str = "gaussian",
         object_noise_level: float = 0.0,
     ):
+        assert len(image_shape) >= 2
+        "All images must be in at least 2d."
+
         self.image_shape = image_shape
         self.image = np.zeros(self.image_shape)
 
@@ -43,14 +46,10 @@ def generate_noise(self, seed=42):
         }
 
         if self.object_noise_type not in noise_map.keys():
-            raise NotImplementedError(
-                f"{self.object_noise_type} noise type not available"
-            )
-
-        assert self.image is not None, "Image not generated, please run combine_objects"
+            raise NotImplementedError(f"{self.object_noise_type} noise not available")
 
         noise = noise_map[self.object_noise_type](seed)
-        self.image += noise
+        return noise
 
     def save_image(self, save_dir="results", image_name="image_1", image_format="jpg"):
         """

src/deepbench/image/shape_image.py

@@ -13,7 +13,6 @@ def __init__(
     ):
         self.shapes = ShapeGenerator(image_shape=image_shape)
         self.method_map = self._get_methods()
-        print(self.method_map)
         super().__init__(
             image_shape=image_shape,
             object_noise_level=object_noise_level,
@@ -32,7 +31,9 @@ def _get_methods(self):
         return {method[0].split("_")[-1]: method[1] for method in methods}
 
     def _create_object(self, shape, shape_params):
-        assert shape in self.method_map.keys()
+
+        if shape not in self.method_map.keys():
+            raise NotImplementedError()
         return self.method_map[shape](self.shapes, **shape_params)
 
     def combine_objects(self, objects, instance_params, object_params, seed=42):
@@ -50,10 +51,18 @@ def combine_objects(self, objects, instance_params, object_params, seed=42):
         }]
 
         """
-        self.image = self.shapes.create_empty_shape()
+        image = self.shapes.create_empty_shape()
+
+        if type(objects) == str:
+            objects = [objects]
+        if type(instance_params) == dict:
+            instance_params = [instance_params]
+        if type(object_params) == dict:
+            object_params = [object_params]
 
         for shape, _ in zip(objects, instance_params):
             for object in object_params:
-                self.image += self._create_object(shape, object)
-        self.generate_noise(seed)
-        return self.image
+                image += self._create_object(shape, object)
+        noise = self.generate_noise(seed)
+        image += noise
+        return image