Skip to content

anypinn.core.samplers

Collocation point sampling strategies for PINN training.

AdaptiveSampler

Residual-weighted adaptive collocation sampler.

Draws an oversample of candidate points, scores them using a ResidualScorer, and retains the top-scoring subset. A configurable exploration_ratio ensures a fraction of purely random points to prevent mode collapse.

Parameters:

Name Type Description Default
scorer ResidualScorer

Callable returning per-point residual scores (n,).

 required
oversample_factor int

Multiplier on n for candidate generation.

 4
exploration_ratio float

Fraction of the budget reserved for random points.

 0.2
seed int | None

Optional seed for reproducible sampling.

 None
Source code in src/anypinn/core/samplers.py 
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
class AdaptiveSampler:
    """Residual-weighted adaptive collocation sampler.

    Draws an oversample of candidate points, scores them using a
    ``ResidualScorer``, and retains the top-scoring subset. A configurable
    ``exploration_ratio`` ensures a fraction of purely random points to prevent
    mode collapse.

    Args:
        scorer: Callable returning per-point residual scores ``(n,)``.
        oversample_factor: Multiplier on ``n`` for candidate generation.
        exploration_ratio: Fraction of the budget reserved for random points.
        seed: Optional seed for reproducible sampling.
    """

    def __init__(
        self,
        scorer: ResidualScorer,
        oversample_factor: int = 4,
        exploration_ratio: float = 0.2,
        seed: int | None = None,
    ) -> None:
        if oversample_factor < 1:
            raise ValueError(f"oversample_factor must be >= 1, got {oversample_factor}.")
        if not (0.0 <= exploration_ratio <= 1.0):
            raise ValueError(f"exploration_ratio must be in [0, 1], got {exploration_ratio}.")
        self._scorer = scorer
        self._oversample = oversample_factor
        self._explore = exploration_ratio
        self._random = RandomSampler(seed=seed)

    def sample(self, n: int, domain: Domain) -> Tensor:
        n_explore = max(1, int(n * self._explore))
        n_exploit = n - n_explore

        explore_pts = self._random.sample(n_explore, domain)

        if n_exploit <= 0:
            return explore_pts

        n_candidates = n_exploit * self._oversample
        candidates = self._random.sample(n_candidates, domain)

        with torch.no_grad():
            scores = self._scorer.residual_score(candidates)

        _, top_idx = scores.topk(min(n_exploit, len(scores)))
        exploit_pts = candidates[top_idx]

        return torch.cat([explore_pts, exploit_pts], dim=0)

__init__(scorer: ResidualScorer, oversample_factor: int = 4, exploration_ratio: float = 0.2, seed: int | None = None) -> None

Source code in src/anypinn/core/samplers.py 
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def __init__(
    self,
    scorer: ResidualScorer,
    oversample_factor: int = 4,
    exploration_ratio: float = 0.2,
    seed: int | None = None,
) -> None:
    if oversample_factor < 1:
        raise ValueError(f"oversample_factor must be >= 1, got {oversample_factor}.")
    if not (0.0 <= exploration_ratio <= 1.0):
        raise ValueError(f"exploration_ratio must be in [0, 1], got {exploration_ratio}.")
    self._scorer = scorer
    self._oversample = oversample_factor
    self._explore = exploration_ratio
    self._random = RandomSampler(seed=seed)

sample(n: int, domain: Domain) -> Tensor

Source code in src/anypinn/core/samplers.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def sample(self, n: int, domain: Domain) -> Tensor:
    n_explore = max(1, int(n * self._explore))
    n_exploit = n - n_explore

    explore_pts = self._random.sample(n_explore, domain)

    if n_exploit <= 0:
        return explore_pts

    n_candidates = n_exploit * self._oversample
    candidates = self._random.sample(n_candidates, domain)

    with torch.no_grad():
        scores = self._scorer.residual_score(candidates)

    _, top_idx = scores.topk(min(n_exploit, len(scores)))
    exploit_pts = candidates[top_idx]

    return torch.cat([explore_pts, exploit_pts], dim=0)

