# /// script # requires-python = ">=3.11,<3.15" # dependencies = ["itk==5.4.7", "SimpleITK==2.5.5", "numpy==2.4.6", "psutil==7.2.2"] # /// """Linux/glibc: uv run registration_memory_mwe.py No benchmark files are written. uv manages dependencies in its cache. Direct generated-wrapper imports are diagnostic-only, not a supported API. """ import ctypes import gc import hashlib import importlib import json import os import select import subprocess import sys import numpy as np import psutil def wrapped(name, suffix): return getattr(importlib.import_module(f"itk.itk{name}Python"), f"itk{name}{suffix}") def register(mode): k, j, i = np.mgrid[:4, :112, :128] x = (np.exp(-((j-47)**2+(i-51)**2)/180)*(1+k/10) + .65*np.exp(-((j-70)**2+(i-80)**2)/80)).astype("float64") iterations = [] if mode == "simpleitk": import SimpleITK as s s.ProcessObject.SetGlobalDefaultNumberOfThreads(1) fixed, moving = s.GetImageFromArray(x), s.GetImageFromArray(np.roll(x, 2, axis=2)) t = s.VersorRigid3DTransform(s.CenteredTransformInitializer( fixed, moving, s.VersorRigid3DTransform(), s.CenteredTransformInitializerFilter.GEOMETRY)) r = s.ImageRegistrationMethod() r.SetMetricAsCorrelation() r.SetInterpolator(s.sitkLinear) r.SetInitialTransform(t, inPlace=True) r.SetOptimizerAsGradientDescent(.01, 40, 1e-30, 1000, estimateLearningRate=r.Never) r.SetOptimizerScalesFromPhysicalShift() r.SetShrinkFactorsPerLevel([2, 1]) r.SetSmoothingSigmasPerLevel([1, 0]) r.AddCommand(s.sitkIterationEvent, lambda: iterations.append(int(r.GetOptimizerIteration()))) r.Execute(fixed, moving) r.RemoveAllCommands() value = r.GetMetricValue() y = s.GetArrayFromImage(s.Resample(moving, fixed, t, s.sitkLinear, 0)) version, path = s.Version_VersionString(), s.__file__ else: import itk itk.MultiThreaderBase.SetGlobalDefaultNumberOfThreads(1) I, T = itk.Image[itk.D, 3], itk.VersorRigid3DTransform[itk.D] if mode == "itk-public": M = itk.CorrelationImageToImageMetricv4[I, I] O, S = itk.GradientDescentOptimizerv4, itk.RegistrationParameterScalesFromPhysicalShift[M] R, C = itk.ImageRegistrationMethodv4[I, I], itk.CenteredTransformInitializer[T, I, I] F, L = itk.ResampleImageFilter[I, I], itk.LinearInterpolateImageFunction[I, itk.D] else: M, O, S, R, C, F, L = (wrapped(name, suffix) for name, suffix in [ ("CorrelationImageToImageMetricv4", "ID3ID3"), ("GradientDescentOptimizerv4", "TemplateD"), ("RegistrationParameterScalesFromPhysicalShift", "CITIMD3"), ("ImageRegistrationMethodv4", "REGv4D3D3"), ("CenteredTransformInitializer", "VR3DTDID3ID3"), ("ResampleImageFilter", "ID3ID3"), ("LinearInterpolateImageFunction", "ID3D")]) fixed, moving = itk.image_from_array(x), itk.image_from_array(np.roll(x, 2, axis=2)) t = T.New() c = C.New(Transform=t, FixedImage=fixed, MovingImage=moving) c.GeometryOn() c.InitializeTransform() metric = M.New() metric.SetMaximumNumberOfWorkUnits(1) o = O.New(LearningRate=.01, NumberOfIterations=40, MinimumConvergenceValue=1e-30, ConvergenceWindowSize=1000, DoEstimateLearningRateOnce=False, DoEstimateLearningRateAtEachIteration=False) scales = S.New(Metric=metric) o.SetScalesEstimator(scales) r = R.New(FixedImage=fixed, MovingImage=moving, InitialTransform=t, Metric=metric, Optimizer=o, InPlace=True) r.SetNumberOfWorkUnits(1) r.SetNumberOfLevels(2) r.SetShrinkFactorsPerLevel([2, 1]) r.SetSmoothingSigmasPerLevel([1, 0]) observer = o.AddObserver(itk.IterationEvent(), lambda: iterations.append(int(o.GetCurrentIteration()))) r.Update() o.RemoveObserver(observer) value = o.GetValue() f = F.New(Input=moving, Transform=t, Interpolator=L.New(), ReferenceImage=fixed, UseReferenceImage=True, DefaultPixelValue=0) f.Update() y = itk.array_from_image(f.GetOutput()) version, path = itk.Version.GetITKVersion(), itk.__file__ assert iterations == list(range(40))*2 assert np.isfinite(y).all() assert np.mean((y-x)**2) < np.mean((np.roll(x, 2, axis=2)-x)**2) return dict(output_sha256=hashlib.sha256(y.tobytes()).hexdigest(), parameters=list(t.GetParameters()), metric=value, iterations=iterations, version=version, package=path) def memory(pid): p = psutil.Process(pid) m = p.memory_full_info() libraries = [x for x in p.memory_maps() if ".so" in x.path] return [v/2**20 for v in (m.uss, m.pss, sum(x.private_dirty for x in libraries), sum(x.private_clean for x in libraries))] def main(): if len(sys.argv) == 2: result = register(sys.argv[1]) gc.collect() libc = ctypes.CDLL(None) libc.malloc_trim.argtypes = [ctypes.c_size_t] libc.malloc_trim(0) print(json.dumps(result), flush=True) sys.stdin.readline() # Keep libraries loaded until the controller finishes sampling. return env = dict(os.environ, PYTHONDONTWRITEBYTECODE="1", OPENBLAS_NUM_THREADS="1", OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS="1", ITK_EAGER_IMPORT="0", ITK_PYTHON_LAZYLOADING="1", ITK_PYTHON_DEFAULTFACTORYLOADING="1") reference = None print("Approach USS PSS library private dirty / clean (MiB)") for mode in ("itk-public", "itk-direct", "simpleitk"): children, results = [], [] try: for _ in range(2): child = subprocess.Popen([sys.executable, "-B", os.path.abspath(__file__), mode], env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True) children.append(child) if not select.select([child.stdout], [], [], 180)[0]: raise TimeoutError(f"{mode}: no result within 180 seconds") line = child.stdout.readline() if not line: raise RuntimeError(f"{mode} failed; see stderr above") results.append(json.loads(line)) for child, result in zip(children, results): if reference is None: reference = result assert result["output_sha256"] == reference["output_sha256"], "Output arrays differ" assert result["iterations"] == reference["iterations"] np.testing.assert_allclose(result["parameters"], reference["parameters"], rtol=1e-12, atol=1e-12) np.testing.assert_allclose(result["metric"], reference["metric"], rtol=1e-12, atol=1e-12) print(f"{mode:14s} " + " ".join(f"{v:9.2f}" for v in memory(child.pid))) print(f" version {result['version']}; {result['package']}") finally: for child in children: child.stdin.close() for child in children: try: child.wait(timeout=10) except subprocess.TimeoutExpired: child.kill() child.wait() child.stdout.close() print("PASS: identical output hashes, matching parameters/metric, 80 iterations per run.") if __name__ == "__main__": main()