The class facade#
A bare @sweep function has nowhere clean to put state built once and reused
across calls: an engine handle, a loaded lookup table, an open connection. A
closure breaks the pickling the process executor needs, and a module global is
not much better.
SweepModule gives that state a home, modelled on the torch.nn.Module
idiom:
class RhoAtm(SweepModule):
contract = "loop(aot, rh) vec(wl @ 8) -> rho_atm(wl)"
def __init__(self, policy=None, engine=None):
super().__init__(policy)
self.engine = engine # built once, reused by every call
def forward(self, aot, rh, wl, *, n_ph):
return self.engine.run(aot, rh, wl, n_ph)
mod = RhoAtm(
SweepPolicy(store="runs/rho.zarr", executor="process"), engine=load_engine()
)
result = mod(space, n_ph=int(1e6))
contract is a class attribute, so a malformed one raises when the module is
imported rather than after twenty minutes of engine time.
forward stays pure physics and is testable with no sweep involved:
RhoAtm(engine=engine).forward(0.1, 50.0, wl, n_ph=100)
__call__ does the orchestration, and explain works the same as on a
decorated function.
When to reach for it#
This is also the combination that most needs executor="process": something
expensive, built once, then run in parallel. The instance’s state travels
into the worker processes intact rather than being rebuilt per worker.
If a bare function is enough, the decorator is simpler and gets you the same
contract and policy machinery. Reach for the class when __init__ earns its
keep.
Worked in SweepModule.