nd2py 3.0.0__tar.gz

nd2py-3.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 YuMeow
+
+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.

nd2py-3.0.0/PKG-INFO

@@ -0,0 +1,35 @@
+Metadata-Version: 2.4
+Name: nd2py
+Version: 3.0.0
+Summary: nd2py (Neural Discovery of Network Dynamics) symbolic regression
+Author-email: YuMeow <yuzh19@tsinghua.org.cn>
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tqdm>=4.66
+Requires-Dist: pytest>=8.3
+Requires-Dist: numpy>=1.20
+Requires-Dist: scipy>=1.14
+Requires-Dist: pandas>=2.2
+Requires-Dist: scikit-learn>=1.5
+Requires-Dist: setproctitle
+Requires-Dist: matplotlib
+Requires-Dist: seaborn
+Requires-Dist: pyyaml
+Requires-Dist: rich
+Requires-Dist: dotenv
+Requires-Dist: requests
+Requires-Dist: pyperclip
+Provides-Extra: nn
+Requires-Dist: torch>=1.12; extra == "nn"
+Requires-Dist: torch_geometric>=2.0; extra == "nn"
+Provides-Extra: all
+Requires-Dist: nd2py[nn]; extra == "all"
+Dynamic: license-file
+
+
+
+```shell
+conda create --prefix ./venv python=3.12
+
+```

nd2py-3.0.0/README.md

@@ -0,0 +1,6 @@
+
+
+```shell
+conda create --prefix ./venv python=3.12
+
+```

nd2py-3.0.0/nd2py/__init__.py

@@ -0,0 +1,19 @@
+from .core import *
+from . import dataset, generator, search, utils
+
+# 定义模块级别的 __getattr__ 拦截
+def __getattr__(name):
+    if name == '__all__':
+        # Triggered when 'from nd2py import *' is used
+        import warnings
+        warnings.warn(
+            "Detected 'from nd2py import *'. WARNING: This will shadow standard Python built-in functions (such as 'sum').\n"
+            "It is strongly recommended to use explicit imports, e.g., 'from nd2py import sum' or 'import nd2py as nd'.",
+            category=UserWarning,
+            stacklevel=2
+        )
+        # 动态获取当前模块（__init__.py）中所有不以 '_' 开头的全局变量和模块
+        return [n for n in globals() if not n.startswith('_')]
+
+    # 对于其他不存在的属性，保持默认的报错行为
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

nd2py-3.0.0/nd2py/core/__init__.py

@@ -0,0 +1,5 @@
+from .calc import *
+from .symbols import *
+from .context import *
+from .converter import *
+from .transform import *

nd2py-3.0.0/nd2py/core/base_visitor.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2024-present, Yumeow. Licensed under the MIT License.
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING: # 避免循环引用，仅用于类型检查
+    from .symbols import Symbol
+
+
+def yield_nothing():
+    """A generator that yields nothing, used as a placeholder for methods that do not yield."""
+    if False:
+        yield
+
+
+class Visitor(ABC):
+    def __call__(self, node: Symbol, *args, **kwargs):
+        """
+        1) call Visitor.visit_<ClassName> based on type(node)
+        2) call Visitor.generic_visit if no Visitor.visit_<ClassName> defined
+        """
+        stack = [("start", node, args, kwargs, None)]
+        result = None
+        while stack:
+            state, node, args, kwargs, gen = stack.pop()
+            if state == "start":
+                method = getattr(
+                    self, "visit_" + type(node).__name__, self.generic_visit
+                )
+                gen = method(node, *args, **kwargs)
+                if not hasattr(gen, "__next__"):
+                    raise TypeError(
+                        f"Expected a generator but got {type(gen).__name__}, please add `yield from yield_nothing()` in {type(self).__name__}.{method.__name__}."
+                    )
+                try:
+                    child, args, kwargs = next(gen)
+                    stack.append(("resume", node, None, None, gen))
+                    stack.append(("start", child, args, kwargs, None))
+                except StopIteration as e:
+                    result = e.value
+            elif state == "resume":
+                try:
+                    child, args, kwargs = gen.send(result)
+                    stack.append(("resume", node, None, None, gen))
+                    stack.append(("start", child, args, kwargs, None))
+                except StopIteration as e:
+                    result = e.value
+            else:
+                raise ValueError(f"Unknown state: {state}")
+        return result
+
+    @abstractmethod
+    def generic_visit(self, node: Symbol, *args, **kwargs):
+        msg = f"generic_visit not implemented for {type(self).__name__}"
+        raise NotImplementedError(msg)
+
+    _SYMBOL_CACHE = {}
+
+    @classmethod
+    def _get_symbol(cls, name):
+        if name not in cls._SYMBOL_CACHE:
+            from . import symbols
+            cls._SYMBOL_CACHE[name] = getattr(symbols, name)
+        return cls._SYMBOL_CACHE[name]

