engram

Personal research utilities for differentiable rendering (Mitsuba 3 / Dr.Jit), NN training and vision science. Engram in this context is more a "memory trace" that I upload here. It will be updated as I go, and may not be fully functional at any given time. The goal however is to provide a stable package at post-alpha time with a clean API making it easier for vision/neuro researchers to use Mitsuba 3 and Dr.Jit in their experiments.

Moves with the experiments I am running. The API is still unstable. Consider this pre-alpha.

uv add engram            # core: numpy, torch, matplotlib, mitsuba, drjit, mitsuba-scene-description
uv add "engram[all]"     # + scipy, wandb/pandas, flip-evaluator, moviepy, gpytoolbox, transformers/accelerate

Python >= 3.12. import engram is side-effect free: no Mitsuba variant is selected, nothing heavy is imported until you touch the module that needs it.

Module tour

Training framework

engram.train ships an HF-Trainer-shaped framework that covers inverse rendering and NN training with one interface: OptimizationConfig mirrors TrainingArguments (output dir, evaluation/save strategies, gradient accumulation, logging frequency), and Optimization subclasses implement the step/loop methods. With use_accelerator=True the run goes through Hugging Face accelerate (device placement, DDP, mixed precision — engram[hf]).

import mitsuba as mi
from engram.mi import set_variant

set_variant()  # cuda_ad_rgb -> llvm_ad_rgb -> scalar_rgb, or $ENGRAM_MI_VARIANT

from engram.mi.scene import SceneDescription
from engram.train import DROptimization, OptimizationConfig
from engram.train.render import RenderPipeline
from engram.train.optimization import GroundTruth

scene = SceneDescription.from_dict(mi.cornell_box())      # dicts or msd.Plugin objects
target_render = scene.render()
scene["red.reflectance.value"] = mi.Color3f(0.5, 0.5, 0.5)

optim = DROptimization(
    config=OptimizationConfig(output_dir="albedo-recovery", epochs=15, learning_rate=0.05),
    pipeline=RenderPipeline(spp=16),
    scene=scene,
)

report = optim.fit(GroundTruth(image=target_render))
report.plot_losses()

Scene dictionaries can be authored with mitsuba-scene-description (typed dataclasses per Mitsuba plugin); SceneDescription.from_dict accepts anything exposing to_dict().

Install groups (just pick what you need)

Development

uv sync --group dev
uv run pytest          # mitsuba tests use llvm_ad_rgb/scalar_rgb and skip when unavailable
uv run ruff check src tests

# package up for PyPI
uv build