Skip to content

Base Conditioner

Conditioning mechanisms for injecting PDE parameters into flow models.

base_conditioner

Base Conditioner Abstract Class

Defines the interface for conditioning mechanisms in FlowPDE.

BaseConditioner

Bases: ABC, Module

Abstract base class for conditioning mechanisms.

Conditioners process and inject conditioning information (e.g., PDE coefficients, boundary conditions, observations) into the flow model.

Examples: - Concatenation: Simply concatenate condition with input - Cross-attention: Use attention mechanism to condition - FiLM: Feature-wise Linear Modulation - Adaptive Instance Normalization

Source code in flowpde/core/base_conditioner.py 
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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
class BaseConditioner(ABC, nn.Module):
    """
    Abstract base class for conditioning mechanisms.

    Conditioners process and inject conditioning information (e.g., PDE coefficients,
    boundary conditions, observations) into the flow model.

    Examples:
    - Concatenation: Simply concatenate condition with input
    - Cross-attention: Use attention mechanism to condition
    - FiLM: Feature-wise Linear Modulation
    - Adaptive Instance Normalization
    """

    def __init__(self, **kwargs: Any):
        """
        Initialize conditioner.

        Args:
            **kwargs: Conditioner-specific parameters
        """
        super().__init__()
        self._config = kwargs

    @abstractmethod
    def forward(
        self,
        x: Tensor,
        condition: Tensor,
        **kwargs: Any
    ) -> Union[Tensor, Dict[str, Tensor]]:
        """
        Apply conditioning to input.

        Args:
            x: Input tensor (batch_size, dim) or (batch_size, channels, H, W)
            condition: Conditioning information
            **kwargs: Additional parameters

        Returns:
            conditioned_x: Conditioned tensor (same shape as x or transformed)
            OR
            Dictionary with conditioned features and auxiliary outputs
        """
        raise NotImplementedError

    def preprocess_condition(self, condition: Tensor) -> Tensor:
        """
        Preprocess conditioning information before use.

        Override this to implement custom preprocessing (e.g., normalization,
        embedding, feature extraction).

        Args:
            condition: Raw conditioning tensor

        Returns:
            Preprocessed conditioning tensor
        """
        return condition

    def get_config(self) -> Dict[str, Any]:
        """Get conditioner configuration."""
        return {
            'type': self.__class__.__name__,
            **self._config
        }

    def extra_repr(self) -> str:
        """Extra information for repr."""
        config_str = ', '.join(f'{k}={v}' for k, v in self._config.items())
        return config_str
__init__(**kwargs)

Initialize conditioner.

Parameters:

Name Type Description Default
**kwargs Any

Conditioner-specific parameters

 {}
Source code in flowpde/core/base_conditioner.py 
28
29
30
31
32
33
34
35
36
def __init__(self, **kwargs: Any):
    """
    Initialize conditioner.

    Args:
        **kwargs: Conditioner-specific parameters
    """
    super().__init__()
    self._config = kwargs
forward(x, condition, **kwargs) abstractmethod

Apply conditioning to input.

Parameters:

Name Type Description Default
x Tensor

Input tensor (batch_size, dim) or (batch_size, channels, H, W)

 required
condition Tensor

Conditioning information

 required
**kwargs Any

Additional parameters

 {}

Returns:

Name Type Description
conditioned_x Union[Tensor, Dict[str, Tensor]]

Conditioned tensor (same shape as x or transformed)
Union[Tensor, Dict[str, Tensor]]

OR
Union[Tensor, Dict[str, Tensor]]

Dictionary with conditioned features and auxiliary outputs
Source code in flowpde/core/base_conditioner.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@abstractmethod
def forward(
    self,
    x: Tensor,
    condition: Tensor,
    **kwargs: Any
) -> Union[Tensor, Dict[str, Tensor]]:
    """
    Apply conditioning to input.

    Args:
        x: Input tensor (batch_size, dim) or (batch_size, channels, H, W)
        condition: Conditioning information
        **kwargs: Additional parameters

    Returns:
        conditioned_x: Conditioned tensor (same shape as x or transformed)
        OR
        Dictionary with conditioned features and auxiliary outputs
    """
    raise NotImplementedError
preprocess_condition(condition)

Preprocess conditioning information before use.

Override this to implement custom preprocessing (e.g., normalization, embedding, feature extraction).

Parameters:

Name Type Description Default
condition Tensor

Raw conditioning tensor

 required

Returns:

Type Description
Tensor

Preprocessed conditioning tensor
Source code in flowpde/core/base_conditioner.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def preprocess_condition(self, condition: Tensor) -> Tensor:
    """
    Preprocess conditioning information before use.

    Override this to implement custom preprocessing (e.g., normalization,
    embedding, feature extraction).

    Args:
        condition: Raw conditioning tensor

    Returns:
        Preprocessed conditioning tensor
    """
    return condition
get_config()

Get conditioner configuration.