nd2py-3.0.0/nd2py/core/basic/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2024-present, Yumeow. Licensed under the MIT License.
+"""
+The implementation of basic functionality of Symbol as an object, including:
+- GetCopy: Deep copy a Symbol tree.
+- GetLength: Count the number of nodes in a Symbol tree.
+"""
+
+from .get_copy import GetCopy
+from .get_length import GetLength

nd2py-3.0.0/nd2py/core/basic/get_copy.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024-present, Yumeow. Licensed under the MIT License.
+from copy import deepcopy
+from ..base_visitor import Visitor, yield_nothing
+from ..context.copy_value import get_copy_value
+
+
+class GetCopy(Visitor):
+    def __call__(self, node, *args, **kwargs):
+        """Create a copy of the node."""
+        return super().__call__(node, *args, **kwargs)
+
+    def generic_visit(self, node, *args, **kwargs):
+        yield from yield_nothing()
+        children = []
+        for child in node.operands:
+            child = yield child, args, kwargs
+            children.append(child)
+        return node.__class__(*children, nettype=node._assigned_nettypes)
+
+    def visit_Number(self, node, *args, **kwargs):
+        yield from yield_nothing()
+        if get_copy_value():
+            return node.__class__(
+                deepcopy(node.value), nettype=node.nettype, fitable=node.fitable
+            )
+        else:
+            return node.__class__(
+                node.value, nettype=node._assigned_nettypes, fitable=node.fitable
+            )
+
+    def visit_Variable(self, node, *args, **kwargs):
+        yield from yield_nothing()
+        return node.__class__(node.name, nettype=node._assigned_nettypes)
+
+    def visit_Empty(self, node, *args, **kwargs):
+        yield from yield_nothing()
+        return node.__class__(nettype=node._assigned_nettypes)

nd2py-3.0.0/nd2py/core/basic/get_length.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2024-present, Yumeow. Licensed under the MIT License.
+from ..base_visitor import Visitor, yield_nothing
+
+
+class GetLength(Visitor):
+    def __call__(self, node, *args, **kwargs):
+        """Count the number of nodes in the tree."""
+        return super().__call__(node, *args, **kwargs)
+
+    def generic_visit(self, node, *args, **kwargs):
+        yield from yield_nothing()
+        children = []
+        for child in node.operands:
+            child = yield child, args, kwargs
+            children.append(child)
+        return 1 + sum(children)

nd2py-3.0.0/nd2py/core/calc/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2024-present, Yumeow. Licensed under the MIT License.
+"""
+The implementation of calculation functionality of Symbol, including:
+- NumpyCalc: Calculate the value of a Symbol tree using numpy.
+- TorchCalc: Calculate the value of a Symbol tree using torch.
+"""
+
+from .numpy_calc import NumpyCalc
+from .torch_calc import TorchCalc

nd2py-3.0.0/nd2py/core/calc/numpy_calc.py

