caskade 0.6.1__py3-none-any.whl

caskade/__init__.py

@@ -0,0 +1,48 @@
+from ._version import version as VERSION  # noqa
+
+from .base import Node
+from .context import ActiveContext, ValidContext, OverrideParam
+from .decorators import forward
+from .module import Module
+from .param import Param
+from .tests import test
+from .errors import (
+    CaskadeException,
+    GraphError,
+    NodeConfigurationError,
+    ParamConfigurationError,
+    ParamTypeError,
+    ActiveStateError,
+    FillDynamicParamsError,
+    FillDynamicParamsTensorError,
+    FillDynamicParamsSequenceError,
+    FillDynamicParamsMappingError,
+)
+from .warnings import CaskadeWarning, InvalidValueWarning
+
+
+__version__ = VERSION
+__author__ = "Connor Stone and Alexandre Adam"
+
+__all__ = (
+    "Node",
+    "Module",
+    "Param",
+    "ActiveContext",
+    "ValidContext",
+    "OverrideParam",
+    "forward",
+    "test",
+    "CaskadeException",
+    "GraphError",
+    "NodeConfigurationError",
+    "ParamConfigurationError",
+    "ParamTypeError",
+    "ActiveStateError",
+    "FillDynamicParamsError",
+    "FillDynamicParamsTensorError",
+    "FillDynamicParamsSequenceError",
+    "FillDynamicParamsMappingError",
+    "CaskadeWarning",
+    "InvalidValueWarning",
+)

caskade/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.6.1'
+__version_tuple__ = version_tuple = (0, 6, 1)

caskade/base.py

