Spaces

ConstantSpace dataclass

Bases: SpaceType

Represents a constant value repeated multiple times.

This class generates a tensor containing the same constant value repeated num times.

Attributes:

Name	Type	Description
value	float	The constant float value to be used.
num	int | None	The number of times the constant value should be repeated. Defaults to 1.

Source code in src/einmesh/spaces.py

@dataclass
class ConstantSpace(SpaceType):
    """
    Represents a constant value repeated multiple times.

    This class generates a tensor containing the same constant value repeated
    `num` times.

    Attributes:
        value: The constant float value to be used.
        num: The number of times the constant value should be repeated. Defaults to 1.
    """

    value: float
    num: int | None = None

    def __post_init__(self) -> None:
        if self.num is None:
            object.__setattr__(self, "num", 1)

    def _generate_samples(self, backend: AbstractBackend):
        """Generates a tensor with the constant value repeated."""
        return backend.full((self.num,), self.value)

_generate_samples(backend)

Generates a tensor with the constant value repeated.

Source code in src/einmesh/spaces.py

def _generate_samples(self, backend: AbstractBackend):
    """Generates a tensor with the constant value repeated."""
    return backend.full((self.num,), self.value)

Distribution dataclass

Bases: SpaceType

Base class for all distribution spaces.

Source code in src/einmesh/spaces.py

@dataclass
class Distribution(SpaceType):
    """Base class for all distribution spaces."""

    ...

LgSpace dataclass

Bases: LogSpace

Represents a sequence of points spaced in base 2.

Source code in src/einmesh/spaces.py

@dataclass
class LgSpace(LogSpace):
    """
    Represents a sequence of points spaced in base 2.
    """

    start: float
    end: float
    num: int
    base: float = field(init=False, default=2)

    def __post_init__(self) -> None:
        object.__setattr__(self, "base", 2)

LinSpace dataclass

Bases: SpaceType

Represents a sequence of points spaced linearly.

This class generates a tensor of num points evenly spaced between start and end (inclusive).

Attributes:

Name	Type	Description
start	float	The starting value of the sequence.
end	float	The ending value of the sequence.
num	int	The number of points to generate.

Source code in src/einmesh/spaces.py

@dataclass
class LinSpace(SpaceType):
    """
    Represents a sequence of points spaced linearly.

    This class generates a tensor of `num` points evenly spaced between `start`
    and `end` (inclusive).

    Attributes:
        start: The starting value of the sequence.
        end: The ending value of the sequence.
        num: The number of points to generate.
    """

    start: float
    end: float
    num: int

    def _generate_samples(self, backend: AbstractBackend):
        """Generates the linearly spaced points."""
        return backend.linspace(self.start, self.end, self.num)

_generate_samples(backend)

Generates the linearly spaced points.

Source code in src/einmesh/spaces.py

def _generate_samples(self, backend: AbstractBackend):
    """Generates the linearly spaced points."""
    return backend.linspace(self.start, self.end, self.num)

ListSpace dataclass

Bases: SpaceType

Represents a predefined list of values.

This class generates a tensor directly from a provided list of float values. The number of points generated is equal to the length of the input list.

Attributes:

Name	Type	Description
values	list[float]	A list of float values to be converted into a tensor.
num	int	The number of points, automatically set to length of values.

Source code in src/einmesh/spaces.py

@dataclass
class ListSpace(SpaceType):
    """
    Represents a predefined list of values.

    This class generates a tensor directly from a provided list of float values.
    The number of points generated is equal to the length of the input list.

    Attributes:
        values: A list of float values to be converted into a tensor.
        num: The number of points, automatically set to length of values.
    """

    values: list[float]
    num: int = field(init=False)

    def __post_init__(self) -> None:
        object.__setattr__(self, "num", len(self.values))

    def _generate_samples(self, backend: AbstractBackend):
        """Generates a tensor from the provided list of values."""
        return backend.tensor(self.values)

_generate_samples(backend)

Generates a tensor from the provided list of values.

Source code in src/einmesh/spaces.py

def _generate_samples(self, backend: AbstractBackend):
    """Generates a tensor from the provided list of values."""
    return backend.tensor(self.values)

LnSpace dataclass

Bases: LogSpace

Represents a sequence of points spaced in base e.

Source code in src/einmesh/spaces.py

@dataclass
class LnSpace(LogSpace):
    """
    Represents a sequence of points spaced in base e.
    """

    start: float
    end: float
    num: int
    base: float = field(init=False, default=math.e)

    def __post_init__(self) -> None:
        object.__setattr__(self, "base", math.e)

LogSpace dataclass

Bases: SpaceType

Represents a sequence of points spaced logarithmically.

This class generates a tensor of num points between 10**start and 10**end (or base**start and base**end if base is specified), spaced logarithmically.

Attributes:

Name	Type	Description
start	float	The starting exponent of the sequence.
end	float	The ending exponent of the sequence.
num	int	The number of points to generate.
base	float	The base of the logarithm. Defaults to 10.

Source code in src/einmesh/spaces.py

@dataclass
class LogSpace(SpaceType):
    """
    Represents a sequence of points spaced logarithmically.

    This class generates a tensor of `num` points between `10**start` and `10**end`
    (or `base**start` and `base**end` if `base` is specified), spaced
    logarithmically.

    Attributes:
        start: The starting exponent of the sequence.
        end: The ending exponent of the sequence.
        num: The number of points to generate.
        base: The base of the logarithm. Defaults to 10.
    """

    start: float
    end: float
    num: int
    base: float = 10

    def _generate_samples(self, backend: AbstractBackend):
        """Generates the logarithmically spaced points."""
        sample = backend.logspace(self.start, self.end, self.num, base=self.base)
        return self._apply_operators(sample, backend)

_generate_samples(backend)

Generates the logarithmically spaced points.

Source code in src/einmesh/spaces.py

def _generate_samples(self, backend: AbstractBackend):
    """Generates the logarithmically spaced points."""
    sample = backend.logspace(self.start, self.end, self.num, base=self.base)
    return self._apply_operators(sample, backend)

NormalDistribution dataclass

Bases: Distribution

Represents a sampling from a normal (Gaussian) distribution.

This class generates a tensor of num random numbers sampled from a normal distribution with the specified mean and standard deviation std.

Attributes:

Name	Type	Description
mean	float	The mean (center) of the normal distribution.
std	float	The standard deviation (spread or width) of the normal distribution.
num	int	The number of samples to generate.

Source code in src/einmesh/spaces.py

@dataclass
class NormalDistribution(Distribution):
    """
    Represents a sampling from a normal (Gaussian) distribution.

    This class generates a tensor of `num` random numbers sampled from a normal
    distribution with the specified `mean` and standard deviation `std`.

    Attributes:
        mean: The mean (center) of the normal distribution.
        std: The standard deviation (spread or width) of the normal distribution.
        num: The number of samples to generate.
    """

    mean: float
    std: float
    num: int

    def _generate_samples(self, backend: AbstractBackend):
        """Generates samples from the normal distribution."""
        return backend.normal(mean=self.mean, std=self.std, size=(self.num,))

_generate_samples(backend)

Generates samples from the normal distribution.

Source code in src/einmesh/spaces.py

def _generate_samples(self, backend: AbstractBackend):
    """Generates samples from the normal distribution."""
    return backend.normal(mean=self.mean, std=self.std, size=(self.num,))

RangeSpace dataclass

Bases: SpaceType

Return evenly spaced values within a given interval.

This class generates a tensor of num points evenly spaced between start and end (inclusive). It works like a arange function of numpy, torch or jax.

Attributes:

Name	Type	Description
start	float	The starting value of the sequence.
end	float	The ending value of the sequence.
num	float	The number of points to generate.

Source code in src/einmesh/spaces.py

@dataclass
class RangeSpace(SpaceType):
    """
    Return evenly spaced values within a given interval.

    This class generates a tensor of `num` points evenly spaced between `start`
    and `end` (inclusive). It works like a arange function of numpy, torch or jax.

    Attributes:
        start: The starting value of the sequence.
        end: The ending value of the sequence.
        num: The number of points to generate.
    """

    start: float
    stop: float
    step: int = 1

    def __post_init__(self) -> None:
        self.num: int = int((self.stop - self.start) / self.step)

    def _generate_samples(self, backend: AbstractBackend) -> Any:
        """Generates the linearly spaced points using arange."""
        return backend.arange(self.start, self.stop, self.step)

_generate_samples(backend)

Generates the linearly spaced points using arange.

Source code in src/einmesh/spaces.py

def _generate_samples(self, backend: AbstractBackend) -> Any:
    """Generates the linearly spaced points using arange."""
    return backend.arange(self.start, self.stop, self.step)

SpaceType dataclass

Bases: ABC

Base class for all space types.

Source code in src/einmesh/spaces.py