CollocationSampler

Bases: Protocol

Protocol for collocation point samplers.

Implementations must return a tensor of shape (n, domain.ndim) with all points inside the domain bounds.

Source code in src/anypinn/core/samplers.py 
15
16
17
18
19
20
21
22
class CollocationSampler(Protocol):
    """Protocol for collocation point samplers.

    Implementations must return a tensor of shape ``(n, domain.ndim)`` with all
    points inside the domain bounds.
    """

    def sample(self, n: int, domain: Domain) -> Tensor: ...

sample(n: int, domain: Domain) -> Tensor

Source code in src/anypinn/core/samplers.py 
22
def sample(self, n: int, domain: Domain) -> Tensor: ...

LatinHypercubeSampler

Latin Hypercube sampler (pure-PyTorch, no SciPy dependency).

Stratifies each dimension into n equal intervals and places one sample per interval, then shuffles columns independently.

Parameters:

Name Type Description Default
seed int | None

Optional seed for reproducible sampling.

 None
Source code in src/anypinn/core/samplers.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
class LatinHypercubeSampler:
    """Latin Hypercube sampler (pure-PyTorch, no SciPy dependency).

    Stratifies each dimension into ``n`` equal intervals and places one sample
    per interval, then shuffles columns independently.

    Args:
        seed: Optional seed for reproducible sampling.
    """

    def __init__(self, seed: int | None = None) -> None:
        self._gen = torch.Generator()
        if seed is not None:
            self._gen.manual_seed(seed)

    def sample(self, n: int, domain: Domain) -> Tensor:
        d = domain.ndim
        result = torch.empty(n, d)

        for i, (lo, hi) in enumerate(domain.bounds):
            perm = torch.randperm(n, generator=self._gen)
            base = (perm.float() + torch.rand(n, generator=self._gen)) / n
            result[:, i] = base * (hi - lo) + lo

        return result

__init__(seed: int | None = None) -> None

Source code in src/anypinn/core/samplers.py 
78
79
80
81
def __init__(self, seed: int | None = None) -> None:
    self._gen = torch.Generator()
    if seed is not None:
        self._gen.manual_seed(seed)

sample(n: int, domain: Domain) -> Tensor

Source code in src/anypinn/core/samplers.py 
83
84
85
86
87
88
89
90
91
92
def sample(self, n: int, domain: Domain) -> Tensor:
    d = domain.ndim
    result = torch.empty(n, d)

    for i, (lo, hi) in enumerate(domain.bounds):
        perm = torch.randperm(n, generator=self._gen)
        base = (perm.float() + torch.rand(n, generator=self._gen)) / n
        result[:, i] = base * (hi - lo) + lo

    return result

LogUniform1DSampler

Log-uniform sampler for 1-D domains (reproduces SIR collocation behavior).

Samples uniformly in log1p space and maps back via expm1, producing a distribution that is denser near the lower bound — useful for epidemic models where early dynamics are most informative.

Parameters:

Name Type Description Default
seed int | None

Optional seed for reproducible sampling.

 None

Raises:

Type Description
ValueError

If the domain is not 1-D or x0 <= -1.
Source code in src/anypinn/core/samplers.py 
 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
class LogUniform1DSampler:
    """Log-uniform sampler for 1-D domains (reproduces SIR collocation behavior).

    Samples uniformly in ``log1p`` space and maps back via ``expm1``, producing
    a distribution that is denser near the lower bound — useful for epidemic
    models where early dynamics are most informative.

    Args:
        seed: Optional seed for reproducible sampling.

    Raises:
        ValueError: If the domain is not 1-D or ``x0 <= -1``.
    """

    def __init__(self, seed: int | None = None) -> None:
        self._gen = torch.Generator()
        if seed is not None:
            self._gen.manual_seed(seed)

    def sample(self, n: int, domain: Domain) -> Tensor:
        if domain.ndim != 1:
            raise ValueError(
                f"log_uniform_1d sampler supports only 1-D domains, got ndim={domain.ndim}."
            )
        x0, x1 = domain.x0, domain.x1
        if x0 <= -1.0:
            raise ValueError(f"log_uniform_1d requires x0 > -1 for log1p, got x0={x0}.")
        log_lo = torch.tensor(x0, dtype=torch.float32).log1p()
        log_hi = torch.tensor(x1, dtype=torch.float32).log1p()
        u = torch.rand((n, 1), generator=self._gen)
        return torch.expm1(u * (log_hi - log_lo) + log_lo)