@@ -0,0 +1,224 @@
+from typing import Optional, Union
+
+from .errors import GraphError, NodeConfigurationError
+
+
+class Node(object):
+    """
+    Base graph node class for ``caskade`` objects.
+
+    The ``Node`` object is the base class for all ``caskade`` objects. It is used to
+    construct the directed acyclic graph (DAG). The primary function of the
+    ``Node`` object is to manage the parent-child relationships between nodes in
+    the graph. There is limited functionality for the ``Node`` object, though it
+    implements the base versions of the ``active`` state and ``to`` /
+    ``update_graph`` methods. The ``active`` state is used to communicate
+    through the graph that the simulator is currently running. The ``to`` method
+    is used to move and/or cast the values of the parameter. The ``update_graph``
+    method is used signal all parents that the graph below them has changed.
+
+    Examples
+    --------
+
+    Example making some ``Node`` objects and then linking/unlinking them::
+
+       n1 = Node()
+       n2 = Node()
+       n1.link("subnode", n2) # link n2 as a child of n1, may use any str as the key
+       n1.unlink("subnode") # alternately n1.unlink(n2) to unlink by object
+    """
+
+    graphviz_types = {"node": {"style": "solid", "color": "black", "shape": "circle"}}
+
+    def __init__(self, name: Optional[str] = None):
+        if name is None:
+            name = self.__class__.__name__
+        if not isinstance(name, str):
+            raise NodeConfigurationError(f"{self.__class__.__name__} name must be a string")
+        if "|" in name:
+            raise NodeConfigurationError(f"{self.__class__.__name__} cannot contain '|'")
+        self._name = name
+        self._children = {}
+        self._parents = set()
+        self._active = False
+        self._type = "node"
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def children(self) -> dict:
+        return self._children
+
+    @property
+    def parents(self) -> set:
+        return self._parents
+
+    def link(self, key: Union[str, "Node"], child: Optional["Node"] = None):
+        """Link the current ``Node`` object to another ``Node`` object as a child.
+
+        Parameters
+        ----------
+        key: (Union[str, Node])
+            The key to link the child node with.
+        child: (Optional[Node], optional)
+            The child ``Node`` object to link to. Defaults to None in which
+            case the key is used as the child.
+
+        Examples
+        --------
+
+        Example making some ``Node`` objects and then linking/unlinking them. demonstrating multiple ways to link/unlink::
+
+            n1 = Node()
+            n2 = Node()
+
+            n1.link("subnode", n2) # may use any str as the key
+            n1.unlink("subnode")
+
+            # Alternately, link by object
+            n1.link(n2)
+            n1.unlink(n2)
+        """
+        if child is None:
+            child = key
+            key = child.name
+        # Avoid double linking to the same object
+        if key in self.children:
+            raise GraphError(f"Child key {key} already linked to parent {self.name}")
+        if child in self.children.values():
+            raise GraphError(f"Child {child.name} already linked to parent {self.name}")
+        # avoid cycles
+        if self in child.topological_ordering():
+            raise GraphError(
+                f"Linking {child.name} to {self.name} would create a cycle in the graph"
+            )
+
+        self._children[key] = child
+        child._parents.add(self)
+        self.update_graph()
+
+    def unlink(self, key: Union[str, "Node"]):
+        """Unlink the current ``Node`` object from another ``Node`` object which is a child."""
+        if isinstance(key, Node):
+            for node in self.children:
+                if self.children[node] == key:
+                    key = node
+                    break
+        self._children[key]._parents.remove(self)
+        self._children[key].update_graph()
+        del self._children[key]
+        self.update_graph()
+
+    def topological_ordering(self, with_type: Optional[str] = None) -> tuple["Node"]:
+        """Return a topological ordering of the graph below the current node."""
+        ordering = [self]
+        for node in self.children.values():
+            for subnode in node.topological_ordering():
+                if subnode not in ordering:
+                    ordering.append(subnode)
+        if with_type is None:
+            return tuple(ordering)
+        return tuple(filter(lambda n: n._type == with_type, ordering))
+
+    def update_graph(self):
+        """Triggers a call to all parents that the graph below them has been
+        updated. The base ``Node`` object does nothing with this information, but
+        other node types may use this to update internal state."""
+        for parent in self.parents:
+            parent.update_graph()
+
+    @property
+    def active(self) -> bool:
+        return self._active
+
+    @active.setter
+    def active(self, value: bool):
+        # Avoid unnecessary updates
+        if self._active == value:
+            return
+
+        # Set self active level
+        self._active = value
+
+        # Propagate active level to children
+        for child in self._children.values():
+            child.active = value
+
+    def to(self, device=None, dtype=None):
+        """
+        Moves and/or casts the PyTorch values of the ``Node``.
+
+        Parameters
+        ----------
+        device: (Optional[torch.device], optional)
+            The device to move the values to. Defaults to None.
+        dtype: (Optional[torch.dtype], optional)
+            The desired data type. Defaults to None.
+        """
+
+        for child in self.children.values():
+            child.to(device=device, dtype=dtype)
+
+        return self
+
+    def graphviz(self, top_down=True) -> "graphviz.Digraph":
+        """Return a graphviz object representing the graph below the current
+        node in the DAG.
+
+        Parameters
+        ----------
+        top_down: (bool, optional)
+            Whether to draw the graph top-down (current node at top) or
+            bottom-up (current node at bottom). Defaults to True.
+        """
+        import graphviz
+
+        components = set()
+
+        def add_node(node, dot):
+            if node in components:
+                return
+            dot.attr("node", **node.graphviz_types[node._type])
+            dot.node(str(id(node)), repr(node))
+            components.add(node)
+
+            for child in node.children.values():
+                add_node(child, dot)
+                if top_down:
+                    dot.edge(str(id(node)), str(id(child)))
+                else:
+                    dot.edge(str(id(child)), str(id(node)))
+
+        dot = graphviz.Digraph(strict=True)
+        add_node(self, dot)
+        return dot
+
+    def graph_dict(self) -> dict[str, dict]:
+        """Return a dictionary representation of the graph below the current
+        node."""
+        graph = {
+            f"{self.name}|{self._type}": {},
+        }
+        for node in self.children.values():
+            graph[f"{self.name}|{self._type}"].update(node.graph_dict())
+        return graph
+
+    def graph_print(self, dag: dict, depth: int = 0, indent: int = 4, result: str = "") -> str:
+        """Print the graph dictionary in a human-readable format."""
+        for key in dag:
+            result = f"{result}{' ' * indent * depth}{key}\n"
+            result = self.graph_print(dag[key], depth + 1, indent, result) + "\n"
+        if result:  # remove trailing newline
+            result = result[:-1]
+        return result
+
+    def __str__(self) -> str:
+        return self.graph_print(self.graph_dict())
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.name})"
+
+    def __getitem__(self, key: str) -> "Node":
+        return self.children[key]

