pyect 0.1.0__py3-none-any.whl

pyect/__init__.py

@@ -0,0 +1,9 @@
+from .wect import WECT
+from .tensor_complex import Complex
+from .directions import sample_directions_2d, sample_directions_3d
+from .image_ecf import Image_ECF_2D, Image_ECF_3D
+from .preprocessing.image_processing import (
+    weighted_freudenthal,
+    weighted_cubical,
+    image_to_grayscale_tensor
+)

pyect/directions.py

@@ -0,0 +1,28 @@
+import math
+import torch
+
+golden_angle = math.pi * (3.0 - math.sqrt(5.0))
+
+def sample_directions_2d(num_dirs: int, *, device=None):
+    """
+    Sample num_dirs directions evenly from S^1.
+    """
+
+    angles = 2 * math.pi * torch.arange(num_dirs, dtype=torch.float32, device=device) / num_dirs
+    directions = torch.stack([torch.cos(angles), torch.sin(angles)], dim=-1)
+    return directions.contiguous()
+
+def sample_directions_3d(num_dirs: int, *, device=None):
+    """
+    Sample num_dirs directions from S^2 using the Fibonacci spiral method.
+    """
+
+    i = torch.arange(num_dirs, dtype=torch.float32, device=device)
+    theta = golden_angle * i
+    y = 1.0 - (2.0 * (i + 0.5) / num_dirs)
+    r = torch.sqrt(torch.clamp(1.0 - y * y, min=0.0))
+    x = torch.cos(theta) * r
+    z = torch.sin(theta) * r
+    directions = torch.stack([x, y, z], dim=-1)
+
+    return directions.contiguous()

pyect/dtypes.py

@@ -0,0 +1,7 @@
+"""Export types for the PyECT package."""
+
+import torch
+
+COORDS_DTYPE = torch.float32
+INDICES_DTYPE = torch.int64
+WEIGHTS_DTYPE = torch.float32

pyect/image_ecf.py