__init__(seed: int | None = None) -> None

Source code in src/anypinn/core/samplers.py 
109
110
111
112
def __init__(self, seed: int | None = None) -> None:
    self._gen = torch.Generator()
    if seed is not None:
        self._gen.manual_seed(seed)

sample(n: int, domain: Domain) -> Tensor

Source code in src/anypinn/core/samplers.py 
114
115
116
117
118
119
120
121
122
123
124
125
def sample(self, n: int, domain: Domain) -> Tensor:
    if domain.ndim != 1:
        raise ValueError(
            f"log_uniform_1d sampler supports only 1-D domains, got ndim={domain.ndim}."
        )
    x0, x1 = domain.x0, domain.x1
    if x0 <= -1.0:
        raise ValueError(f"log_uniform_1d requires x0 > -1 for log1p, got x0={x0}.")
    log_lo = torch.tensor(x0, dtype=torch.float32).log1p()
    log_hi = torch.tensor(x1, dtype=torch.float32).log1p()
    u = torch.rand((n, 1), generator=self._gen)
    return torch.expm1(u * (log_hi - log_lo) + log_lo)

RandomSampler

Uniform random sampler inside domain bounds.

Parameters:

Name Type Description Default
seed int | None

Optional seed for reproducible sampling.

 None
Source code in src/anypinn/core/samplers.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class RandomSampler:
    """Uniform random sampler inside domain bounds.

    Args:
        seed: Optional seed for reproducible sampling.
    """

    def __init__(self, seed: int | None = None) -> None:
        self._gen = torch.Generator()
        if seed is not None:
            self._gen.manual_seed(seed)

    def sample(self, n: int, domain: Domain) -> Tensor:
        d = domain.ndim
        u = torch.rand((n, d), generator=self._gen)
        for i, (lo, hi) in enumerate(domain.bounds):
            u[:, i] = u[:, i] * (hi - lo) + lo
        return u

__init__(seed: int | None = None) -> None

Source code in src/anypinn/core/samplers.py 
55
56
57
58
def __init__(self, seed: int | None = None) -> None:
    self._gen = torch.Generator()
    if seed is not None:
        self._gen.manual_seed(seed)

sample(n: int, domain: Domain) -> Tensor

Source code in src/anypinn/core/samplers.py 
60
61
62
63
64
65
def sample(self, n: int, domain: Domain) -> Tensor:
    d = domain.ndim
    u = torch.rand((n, d), generator=self._gen)
    for i, (lo, hi) in enumerate(domain.bounds):
        u[:, i] = u[:, i] * (hi - lo) + lo
    return u

ResidualScorer

Bases: Protocol

Protocol for scoring candidate collocation points by PDE residual magnitude.

Source code in src/anypinn/core/samplers.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
class ResidualScorer(Protocol):
    """Protocol for scoring candidate collocation points by PDE residual magnitude."""

    def residual_score(self, x: Tensor) -> Tensor:
        """Return per-point non-negative residual score of shape ``(n,)``.

        Args:
            x: Candidate collocation points ``(n, d)``.

        Returns:
            Scores ``(n,)`` — higher means larger residual.
        """
        ...

residual_score(x: Tensor) -> Tensor

Return per-point non-negative residual score of shape (n,).

Parameters:

Name Type Description Default
x Tensor

Candidate collocation points (n, d).

 required

Returns:

Type Description
Tensor

