Extending the Object Model#

Use this page when an extension changes a model class, introduces runtime-visible state, or adds a registered or polymorphic object type. It assumes the hierarchy and data representation described in Simulation Compilation and Runtime Data Layout and applies those designs as implementation recipes.

When the new state is consumed during particle transport, continue with Writing Numba-Compatible Transport Code for type, dispatch, allocation, CPU/GPU compatibility, and verification guidance.

Choose the Extension Type#

Choose the narrowest extension that represents the new concept:

Change

Starting point

Use when

Example

Add a field

Existing class

The concept already belongs to an existing model or configuration object.

Add a new source parameter to Source.

Add embedded state

MCDCBase

The state belongs to one parent and does not need an independently addressable entry in a simulation registry.

Add technique settings owned by Simulation.

Add a registered category

MCDCObject

Transport must refer to independently registered instances by simulation-local ID.

Add a new Surface-like category with its own collection.

Add a representation to an existing category

MCDCPolymorphic

The extension shares a category interface but needs a distinct packed layout and dispatch code.

Add a concrete MeshBase representation.

Prefer adding a subtype to an existing polymorphic family over creating a new registered category when the new object has the same conceptual role. For example, implement a new tally estimator as a Tally subtype.

The Common Class Contract#

Every MC/DC model class must follow the conventions used by the compiler and Numba-layer generator.

label

Provide a unique, stable, lower-case label such as structured_mesh. The label names generated structured layouts and the corresponding modules under mcdc_get and mcdc_set.

class MeshStructured(MeshBase):
    label = "structured_mesh"
sub_type

Give every concrete MCDCPolymorphic subclass a unique named integer constant within its family. The shared base uses sub_type = -1.

class MeshStructured(MeshBase):
    sub_type = MESH_STRUCTURED
Type annotations

Annotate every field that must be represented at runtime. Annotations define scalar fields, embedded structures, object-ID references, and variable-length payloads.

active: bool
translation: Annotated[NDArray[float64], (3,)]
move_velocities: Annotated[NDArray[float64], ("N_move", 3)]
surfaces: list[Surface]
Initialization

Assign every runtime-visible field a valid initial value. Subclasses of MCDCObject must call super().__init__() so ID is initialized; subclasses of MCDCPolymorphic must do the same so both ID and sub_ID are initialized.

def __init__(self, name, boundaries):
    super().__init__()
    self.name = name
    self.boundaries = np.asarray(boundaries, dtype=float64)
non_numba

List Python-only fields that should not be traversed or packed automatically. The class must explicitly convert any required information from those fields into annotated runtime-visible fields before packing.

For example, Cell keeps its expressive region and fill objects on the Python side, then derives RPN tokens, a fill-type code, and a fill ID for transport:

non_numba = ["region", "fill"]

region: Region
fill: Material | Universe | Lattice | None
region_RPN_tokens: list[int]
fill_type: int
fill_ID: int
Compilation hook

Use the inherited _compile_into_simulation implementation unless the class must canonicalize objects, compile excluded references, derive fields from assigned object IDs, or otherwise finalize state owned by that object.

def _compile_into_simulation(self, simulation):
    if not super()._compile_into_simulation(simulation):
        return False

    self.reference._compile_into_simulation(simulation)
    self.reference_ID = self.reference.ID
    return True

If Python-only members must be canonicalized before ordinary traversal, guard the work with compile_ID and then call super exactly once:

def _compile_into_simulation(self, simulation):
    if self.compile_ID == simulation.compile_ID:
        return False

    self._resolve_python_inputs(simulation)

    if not super()._compile_into_simulation(simulation):
        return False

    self._derive_post_registration_fields()
    return True

Do not assign compile_ID, ID, or sub_ID manually.

Model-wide finalization

Object hooks should not normalize or coordinate unrelated registries. When a value requires the complete discovered model, coordinate it once in Simulation._finalize_compilation instead. Source-probability normalization, particle-bank capacities, and settings derived from the complete material or tally collections are examples of model-wide finalization. Explicitly compile any new runtime-visible object introduced during this phase because ordinary recursive discovery has already occurred.

