union domain implementation (#109)

* Union Class
* Implemented Union Tests

---------

Co-authored-by: Dario Coscia <93731561+dario-coscia@users.noreply.github.com>

pina/geometry/__init__.py

@@ -2,9 +2,11 @@ __all__ = [
     'Location',
     'CartesianDomain',
     'EllipsoidDomain',
+    'Union'
 ]
 
 from .location import Location
 from .cartesian import CartesianDomain
 from .ellipsoid import EllipsoidDomain
-from .difference_domain import Difference
+from .difference_domain import Difference
+from .union_domain import Union

pina/geometry/union_domain.py

@@ -0,0 +1,137 @@
+import torch
+from .location import Location
+from ..utils import check_consistency
+from ..label_tensor import LabelTensor
+
+
+class Union(Location):
+    """ PINA implementation of Unions of Domains."""
+
+    def __init__(self, geometries):
+        """ PINA implementation of Unions of Domains.
+
+        :param list geometries: A list of geometries from 'pina.geometry' 
+            such as 'EllipsoidDomain' or 'CartesianDomain'.
+
+        :Example:
+            # Create two ellipsoid domains
+            >>> ellipsoid1 = EllipsoidDomain({'x': [-1, 1], 'y': [-1, 1]})
+            >>> ellipsoid2 = EllipsoidDomain({'x': [0, 2], 'y': [0, 2]})
+
+            # Create a union of the ellipsoid domains
+            >>> union = GeometryUnion([ellipsoid1, ellipsoid2])
+
+        """
+        super().__init__()
+
+        # union checks
+        self._check_union_inheritance(geometries)
+        self._check_union_consistency(geometries)
+        
+        # assign geometries
+        self._geometries = geometries
+
+    @property
+    def geometries(self):
+        """ 
+        The geometries."""
+        return self._geometries
+        
+    @property
+    def variables(self):
+        """
+        Spatial variables.
+
+        :return: All the spatial variables defined in '__init__()' in order.
+        :rtype: list[str]
+        """
+        all_variables = []
+        seen_variables = set()
+        for geometry in self.geometries:
+            for variable in geometry.variables:
+                if variable not in seen_variables:
+                    all_variables.append(variable)
+                    seen_variables.add(variable)
+        return all_variables
+
+    def is_inside(self, point, check_border=False):
+        """Check if a point is inside the union domain.
+
+        :param point: Point to be checked.
+        :type point: LabelTensor
+        :param check_border: Check if the point is also on the frontier
+            of the ellipsoid, default False.
+        :type check_border: bool
+        :return: Returning True if the point is inside, False otherwise.
+        :rtype: bool
+        """
+        for geometry in self.geometries:
+            if geometry.is_inside(point, check_border):
+                return True
+        return False
+
+    def sample(self, n, mode='random', variables='all'):
+        """Sample routine.
+
+        :param n: Number of points to sample in the shape.
+        :type n: int
+        :param mode: Mode for sampling, defaults to 'random'.
+            Available modes include: random sampling, 'random'.
+        :type mode: str, optional
+        :param variables: pinn variable to be sampled, defaults to 'all'.
+        :type variables: str or list[str], optional
+
+        :Example:
+            # Create two ellipsoid domains
+            >>> ellipsoid1 = EllipsoidDomain({'x': [-1, 1], 'y': [-1, 1]})
+            >>> ellipsoid2 = EllipsoidDomain({'x': [0, 2], 'y': [0, 2]})
+
+            # Create a union of the ellipsoid domains
+            >>> union = GeometryUnion([ellipsoid1, ellipsoid2])
+
+            >>> union.sample(n=1000)
+                LabelTensor([[-0.2025,  0.0072],
+                            [ 0.0358,  0.5748],
+                            [ 0.5083,  0.0482],
+                            ...,
+                            [ 0.5857,  0.9279],
+                            [ 1.1496,  1.7339],
+                            [ 0.7650,  1.0469]])
+
+            >>> len(union.sample(n=1000)
+                1000
+        """
+        sampled_points = []
+
+        # calculate the number of points to sample for each geometry and the remainder
+        remainder = n % len(self.geometries)
+        num_points = n // len(self.geometries)
+
+        # sample the points
+        for i, geometry in enumerate(self.geometries):
+            # add to sample total if remainder is not 0
+            if i < remainder:
+                num_points += 1
+            sampled_points.append(geometry.sample(num_points, mode, variables))
+
+        return LabelTensor(torch.cat(sampled_points), labels=[f'{i}' for i in self.variables])
+
+    def _check_union_consistency(self, geometries):
+        """Check if the dimensions of the geometries are consistent.
+
+        :param geometries: Geometries to be checked.
+        :type geometries: list[Location]
+        """
+        for geometry in geometries:
+            if geometry.variables != geometries[0].variables:
+                raise NotImplementedError(
+                    f'The geometries need to be the same dimensions. {geometry.variables} is not equal to {geometries[0].variables}')
+
+    def _check_union_inheritance(self, geometries):
+        """Check if the geometries are inherited from 'pina.geometry.Location'.
+
+        param geometries: Geometries to be checked.
+        :type geometries: list[Location]
+        """
+        for idx, geometry in enumerate(geometries):
+            check_consistency(geometry, Location, f'geometry[{idx}]')

tests/test_union.py

@@ -0,0 +1,46 @@
+import torch
+
+from pina import LabelTensor
+from pina.geometry import Union, EllipsoidDomain, CartesianDomain
+
+
+def test_constructor_two_CartesianDomains():
+    Union([CartesianDomain({'x': [0, 1], 'y': [0, 1]}),
+           CartesianDomain({'x': [0.5, 2], 'y': [-1, 0.1]})])
+
+
+def test_constructor_two_EllipsoidDomains():
+    Union([EllipsoidDomain({'x': [-1, 1], 'y': [-1, 1], 'z': [-1, 1]}),
+           EllipsoidDomain({'x': [-0.5, 0.5], 'y': [-0.5, 0.5], 'z': [-0.5, 0.5]})])
+
+
+def test_constructor_EllipsoidDomain_CartesianDomain():
+    Union([EllipsoidDomain({'x': [-1, 1], 'y': [-1, 1]}),
+           CartesianDomain({'x': [-0.5, 0.5], 'y': [-0.5, 0.5]})])
+
+
+def test_is_inside_two_CartesianDomains():
+    pt_1 = LabelTensor(torch.tensor([[0.5, 0.5]]), ['x', 'y'])
+    pt_2 = LabelTensor(torch.tensor([[-1, -1]]), ['x', 'y'])
+    domain = Union([CartesianDomain({'x': [0, 1], 'y': [0, 1]}),
+                    CartesianDomain({'x': [0.5, 2], 'y': [-1, 0.1]})])
+    assert domain.is_inside(pt_1) == True
+    assert domain.is_inside(pt_2) == False
+
+
+def test_is_inside_two_EllipsoidDomains():
+    pt_1 = LabelTensor(torch.tensor([[0.5, 0.5, 0.5]]), ['x', 'y', 'z'])
+    pt_2 = LabelTensor(torch.tensor([[-1, -1, -1]]), ['x', 'y', 'z'])
+    domain = Union([EllipsoidDomain({'x': [-1, 1], 'y': [-1, 1], 'z': [-1, 1]}),
+                    EllipsoidDomain({'x': [-0.5, 0.5], 'y': [-0.5, 0.5], 'z': [-0.5, 0.5]})])
+    assert domain.is_inside(pt_1) == True
+    assert domain.is_inside(pt_2) == False
+
+
+def test_is_inside_EllipsoidDomain_CartesianDomain():
+    pt_1 = LabelTensor(torch.tensor([[0.5, 0.5]]), ['x', 'y'])
+    pt_2 = LabelTensor(torch.tensor([[-1, -1]]), ['x', 'y'])
+    domain = Union([EllipsoidDomain({'x': [-1, 1], 'y': [-1, 1], }),
+                    CartesianDomain({'x': [0.6, 1.5], 'y': [-2, 0]})])
+    assert domain.is_inside(pt_1) == True
+    assert domain.is_inside(pt_2) == False