Skip to content

base_generator

BaseGenerator

Bases: ABC

An abstract base class for writing structure generators.

Source code in src/torch_structure/generators/base_generator.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class BaseGenerator(ABC):
    r"""An abstract base class for writing structure generators."""

    def __init__(self, **overrides):
        self.overrides = overrides

        # Optional attempt logic for handling invalid samples
        self.max_attempts = 1
        self.attempt_count = 0
        self.success_count = 0
        self.verbose = False

    def __call__(self):
        return self._build()

    def _build(self) -> StructData:
        if self.max_attempts == 1:
            self.input = self.sample_input(**self.overrides)
            self.validate_input(**self.input)
            return self.generate(**self.input)
        else:
            for attempt in range(self.max_attempts):
                self.attempt_count += 1
                self.input = self.sample_input(**self.overrides)
                self.validate_input(**self.input)
                try:
                    result = self.generate(**self.input)
                    self.success_count += 1
                    return result
                except InvalidSampleError as e:
                    error_message = f"Attempt {attempt+1}/{self.max_attempts} failed: {str(e)}"
                    if self.verbose:
                        print(error_message)
                    continue

            raise RuntimeError(
                f"Failed to generate structure after {self.max_attempts} attempts. {error_message}"
            )

    @abstractmethod
    def sample_input(self, **overrides) -> dict:
        """Samples missing input parameters."""
        pass

    def validate_input(self, **input) -> None:
        """Raises errors if inputs are invalid."""
        pass

    @abstractmethod
    def generate(self, **input) -> StructData:
        """Generates a StructData object."""
        pass

generate(**input) abstractmethod

Generates a StructData object.

Source code in src/torch_structure/generators/base_generator.py 
57
58
59
60
@abstractmethod
def generate(self, **input) -> StructData:
    """Generates a StructData object."""
    pass

sample_input(**overrides) abstractmethod

Samples missing input parameters.

Source code in src/torch_structure/generators/base_generator.py 
48
49
50
51
@abstractmethod
def sample_input(self, **overrides) -> dict:
    """Samples missing input parameters."""
    pass

validate_input(**input)

Raises errors if inputs are invalid.

Source code in src/torch_structure/generators/base_generator.py 
53
54
55
def validate_input(self, **input) -> None:
    """Raises errors if inputs are invalid."""
    pass

InvalidSampleError

Bases: Exception

Sample violated generator constraints.

Source code in src/torch_structure/generators/base_generator.py 
5
6
7
class InvalidSampleError(Exception):
    """Sample violated generator constraints."""
    pass