Scores (n,) — higher means larger residual.
Source code in src/anypinn/core/samplers.py 
131
132
133
134
135
136
137
138
139
140
def residual_score(self, x: Tensor) -> Tensor:
    """Return per-point non-negative residual score of shape ``(n,)``.

    Args:
        x: Candidate collocation points ``(n, d)``.

    Returns:
        Scores ``(n,)`` — higher means larger residual.
    """
    ...

UniformSampler

Cartesian grid sampler that distributes points evenly across the domain.

For d-dimensional domains, places ceil(n^(1/d)) points per axis then takes the first n points of the resulting grid.

Parameters:

Name Type Description Default
seed int | None

Optional seed (unused — grid is deterministic).

 None
Source code in src/anypinn/core/samplers.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class UniformSampler:
    """Cartesian grid sampler that distributes points evenly across the domain.

    For d-dimensional domains, places ``ceil(n^(1/d))`` points per axis then
    takes the first ``n`` points of the resulting grid.

    Args:
        seed: Optional seed (unused — grid is deterministic).
    """

    def __init__(self, seed: int | None = None) -> None:
        pass

    def sample(self, n: int, domain: Domain) -> Tensor:
        d = domain.ndim
        pts_per_dim = math.ceil(n ** (1.0 / d))

        linspaces = [torch.linspace(lo, hi, pts_per_dim) for lo, hi in domain.bounds]
        grids = torch.meshgrid(*linspaces, indexing="ij")
        flat = torch.stack([g.reshape(-1) for g in grids], dim=-1)
        return flat[:n]

__init__(seed: int | None = None) -> None

Source code in src/anypinn/core/samplers.py 
35
36
def __init__(self, seed: int | None = None) -> None:
    pass

sample(n: int, domain: Domain) -> Tensor

Source code in src/anypinn/core/samplers.py 
38
39
40
41
42
43
44
45
def sample(self, n: int, domain: Domain) -> Tensor:
    d = domain.ndim
    pts_per_dim = math.ceil(n ** (1.0 / d))

    linspaces = [torch.linspace(lo, hi, pts_per_dim) for lo, hi in domain.bounds]
    grids = torch.meshgrid(*linspaces, indexing="ij")
    flat = torch.stack([g.reshape(-1) for g in grids], dim=-1)
    return flat[:n]

build_sampler(strategy: CollocationStrategies, seed: int | None = None, scorer: ResidualScorer | None = None) -> CollocationSampler

Construct a collocation sampler from a strategy name.

Parameters:

Name Type Description Default
strategy CollocationStrategies

One of the CollocationStrategies literals.

 required
seed int | None

Optional seed for reproducible sampling.

 None
scorer ResidualScorer | None

Required when strategy="adaptive".

 None

Returns:

Type Description
CollocationSampler

A sampler instance satisfying the CollocationSampler protocol.

Raises:

Type Description
ValueError

If strategy="adaptive" but no scorer is provided.
Source code in src/anypinn/core/samplers.py 
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
def build_sampler(
    strategy: CollocationStrategies,
    seed: int | None = None,
    scorer: ResidualScorer | None = None,
) -> CollocationSampler:
    """Construct a collocation sampler from a strategy name.

    Args:
        strategy: One of the ``CollocationStrategies`` literals.
        seed: Optional seed for reproducible sampling.
        scorer: Required when ``strategy="adaptive"``.

    Returns:
        A sampler instance satisfying the ``CollocationSampler`` protocol.

    Raises:
        ValueError: If ``strategy="adaptive"`` but no scorer is provided.
    """
    if strategy == "adaptive":
        if scorer is None:
            raise ValueError(
                "AdaptiveSampler requires a ResidualScorer. "
                "Pass a scorer via PINNDataModule or use a different strategy."
            )
        return AdaptiveSampler(scorer=scorer, seed=seed)

    cls = _SAMPLER_REGISTRY.get(strategy)
    if cls is None:
        raise ValueError(
            f"Unknown collocation strategy '{strategy}'. "
            f"Choose from: {', '.join(_SAMPLER_REGISTRY)} or 'adaptive'."
        )
    return cls(seed=seed)