Skip to content

Model Interface

openlithohub.models.base

Abstract base class for lithography optimization models.

PredictionResult dataclass

Result from a model prediction.

Output contract for mask:

  • shape: same as the input design ((H, W), (C, H, W), or (B, C, H, W)).
  • dtype: torch.float32.
  • range: values in [0, 1]. Models that emit logits MUST apply sigmoid (and any project-required binarization) before populating this field.
  • binarization: implementations should return binarized masks (mask > 0.5) by default so downstream metrics (DRC/MRC, PV-band, shot-count) see the same contract across baselines. Models that benefit from emitting a soft mask for further processing should document the deviation in their predict() docstring; the default Neural-ILT and GAN-OPC adapters both binarize.

contour is optional and used for vector-mode visualization / GDS export. metadata carries per-model side-channel info (weights provenance, iteration count, ...).

Source code in src/openlithohub/models/base.py 
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
61
62
63
64
@dataclass
class PredictionResult:
    """Result from a model prediction.

    Output contract for ``mask``:

    - **shape**: same as the input ``design`` (``(H, W)``, ``(C, H, W)``,
      or ``(B, C, H, W)``).
    - **dtype**: ``torch.float32``.
    - **range**: values in ``[0, 1]``. Models that emit logits MUST apply
      ``sigmoid`` (and any project-required binarization) before populating
      this field.
    - **binarization**: implementations should return *binarized* masks
      (``mask > 0.5``) by default so downstream metrics (DRC/MRC, PV-band,
      shot-count) see the same contract across baselines. Models that
      benefit from emitting a soft mask for further processing should
      document the deviation in their ``predict()`` docstring; the
      default Neural-ILT and GAN-OPC adapters both binarize.

    ``contour`` is optional and used for vector-mode visualization /
    GDS export. ``metadata`` carries per-model side-channel info
    (weights provenance, iteration count, ...).
    """

    mask: torch.Tensor
    contour: torch.Tensor | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def _repr_html_(self) -> str:
        from openlithohub.jupyter._html import (
            kv_table,
            mask_thumbnail_png_b64,
            panel,
            png_b64_to_img_tag,
        )

        shape = "x".join(str(d) for d in tuple(self.mask.shape)) if self.mask is not None else "—"
        rows = [
            ("mask shape", shape),
            ("mask dtype", str(self.mask.dtype) if self.mask is not None else "—"),
            ("contour", "yes" if self.contour is not None else "no"),
        ]
        for k, v in self.metadata.items():
            rows.append((f"meta:{k}", str(v)))

        img = png_b64_to_img_tag(mask_thumbnail_png_b64(self.mask), alt="mask")
        body = (
            f'<div style="display:flex;gap:10px;align-items:flex-start;">'
            f'<div style="flex:0 0 auto;">{img}</div>'
            f'<div style="flex:1 1 auto;">{kv_table(rows)}</div>'
            f"</div>"
        )
        return panel(title="PredictionResult", header_html="", body_html=body)

LithographyModel

Bases: ABC

Abstract interface for lithography optimization models.

Any model (heuristic OPC, U-Net, diffusion-based ILT, curvyILT) can join the evaluation pipeline by implementing predict().

Subclasses MUST set the class-level NAME attribute. The registry reads it without instantiating the class, so it cannot be set in __init__.

Source code in src/openlithohub/models/base.py 
 67
 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
 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
132
133
134
135
136
137
138
139
140
141
class LithographyModel(ABC):
    """Abstract interface for lithography optimization models.

    Any model (heuristic OPC, U-Net, diffusion-based ILT, curvyILT)
    can join the evaluation pipeline by implementing predict().

    Subclasses MUST set the class-level ``NAME`` attribute. The registry
    reads it without instantiating the class, so it cannot be set in
    ``__init__``.
    """

    NAME: ClassVar[str]
    SUPPORTS_CURVILINEAR: ClassVar[bool] = False
    RECEPTIVE_FIELD_PX: ClassVar[int] = 0
    """Half-width of the model's receptive field in pixels.

    Tile inference adds at least this many pixels of halo on every side
    so the model sees real layout context, not zero-padding, at tile
    boundaries. Models with a static convolutional receptive field
    should set this on the subclass; iterative optimizers that consume
    their entire input (e.g. level-set ILT) can leave it 0 — their halo
    comes from the optical interaction radius of the process node.
    """

    @property
    def name(self) -> str:
        """Human-readable model name for leaderboard display."""
        return type(self).NAME

    @property
    def supports_curvilinear(self) -> bool:
        """Whether this model produces curvilinear (non-Manhattan) output."""
        return type(self).SUPPORTS_CURVILINEAR

    @property
    def receptive_field_px(self) -> int:
        """Per-instance accessor for the class-level ``RECEPTIVE_FIELD_PX``."""
        return type(self).RECEPTIVE_FIELD_PX

    @abstractmethod
    def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
        """Run model inference on a design layout tensor.

        Args:
            design: Input design tensor of shape (H, W) or (B, C, H, W).
            **kwargs: Model-specific parameters (process node, dose, etc.)

        Returns:
            PredictionResult with the optimized mask and optional contour.
        """
        ...

    def setup(self) -> None:
        """Optional setup hook (load weights, initialize GPU, etc.)."""

    def teardown(self) -> None:
        """Optional cleanup hook."""

    def to_torch_module(self) -> torch.nn.Module:
        """Return a single ``nn.Module`` that maps an input mask/design tensor
        to an output mask tensor, suitable for ONNX / TorchScript export.

        Default raises ``NotImplementedError``. Override this on models that
        wrap a static ``nn.Module`` (e.g. Neural-ILT). Iterative optimizers
        like Level-Set ILT do not have a single forward graph and cannot be
        exported — those should keep the default.

        The returned module must be in ``eval()`` mode and accept tensors
        of shape ``(B, 1, H, W)`` in ``[0, 1]``, returning the same shape.
        """
        raise NotImplementedError(
            f"Model {type(self).NAME!r} does not support export to a single "
            "nn.Module — override LithographyModel.to_torch_module() if it "
            "wraps a static forward graph."
        )

RECEPTIVE_FIELD_PX = 0 class-attribute

Half-width of the model's receptive field in pixels.

Tile inference adds at least this many pixels of halo on every side so the model sees real layout context, not zero-padding, at tile boundaries. Models with a static convolutional receptive field should set this on the subclass; iterative optimizers that consume their entire input (e.g. level-set ILT) can leave it 0 — their halo comes from the optical interaction radius of the process node.

name property

Human-readable model name for leaderboard display.

supports_curvilinear property

Whether this model produces curvilinear (non-Manhattan) output.

receptive_field_px property

Per-instance accessor for the class-level RECEPTIVE_FIELD_PX.

predict(design, **kwargs) abstractmethod

Run model inference on a design layout tensor.

Parameters:

Name Type Description Default
design Tensor

Input design tensor of shape (H, W) or (B, C, H, W).

 required
**kwargs Any

Model-specific parameters (process node, dose, etc.)

 {}

Returns:

Type Description
PredictionResult

PredictionResult with the optimized mask and optional contour.
Source code in src/openlithohub/models/base.py 
106
107
108
109
110
111
112
113
114
115
116
117
@abstractmethod
def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
    """Run model inference on a design layout tensor.

    Args:
        design: Input design tensor of shape (H, W) or (B, C, H, W).
        **kwargs: Model-specific parameters (process node, dose, etc.)

    Returns:
        PredictionResult with the optimized mask and optional contour.
    """
    ...

setup()

Optional setup hook (load weights, initialize GPU, etc.).

Source code in src/openlithohub/models/base.py 
119
120
def setup(self) -> None:
    """Optional setup hook (load weights, initialize GPU, etc.)."""

teardown()

Optional cleanup hook.

Source code in src/openlithohub/models/base.py 
122
123
def teardown(self) -> None:
    """Optional cleanup hook."""

to_torch_module()

Return a single nn.Module that maps an input mask/design tensor to an output mask tensor, suitable for ONNX / TorchScript export.

Default raises NotImplementedError. Override this on models that wrap a static nn.Module (e.g. Neural-ILT). Iterative optimizers like Level-Set ILT do not have a single forward graph and cannot be exported — those should keep the default.

The returned module must be in eval() mode and accept tensors of shape (B, 1, H, W) in [0, 1], returning the same shape.

Source code in src/openlithohub/models/base.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def to_torch_module(self) -> torch.nn.Module:
    """Return a single ``nn.Module`` that maps an input mask/design tensor
    to an output mask tensor, suitable for ONNX / TorchScript export.

    Default raises ``NotImplementedError``. Override this on models that
    wrap a static ``nn.Module`` (e.g. Neural-ILT). Iterative optimizers
    like Level-Set ILT do not have a single forward graph and cannot be
    exported — those should keep the default.

    The returned module must be in ``eval()`` mode and accept tensors
    of shape ``(B, 1, H, W)`` in ``[0, 1]``, returning the same shape.
    """
    raise NotImplementedError(
        f"Model {type(self).NAME!r} does not support export to a single "
        "nn.Module — override LithographyModel.to_torch_module() if it "
        "wraps a static forward graph."
    )

openlithohub.models.registry

Model registry — discover and instantiate lithography models.

ModelRegistry

Registry for discovering and instantiating lithography models.

Source code in src/openlithohub/models/registry.py 
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
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 ModelRegistry:
    """Registry for discovering and instantiating lithography models."""

    def __init__(self) -> None:
        self._models: dict[str, type[LithographyModel]] = {}

    def register(self, model_cls: type[LithographyModel]) -> type[LithographyModel]:
        """Register a model class. Can be used as a decorator.

        The model class must define ``NAME`` directly on itself (not inherit
        it). The registry reads ``vars(model_cls)`` so that a default ``NAME``
        on a future base class cannot cause every concrete subclass that
        forgets to override it to silently collide on the same key.

        Re-registering an identical class (same module + qualname) is a
        no-op — that happens when ``register_builtin_models`` is called
        multiple times. A *different* class with the same NAME emits a
        ``UserWarning`` and overrides; users who shadow a built-in by
        accident will see the warning instead of silent replacement.
        """
        name = vars(model_cls).get("NAME")
        if not isinstance(name, str) or not name:
            raise TypeError(
                f"Model {model_cls.__name__} must define a class-level "
                f"`NAME: ClassVar[str]` attribute on itself (not inherited) "
                f"to be registered."
            )
        existing = self._models.get(name)
        if existing is not None and existing is not model_cls:
            same_qualname = (
                existing.__module__ == model_cls.__module__
                and existing.__qualname__ == model_cls.__qualname__
            )
            if not same_qualname:
                warnings.warn(
                    f"Model NAME {name!r} is being re-registered: "
                    f"{existing.__module__}.{existing.__qualname__} -> "
                    f"{model_cls.__module__}.{model_cls.__qualname__}. "
                    f"The previous registration will be replaced.",
                    UserWarning,
                    stacklevel=2,
                )
        self._models[name] = model_cls
        return model_cls

    def get(self, name: str, **kwargs: Any) -> LithographyModel:
        """Instantiate a registered model by name.

        Kwargs that the target model's ``__init__`` does not accept are
        silently dropped, so optional CLI flags like ``--pretrained`` work
        across the whole registry without each call site needing to know
        which models support which options. Real bugs in the model's
        ``__init__`` (mistyped args, missing required positionals) still
        propagate as ``TypeError``.
        """
        if name not in self._models:
            available = ", ".join(sorted(self._models.keys()))
            raise KeyError(f"Model '{name}' not found. Available: [{available}]")
        cls = self._models[name]
        return cls(**_filter_supported_kwargs(cls, kwargs))

    def supports_kwargs(self, name: str, kwargs: dict[str, Any]) -> dict[str, bool]:
        """Return a per-key flag indicating whether the named model accepts each kwarg."""
        if name not in self._models:
            raise KeyError(f"Model '{name}' not found.")
        cls = self._models[name]
        accepted = _accepted_kwargs(cls)
        if accepted is None:
            return {k: True for k in kwargs}
        return {k: k in accepted for k in kwargs}

    def list_models(self) -> list[str]:
        """Return names of all registered models."""
        return sorted(self._models.keys())

register(model_cls)

Register a model class. Can be used as a decorator.

The model class must define NAME directly on itself (not inherit it). The registry reads vars(model_cls) so that a default NAME on a future base class cannot cause every concrete subclass that forgets to override it to silently collide on the same key.

Re-registering an identical class (same module + qualname) is a no-op — that happens when register_builtin_models is called multiple times. A different class with the same NAME emits a UserWarning and overrides; users who shadow a built-in by accident will see the warning instead of silent replacement.

