From 438a484c71071d0cf2dda88500c52ccf35202e99 Mon Sep 17 00:00:00 2001
From: Connor Ferster <connorferster@gmail.com>
Date: Mon, 14 Jul 2025 21:42:54 -0700
Subject: [PATCH 1/2] feat: add single-level aggregation

---
 src/pynite_tools/__init__.py |  8 ++++
 src/pynite_tools/envelope.py | 79 ++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 src/pynite_tools/envelope.py

diff --git a/src/pynite_tools/__init__.py b/src/pynite_tools/__init__.py
index 5e51f88..1bb042e 100644
--- a/src/pynite_tools/__init__.py
+++ b/src/pynite_tools/__init__.py
@@ -40,4 +40,12 @@
 from .visualize import (
     plot_model,
     Renderer
+)
+
+from .envelope import (
+    envelope_tree,
+    absmax,
+    absmin,
+    max,
+    min
 )
\ No newline at end of file
diff --git a/src/pynite_tools/envelope.py b/src/pynite_tools/envelope.py
new file mode 100644
index 0000000..1a24bf4
--- /dev/null
+++ b/src/pynite_tools/envelope.py
@@ -0,0 +1,79 @@
+from copy import copy
+from typing import Hashable
+
+PYMAX = max
+PYMIN = min
+
+def envelope_tree(tree: dict | list, levels: list[Hashable | None], leaf: Hashable | None, agg_func: callable) -> dict:
+    """
+    Envelopes the tree at the leaf node with 'agg_func'.
+    """
+    env_acc = {}
+    # If we are at the last branch...
+    if not levels:
+        if isinstance(tree, list):
+            tree = {idx: leaf for idx, leaf in enumerate(tree)}
+        leaf_acc = {}
+        for k, leaves in tree.items():
+            if leaf is not None:
+                leaf_acc.update({k: leaves[leaf]})
+            else:
+                leaf_acc = tree
+        # ...create a dict of the enveloped value and the key
+        # that it belongs to and return it.
+        env_values = list(leaf_acc.values())
+        env_value = agg_func(env_values)
+        env_keys = list(leaf_acc.keys())
+        try:
+            env_key = env_keys[env_values.index(env_value)]
+        except ValueError: # The value was transformed, likely due to abs()
+            env_key = env_keys[env_values.index(-1 * env_value)]
+        env_acc.update({"key": env_key, "value": env_value})
+        return env_acc
+    else:
+        # Otherwise, pop the next level and dive into the tree on that branch
+        level = levels[0]
+        if level is not None:
+            env_acc.update({level: envelope_tree(tree[level], levels[1:], leaf, agg_func)})
+            return env_acc
+        else:
+            # If None, then walk all branches of this node of the tree
+            if isinstance(tree, list):
+                tree = {idx: leaf for idx, leaf in enumerate(tree)}
+            for k, v in tree.items():
+                env_acc.update({k: envelope_tree(v, levels[1:], leaf, agg_func)})
+            return env_acc
+
+    
+def absmax(x: list[float | int]) -> float | int:
+    """
+    Returns the absolute maximum value in 'x'.
+    """
+    abs_acc = [abs(y) for y in x]
+    return max(abs_acc)
+
+
+def absmin(x: list[float | int]) -> float | int:
+    """
+    Returns the absolute minimum value in 'x'.
+    """
+    abs_acc = [abs(y) for y in x]
+    return min(abs_acc)
+
+
+def max(x: list[float | int | None]) -> float | int | None:
+    """
+    Returns the max value in 'x' while ignoring "None".
+    If all values in 'x' are None, then None is returned.
+    """
+    return PYMAX([y for y in x if y is not None])
+
+
+def min(x: list[float | int | None]) -> float | int | None:
+    """
+    Returns the max value in 'x' while ignoring "None".
+    If all values in 'x' are None, then None is returned.
+    """
+    return PYMIN([y for y in x if y is not None])
+
+

From cf4ab8c2ba5d56c510e6c0925030426336d8ccdb Mon Sep 17 00:00:00 2001
From: Connor Ferster <connorferster@gmail.com>
Date: Mon, 14 Jul 2025 22:04:33 -0700
Subject: [PATCH 2/2] feat: can now return dicts that can be double enveloped

---
 src/pynite_tools/envelope.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/pynite_tools/envelope.py b/src/pynite_tools/envelope.py
index 1a24bf4..ff33535 100644
--- a/src/pynite_tools/envelope.py
+++ b/src/pynite_tools/envelope.py
@@ -4,7 +4,7 @@
 PYMAX = max
 PYMIN = min
 
-def envelope_tree(tree: dict | list, levels: list[Hashable | None], leaf: Hashable | None, agg_func: callable) -> dict:
+def envelope_tree(tree: dict | list, levels: list[Hashable | None], leaf: Hashable | None, agg_func: callable, with_trace: bool = True) -> dict:
     """
     Envelopes the tree at the leaf node with 'agg_func'.
     """
@@ -29,19 +29,22 @@ def envelope_tree(tree: dict | list, levels: list[Hashable | None], leaf: Hashab
         except ValueError: # The value was transformed, likely due to abs()
             env_key = env_keys[env_values.index(-1 * env_value)]
         env_acc.update({"key": env_key, "value": env_value})
-        return env_acc
+        if with_trace:
+            return env_acc
+        else:
+            return env_value
     else:
         # Otherwise, pop the next level and dive into the tree on that branch
         level = levels[0]
         if level is not None:
-            env_acc.update({level: envelope_tree(tree[level], levels[1:], leaf, agg_func)})
+            env_acc.update({level: envelope_tree(tree[level], levels[1:], leaf, agg_func, with_trace)})
             return env_acc
         else:
             # If None, then walk all branches of this node of the tree
             if isinstance(tree, list):
                 tree = {idx: leaf for idx, leaf in enumerate(tree)}
             for k, v in tree.items():
-                env_acc.update({k: envelope_tree(v, levels[1:], leaf, agg_func)})
+                env_acc.update({k: envelope_tree(v, levels[1:], leaf, agg_func, with_trace)})
             return env_acc