caskade 0.0.1__py3-none-any.whl

caskade/__init__.py

@@ -0,0 +1,14 @@
+from ._version import version as VERSION  # noqa
+
+from .base import Node
+from .context import ActiveContext
+from .decorators import forward
+from .module import Module
+from .param import Param, LiveParam
+from .tests import test
+
+
+__version__ = VERSION
+__author__ = "Connor and Alexandre"
+
+__all__ = ("Node", "Module", "Param", "LiveParam", "ActiveContext", "forward")

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.0.1'
+__version_tuple__ = version_tuple = (0, 0, 1)

caskade/base.py

@@ -0,0 +1,132 @@
+from typing import Optional
+
+import torch
+
+
+class Node:
+    """
+    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_dynamic_params` 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_dynamic_params` method is used by `Module` objects to keep track of
+    all dynamic `Param` objects below them in the graph.
+
+    Examples
+    --------
+    ``` python
+    n1 = Node("node1")
+    n2 = Node("node2")
+    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
+    """
+
+    def __init__(self, name):
+        assert isinstance(name, str), f"{self.__class__.__name__} name must be a string"
+        assert "|" not in name, 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, child):
+        # Avoid double linking to the same object
+        if key in self.children:
+            raise ValueError(f"Child key {key} already linked to parent {self.name}")
+        for ownchild in self.children.values():
+            if ownchild == child:
+                raise ValueError(f"Child {child.name} already linked to parent {self.name}")
+
+        self._children[key] = child
+        child._parents.add(self)
+        self.update_dynamic_params()
+
+    def unlink(self, key):
+        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_dynamic_params()
+        del self._children[key]
+        self.update_dynamic_params()
+
+    def topological_ordering(self, with_type=None) -> tuple:
+        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_dynamic_params(self):
+        for parent in self.parents:
+            parent.update_dynamic_params()
+
+    @property
+    def active(self) -> bool:
+        return self._active
+
+    @active.setter
+    def active(self, value):
+        # 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: Optional[torch.device] = None, dtype: Optional[torch.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)
+
+    def graph_dict(self) -> dict:
+        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 __str__(self) -> str:
+        return str(self.graph_dict())
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.name})"

caskade/context.py

@@ -0,0 +1,21 @@
+from typing import Union, Mapping, Sequence
+
+from torch import Tensor
+
+from .module import Module
+
+
+class ActiveContext:
+    def __init__(
+        self, module: Module, params: Union[Sequence[Tensor], Mapping[str, Tensor], Tensor]
+    ):
+        self.module = module
+        self.params = params
+
+    def __enter__(self):
+        self.module.active = True
+        self.module.fill_params(self.params)
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.module.clear_params()
+        self.module.active = False

caskade/decorators.py