Source code in src/openlithohub/models/registry.py 
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
def register(self, model_cls: type[LithographyModel]) -> type[LithographyModel]:
    """Register a model class. Can be used as a decorator.

    The model class must define ``NAME`` directly on itself (not inherit
    it). The registry reads ``vars(model_cls)`` so that a default ``NAME``
    on a future base class cannot cause every concrete subclass that
    forgets to override it to silently collide on the same key.

    Re-registering an identical class (same module + qualname) is a
    no-op — that happens when ``register_builtin_models`` is called
    multiple times. A *different* class with the same NAME emits a
    ``UserWarning`` and overrides; users who shadow a built-in by
    accident will see the warning instead of silent replacement.
    """
    name = vars(model_cls).get("NAME")
    if not isinstance(name, str) or not name:
        raise TypeError(
            f"Model {model_cls.__name__} must define a class-level "
            f"`NAME: ClassVar[str]` attribute on itself (not inherited) "
            f"to be registered."
        )
    existing = self._models.get(name)
    if existing is not None and existing is not model_cls:
        same_qualname = (
            existing.__module__ == model_cls.__module__
            and existing.__qualname__ == model_cls.__qualname__
        )
        if not same_qualname:
            warnings.warn(
                f"Model NAME {name!r} is being re-registered: "
                f"{existing.__module__}.{existing.__qualname__} -> "
                f"{model_cls.__module__}.{model_cls.__qualname__}. "
                f"The previous registration will be replaced.",
                UserWarning,
                stacklevel=2,
            )
    self._models[name] = model_cls
    return model_cls

get(name, **kwargs)

Instantiate a registered model by name.

Kwargs that the target model's __init__ does not accept are silently dropped, so optional CLI flags like --pretrained work across the whole registry without each call site needing to know which models support which options. Real bugs in the model's __init__ (mistyped args, missing required positionals) still propagate as TypeError.

Source code in src/openlithohub/models/registry.py 
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def get(self, name: str, **kwargs: Any) -> LithographyModel:
    """Instantiate a registered model by name.

    Kwargs that the target model's ``__init__`` does not accept are
    silently dropped, so optional CLI flags like ``--pretrained`` work
    across the whole registry without each call site needing to know
    which models support which options. Real bugs in the model's
    ``__init__`` (mistyped args, missing required positionals) still
    propagate as ``TypeError``.
    """
    if name not in self._models:
        available = ", ".join(sorted(self._models.keys()))
        raise KeyError(f"Model '{name}' not found. Available: [{available}]")
    cls = self._models[name]
    return cls(**_filter_supported_kwargs(cls, kwargs))

supports_kwargs(name, kwargs)

Return a per-key flag indicating whether the named model accepts each kwarg.

Source code in src/openlithohub/models/registry.py 
73
74
75
76
77
78
79
80
81
def supports_kwargs(self, name: str, kwargs: dict[str, Any]) -> dict[str, bool]:
    """Return a per-key flag indicating whether the named model accepts each kwarg."""
    if name not in self._models:
        raise KeyError(f"Model '{name}' not found.")
    cls = self._models[name]
    accepted = _accepted_kwargs(cls)
    if accepted is None:
        return {k: True for k in kwargs}
    return {k: k in accepted for k in kwargs}

list_models()

Return names of all registered models.

Source code in src/openlithohub/models/registry.py 
83
84
85
def list_models(self) -> list[str]:
    """Return names of all registered models."""
    return sorted(self._models.keys())

register_builtin_models()

Side-effect import the in-tree models so the registry is populated.

Idempotent — Python caches modules in sys.modules so repeated calls are cheap. Both the optimize CLI and the multiprocessing workers call this so workers populate their registry the same way the parent does.

Source code in src/openlithohub/models/registry.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def register_builtin_models() -> None:
    """Side-effect import the in-tree models so the registry is populated.

    Idempotent — Python caches modules in ``sys.modules`` so repeated calls
    are cheap. Both the optimize CLI and the multiprocessing workers call
    this so workers populate their registry the same way the parent does.
    """
    import openlithohub.models.examples.dummy_model  # noqa: F401
    import openlithohub.models.gan_opc  # noqa: F401
    import openlithohub.models.levelset_ilt  # noqa: F401
    import openlithohub.models.neural_ilt  # noqa: F401
    import openlithohub.models.openilt  # noqa: F401
    import openlithohub.models.rule_based_opc  # noqa: F401
    import openlithohub.models.surrogate_ilt  # noqa: F401
    import openlithohub.models.vae_ilt  # noqa: F401

openlithohub.models.levelset_ilt

LevelSet-ILT: Iterative mask optimization via gradient descent.

The level-set / continuous-mask formulation of Inverse Lithography Technology dates to Pang, Liu & Abrams, Inverse lithography technology principles in practice: unintuitive patterns (Proc. SPIE 5992, 2005) and Poonawala & Milanfar, Mask design for optical microlithography — an inverse imaging problem (IEEE TIP 16(3), 2007). This implementation follows the SimpleILT-style L2 + total-variation formulation surveyed in [Yang2023_LithoBench, §3.3, p.5] (open-access substitute for the paywalled Granik / Pang journal write-ups).

Confidence B — algorithmic intent is matched against the LithoBench narrative; specific hyperparameters here (lr, sigma_px, tv_weight) are this project's defaults, not literal paper values.

LevelSetILTModel

Bases: LithographyModel

Inverse Lithography Technology via level-set gradient descent.

Optimizes a continuous mask representation to minimize the difference between the simulated resist image and the target design pattern. Supports two forward models:

  • gaussian (default): a single Gaussian PSF — fast, used in tests.
  • hopkins: SOCS-truncated partial-coherence Hopkins imaging — physically faithful, suitable for end-to-end AI-OPC research.

When helmholtz_radius > 0, a differentiable Helmholtz PDE filter is applied to the continuous mask before the aerial image simulation, suppressing sub-resolution features and enforcing a minimum manufacturable feature size (~2*radius).

Source code in src/openlithohub/models/levelset_ilt.py 
 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
 86
 87
 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
