python-drs 0.1.0__py3-none-any.whl

drs/__init__.py

@@ -0,0 +1,49 @@
+import logging
+
+__version__ = "0.1.0"
+
+# Configure a NullHandler to prevent "No handler found" warnings
+# Users of the library can configure their own logging handlers
+logging.getLogger(__name__).addHandler(logging.NullHandler())
+
+from .module import Module, DataSource
+from .engine import DRSEngine, SimulationResult
+from .variables import Variable, Level, Timer, Expression
+from .data_source import DataPoint
+from .flow import Flow
+from .telemetry import Telemetry
+from .exceptions import StateMutationError, DeadlockError
+from .config import DRSConfig, EngineConfig
+from .callbacks import Callback, ProgressBarCallback
+from .serialize import (
+    save_state,
+    load_state,
+    export_architecture,
+    save_checkpoint,
+    load_checkpoint,
+)
+
+__all__ = [
+    "DRSEngine",
+    "SimulationResult",
+    "Callback",
+    "ProgressBarCallback",
+    "Variable",
+    "Level",
+    "Timer",
+    "Expression",
+    "DataPoint",
+    "DataSource",
+    "Module",
+    "Flow",
+    "Telemetry",
+    "StateMutationError",
+    "DeadlockError",
+    "DRSConfig",
+    "EngineConfig",
+    "save_state",
+    "load_state",
+    "export_architecture",
+    "save_checkpoint",
+    "load_checkpoint",
+]

drs/_execution_context.py

@@ -0,0 +1,89 @@
+import threading
+
+
+class ExecutionContext:
+    """
+    [INTERNAL] Thread-local context tracking active modules and engines during evaluation.
+
+    Power User Note: Enforces state ownership and builds the dependency graph implicitly
+    as the simulation steps.
+    """
+
+    _local = threading.local()
+
+    @classmethod
+    def push(cls, module):
+        """
+        [INTERNAL] Push a module onto the active execution stack.
+
+        Power User Note: Used when evaluating a module's forward pass to record active scope.
+        """
+        if not hasattr(cls._local, "stack"):
+            cls._local.stack = []
+            cls._local.flow_edges = []
+        cls._local.stack.append(module)
+
+    @classmethod
+    def set_tracing(cls, enabled: bool = True):
+        """
+        [INTERNAL] Toggle symbolic tracing mode.
+
+        Power User Note: Enables AST expression recording instead of eager evaluation.
+        """
+        cls._local.tracing = enabled
+
+    @classmethod
+    def is_tracing(cls) -> bool:
+        """
+        [INTERNAL] Check if symbolic tracing mode is active.
+
+        Power User Note: Used by Variable to decide whether to return expressions.
+        """
+        return getattr(cls._local, "tracing", False)
+
+    @classmethod
+    def pop(cls):
+        """
+        [INTERNAL] Pop the top module from the execution stack.
+
+        Power User Note: Restores the caller's context after a module's execution completes.
+        """
+        cls._local.stack.pop()
+
+    @classmethod
+    def get_current(cls):
+        """
+        [INTERNAL] Retrieve the currently executing module.
+
+        Power User Note: Returns the module at the top of the stack, or None if executing externally.
+        """
+        stack = getattr(cls._local, "stack", [])
+        return stack[-1] if stack else None
+
+    @classmethod
+    def set_engine(cls, engine):
+        """
+        [INTERNAL] Bind the current DRSEngine to the context.
+
+        Power User Note: Used to provide modules access to simulation clock and configs.
+        """
+        cls._local.engine = engine
+
+    @classmethod
+    def get_engine(cls):
+        """
+        [INTERNAL] Retrieve the currently active DRSEngine.
+
+        Power User Note: Accesses thread-local storage to retrieve the running engine instance.
+        """
+        return getattr(cls._local, "engine", None)
+
+    @classmethod
+    def record_flow_edge(cls, source, target):
+        """
+        [INTERNAL] Record a transient flow edge between a source and target module.
+        """
+        if source is not None:
+            if not hasattr(cls._local, "flow_edges"):
+                cls._local.flow_edges = []
+            cls._local.flow_edges.append((source, target))

drs/callbacks.py

