File size: 1,271 Bytes
bfdbd33 0a16898 bfdbd33 0a16898 bfdbd33 | 1 2 3 4 5 6 7 8 9 10 11 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 | """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
|