Source code in flowpde/core/base_conditioner.py 
75
76
77
78
79
80
def get_config(self) -> Dict[str, Any]:
    """Get conditioner configuration."""
    return {
        'type': self.__class__.__name__,
        **self._config
    }
extra_repr()

Extra information for repr.

Source code in flowpde/core/base_conditioner.py 
82
83
84
85
def extra_repr(self) -> str:
    """Extra information for repr."""
    config_str = ', '.join(f'{k}={v}' for k, v in self._config.items())
    return config_str

ConcatConditioner

Bases: BaseConditioner

Simple concatenation-based conditioning.

Concatenates condition with input along specified dimension.

Source code in flowpde/core/base_conditioner.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class ConcatConditioner(BaseConditioner):
    """
    Simple concatenation-based conditioning.

    Concatenates condition with input along specified dimension.
    """

    def __init__(self, dim: int = 1):
        """
        Initialize concatenation conditioner.

        Args:
            dim: Dimension along which to concatenate (default: 1 for feature dimension)
        """
        super().__init__(dim=dim)
        self.dim = dim

    def forward(
        self,
        x: Tensor,
        condition: Tensor,
    ) -> Tensor:
        """
        Concatenate condition with input.

        Args:
            x: Input tensor
            condition: Conditioning tensor

        Returns:
            Concatenated tensor
        """
        condition = self.preprocess_condition(condition)

        # Ensure compatible shapes
        if x.dim() != condition.dim():
            # Try to reshape condition to match x
            if condition.dim() == 2 and x.dim() == 4:
                # Condition is (B, C), x is (B, C, H, W)
                # Reshape condition to (B, C, 1, 1) and broadcast
                condition = condition.view(condition.shape[0], condition.shape[1], 1, 1)
                condition = condition.expand(-1, -1, x.shape[2], x.shape[3])

        return torch.cat([x, condition], dim=self.dim)
__init__(dim=1)

Initialize concatenation conditioner.

Parameters:

Name Type Description Default
dim int

Dimension along which to concatenate (default: 1 for feature dimension)

 1
Source code in flowpde/core/base_conditioner.py 
 95
 96
 97
 98
 99
100
101
102
103
def __init__(self, dim: int = 1):
    """
    Initialize concatenation conditioner.

    Args:
        dim: Dimension along which to concatenate (default: 1 for feature dimension)
    """
    super().__init__(dim=dim)
    self.dim = dim
forward(x, condition)

Concatenate condition with input.

Parameters:

Name Type Description Default
x Tensor

Input tensor

 required
condition Tensor

Conditioning tensor

 required

Returns:

Type Description
Tensor

Concatenated tensor
Source code in flowpde/core/base_conditioner.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def forward(
    self,
    x: Tensor,
    condition: Tensor,
) -> Tensor:
    """
    Concatenate condition with input.

    Args:
        x: Input tensor
        condition: Conditioning tensor

    Returns:
        Concatenated tensor
    """
    condition = self.preprocess_condition(condition)

    # Ensure compatible shapes
    if x.dim() != condition.dim():
        # Try to reshape condition to match x
        if condition.dim() == 2 and x.dim() == 4:
            # Condition is (B, C), x is (B, C, H, W)
            # Reshape condition to (B, C, 1, 1) and broadcast
            condition = condition.view(condition.shape[0], condition.shape[1], 1, 1)
            condition = condition.expand(-1, -1, x.shape[2], x.shape[3])

    return torch.cat([x, condition], dim=self.dim)

FiLMConditioner

Bases: BaseConditioner

Feature-wise Linear Modulation (FiLM) conditioning.

Generates scale and shift parameters from condition to modulate features as (1 + scale) * x + shift.

Source code in flowpde/core/base_conditioner.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
class FiLMConditioner(BaseConditioner):
    """
    Feature-wise Linear Modulation (FiLM) conditioning.

    Generates scale and shift parameters from condition to modulate
    features as ``(1 + scale) * x + shift``.
    """

    def __init__(
        self,
        condition_dim: int,
        feature_dim: int,
        hidden_dim: Optional[int] = None
    ):
        """
        Initialize FiLM conditioner.

        Args:
            condition_dim: Dimension of conditioning input
            feature_dim: Dimension of features to modulate
            hidden_dim: Hidden dimension for parameter generation (default: condition_dim)
        """
        super().__init__(
            condition_dim=condition_dim,
            feature_dim=feature_dim,
            hidden_dim=hidden_dim
        )

        hidden_dim = hidden_dim or condition_dim

        # Network to generate scale and shift parameters
        self.film_generator = nn.Sequential(
            nn.Linear(condition_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 2 * feature_dim)  # 2x for scale and shift
        )

    def forward(
        self,
        x: Tensor,
        condition: Tensor,
        **kwargs: Any
    ) -> Tensor:
        """
        Apply FiLM conditioning.

        Args:
            x: Input features (batch_size, feature_dim, ...)
            condition: Conditioning information (batch_size, condition_dim)

        Returns:
            Modulated features: (1 + scale) * x + shift
        """
        condition = self.preprocess_condition(condition)

        # Flatten condition if needed
        if condition.dim() > 2:
            condition = condition.flatten(start_dim=1)

        # Generate scale and shift parameters
        params = self.film_generator(condition)  # (batch_size, 2 * feature_dim)
        scale, shift = params.chunk(2, dim=1)    # Each (batch_size, feature_dim)

        # Reshape for broadcasting if x is spatial
        if x.dim() == 4:  # (B, C, H, W)
            scale = scale.view(scale.shape[0], scale.shape[1], 1, 1)
            shift = shift.view(shift.shape[0], shift.shape[1], 1, 1)
        elif x.dim() == 3:  # (B, L, C)
            scale = scale.unsqueeze(1)
            shift = shift.unsqueeze(1)

        # Apply FiLM as (1 + scale), the standard formulation: the
        # generator starts near zero, so the block starts as the identity
        # and learns a modulation, instead of multiplying features by a
        # random zero-mean gain before training has begun.
        return (1 + scale) * x + shift