@@ -0,0 +1,97 @@
+from typing import TYPE_CHECKING, Optional
+import time
+
+if TYPE_CHECKING:
+    from .engine import DRSEngine, SimulationResult
+    from .variables import Variable
+
+# TODO: Should this file be a part of the public api?
+# TODO: Should we replace the "forward" system with a callback system on the levels?
+
+
+class Callback:
+    """Base class for DRS Engine callbacks.
+
+    Subclass this to hook into the simulation lifecycle.
+    """
+
+    def on_simulation_start(self, engine: "DRSEngine") -> None:
+        """Called before the simulation loop begins."""
+        pass
+
+    def on_step_start(self, engine: "DRSEngine") -> None:
+        """Called at the beginning of each simulation step, after rates are zeroed and models evaluated."""
+        pass
+
+    def on_threshold(
+        self, engine: "DRSEngine", trigger_var: "Variable", is_upper: bool
+    ) -> None:
+        """Called when a variable's threshold is the trigger for the next time step."""
+        pass
+
+    def on_deadlock(self, engine: "DRSEngine") -> None:
+        """Called immediately before a DeadlockError is raised."""
+        pass
+
+    def on_complete(self, engine: "DRSEngine", result: "SimulationResult") -> None:
+        """Called when the simulation loop has completely finished."""
+        pass
+
+
+class ProgressBarCallback(Callback):
+    """A built-in callback that uses rich to display a progress bar for the simulation."""
+
+    def __init__(self):
+        try:
+            from rich.progress import (
+                Progress,
+                TextColumn,
+                BarColumn,
+                TaskProgressColumn,
+                TimeElapsedColumn,
+                TimeRemainingColumn,
+            )
+        except ImportError:
+            raise ImportError(
+                "The 'rich' package is required for the ProgressBarCallback. Install it with 'pip install rich'."
+            )
+
+        self.progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            redirect_stdout=False,
+            redirect_stderr=False,
+        )
+        self.task_id = None
+        self.max_time = None
+
+    def on_simulation_start(self, engine: "DRSEngine") -> None:
+        self.progress.start()
+        # the engine doesn't technically know its max_time until run() is called,
+        # but max_time is passed to run(). We can extract it if we attach it to the engine temporarily,
+        # or we just assume we'll get it from engine if we store it.
+        # Wait, DRSEngine.run(max_time) takes max_time. So it's not known here unless we store it on the engine.
+        # Let's assume engine._current_max_time is set before calling this.
+        self.max_time = getattr(engine, "_current_max_time", None)
+
+        if self.max_time:
+            self.task_id = self.progress.add_task(
+                "[cyan]Simulating...", total=self.max_time
+            )
+        else:
+            self.task_id = self.progress.add_task("[cyan]Simulating...", total=None)
+
+    def on_step_start(self, engine: "DRSEngine") -> None:
+        if self.task_id is not None and self.max_time:
+            self.progress.update(self.task_id, completed=engine.current_time)
+
+    def on_complete(self, engine: "DRSEngine", result: "SimulationResult") -> None:
+        if self.task_id is not None and self.max_time:
+            self.progress.update(self.task_id, completed=self.max_time)
+        self.progress.stop()
+
+    def on_deadlock(self, engine: "DRSEngine") -> None:
+        self.progress.stop()

drs/config.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict
+
+
+@dataclass
+class DRSConfig:
+    """Base configuration class for DRS Modules.
+
+    Subclass this using @dataclass to define strictly-typed configuration
+    blocks for your models, allowing IDE autocomplete and clear parameterization.
+    """
+
+    pass
+
+
+@dataclass
+class EngineConfig(DRSConfig):
+    """Configuration for the DRS Engine."""
+
+    max_step_size: float = (
+        0.5  # TODO: should change this to float("inf") but for parity dont, it seems to cause behaviour changes, theory is that its because of precision propogation over longer time horizons. Why do we even need a max step size?
+    )
+    max_deadlock_steps: int = 20
+    max_time: float = None
+    strict_mode: bool = False

drs/data_source.py

@@ -0,0 +1,49 @@
+from typing import Any
+
+"""
+Generic data point for data yielded by DataSource.
+
+Note on alternatives:
+If you prefer standard Python semantics, you could potentially bypass `DataSource` entirely 
+and use a plain Python generator (e.g. `yield DataPoint(...)`). However, the `DataSource` 
+class is retained to seamlessly integrate with `Module`, making it a first-class citizen 
+in the simulation graph and tracking execution contexts and telemetry.
+"""
+
+
+class DataPoint:
+    """A single batch of data yielded by a DataSource.
+
+    Fields are accessed as attributes. Any keyword argument passed to
+    ``__init__`` is stored and accessible via ``.name``::
+
+        point = DataPoint(weight=40000.0, priority=1)
+        point.weight    # → 40000.0
+        point.priority  # → 1
+    """
+
+    def __init__(self, **kwargs: Any):
+        """
+        Initialize a new DataPoint with arbitrary attributes.
+
+        Args:
+            **kwargs: Arbitrary keyword arguments that will be stored and exposed 
+                as attributes on the DataPoint instance.
+        """
+        self._data = dict(kwargs)
+        self._source = None
+
+    def __getattr__(self, name):
+        if name.startswith("_"):
+            raise AttributeError(name)
+        try:
+            return self._data[name]
+        except KeyError:
+            raise AttributeError(
+                f"'{type(self).__name__}' has no field '{name}'. "
+                f"Available fields: {list(self._data)}"
+            )
+
+    def __repr__(self):
+        items = ", ".join(f"{k}={v}" for k, v in self._data.items())
+        return f"DataPoint({items})"