caskade/context.py

@@ -0,0 +1,73 @@
+from .module import Module
+from .param import Param
+
+
+class ActiveContext:
+    """
+    Context manager to activate a module for a simulation. Only inside an
+    ActiveContext is it possible to fill/clear the dynamic and live parameters.
+    """
+
+    def __init__(self, module: Module, active: bool = True):
+        self.module = module
+        self.active = active
+
+    def __enter__(self):
+        self.outer_active = self.module.active
+        if self.outer_active and not self.active:
+            self.outer_params = list(p.value for p in self.module.dynamic_params)
+            self.module.clear_params()
+        self.module.active = self.active
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if not self.outer_active and self.active:
+            self.module.clear_params()
+        self.module.active = self.outer_active
+        if self.outer_active and not self.active:
+            self.module.fill_params(self.outer_params)
+
+
+class ValidContext:
+    """
+    Context manager to set valid values for parameters. Only inside a
+    ValidContext will parameters automatically be assumed valid.
+    """
+
+    def __init__(self, module: Module):
+        self.module = module
+
+    def __enter__(self):
+        self.module.valid_context = True
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.module.valid_context = False
+
+
+class OverrideParam:
+    """
+    Context manager to override a parameter value. Only inside an
+    OverrideParam will the parameter be set to the new value.
+    """
+
+    def __init__(self, param, value):
+        self.param = param
+        self.value = value
+
+    def __enter__(self):
+        # Store the old value
+        self.old_values = {str(id(self.param)): self.param._value}
+        # Set the new value
+        self.param._value = self.value
+        # Clear the pointer values as they may have updated
+        for node in self.param.parents:
+            if isinstance(node, Param) and node.pointer:
+                self.old_values[str(id(node))] = node._value
+                node._value = None
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Reset the old value
+        self.param._value = self.old_values[str(id(self.param))]
+        # Clear the pointer values as they may have updated
+        for node in self.param.parents:
+            if isinstance(node, Param) and node.pointer:
+                node._value = self.old_values[str(id(node))]

caskade/decorators.py

@@ -0,0 +1,75 @@
+import inspect
+import functools
+
+from .context import ActiveContext
+
+__all__ = ("forward",)
+
+
+def _get_arguments(method):
+    sig = inspect.signature(method)
+    return tuple(sig.parameters.keys())
+
+
+def forward(method):
+    """
+    Decorator to define a forward method for a module.
+
+    Parameters
+    ----------
+    method: (Callable)
+        The forward method to be decorated.
+
+    Examples
+    --------
+    Standard usage of the forward decorator::
+
+        class ExampleSim(Module):
+            def __init__(self, a, b, c):
+                super().__init__("example_sim")
+                self.a = a
+                self.b = Param("b", b)
+                self.c = Param("c", c)
+
+            @forward
+            def example_func(self, x, b=None):
+                return x + self.a + b
+
+        E = ExampleSim(a=1, b=None, c=3)
+        print(E.example_func(4, params=[5]))
+        # Output: 10
+
+    Returns
+    -------
+    Callable
+        The decorated forward method.
+    """
+
+    # Get arguments from function signature
+    method_params = _get_arguments(method)
+
+    @functools.wraps(method)
+    def wrapped(self, *args, **kwargs):
+        if self.active:
+            kwargs = {**self.fill_kwargs(method_params), **kwargs}
+            return method(self, *args, **kwargs)
+
+        # Extract params from the arguments
+        if len(self.dynamic_params) == 0:
+            params = {}
+        elif "params" in kwargs:
+            params = kwargs.pop("params")
+        elif args:
+            params = args[-1]
+            args = args[:-1]
+        else:
+            raise ValueError(
+                f"Params must be provided for a top level @forward method. Either by keyword 'method(params=params)' or as the last positional argument 'method(a, b, c, params)'"
+            )
+
+        with ActiveContext(self):
+            self.fill_params(params)
+            kwargs = {**self.fill_kwargs(method_params), **kwargs}
+            return method(self, *args, **kwargs)
+
+    return wrapped

