| """Load physarum natively on hosts without a published variant (e.g. a Pi). |
| |
| JIT-builds the local source with `torch.utils.cpp_extension` and exposes the |
| identical `physarum` module API. |
| |
| import load_local |
| physarum = load_local.load() |
| sim = physarum.Physarum(width=1024, height=1024, agents=200000) |
| """ |
| import os |
| import sys |
| import types |
|
|
| _cached = None |
|
|
|
|
| def load(verbose=False): |
| global _cached |
| if _cached is not None: |
| return _cached |
| from torch.utils.cpp_extension import load as _jit |
|
|
| root = os.path.dirname(os.path.abspath(__file__)) |
| ext = _jit(name="physarum_ext", |
| sources=[os.path.join(root, "local_bind.cpp"), |
| os.path.join(root, "physarum_csrc", "physarum_cpu.cpp"), |
| os.path.join(root, "physarum_csrc", "flow_cpu.cpp")], |
| extra_cflags=["-O3"], |
| verbose=verbose) |
| ops_mod = types.ModuleType("physarum._ops") |
| ops_mod.ops = type("_Ops", (), {"physarum_step": staticmethod(ext.physarum_step), |
| "flow_cg": staticmethod(ext.flow_cg)})() |
| sys.modules["physarum._ops"] = ops_mod |
| sys.path.insert(0, os.path.join(root, "torch-ext")) |
| import physarum |
| _cached = physarum |
| return physarum |
|
|