@dataclass
class SpaceType(ABC):
    """Base class for all space types."""

    operators: list[BackendOperator] = field(init=False, default_factory=list)

    def _with_operator(self, operator: BackendOperator, prepend: bool = False) -> SpaceType:
        """Return a copy of this space with the given operator applied."""
        new_space = deepcopy(self)
        if prepend:
            new_space.operators.insert(0, operator)
        else:
            new_space.operators.append(operator)
        return new_space

    def _with_operators(self, operators: list[BackendOperator], prepend: bool = False) -> SpaceType:
        """Return a copy of this space with multiple operators applied."""
        new_space = deepcopy(self)
        if prepend:
            new_space.operators = operators + new_space.operators
        else:
            new_space.operators.extend(operators)
        return new_space

    def __abs__(self) -> SpaceType:
        return self._with_operator(OperatorFactory.abs())

    def __add__(self, other) -> SpaceType:
        return self._with_operator(OperatorFactory.add(value=other))

    def __radd__(self, other) -> SpaceType:
        return self.__add__(other)

    def __sub__(self, other) -> SpaceType:
        return self._with_operator(OperatorFactory.sub(value=other))

    def __rsub__(self, other) -> SpaceType:
        return self._with_operators(
            operators=[
                OperatorFactory.neg(),
                OperatorFactory.add(value=other),
            ],
            prepend=True,
        )

    def __mul__(self, other) -> SpaceType:
        return self._with_operator(OperatorFactory.mul(value=other))

    def __rmul__(self, other) -> SpaceType:
        return self.__mul__(other)

    def __mod__(self, other) -> SpaceType:
        return self._with_operator(OperatorFactory.mod(value=other))

    def __floordiv__(self, other) -> SpaceType:
        return self._with_operator(OperatorFactory.floor_div(value=other))

    def __truediv__(self, other) -> SpaceType:
        return self._with_operator(OperatorFactory.div(value=other))

    def __neg__(self) -> SpaceType:
        return self._with_operator(OperatorFactory.neg())

    def __pos__(self) -> SpaceType:
        return self._with_operator(OperatorFactory.pos())

    def __pow__(self, other) -> SpaceType:
        return self._with_operator(OperatorFactory.pow(exponent=other))

    @abstractmethod
    def _generate_samples(self, backend: AbstractBackend):
        """Create a new space type."""
        ...

    def _sample(self, backend: AbstractBackend) -> Any:
        """Sample the space type."""
        sample = self._generate_samples(backend)
        return self._apply_operators(sample, backend)

    def _apply_operators(self, sample, backend: AbstractBackend) -> Any:
        """Apply the operators to the space type."""
        for operator in self.operators:
            sample = operator._apply(x=sample, backend=backend)
        return sample

_apply_operators(sample, backend)

Apply the operators to the space type.

Source code in src/einmesh/spaces.py

def _apply_operators(self, sample, backend: AbstractBackend) -> Any:
    """Apply the operators to the space type."""
    for operator in self.operators:
        sample = operator._apply(x=sample, backend=backend)
    return sample

_generate_samples(backend) abstractmethod

Create a new space type.

Source code in src/einmesh/spaces.py

@abstractmethod
def _generate_samples(self, backend: AbstractBackend):
    """Create a new space type."""
    ...

_sample(backend)

Sample the space type.

Source code in src/einmesh/spaces.py

def _sample(self, backend: AbstractBackend) -> Any:
    """Sample the space type."""
    sample = self._generate_samples(backend)
    return self._apply_operators(sample, backend)

_with_operator(operator, prepend=False)

Return a copy of this space with the given operator applied.

Source code in src/einmesh/spaces.py

def _with_operator(self, operator: BackendOperator, prepend: bool = False) -> SpaceType:
    """Return a copy of this space with the given operator applied."""
    new_space = deepcopy(self)
    if prepend:
        new_space.operators.insert(0, operator)
    else:
        new_space.operators.append(operator)
    return new_space

_with_operators(operators, prepend=False)

Return a copy of this space with multiple operators applied.

Source code in src/einmesh/spaces.py

def _with_operators(self, operators: list[BackendOperator], prepend: bool = False) -> SpaceType:
    """Return a copy of this space with multiple operators applied."""
    new_space = deepcopy(self)
    if prepend:
        new_space.operators = operators + new_space.operators
    else:
        new_space.operators.extend(operators)
    return new_space

UniformDistribution dataclass

Bases: Distribution

Represents a sampling from a uniform distribution.

This class generates a tensor of num random numbers sampled from a uniform distribution over the interval [low, high).

Attributes:

Name	Type	Description
low	float	The lower boundary of the output interval.
high	float	The upper boundary of the output interval.
num	int	The number of samples to generate.

Source code in src/einmesh/spaces.py

@dataclass
class UniformDistribution(Distribution):
    """
    Represents a sampling from a uniform distribution.

    This class generates a tensor of `num` random numbers sampled from a uniform
    distribution over the interval [`low`, `high`).

    Attributes:
        low: The lower boundary of the output interval.
        high: The upper boundary of the output interval.
        num: The number of samples to generate.
    """

    low: float
    high: float
    num: int

    def _generate_samples(self, backend: AbstractBackend):
        """Generates samples from the uniform distribution."""
        # torch.rand samples from [0, 1), so we scale and shift.
        return backend.rand(size=(self.num,)) * (self.high - self.low) + self.low

_generate_samples(backend)

Generates samples from the uniform distribution.

Source code in src/einmesh/spaces.py

def _generate_samples(self, backend: AbstractBackend):
    """Generates samples from the uniform distribution."""
    # torch.rand samples from [0, 1), so we scale and shift.
    return backend.rand(size=(self.num,)) * (self.high - self.low) + self.low