caskade/errors.py

@@ -0,0 +1,99 @@
+from math import prod
+from textwrap import dedent
+
+
+class CaskadeException(Exception):
+    """Base class for all exceptions in ``caskade``."""
+
+
+class GraphError(CaskadeException):
+    """Class for graph exceptions in ``caskade``."""
+
+
+class NodeConfigurationError(CaskadeException):
+    """Class for node configuration exceptions in ``caskade``."""
+
+
+class ParamConfigurationError(NodeConfigurationError):
+    """Class for parameter configuration exceptions in ``caskade``."""
+
+
+class ParamTypeError(CaskadeException):
+    """Class for exceptions related to the type of a parameter in ``caskade``."""
+
+
+class ActiveStateError(CaskadeException):
+    """Class for exceptions related to the active state of a node in ``caskade``."""
+
+
+class FillDynamicParamsError(CaskadeException):
+    """Class for exceptions related to filling dynamic parameters in ``caskade``."""
+
+
+class FillDynamicParamsTensorError(FillDynamicParamsError):
+    """Class for exceptions related to filling dynamic parameters with a tensor in ``caskade``."""
+
+    def __init__(self, name, input_params, dynamic_params):
+        fullnumel = sum(max(1, prod(p.shape)) for p in dynamic_params)
+        message = dedent(
+            f"""
+            For flattened Tensor input, the (last) dim of the Tensor should
+            equal the sum of all flattened dynamic params ({fullnumel}).
+            Input params shape {input_params.shape} does not match dynamic
+            params shape of: {name}. 
+            
+            Registered dynamic params (name: shape):
+            {', '.join(f"{repr(p)}: {str(p.shape)}" for p in dynamic_params)}"""
+        )
+        super().__init__(message)
+
+
+class FillDynamicParamsSequenceError(FillDynamicParamsError):
+    """Class for exceptions related to filling dynamic parameters with a sequence (list, tuple, etc.) in ``caskade``."""
+
+    def __init__(self, name, input_params, dynamic_params, dynamic_modules):
+        message = dedent(
+            f"""
+            Input params length ({len(input_params)}) does not match dynamic
+            params length ({len(dynamic_params)}) or number of dynamic
+            modules ({len(dynamic_modules)}) of: {name}.
+            
+            Registered dynamic modules: 
+            {', '.join(repr(m) for m in dynamic_modules)}
+
+            Registered dynamic params:
+            {', '.join(repr(p) for p in dynamic_params)}"""
+        )
+        super().__init__(message)
+
+
+class FillDynamicParamsMappingError(FillDynamicParamsError):
+    """Class for exceptions related to filling dynamic parameters with a mapping (dict) in ``caskade``."""
+
+    def __init__(self, name, children, dynamic_modules, missing_key=None, missing_param=None):
+        if missing_key is not None:
+            message = dedent(
+                f"""
+                Input params key "{missing_key}" not found in dynamic modules or children of: {name}. 
+                
+                Registered dynamic modules: 
+                {', '.join(repr(m) for m in dynamic_modules)}
+
+                Registered dynamic children:
+                {', '.join(repr(c) for c in children.values() if c.dynamic)}"""
+            )
+        else:
+            message = dedent(
+                f"""
+                Dynamic param "{missing_param.name}" not filled with given input params dict passed to {name}.
+
+                Dynamic param parent(s):
+                {', '.join(repr(p) for p in missing_param.parents)}
+                
+                Registered dynamic modules: 
+                {', '.join(repr(m) for m in dynamic_modules)}
+
+                Registered dynamic children:
+                {', '.join(repr(c) for c in children.values() if c.dynamic)}"""
+            )
+        super().__init__(message)