132
133
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
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
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
@registry.register
class LevelSetILTModel(LithographyModel):
    """Inverse Lithography Technology via level-set gradient descent.

    Optimizes a continuous mask representation to minimize the difference
    between the simulated resist image and the target design pattern.
    Supports two forward models:

    - ``gaussian`` (default): a single Gaussian PSF — fast, used in tests.
    - ``hopkins``: SOCS-truncated partial-coherence Hopkins imaging — physically
      faithful, suitable for end-to-end AI-OPC research.

    When ``helmholtz_radius > 0``, a differentiable Helmholtz PDE filter
    is applied to the continuous mask before the aerial image simulation,
    suppressing sub-resolution features and enforcing a minimum
    manufacturable feature size (~2*radius).
    """

    NAME = "levelset-ilt"
    SUPPORTS_CURVILINEAR = True
    # Issue #75: see openilt.py — iterative ILT gradient flow propagates
    # across many OIR's worth of pixels, so 0 px under-shoots the seam-free
    # halo at tile boundaries. 64 px is the same conservative bound used by
    # the U-Net-based Neural-ILT / GAN-OPC.
    RECEPTIVE_FIELD_PX = 64

    def __init__(
        self,
        iterations: int = 200,
        lr: float = 0.1,
        sigma_px: float = 2.0,
        tv_weight: float = 0.01,
        dose: float = 1.0,
        resist_steepness: float = 50.0,
        resist_diffusion_nm: float = 0.0,
        quencher: float = 0.0,
        pixel_size_nm: float = 1.0,
        forward_model: ForwardModelKind = "gaussian",
        hopkins_params: HopkinsParams | None = None,
        helmholtz_radius: float = 0.0,
    ) -> None:
        self._iterations = iterations
        self._lr = lr
        self._sigma_px = sigma_px
        self._tv_weight = tv_weight
        self._dose = dose
        self._resist_steepness = resist_steepness
        self._resist_diffusion_nm = resist_diffusion_nm
        self._quencher = quencher
        self._pixel_size_nm = pixel_size_nm
        self._forward_model = forward_model
        self._hopkins_params = hopkins_params or HopkinsParams()
        self._helmholtz_radius = helmholtz_radius
        self._cached_kernels: torch.Tensor | None = None
        self._cached_weights: torch.Tensor | None = None
        self._cached_grid: int | None = None
        self._compiled_hopkins_cache: dict[tuple[Any, ...], Any] = {}
        # Guards the kernel cache so concurrent ``predict()`` callers (e.g. a
        # FastAPI handler under load) cannot race on the check-then-rebuild
        # path and leave the cache in a half-populated state. Mirrors the
        # pattern in ``OpenILTModel``.
        self._cache_lock = threading.Lock()

    def _ensure_hopkins_kernels(
        self, grid_size: int, device: torch.device
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # The lock serializes the check-then-rebuild so two concurrent
        # callers cannot both observe a stale-or-empty cache and race on
        # the write.
        with self._cache_lock:
            if (
                self._cached_kernels is None
                or self._cached_weights is None
                or self._cached_grid != grid_size
                or self._cached_kernels.device != device
            ):
                kernels, weights = compute_socs_kernels(self._hopkins_params, grid_size, device)
                self._cached_kernels = kernels
                self._cached_weights = weights
                self._cached_grid = grid_size
            assert self._cached_kernels is not None
            assert self._cached_weights is not None
            return self._cached_kernels, self._cached_weights

    def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
        """Optimize a mask to reproduce the target design under lithography simulation.

        Args:
            design: Target design pattern (H, W), binary.
            **kwargs: Optional overrides — iterations, lr, sigma_px, tv_weight,
                forward_model, hopkins_params, helmholtz_radius, device, dtype,
                compile_forward, process_window, pw_corners, checkpoint_dir,
                save_freq, resume_from.

                ``process_window=True`` swaps the nominal-only fidelity loss
                for ``workflow.process_window.pw_fidelity_loss`` evaluated
                across ``pw_corners`` (defaults to
                ``workflow.process_window.DEFAULT_PW_CORNERS``). Currently only
                supported on the ``gaussian`` forward model — pass
                ``forward_model="gaussian"`` (the default) when enabling PW.

                ``checkpoint_dir`` (``str | Path | None``) + ``save_freq``
                (int, default 0 = off) write the running ``mask_logit`` and
                Adam state to ``{checkpoint_dir}/ckpt_iter{N}.pt`` every
                ``save_freq`` iterations. ``resume_from`` (``str | Path``)
                loads such a file and continues from its iteration counter,
                so a SLURM preemption or CUDA crash does not erase prior
                progress.

                ``helmholtz_radius`` (float, default ``0.0`` = disabled)
                applies a differentiable Helmholtz PDE filter to
                ``mask_continuous`` before the aerial image simulation.
                When > 0, features smaller than ~2*radius are suppressed,
                enforcing a minimum manufacturable feature size.  The
                filter is fully differentiable — gradients flow through to
                ``mask_logit`` without approximation.
        """
        target = design.detach().float()
        if target.ndim > 2:
            target = target.squeeze()

        iterations = kwargs.get("iterations", self._iterations)
        lr = kwargs.get("lr", self._lr)
        sigma_px = kwargs.get("sigma_px", self._sigma_px)
        tv_weight = kwargs.get("tv_weight", self._tv_weight)
        forward_model = kwargs.get("forward_model", self._forward_model)
        hopkins_params = kwargs.get("hopkins_params", self._hopkins_params)
        helmholtz_radius = kwargs.get("helmholtz_radius", self._helmholtz_radius)
        dtype = kwargs.get("dtype", torch.float32)
        compile_forward = kwargs.get("compile_forward", False)
        device = kwargs.get("device")
        process_window = kwargs.get("process_window", False)
        pw_corners = kwargs.get("pw_corners", DEFAULT_PW_CORNERS)
        checkpoint_dir = kwargs.get("checkpoint_dir")
        save_freq = int(kwargs.get("save_freq", 0))
        resume_from = kwargs.get("resume_from")
        if save_freq < 0:
            raise ValueError(f"save_freq must be >= 0, got {save_freq}")
        if save_freq > 0 and checkpoint_dir is None:
            raise ValueError("save_freq > 0 requires checkpoint_dir to be set.")
        ckpt_dir_path: Path | None = None
        if checkpoint_dir is not None:
            ckpt_dir_path = Path(checkpoint_dir)
            ckpt_dir_path.mkdir(parents=True, exist_ok=True)
        if process_window and forward_model != "gaussian":
            raise ValueError(
                "process_window=True currently only supports forward_model='gaussian'; "
                "Hopkins corner sweeps will land in a follow-up. Set "
                "forward_model='gaussian' or process_window=False."
            )
        if device is not None:
            target = target.to(device)

        if forward_model == "hopkins":
            if hopkins_params != self._hopkins_params:
                with self._cache_lock:
                    if hopkins_params != self._hopkins_params:
                        self._hopkins_params = hopkins_params
                        self._cached_kernels = None
                        self._cached_weights = None
                        self._cached_grid = None
            kernels, weights = self._ensure_hopkins_kernels(target.shape[0], target.device)
            hopkins_fn: Callable[..., torch.Tensor] | None = simulate_aerial_image_hopkins
            if compile_forward:
                cache_key = (
                    target.shape[0],
                    target.shape[-1],
                    str(target.device),
                    str(dtype),
                    forward_model,
                )
                compiled = self._compiled_hopkins_cache.get(cache_key)
                if compiled is None:
                    # Lift dynamo's per-function recompile ceiling so a run that
                    # legitimately encounters several tile sizes (e.g. ORFS with
                    # mixed-pitch designs) keeps cache hits instead of evicting.
                    try:
                        import torch._dynamo as _dynamo

                        _dynamo.config.cache_size_limit = max(_dynamo.config.cache_size_limit, 64)
                    except Exception:  # noqa: BLE001, S110 — best-effort tuning; old PyTorch lacks the symbol
                        pass
                    try:
                        compiled = torch.compile(hopkins_fn, mode="reduce-overhead", dynamic=False)
                    except Exception:
                        # torch.compile may be unavailable (Windows without Triton,
                        # some MPS configs, torch < 2.0). Falling back to eager
                        # keeps the run alive — caller can pass --no-compile to
                        # silence the attempt.
                        compiled = hopkins_fn
                    self._compiled_hopkins_cache[cache_key] = compiled
                hopkins_fn = compiled  # type: ignore[assignment]
        else:
            kernels = None
            weights = None
            hopkins_fn = None

        mask_logit = torch.zeros_like(target, requires_grad=True)
        with torch.no_grad():
            mask_logit.copy_(target * 4.0 - 2.0)
        mask_logit = mask_logit.clone().detach().requires_grad_(True)

        optimizer = torch.optim.Adam([mask_logit], lr=lr)

        best_loss = float("inf")
        best_mask: torch.Tensor = target.clone()
        start_iter = 0

        if resume_from is not None:
            # ``weights_only=False`` because we serialise an Adam state_dict
            # alongside the tensor — pickled Python dicts that torch refuses
            # to load under the 2.6+ default. Checkpoints are written by
            # this code path only, so the trust boundary is the user's own
            # filesystem.
            state = torch.load(  # nosec B614 — trust boundary is the user's own filesystem; checkpoint written by this code path
                str(resume_from), map_location=target.device, weights_only=False
            )
            with torch.no_grad():
                mask_logit.copy_(state["mask_logit"])
            optimizer.load_state_dict(state["optimizer"])
            start_iter = int(state.get("iteration", 0))
            best_loss = float(state.get("best_loss", best_loss))
            if "best_mask" in state:
                best_mask = state["best_mask"].to(target.device)

        for it in range(start_iter, iterations):
            optimizer.zero_grad()

            mask_continuous = torch.sigmoid(mask_logit)
            if helmholtz_radius > 0.0:
                mask_continuous = apply_helmholtz_filter(
                    mask_continuous,
                    radius=helmholtz_radius,
                )
            if forward_model == "hopkins":
                assert hopkins_fn is not None  # narrowed by forward_model == "hopkins" branch
                aerial = hopkins_fn(
                    mask_continuous,
                    kernels=kernels,
                    weights=weights,
                    dose=self._dose,
                    dtype=dtype,
                )
            else:
                aerial = simulate_aerial_image(mask_continuous, sigma_px=sigma_px, dose=self._dose)
            if aerial.dtype != torch.float32:
                aerial = aerial.float()

            resist = apply_differentiable_resist(
                aerial,
                threshold=0.5,
                steepness=self._resist_steepness,
                resist_diffusion_nm=self._resist_diffusion_nm,
                pixel_size_nm=self._pixel_size_nm,
                quencher=self._quencher,
            )

            if process_window:
                fidelity_loss = pw_fidelity_loss(
                    mask_continuous,
                    target,
                    corners=pw_corners,
                    threshold=0.5,
                    steepness=self._resist_steepness,
                )
            else:
                fidelity_loss = torch.nn.functional.mse_loss(resist, target)
            tv_loss = _total_variation(mask_continuous)
            loss = fidelity_loss + tv_weight * tv_loss

            loss.backward()
            optimizer.step()

            loss_val = loss.item()
            if loss_val < best_loss:
                best_loss = loss_val
                best_mask = (mask_continuous > 0.5).float().detach()

            if save_freq > 0 and ckpt_dir_path is not None and (it + 1) % save_freq == 0:
                # ``it + 1`` is the count of *completed* steps, so a resume
                # from this file restarts iteration index at it+1 and runs
                # exactly ``iterations - (it+1)`` more steps — total work
                # equals an uninterrupted run.
                torch.save(
                    {
                        "iteration": it + 1,
                        "mask_logit": mask_logit.detach().clone(),
                        "optimizer": optimizer.state_dict(),
                        "best_loss": best_loss,
                        "best_mask": best_mask.detach().clone(),
                    },
                    str(ckpt_dir_path / f"ckpt_iter{it + 1}.pt"),
                )

        return PredictionResult(
            mask=best_mask,
            metadata={
                "final_loss": best_loss,
                "iterations": iterations,
                "sigma_px": sigma_px,
                "forward_model": forward_model,
                "helmholtz_radius": helmholtz_radius,
                "process_window": process_window,
                "pw_corner_count": len(pw_corners) if process_window else 0,
            },
        )

predict(design, **kwargs)

Optimize a mask to reproduce the target design under lithography simulation.

Parameters:

Name Type Description Default
design Tensor

Target design pattern (H, W), binary.

 required
**kwargs Any

Optional overrides — iterations, lr, sigma_px, tv_weight, forward_model, hopkins_params, helmholtz_radius, device, dtype, compile_forward, process_window, pw_corners, checkpoint_dir, save_freq, resume_from.

process_window=True swaps the nominal-only fidelity loss for workflow.process_window.pw_fidelity_loss evaluated across pw_corners (defaults to workflow.process_window.DEFAULT_PW_CORNERS). Currently only supported on the gaussian forward model — pass forward_model="gaussian" (the default) when enabling PW.

checkpoint_dir (str | Path | None) + save_freq (int, default 0 = off) write the running mask_logit and Adam state to {checkpoint_dir}/ckpt_iter{N}.pt every save_freq iterations. resume_from (str | Path) loads such a file and continues from its iteration counter, so a SLURM preemption or CUDA crash does not erase prior progress.

helmholtz_radius (float, default 0.0 = disabled) applies a differentiable Helmholtz PDE filter to mask_continuous before the aerial image simulation. When > 0, features smaller than ~2*radius are suppressed, enforcing a minimum manufacturable feature size. The filter is fully differentiable — gradients flow through to mask_logit without approximation.

 {}
Source code in src/openlithohub/models/levelset_ilt.py 
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
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
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
    """Optimize a mask to reproduce the target design under lithography simulation.

    Args:
        design: Target design pattern (H, W), binary.
        **kwargs: Optional overrides — iterations, lr, sigma_px, tv_weight,
            forward_model, hopkins_params, helmholtz_radius, device, dtype,
            compile_forward, process_window, pw_corners, checkpoint_dir,
            save_freq, resume_from.

            ``process_window=True`` swaps the nominal-only fidelity loss
            for ``workflow.process_window.pw_fidelity_loss`` evaluated
            across ``pw_corners`` (defaults to
            ``workflow.process_window.DEFAULT_PW_CORNERS``). Currently only
            supported on the ``gaussian`` forward model — pass
            ``forward_model="gaussian"`` (the default) when enabling PW.

            ``checkpoint_dir`` (``str | Path | None``) + ``save_freq``
            (int, default 0 = off) write the running ``mask_logit`` and
            Adam state to ``{checkpoint_dir}/ckpt_iter{N}.pt`` every
            ``save_freq`` iterations. ``resume_from`` (``str | Path``)
            loads such a file and continues from its iteration counter,
            so a SLURM preemption or CUDA crash does not erase prior
            progress.

            ``helmholtz_radius`` (float, default ``0.0`` = disabled)
            applies a differentiable Helmholtz PDE filter to
            ``mask_continuous`` before the aerial image simulation.
            When > 0, features smaller than ~2*radius are suppressed,
            enforcing a minimum manufacturable feature size.  The
            filter is fully differentiable — gradients flow through to
            ``mask_logit`` without approximation.
    """
    target = design.detach().float()
    if target.ndim > 2:
        target = target.squeeze()

    iterations = kwargs.get("iterations", self._iterations)
    lr = kwargs.get("lr", self._lr)
    sigma_px = kwargs.get("sigma_px", self._sigma_px)
    tv_weight = kwargs.get("tv_weight", self._tv_weight)
    forward_model = kwargs.get("forward_model", self._forward_model)
    hopkins_params = kwargs.get("hopkins_params", self._hopkins_params)
    helmholtz_radius = kwargs.get("helmholtz_radius", self._helmholtz_radius)
    dtype = kwargs.get("dtype", torch.float32)
    compile_forward = kwargs.get("compile_forward", False)
    device = kwargs.get("device")
    process_window = kwargs.get("process_window", False)
    pw_corners = kwargs.get("pw_corners", DEFAULT_PW_CORNERS)
    checkpoint_dir = kwargs.get("checkpoint_dir")
    save_freq = int(kwargs.get("save_freq", 0))
    resume_from = kwargs.get("resume_from")
    if save_freq < 0:
        raise ValueError(f"save_freq must be >= 0, got {save_freq}")
    if save_freq > 0 and checkpoint_dir is None:
        raise ValueError("save_freq > 0 requires checkpoint_dir to be set.")
    ckpt_dir_path: Path | None = None
    if checkpoint_dir is not None:
        ckpt_dir_path = Path(checkpoint_dir)
        ckpt_dir_path.mkdir(parents=True, exist_ok=True)
    if process_window and forward_model != "gaussian":
        raise ValueError(
            "process_window=True currently only supports forward_model='gaussian'; "
            "Hopkins corner sweeps will land in a follow-up. Set "
            "forward_model='gaussian' or process_window=False."
        )
    if device is not None:
        target = target.to(device)

    if forward_model == "hopkins":
        if hopkins_params != self._hopkins_params:
            with self._cache_lock:
                if hopkins_params != self._hopkins_params:
                    self._hopkins_params = hopkins_params
                    self._cached_kernels = None
                    self._cached_weights = None
                    self._cached_grid = None
        kernels, weights = self._ensure_hopkins_kernels(target.shape[0], target.device)
        hopkins_fn: Callable[..., torch.Tensor] | None = simulate_aerial_image_hopkins
        if compile_forward:
            cache_key = (
                target.shape[0],
                target.shape[-1],
                str(target.device),
                str(dtype),
                forward_model,
            )
            compiled = self._compiled_hopkins_cache.get(cache_key)
            if compiled is None:
                # Lift dynamo's per-function recompile ceiling so a run that
                # legitimately encounters several tile sizes (e.g. ORFS with
                # mixed-pitch designs) keeps cache hits instead of evicting.
                try:
                    import torch._dynamo as _dynamo

                    _dynamo.config.cache_size_limit = max(_dynamo.config.cache_size_limit, 64)
                except Exception:  # noqa: BLE001, S110 — best-effort tuning; old PyTorch lacks the symbol
                    pass
                try:
                    compiled = torch.compile(hopkins_fn, mode="reduce-overhead", dynamic=False)
                except Exception:
                    # torch.compile may be unavailable (Windows without Triton,
                    # some MPS configs, torch < 2.0). Falling back to eager
                    # keeps the run alive — caller can pass --no-compile to
                    # silence the attempt.
                    compiled = hopkins_fn
                self._compiled_hopkins_cache[cache_key] = compiled
            hopkins_fn = compiled  # type: ignore[assignment]
    else:
        kernels = None
        weights = None
        hopkins_fn = None

    mask_logit = torch.zeros_like(target, requires_grad=True)
    with torch.no_grad():
        mask_logit.copy_(target * 4.0 - 2.0)
    mask_logit = mask_logit.clone().detach().requires_grad_(True)

    optimizer = torch.optim.Adam([mask_logit], lr=lr)

    best_loss = float("inf")
    best_mask: torch.Tensor = target.clone()
    start_iter = 0

    if resume_from is not None:
        # ``weights_only=False`` because we serialise an Adam state_dict
        # alongside the tensor — pickled Python dicts that torch refuses
        # to load under the 2.6+ default. Checkpoints are written by
        # this code path only, so the trust boundary is the user's own
        # filesystem.
        state = torch.load(  # nosec B614 — trust boundary is the user's own filesystem; checkpoint written by this code path
            str(resume_from), map_location=target.device, weights_only=False
        )
        with torch.no_grad():
            mask_logit.copy_(state["mask_logit"])
        optimizer.load_state_dict(state["optimizer"])
        start_iter = int(state.get("iteration", 0))
        best_loss = float(state.get("best_loss", best_loss))
        if "best_mask" in state:
            best_mask = state["best_mask"].to(target.device)

    for it in range(start_iter, iterations):
        optimizer.zero_grad()

        mask_continuous = torch.sigmoid(mask_logit)
        if helmholtz_radius > 0.0:
            mask_continuous = apply_helmholtz_filter(
                mask_continuous,
                radius=helmholtz_radius,
            )
        if forward_model == "hopkins":
            assert hopkins_fn is not None  # narrowed by forward_model == "hopkins" branch
            aerial = hopkins_fn(
                mask_continuous,
                kernels=kernels,
                weights=weights,
                dose=self._dose,
                dtype=dtype,
            )
        else:
            aerial = simulate_aerial_image(mask_continuous, sigma_px=sigma_px, dose=self._dose)
        if aerial.dtype != torch.float32:
            aerial = aerial.float()

        resist = apply_differentiable_resist(
            aerial,
            threshold=0.5,
            steepness=self._resist_steepness,
            resist_diffusion_nm=self._resist_diffusion_nm,
            pixel_size_nm=self._pixel_size_nm,
            quencher=self._quencher,
        )

        if process_window:
            fidelity_loss = pw_fidelity_loss(
                mask_continuous,
                target,
                corners=pw_corners,
                threshold=0.5,
                steepness=self._resist_steepness,
            )
        else:
            fidelity_loss = torch.nn.functional.mse_loss(resist, target)
        tv_loss = _total_variation(mask_continuous)
        loss = fidelity_loss + tv_weight * tv_loss

        loss.backward()
        optimizer.step()

        loss_val = loss.item()
        if loss_val < best_loss:
            best_loss = loss_val
            best_mask = (mask_continuous > 0.5).float().detach()

        if save_freq > 0 and ckpt_dir_path is not None and (it + 1) % save_freq == 0:
            # ``it + 1`` is the count of *completed* steps, so a resume
            # from this file restarts iteration index at it+1 and runs
            # exactly ``iterations - (it+1)`` more steps — total work
            # equals an uninterrupted run.
            torch.save(
                {
                    "iteration": it + 1,
                    "mask_logit": mask_logit.detach().clone(),
                    "optimizer": optimizer.state_dict(),
                    "best_loss": best_loss,
                    "best_mask": best_mask.detach().clone(),
                },
                str(ckpt_dir_path / f"ckpt_iter{it + 1}.pt"),
            )

    return PredictionResult(
        mask=best_mask,
        metadata={
            "final_loss": best_loss,
            "iterations": iterations,
            "sigma_px": sigma_px,
            "forward_model": forward_model,
            "helmholtz_radius": helmholtz_radius,
            "process_window": process_window,
            "pw_corner_count": len(pw_corners) if process_window else 0,
        },
    )

openlithohub.models.neural_ilt

Neural-ILT: U-Net based mask prediction model.

NeuralILTModel

Bases: LithographyModel

U-Net mask predictor (Neural-ILT-style — see audit caveats below).

Predicts an optimised mask directly from the design layout in a single forward pass. Much faster than iterative methods at inference time, but requires pretrained weights for good results.

What this adapter is not. This is not a paper-faithful re-implementation of Jiang2020_NeuralILT (ICCAD'20). It is a U-Net mask predictor whose architecture is inspired by that paper. The headline contribution of Jiang2020 — the differentiable ILT correction layer that lets the L2/PVB loss flow back through a forward simulator at training time — is not implemented here. Eval-time forward simulation lives in the leaderboard scoring pipeline (see benchmark/metrics/l2_error.py, which already accepts simulator=); training-time loss-through-simulator is a separate piece of code that is not in this adapter and has no open issue or RFC pinned to it.

Lineage references (cite for architecture inspiration, not for expected metric): - Jiang2020_NeuralILT (ICCAD'20, paywalled) — original paper. - Yang2023_LithoBench §4.3 (NeurIPS 2023, open access) — restates the architecture and training schedule.

Architecture divergences vs. Jiang2020: - Backbone: 3-level U-Net (3 down / 3 up), paper uses 4 levels. - Channel widths: 32→64→128→256 (paper: 64→128→256→512). - Loss / training: not part of this adapter.

See docs/audits/neural-ilt-architecture.md for the full audit, the 2026-05-23 decision record (option (b): downgrade naming instead of implementing the correction layer), and re-audit triggers.

Source code in src/openlithohub/models/neural_ilt.py 
 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
 86
 87
 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
132
133
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
@registry.register
class NeuralILTModel(LithographyModel):
    """U-Net mask predictor (Neural-ILT-style — see audit caveats below).

    Predicts an optimised mask directly from the design layout in a single
    forward pass. Much faster than iterative methods at inference time,
    but requires pretrained weights for good results.

    **What this adapter is not.** This is *not* a paper-faithful
    re-implementation of ``Jiang2020_NeuralILT`` (ICCAD'20). It is a
    U-Net mask predictor whose architecture is *inspired by* that paper.
    The headline contribution of Jiang2020 — the differentiable ILT
    correction layer that lets the L2/PVB loss flow back through a
    forward simulator at training time — is **not implemented here**.
    Eval-time forward simulation lives in the leaderboard scoring
    pipeline (see ``benchmark/metrics/l2_error.py``, which already
    accepts ``simulator=``); training-time loss-through-simulator is a
    separate piece of code that is not in this adapter and has no
    open issue or RFC pinned to it.

    **Lineage references** (cite for *architecture inspiration*, not for
    *expected metric*):
    - ``Jiang2020_NeuralILT`` (ICCAD'20, paywalled) — original paper.
    - ``Yang2023_LithoBench`` §4.3 (NeurIPS 2023, open access) —
      restates the architecture and training schedule.

    **Architecture divergences vs. Jiang2020:**
    - Backbone: 3-level U-Net (3 down / 3 up), paper uses 4 levels.
    - Channel widths: 32→64→128→256 (paper: 64→128→256→512).
    - Loss / training: not part of this adapter.

    See ``docs/audits/neural-ilt-architecture.md`` for the full audit,
    the 2026-05-23 decision record (option (b): downgrade naming
    instead of implementing the correction layer), and re-audit
    triggers.
    """

    NAME = "neural-ilt"
    SUPPORTS_CURVILINEAR = True
    RECEPTIVE_FIELD_PX = 64

    def __init__(
        self,
        weights: str | Path | None = None,
        pretrained: bool = False,
        device: str = "cpu",
        repo_id: str = "openlithohub/neural-ilt-v0.1",
        repo_filename: str = "model.pt",
        url_sha256: str | None = None,
    ) -> None:
        self._weights_path = weights
        self._pretrained = pretrained
        self._device = device
        self._repo_id = repo_id
        self._repo_filename = repo_filename
        self._url_sha256 = url_sha256
        self._net: torch.nn.Module | None = None

    def setup(self) -> None:
        from openlithohub.models._unet import UNet

        self._net = UNet(in_channels=1, out_channels=1).to(self._device)

        weights_loaded = False
        if self._weights_path is not None:
            weights_path = Path(self._weights_path)
            if weights_path.exists():
                state_dict = torch.load(
                    str(weights_path), map_location=self._device, weights_only=True
                )
                self._net.load_state_dict(state_dict)
                weights_loaded = True
        elif self._pretrained:
            from openlithohub.models.hub import ModelHub

            hub = ModelHub()
            path = hub.download_weights(
                self._repo_id, filename=self._repo_filename, sha256=self._url_sha256
            )
            state_dict = torch.load(str(path), map_location=self._device, weights_only=True)
            self._net.load_state_dict(state_dict)
            weights_loaded = True

        if not weights_loaded:
            # BatchNorm in eval() mode with default running stats produces
            # near-arbitrary outputs that are then thresholded to a binary
            # mask. Surface this so users know predictions are meaningless.
            warnings.warn(
                "NeuralILTModel is running with random-initialized weights. "
                "Predictions are not meaningful. Pass --pretrained or --weights "
                "to load trained parameters.",
                UserWarning,
                stacklevel=2,
            )

        self._net.eval()

    def teardown(self) -> None:
        self._net = None

    def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
        if self._net is None:
            self.setup()
        assert self._net is not None

        inp = design.detach().float()
        original_ndim = inp.ndim
        if inp.ndim == 2:
            inp = inp.unsqueeze(0).unsqueeze(0)
        elif inp.ndim == 3:
            inp = inp.unsqueeze(0)
        elif inp.ndim == 4:
            pass  # already (B, C, H, W)
        else:
            raise ValueError(f"expected 2D–4D input, got {inp.ndim}D")

        inp = inp.to(self._device)

        with torch.no_grad():
            logits = self._net(inp)
            mask = torch.sigmoid(logits)

        # Preserve the original dimensionality contract: 2D/3D input → 2D output.
        if original_ndim <= 3:
            mask = mask.squeeze()

        mask_binary = (mask > 0.5).float()
        return PredictionResult(
            mask=mask_binary,
            contour=None,
            metadata={"logits_range": (logits.min().item(), logits.max().item())},
        )

    def to_torch_module(self) -> torch.nn.Module:
        """Return the U-Net wrapped so the export forward emits a sigmoid mask
        in ``[0, 1]`` directly — what a downstream production pipeline expects.
        """
        if self._net is None:
            self.setup()
        assert self._net is not None
        unet = self._net

        class _NeuralILTExportWrapper(torch.nn.Module):
            def __init__(self, net: torch.nn.Module) -> None:
                super().__init__()
                self.net = net

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                return torch.sigmoid(self.net(x))

        wrapper = _NeuralILTExportWrapper(unet).eval()
        return wrapper

to_torch_module()

Return the U-Net wrapped so the export forward emits a sigmoid mask in [0, 1] directly — what a downstream production pipeline expects.

Source code in src/openlithohub/models/neural_ilt.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def to_torch_module(self) -> torch.nn.Module:
    """Return the U-Net wrapped so the export forward emits a sigmoid mask
    in ``[0, 1]`` directly — what a downstream production pipeline expects.
    """
    if self._net is None:
        self.setup()
    assert self._net is not None
    unet = self._net

    class _NeuralILTExportWrapper(torch.nn.Module):
        def __init__(self, net: torch.nn.Module) -> None:
            super().__init__()
            self.net = net

        def forward(self, x: torch.Tensor) -> torch.Tensor:
            return torch.sigmoid(self.net(x))

    wrapper = _NeuralILTExportWrapper(unet).eval()
    return wrapper

openlithohub.models.rule_based_opc

Rule-based geometric OPC baseline.

Classic non-AI optical proximity correction by morphological edge biasing. Serves as the geometric reference point that any learning-based method must beat: shows what plain dilation buys you and what it leaves on the table (line-end pull-back, corner rounding, MRC violations on tight pitch).

This implementation deliberately stays in the "no learning, no simulation" regime. It augments uniform dilation with three context-aware bias modes that are still pure geometry — directional hammerheads at line ends, serifs at concave inner corners, and an iso/dense bias split — plus an MRC self-check that reports (and optionally retreats from) violations caused by the bias itself. Everything runs on torch.nn.functional pooling; no scipy, no distance transforms.

Future work (intentionally out of scope here, to keep this a "minimal geometric baseline"): - SRAF placement (sub-resolution assist features) — needs distance transforms and per-process-node assist rules. - Segment-based fragmentation — break edges into segments and move each independently, the standard production OPC primitive. - Iterative bias loop (model-based OPC) — dilate → simulate → measure EPE → adjust bias. Crossing this line stops being "rule-based" by definition. - Per-layer rule tables — different metals / poly / via layers carry different bias tables in real flows.

RuleBasedOPCModel

Bases: LithographyModel

Geometric edge-bias OPC with directional hammerheads and inner-corner serifs.

Default behaviour is conservative: uniform 1px dilation plus directional line-end extension. Inner-corner serifs, iso/dense split, and MRC retreat are all opt-in via constructor or per-call kwargs.

Source code in src/openlithohub/models/rule_based_opc.py 
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
@registry.register
class RuleBasedOPCModel(LithographyModel):
    """Geometric edge-bias OPC with directional hammerheads and inner-corner serifs.

    Default behaviour is conservative: uniform 1px dilation plus directional
    line-end extension. Inner-corner serifs, iso/dense split, and MRC retreat
    are all opt-in via constructor or per-call kwargs.
    """

    NAME = "rule-based-opc"
    SUPPORTS_CURVILINEAR = False
    RECEPTIVE_FIELD_PX = 16

    def __init__(
        self,
        bias_radius_px: int = 1,
        line_end_extra_px: int = 1,
        inner_corner_extra_px: int = 0,
        directional_line_end: bool = True,
        iso_radius_px: int | None = None,
        dense_radius_px: int | None = None,
        density_window_px: int = 9,
        density_threshold: float = 0.25,
        mrc_min_space_px: int = 0,
    ) -> None:
        self._bias_radius_px = bias_radius_px
        self._line_end_extra_px = line_end_extra_px
        self._inner_corner_extra_px = inner_corner_extra_px
        self._directional_line_end = directional_line_end
        self._iso_radius_px = iso_radius_px
        self._dense_radius_px = dense_radius_px
        self._density_window_px = density_window_px
        self._density_threshold = density_threshold
        self._mrc_min_space_px = mrc_min_space_px

    def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
        bias_radius = int(kwargs.get("bias_radius_px", self._bias_radius_px))
        line_end_extra = int(kwargs.get("line_end_extra_px", self._line_end_extra_px))
        inner_corner_extra = int(kwargs.get("inner_corner_extra_px", self._inner_corner_extra_px))
        directional = bool(kwargs.get("directional_line_end", self._directional_line_end))
        iso_radius = kwargs.get("iso_radius_px", self._iso_radius_px)
        dense_radius = kwargs.get("dense_radius_px", self._dense_radius_px)
        density_window = int(kwargs.get("density_window_px", self._density_window_px))
        density_threshold = float(kwargs.get("density_threshold", self._density_threshold))
        mrc_min_space = int(kwargs.get("mrc_min_space_px", self._mrc_min_space_px))

        bias_radius_nm = kwargs.get("bias_radius_nm")
        pixel_size_nm = kwargs.get("pixel_size_nm")
        if bias_radius_nm is not None or pixel_size_nm is not None:
            if bias_radius_nm is None or pixel_size_nm is None:
                raise ValueError("bias_radius_nm and pixel_size_nm must be provided together")
            if pixel_size_nm <= 0:
                raise ValueError("pixel_size_nm must be positive")
            bias_radius = int(round(float(bias_radius_nm) / float(pixel_size_nm)))

        target = design.float()
        if target.ndim > 2:
            target = target.squeeze()

        if iso_radius is not None and dense_radius is not None:
            mask = self._iso_dense_dilate(
                target,
                int(iso_radius),
                int(dense_radius),
                density_window,
                density_threshold,
            )
        else:
            mask = _binary_dilate(target, bias_radius)

        n_line_end_tips = 0
        if line_end_extra > 0 and target.sum() > 0:
            if directional:
                up_tips, down_tips, left_tips, right_tips = _directional_line_end(target)
                tip_total = up_tips + down_tips + left_tips + right_tips
                n_line_end_tips = int((tip_total > 0).sum().item())
                if n_line_end_tips > 0:
                    extension = torch.maximum(
                        torch.maximum(
                            _directional_dilate(up_tips, line_end_extra, "up"),
                            _directional_dilate(down_tips, line_end_extra, "down"),
                        ),
                        torch.maximum(
                            _directional_dilate(left_tips, line_end_extra, "left"),
                            _directional_dilate(right_tips, line_end_extra, "right"),
                        ),
                    )
                    mask = torch.maximum(mask, extension)
            else:
                tips = _line_end_mask(target)
                n_line_end_tips = int(tips.sum().item())
                if n_line_end_tips > 0:
                    mask = torch.maximum(mask, _binary_dilate(tips, line_end_extra))

        n_inner_corners = 0
        if inner_corner_extra > 0 and target.sum() > 0:
            corners = _inner_corner_mask(target)
            n_inner_corners = int(corners.sum().item())
            if n_inner_corners > 0:
                mask = torch.maximum(mask, _binary_dilate(corners, inner_corner_extra))

        binary = (mask > 0.5).float()

        min_space = _min_space_px(binary)
        min_width = _min_width_px(binary)
        mrc_violated = mrc_min_space > 0 and 0 < min_space < mrc_min_space

        if mrc_violated:
            binary = self._retreat_until_mrc_clean(binary, mrc_min_space)
            min_space = _min_space_px(binary)
            min_width = _min_width_px(binary)
            mrc_violated = 0 < min_space < mrc_min_space

        design_sum = float(target.sum().item())
        mask_area_growth = float(binary.sum().item()) / design_sum if design_sum > 0 else 0.0
        bias_radius_nm_meta = float(bias_radius_nm) if bias_radius_nm is not None else None

        return PredictionResult(
            mask=binary,
            metadata={
                "bias_radius_px": bias_radius,
                "line_end_extra_px": line_end_extra,
                "inner_corner_extra_px": inner_corner_extra,
                "directional_line_end": directional,
                "iso_radius_px": iso_radius,
                "dense_radius_px": dense_radius,
                "bias_radius_nm": bias_radius_nm_meta,
                "mask_area_growth": mask_area_growth,
                "n_line_end_tips": n_line_end_tips,
                "n_inner_corners": n_inner_corners,
                "min_space_px": min_space,
                "min_width_px": min_width,
                "mrc_min_space_px": mrc_min_space,
                "mrc_violated": mrc_violated,
                "output_geometry": "manhattan",
            },
        )

    def _iso_dense_dilate(
        self,
        target: torch.Tensor,
        iso_radius: int,
        dense_radius: int,
        window_px: int,
        threshold: float,
    ) -> torch.Tensor:
        density = _compute_density(target, window_px)
        is_iso = density < threshold
        iso_mask = _binary_dilate(target, iso_radius)
        dense_mask = _binary_dilate(target, dense_radius)
        combined = torch.where(is_iso, iso_mask, dense_mask)
        return combined

    def _retreat_until_mrc_clean(
        self, mask: torch.Tensor, min_space_px: int, max_iters: int = 8
    ) -> torch.Tensor:
        current = mask
        for _ in range(max_iters):
            space = _min_space_px(current)
            if space == 0 or space >= min_space_px:
                return current
            current = _erode(current, 1)
            if current.sum().item() == 0:
                return current
        return current

openlithohub.models.hub

Model hub for downloading and caching pretrained weights.

Remote downloads (direct HTTPS URLs) require a known SHA256 — torch.load on attacker-controlled bytes is a known RCE vector even with weights_only=True. The HuggingFace path relies on the Hub's own content-addressed verification.

ChecksumMismatchError

Bases: RuntimeError

Raised when a downloaded file's SHA256 does not match the expected value.

Source code in src/openlithohub/models/hub.py 
29
30
class ChecksumMismatchError(RuntimeError):
    """Raised when a downloaded file's SHA256 does not match the expected value."""

ModelHub

Manages download and caching of pretrained model weights.

Supports HuggingFace Hub (if installed) and direct URL downloads. Direct URL downloads MUST come with a SHA256 checksum.

Cache-key contract

Three identifier shapes flow through this class. Knowing which shape goes where prevents the kind of round-trip bug fixed in the May 2026 review:

  • owner/repo — the public form a caller passes to :meth:download_weights for a HuggingFace-style model.
  • owner--repo — the on-disk form, with the path separator rewritten to a double-dash so the segment is filesystem-safe. :meth:list_cached decodes this back to owner/repo.
  • url--<hex> — the on-disk form for direct-URL downloads, where <hex> is the first 32 hex chars of sha256(url.encode()). :meth:list_cached returns this verbatim (there is no public owner/repo-shaped name to decode back to), and :meth:clear_cache accepts both the original URL and the verbatim url--<hex> segment so the list_cached → clear_cache pipeline composes cleanly.

Every caller-supplied identifier passes through _safe_cache_segment (or the URL hash) before it touches the filesystem; .., embedded slashes, NUL bytes, and absolute paths are rejected. URL-keyed segments accept only hex suffixes.

Source code in src/openlithohub/models/hub.py 
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
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
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
class ModelHub:
    """Manages download and caching of pretrained model weights.

    Supports HuggingFace Hub (if installed) and direct URL downloads.
    Direct URL downloads MUST come with a SHA256 checksum.

    Cache-key contract
    ------------------
    Three identifier shapes flow through this class. Knowing which
    shape goes where prevents the kind of round-trip bug fixed in the
    May 2026 review:

    - ``owner/repo`` — the **public** form a caller passes to
      :meth:`download_weights` for a HuggingFace-style model.
    - ``owner--repo`` — the **on-disk** form, with the path separator
      rewritten to a double-dash so the segment is filesystem-safe.
      :meth:`list_cached` decodes this back to ``owner/repo``.
    - ``url--<hex>`` — the **on-disk** form for direct-URL downloads,
      where ``<hex>`` is the first 32 hex chars of
      ``sha256(url.encode())``. :meth:`list_cached` returns this
      verbatim (there is no public ``owner/repo``-shaped name to decode
      back to), and :meth:`clear_cache` accepts both the original URL
      and the verbatim ``url--<hex>`` segment so the
      ``list_cached → clear_cache`` pipeline composes cleanly.

    Every caller-supplied identifier passes through
    ``_safe_cache_segment`` (or the URL hash) before it touches the
    filesystem; ``..``, embedded slashes, NUL bytes, and absolute
    paths are rejected. URL-keyed segments accept only hex suffixes.
    """

    def __init__(self, cache_dir: Path | None = None, timeout: float = 300.0) -> None:
        self.cache_dir = cache_dir or _DEFAULT_CACHE_DIR
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.timeout = timeout

    def download_weights(
        self,
        model_id: str,
        filename: str = "model.pt",
        revision: str = _DEFAULT_REVISION,
        sha256: str | None = None,
    ) -> Path:
        """Download model weights, returning the cached file path.

        Args:
            model_id: A HuggingFace repo ID (``owner/repo``) or an HTTPS URL.
            filename: File name within the repo (HF Hub only).
            revision: Git revision (HF Hub only). Pinning to a commit hash
                (40-hex) makes the download exactly reproducible — the
                default ``"main"`` is mutable and a publisher push can
                change the bytes you receive.
            sha256: Hex digest of the expected file contents. **Verified
                on both paths** — direct HTTPS downloads and HF Hub
                downloads. Issue #20: previously the HF Hub branch
                silently ignored this argument, so a malicious or
                accidentally-mutated repo at the configured revision
                would have produced bytes whose hash did not match what
                the caller pinned. ``None`` skips verification (HF
                Hub-only — the URL branch still requires it).
        """
        if model_id.startswith("http://") or model_id.startswith("https://"):
            cache_segment = "url--" + hashlib.sha256(model_id.encode("utf-8")).hexdigest()[:32]
        else:
            cache_segment = _safe_cache_segment(model_id, kind="model_id")
        safe_filename = _safe_cache_segment(filename, kind="filename")
        cached_path = self.cache_dir / cache_segment / safe_filename
        if cached_path.exists():
            if sha256 is not None and self.get_checksum(cached_path) != sha256.lower():
                raise ChecksumMismatchError(
                    f"Cached file {cached_path} does not match expected SHA256."
                )
            return cached_path

        if model_id.startswith("http://") or model_id.startswith("https://"):
            if sha256 is None:
                raise ValueError(
                    "Direct URL downloads require a sha256= argument. "
                    "Pass the expected hex digest of the weights file."
                )
            return self._download_url(model_id, cached_path, sha256)

        try:
            hf_path = self._download_hf_hub(model_id, filename, revision)
        except ImportError:
            raise ImportError(
                f"Cannot download model '{model_id}': huggingface_hub not installed. "
                "Install with: pip install openlithohub[models]"
            ) from None
        if sha256 is not None:
            actual = self.get_checksum(hf_path)
            if actual != sha256.lower():
                raise ChecksumMismatchError(
                    f"SHA256 mismatch for HF Hub file {model_id}/{filename} "
                    f"@ {revision}: expected {sha256.lower()}, got {actual}. "
                    f"The file at this revision differs from the digest you "
                    f"pinned — refuse to load."
                )
        elif revision in {"main", "master"} or revision is None:
            # Mutable revision + no sha256 = the bytes you receive today
            # may differ from the bytes you receive tomorrow. Warn so
            # downstream reproducibility claims aren't silently false.
            warnings.warn(
                f"download_weights({model_id!r}, revision={revision!r}) "
                f"with no sha256= and a mutable revision: the publisher "
                f"can change the bytes at any time. Pass either a "
                f"40-hex commit hash as `revision=` or a `sha256=` "
                f"digest for reproducibility.",
                UserWarning,
                stacklevel=2,
            )
        return hf_path

    def _download_hf_hub(self, repo_id: str, filename: str, revision: str) -> Path:
        """Download from HuggingFace Hub."""
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            revision=revision,
            cache_dir=str(self.cache_dir),
        )
        return Path(path)

    def _download_url(self, url: str, target: Path, sha256: str) -> Path:
        """Download from a direct URL with timeout, size limit, and SHA256 verification.

        Resolves the host once, refuses private/non-routable addresses, and
        then forces the TLS connection to that exact IP. This closes the DNS
        TOCTOU window that a plain ``urlopen(url)`` leaves open: a rebinder
        cannot return a public IP for the vetting query and a private IP for
        the actual fetch, because the fetch never re-resolves.
        """
        if not url.startswith("https://"):
            raise ValueError("Only HTTPS URLs are supported for model downloads")

        parsed = urllib.parse.urlparse(url)
        host = parsed.hostname
        if not host:
            raise ValueError(f"URL has no host component: {url}")
        port = parsed.port or 443
        # Build the request target by joining the (already URL-encoded) path
        # and query verbatim — re-quoting either part would corrupt
        # already-encoded characters. ``parsed.path`` is empty for URLs like
        # ``https://host?x=1``; default to ``/`` per HTTP/1.1 spec.
        request_target = parsed.path or "/"
        if parsed.query:
            request_target = f"{request_target}?{parsed.query}"
        # Preserve the port in the Host header for non-default ports — some
        # virtual-host-aware origins use it for routing.
        host_header = host if port == 443 else f"{host}:{port}"
        ips = _resolve_and_vet(host)

        target.parent.mkdir(parents=True, exist_ok=True)
        max_size = 2 * 1024 * 1024 * 1024  # 2 GB

        ctx = ssl.create_default_context()
        # Try each vetted IP in order — falls back from a broken IPv6 to a
        # working IPv4 record. SNI/host header still use the original hostname
        # so cert validation works.
        conn: _PinnedHTTPSConnection | None = None
        last_err: Exception | None = None
        for ip in ips:
            candidate = _PinnedHTTPSConnection(
                host, ip, port=port, timeout=self.timeout, context=ctx
            )
            try:
                candidate.connect()
            except OSError as exc:
                candidate.close()
                last_err = exc
                continue
            conn = candidate
            break
        if conn is None:
            raise ValueError(f"Could not connect to any vetted address for {host!r}: {last_err}")
        try:
            conn.request("GET", request_target, headers={"Host": host_header})
            response = conn.getresponse()
            # Reject redirects — we vetted only the original host, and a
            # 30x to a different host would leak the SSRF guard.
            if response.status in (301, 302, 303, 307, 308):
                raise ValueError(
                    f"Refusing to follow redirect from {url} to {response.getheader('Location')!r}"
                )
            if response.status != 200:
                raise ValueError(f"Unexpected HTTP status {response.status} fetching {url}")
            content_length = response.getheader("Content-Length")
            if content_length and int(content_length) > max_size:
                raise ValueError(
                    f"File size {int(content_length)} bytes exceeds limit of {max_size} bytes"
                )
            sha = hashlib.sha256()
            downloaded = 0
            with open(target, "wb") as f:
                while chunk := response.read(8192):
                    downloaded += len(chunk)
                    if downloaded > max_size:
                        target.unlink(missing_ok=True)
                        raise ValueError(f"Download exceeded size limit of {max_size} bytes")
                    sha.update(chunk)
                    f.write(chunk)
        finally:
            conn.close()
        actual = sha.hexdigest()
        if actual != sha256.lower():
            target.unlink(missing_ok=True)
            raise ChecksumMismatchError(
                f"SHA256 mismatch for {url}: expected {sha256.lower()}, got {actual}"
            )
        return target

    def list_cached(self) -> list[str]:
        """List model IDs that have cached weights.

        HF-Hub-style entries are decoded from the on-disk ``owner--repo`` form
        back to ``owner/repo``. URL-keyed entries (stored under
        ``url--<sha256>``) are returned verbatim so a caller can hand the
        result straight back to :meth:`clear_cache` without re-keying.
        """
        if not self.cache_dir.exists():
            return []
        names: list[str] = []
        for d in self.cache_dir.iterdir():
            if not d.is_dir():
                continue
            if d.name.startswith("url--"):
                names.append(d.name)
            else:
                names.append(d.name.replace("--", "/"))
        return sorted(names)

    def clear_cache(self, model_id: str | None = None) -> None:
        """Remove cached weights for a model (or all if model_id is None).

        ``model_id`` is routed through the same per-segment validator used by
        ``download_weights`` so a caller-supplied ``..`` cannot escape the
        cache directory and rmtree something it shouldn't. URL-keyed entries
        (cached under the ``url--<sha256>`` segment that ``download_weights``
        creates) are also accepted.
        """
        import shutil

        if model_id is None:
            if self.cache_dir.exists():
                shutil.rmtree(self.cache_dir)
                self.cache_dir.mkdir(parents=True, exist_ok=True)
            return

        if model_id.startswith("http://") or model_id.startswith("https://"):
            cache_segment = "url--" + hashlib.sha256(model_id.encode("utf-8")).hexdigest()[:32]
        elif model_id.startswith("url--"):
            # The exact on-disk segment as returned by `list_cached` for
            # URL-keyed entries. Validate the suffix is a clean hex segment
            # so a caller can't sneak `..` past us via this branch.
            suffix = model_id[len("url--") :]
            if not suffix or any(c not in "0123456789abcdef" for c in suffix):
                raise ValueError(f"Refusing unsafe url-keyed cache id: {model_id!r}")
            cache_segment = model_id
        else:
            cache_segment = _safe_cache_segment(model_id, kind="model_id")
        model_dir = self.cache_dir / cache_segment
        if model_dir.exists():
            shutil.rmtree(model_dir)

    def get_checksum(self, path: Path) -> str:
        """Compute SHA256 checksum of a file."""
        sha = hashlib.sha256()
        with open(path, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                sha.update(chunk)
        return sha.hexdigest()

download_weights(model_id, filename='model.pt', revision=_DEFAULT_REVISION, sha256=None)

Download model weights, returning the cached file path.

Parameters:

Name Type Description Default
model_id str

A HuggingFace repo ID (owner/repo) or an HTTPS URL.

 required
filename str

File name within the repo (HF Hub only).

 'model.pt'
revision str

Git revision (HF Hub only). Pinning to a commit hash (40-hex) makes the download exactly reproducible — the default "main" is mutable and a publisher push can change the bytes you receive.

 _DEFAULT_REVISION
sha256 str | None

Hex digest of the expected file contents. Verified on both paths — direct HTTPS downloads and HF Hub downloads. Issue #20: previously the HF Hub branch silently ignored this argument, so a malicious or accidentally-mutated repo at the configured revision would have produced bytes whose hash did not match what the caller pinned. None skips verification (HF Hub-only — the URL branch still requires it).

 None
Source code in src/openlithohub/models/hub.py 
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
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
236
237
238
239
240
241
242
243
244
245
246
247
248
def download_weights(
    self,
    model_id: str,
    filename: str = "model.pt",
    revision: str = _DEFAULT_REVISION,
    sha256: str | None = None,
) -> Path:
    """Download model weights, returning the cached file path.

    Args:
        model_id: A HuggingFace repo ID (``owner/repo``) or an HTTPS URL.
        filename: File name within the repo (HF Hub only).
        revision: Git revision (HF Hub only). Pinning to a commit hash
            (40-hex) makes the download exactly reproducible — the
            default ``"main"`` is mutable and a publisher push can
            change the bytes you receive.
        sha256: Hex digest of the expected file contents. **Verified
            on both paths** — direct HTTPS downloads and HF Hub
            downloads. Issue #20: previously the HF Hub branch
            silently ignored this argument, so a malicious or
            accidentally-mutated repo at the configured revision
            would have produced bytes whose hash did not match what
            the caller pinned. ``None`` skips verification (HF
            Hub-only — the URL branch still requires it).
    """
    if model_id.startswith("http://") or model_id.startswith("https://"):
        cache_segment = "url--" + hashlib.sha256(model_id.encode("utf-8")).hexdigest()[:32]
    else:
        cache_segment = _safe_cache_segment(model_id, kind="model_id")
    safe_filename = _safe_cache_segment(filename, kind="filename")
    cached_path = self.cache_dir / cache_segment / safe_filename
    if cached_path.exists():
        if sha256 is not None and self.get_checksum(cached_path) != sha256.lower():
            raise ChecksumMismatchError(
                f"Cached file {cached_path} does not match expected SHA256."
            )
        return cached_path

    if model_id.startswith("http://") or model_id.startswith("https://"):
        if sha256 is None:
            raise ValueError(
                "Direct URL downloads require a sha256= argument. "
                "Pass the expected hex digest of the weights file."
            )
        return self._download_url(model_id, cached_path, sha256)

    try:
        hf_path = self._download_hf_hub(model_id, filename, revision)
    except ImportError:
        raise ImportError(
            f"Cannot download model '{model_id}': huggingface_hub not installed. "
            "Install with: pip install openlithohub[models]"
        ) from None
    if sha256 is not None:
        actual = self.get_checksum(hf_path)
        if actual != sha256.lower():
            raise ChecksumMismatchError(
                f"SHA256 mismatch for HF Hub file {model_id}/{filename} "
                f"@ {revision}: expected {sha256.lower()}, got {actual}. "
                f"The file at this revision differs from the digest you "
                f"pinned — refuse to load."
            )
    elif revision in {"main", "master"} or revision is None:
        # Mutable revision + no sha256 = the bytes you receive today
        # may differ from the bytes you receive tomorrow. Warn so
        # downstream reproducibility claims aren't silently false.
        warnings.warn(
            f"download_weights({model_id!r}, revision={revision!r}) "
            f"with no sha256= and a mutable revision: the publisher "
            f"can change the bytes at any time. Pass either a "
            f"40-hex commit hash as `revision=` or a `sha256=` "
            f"digest for reproducibility.",
            UserWarning,
            stacklevel=2,
        )
    return hf_path

list_cached()

List model IDs that have cached weights.

HF-Hub-style entries are decoded from the on-disk owner--repo form back to owner/repo. URL-keyed entries (stored under url--<sha256>) are returned verbatim so a caller can hand the result straight back to :meth:clear_cache without re-keying.

Source code in src/openlithohub/models/hub.py 
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
def list_cached(self) -> list[str]:
    """List model IDs that have cached weights.

    HF-Hub-style entries are decoded from the on-disk ``owner--repo`` form
    back to ``owner/repo``. URL-keyed entries (stored under
    ``url--<sha256>``) are returned verbatim so a caller can hand the
    result straight back to :meth:`clear_cache` without re-keying.
    """
    if not self.cache_dir.exists():
        return []
    names: list[str] = []
    for d in self.cache_dir.iterdir():
        if not d.is_dir():
            continue
        if d.name.startswith("url--"):
            names.append(d.name)
        else:
            names.append(d.name.replace("--", "/"))
    return sorted(names)

clear_cache(model_id=None)

Remove cached weights for a model (or all if model_id is None).

model_id is routed through the same per-segment validator used by download_weights so a caller-supplied .. cannot escape the cache directory and rmtree something it shouldn't. URL-keyed entries (cached under the url--<sha256> segment that download_weights creates) are also accepted.

Source code in src/openlithohub/models/hub.py 
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
def clear_cache(self, model_id: str | None = None) -> None:
    """Remove cached weights for a model (or all if model_id is None).

    ``model_id`` is routed through the same per-segment validator used by
    ``download_weights`` so a caller-supplied ``..`` cannot escape the
    cache directory and rmtree something it shouldn't. URL-keyed entries
    (cached under the ``url--<sha256>`` segment that ``download_weights``
    creates) are also accepted.
    """
    import shutil

    if model_id is None:
        if self.cache_dir.exists():
            shutil.rmtree(self.cache_dir)
            self.cache_dir.mkdir(parents=True, exist_ok=True)
        return

    if model_id.startswith("http://") or model_id.startswith("https://"):
        cache_segment = "url--" + hashlib.sha256(model_id.encode("utf-8")).hexdigest()[:32]
    elif model_id.startswith("url--"):
        # The exact on-disk segment as returned by `list_cached` for
        # URL-keyed entries. Validate the suffix is a clean hex segment
        # so a caller can't sneak `..` past us via this branch.
        suffix = model_id[len("url--") :]
        if not suffix or any(c not in "0123456789abcdef" for c in suffix):
            raise ValueError(f"Refusing unsafe url-keyed cache id: {model_id!r}")
        cache_segment = model_id
    else:
        cache_segment = _safe_cache_segment(model_id, kind="model_id")
    model_dir = self.cache_dir / cache_segment
    if model_dir.exists():
        shutil.rmtree(model_dir)

get_checksum(path)

Compute SHA256 checksum of a file.

Source code in src/openlithohub/models/hub.py 
403
404
405
406
407
408
409
def get_checksum(self, path: Path) -> str:
    """Compute SHA256 checksum of a file."""
    sha = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            sha.update(chunk)
    return sha.hexdigest()

openlithohub.models.examples.dummy_model

Dummy identity model for testing the evaluation pipeline.

DummyModel

Bases: LithographyModel

Trivial model that returns the input design as the mask (identity).

Source code in src/openlithohub/models/examples/dummy_model.py 
13
14
15
16
17
18
19
20
21
22
@registry.register
class DummyModel(LithographyModel):
    """Trivial model that returns the input design as the mask (identity)."""

    NAME = "dummy-identity"
    SUPPORTS_CURVILINEAR = False
    RECEPTIVE_FIELD_PX = 0

    def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
        return PredictionResult(mask=design.clone())

FailingDummyModel

Bases: LithographyModel

Dummy model whose predict always raises.

Exists so the multi-GPU tile pipeline can exercise worker error propagation across processes — pickling a closure that raises does not survive the spawn boundary, but a registered model name does.

Source code in src/openlithohub/models/examples/dummy_model.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
@registry.register
class FailingDummyModel(LithographyModel):
    """Dummy model whose ``predict`` always raises.

    Exists so the multi-GPU tile pipeline can exercise worker error
    propagation across processes — pickling a closure that raises does
    not survive the spawn boundary, but a registered model name does.
    """

    NAME = "dummy-failing"
    SUPPORTS_CURVILINEAR = False
    RECEPTIVE_FIELD_PX = 0

    def predict(self, design: torch.Tensor, **kwargs: Any) -> PredictionResult:
        raise RuntimeError("dummy-failing: deliberate predict failure")

Differentiable forward models

openlithohub._utils.forward_model

Simplified aerial image forward model using Gaussian PSF convolution.

Padding contract

All convolutions in this module MUST use circular (periodic) padding, not zero-padding. This is not stylistic — the Hopkins source-mask formulation treats the mask as a tile in a periodic illumination, and zero-padding the input introduces spurious dim-aerial fringes near the frame edge that silently degrade EPE / PV-band metrics on layouts with features close to the boundary. If you add a new conv-based simulator here, route it through _circular_pad_clamped (or an equivalent mode="circular" call) — do not switch to F.conv2d's default zero-pad as a "simplification".

simulate_aerial_image(mask, sigma_px, dose=1.0)

Simulate aerial image via Gaussian PSF convolution.

Approximates Hopkins diffraction with a single Gaussian point spread function.

Accepts (H, W) for single-image use and (B, 1, H, W) for batched forward passes — the output preserves the input rank.

Uses circular (periodic) padding to match the Hopkins forward model's convention. OPC treats the mask as a tile of an infinite layout, so zero-padding at the border would introduce spurious dim-aerial fringes that the Hopkins path does not have.

Source code in src/openlithohub/_utils/forward_model.py 
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
def simulate_aerial_image(
    mask: torch.Tensor,
    sigma_px: float,
    dose: float = 1.0,
) -> torch.Tensor:
    """Simulate aerial image via Gaussian PSF convolution.

    Approximates Hopkins diffraction with a single Gaussian point spread function.

    Accepts ``(H, W)`` for single-image use and ``(B, 1, H, W)`` for batched
    forward passes — the output preserves the input rank.

    Uses circular (periodic) padding to match the Hopkins forward model's
    convention. OPC treats the mask as a tile of an infinite layout, so
    zero-padding at the border would introduce spurious dim-aerial fringes
    that the Hopkins path does not have.
    """
    if sigma_px < 1e-6:
        return mask.float() * dose

    kernel = _build_gaussian_kernel(sigma_px, mask.device)
    padding = kernel.shape[-1] // 2

    squeezed = False
    if mask.ndim == 2:
        inp = mask.float().unsqueeze(0).unsqueeze(0)
        squeezed = True
    elif mask.ndim == 4 and mask.shape[1] == 1:
        inp = mask.float()
    else:
        raise ValueError(f"Expected mask shape (H,W) or (B,1,H,W); got {tuple(mask.shape)}")

    inp_padded = _circular_pad_clamped(inp, padding)
    aerial = functional.conv2d(inp_padded, kernel)
    if squeezed:
        aerial = aerial.squeeze(0).squeeze(0)
    return aerial * dose

apply_resist_threshold(aerial_image, threshold=THRESHOLD_GENERIC, *, resist_diffusion_nm=0.0, pixel_size_nm=1.0, quencher=0.0)

Apply a hard resist threshold to produce a binary resist pattern.

The 0.5 default is a generic mid-intensity cutoff for ad-hoc use; the canonical ICCAD16 / LithoBench cutoff is 0.225 (see [Yang2023_LithoBench, §3.2, p.5] and SimulatorConfig.threshold). Pass threshold=0.225 when reproducing benchmark numbers.

With resist_diffusion_nm=0.0 and quencher=0.0 (default) this is constant threshold resist (CTR) without diffusion — the sigmoid-on-aerial simplification documented in docs/architecture.md → Resist Model Simplification. The output is bit-identical to the legacy (aerial >= threshold).float() path.

With a positive resist_diffusion_nm the aerial image is first blurred by a Gaussian whose sigma matches the acid diffusion length, then quencher is subtracted, and finally the hard threshold is applied. This models a simplified chemically-amplified resist (CAR).

Real per-node CTR parameters are foundry-confidential and cannot ship in an open-source repo; benchmark-relative comparison is unaffected, but absolute wafer prediction is not in scope.

Returns a hard 0/1 tensor — gradients do not flow back through this function. The README's "end-to-end differentiable" claim refers to the ILT optimizer path, which uses :func:openlithohub._utils.resist_model.differentiable_threshold (a temperature-controlled sigmoid). Use that helper for any gradient-bearing forward; reserve this hard threshold for measurement / scoring code (PVB envelopes, stochastic comparisons, leaderboard pass/fail).

Source code in src/openlithohub/_utils/forward_model.py 
129
130
131
132
133
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
def apply_resist_threshold(
    aerial_image: torch.Tensor,
    threshold: float = THRESHOLD_GENERIC,
    *,
    resist_diffusion_nm: float = 0.0,
    pixel_size_nm: float = 1.0,
    quencher: float = 0.0,
) -> torch.Tensor:
    """Apply a hard resist threshold to produce a binary resist pattern.

    The 0.5 default is a generic mid-intensity cutoff for ad-hoc use; the
    canonical ICCAD16 / LithoBench cutoff is 0.225 (see
    [Yang2023_LithoBench, §3.2, p.5] and ``SimulatorConfig.threshold``).
    Pass ``threshold=0.225`` when reproducing benchmark numbers.

    With ``resist_diffusion_nm=0.0`` and ``quencher=0.0`` (default) this
    is **constant threshold resist (CTR) without diffusion** — the
    sigmoid-on-aerial simplification documented in
    ``docs/architecture.md → Resist Model Simplification``. The output is
    bit-identical to the legacy ``(aerial >= threshold).float()`` path.

    With a positive ``resist_diffusion_nm`` the aerial image is first
    blurred by a Gaussian whose sigma matches the acid diffusion length,
    then quencher is subtracted, and finally the hard threshold is
    applied. This models a simplified chemically-amplified resist (CAR).

    Real per-node CTR parameters are foundry-confidential and cannot ship
    in an open-source repo; benchmark-relative comparison is unaffected,
    but absolute wafer prediction is not in scope.

    Returns a hard 0/1 tensor — gradients do **not** flow back through
    this function. The README's "end-to-end differentiable" claim refers
    to the ILT optimizer path, which uses
    :func:`openlithohub._utils.resist_model.differentiable_threshold`
    (a temperature-controlled sigmoid). Use that helper for any
    gradient-bearing forward; reserve this hard threshold for
    measurement / scoring code (PVB envelopes, stochastic comparisons,
    leaderboard pass/fail).
    """
    if resist_diffusion_nm <= 0.0 and quencher <= 0.0:
        return (aerial_image >= threshold).float()

    acid = aerial_image.clone()
    sigma_px = resist_diffusion_nm / max(pixel_size_nm, 1e-6)
    if sigma_px > 0.1:
        acid = _gaussian_diffuse(acid, sigma_px)
    acid = (acid - quencher).clamp(min=0.0)
    return (acid >= threshold).float()

openlithohub._utils.hopkins

Differentiable Hopkins partial-coherence aerial image model via SOCS.

Hopkins formulation describes partial-coherent imaging through a Transmission Cross Coefficient (TCC):

I(x) = ∫∫ TCC(f1, f2) * M~(f1) * conj(M~(f2)) * exp(2πi (f1-f2) x) df1 df2

where M~ is the mask Fourier transform, J is the source intensity, P is the pupil function, and

TCC(f1, f2) = ∫ J(f) * P(f + f1) * conj(P(f + f2)) df

The Sum Of Coherent Systems (SOCS) decomposition is the eigendecomposition of TCC viewed as a Hermitian operator:

TCC = Σ_k w_k * φ_k(f1) * conj(φ_k(f2))           with w_k descending

Truncating at K kernels yields the standard fast OPC forward model:

I(x) ≈ Σ_k w_k * | (mask * φ_k)(x) |^2

This module implements both steps in pure PyTorch so the entire chain is auto-differentiable (the mask is the optimization variable in ILT).

HopkinsParams dataclass

Optical parameters for Hopkins partial-coherence imaging.

Attributes:

Name Type Description
wavelength_nm float

Exposure wavelength (193 nm = ArF, 13.5 nm = EUV).
na float

Numerical aperture (image-side). 1.35 = ArF immersion, 0.33 = EUV NXE.
sigma float

Partial-coherence factor for circular illumination, or outer sigma for annular/dipole/quasar.
sigma_inner float

Inner sigma for annular/dipole/quasar (ignored for circular).
pixel_size_nm float

Physical size of one mask pixel.
num_kernels int

SOCS truncation order. Defaults to 24 to match the Yang2023_LithoBench Table II benchmark; that paper does not publish a truncation-error vs K curve, so the underlying defensibility chain is Cobb 1995, §IV (the original SOCS construction) plus accumulated practice. Production deployments at a different node should re-sweep K against their own EPE noise floor before pinning this value.
illumination IlluminationKind

Source shape — circular, annular, dipole (X-direction poles), or quasar (4-pole, CQuad).
dipole_angle_deg float

Pole-pair orientation for dipole/quasar (degrees).
pole_opening_deg float

Half-angle of each pole wedge for dipole/quasar (degrees). 30° is a common production CQuad value.
defocus_nm float

Defocus offset; affects the pupil phase only.
Source code in src/openlithohub/_utils/hopkins.py 
 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
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
@dataclass(frozen=True)
class HopkinsParams:
    """Optical parameters for Hopkins partial-coherence imaging.

    Attributes:
        wavelength_nm: Exposure wavelength (193 nm = ArF, 13.5 nm = EUV).
        na: Numerical aperture (image-side). 1.35 = ArF immersion, 0.33 = EUV NXE.
        sigma: Partial-coherence factor for circular illumination, or
            outer sigma for annular/dipole/quasar.
        sigma_inner: Inner sigma for annular/dipole/quasar (ignored for circular).
        pixel_size_nm: Physical size of one mask pixel.
        num_kernels: SOCS truncation order. Defaults to 24 to match the
            ``Yang2023_LithoBench`` Table II benchmark; that paper does
            not publish a truncation-error vs K curve, so the underlying
            defensibility chain is Cobb 1995, §IV (the original SOCS
            construction) plus accumulated practice. Production
            deployments at a different node should re-sweep K against
            their own EPE noise floor before pinning this value.
        illumination: Source shape — circular, annular, dipole (X-direction
            poles), or quasar (4-pole, CQuad).
        dipole_angle_deg: Pole-pair orientation for dipole/quasar (degrees).
        pole_opening_deg: Half-angle of each pole wedge for dipole/quasar
            (degrees). 30° is a common production CQuad value.
        defocus_nm: Defocus offset; affects the pupil phase only.
    """

    wavelength_nm: float = WAVELENGTH_ARF_NM
    na: float = NA_IMMERSION
    sigma: float = SIGMA_OUTER_DEFAULT
    sigma_inner: float = SIGMA_INNER_DEFAULT
    pixel_size_nm: float = PIXEL_SIZE_NM_DEFAULT
    num_kernels: int = NUM_KERNELS_DEFAULT
    illumination: IlluminationKind = "circular"
    dipole_angle_deg: float = 0.0
    pole_opening_deg: float = POLE_OPENING_DEG_DEFAULT
    defocus_nm: float = DEFOCUS_NM_DEFAULT

    def cache_key(self, grid_size: int, device: str, kernel_dtype: str) -> Hashable:
        return (
            self.wavelength_nm,
            self.na,
            self.sigma,
            self.sigma_inner,
            self.pixel_size_nm,
            self.num_kernels,
            self.illumination,
            self.dipole_angle_deg,
            self.pole_opening_deg,
            self.defocus_nm,
            grid_size,
            device,
            kernel_dtype,
        )

compute_socs_kernels(params, grid_size, device='cpu', dtype=torch.complex64)

Compute SOCS kernels and their weights for a square grid.

Parameters:

Name Type Description Default
params HopkinsParams

Optical parameters.

 required
grid_size int

Square grid edge length (pixels).

 required
device device | str

PyTorch device.

 'cpu'
dtype dtype

Complex dtype of the returned kernels (complex64 or complex128). The internal FFT/SVD always runs in complex64 — this dtype only controls the cached output.

 complex64

Returns:

Name Type Description
kernels Tensor

complex tensor of shape (K, H, W) with the requested dtype. Each kernel is in the spatial domain, ready for FFT-based convolution.
weights Tensor

real float32 tensor of shape (K,). Sorted descending.

The returned kernels are zero-centered (fftshift-style) with the kernel support concentrated near the origin, which is what simulate_aerial_image_hopkins expects.

Source code in src/openlithohub/_utils/hopkins.py 
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
def compute_socs_kernels(
    params: HopkinsParams,
    grid_size: int,
    device: torch.device | str = "cpu",
    dtype: torch.dtype = torch.complex64,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Compute SOCS kernels and their weights for a square grid.

    Args:
        params: Optical parameters.
        grid_size: Square grid edge length (pixels).
        device: PyTorch device.
        dtype: Complex dtype of the returned kernels (``complex64`` or
            ``complex128``). The internal FFT/SVD always runs in
            ``complex64`` — this dtype only controls the cached output.

    Returns:
        kernels: complex tensor of shape (K, H, W) with the requested dtype.
            Each kernel is in the spatial domain, ready for FFT-based
            convolution.
        weights: real float32 tensor of shape (K,). Sorted descending.

    The returned kernels are zero-centered (fftshift-style) with the kernel
    support concentrated near the origin, which is what
    `simulate_aerial_image_hopkins` expects.
    """
    dev = torch.device(device)
    cache_key = params.cache_key(grid_size, str(dev), str(dtype))
    cached = _KERNEL_CACHE.get(cache_key)
    if cached is not None:
        _KERNEL_CACHE.move_to_end(cache_key)
        return cached

    f = _frequency_grid(grid_size, params.pixel_size_nm, dev)
    fy, fx = torch.meshgrid(f, f, indexing="ij")

    pupil = _pupil(fx, fy, params)

    src_shifts, src_weights = _illumination_samples(params, grid_size, dev)
    n_src = src_shifts.shape[0]

    n_freq = grid_size * grid_size
    # SVD matrix is (n_src, n_freq) complex64 — warn on large grids where
    # memory and compute time grow as O(n_src * grid_size^2).
    mem_gb = n_src * n_freq * 8 / 1e9
    if mem_gb > 4.0:
        warnings.warn(
            f"SOCS kernel computation for grid_size={grid_size} allocates ~{mem_gb:.1f} GB. "
            f"Consider using a smaller grid or reducing num_kernels.",
            stacklevel=2,
        )

    yy, xx = torch.meshgrid(
        torch.arange(grid_size, device=dev),
        torch.arange(grid_size, device=dev),
        indexing="ij",
    )
    H = torch.zeros((n_src, n_freq), dtype=torch.complex64, device=dev)  # noqa: N806
    for k in range(n_src):
        sy = int(src_shifts[k, 0].item())
        sx = int(src_shifts[k, 1].item())
        idx_y = (yy - sy) % grid_size
        idx_x = (xx - sx) % grid_size
        shifted = pupil[idx_y, idx_x]
        weight = torch.sqrt(src_weights[k])
        H[k] = (shifted * weight).reshape(-1)

    K = max(1, min(params.num_kernels, n_src))  # noqa: N806
    u, s, vh = torch.linalg.svd(H, full_matrices=False)
    s2 = (s**2)[:K]
    eigvecs = vh[:K]

    kernels_freq = eigvecs.reshape(K, grid_size, grid_size)
    weights = s2.to(torch.float32)

    kernels_spatial = torch.fft.ifft2(kernels_freq, norm="backward")
    kernels_spatial = torch.fft.fftshift(kernels_spatial, dim=(-2, -1))

    # Calibrate so that an open-frame (all-ones) mask produces aerial ≈ 1.
    # For a constant mask, coherent_k = sum(kernel_k); aerial_open = Σ_k w_k |sum(k_k)|².
    open_frame = torch.zeros((), dtype=torch.float32, device=dev)
    for k_idx in range(K):
        coherent_dc = kernels_spatial[k_idx].sum()
        open_frame = open_frame + weights[k_idx] * (coherent_dc.real**2 + coherent_dc.imag**2)
    if float(open_frame) > 0.0:
        weights = weights / open_frame
        # Clamp rescaled weights to prevent overflow in downstream aerial
        # accumulation (extremely small open_frame values can produce huge weights).
        weights = weights.clamp(max=1e6)

    kernels_spatial = kernels_spatial.to(dtype)
    _KERNEL_CACHE[cache_key] = (kernels_spatial.detach(), weights.detach())
    while len(_KERNEL_CACHE) > _KERNEL_CACHE_MAXSIZE:
        _KERNEL_CACHE.popitem(last=False)
    return _KERNEL_CACHE[cache_key]

simulate_aerial_image_hopkins(mask, params=None, kernels=None, weights=None, dose=1.0, dtype=torch.float32, precomputed_kernels_f=None)

Simulate aerial image via SOCS-truncated Hopkins imaging.

Parameters:

Name Type Description Default
mask Tensor

Real-valued mask (H, W) or (B, 1, H, W), values in [0, 1]. Differentiable: gradients flow back through the kernels.

 required
params HopkinsParams | None

Optical parameters. Required if kernels/weights are None.

 None
kernels Tensor | None

Pre-computed complex SOCS kernels (K, H, W). If provided, params is only used for dose (and may be None).

 None
weights Tensor | None

Pre-computed real weights (K,). Must accompany kernels.

 None
dose float

Linear dose multiplier on the resulting intensity.

 1.0
dtype dtype

Real dtype of the returned aerial image (float32 or bfloat16). The internal FFT is always done in complex64 because PyTorch's fft2 does not support bfloat16-complex; the cast happens before squaring and at the output.

 float32
precomputed_kernels_f Tensor | None

Optional pre-FFT'd kernels of shape (K, H, W), complex64. When provided, the inner loop skips the per-kernel ifftshift + fft2 cost. Must be the FFT of ifftshift(kernels, dim=(-2,-1)) for numerical equivalence. Coerced to complex64 if a different complex dtype is passed.

 None

Returns:

Type Description
Tensor

Real-valued aerial image with the same spatial shape as mask.
Source code in src/openlithohub/_utils/hopkins.py 
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
def simulate_aerial_image_hopkins(
    mask: torch.Tensor,
    params: HopkinsParams | None = None,
    kernels: torch.Tensor | None = None,
    weights: torch.Tensor | None = None,
    dose: float = 1.0,
    dtype: torch.dtype = torch.float32,
    precomputed_kernels_f: torch.Tensor | None = None,
) -> torch.Tensor:
    """Simulate aerial image via SOCS-truncated Hopkins imaging.

    Args:
        mask: Real-valued mask (H, W) or (B, 1, H, W), values in [0, 1].
            Differentiable: gradients flow back through the kernels.
        params: Optical parameters. Required if `kernels`/`weights` are None.
        kernels: Pre-computed complex SOCS kernels (K, H, W). If provided,
            `params` is only used for `dose` (and may be None).
        weights: Pre-computed real weights (K,). Must accompany `kernels`.
        dose: Linear dose multiplier on the resulting intensity.
        dtype: Real dtype of the returned aerial image (``float32`` or
            ``bfloat16``). The internal FFT is always done in
            ``complex64`` because PyTorch's ``fft2`` does not support
            ``bfloat16``-complex; the cast happens before squaring and at
            the output.
        precomputed_kernels_f: Optional pre-FFT'd kernels of shape
            (K, H, W), complex64. When provided, the inner loop skips the
            per-kernel ``ifftshift + fft2`` cost. Must be the FFT of
            ``ifftshift(kernels, dim=(-2,-1))`` for numerical equivalence.
            Coerced to complex64 if a different complex dtype is passed.

    Returns:
        Real-valued aerial image with the same spatial shape as `mask`.
    """
    squeezed = False
    if mask.ndim == 2:
        mask4d = mask.unsqueeze(0).unsqueeze(0)
        squeezed = True
    elif mask.ndim == 4 and mask.shape[1] == 1:
        mask4d = mask
    else:
        raise ValueError(f"Expected mask shape (H,W) or (B,1,H,W); got {tuple(mask.shape)}")

    B, _, H, W = mask4d.shape  # noqa: N806
    if H != W:
        raise ValueError(f"Hopkins forward model expects a square grid; got {H}x{W}")

    if kernels is None or weights is None:
        if params is None:
            raise ValueError("Provide either (params) or (kernels and weights).")
        kernels, weights = compute_socs_kernels(params, H, mask4d.device)

    if params is not None:
        # Tile must be wider than a few Rayleigh units, otherwise the
        # circular FFT convolution wraps optical energy from one edge to
        # the other and contaminates the interior. ~4 Rayleigh
        # (lambda/NA) gives the kernel room to decay before wrapping.
        rayleigh_nm = params.wavelength_nm / max(params.na, 1e-6)
        tile_extent_nm = H * params.pixel_size_nm
        if tile_extent_nm < 4.0 * rayleigh_nm:
            warnings.warn(
                f"Hopkins forward: tile extent {tile_extent_nm:.0f} nm "
                f"({H} px x {params.pixel_size_nm} nm) is smaller than "
                f"4*lambda/NA={4 * rayleigh_nm:.0f} nm; circular FFT "
                f"wraparound will pollute tile edges. Use larger tiles "
                f"or pad before calling.",
                UserWarning,
                stacklevel=2,
            )

    image = mask4d.to(torch.float32).squeeze(1)  # (B, H, W)
    K = kernels.shape[0]  # noqa: N806
    aerial = torch.zeros_like(image)
    kernels_c64 = kernels.to(torch.complex64) if kernels.dtype != torch.complex64 else kernels

    if precomputed_kernels_f is not None:
        if precomputed_kernels_f.dtype != torch.complex64:
            precomputed_kernels_f = precomputed_kernels_f.to(torch.complex64)
        if precomputed_kernels_f.shape[0] != K:
            raise ValueError(
                f"precomputed_kernels_f has K={precomputed_kernels_f.shape[0]}; "
                f"expected {K} matching kernels."
            )

    image_c = image.to(torch.complex64)
    image_f = torch.fft.fft2(image_c)  # (B, H, W)
    for k in range(K):
        if precomputed_kernels_f is not None:
            kernel_f = precomputed_kernels_f[k]
        else:
            kernel_shifted = torch.fft.ifftshift(kernels_c64[k], dim=(-2, -1))
            kernel_f = torch.fft.fft2(kernel_shifted)  # (H, W)
        coherent = torch.fft.ifft2(image_f * kernel_f.unsqueeze(0))  # (B, H, W)
        aerial = aerial + weights[k] * (coherent.real**2 + coherent.imag**2)

    aerial = aerial * dose
    if dtype != torch.float32:
        aerial = aerial.to(dtype)
    if squeezed:
        return aerial.squeeze(0)
    return aerial.unsqueeze(1)

clear_kernel_cache()

Drop all cached SOCS kernels. Useful in tests and long-running services.

Source code in src/openlithohub/_utils/hopkins.py 
474
475
476
def clear_kernel_cache() -> None:
    """Drop all cached SOCS kernels. Useful in tests and long-running services."""
    _KERNEL_CACHE.clear()

openlithohub._utils.resist_model

Chemically-amplified resist simulation with acid diffusion.

differentiable_threshold(aerial_image, threshold=THRESHOLD_GENERIC, steepness=STEEPNESS_DEFAULT)

Smooth, differentiable substitute for a hard resist threshold.

Returns sigmoid(steepness * (aerial - threshold)). As steepness increases the output approaches a step function while remaining differentiable everywhere — required for gradient-based ILT.

Source code in src/openlithohub/_utils/resist_model.py 
12
13
14
15
16
17
18
19
20
21
22
23
def differentiable_threshold(
    aerial_image: torch.Tensor,
    threshold: float = THRESHOLD_GENERIC,
    steepness: float = STEEPNESS_DEFAULT,
) -> torch.Tensor:
    """Smooth, differentiable substitute for a hard resist threshold.

    Returns ``sigmoid(steepness * (aerial - threshold))``. As ``steepness``
    increases the output approaches a step function while remaining
    differentiable everywhere — required for gradient-based ILT.
    """
    return torch.sigmoid(steepness * (aerial_image - threshold))

simulate_resist(aerial_image, acid_diffusion_length_nm=5.0, pixel_size_nm=1.0, threshold=THRESHOLD_GENERIC, quencher_concentration=0.1)

Simulate chemically-amplified resist response with acid diffusion.

Models a physically-motivated resist development process: 1. Photoacid generation proportional to aerial image intensity 2. Acid diffusion via Gaussian blur (diffusion length determines spread) 3. Quencher neutralization (constant subtraction) 4. Threshold to binary resist pattern

Parameters:

Name Type Description Default
aerial_image Tensor

Aerial image intensity (H, W), values in [0, 1].

 required
acid_diffusion_length_nm float

Acid diffusion length in nanometers.

 5.0
pixel_size_nm float

Physical pixel size for unit conversion.

 1.0
threshold float

Development threshold applied to the post-quencher acid field. The quencher is subtracted before thresholding, so an aerial intensity of threshold + quencher_concentration is the dose where development just kicks in. This is the standard CAR convention (more quencher → more acid needed) — not a normalized intensity threshold. To keep a fixed effective dose cutoff regardless of quencher, set threshold = nominal_threshold - quencher_concentration.

 THRESHOLD_GENERIC
quencher_concentration float

Base quencher level subtracted from acid.

 0.1

Returns:

Type Description
Tensor

Binary resist pattern (H, W), 1.0 where resist remains.
Source code in src/openlithohub/_utils/resist_model.py 
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
def simulate_resist(
    aerial_image: torch.Tensor,
    acid_diffusion_length_nm: float = 5.0,
    pixel_size_nm: float = 1.0,
    threshold: float = THRESHOLD_GENERIC,
    quencher_concentration: float = 0.1,
) -> torch.Tensor:
    """Simulate chemically-amplified resist response with acid diffusion.

    Models a physically-motivated resist development process:
    1. Photoacid generation proportional to aerial image intensity
    2. Acid diffusion via Gaussian blur (diffusion length determines spread)
    3. Quencher neutralization (constant subtraction)
    4. Threshold to binary resist pattern

    Args:
        aerial_image: Aerial image intensity (H, W), values in [0, 1].
        acid_diffusion_length_nm: Acid diffusion length in nanometers.
        pixel_size_nm: Physical pixel size for unit conversion.
        threshold: Development threshold *applied to the post-quencher acid
            field*. The quencher is subtracted before thresholding, so an
            aerial intensity of ``threshold + quencher_concentration`` is
            the dose where development just kicks in. This is the standard
            CAR convention (more quencher → more acid needed) — not a
            normalized intensity threshold. To keep a fixed effective dose
            cutoff regardless of quencher, set
            ``threshold = nominal_threshold - quencher_concentration``.
        quencher_concentration: Base quencher level subtracted from acid.

    Returns:
        Binary resist pattern (H, W), 1.0 where resist remains.
    """
    acid = aerial_image.clone()

    sigma_diffusion_px = acid_diffusion_length_nm / max(pixel_size_nm, 1e-6)
    if sigma_diffusion_px > 0.1:
        acid = _diffuse_acid(acid, sigma_diffusion_px)

    acid = (acid - quencher_concentration).clamp(min=0.0)
    return (acid >= threshold).float()

simulate_resist_soft(aerial_image, acid_diffusion_length_nm=5.0, pixel_size_nm=1.0, threshold=THRESHOLD_GENERIC, quencher_concentration=0.1, steepness=STEEPNESS_DEFAULT)

Differentiable resist simulation using sigmoid instead of hard threshold.

Same physics as simulate_resist but uses a smooth sigmoid for the development step, making it suitable for gradient-based optimization.

Source code in src/openlithohub/_utils/resist_model.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def simulate_resist_soft(
    aerial_image: torch.Tensor,
    acid_diffusion_length_nm: float = 5.0,
    pixel_size_nm: float = 1.0,
    threshold: float = THRESHOLD_GENERIC,
    quencher_concentration: float = 0.1,
    steepness: float = STEEPNESS_DEFAULT,
) -> torch.Tensor:
    """Differentiable resist simulation using sigmoid instead of hard threshold.

    Same physics as `simulate_resist` but uses a smooth sigmoid for the
    development step, making it suitable for gradient-based optimization.
    """
    acid = aerial_image.clone()

    sigma_diffusion_px = acid_diffusion_length_nm / max(pixel_size_nm, 1e-6)
    if sigma_diffusion_px > 0.1:
        acid = _diffuse_acid(acid, sigma_diffusion_px)

    acid = (acid - quencher_concentration).clamp(min=0.0)
    return differentiable_threshold(acid, threshold=threshold, steepness=steepness)