@@ -0,0 +1,391 @@
+# Copyright (c) 2024-present, Yumeow. Licensed under the MIT License.
+from __future__ import annotations
+import numbers
+import warnings
+import functools
+import traceback
+import numpy as np
+from typing import List, Tuple, TYPE_CHECKING
+from ..base_visitor import Visitor, yield_nothing
+if TYPE_CHECKING:
+    from ..symbols import *
+
+
+# Decorator to unpack operands for operations
+# This allows us to handle operations with multiple operands in a clean way
+# We can also use this decorator to suppress numpy errors
+def unpack_operands():
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(self, node, *args, **kwargs):
+            # Calculate the values of the operands
+            yield from yield_nothing()
+            X = []
+            for op in node.operands:
+                x = yield (op, args, kwargs)
+                X.append(x)
+            # Use the defined 'visit_<Operation>' as 'func' to process the operands
+            with np.errstate(divide="ignore", invalid="ignore", over="ignore"):
+                return func(self, node, *X, *args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+class NumpyCalc(Visitor):
+    def __call__(
+        self,
+        node: Symbol,
+        vars: dict = {},
+        edge_list: Tuple[List[int], List[int]] = None,
+        num_nodes: int = None,
+        use_eps: float = 0.0,
+    ):
+        """
+        Args:
+        - vars: a dictionary of variable names and their values
+        - edge_list: edges (i,j) in the graph
+            i and j are the indices of the nodes (starting from 0)
+        - num_nodes: the number of nodes in the graph
+            if not provided, it will be inferred from edge_list
+        - use_eps: a small value to avoid division by zero
+        """
+        if num_nodes is None and edge_list is not None:
+            nodes = np.unique(np.array(edge_list).reshape(-1))
+            num_nodes = max(nodes) + 1
+
+        return super().__call__(
+            node,
+            vars=vars,
+            edge_list=edge_list,
+            num_nodes=num_nodes,
+            use_eps=use_eps,
+        )
+
+    def generic_visit(self, node: Symbol, *args, **kwargs):
+        raise NotImplementedError(
+            f"{type(self).__name__}.visit_{type(node).__name__} not implemented"
+        )
+
+    def visit_Empty(self, node: Empty, *args, **kwargs):
+        raise ValueError(
+            f"Incomplete expression with Empty node is not allowed to evaluate: {node}"
+        )
+
+    def visit_Number(self, node: Number, *args, **kwargs):
+        yield from yield_nothing()
+        return np.asarray(node.value)
+
+    def visit_Variable(self, node: Variable, *args, **kwargs):
+        yield from yield_nothing()
+        return np.asarray(kwargs["vars"][node.name])
+
+    @unpack_operands()
+    def visit_Add(self, node: Add, x1, x2, *args, **kwargs):
+        return x1 + x2
+
+    @unpack_operands()
+    def visit_Sub(self, node: Sub, x1, x2, *args, **kwargs):
+        return x1 - x2
+
+    @unpack_operands()
+    def visit_Mul(self, node: Mul, x1, x2, *args, **kwargs):
+        return x1 * x2
+
+    @unpack_operands()
+    def visit_Div(self, node: Div, x1, x2, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return x1 / (x2 + eps * (x2 == 0))
+
+    @unpack_operands()
+    def visit_Pow(self, node: Pow, x1, x2, *args, **kwargs):
+        return x1**x2
+
+    @unpack_operands()
+    def visit_Max(self, node: Max, x1, x2, *args, **kwargs):
+        return np.maximum(x1, x2)
+
+    @unpack_operands()
+    def visit_Min(self, node: Min, x1, x2, *args, **kwargs):
+        return np.minimum(x1, x2)
+    
+    @unpack_operands()
+    def visit_Identity(self, node: Identity, x, *args, **kwargs):
+        return x
+
+    @unpack_operands()
+    def visit_Sin(self, node: Sin, x, *args, **kwargs):
+        return np.sin(x)
+
+    @unpack_operands()
+    def visit_Cos(self, node: Cos, x, *args, **kwargs):
+        return np.cos(x)
+
+    @unpack_operands()
+    def visit_Tan(self, node: Tan, x, *args, **kwargs):
+        return np.tan(x)
+
+    @unpack_operands()
+    def visit_Sec(self, node: Sec, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (np.cos(x) + eps * (np.cos(x) == 0))
+
+    @unpack_operands()
+    def visit_Csc(self, node: Csc, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (np.sin(x) + eps * (np.sin(x) == 0))
+
+    @unpack_operands()
+    def visit_Cot(self, node: Cot, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (np.tan(x) + eps * (np.tan(x) == 0))
+
+    @unpack_operands()
+    def visit_Log(self, node: Log, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return np.log(x + eps * (x == 0))
+
+    @unpack_operands()
+    def visit_LogAbs(self, node: LogAbs, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return np.log(np.abs(x + eps * (x == 0)))
+
+    @unpack_operands()
+    def visit_Exp(self, node: Exp, x, *args, **kwargs):
+        return np.exp(x)
+
+    @unpack_operands()
+    def visit_Abs(self, node: Abs, x, *args, **kwargs):
+        return np.abs(x)
+
+    @unpack_operands()
+    def visit_Neg(self, node: Neg, x, *args, **kwargs):
+        return -x
+
+    @unpack_operands()
+    def visit_Inv(self, node: Inv, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (x + eps * (x == 0))
+
+    @unpack_operands()
+    def visit_Sqrt(self, node: Sqrt, x, *args, **kwargs):
+        return np.sqrt(x)
+
+    @unpack_operands()
+    def visit_SqrtAbs(self, node: SqrtAbs, x, *args, **kwargs):
+        return np.sqrt(np.abs(x))
+
+    @unpack_operands()
+    def visit_Pow2(self, node: Pow2, x, *args, **kwargs):
+        return x**2
+
+    @unpack_operands()
+    def visit_Pow3(self, node: Pow3, x, *args, **kwargs):
+        return x**3
+
+    @unpack_operands()
+    def visit_Arcsin(self, node: Arcsin, x, *args, **kwargs):
+        return np.arcsin(x)
+
+    @unpack_operands()
+    def visit_Arccos(self, node: Arccos, x, *args, **kwargs):
+        return np.arccos(x)
+
+    @unpack_operands()
+    def visit_Arctan(self, node: Arctan, x, *args, **kwargs):
+        return np.arctan(x)
+
+    @unpack_operands()
+    def visit_Sinh(self, node: Sinh, x, *args, **kwargs):
+        return np.sinh(x)
+
+    @unpack_operands()
+    def visit_Cosh(self, node: Cosh, x, *args, **kwargs):
+        return np.cosh(x)
+
+    @unpack_operands()
+    def visit_Tanh(self, node: Tanh, x, *args, **kwargs):
+        return np.tanh(x)
+
+    @unpack_operands()
+    def visit_Sech(self, node: Sech, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (np.cosh(x) + eps * (np.cosh(x) == 0))
+
+    @unpack_operands()
+    def visit_Csch(self, node: Csch, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (np.sinh(x) + eps * (np.sinh(x) == 0))
+
+    @unpack_operands()
+    def visit_Coth(self, node: Coth, x, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (np.tanh(x) + eps * (np.tanh(x) == 0))
+
+    @unpack_operands()
+    def visit_Sigmoid(self, node: Sigmoid, x, *args, **kwargs):
+        return 1 / (1 + np.exp(-x))
+
+    @unpack_operands()
+    def visit_Regular(self, node: Regular, x1, x2, *args, **kwargs):
+        eps = kwargs.get("use_eps")
+        return 1 / (1 + (np.abs(x1) + eps * (x1 == 0)) ** (-x2))
+
+    @unpack_operands()
+    def visit_Sour(self, node: Sour, x, *args, **kwargs):
+        """(*, n_nodes or 1) -> (*, n_edges or 1)"""
+        edge_list = kwargs.get("edge_list", ([], []))
+
+        if isinstance(x, numbers.Number) or x.size == 1:
+            return x  # (1,) -> (1,)
+        elif node.operands[0].nettype == "scalar" or x.shape[-1] == 1:
+            if x.shape[-1] != 1:
+                x = x[..., np.newaxis]
+            return x  # (*, 1) -> (*, 1)
+        else:
+            return x[..., edge_list[0]]  # (*, V) -> (*, E)
+
+    @unpack_operands()
+    def visit_Targ(self, node: Targ, x, *args, **kwargs):
+        """(*, n_nodes or 1) -> (*, n_edges or 1)"""
+        edge_list = kwargs.get("edge_list", ([], []))
+
+        if isinstance(x, numbers.Number) or x.size == 1:
+            return x  # (1,) -> (1,)
+        elif node.operands[0].nettype == "scalar" or x.shape[-1] == 1:
+            if x.shape[-1] != 1:
+                x = x[..., np.newaxis]
+            return x  # (*, 1) -> (*, 1)
+        else:
+            return x[..., edge_list[1]]  # (*, V) -> (*, E)
+
+    @unpack_operands()
+    def visit_Aggr(self, node: Aggr, x, *args, **kwargs):
+        """(*, n_edges or 1) -> (*, n_nodes)"""
+        edge_list = kwargs.get("edge_list", ([], []))
+        num_nodes = kwargs.get("num_nodes")
+
+        if isinstance(x, numbers.Number) or x.size == 1:
+            y = np.zeros((num_nodes,))
+            np.add.at(y, edge_list[1], float(x))
+            return y
+        elif node.operands[0].nettype == "scalar" or x.shape[-1] == 1:
+            if x.shape[-1] != 1:
+                x = x[..., np.newaxis]
+            y = np.zeros((num_nodes,))
+            np.add.at(y, edge_list[1], 1)
+            y = y * x
+            return y
+        else:
+            y = np.zeros((*x.shape[:-1], num_nodes))
+            for k, j in enumerate(edge_list[1]):
+                y[..., j] += x[..., k]
+            return y
+
+    @unpack_operands()
+    def visit_Rgga(self, node: Rgga, x, *args, **kwargs):
+        """(*, n_edges or 1) -> (*, n_nodes)"""
+        edge_list = kwargs.get("edge_list", ([], []))
+        num_nodes = kwargs.get("num_nodes")
+
+        if isinstance(x, numbers.Number) or x.size == 1:
+            y = np.zeros((num_nodes,))
+            np.add.at(y, edge_list[1], x)
+            return y
+        elif node.operands[0].nettype == "scalar" or x.shape[-1] == 1:
+            if x.shape[-1] != 1:
+                x = x[..., np.newaxis]
+            y = np.zeros((num_nodes,))
+            np.add.at(y, edge_list[0], 1)
+            y = y * x
+            return y
+        else:
+            y = np.zeros((*x.shape[:-1], num_nodes))
+            for k, i in enumerate(edge_list[0]):
+                y[..., i] += x[..., k]
+            return y
+
+    @unpack_operands()
+    def visit_Readout(self, node: Readout, x, *args, **kwargs):
+        """(*, n_nodes or n_edges or 1) -> (*, 1)"""
+        return np.sum(x, axis=-1, keepdims=True)
+
+
+"""
+# 比较 aggr 不同实现方式的性能
+
+import numpy as np
+
+T = 100
+V = 10
+E = 10
+x = np.random.rand(T, E)
+index = np.random.randint(0, V, size=(E,))
+
+def aggr0(x, index, V):
+    y = np.zeros((*x.shape[:-1], V), dtype=x.dtype)
+    for k, j in enumerate(index):
+        y[..., j] += x[..., k]
+    return y
+    
+
+def aggr1(x, index, V):
+    x = np.asarray(x)
+    orig_shape = x.shape
+    x_flat = x.reshape(-1, x.shape[-1])  # shape: (B, N), B = product of leading dims
+
+    T = x_flat.shape[0]
+    y = np.zeros((T, V), dtype=x.dtype)
+    for i in range(T):
+        np.add.at(y[i, :], index, x_flat[i, :])
+    y = y.reshape(*orig_shape[:-1], V)
+    return y
+
+def aggr2(x, index, V):
+    x = np.asarray(x)
+    orig_shape = x.shape
+    x_flat = x.reshape(-1, x.shape[-1])  # shape: (B, N), B = product of leading dims
+
+    T = x_flat.shape[0]
+    y = np.zeros((T, V), dtype=x.dtype)
+    for i in range(T):
+        y[i, :] = np.bincount(index, weights=x[i, :], minlength=V)
+    y = y.reshape(*orig_shape[:-1], V)
+    return y
+
+def aggr3(x, index, V):
+    x = np.asarray(x)
+    orig_shape = x.shape
+    x_flat = x.reshape(-1, x.shape[-1])  # shape: (B, E)
+
+    T = x_flat.shape[0]
+    # 扩展 index 为 shape (T, E) 以对齐每个 batch
+    index_broadcast = np.broadcast_to(index, (T, E))
+
+    # 准备目标数组
+    y_flat = np.zeros((T, V), dtype=x.dtype)
+
+    # 扁平化批次索引：行号 (0,0,...,1,1,1,...T-1) 与 index 构成 2D 索引
+    row_idx = np.repeat(np.arange(T), E)
+    col_idx = index_broadcast.ravel()
+    values = x_flat.ravel()
+
+    # 聚合加和
+    np.add.at(y_flat, (row_idx, col_idx), values)
+
+    # reshape 回原来的批次形状
+    return y_flat.reshape(*orig_shape[:-1], V)
+
+
+y0 = aggr0(x, index, V)
+
+y1 = aggr1(x, index, V)
+assert (y0 == y1).all()
+
+y2 = aggr2(x, index, V)
+assert (y0 == y2).all()
+
+y3 = aggr3(x, index, V)
+assert (y0 == y3).all()
+"""