compile_simulation orchestrates recursive discovery and calls this model-wide finalization phase. Change the compiler orchestration only when adding a new compilation phase or registered category. Do not place model-specific finalization in mcdc.main.prepare; that function is reserved for framework-level packing, resource allocation, backend configuration, and external runtime state.

Represent Fields Deliberately#

The layer generator interprets annotations according to the field’s role:

Example annotation

Runtime representation

Transport access

active: bool

Scalar structured field

simulation["technique"]["new_technique"]["active"]

translation: Annotated[NDArray[float64], (3,)]

Fixed-size embedded array

cell["translation"][axis]

energy: NDArray[float64]

Offset and length plus values in data

mcdc_get.tally.energy(index, tally, data)

move_velocities: Annotated[NDArray[float64], ("N_move", 3)]

Offset, length, and shape metadata plus flattened values

mcdc_get.surface.move_velocities(move, axis, surface, data)

energy_pmf: DistributionPMF

Simulation-local object ID

simulation["distributions"][source["energy_pmf_ID"]]

collision_tallies: list[TallyCollision]

Count and offset to IDs stored in data

mcdc_get.cell.collision_tally_IDs(index, cell, data)

Use an integer-only shape when an array is always the same size. Use symbolic dimensions when a shape depends on the model. Do not store a Python reference in transport-visible state; annotate it as an MCDCObject or polymorphic base so the packed layer records an ID.

Adding a Field to an Existing Class#

  1. Add the annotation to the class that owns the concept.

  2. Initialize the field for every construction path.

  3. Decide whether it is fixed-size, variable-length, an embedded MCDCBase, or an MCDCObject reference.

  4. Update any compile hook that derives the field or converts a Python-only input into its runtime representation.

  5. Consume the field through direct structured access or its generated mcdc_get and mcdc_set helpers.

  6. Update the public docstring and API documentation when users can configure the field.

For example, a new variable-length energy_bias field on Source requires an annotation and initialized array on the model class:

# In Source annotations
energy_bias: NDArray[float64]

# In Source.__init__
self.energy_bias = np.asarray(energy_bias, dtype=float64)

Transport then reads one value through the generated accessor:

bias = mcdc_get.source.energy_bias(index, source, data)

Avoid adding parallel state in several classes. If a value belongs to the simulation as a whole, place it in Simulation or one of its embedded configuration objects and pass or access that representation consistently.

Adding Embedded MCDCBase State#

Use MCDCBase for configuration or runtime state that is owned by one parent. A minimal class has a label, annotated fields, and initialized values:

class NewTechnique(MCDCBase):
    label = "new_technique"

    active: bool
    strength: float

    def __init__(self) -> None:
        super().__init__()
        self.active = False
        self.strength = 1.0

Add an annotated field for the object to its owner and instantiate it with the owner. Simulation-wide transport techniques belong to the Technique aggregate, which is itself owned by Simulation. The embedded object participates in recursive compilation but does not require a registry branch or object ID. This Python ownership hierarchy is preserved under the packed simulation["technique"] record.

class Technique(MCDCBase):
    new_technique: NewTechnique

    def __init__(self):
        self.new_technique = NewTechnique()

The corresponding user and runtime interfaces are simulation.technique.new_technique(...) and simulation["technique"]["new_technique"].

If the embedded object refers to an MCDCObject, annotate that reference. The default traversal will register the referenced object when compilation reaches the embedded configuration.

Adding a New MCDCObject Category#

A genuinely new registered category requires coordinated changes:

  1. Define the class with a unique label, annotations, initialized fields, and a call to super().__init__().

  2. Add the category collection to Simulation annotations and initialize or reset it in Simulation._reset_model.

  3. Import the category in mcdc/code_factory/python_objects_compiler.py and add an isinstance branch in register_object that selects its simulation collection.

  4. Ensure the object is reachable from an existing simulation root or add an explicit root and compilation step.

  5. Ensure the class’s module is imported before numba_layers_generator.py discovers the classes. A public class is normally imported through mcdc/__init__.py; an internal class must be imported by another module in the compilation path.

  6. Add the collection and lookup behavior required by transport.

For example, the class and its simulation collection begin with:

class Detector(MCDCObject):
    label = "detector"

    name: str
    response: NDArray[float64]

    def __init__(self, name, response):
        super().__init__()
        self.name = name
        self.response = np.asarray(response, dtype=float64)


