diff --git a/deeptrack/tests/test_scatterers.py b/deeptrack/tests/test_scatterers.py index de82c6552..ca926d855 100644 --- a/deeptrack/tests/test_scatterers.py +++ b/deeptrack/tests/test_scatterers.py @@ -1,17 +1,33 @@ -import sys - -sys.path.append(".") # Adds the module to path +# Use this only when running the test locally. +# import sys +# sys.path.append(".") # Adds the module to path import unittest -from deeptrack import scatterers - import numpy as np + from deeptrack.optics import Fluorescence, Brightfield -from deeptrack.image import Image +from deeptrack import scatterers +from deeptrack.backend import TORCH_AVAILABLE, xp +from deeptrack.tests import BackendTestBase + +if TORCH_AVAILABLE: + import torch + + +class TestScatterers_NumPy(BackendTestBase): + BACKEND = "numpy" + + @property + def array_type(self): + if self.BACKEND == "numpy": + return np.ndarray + elif self.BACKEND == "torch": + return torch.Tensor + else: + raise ValueError(f"Unsupported backend: {self.BACKEND}") -class TestScatterers(unittest.TestCase): def test_PointParticle(self): optics = Fluorescence( NA=0.7, @@ -27,7 +43,7 @@ def test_PointParticle(self): ) imaged_scatterer = optics(scatterer) output_image = imaged_scatterer.resolve() - self.assertIsInstance(output_image, np.ndarray) + self.assertIsInstance(output_image, self.array_type) self.assertEqual(output_image.shape, (64, 64, 1)) def test_Ellipse(self): @@ -48,7 +64,7 @@ def test_Ellipse(self): ) imaged_scatterer = optics(scatterer) output_image = imaged_scatterer.resolve() - self.assertIsInstance(output_image, np.ndarray) + self.assertIsInstance(output_image, self.array_type) self.assertEqual(output_image.shape, (64, 64, 1)) def test_EllipseUpscale(self): @@ -146,7 +162,7 @@ def test_Sphere(self): ) imaged_scatterer = optics(scatterer) output_image = imaged_scatterer.resolve() - self.assertIsInstance(output_image, np.ndarray) + self.assertIsInstance(output_image, self.array_type) self.assertEqual(output_image.shape, (64, 64, 1)) def test_SphereUpscale(self): @@ -188,7 +204,7 @@ def test_Ellipsoid(self): ) imaged_scatterer = optics(scatterer) output_image = imaged_scatterer.resolve() - self.assertIsInstance(output_image, np.ndarray) + self.assertIsInstance(output_image, self.array_type) self.assertEqual(output_image.shape, (64, 64, 1)) def test_EllipsoidUpscale(self): @@ -343,6 +359,12 @@ def test_MieStratifiedSphere(self): imaged_scatterer_1 = optics_1(scatterer) imaged_scatterer_1.update().resolve() +# TODO: Extending the test and setting the backend to torch +# @unittest.skipUnless(TORCH_AVAILABLE, "PyTorch is not installed.") +# class TestScatterers_PyTorch(TestScatterers_NumPy): +# BACKEND = "torch" +# pass + if __name__ == "__main__": unittest.main()