__init__(condition_dim, feature_dim, hidden_dim=None)

Initialize FiLM conditioner.

Parameters:

Name Type Description Default
condition_dim int

Dimension of conditioning input

 required
feature_dim int

Dimension of features to modulate

 required
hidden_dim Optional[int]

Hidden dimension for parameter generation (default: condition_dim)

 None
Source code in flowpde/core/base_conditioner.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def __init__(
    self,
    condition_dim: int,
    feature_dim: int,
    hidden_dim: Optional[int] = None
):
    """
    Initialize FiLM conditioner.

    Args:
        condition_dim: Dimension of conditioning input
        feature_dim: Dimension of features to modulate
        hidden_dim: Hidden dimension for parameter generation (default: condition_dim)
    """
    super().__init__(
        condition_dim=condition_dim,
        feature_dim=feature_dim,
        hidden_dim=hidden_dim
    )

    hidden_dim = hidden_dim or condition_dim

    # Network to generate scale and shift parameters
    self.film_generator = nn.Sequential(
        nn.Linear(condition_dim, hidden_dim),
        nn.ReLU(),
        nn.Linear(hidden_dim, 2 * feature_dim)  # 2x for scale and shift
    )
forward(x, condition, **kwargs)

Apply FiLM conditioning.

Parameters:

Name Type Description Default
x Tensor

Input features (batch_size, feature_dim, ...)

 required
condition Tensor

Conditioning information (batch_size, condition_dim)

 required

Returns:

Type Description
Tensor

Modulated features: (1 + scale) * x + shift
Source code in flowpde/core/base_conditioner.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def forward(
    self,
    x: Tensor,
    condition: Tensor,
    **kwargs: Any
) -> Tensor:
    """
    Apply FiLM conditioning.

    Args:
        x: Input features (batch_size, feature_dim, ...)
        condition: Conditioning information (batch_size, condition_dim)

    Returns:
        Modulated features: (1 + scale) * x + shift
    """
    condition = self.preprocess_condition(condition)

    # Flatten condition if needed
    if condition.dim() > 2:
        condition = condition.flatten(start_dim=1)

    # Generate scale and shift parameters
    params = self.film_generator(condition)  # (batch_size, 2 * feature_dim)
    scale, shift = params.chunk(2, dim=1)    # Each (batch_size, feature_dim)

    # Reshape for broadcasting if x is spatial
    if x.dim() == 4:  # (B, C, H, W)
        scale = scale.view(scale.shape[0], scale.shape[1], 1, 1)
        shift = shift.view(shift.shape[0], shift.shape[1], 1, 1)
    elif x.dim() == 3:  # (B, L, C)
        scale = scale.unsqueeze(1)
        shift = shift.unsqueeze(1)

    # Apply FiLM as (1 + scale), the standard formulation: the
    # generator starts near zero, so the block starts as the identity
    # and learns a modulation, instead of multiplying features by a
    # random zero-mean gain before training has begun.
    return (1 + scale) * x + shift

NullConditioner

Bases: BaseConditioner

No-op conditioner that returns input unchanged.

Useful for unconditional models or as a placeholder.

Source code in flowpde/core/base_conditioner.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
class NullConditioner(BaseConditioner):
    """
    No-op conditioner that returns input unchanged.

    Useful for unconditional models or as a placeholder.
    """

    def __init__(self):
        super().__init__()

    def forward(
        self,
        x: Tensor,
        condition: Optional[Tensor] = None,
        **kwargs: Any
    ) -> Tensor:
        """Return input unchanged."""
        return x
forward(x, condition=None, **kwargs)

Return input unchanged.

Source code in flowpde/core/base_conditioner.py 
222
223
224
225
226
227
228
229
def forward(
    self,
    x: Tensor,
    condition: Optional[Tensor] = None,
    **kwargs: Any
) -> Tensor:
    """Return input unchanged."""
    return x