class Simulation(MCDCBase):
    detectors: list[Detector]

The compiler then selects that collection explicitly:

elif isinstance(object_, Detector):
    object_list = simulation.detectors

Do not add a fallback registration branch that silently accepts unknown objects. An explicit category branch keeps registry ownership and runtime layout reviewable.

Adding a Polymorphic Subtype#

Adding a concrete subtype to an existing family is more localized than adding a category:

  1. Add a unique integer constant for the subtype.

  2. Inherit from the existing polymorphic base, such as MeshBase or Tally.

  3. Set a unique concrete label and the new sub_type constant.

  4. Call the base initializer so shared fields, ID, and sub_ID are initialized.

  5. Annotate and initialize subtype-specific fields.

  6. Import the subtype before layer generation and expose it from mcdc/__init__.py when it is public.

  7. Add transport dispatch for the new sub_type and implement the subtype-specific behavior.

For example, the structural part of a mesh subtype follows this pattern:

class MeshNew(MeshBase):
    label = "new_mesh"
    sub_type = MESH_NEW

    boundaries: NDArray[float64]

    def __init__(self, boundaries, name="") -> None:
        super().__init__(name)
        self.boundaries = np.asarray(boundaries, dtype=float64)

No new register_object branch is needed for a subtype of an already registered family. The existing isinstance(..., MeshBase) or corresponding category check places it in the base collection, while sub_type and sub_ID connect it to its concrete packed collection.

Rebuilding Numba Support#

Changes to runtime-visible annotations or object types under mcdc/object_ require rebuilding the generated Numba support. The annotations are the source of truth; do not edit mcdc/numba_types.py, mcdc_get, or mcdc_set directly to introduce a field.

Run the rebuild script after changing the object model:

python mcdc/code_factory/rebuild_numba_support.py

Then verify the access pattern predicted by the field representation chosen above, and commit the regenerated files together with the object model change.

During an active mcdc/object_ edit-test cycle, add -r (or --rebuild) to the test input-deck command instead. MC/DC then rebuilds the generated Numba support during package initialization, after loading the full object model and before importing the generated files. Developers who are not changing the object model do not need this option.

Important

Within one MPI launch, rank zero performs the rebuild and the other ranks wait. Independent launches do not share that barrier, so do not use -r or --rebuild from concurrent jobs that share an MC/DC source tree.

Rebuilding refreshes the shared runtime schema. Problem-dependent dtypes and prepared runtime state are still created separately for each simulation. See Generated Numba Support and Problem-Dependent Dtypes for those lifetimes and the import order.

After rebuilding, a variable-length Detector.response field should produce element accessors associated with the detector label:

value = mcdc_get.detector.response(index, detector, data)
mcdc_set.detector.response(index, detector, data, new_value)

A fixed-size field such as Cell.translation remains embedded and is accessed directly:

value = cell["translation"][axis]

Public API and Documentation#

For a user-facing class or constructor:

  1. Export the class from mcdc/__init__.py.

  2. Add it to the appropriate autosummary group in docs/source/reference/python_api/index.rst.

  3. Document parameters, units, defaults, constraints, and at least one usable example in the class docstring.

  4. Update the User Guide when the extension changes how users construct or run a model.

Keep internal helper classes under mcdc.object_ and import them explicitly in the compilation path. Export only classes that form part of the public API.

For example, a public Detector is re-exported from the package and listed by its qualified name in the API autosummary:

from mcdc.object_.detector import Detector
~mcdc.Detector

Verification Checklist#

An object model extension should verify all affected layers:

  • Construction accepts valid input and rejects invalid shapes or types.

  • Compilation discovers the object from the intended root.

  • Shared references register once, and recompilation produces a valid new snapshot.

  • Packed fields, object IDs, offsets, and generated accessors contain the expected values.

  • Python and Numba-CPU modes produce equivalent behavior.

  • GPU execution is covered when the changed transport path supports GPUs.

  • Public examples compile under the example validator when the API changes.

  • API and developer documentation build without warnings.

Add focused unit tests near test/unit/test_object_compilation.py for compilation behavior and near the relevant transport tests for runtime behavior. Use Example Validation when an extension changes public examples.