@@ -0,0 +1,50 @@
+import inspect
+import functools
+
+from .context import ActiveContext
+
+
+def forward(method):
+    """
+    Decorator to define a forward method for a module.
+
+    Parameters
+    ----------
+    method: (Callable)
+        The forward method to be decorated.
+
+    Returns
+    -------
+    Callable
+        The decorated forward method.
+    """
+
+    # Get kwargs from function signature
+    method_kwargs = []
+    for arg in inspect.signature(method).parameters.values():
+        if arg.default is not arg.empty:
+            method_kwargs.append(arg.name)
+
+    @functools.wraps(method)
+    def wrapped(self, *args, **kwargs):
+        if self.active:
+            kwargs.update(self.fill_kwargs(method_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.pop(0)
+        else:
+            raise ValueError(
+                f"Params must be provided for dynamic modules. Expected {len(self.dynamic_params)} params."
+            )
+
+        with ActiveContext(self, params):
+            kwargs.update(self.fill_kwargs(method_kwargs))
+            return method(self, *args, **kwargs)
+
+    return wrapped

caskade/module.py

@@ -0,0 +1,99 @@
+from typing import Sequence, Mapping
+
+from torch import Tensor
+
+from .base import Node
+from .param import Param, LiveParam
+
+
+class Module(Node):
+
+    def __init__(self, name):
+        super().__init__(name=name)
+        self.dynamic_params = ()
+        self.live_params = ()
+        self._type = "module"
+        self._batch = False
+
+    @property
+    def batch(self) -> bool:
+        return self._batch
+
+    @batch.setter
+    def batch(self, value):
+        assert isinstance(value, bool)
+        self._batch = value
+
+    def update_dynamic_params(self):
+        super().update_dynamic_params()
+        self.dynamic_params = tuple(self.topological_ordering("dynamic"))
+        self.live_params = tuple(self.topological_ordering("live"))
+
+    def fill_params(self, params):
+        assert self.active, "Module must be active to fill params"
+
+        if isinstance(params, Tensor):
+            if self.batch:
+                B = params.shape[0]
+            pos = 0
+            for param in self.dynamic_params:
+                try:
+                    size = param.shape.numel()
+                except AttributeError:
+                    raise ValueError(
+                        f"Param {param.name} has no shape. dynamic parameters must have a shape to use Tensor input."
+                    )
+                if self.batch:
+                    param.value = params[:, pos : pos + size].view((B,) + param.shape)
+                    pos += size * B
+                else:
+                    param.value = params[pos : pos + size].view(param.shape)
+                    pos += size
+        elif isinstance(params, Sequence):
+            if len(params) == len(self.dynamic_params):
+                for param, value in zip(self.dynamic_params, params):
+                    param.value = value
+            else:
+                raise ValueError(
+                    f"Input params length ({len(params)}) does not match dynamic params length ({len(self.dynamic_params)})"
+                )
+        elif isinstance(params, Mapping):
+            for key in params:
+                if key in self.children:
+                    if isinstance(self.children[key], Param):
+                        self.children[key].value = params[key]
+                    elif isinstance(self.children[key], Module):
+                        self.children[key].fill_params(params[key])
+                    else:
+                        raise ValueError(f"Key {key} type {type(self.children[key])} not supported")
+                else:
+                    raise ValueError(f"Key {key} not found in {self.name} children")
+        else:
+            raise ValueError(
+                f"Input params type {type(params)} not supported. Should be Tensor, Sequence or Mapping."
+            )
+
+    def clear_params(self):
+        assert self.active, "Module must be active to clear params"
+
+        for param in self.dynamic_params:
+            param.value = None
+
+        for param in self.live_params:
+            param.value = LiveParam
+
+    def fill_kwargs(self, keys) -> dict[str, Tensor]:
+        return {key: getattr(self, key).value for key in keys}
+
+    def __setattr__(self, key, value):
+        try:
+            if key in self.children and isinstance(self.children[key], Param):
+                self.children[key].value = value
+                return
+            if isinstance(value, Node):
+                self.link(key, value)
+                self.update_dynamic_params()
+
+            super().__setattr__(key, value)
+        except AttributeError:
+            super().__setattr__(key, value)

caskade/param.py

@@ -0,0 +1,147 @@
+from typing import Optional, Union, Callable
+
+import torch
+from torch import Tensor
+
+from .base import Node
+
+__all__ = ("Param", "LiveParam")
+
+
+class LiveParamBase:
+    """Placeholder to identify a parameter as live updating. Like `None` there
+    exists only one instance of this class."""
+
+    pass
+
+
+LiveParam = LiveParamBase()
+
+
+class Param(Node):
+    """
+    Node to represent a parameter in the graph.
+
+    The `Param` object is used to represent a parameter in the graph. During
+    runtime this will represent a tensor value which can be used in various
+    calculations. The `Param` object can be set to a constant value (`value`);
+    `None` meaning the value is to be provided at runtime (`dynamic`);
+    `LiveParam` meaning the value will be computed internally in the simulator
+    during runtime (`live`); another `Param` object meaning it will take on that
+    value at runtime (`pointer`); or a function of other `Param` objects to be
+    computed at runtime (`function`). These options allow users to flexibly set
+    the behavior of the simulator.
+
+    Examples
+    --------
+    ``` python
+    p1 = Param("test", (1.0, 2.0)) # constant value, length 2 vector
+    p2 = Param("test", None, (2,2)) # dynamic 2x2 matrix value
+    p3 = Param("test", LiveParam) # live updating value
+    p4 = Param("test", p1) # pointer to another parameter
+    p5 = Param("test", lambda p: p.children["other"].value * 2) # function of another parameter
+    p5.link("other", p2) # link the other parameter needed for the function
+    ```
+
+    Parameters
+    ----------
+    name: (str)
+        The name of the parameter.
+    value: (Optional[Union[Tensor, float, int]], optional)
+        The value of the parameter. Defaults to None meaning dynamic.
+    shape: (Optional[tuple[int, ...]], optional)
+        The shape of the parameter. Defaults to () meaning scalar.
+    """
+
+    def __init__(
+        self,
+        name,
+        value: Optional[Union[Tensor, float, int]] = None,
+        shape: Optional[tuple[int, ...]] = (),
+    ):
+        super().__init__(name=name)
+        if value is None:
+            if shape is None:
+                raise ValueError("Either value or shape must be provided")
+            if not isinstance(shape, tuple):
+                raise ValueError("Shape must be a tuple")
+            self.shape = shape
+        elif not isinstance(value, (Param, Callable, LiveParamBase)):
+            value = torch.as_tensor(value)
+            assert (
+                shape == () or shape == value.shape
+            ), f"Shape {shape} does not match value shape {value.shape}"
+        self.value = value
+
+    @property
+    def dynamic(self) -> bool:
+        return self._type == "dynamic"
+
+    @property
+    def live(self) -> bool:
+        return self._type == "live"
+
+    @property
+    def shape(self) -> tuple:
+        return self._shape
+
+    @shape.setter
+    def shape(self, shape):
+        if self._type in ["pointer", "function"]:
+            raise RuntimeError("Cannot set shape of parameter with type 'pointer' or 'function'")
+        self._shape = shape
+
+    @property
+    def value(self) -> Union[Tensor, None]:
+        if self._type == "pointer":
+            return self._value.value
+        if self._type == "function":
+            return self._value(self)
+        return self._value
+
+    @value.setter
+    def value(self, value):
+        # While active, update silently
+        if self.active:
+            if self.dynamic or self.live:
+                self._value = value
+                return
+            raise RuntimeError(f"Cannot set value of non-live parameter {self.name} while active")
+
+        # unlink if pointer to avoid floating references
+        if self._type == "pointer":
+            self.unlink(self._value)
+
+        if value is None:
+            self._type = "dynamic"
+        elif isinstance(value, LiveParamBase):
+            self._type = "live"
+        elif isinstance(value, Param):
+            self._type = "pointer"
+            self.link(value.name, value)
+            self._shape = None
+        elif callable(value):
+            self._type = "function"
+            self._shape = None
+        else:
+            self._type = "value"
+            value = torch.as_tensor(value)
+            self.shape = value.shape
+
+        self._value = value
+        self.update_dynamic_params()
+
+    def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
+        """
+        Moves and/or casts the values of the parameter.
+
+        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.
+        """
+        super().to(device=device, dtype=dtype)
+        if self._type == "value":
+            self._value = self._value.to(device=device, dtype=dtype)

caskade/tests.py

@@ -0,0 +1,47 @@
+import torch
+
+from caskade import Module, Param, forward, LiveParam
+
+__all__ = ("test",)
+
+
+def _test_full_integration():
+
+    class TestSim(Module):
+        def __init__(self, a, b, c, c_shape, m1):
+            super().__init__("test_sim")
+            self.a = a
+            self.b = Param("b", b)
+            self.c = Param("c", c, c_shape)
+            self.m1 = m1
+
+        @forward
+        def testfun(self, x, b=None):
+            self.c.value = b + x
+            y = self.m1()
+            return x + self.a + b + y
+
+    class TestSubSim(Module):
+        def __init__(self, d, e, f):
+            super().__init__("test_sub_sim")
+            self.d = Param("d", d)
+            self.e = Param("e", e)
+            self.f = Param("f", f)
+
+        @forward
+        def __call__(self, d=None, e=None, f=None):
+            return d + e + f
+
+    sub1 = TestSubSim(d=1.0, e=lambda s: s.children["flink"].value, f=None)
+    sub1.e.link("flink", sub1.f)
+    main1 = TestSim(a=2.0, b=None, c=LiveParam, c_shape=(), m1=sub1)
+    sub1.f = main1.c
+
+    b_value = torch.tensor(3.0)
+    res = main1.testfun(1.0, params=[b_value])
+    assert res.item() == 15.0
+
+
+def test():
+    _test_full_integration()
+    print("Success!")

caskade-0.0.1.dist-info/METADATA

@@ -0,0 +1,99 @@
+Metadata-Version: 2.3
+Name: caskade
+Version: 0.0.1
+Summary: Package for building scientific simulators, with dynamic arguments arranged in a directed acyclic graph.
+Project-URL: Homepage, https://github.com/ConnorStoneAstro/caskade
+Project-URL: Documentation, https://github.com/ConnorStoneAstro/caskade
+Project-URL: Repository, https://github.com/ConnorStoneAstro/caskade
+Project-URL: Issues, https://github.com/ConnorStoneAstro/caskade/issues
+Author-email: Connor Stone <connorstone628@gmail.com>, Alexandre Adam <alexandre.adam@mila.quebec>
+License: MIT License
+        
+        Copyright (c) 2024 Connor Stone, PhD
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: DAG,caskade,differentiable programming,pytorch,scientific python
+Classifier: Development Status :: 1 - Planning
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.9
+Requires-Dist: torch
+Provides-Extra: dev
+Requires-Dist: pre-commit<4,>=3.6; extra == 'dev'
+Requires-Dist: pytest-cov<5,>=4.1; extra == 'dev'
+Requires-Dist: pytest-mock<4,>=3.12; extra == 'dev'
+Requires-Dist: pytest<9,>=8.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# caskade
+
+Build scientific simulators, treating them as a directed acyclic graph. Handles
+argument passing for complex nested simulators.
+
+## Install
+
+``` bash
+pip install caskade
+```
+
+## Usage
+
+Make a `Module` object which may have some `Param`s. Define a `forward` method
+using the decorator.
+
+``` python
+from caskade import Module, Param, forward
+
+class MySim(Module):
+    def __init__(self, a, b=None):
+        super().__init__()
+        self.a = a
+        self.b = Param("b", b)
+
+    @forward
+    def myfun(self, x, b=None):
+        return x + self.a + b
+```
+
+We may now create instances of the simulator and pass the dynamic parameters.
+
+``` python
+import torch
+
+sim = MySim(1.0)
+
+params = [torch.tensor(2.0)]
+
+print(sim.myfun(3.0, params=params))
+```
+
+Which will print `6` by automatically filling `b` with the value from `params`.
+
+### Why do this?
+
+The above example is not very impressive, the real power comes from the fact
+that `Module` objects can be nested arbitrarily making a much more complicated
+analysis graph. Further, the `Param` objects can be linked or have other complex
+relationships. All of the complexity of the nested structure and argument
+passing is abstracted away so that at the top one need only pass a list of
+tensors for each parameter, a single large 1d tensor, or a dictionary with the
+same structure as the graph.

caskade-0.0.1.dist-info/RECORD

@@ -0,0 +1,12 @@
+caskade/__init__.py,sha256=17BGDM0TQqKWV_7biXfFwgVCs2uy_tLwuMSzrGvz6K8,367
+caskade/_version.py,sha256=pMnmqZnpVmaqR5nqHztNWzbbtb1oy5bPN_v7uhOH8K8,411
+caskade/base.py,sha256=LrzdRvTXbPO161Xkr2JzczmhZb90NX9KqM_eZk3VdXA,4383
+caskade/context.py,sha256=cxAVthi1Btm8JIf410AfPfVd4E7VjQy3bi3wNQ_i1io,528
+caskade/decorators.py,sha256=LRcjKZVWYicMa2FSejz0xjqO5Iu2A17EYKcbgmarALI,1324
+caskade/module.py,sha256=kXvkfylqRer9LTNHPMIrZCC7Fy6rLE27RxUn2utMR9E,3566
+caskade/param.py,sha256=_xnRzyu_CxVjxQrvNYOifOWBHUF2Dt140zAyzv5_u00,4908
+caskade/tests.py,sha256=OT1lzOb1fSzLJBoZqZWBdNDYzPanJJgFoA4BjXcsad4,1212
+caskade-0.0.1.dist-info/METADATA,sha256=DlxkQy_G4SDAL9jn3sfllHZiTd29eIfsIaAPo80vrVY,3782
+caskade-0.0.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+caskade-0.0.1.dist-info/licenses/LICENSE,sha256=LSU9RKlizfhpJiFNnlbEIxgtbq0CRfdFGIP-f78Tp30,1074
+caskade-0.0.1.dist-info/RECORD,,

caskade-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.25.0
+Root-Is-Purelib: true
+Tag: py3-none-any

caskade-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Connor Stone, PhD
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.