@@ -0,0 +1,220 @@
+"""For computing the ECF of 2- and 3-dimensional images filtered by pixel intensity"""
+
+import torch
+from typing import List
+
+class Image_ECF_2D(torch.nn.Module):
+    """A torch module for computing the ECF of a 2D image filtered by pixel intensity.
+
+    This module may be used just for computing the ECF of images, or used as a layer in a neural network.
+    Internally, the module stores the number of values used for sampling, so repeated forward calls
+    do not require this parameters to be passed in, and allow streamlined loading/saving of the module for consistent
+    computation.
+
+    This module can also be converted to TorchScript using torch.jit.script for use
+    outside of Python.
+    """
+
+    def __init__(self, num_vals: int) -> None:
+        """Initializes the image_ECF module.
+
+        The initialized module is designed to compute the ECF of a 2D image, discretized by sampling num_vals values.
+
+        Args:
+            num_vals: The number of values to discretize the ECF over.
+        """
+        super().__init__()
+        self.num_vals: int = int(num_vals)
+
+    @staticmethod
+    def cell_values_2D(arr: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Creates a cubical complex with a function on its cells from a 2D tensor.
+        The structure of the cubical complex is ignored with only the function values on the cells
+        being recorded.
+
+        Args:
+            arr (torch.Tensor): A 2D tensor with values between 0 and 1.
+
+        Returns:
+            vertex_values (torch.Tensor): A 1D tensor containing the function values of each vertex.
+            edge_values (torch.Tensor): A 1D tensor containing the function values of each edge.
+            square_values (torch.Tensor): A 1D tensor containing the function values of each square.
+        """
+        arr = arr.float()
+
+        vertex_values = arr.reshape(-1)
+
+        x_edge_values = torch.maximum(arr[1:, :], arr[:-1, :])
+        y_edge_values = torch.maximum(arr[:, 1:], arr[:, :-1])
+        edge_values = torch.cat([
+            x_edge_values.reshape(-1),
+            y_edge_values.reshape(-1)
+        ], dim=0)
+
+        square_values = torch.maximum(y_edge_values[1:, :], y_edge_values[:-1, :]).reshape(-1)
+
+        return [vertex_values, edge_values, square_values]
+    
+    def forward(self, img_arr: torch.Tensor) -> torch.Tensor:
+        """
+        Calculates a discretization of the ECF of a 2D image.
+
+        Args:
+            img_arr (torch.Tensor): a 2D tensor with values between 0 and 1.
+
+        Returns:
+            ecf (torch.Tensor): A 1D tensor of shape (self.num_vals) containing the ECF.
+        """
+
+        device = img_arr.device
+        n = self.num_vals
+        vertex_values, edge_values, square_values = self.cell_values_2D(img_arr)
+
+        vertex_indices = torch.ceil(vertex_values * (n-1)).long()
+        edge_indices = torch.ceil(edge_values * (n-1)).long()
+        square_indices = torch.ceil(square_values * (n-1)).long()
+
+        diff_ecf = torch.zeros(n, dtype=torch.int32, device=device)
+
+        # Add the contribution of the vertices
+        diff_ecf.scatter_add_(
+            0,
+            vertex_indices,
+            torch.ones_like(vertex_indices, dtype=torch.int32)
+        )
+
+        # Add the contribution of the edges
+        diff_ecf.scatter_add_(
+            0,
+            edge_indices,
+            -1 * torch.ones_like(edge_indices, dtype=torch.int32)
+        )
+
+        # Add the contribution of the squares
+        diff_ecf.scatter_add_(
+            0,
+            square_indices,
+            torch.ones_like(square_indices, dtype=torch.int32)
+        )
+
+        return torch.cumsum(diff_ecf, dim=0)
+
+
+class Image_ECF_3D(torch.nn.Module):
+    """A torch module for computing the ECF of a 3D image filtered by pixel intensity.
+
+    This module may be used just for computing the ECF of images, or used as a layer in a neural network.
+    Internally, the module stores the number of values used for sampling, so repeated forward calls
+    do not require this parameters to be passed in, and allow streamlined loading/saving of the module for consistent
+    computation.
+
+    This module can also be converted to TorchScript using torch.jit.script for use
+    outside of Python.
+    """
+
+    def __init__(self, num_vals: int) -> None:
+        """Initializes the image_ECF module.
+
+        The initialized module is designed to compute the ECF of a 3D image, discretized by sampling num_vals values.
+
+        Args:
+            num_vals: The number of values to discretize the ECF over.
+        """
+        super().__init__()
+        self.num_vals: int = int(num_vals)
+
+    @staticmethod
+    def cell_values_3D(arr: torch.Tensor) -> List[torch.Tensor]:
+        """
+        Creates a cubical complex with a function on its cells from a 3D tensor.
+        The structure of the cubical complex is ignored with only the function values on the cells
+        being recorded.
+
+        Args:
+            arr (torch.Tensor): A 3D tensor with values between 0 and 1.
+
+        Returns:
+            vertex_values (torch.Tensor): A 1D tensor containing the function values of each vertex.
+            edge_values (torch.Tensor): A 1D tensor containing the function values of each edge.
+            square_values (torch.Tensor): A 1D tensor containing the function values of each square.
+            cube_values (torch.Tensor): A 1D tensor containing the function values of each cube.
+        """
+        arr = arr.float()
+
+        vertex_values = arr.reshape(-1)
+
+        x_edge_values = torch.maximum(arr[1:, ...], arr[:-1, ...])
+        y_edge_values = torch.maximum(arr[:, 1:, :], arr[:, :-1, :])
+        z_edge_values = torch.maximum(arr[..., 1:], arr[..., :-1])
+        edge_values = torch.cat([
+            x_edge_values.reshape(-1),
+            y_edge_values.reshape(-1),
+            z_edge_values.reshape(-1)
+            ], dim=0)
+
+        x_square_values = torch.maximum(y_edge_values[..., 1:], y_edge_values[..., :-1])
+        y_square_values = torch.maximum(z_edge_values[1:, ...], z_edge_values[:-1, ...])
+        z_square_values = torch.maximum(x_edge_values[:, 1:, :], x_edge_values[:, :-1, :])
+        square_values = torch.cat([
+            x_square_values.reshape(-1),
+            y_square_values.reshape(-1),
+            z_square_values.reshape(-1)
+        ], dim=0)
+
+        cube_values = torch.maximum(x_square_values[1:, ...], x_square_values[:-1, ...]).reshape(-1)
+
+        return [vertex_values, edge_values, square_values, cube_values]
+
+    
+    def forward(self, img_arr: torch.Tensor) -> torch.Tensor:
+        """
+        Calculates a discretization of the ECF of a 3D image.
+
+        Args:
+            img_arr (torch.Tensor): A 3D tensor with values between 0 and 1.
+
+        Returns:
+            ecf (torch.Tensor): A 1D tensor of shape (self.num_vals) containing the sublevel set ECF.
+        """
+
+        device = img_arr.device
+        n = self.num_vals
+        vertex_values, edge_values, square_values, cube_values = self.cell_values_3D(img_arr)
+
+        vertex_indices = torch.ceil(vertex_values * (n-1)).long()
+        edge_indices = torch.ceil(edge_values * (n-1)).long()
+        square_indices = torch.ceil(square_values * (n-1)).long()
+        cube_indices = torch.ceil(cube_values * (n-1)).long()
+
+        diff_ecf = torch.zeros(n, dtype=torch.int32, device=device)
+
+        # Add the contribution of the vertices
+        diff_ecf.scatter_add_(
+            0,
+            vertex_indices,
+            torch.ones_like(vertex_indices, dtype=torch.int32)
+        )
+
+        # Add the contribution of the edges
+        diff_ecf.scatter_add_(
+            0,
+            edge_indices,
+            -1 * torch.ones_like(edge_indices, dtype=torch.int32)
+        )
+
+        # Add the contribution of the squares
+        diff_ecf.scatter_add_(
+            0,
+            square_indices,
+            torch.ones_like(square_indices, dtype=torch.int32)
+        )
+
+        # Add the contribution of the cubes
+        diff_ecf.scatter_add_(
+            0,
+            cube_indices,
+            -1 * torch.ones_like(cube_indices, dtype=torch.int32)
+        )
+
+        return torch.cumsum(diff_ecf, dim=0)

pyect/preprocessing/image_processing.py

@@ -0,0 +1,242 @@
+from typing import Optional
+import torch
+import torchvision.transforms as transforms
+from pyect import Complex
+from PIL import Image
+
+
+def image_to_grayscale_tensor(image_path: str, device: torch.device) -> torch.Tensor:
+    # Open the image using PIL
+    image = Image.open(image_path)
+    # Convert the image to grayscale (mode 'L')
+    grayscale_image = image.convert("L")
+    # Convert the grayscale image to a tensor with values in [0,1]
+    tensor = transforms.ToTensor()(grayscale_image).squeeze(dim=0)
+    # The resulting tensor will have shape (H, W)
+    return tensor.to(device)
+
+
+def weighted_freudenthal(
+    img_arr: torch.Tensor, device: Optional[torch.device] = None
+) -> Complex:
+    """
+    Creates the weighted Freudenthal complex of an image array using a max function extension.
+    Discards edges and triangles that have a vertex with a zero weight.
+    By default, the device of the input tensor is used unless a different device is specified.
+
+    The vertices are a (h*w, 2) tensor with recentered pixel coordinates.
+    The vertex weights are a (h*w,) tensor containing the pixel intensities.
+    The edges are a (num_valid_edges, 2) tensor of vertex indices.
+    The edge weights are a (num_valid_edges,) tensor with the maximum weight on the edge.
+    The triangles are a (num_valid_triangles, 3) tensor of vertex indices.
+    The triangle weights are a (num_valid_triangles,) tensor with the maximum weight on the triangle.
+
+    Args:
+        img_arr (torch.Tensor): A grayscale image of shape (h, w).
+        device (torch.device, optional): The device to create tensors on.
+                If None, the device of the input tensor is used.
+
+    Returns:
+        Complex: A complex containing the weighted vertices, weighted edges, and weighted triangles.
+    """
+
+    device = img_arr.device if device is None else device
+    img_arr = img_arr.float().to(device)
+    h, w = img_arr.shape
+
+    # Create a mask of the nonzero pixels
+    img_mask = img_arr != 0 
+
+    # Indices of nonzero pixels (vertices)
+    nonzero_vertices = torch.nonzero(img_mask, as_tuple=True)
+
+    # Enumerate the nonzero vertices in the index array with all other values set to 0
+    vertex_numbers = torch.zeros_like(img_arr, dtype=torch.int64, device=device)
+    vertex_numbers[nonzero_vertices] = torch.arange(
+        nonzero_vertices[0].size(0), dtype=torch.int64, device=device
+    )
+
+    # Construct the vertex coords and weights
+    vertex_coords = torch.stack([
+        nonzero_vertices[1] - (w - 1) / 2.0,
+        (h - 1) / 2.0 - nonzero_vertices[0]
+    ], dim=1)
+    vertex_weights = img_arr[nonzero_vertices]
+    vertices = (vertex_coords, vertex_weights)
+
+    ### Horizontal Edges
+    # Remove the first and last columns of img_mask and check where the resulting arrays are both nonzero
+    horizontal_edge_mask = img_mask[:, :-1] & img_mask[:, 1:] 
+    horizontal_edge_indices = torch.nonzero(horizontal_edge_mask, as_tuple=True)
+
+    # Get the vertex numbers of the endpoints of each horizontal edge
+    horizontal_edge_vertices = torch.stack([
+        vertex_numbers[horizontal_edge_indices],
+        vertex_numbers[:, 1:][horizontal_edge_indices]
+    ], dim=1)
+    horizontal_edge_weights = vertex_weights[horizontal_edge_vertices].amax(dim=1)
+
+    ### Vertical Edges
+    # Remove the first and last rows of img_mask and check where the resulting arrays are both nonzero
+    vertical_edge_mask = img_mask[:-1, :] & img_mask[1:, :]
+    vertical_edge_indices = torch.nonzero(vertical_edge_mask, as_tuple=True)
+
+    # Get the vertex numbers of the endpoints of each vertical edge
+    vertical_edge_vertices = torch.stack([
+        vertex_numbers[vertical_edge_indices],
+        vertex_numbers[1:, :][vertical_edge_indices]
+    ], dim=1)
+    vertical_edge_weights = vertex_weights[vertical_edge_vertices].amax(dim=1)
+
+    ### Diagonal Edges
+    diagonal_edge_mask = img_mask[:-1, :-1] & img_mask[1:, 1:]
+    diagonal_edge_indices = torch.nonzero(diagonal_edge_mask, as_tuple=True)
+    diagonal_edge_vertices = torch.stack([
+        vertex_numbers[diagonal_edge_indices],
+        vertex_numbers[1:, 1:][diagonal_edge_indices]
+    ], dim=1)
+    diagonal_edge_weights = vertex_weights[diagonal_edge_vertices].amax(dim=1)
+
+    # Concatenate the horizontal, vertical, and diagonal edges
+    edge_vertices = torch.cat([
+        horizontal_edge_vertices,
+        vertical_edge_vertices,
+        diagonal_edge_vertices
+    ], dim=0)
+    edge_weights = torch.cat([
+        horizontal_edge_weights,
+        vertical_edge_weights,
+        diagonal_edge_weights
+    ], dim=0)
+    edges = (edge_vertices, edge_weights)
+
+    ### Upper Triangles
+    upper_triangle_mask = img_mask[:-1, :-1] & img_mask[:-1, 1:] & img_mask[1:, 1:]
+    upper_triangle_indices = torch.nonzero(upper_triangle_mask, as_tuple=True)
+    upper_triangle_vertices = torch.stack([
+        vertex_numbers[upper_triangle_indices],
+        vertex_numbers[:, 1:][upper_triangle_indices],
+        vertex_numbers[1:, 1:][upper_triangle_indices]
+    ], dim=1)
+    upper_triangle_weights = vertex_weights[upper_triangle_vertices].amax(dim=1)
+
+    ### Lower Triangles
+    lower_triangle_mask = img_mask[:-1, :-1] & img_mask[1:, :-1] & img_mask[1:, 1:]
+    lower_triangle_indices = torch.nonzero(lower_triangle_mask, as_tuple=True)
+    lower_triangle_vertices = torch.stack([
+        vertex_numbers[lower_triangle_indices],
+        vertex_numbers[1:, :][lower_triangle_indices],
+        vertex_numbers[1:, 1:][lower_triangle_indices]
+    ], dim=1)
+    lower_triangle_weights = vertex_weights[lower_triangle_vertices].amax(dim=1)
+
+    ### Concatenate the upper and lower triangles
+    triangle_vertices = torch.cat([
+        upper_triangle_vertices,
+        lower_triangle_vertices
+    ], dim=0)
+    triangle_weights = torch.cat([
+        upper_triangle_weights,
+        lower_triangle_weights
+    ], dim=0)
+    triangles = (triangle_vertices, triangle_weights)
+
+    return Complex(vertices, edges, triangles, device=device)
+
+
+def weighted_cubical(
+    img_arr: torch.Tensor, device: Optional[torch.device] = None
+) -> Complex:
+    """
+    Creates the weighted cubical complex of an image array.
+    Discards edges and squares that have a vertex with zero weight.
+
+    The vertices are a (h*w, 2) tensor with recentered pixel coordinates.
+    The vertex weights are a (h*w,) tensor containing the pixel intensities.
+    The edges are a (num_valid_edges, 2) tensor of vertex indices.
+    The edge weights are a (num_valid_edges,) tensor with the maximum weight on the edge.
+    The squares are a (num_valid_squares, 4) tensor of vertex indices.
+    The square weights are a (num_valid_squares,) tensor with the maximum weight on
+    the square.
+
+    Args:
+        img_arr (torch.Tensor): A grayscale image of shape (h, w).
+        device (torch.device, optional): The device to create tensors on.
+                If None, the device of the input tensor is used.
+
+    Returns:
+        Complex: A complex containing the weighted vertices, weighted edges, and weighted squares.
+    """
+
+    device = img_arr.device if device is None else device
+    img_arr = img_arr.float().to(device)
+    h, w = img_arr.shape
+
+    # Create a mask of the nonzero pixels
+    img_mask = img_arr != 0
+
+    # Indices of nonzero pixels (vertices)
+    nonzero_vertices = torch.nonzero(img_mask, as_tuple=True)
+
+    # Create an array enumerating the nonzero vertices with all other values 0
+    vertex_numbers = torch.zeros_like(img_arr, dtype=torch.int64, device=device)
+    vertex_numbers[nonzero_vertices] = torch.arange(
+        nonzero_vertices[0].size(0), dtype=torch.int64, device=device
+    )
+
+    # Construct the vertex coords and weights
+    vertex_coords = torch.stack([
+        nonzero_vertices[1] - (w - 1) / 2.0,
+        (h - 1) / 2.0 - nonzero_vertices[0]
+    ], dim=1)
+    vertex_weights = img_arr[nonzero_vertices]
+    vertices = (vertex_coords, vertex_weights)
+
+    ### Horizontal Edges
+    # Remove the first and last columns of img_mask and check where the resulting arrays are both nonzero
+    horizontal_edge_mask = img_mask[:, :-1] & img_mask[:, 1:]
+    horizontal_edge_indices = torch.nonzero(horizontal_edge_mask, as_tuple=True)
+
+    # Get the vertex numbers of the endpoints of each horizontal edge
+    horizontal_edge_vertices = torch.stack([
+        vertex_numbers[horizontal_edge_indices],
+        vertex_numbers[:, 1:][horizontal_edge_indices]
+    ], dim=1)
+    horizontal_edge_weights = vertex_weights[horizontal_edge_vertices].amax(dim=1)
+
+    ### Vertical Edges
+    # Remove the first and last rows of img_mask and check where the resulting arrays are both nonzero
+    vertical_edge_mask = img_mask[:-1, :] & img_mask[1:, :]
+    vertical_edge_indices = torch.nonzero(vertical_edge_mask, as_tuple=True)
+
+    # Get the vertex numbers of the endpoints of each vertical edge
+    vertical_edge_vertices = torch.stack([
+        vertex_numbers[vertical_edge_indices],
+        vertex_numbers[1:, :][vertical_edge_indices]
+    ], dim=1)
+    vertical_edge_weights = vertex_weights[vertical_edge_vertices].amax(dim=1)
+
+    # Concatenate the horizontal and vertical edges
+    edge_vertices = torch.cat([
+        horizontal_edge_vertices,
+        vertical_edge_vertices
+    ], dim=0)
+    edge_weights = torch.cat([
+        horizontal_edge_weights,
+        vertical_edge_weights
+    ], dim=0)
+    edges = (edge_vertices, edge_weights)
+
+    ###Squares
+    square_mask = horizontal_edge_mask[:-1, :] & horizontal_edge_mask[1:, :]
+    square_indices = torch.nonzero(square_mask, as_tuple=True)
+    square_vertices = torch.stack([
+        vertex_numbers[square_indices],
+        vertex_numbers[1:, :][square_indices],
+        vertex_numbers[:, 1:][square_indices],
+        vertex_numbers[1:, 1:][square_indices]
+    ], dim=1)
+    square_weights = vertex_weights[square_vertices].amax(dim=1)
+    squares = (square_vertices, square_weights)
+
+    return Complex(vertices, edges, squares, n_type="cubical", device=device)

pyect/tensor_complex.py

@@ -0,0 +1,204 @@
+"""Tools for working with simplicial complexes.
+
+The Complex class is a collection of simplices, each of which is represented by a
+tensor of coordinates and a tensor of weights.
+"""
+
+from typing import Tuple, Optional
+
+import torch
+import warnings
+import numpy.typing as npt
+
+from .dtypes import COORDS_DTYPE, INDICES_DTYPE, WEIGHTS_DTYPE
+
+
+class Complex:
+    """A simplicial complex of arbitrary dimension.
+
+    The representation is as a collection of simplices (or cubical cells) using tensors.
+    """
+
+    def __init__(
+        self,
+        *args: Tuple[torch.Tensor, torch.Tensor],
+        vertex_dtype: torch.dtype = COORDS_DTYPE,
+        index_dtype: torch.dtype = INDICES_DTYPE,
+        weights_dtype: torch.dtype = WEIGHTS_DTYPE,
+        device: Optional[torch.device] = None,
+        n_type: str = "simplicial",
+    ) -> None:
+        """Initializes a complex.
+
+        All tensors are cast to the given types.
+
+        Args:
+            *args: A variable number of tuples, each containing the simplices of a given
+                dimension. Each tuple should contain two tensors.
+                The first tensor contains the coordinates of the simplices
+                The second tensor contains the weights of the simplices.
+
+                The first tuple should contain the vertices of the complex, and
+                therefore must be a tensor of shape [num_vertices, d].
+
+                Any following tuples should contain indices into the vertices tensor,
+                and therefore must be a tensor of shape [num_simplices, k], where k+1 is the
+                dimension of the simplex.
+
+            vertex_dtype: The data type to use for the vertex coordinates.
+            index_dtype: The data type to use for the simplex indices.
+            weights_dtype: The data type to use for the simplex weights.
+            device: The device to use for the tensors.
+            n_type: The type of complex. Currently only "simplicial" and "cubical"
+                are supported.
+        """
+        # Verify the dimensions of the simplices, and raise a UserError if
+        # there is a mismatch.
+        self._validate_dimensions(*args, n_type=n_type)
+
+        # Call .to on each tensor to cast to the given type and device.
+        types = [vertex_dtype] + [index_dtype] * (len(args) - 1)
+        self.dimensions = tuple(
+            (
+                (
+                    coords.to(dtype=types[dim], device=device),
+                    weights.to(dtype=weights_dtype, device=device),
+                )
+            )
+            for dim, (coords, weights) in enumerate(args)
+        )
+        self.n_type = n_type
+
+    @staticmethod
+    def from_numpy(
+        *args: Tuple[npt.NDArray, npt.NDArray],
+        vertex_dtype: torch.dtype = COORDS_DTYPE,
+        index_dtype: torch.dtype = INDICES_DTYPE,
+        weights_dtype: torch.dtype = WEIGHTS_DTYPE,
+        device: Optional[torch.device] = None,
+        n_type: str = "simplicial",
+    ) -> "Complex":
+        """Initializes a simplicial complex from numpy arrays.
+
+        Args:
+            *args: A variable number of tuples, each containing the simplices of a given
+                dimension. Each tuple should contain two numpy arrays.
+                The first array contains the coordinates of the simplices
+                The second array contains the weights of the simplices.
+
+                The first tuple should contain the vertices of the complex, and
+                therefore must be a tensor of shape [num_vertices, d].
+
+                Any following tuples should contain indices into the vertices tensor,
+                and therefore must be a tensor of shape [num_simplices, k], where k+1 is the
+                dimension of the simplex.
+
+            vertex_dtype: The data type to use for the vertex coordinates.
+            index_dtype: The data type to use for the simplex indices.
+            weights_dtype: The data type to use for the simplex weights.
+            device:
+                The device to use for the tensors.
+            n_type: The type of the simplicial complex. Currently only "simplicial" and "cubical"
+                are supported.
+
+        """
+        if device is None:
+            device = (
+                torch.device("cuda")
+                if torch.cuda.is_available()
+                else torch.device("cpu")
+            )
+
+        typematch = [vertex_dtype] + [index_dtype] * (len(args) - 1)
+        dimensions = tuple(
+            (
+                torch.as_tensor(coords, device=device, dtype=typematch[i]),
+                torch.as_tensor(weights, device=device, dtype=weights_dtype),
+            )
+            for i, (coords, weights) in enumerate(args)
+        )
+        return Complex(*dimensions, device=device, n_type=n_type)
+
+    def to(self, device: torch.device) -> "Complex":
+        """Moves the complex to the given device."""
+        return Complex(*self.dimensions, device=device, n_type=self.n_type)
+
+    def __getitem__(self, dim: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Returns the simplices of the given dimension."""
+        return self.dimensions[dim]
+
+    def get_coords(self, dim: int) -> torch.Tensor:
+        """Returns the coordinates of the simplices of the given dimension."""
+        return self.dimensions[dim][0]
+
+    def get_weights(self, dim: int) -> torch.Tensor:
+        """Returns the weights of the simplices of the given dimension."""
+        return self.dimensions[dim][1]
+
+    def top_dim(self) -> int:
+        """Returns the top dimension of the complex."""
+        return len(self) - 1
+
+    def __len__(self) -> int:
+        """Returns the number of dimensions in the complex."""
+        return len(self.dimensions)
+
+    def space_dim(self) -> int:
+        """Returns the dimension of the space the complex is embedded in."""
+        return self.dimensions[0][0].shape[1]
+    
+    def center_(self) -> "Complex":
+        """
+        Re-center the complex in-place so that the average vertex coordinate is at the origin.
+        """
+        if len(self.dimensions) == 0:
+            return self
+
+        v_coords, v_weights = self.dimensions[0]
+        if v_coords.numel() == 0:
+            return self
+
+        center = v_coords.mean(dim=0)
+        new_v_coords = (v_coords - center).contiguous()
+
+        dims: list[Tuple[torch.Tensor, torch.Tensor]] = list(self.dimensions)
+        dims[0] = (new_v_coords, v_weights)
+        self.dimensions = tuple(dims)
+        return self
+
+    @staticmethod
+    def _validate_dimensions(
+        *args: Tuple[torch.Tensor, torch.Tensor], n_type: str
+    ) -> None:
+        for dim, simplex_list in enumerate(args):
+            if simplex_list[0].dim() != 2:
+                raise ValueError(
+                    f"Dimension {dim} simplices must be a 2d tensor."
+                    + f" Got {simplex_list[0].dim()} dimensions."
+                )
+            if simplex_list[1].dim() != 1:
+                raise ValueError(
+                    f"Dimension {dim} weights must be a 1d tensor."
+                    + f" Got {simplex_list[1].dim()} dimensions."
+                )
+            if simplex_list[0].shape[0] != simplex_list[1].shape[0]:
+                raise ValueError(
+                    f"Dimension {dim} coordinates and weights must have the same number of simplices."
+                    + f" Got {simplex_list[0].shape[0]} simplices and {simplex_list[1].shape[0]} weights."
+                )
+
+            if dim > 0:  # simplices, k > 0
+                if n_type == "simplicial":
+                    if simplex_list[0].shape[1] != dim + 1:
+                        raise ValueError(
+                            f"Dimension {dim} simplices must have {dim + 1} columns."
+                            + f" Got {simplex_list[0].shape[1]} columns."
+                        )
+                elif n_type == "cubical":
+                    if simplex_list[0].shape[1] != 2 ** dim:
+                        raise ValueError(
+                            f"Dimension {dim} simplices must have {2 ** dim} columns."
+                            + f" Got {simplex_list[0].shape[1]} columns."
+                        )
+                else:  # warn that validation not implementod for n_type, but no error
+                    warnings.warn(f"Validation not implemented for n_type {n_type}. Proceed with caution.")

pyect/wecfs.py

@@ -0,0 +1,65 @@
+""" For computing the WECFs of lower-star filtrations of
+weighted simplicial/cubical complex with respect to a set of filter functions."""
+
+import torch
+from typing import List, Tuple
+
+def compute_wecfs(
+    complex_data: List[Tuple[torch.Tensor, torch.Tensor]],
+    num_vals: int
+) -> torch.Tensor:
+    """Calculates a discretization of the WECFs of a weighted complex with respect to a set of filter functions.
+
+    Args:
+        complex_data: A weighted simplicial or cubical complex with a collection of filter functions,
+        represented as a list of pairs of tensors.
+            complex_data[0] = (filters, v_weights):
+                filters (torch.Tensor): A tensor of shape (k_0, m) where k_0 is the
+                    number of vertices and m is the number of filter functions.
+                    Each column contains the values of a filter function on the vertices.
+
+                v_weights (torch.Tensor): A tensor of shape (k_0). Values are the weights of the vertices.
+
+            for i > 0:
+                complex_data[i] = (simp_verts, simp_weights):
+                    simp_verts (torch.Tensor): A tensor of shape (k_i, i+1) where k_i is the number of i-simplices.
+                        Rows are the vertex sets of the i-simplices.
+
+                    simp_weights (torch.Tensor): A tensor of shape (k_i). Values are the weights of the i-simplices.
+
+    Returns:
+        wecfs (torch.Tensor): A 2d tensor of shape (m, num_vals)
+            containing the WECFs.
+    """
+
+    filters = complex_data[0][0].float()
+    m = filters.size(dim=1)
+    device = filters.device
+    v_weights = complex_data[0][1].to(device=device, dtype=torch.float32)
+
+    expanded_v_weights = v_weights.unsqueeze(0).expand(m, -1)  # Expand to shape (m, k_0)
+
+    # Map the values of the filter functions to indices in range(num_vals)
+    max_val = filters.abs().amax()
+    v_indices = torch.ceil(
+        (num_vals - 1) * (max_val + filters) / (2.0 * max_val)
+    ).clamp(0, num_vals-1).long()
+
+    # Initialize the differentiated WECFs
+    diff_wecfs = torch.zeros((m, num_vals), dtype=torch.float32, device=device)
+
+    # Add the contribution of the vertices to the differentiated WECFs
+    diff_wecfs.scatter_add_(1, v_indices.T, expanded_v_weights)
+
+    for i in range(1, len(complex_data)):
+        simp_verts = complex_data[i][0].to(device=device, dtype=torch.long)
+        simp_weights = complex_data[i][1].to(device=device, dtype=torch.float32)
+
+        expanded_simp_weights = (-1) ** i * simp_weights.unsqueeze(0).expand(m, -1)
+
+        simp_indices = v_indices[simp_verts]
+        max_simp_indices = torch.amax(simp_indices, dim=1)
+
+        diff_wecfs.scatter_add_(1, max_simp_indices.T, expanded_simp_weights)
+
+    return torch.cumsum(diff_wecfs, dim=1)

pyect/wect.py

@@ -0,0 +1,130 @@
+"""For computing the WECT of a weighted geometric simplicial/cubical complex embedded in R^n."""
+
+import torch
+from typing import List, Tuple
+
+
+class WECT(torch.nn.Module):
+    """A torch module for computing the Weighted Euler Characteristic Transform (WECT) of a simplicial complex discretized over a grid.
+
+    This module may be used just for computing the WECT, or used as a layer in a neural network.
+    Internally, the module stores the directions and number of heights used for sampling, so repeated forward calls
+    do not require these parameters to be passed in, and allow streamlined loading/saving of the module for consistent
+    computation.
+
+    This module can also be converted to TorchScript using torch.jit.script for use
+    outside of Python.
+    """
+
+    def __init__(self, dirs: torch.Tensor, num_heights: int) -> None:
+        """Initializes the WECT module.
+
+        The initialized module is designed to compute the WECT of a simplicial complex
+        embedded in R^[dirs.shape[1]], using dirs.shape[0] directions for sampling.
+        The discretization of the WECT is parameterized by num_heights distinct height values.
+
+        Args:
+            dirs: An (d x n) tensor of directions to use for sampling.
+            num_heights: A constant tensor, with the number of distinct height
+                values to round to as an integer
+        """
+        super().__init__()
+        dirs = torch.nn.functional.normalize(dirs, p=2, dim=1, eps=1e-12)
+        self.register_buffer("dirs", dirs)
+        self.num_heights: int = int(num_heights)
+
+    def _vertex_indices(
+        self,
+        vertex_coords: torch.Tensor,
+    ) -> torch.Tensor:
+        """Calculates the height values of each vertex and converts them to an index in range(num_heights).
+
+        Args:
+            vertex_coords (torch.Tensor): A tensor of shape (k_0, n) with rows representing the coordinates of the vertices.
+
+        Returns:
+            torch.Tensor: A tensor of shape (k_0, d) with the height indices of each vertex in each direction.
+        """
+
+        v_norms = torch.norm(vertex_coords, dim=1)
+        max_height = torch.amax(v_norms)
+        v_heights = torch.matmul(vertex_coords, self.dirs.T)
+
+        # The case where all vertices are at the origin
+        if max_height.item() == 0.0:
+            return torch.zeros((v_heights.size(0), self.dirs.size(0)), dtype=torch.long, device=self.dirs.device)
+
+        v_indices = torch.ceil(
+            (self.num_heights - 1) * (max_height + v_heights) / (2.0 * max_height)
+        ).clamp(0, self.num_heights - 1).long()
+
+        return v_indices
+
+    def forward(
+        self,
+        complex_data: List[Tuple[torch.Tensor, torch.Tensor]],
+    ) -> torch.Tensor:
+        """Calculates a discretization of the WECT of a complex embedded in n-dimensional space.
+
+        Args:
+            complex_data: A weighted simplicial or cubical complex, represented as a list of pairs of tensors.
+                complex_data[0] = (v_coords, v_weights):
+                    v_coords (torch.Tensor): A tensor of shape (k_0, n) where k_0 is the number of vertices.
+                    Rows are the coordinates of the vertices.
+
+                    v_weights (torch.Tensor): A tensor of shape (k_0). Values are the weights of the vertices.
+
+                for i > 0:
+                    complex_data[i] = (simp_verts, simp_weights):
+                        simp_verts (torch.Tensor): A tensor of shape (k_i, i+1) where k_i is the number of i-simplices.
+                        Rows are the vertex sets of the i-simplices.
+
+                        simp_weights (torch.Tensor): A tensor of shape (k_i). Values are the weights of the i-simplices.
+
+        Returns:
+            wect (torch.Tensor): A 2d tensor of shape (self.dirs.shape[0], self.num_heights)
+                containing the WECT.
+        """
+
+        d = self.dirs.size(dim=0)
+        h = self.num_heights
+
+        if h <= 0:
+            raise ValueError("num_heights must be positive.")
+        
+        device = self.dirs.device
+        v_coords  = complex_data[0][0].to(device=device, dtype=torch.float32)
+        v_weights = complex_data[0][1].to(device=device, dtype=torch.float32)
+
+        # Check for empty inputs
+        if v_coords.size(0) == 0:
+            return torch.zeros((d, h), dtype=torch.float32, device=device)
+
+        expanded_v_weights = v_weights.unsqueeze(0).expand(
+            d, -1
+        )  # Expand to shape (d, k_0)
+
+        # Initialize the differentiated WECT
+        diff_wect = torch.zeros((d, h), dtype=torch.float32, device=device)
+
+        # Compute the height index of each vertex
+        v_indices = self._vertex_indices(v_coords)
+
+        # Add the contribution of the vertices to the differentiated WECT
+        diff_wect.scatter_add_(1, v_indices.T, expanded_v_weights)
+
+        for i in range(1, len(complex_data)):
+            simp_verts = complex_data[i][0].to(device=device, dtype=torch.long)
+            simp_weights = complex_data[i][1].to(device=device, dtype=torch.float32)
+
+            # Expand to shape (d, k_i)
+            expanded_simp_weights = (-1) ** i * simp_weights.unsqueeze(0).expand(d, -1)
+
+            # Compute the maximum index for each simplex's vertices
+            simp_indices = v_indices[simp_verts]
+            max_simp_indices = torch.amax(simp_indices, dim=1)
+
+            # Add the contribution of the i-simplices to the differentiated WECT
+            diff_wect.scatter_add_(1, max_simp_indices.T, expanded_simp_weights)
+
+        return torch.cumsum(diff_wect, dim=1)

pyect-0.1.0.dist-info/METADATA

@@ -0,0 +1,80 @@
+Metadata-Version: 2.4
+Name: pyect
+Version: 0.1.0
+Summary: Generalized computation of the WECT using PyTorch.
+Home-page: https://github.com/compTAG/pyECT
+Author: Alex McCleary, Eli Quist, Jack Ruder, Jacob Sriraman
+Author-email: eli.quist@student.montana.edu
+License: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: Pillow
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# pyECT
+
+The Weighted Euler Characteristic Transform (WECT) is a mathematical tool
+used to analyze and summarize geometric and topological features of data.
+This package provides an efficient and simple implementation of the WECT using
+PyTorch.
+
+This codebase accompanies the following paper (and should be cited if you use
+this package):
+
+```
+TODO: Add Citation
+```
+
+## Installation
+
+To install `pyECT`, use pip:
+
+```bash
+pip install pyect 
+```
+
+## Usage
+
+Here's a simple example of how to use `pyECT`:
+
+```python
+from pyect import WECT
+
+# Example data and weight function
+data = [...]  # Replace with your data
+weight_function = lambda x: x**2  # Replace with your weight function
+
+# Compute the WECT
+wect = WECT(data, weight_function)
+result = wect.compute()
+
+print("WECT result:", result)
+```
+
+For more detailed examples, please see the `/examples` directory.
+
+## Contributing
+
+Contributions are welcome! If you'd like to contribute, please fork the
+repository and submit a pull request. For major changes, please open an issue
+first to discuss what you'd like to change.
+
+## License
+
+This project is licensed under the MIT License. See the [LICENSE](LICENSE)
+file for details.

pyect-0.1.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+pyect/__init__.py,sha256=roSWMLV7yFltMExesaIqRyEZANwYaKuzBJOPlDuQa7E,302
+pyect/directions.py,sha256=XmCyfc6GrRHmgeeeRVVTCJs2QvE4li0pU_pc98TDek0,890
+pyect/dtypes.py,sha256=_GxQ8cm3UkRo-CU3v--4D_DFd0NhTIZl_u0yuIADMIY,144
+pyect/image_ecf.py,sha256=2C_ga6sgMOO8XAua8tjbsFXTelmwj5arIN_2sCVwgc4,8175
+pyect/tensor_complex.py,sha256=R_lgSOmTkK5PV8Q5BZF7MO73PU6ROGDew5mWwftDcE0,8345
+pyect/wecfs.py,sha256=C__WLyuELs2wmhLIK2aZ-suNffotUmWyL_Q1S7_SYJA,2771
+pyect/wect.py,sha256=g1wXsnkUx3m7NN9n3Tu6eI7hFWxJC9w_1htO0p98lAg,5560
+pyect/preprocessing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pyect/preprocessing/image_processing.py,sha256=hMWoZgytBruwdM6Op4HUWlbqoYJKsT2rs_hp7U7lULM,9967
+pyect-0.1.0.dist-info/licenses/LICENSE,sha256=pk1cMbtWYROTtwmnpV4fWVheiH-vqVq_kzP4mtkCywY,1144
+pyect-0.1.0.dist-info/METADATA,sha256=Dhtj726KbcfZTCY59OJlBpehDkjF4CSxeMt4aR6pcnI,1964
+pyect-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pyect-0.1.0.dist-info/top_level.txt,sha256=0kDTwnLAArCjbb33d93Hg0IvoHhMKV9Fv9P-rQ-p6FY,6
+pyect-0.1.0.dist-info/RECORD,,

pyect-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

pyect-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,23 @@
+MIT License
+
+Copyright (c) 2025 Montana State University Computational Topology and
+Geometry (CompTaG) Research Group.
+
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

pyect-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+pyect